From 7bcc157f282390c5461be0116bf0847f87749ab4 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Thu, 18 Apr 2024 16:16:56 +0200 Subject: [PATCH 01/11] MAINT: capitalize QRules `REACTION` variables --- docs/comparison/d2kkk.ipynb | 10 +++++----- docs/comparison/jpsi2phipipi.ipynb | 10 +++++----- docs/comparison/jpsi2pipipi.ipynb | 10 +++++----- 3 files changed, 15 insertions(+), 15 deletions(-) diff --git a/docs/comparison/d2kkk.ipynb b/docs/comparison/d2kkk.ipynb index 79fc5055..42cafc05 100644 --- a/docs/comparison/d2kkk.ipynb +++ b/docs/comparison/d2kkk.ipynb @@ -154,14 +154,14 @@ }, "outputs": [], "source": [ - "reaction = qrules.generate_transitions(\n", + "REACTION = qrules.generate_transitions(\n", " initial_state=INITIAL_STATE.name,\n", " final_state=[p.name for p in FINAL_STATE],\n", " allowed_intermediate_particles=[\"a(0)\", \"f(0)(980)\", \"pi(1)\"],\n", " mass_conservation_factor=0.2,\n", " formalism=\"helicity\",\n", ")\n", - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", + "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, @@ -179,7 +179,7 @@ "outputs": [], "source": [ "qrules_resonances = sorted(\n", - " reaction.get_intermediate_particles(),\n", + " REACTION.get_intermediate_particles(),\n", " key=lambda p: (p.charge, p.mass, p.name),\n", ")\n", "resonance_names = [p.name for p in qrules_resonances]\n", @@ -415,7 +415,7 @@ }, "outputs": [], "source": [ - "model_builder = ampform.get_builder(reaction)\n", + "model_builder = ampform.get_builder(REACTION)\n", "model_builder.use_helicity_couplings = False\n", "ampform_model = model_builder.formulate()\n", "ampform_model.intensity" @@ -1022,7 +1022,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.9.19" } }, "nbformat": 4, diff --git a/docs/comparison/jpsi2phipipi.ipynb b/docs/comparison/jpsi2phipipi.ipynb index 368767af..4348a821 100644 --- a/docs/comparison/jpsi2phipipi.ipynb +++ b/docs/comparison/jpsi2phipipi.ipynb @@ -154,14 +154,14 @@ }, "outputs": [], "source": [ - "reaction = qrules.generate_transitions(\n", + "REACTION = qrules.generate_transitions(\n", " initial_state=INITIAL_STATE.name,\n", " final_state=[p.name for p in FINAL_STATE],\n", " allowed_intermediate_particles=[\"a(0)(1450\", \"rho(1450)\"],\n", " mass_conservation_factor=0,\n", " formalism=\"helicity\",\n", ")\n", - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", + "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, @@ -179,7 +179,7 @@ "outputs": [], "source": [ "qrules_resonances = sorted(\n", - " reaction.get_intermediate_particles(),\n", + " REACTION.get_intermediate_particles(),\n", " key=lambda p: (p.charge, p.mass, p.name),\n", ")\n", "resonance_names = [p.name for p in qrules_resonances]\n", @@ -398,7 +398,7 @@ }, "outputs": [], "source": [ - "model_builder = ampform.get_builder(reaction)\n", + "model_builder = ampform.get_builder(REACTION)\n", "model_builder.use_helicity_couplings = False\n", "ampform_model = model_builder.formulate()\n", "ampform_model.intensity" @@ -1007,7 +1007,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.9.19" } }, "nbformat": 4, diff --git a/docs/comparison/jpsi2pipipi.ipynb b/docs/comparison/jpsi2pipipi.ipynb index caccab7b..b0c210e8 100644 --- a/docs/comparison/jpsi2pipipi.ipynb +++ b/docs/comparison/jpsi2pipipi.ipynb @@ -154,14 +154,14 @@ }, "outputs": [], "source": [ - "reaction = qrules.generate_transitions(\n", + "REACTION = qrules.generate_transitions(\n", " initial_state=INITIAL_STATE.name,\n", " final_state=[p.name for p in FINAL_STATE],\n", " allowed_intermediate_particles=[\"a(0)(980)\", \"rho(770)\"],\n", " mass_conservation_factor=0,\n", " formalism=\"helicity\",\n", ")\n", - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", + "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, @@ -179,7 +179,7 @@ "outputs": [], "source": [ "qrules_resonances = sorted(\n", - " reaction.get_intermediate_particles(),\n", + " REACTION.get_intermediate_particles(),\n", " key=lambda p: (p.charge, p.mass, p.name),\n", ")\n", "resonance_names = [p.name for p in qrules_resonances]\n", @@ -398,7 +398,7 @@ }, "outputs": [], "source": [ - "model_builder = ampform.get_builder(reaction)\n", + "model_builder = ampform.get_builder(REACTION)\n", "model_builder.use_helicity_couplings = False\n", "ampform_model = model_builder.formulate()\n", "ampform_model.intensity" @@ -1005,7 +1005,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.9.19" } }, "nbformat": 4, From 6cfa36cd6dcf57b71b9b16375db7408111a87789 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Thu, 18 Apr 2024 16:20:08 +0200 Subject: [PATCH 02/11] DOC: relabel QRules final state IDs to 1,2,3 --- docs/comparison/d2kkk.ipynb | 4 +++- docs/comparison/jpsi2phipipi.ipynb | 4 +++- docs/comparison/jpsi2pipipi.ipynb | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/docs/comparison/d2kkk.ipynb b/docs/comparison/d2kkk.ipynb index 42cafc05..1bed9082 100644 --- a/docs/comparison/d2kkk.ipynb +++ b/docs/comparison/d2kkk.ipynb @@ -45,6 +45,7 @@ "import matplotlib.pyplot as plt\n", "import qrules\n", "import sympy as sp\n", + "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n", "from ampform.sympy import perform_cached_doit\n", "from IPython.display import SVG, Latex, Markdown, clear_output, display\n", @@ -161,7 +162,8 @@ " mass_conservation_factor=0.2,\n", " formalism=\"helicity\",\n", ")\n", - "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", + "REACTION123 = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, diff --git a/docs/comparison/jpsi2phipipi.ipynb b/docs/comparison/jpsi2phipipi.ipynb index 4348a821..99f88c28 100644 --- a/docs/comparison/jpsi2phipipi.ipynb +++ b/docs/comparison/jpsi2phipipi.ipynb @@ -45,6 +45,7 @@ "import matplotlib.pyplot as plt\n", "import qrules\n", "import sympy as sp\n", + "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n", "from ampform.sympy import perform_cached_doit\n", "from IPython.display import SVG, Latex, Markdown, clear_output, display\n", @@ -161,7 +162,8 @@ " mass_conservation_factor=0,\n", " formalism=\"helicity\",\n", ")\n", - "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", + "REACTION123 = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, diff --git a/docs/comparison/jpsi2pipipi.ipynb b/docs/comparison/jpsi2pipipi.ipynb index b0c210e8..6c6ff454 100644 --- a/docs/comparison/jpsi2pipipi.ipynb +++ b/docs/comparison/jpsi2pipipi.ipynb @@ -45,6 +45,7 @@ "import matplotlib.pyplot as plt\n", "import qrules\n", "import sympy as sp\n", + "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n", "from ampform.sympy import perform_cached_doit\n", "from IPython.display import SVG, Latex, Markdown, clear_output, display\n", @@ -161,7 +162,8 @@ " mass_conservation_factor=0,\n", " formalism=\"helicity\",\n", ")\n", - "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", + "REACTION123 = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", "graphviz.Source(dot)" ] }, From e6ca46587a898fa90b1915ace94625fcd76b908d Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Thu, 18 Apr 2024 15:45:12 +0200 Subject: [PATCH 03/11] DOC: generate J/psi decay figure with QRules --- docs/jpsi2ksp.ipynb | 44 ++++++++++++++++++++++++++++++++------------ pyproject.toml | 3 ++- 2 files changed, 34 insertions(+), 13 deletions(-) diff --git a/docs/jpsi2ksp.ipynb b/docs/jpsi2ksp.ipynb index 18e76cc4..c5a67730 100644 --- a/docs/jpsi2ksp.ipynb +++ b/docs/jpsi2ksp.ipynb @@ -29,13 +29,16 @@ "import itertools\n", "import logging\n", "import os\n", + "import warnings\n", "from typing import TYPE_CHECKING, Any, Iterable\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.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.phasespace import compute_third_mandelstam\n", "from ampform.sympy import perform_cached_doit, unevaluated\n", "from IPython.display import Latex, Markdown\n", @@ -52,14 +55,15 @@ ")\n", "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n", "\n", - "if TYPE_CHECKING:\n", - " from tensorwaves.interface import DataSample, ParametrizedFunction\n", - "\n", "simplify_latex_rendering()\n", "logging.getLogger(\"absl\").setLevel(logging.ERROR) # mute JAX\n", + "warnings.simplefilter(\"ignore\")\n", + "\n", "NO_TQDM = \"EXECUTE_NB\" in os.environ\n", "if NO_TQDM:\n", - " logging.getLogger(\"ampform_dpd.io\").setLevel(logging.ERROR)" + " logging.getLogger(\"ampform_dpd.io\").setLevel(logging.ERROR)\n", + "if TYPE_CHECKING:\n", + " from tensorwaves.interface import DataSample, ParametrizedFunction" ] }, { @@ -73,13 +77,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "We follow [this example](https://qrules.readthedocs.io/en/0.9.7/usage.html#investigate-intermediate-resonances), which was generated with QRules, and leave out the $K$-resonances and the resonances that lie far outside of phase space:\n", - "\n", - "![](https://qrules.readthedocs.io/en/0.9.7/_images/usage_9_0.svg)\n", - "\n", - ":::{warning}\n", - "In the above figure, the final states are labeled `0`, `1`, `2`, but in the DPD formalism, the final states are labeled `1`, `2`, `3`.\n", - ":::" + "We follow [this example](https://qrules.readthedocs.io/en/0.10.1/usage.html#investigate-intermediate-resonances), which was generated with QRules, and leave out the $K$-resonances and the resonances that lie far outside of phase space." ] }, { @@ -94,6 +92,28 @@ ] }, "outputs": [], + "source": [ + "REACTION = qrules.generate_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"K0\", \"Sigma+\", \"p~\"],\n", + " allowed_interaction_types=\"strong\",\n", + " formalism=\"helicity\", # AmpForm-DPD already generates LS couplings\n", + " mass_conservation_factor=0.05,\n", + ")\n", + "REACTION = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", + "graphviz.Source(dot)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], "source": [ "PDG = qrules.load_pdg()\n", "PARTICLE_DB = {\n", @@ -960,7 +980,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.18" + "version": "3.9.19" } }, "nbformat": 4, diff --git a/pyproject.toml b/pyproject.toml index ff541f57..a6174d0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,9 +26,10 @@ classifiers = [ "Typing :: Typed", ] dependencies = [ - "ampform >=0.14.8", # Kibble and Kallen functions, perform_cached_doit, @unevaluated + "ampform >=0.15.0", # relabel_edge_ids "attrs >=20.1.0", # on_setattr and https://www.attrs.org/en/stable/api.html#next-gen "cloudpickle", + "qrules >=0.10.0", "sympy >=1.10", # module sympy.printing.numpy and array expressions with shape kwarg "tensorwaves[jax]", ] From ada29c57e9c8604443b50ccd2df3d3ce62e924df Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Thu, 18 Apr 2024 17:52:41 +0200 Subject: [PATCH 04/11] ENH: make `ThreeBodyDecay` sortable --- docs/jpsi2ksp.ipynb | 4 ++-- src/ampform_dpd/decay.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/docs/jpsi2ksp.ipynb b/docs/jpsi2ksp.ipynb index c5a67730..7d10d2f3 100644 --- a/docs/jpsi2ksp.ipynb +++ b/docs/jpsi2ksp.ipynb @@ -159,7 +159,7 @@ " \"N(1710)+\",\n", " \"N(1720)+\",\n", "]\n", - "resonances = [PARTICLE_DB[name] for name in resonance_names]\n", + "resonances = sorted(PARTICLE_DB[name] for name in resonance_names)\n", "Markdown(as_markdown_table(resonances))" ] }, @@ -187,7 +187,7 @@ " chains.extend(_create_isobar(res, min_ls))\n", " return ThreeBodyDecay(\n", " states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n", - " chains=tuple(chains),\n", + " chains=tuple(sorted(chains)),\n", " )\n", "\n", "\n", diff --git a/src/ampform_dpd/decay.py b/src/ampform_dpd/decay.py index d5fed3ce..f948a2f7 100644 --- a/src/ampform_dpd/decay.py +++ b/src/ampform_dpd/decay.py @@ -13,7 +13,7 @@ import sympy as sp -@frozen +@frozen(order=True) class Particle: name: str latex: str @@ -23,7 +23,7 @@ class Particle: width: float -@frozen +@frozen(order=True) class IsobarNode: parent: Particle child1: Particle | IsobarNode @@ -104,7 +104,7 @@ def get_decay_product_ids(spectator_id: Literal[1, 2, 3]) -> tuple[int, int]: """Mapping of the initial and final state IDs to their `.Particle` definition.""" -@frozen +@frozen(order=True) class ThreeBodyDecayChain: decay: IsobarNode = field(validator=instance_of(IsobarNode)) @@ -156,7 +156,7 @@ def outgoing_ls(self) -> LSCoupling: return self.decay.child1.interaction -@frozen +@frozen(order=True) class LSCoupling: L: int S: sp.Rational = field(converter=to_rational, validator=assert_spin_value) From b5eabe2a5fe35b9fd1a69b62e7bb667ad65a793c Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 22 Apr 2024 14:46:53 +0200 Subject: [PATCH 05/11] FEAT: implement `filter_min_ls()` --- src/ampform_dpd/adapter/__init__.py | 0 src/ampform_dpd/adapter/qrules.py | 28 +++++++++++++++ tests/adapter/__init__.py | 0 tests/adapter/test_qrules.py | 56 +++++++++++++++++++++++++++++ 4 files changed, 84 insertions(+) create mode 100644 src/ampform_dpd/adapter/__init__.py create mode 100644 src/ampform_dpd/adapter/qrules.py create mode 100644 tests/adapter/__init__.py create mode 100644 tests/adapter/test_qrules.py diff --git a/src/ampform_dpd/adapter/__init__.py b/src/ampform_dpd/adapter/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/ampform_dpd/adapter/qrules.py b/src/ampform_dpd/adapter/qrules.py new file mode 100644 index 00000000..5e41cdfd --- /dev/null +++ b/src/ampform_dpd/adapter/qrules.py @@ -0,0 +1,28 @@ +from __future__ import annotations + +from collections import defaultdict +from typing import Iterable + +from qrules.topology import EdgeType, FrozenTransition, NodeType + + +def filter_min_ls( + transitions: Iterable[FrozenTransition[EdgeType, NodeType]], +) -> tuple[FrozenTransition[EdgeType, NodeType], ...]: + grouped_transitions = defaultdict(list) + for transition in transitions: + resonances = tuple(transition.intermediate_states.values()) + grouped_transitions[resonances].append(transition) + min_transitions = [] + for group in grouped_transitions.values(): + transition, *_ = group + min_transition = FrozenTransition( + topology=transition.topology, + states=transition.states, + interactions={ + i: min(t.interactions[i] for t in group) + for i in transition.interactions + }, + ) + min_transitions.append(min_transition) + return tuple(min_transitions) diff --git a/tests/adapter/__init__.py b/tests/adapter/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/adapter/test_qrules.py b/tests/adapter/test_qrules.py new file mode 100644 index 00000000..a4740bca --- /dev/null +++ b/tests/adapter/test_qrules.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, TypedDict + +import qrules + +from ampform_dpd.adapter.qrules import filter_min_ls + +if TYPE_CHECKING: + from qrules.transition import StateTransition + + +def test_filter_min_ls(): + reaction = qrules.generate_transitions( + initial_state=[("J/psi(1S)", [+1])], + final_state=["K0", ("Sigma+", [+0.5]), ("p~", [+0.5])], + allowed_interaction_types="strong", + allowed_intermediate_particles=["Sigma(1660)"], + formalism="canonical-helicity", + ) + transitions = tuple( + t for t in reaction.transitions if t.states[3].spin_projection == +0.5 + ) + + ls_couplings = [_get_couplings(t) for t in transitions] + assert ls_couplings == [ + ( + {"L": 0, "S": 1.0}, + {"L": 1, "S": 0.5}, + ), + ( + {"L": 2, "S": 1.0}, + {"L": 1, "S": 0.5}, + ), + ] + + min_ls_transitions = filter_min_ls(transitions) + ls_couplings = [_get_couplings(t) for t in min_ls_transitions] + assert ls_couplings == [ + ( + {"L": 0, "S": 1.0}, + {"L": 1, "S": 0.5}, + ), + ] + + +def _get_couplings(transition: StateTransition) -> tuple[LSCouplings, LSCouplings]: + return tuple( + {"L": node.l_magnitude, "S": node.s_magnitude} + for node in transition.interactions.values() + ) + + +class LSCouplings(TypedDict): + L: int + S: float From c177da15bb1c4ac782fba966241810d9f65e1023 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 22 Apr 2024 14:51:30 +0200 Subject: [PATCH 06/11] FEAT: implement `convert_edges_and_nodes()` --- src/ampform_dpd/adapter/qrules.py | 54 ++++++++++++++++++++++++++++++- 1 file changed, 53 insertions(+), 1 deletion(-) diff --git a/src/ampform_dpd/adapter/qrules.py b/src/ampform_dpd/adapter/qrules.py index 5e41cdfd..d7905f90 100644 --- a/src/ampform_dpd/adapter/qrules.py +++ b/src/ampform_dpd/adapter/qrules.py @@ -1,9 +1,61 @@ from __future__ import annotations from collections import defaultdict -from typing import Iterable +from typing import Any, Iterable +from qrules.quantum_numbers import InteractionProperties from qrules.topology import EdgeType, FrozenTransition, NodeType +from qrules.transition import State + +from ampform_dpd.decay import LSCoupling, Particle + + +def convert_edges_and_nodes( + transitions: Iterable[FrozenTransition], +) -> tuple[FrozenTransition[Particle, LSCoupling], ...]: + unique_transitions = { + transition.convert( + state_converter=_convert_edge, + interaction_converter=_convert_node, + ) + for transition in transitions + } + return tuple(sorted(unique_transitions)) + + +def _convert_edge(state: Any) -> Particle: + if isinstance(state, Particle): + return state + if not isinstance(state, State): + msg = f"Cannot convert state of type {type(state)}" + raise NotImplementedError(msg) + particle = state.particle + if particle.parity is None: + msg = f"Cannot convert particle {particle.name} with undefined parity" + raise NotImplementedError(msg) + return Particle( + name=particle.name, + latex=particle.latex, + spin=particle.spin, + parity=particle.parity, + mass=particle.mass, + width=particle.width, + ) + + +def _convert_node(node: Any) -> Particle: + if isinstance(node, LSCoupling): + return node + if not isinstance(node, InteractionProperties): + msg = f"Cannot convert node of type {type(node)}" + raise NotImplementedError(msg) + if node.l_magnitude is None or node.s_magnitude is None: + msg = "Cannot convert node with undefined L or S" + raise NotImplementedError(msg) + return LSCoupling( + L=node.l_magnitude, + S=node.s_magnitude, + ) def filter_min_ls( From d872b92414914e806b3099d2551df12bb00113d2 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 22 Apr 2024 15:05:42 +0200 Subject: [PATCH 07/11] FEAT: implement `to_three_body_decay()` --- src/ampform_dpd/adapter/qrules.py | 54 ++++++++++++++++++++++++++++++- tests/adapter/test_qrules.py | 45 +++++++++++++++++++------- 2 files changed, 87 insertions(+), 12 deletions(-) diff --git a/src/ampform_dpd/adapter/qrules.py b/src/ampform_dpd/adapter/qrules.py index d7905f90..28603e7f 100644 --- a/src/ampform_dpd/adapter/qrules.py +++ b/src/ampform_dpd/adapter/qrules.py @@ -7,7 +7,59 @@ from qrules.topology import EdgeType, FrozenTransition, NodeType from qrules.transition import State -from ampform_dpd.decay import LSCoupling, Particle +from ampform_dpd.decay import ( + IsobarNode, + LSCoupling, + Particle, + ThreeBodyDecay, + ThreeBodyDecayChain, +) + + +def to_three_body_decay( + transitions: Iterable[FrozenTransition[Particle, LSCoupling]], +) -> ThreeBodyDecay: + transitions = tuple(transitions) + if not transitions: + msg = "Need at least one transition object" + raise ValueError(msg) + some_transition = transitions[0] + return ThreeBodyDecay( + states={**some_transition.initial_states, **some_transition.final_states}, + chains=tuple(sorted(to_decay_chain(t) for t in transitions)), + ) + + +def to_decay_chain( + transition: FrozenTransition[Particle, LSCoupling], +) -> ThreeBodyDecayChain: + if len(transition.initial_states) != 1: + msg = f"Can only handle one initial state, but got {len(transition.initial_states)}" + raise ValueError(msg) + if len(transition.final_states) != 3: # noqa: PLR2004 + msg = f"Can only handle three final states, but got {len(transition.final_states)}" + raise ValueError(msg) + if len(transition.interactions) != 2: # noqa: PLR2004 + msg = f"There are {len(transition.interactions)} interaction nodes, so this can't be a three-body decay" + raise ValueError(msg) + topology = transition.topology + spectator_id, resonance_id = sorted(topology.get_edge_ids_outgoing_from_node(0)) + resonance_id, *_ = sorted(topology.get_edge_ids_ingoing_to_node(1)) + child1_id, child2_id = sorted(topology.get_edge_ids_outgoing_from_node(1)) + parent, *_ = transition.initial_states.values() + production_node, decay_node = transition.interactions.values() + isobar = IsobarNode( + parent=parent, + child1=IsobarNode( + parent=transition.states[resonance_id], + child1=transition.states[child1_id], + child2=transition.states[child2_id], + interaction=decay_node, + ), + child2=transition.states[spectator_id], + interaction=production_node, + ) + return ThreeBodyDecayChain(decay=isobar) def convert_edges_and_nodes( diff --git a/tests/adapter/test_qrules.py b/tests/adapter/test_qrules.py index a4740bca..b97b7c3e 100644 --- a/tests/adapter/test_qrules.py +++ b/tests/adapter/test_qrules.py @@ -1,23 +1,29 @@ from __future__ import annotations -from typing import TYPE_CHECKING, TypedDict +from typing import TYPE_CHECKING +import pytest import qrules -from ampform_dpd.adapter.qrules import filter_min_ls +from ampform_dpd.adapter.qrules import filter_min_ls, to_three_body_decay +from ampform_dpd.decay import LSCoupling, Particle if TYPE_CHECKING: - from qrules.transition import StateTransition + from qrules.transition import ReactionInfo, StateTransition -def test_filter_min_ls(): - reaction = qrules.generate_transitions( +@pytest.fixture(scope="session") +def reaction() -> ReactionInfo: + return qrules.generate_transitions( initial_state=[("J/psi(1S)", [+1])], final_state=["K0", ("Sigma+", [+0.5]), ("p~", [+0.5])], allowed_interaction_types="strong", allowed_intermediate_particles=["Sigma(1660)"], formalism="canonical-helicity", ) + + +def test_filter_min_ls(reaction: ReactionInfo): transitions = tuple( t for t in reaction.transitions if t.states[3].spin_projection == +0.5 ) @@ -44,13 +50,30 @@ def test_filter_min_ls(): ] -def _get_couplings(transition: StateTransition) -> tuple[LSCouplings, LSCouplings]: +@pytest.mark.parametrize("min_ls", [False, True]) +def test_to_three_body_decay(reaction: ReactionInfo, min_ls: bool): + decay = to_three_body_decay(reaction.transitions, min_ls) + assert decay.initial_state.name == "J/psi(1S)" + assert {i: p.name for i, p in decay.final_state.items()} == { + 1: "K0", + 2: "Sigma+", + 3: "p~", + } + if min_ls: + assert len(decay.chains) == 1 + assert decay.chains[0].incoming_ls == LSCoupling(L=0, S=1) + assert decay.chains[0].outgoing_ls == LSCoupling(L=1, S=0.5) + else: + assert len(decay.chains) == 2 + assert decay.chains[1].incoming_ls == LSCoupling(L=2, S=1) + assert decay.chains[1].outgoing_ls == LSCoupling(L=1, S=0.5) + for chain in decay.chains: + assert isinstance(chain.resonance, Particle) + assert chain.resonance.name == "Sigma(1660)~-" + + +def _get_couplings(transition: StateTransition) -> tuple[dict, dict]: return tuple( {"L": node.l_magnitude, "S": node.s_magnitude} for node in transition.interactions.values() ) - - -class LSCouplings(TypedDict): - L: int - S: float From bb7b3bf92595633ee8e66c8ea0e468ae7bb71cb0 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 23 Apr 2024 14:08:31 +0200 Subject: [PATCH 08/11] FEAT: implement `load_particles()` --- .cspell.json | 2 ++ .vscode/settings.json | 7 ++++++- src/ampform_dpd/adapter/qrules.py | 24 +++++++++++++++++++++--- src/ampform_dpd/particle-definitions.yml | 14 ++++++++++++++ 4 files changed, 43 insertions(+), 4 deletions(-) create mode 100644 src/ampform_dpd/particle-definitions.yml diff --git a/.cspell.json b/.cspell.json index 0f353a25..a4b54641 100644 --- a/.cspell.json +++ b/.cspell.json @@ -16,6 +16,7 @@ "**/*.bib", "**/.cspell.json", "*.ico", + "*particle*.*ml", ".constraints/*.txt", ".editorconfig", ".gitignore", @@ -138,6 +139,7 @@ "Gordan", "helicities", "helicity", + "isospin", "JPAC", "lambdify", "lambdifying", diff --git a/.vscode/settings.json b/.vscode/settings.json index 8d4db841..0408778f 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -54,5 +54,10 @@ "ruff.enable": true, "ruff.importStrategy": "fromEnvironment", "ruff.organizeImports": true, - "telemetry.telemetryLevel": "off" + "telemetry.telemetryLevel": "off", + "yaml.schemas": { + "https://raw.githubusercontent.com/ComPWA/qrules/0.10.x/src/qrules/particle-validation.json": [ + "*particle*.y*ml" + ] + } } diff --git a/src/ampform_dpd/adapter/qrules.py b/src/ampform_dpd/adapter/qrules.py index 28603e7f..c9e7085e 100644 --- a/src/ampform_dpd/adapter/qrules.py +++ b/src/ampform_dpd/adapter/qrules.py @@ -1,8 +1,10 @@ from __future__ import annotations from collections import defaultdict +from pathlib import Path from typing import Any, Iterable +import qrules from qrules.quantum_numbers import InteractionProperties from qrules.topology import EdgeType, FrozenTransition, NodeType from qrules.transition import State @@ -17,15 +19,23 @@ def to_three_body_decay( - transitions: Iterable[FrozenTransition[Particle, LSCoupling]], + transitions: Iterable[FrozenTransition], + min_ls: bool = False, ) -> ThreeBodyDecay: - transitions = tuple(transitions) + transitions = convert_edges_and_nodes(transitions) + if min_ls: + transitions = filter_min_ls(transitions) if not transitions: msg = "Need at least one transition object" raise ValueError(msg) some_transition = transitions[0] + initial_state, *_ = some_transition.initial_states.values() + final_states = { + i: some_transition.final_states[idx] + for i, idx in enumerate(sorted(some_transition.final_states), 1) + } return ThreeBodyDecay( - states={**some_transition.initial_states, **some_transition.final_states}, + states={0: initial_state, **final_states}, chains=tuple(sorted(to_decay_chain(t) for t in transitions)), ) @@ -130,3 +140,11 @@ def filter_min_ls( ) min_transitions.append(min_transition) return tuple(min_transitions) + + +def load_particles() -> qrules.particle.ParticleCollection: + src_dir = Path(__file__).parent.parent + particle_database = qrules.load_default_particles() + additional_definitions = qrules.io.load(src_dir / "particle-definitions.yml") + particle_database.update(additional_definitions) + return particle_database diff --git a/src/ampform_dpd/particle-definitions.yml b/src/ampform_dpd/particle-definitions.yml new file mode 100644 index 00000000..bfad21b8 --- /dev/null +++ b/src/ampform_dpd/particle-definitions.yml @@ -0,0 +1,14 @@ +particles: + - name: Lambda(2000) + latex: \Lambda(2000) + mass: 2.0 + width: 0.21 + pid: 11111111 + spin: 0.5 + parity: + value: -1 + strangeness: -1 + baryon_number: +1 + isospin: + magnitude: 0 + projection: 0 From ae18e0226fd9e12bf5e94ccfadd95858436ce681 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Wed, 24 Apr 2024 11:14:04 +0200 Subject: [PATCH 09/11] FEAT: make `get_particle()` public --- src/ampform_dpd/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ampform_dpd/__init__.py b/src/ampform_dpd/__init__.py index 918c01bf..f316525a 100644 --- a/src/ampform_dpd/__init__.py +++ b/src/ampform_dpd/__init__.py @@ -290,8 +290,8 @@ def _formulate_clebsch_gordan_factors( raise ValueError(msg) # https://github.com/ComPWA/ampform/blob/65b4efa/src/ampform/helicity/__init__.py#L785-L802 # and supplementary material p.1 (https://cds.cern.ch/record/2824328/files) - child1 = _get_particle(isobar.child1) - child2 = _get_particle(isobar.child2) + child1 = get_particle(isobar.child1) + child2 = get_particle(isobar.child2) child1_helicity = helicities[child1] child2_helicity = helicities[child2] cg_ss = CG( @@ -314,7 +314,7 @@ def _formulate_clebsch_gordan_factors( return sqrt_factor * cg_ll * cg_ss -def _get_particle(isobar: IsobarNode | Particle) -> Particle: +def get_particle(isobar: IsobarNode | Particle) -> Particle: if isinstance(isobar, IsobarNode): return isobar.parent return isobar From fc18dcce8d99b6729f3684fcc67504e0a4503c9b Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Wed, 24 Apr 2024 12:00:08 +0200 Subject: [PATCH 10/11] FEAT: support decays without LS --- src/ampform_dpd/adapter/qrules.py | 11 ++++++----- src/ampform_dpd/decay.py | 10 ++-------- src/ampform_dpd/io.py | 6 ++++-- 3 files changed, 12 insertions(+), 15 deletions(-) diff --git a/src/ampform_dpd/adapter/qrules.py b/src/ampform_dpd/adapter/qrules.py index c9e7085e..d2a83e4f 100644 --- a/src/ampform_dpd/adapter/qrules.py +++ b/src/ampform_dpd/adapter/qrules.py @@ -41,7 +41,7 @@ def to_three_body_decay( def to_decay_chain( - transition: FrozenTransition[Particle, LSCoupling], + transition: FrozenTransition[Particle, LSCoupling | None], ) -> ThreeBodyDecayChain: if len(transition.initial_states) != 1: msg = f"Can only handle one initial state, but got {len(transition.initial_states)}" @@ -74,7 +74,7 @@ def to_decay_chain( def convert_edges_and_nodes( transitions: Iterable[FrozenTransition], -) -> tuple[FrozenTransition[Particle, LSCoupling], ...]: +) -> tuple[FrozenTransition[Particle, LSCoupling | None], ...]: unique_transitions = { transition.convert( state_converter=_convert_edge, @@ -105,15 +105,16 @@ def _convert_edge(state: Any) -> Particle: ) -def _convert_node(node: Any) -> Particle: +def _convert_node(node: Any) -> LSCoupling | None: + if node is None: + return None if isinstance(node, LSCoupling): return node if not isinstance(node, InteractionProperties): msg = f"Cannot convert node of type {type(node)}" raise NotImplementedError(msg) if node.l_magnitude is None or node.s_magnitude is None: - msg = "Cannot convert node with undefined L or S" - raise NotImplementedError(msg) + return None return LSCoupling( L=node.l_magnitude, S=node.s_magnitude, diff --git a/src/ampform_dpd/decay.py b/src/ampform_dpd/decay.py index f948a2f7..2d538f1c 100644 --- a/src/ampform_dpd/decay.py +++ b/src/ampform_dpd/decay.py @@ -121,12 +121,6 @@ def __attrs_post_init__(self) -> None: if not isinstance(self.decay.child2, Particle): msg = f"Child 2 has of type {Particle.__name__} (spectator)" raise TypeError(msg) - if self.incoming_ls is None: # pyright: ignore[reportUnnecessaryComparison] - msg = "LS-coupling for production node required" - raise ValueError(msg) - if self.outgoing_ls is None: # pyright: ignore[reportUnnecessaryComparison] - msg = "LS-coupling for decay node required" - raise ValueError(msg) @property def parent(self) -> Particle: @@ -148,11 +142,11 @@ def spectator(self) -> Particle: return self.decay.child2 @property - def incoming_ls(self) -> LSCoupling: + def incoming_ls(self) -> LSCoupling | None: return self.decay.interaction @property - def outgoing_ls(self) -> LSCoupling: + def outgoing_ls(self) -> LSCoupling | None: return self.decay.child1.interaction diff --git a/src/ampform_dpd/io.py b/src/ampform_dpd/io.py index e737ae63..6b4d1766 100644 --- a/src/ampform_dpd/io.py +++ b/src/ampform_dpd/io.py @@ -200,9 +200,11 @@ def _as_decay_markdown_table(decay_chains: Sequence[ThreeBodyDecayChain]) -> str Rf"${aslatex(chain.resonance, only_jp=True)}$", f"{int(1e3 * chain.resonance.mass):,.0f}", f"{int(1e3 * chain.resonance.width):,.0f}", - chain.outgoing_ls.L, - chain.incoming_ls.L, ] + if chain.outgoing_ls is not None: + row_items.append(chain.outgoing_ls.L) + if chain.incoming_ls is not None: + row_items.append(chain.incoming_ls.L) src += _create_markdown_table_row(row_items) return src From 14cef0a86e01f59e02041fa37ff7097db4b2f740 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Wed, 24 Apr 2024 14:44:13 +0200 Subject: [PATCH 11/11] DOC: build all models with QRules --- .cspell.json | 1 + docs/comparison/d2kkk.ipynb | 253 ++++++++--------------------- docs/comparison/jpsi2phipipi.ipynb | 244 ++++++++-------------------- docs/comparison/jpsi2pipipi.ipynb | 244 ++++++++-------------------- docs/conf.py | 7 +- docs/jpsi2ksp.ipynb | 237 +++++++++------------------ docs/lc2pkpi.ipynb | 215 +++++++++--------------- 7 files changed, 371 insertions(+), 830 deletions(-) diff --git a/.cspell.json b/.cspell.json index a4b54641..17475a4e 100644 --- a/.cspell.json +++ b/.cspell.json @@ -154,6 +154,7 @@ "pytest", "PYTHONHASHSEED", "QRules", + "sympify", "SymPy", "TensorWaves", "Weisskopf" diff --git a/docs/comparison/d2kkk.ipynb b/docs/comparison/d2kkk.ipynb index 1bed9082..c321153c 100644 --- a/docs/comparison/d2kkk.ipynb +++ b/docs/comparison/d2kkk.ipynb @@ -2,7 +2,9 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "# D⁰ → K⁰ K⁺ K⁻\n", "\n", @@ -31,12 +33,11 @@ "source": [ "from __future__ import annotations\n", "\n", - "import itertools\n", "import logging\n", "import os\n", "import warnings\n", "from textwrap import dedent\n", - "from typing import TYPE_CHECKING, Iterable\n", + "from typing import TYPE_CHECKING\n", "\n", "import ampform\n", "import graphviz\n", @@ -48,7 +49,7 @@ "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n", "from ampform.sympy import perform_cached_doit\n", - "from IPython.display import SVG, Latex, Markdown, clear_output, display\n", + "from IPython.display import Latex, Markdown, clear_output, display\n", "from ipywidgets import (\n", " Accordion,\n", " Checkbox,\n", @@ -66,7 +67,8 @@ "from tensorwaves.data.transform import SympyDataTransformer\n", "\n", "from ampform_dpd import DalitzPlotDecompositionBuilder\n", - "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecay, ThreeBodyDecayChain\n", + "from ampform_dpd.adapter.qrules import to_three_body_decay\n", + "from ampform_dpd.decay import Particle\n", "from ampform_dpd.io import (\n", " as_markdown_table,\n", " aslatex,\n", @@ -74,7 +76,6 @@ " perform_cached_lambdify,\n", " simplify_latex_rendering,\n", ")\n", - "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n", "\n", "if TYPE_CHECKING:\n", " from ampform.helicity import HelicityModel\n", @@ -105,11 +106,8 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "mystnb": { - "code_prompt_show": "Define initial and final state" + "code_prompt_show": "Generate transitions" }, "tags": [ "hide-input" @@ -117,26 +115,16 @@ }, "outputs": [], "source": [ - "PDG = qrules.load_pdg()\n", - "PARTICLE_DB = {\n", - " p.name: Particle(\n", - " name=p.name,\n", - " latex=p.latex,\n", - " spin=p.spin,\n", - " parity=int(p.parity),\n", - " mass=p.mass,\n", - " width=p.width,\n", - " )\n", - " for p in PDG\n", - " if p.parity is not None\n", - "}\n", - "INITIAL_STATE = PARTICLE_DB[\"D0\"]\n", - "FS_zero = PARTICLE_DB[\"K0\"]\n", - "FS_neg = PARTICLE_DB[\"K-\"]\n", - "FS_pos = PARTICLE_DB[\"K+\"]\n", - "PARTICLE_TO_ID = {INITIAL_STATE: 0, FS_zero: 1, FS_neg: 2, FS_pos: 3}\n", - "_, *FINAL_STATE = PARTICLE_TO_ID\n", - "Markdown(as_markdown_table(list(PARTICLE_TO_ID)))" + "REACTION = qrules.generate_transitions(\n", + " initial_state=\"D0\",\n", + " final_state=[\"K0\", \"K-\", \"K+\"],\n", + " allowed_intermediate_particles=[\"a(0)\", \"f(0)(980)\", \"pi(1)\"],\n", + " mass_conservation_factor=0.2,\n", + " formalism=\"helicity\",\n", + ")\n", + "REACTION123 = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", + "graphviz.Source(dot)" ] }, { @@ -146,25 +134,14 @@ "jupyter": { "source_hidden": true }, - "mystnb": { - "code_prompt_show": "Generate transitions" - }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ - "REACTION = qrules.generate_transitions(\n", - " initial_state=INITIAL_STATE.name,\n", - " final_state=[p.name for p in FINAL_STATE],\n", - " allowed_intermediate_particles=[\"a(0)\", \"f(0)(980)\", \"pi(1)\"],\n", - " mass_conservation_factor=0.2,\n", - " formalism=\"helicity\",\n", - ")\n", - "REACTION123 = relabel_edge_ids(REACTION)\n", - "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", - "graphviz.Source(dot)" + "DECAY = to_three_body_decay(REACTION123.transitions, min_ls=True)\n", + "Markdown(as_markdown_table([DECAY.initial_state, *DECAY.final_state.values()]))" ] }, { @@ -180,12 +157,11 @@ }, "outputs": [], "source": [ - "qrules_resonances = sorted(\n", - " REACTION.get_intermediate_particles(),\n", - " key=lambda p: (p.charge, p.mass, p.name),\n", + "resonances = sorted(\n", + " {t.resonance for t in DECAY.chains},\n", + " key=lambda p: (p.name[0], p.mass),\n", ")\n", - "resonance_names = [p.name for p in qrules_resonances]\n", - "resonances = [PARTICLE_DB[name] for name in resonance_names]\n", + "resonance_names = [p.name for p in resonances]\n", "Markdown(as_markdown_table(resonances))" ] }, @@ -197,84 +173,11 @@ "source_hidden": true }, "tags": [ - "hide-input", - "scroll-input" + "hide-input" ] }, "outputs": [], "source": [ - "def load_three_body_decay(\n", - " resonance_names: Iterable[str],\n", - " particle_definitions: dict[str, Particle],\n", - " min_ls: bool = True,\n", - ") -> ThreeBodyDecay:\n", - " _resonances = [particle_definitions[name] for name in resonance_names]\n", - " chains: list[ThreeBodyDecayChain] = []\n", - " for res in _resonances:\n", - " chains.extend(_create_isobar(res, min_ls))\n", - " return ThreeBodyDecay(\n", - " states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n", - " chains=tuple(chains),\n", - " )\n", - "\n", - "\n", - "def _create_isobar(resonance: Particle, min_ls: bool) -> list[ThreeBodyDecayChain]:\n", - " if resonance.name.endswith(\"-\"):\n", - " child1, child2, spectator = FS_zero, FS_neg, FS_pos\n", - " elif resonance.name.endswith(\"+\"):\n", - " child1, child2, spectator = FS_pos, FS_zero, FS_neg\n", - " else:\n", - " child1, child2, spectator = FS_neg, FS_pos, FS_zero\n", - " prod_ls_couplings = _generate_ls(\n", - " INITIAL_STATE, resonance, spectator, conserve_parity=False\n", - " )\n", - " dec_ls_couplings = _generate_ls(resonance, child1, child2, conserve_parity=True)\n", - " if min_ls:\n", - " decay = IsobarNode(\n", - " parent=INITIAL_STATE,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=min(dec_ls_couplings),\n", - " ),\n", - " child2=spectator,\n", - " interaction=min(prod_ls_couplings),\n", - " )\n", - " return [ThreeBodyDecayChain(decay)]\n", - " chains = []\n", - " for dec_ls, prod_ls in itertools.product(dec_ls_couplings, prod_ls_couplings):\n", - " decay = IsobarNode(\n", - " parent=INITIAL_STATE,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=dec_ls,\n", - " ),\n", - " child2=spectator,\n", - " interaction=prod_ls,\n", - " )\n", - " chains.append(ThreeBodyDecayChain(decay))\n", - " return chains\n", - "\n", - "\n", - "def _generate_ls(\n", - " parent: Particle, child1: Particle, child2: Particle, conserve_parity: bool\n", - ") -> list[tuple[int, sp.Rational]]:\n", - " ls = generate_ls_couplings(parent.spin, child1.spin, child2.spin)\n", - " if conserve_parity:\n", - " return filter_parity_violating_ls(\n", - " ls, parent.parity, child1.parity, child2.parity\n", - " )\n", - " return ls\n", - "\n", - "\n", - "DECAY = load_three_body_decay(\n", - " resonance_names,\n", - " particle_definitions=PARTICLE_DB,\n", - " min_ls=True,\n", - ")\n", "Latex(aslatex(DECAY, with_jp=True))" ] }, @@ -299,10 +202,7 @@ }, { "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true, - "tags": [] - }, + "metadata": {}, "source": [ "Note that, as opposed to {ref}`Λc⁺ → pπ⁺K⁻` and {ref}`J/ψ → K⁰Σ⁺p̅`, there are no Wigner-$d$ functions, because the final state is spinless." ] @@ -322,8 +222,9 @@ "outputs": [], "source": [ "model_builder = DalitzPlotDecompositionBuilder(DECAY, min_ls=True)\n", - "dpd_model = model_builder.formulate(reference_subsystem=1)\n", - "dpd_model.intensity" + "DPD_MODEL = model_builder.formulate(reference_subsystem=1)\n", + "del model_builder\n", + "DPD_MODEL.intensity.cleanup()" ] }, { @@ -340,12 +241,14 @@ }, "outputs": [], "source": [ - "Latex(aslatex(dpd_model.amplitudes))" + "Latex(aslatex(DPD_MODEL.amplitudes))" ] }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "There is an isobar Wigner-$d$ function, which takes the following helicity angles as argument:" ] @@ -363,7 +266,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(dpd_model.variables))" + "Latex(aslatex(DPD_MODEL.variables))" ] }, { @@ -380,10 +283,11 @@ "outputs": [], "source": [ "masses = {\n", - " sp.Symbol(f\"m{i}\", nonnegative=True): round(p.mass, 7)\n", - " for p, i in PARTICLE_TO_ID.items()\n", + " symbol: sp.sympify(value)\n", + " for symbol, value in DPD_MODEL.parameter_defaults.items()\n", + " if symbol.name.startswith(\"m\")\n", + " if len(symbol.name) == 2\n", "}\n", - "dpd_model.parameter_defaults.update(masses)\n", "Latex(aslatex(masses))" ] }, @@ -399,7 +303,9 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "AmpForm does not formulate alignment Wigner-$D$ functions. For the case of this spinless final state, this means the intensity is the same as that of the [](#dpd-model)." ] @@ -419,8 +325,10 @@ "source": [ "model_builder = ampform.get_builder(REACTION)\n", "model_builder.use_helicity_couplings = False\n", - "ampform_model = model_builder.formulate()\n", - "ampform_model.intensity" + "model_builder.config.scalar_initial_state_mass = True\n", + "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", + "AMPFORM_MODEL = model_builder.formulate()\n", + "AMPFORM_MODEL.intensity" ] }, { @@ -437,7 +345,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(ampform_model.amplitudes))" + "Latex(aslatex(AMPFORM_MODEL.amplitudes))" ] }, { @@ -453,7 +361,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(ampform_model.kinematic_variables))" + "Latex(aslatex(AMPFORM_MODEL.kinematic_variables))" ] }, { @@ -470,6 +378,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "mystnb": { "code_prompt_show": "Formulate kinematic variables in terms of four-momenta" }, @@ -482,6 +393,7 @@ "p1, p2, p3 = tuple(FourMomentumSymbol(f\"p{i}\", shape=[]) for i in (0, 1, 2))\n", "s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", "mass_definitions = {\n", + " **masses,\n", " s1: InvariantMass(p2 + p3) ** 2,\n", " s2: InvariantMass(p1 + p3) ** 2,\n", " s3: InvariantMass(p1 + p2) ** 2,\n", @@ -490,16 +402,13 @@ " sp.Symbol(\"m_12\", nonnegative=True): InvariantMass(p2 + p3),\n", "}\n", "dpd_variables = {\n", - " sp.Symbol(f\"m{i}\", nonnegative=True): sp.Float(p.mass)\n", - " for i, p in enumerate(PARTICLE_TO_ID)\n", + " symbol: expr.doit().xreplace(DPD_MODEL.variables).xreplace(mass_definitions)\n", + " for symbol, expr in DPD_MODEL.variables.items()\n", "}\n", - "for symbol, expr in dpd_model.variables.items():\n", - " expr = expr.doit().xreplace(mass_definitions).xreplace(dpd_variables)\n", - " dpd_variables[symbol] = expr\n", "dpd_transformer = SympyDataTransformer.from_sympy(dpd_variables, backend=\"jax\")\n", "\n", "ampform_transformer = SympyDataTransformer.from_sympy(\n", - " ampform_model.kinematic_variables, backend=\"jax\"\n", + " AMPFORM_MODEL.kinematic_variables, backend=\"jax\"\n", ")" ] }, @@ -520,7 +429,7 @@ " return phsp_generator.generate(size, rng)\n", "\n", "\n", - "phsp = generate_phase_space(ampform_model.reaction_info, size=100_000)\n", + "phsp = generate_phase_space(AMPFORM_MODEL.reaction_info, size=100_000)\n", "ampform_phsp = ampform_transformer(phsp)\n", "dpd_phsp = dpd_transformer(phsp)" ] @@ -551,8 +460,8 @@ " return unfolded_intensity.xreplace(unfolded_amplitudes)\n", "\n", "\n", - "ampform_intensity_expr = unfold_intensity(ampform_model)\n", - "dpd_intensity_expr = unfold_intensity(dpd_model)" + "ampform_intensity_expr = unfold_intensity(AMPFORM_MODEL)\n", + "dpd_intensity_expr = unfold_intensity(DPD_MODEL)" ] }, { @@ -565,11 +474,11 @@ "source": [ "ampform_func = perform_cached_lambdify(\n", " ampform_intensity_expr,\n", - " parameters=ampform_model.parameter_defaults,\n", + " parameters=AMPFORM_MODEL.parameter_defaults,\n", ")\n", "dpd_func = perform_cached_lambdify(\n", " dpd_intensity_expr,\n", - " parameters=dpd_model.parameter_defaults,\n", + " parameters=DPD_MODEL.parameter_defaults,\n", ")" ] }, @@ -643,8 +552,8 @@ "outputs": [], "source": [ "def create_sliders() -> dict[str, ToggleButtons]:\n", - " all_parameters = dict(ampform_model.parameter_defaults.items())\n", - " all_parameters.update(dpd_model.parameter_defaults)\n", + " all_parameters = dict(AMPFORM_MODEL.parameter_defaults.items())\n", + " all_parameters.update(DPD_MODEL.parameter_defaults)\n", " sliders = {}\n", " for symbol, value in all_parameters.items():\n", " value = \"+1\"\n", @@ -729,6 +638,7 @@ "code_prompt_show": "Generate comparison widget" }, "tags": [ + "full-width", "hide-input", "scroll-input" ] @@ -745,17 +655,18 @@ " (ax_s1, ax_s2, ax_s3),\n", " (ax_t1, ax_t2, ax_t3),\n", ") = axes\n", - "final_state = ampform_model.reaction_info.final_state\n", "for ax in axes[:, 0].flatten():\n", " ax.set_ylabel(\"Intensity (a.u.)\")\n", "for ax in axes[:, 1:].flatten():\n", " ax.set_yticks([])\n", - "ax_s1.set_xlabel(f\"$m({FINAL_STATE[1].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_s2.set_xlabel(f\"$m({FINAL_STATE[0].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_s3.set_xlabel(f\"$m({FINAL_STATE[0].latex}, {FINAL_STATE[1].latex})$\")\n", - "ax_t1.set_xlabel(Rf\"$\\theta({FINAL_STATE[1].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_t2.set_xlabel(Rf\"$\\theta({FINAL_STATE[0].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_t3.set_xlabel(Rf\"$\\theta({FINAL_STATE[0].latex}, {FINAL_STATE[1].latex})$\")\n", + "\n", + "final_state = DECAY.final_state\n", + "ax_s1.set_xlabel(f\"$m({final_state[2].latex}, {final_state[3].latex})$\")\n", + "ax_s2.set_xlabel(f\"$m({final_state[1].latex}, {final_state[3].latex})$\")\n", + "ax_s3.set_xlabel(f\"$m({final_state[1].latex}, {final_state[2].latex})$\")\n", + "ax_t1.set_xlabel(Rf\"$\\theta({final_state[2].latex}, {final_state[3].latex})$\")\n", + "ax_t2.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[3].latex})$\")\n", + "ax_t3.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[2].latex})$\")\n", "fig.suptitle(f'Selected resonances: ${\", \".join(resonance_selector.value)}$')\n", "fig.tight_layout()\n", "\n", @@ -950,13 +861,13 @@ " ampform_expr = _simplify(\n", " ampform_intensity_expr,\n", " ampform_pars,\n", - " ampform_model.kinematic_variables,\n", + " AMPFORM_MODEL.kinematic_variables,\n", " selected_resonances,\n", " )\n", " dpd_expr = _simplify(\n", " dpd_intensity_expr,\n", " dpd_pars,\n", - " dpd_model.variables,\n", + " DPD_MODEL.variables,\n", " selected_resonances,\n", " )\n", " else:\n", @@ -983,26 +894,6 @@ ")\n", "display(output, ui)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if NO_TQDM:\n", - " filename = \"d02kkk-comparison.svg\"\n", - " plt.savefig(filename)\n", - " display(SVG(filename))" - ] } ], "metadata": { diff --git a/docs/comparison/jpsi2phipipi.ipynb b/docs/comparison/jpsi2phipipi.ipynb index 99f88c28..819171c6 100644 --- a/docs/comparison/jpsi2phipipi.ipynb +++ b/docs/comparison/jpsi2phipipi.ipynb @@ -2,7 +2,9 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "# J/ψ → φ(1020) π⁺ π⁻\n", "\n", @@ -31,12 +33,11 @@ "source": [ "from __future__ import annotations\n", "\n", - "import itertools\n", "import logging\n", "import os\n", "import warnings\n", "from textwrap import dedent\n", - "from typing import TYPE_CHECKING, Iterable\n", + "from typing import TYPE_CHECKING\n", "\n", "import ampform\n", "import graphviz\n", @@ -48,7 +49,7 @@ "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n", "from ampform.sympy import perform_cached_doit\n", - "from IPython.display import SVG, Latex, Markdown, clear_output, display\n", + "from IPython.display import Latex, Markdown, clear_output, display\n", "from ipywidgets import (\n", " Accordion,\n", " Checkbox,\n", @@ -66,7 +67,8 @@ "from tensorwaves.data.transform import SympyDataTransformer\n", "\n", "from ampform_dpd import DalitzPlotDecompositionBuilder\n", - "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecay, ThreeBodyDecayChain\n", + "from ampform_dpd.adapter.qrules import to_three_body_decay\n", + "from ampform_dpd.decay import Particle\n", "from ampform_dpd.io import (\n", " as_markdown_table,\n", " aslatex,\n", @@ -74,7 +76,6 @@ " perform_cached_lambdify,\n", " simplify_latex_rendering,\n", ")\n", - "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n", "\n", "if TYPE_CHECKING:\n", " from ampform.helicity import HelicityModel\n", @@ -105,11 +106,8 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "mystnb": { - "code_prompt_show": "Define initial and final state" + "code_prompt_show": "Generate transitions" }, "tags": [ "hide-input" @@ -117,26 +115,16 @@ }, "outputs": [], "source": [ - "PDG = qrules.load_pdg()\n", - "PARTICLE_DB = {\n", - " p.name: Particle(\n", - " name=p.name,\n", - " latex=p.latex,\n", - " spin=p.spin,\n", - " parity=int(p.parity),\n", - " mass=p.mass,\n", - " width=p.width,\n", - " )\n", - " for p in PDG\n", - " if p.parity is not None\n", - "}\n", - "INITIAL_STATE = PARTICLE_DB[\"J/psi(1S)\"]\n", - "FS_zero = PARTICLE_DB[\"phi(1020)\"]\n", - "FS_neg = PARTICLE_DB[\"pi-\"]\n", - "FS_pos = PARTICLE_DB[\"pi+\"]\n", - "PARTICLE_TO_ID = {INITIAL_STATE: 0, FS_zero: 1, FS_neg: 2, FS_pos: 3}\n", - "_, *FINAL_STATE = PARTICLE_TO_ID\n", - "Markdown(as_markdown_table(list(PARTICLE_TO_ID)))" + "REACTION = qrules.generate_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"phi(1020)\", \"pi-\", \"pi+\"],\n", + " allowed_intermediate_particles=[\"a(0)(1450\", \"rho(1450)\"],\n", + " mass_conservation_factor=0,\n", + " formalism=\"helicity\",\n", + ")\n", + "REACTION123 = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", + "graphviz.Source(dot)" ] }, { @@ -146,25 +134,14 @@ "jupyter": { "source_hidden": true }, - "mystnb": { - "code_prompt_show": "Generate transitions" - }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ - "REACTION = qrules.generate_transitions(\n", - " initial_state=INITIAL_STATE.name,\n", - " final_state=[p.name for p in FINAL_STATE],\n", - " allowed_intermediate_particles=[\"a(0)(1450\", \"rho(1450)\"],\n", - " mass_conservation_factor=0,\n", - " formalism=\"helicity\",\n", - ")\n", - "REACTION123 = relabel_edge_ids(REACTION)\n", - "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", - "graphviz.Source(dot)" + "DECAY = to_three_body_decay(REACTION123.transitions, min_ls=True)\n", + "Markdown(as_markdown_table([DECAY.initial_state, *DECAY.final_state.values()]))" ] }, { @@ -180,12 +157,11 @@ }, "outputs": [], "source": [ - "qrules_resonances = sorted(\n", - " REACTION.get_intermediate_particles(),\n", - " key=lambda p: (p.charge, p.mass, p.name),\n", + "resonances = sorted(\n", + " {t.resonance for t in DECAY.chains},\n", + " key=lambda p: (p.name[0], p.mass),\n", ")\n", - "resonance_names = [p.name for p in qrules_resonances]\n", - "resonances = [PARTICLE_DB[name] for name in resonance_names]\n", + "resonance_names = [p.name for p in resonances]\n", "Markdown(as_markdown_table(resonances))" ] }, @@ -197,84 +173,11 @@ "source_hidden": true }, "tags": [ - "hide-input", - "scroll-input" + "hide-input" ] }, "outputs": [], "source": [ - "def load_three_body_decay(\n", - " resonance_names: Iterable[str],\n", - " particle_definitions: dict[str, Particle],\n", - " min_ls: bool = True,\n", - ") -> ThreeBodyDecay:\n", - " _resonances = [particle_definitions[name] for name in resonance_names]\n", - " chains: list[ThreeBodyDecayChain] = []\n", - " for res in _resonances:\n", - " chains.extend(_create_isobar(res, min_ls))\n", - " return ThreeBodyDecay(\n", - " states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n", - " chains=tuple(chains),\n", - " )\n", - "\n", - "\n", - "def _create_isobar(resonance: Particle, min_ls: bool) -> list[ThreeBodyDecayChain]:\n", - " if resonance.name.endswith(\"-\"):\n", - " child1, child2, spectator = FS_zero, FS_neg, FS_pos\n", - " elif resonance.name.endswith(\"+\"):\n", - " child1, child2, spectator = FS_pos, FS_zero, FS_neg\n", - " else:\n", - " child1, child2, spectator = FS_neg, FS_pos, FS_zero\n", - " prod_ls_couplings = _generate_ls(\n", - " INITIAL_STATE, resonance, spectator, conserve_parity=False\n", - " )\n", - " dec_ls_couplings = _generate_ls(resonance, child1, child2, conserve_parity=False)\n", - " if min_ls:\n", - " decay = IsobarNode(\n", - " parent=INITIAL_STATE,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=min(dec_ls_couplings),\n", - " ),\n", - " child2=spectator,\n", - " interaction=min(prod_ls_couplings),\n", - " )\n", - " return [ThreeBodyDecayChain(decay)]\n", - " chains = []\n", - " for dec_ls, prod_ls in itertools.product(dec_ls_couplings, prod_ls_couplings):\n", - " decay = IsobarNode(\n", - " parent=INITIAL_STATE,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=dec_ls,\n", - " ),\n", - " child2=spectator,\n", - " interaction=prod_ls,\n", - " )\n", - " chains.append(ThreeBodyDecayChain(decay))\n", - " return chains\n", - "\n", - "\n", - "def _generate_ls(\n", - " parent: Particle, child1: Particle, child2: Particle, conserve_parity: bool\n", - ") -> list[tuple[int, sp.Rational]]:\n", - " ls = generate_ls_couplings(parent.spin, child1.spin, child2.spin)\n", - " if conserve_parity:\n", - " return filter_parity_violating_ls(\n", - " ls, parent.parity, child1.parity, child2.parity\n", - " )\n", - " return ls\n", - "\n", - "\n", - "DECAY = load_three_body_decay(\n", - " resonance_names,\n", - " particle_definitions=PARTICLE_DB,\n", - " min_ls=True,\n", - ")\n", "Latex(aslatex(DECAY, with_jp=True))" ] }, @@ -312,8 +215,9 @@ "outputs": [], "source": [ "model_builder = DalitzPlotDecompositionBuilder(DECAY, min_ls=True)\n", - "dpd_model = model_builder.formulate(reference_subsystem=1)\n", - "dpd_model.intensity" + "DPD_MODEL = model_builder.formulate(reference_subsystem=1)\n", + "del model_builder\n", + "DPD_MODEL.intensity.cleanup()" ] }, { @@ -330,12 +234,14 @@ }, "outputs": [], "source": [ - "Latex(aslatex(dpd_model.amplitudes))" + "Latex(aslatex(DPD_MODEL.amplitudes))" ] }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "There is an isobar Wigner-$d$ function, which takes the following helicity angles as argument:" ] @@ -353,7 +259,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(dpd_model.variables))" + "Latex(aslatex(DPD_MODEL.variables))" ] }, { @@ -370,10 +276,11 @@ "outputs": [], "source": [ "masses = {\n", - " sp.Symbol(f\"m{i}\", nonnegative=True): round(p.mass, 7)\n", - " for p, i in PARTICLE_TO_ID.items()\n", + " symbol: sp.sympify(value)\n", + " for symbol, value in DPD_MODEL.parameter_defaults.items()\n", + " if symbol.name.startswith(\"m\")\n", + " if len(symbol.name) == 2\n", "}\n", - "dpd_model.parameter_defaults.update(masses)\n", "Latex(aslatex(masses))" ] }, @@ -402,8 +309,10 @@ "source": [ "model_builder = ampform.get_builder(REACTION)\n", "model_builder.use_helicity_couplings = False\n", - "ampform_model = model_builder.formulate()\n", - "ampform_model.intensity" + "model_builder.config.scalar_initial_state_mass = True\n", + "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", + "AMPFORM_MODEL = model_builder.formulate()\n", + "AMPFORM_MODEL.intensity" ] }, { @@ -420,7 +329,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(ampform_model.amplitudes))" + "Latex(aslatex(AMPFORM_MODEL.amplitudes))" ] }, { @@ -436,7 +345,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(ampform_model.kinematic_variables))" + "Latex(aslatex(AMPFORM_MODEL.kinematic_variables))" ] }, { @@ -453,6 +362,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "mystnb": { "code_prompt_show": "Formulate kinematic variables in terms of four-momenta" }, @@ -465,6 +377,7 @@ "p1, p2, p3 = tuple(FourMomentumSymbol(f\"p{i}\", shape=[]) for i in (0, 1, 2))\n", "s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", "mass_definitions = {\n", + " **masses,\n", " s1: InvariantMass(p2 + p3) ** 2,\n", " s2: InvariantMass(p1 + p3) ** 2,\n", " s3: InvariantMass(p1 + p2) ** 2,\n", @@ -473,16 +386,13 @@ " sp.Symbol(\"m_12\", nonnegative=True): InvariantMass(p2 + p3),\n", "}\n", "dpd_variables = {\n", - " sp.Symbol(f\"m{i}\", nonnegative=True): sp.Float(p.mass)\n", - " for i, p in enumerate(PARTICLE_TO_ID)\n", + " symbol: expr.doit().xreplace(DPD_MODEL.variables).xreplace(mass_definitions)\n", + " for symbol, expr in DPD_MODEL.variables.items()\n", "}\n", - "for symbol, expr in dpd_model.variables.items():\n", - " expr = expr.doit().xreplace(mass_definitions).xreplace(dpd_variables)\n", - " dpd_variables[symbol] = expr\n", "dpd_transformer = SympyDataTransformer.from_sympy(dpd_variables, backend=\"jax\")\n", "\n", "ampform_transformer = SympyDataTransformer.from_sympy(\n", - " ampform_model.kinematic_variables, backend=\"jax\"\n", + " AMPFORM_MODEL.kinematic_variables, backend=\"jax\"\n", ")" ] }, @@ -503,7 +413,7 @@ " return phsp_generator.generate(size, rng)\n", "\n", "\n", - "phsp = generate_phase_space(ampform_model.reaction_info, size=100_000)\n", + "phsp = generate_phase_space(AMPFORM_MODEL.reaction_info, size=100_000)\n", "ampform_phsp = ampform_transformer(phsp)\n", "dpd_phsp = dpd_transformer(phsp)" ] @@ -534,8 +444,8 @@ " return unfolded_intensity.xreplace(unfolded_amplitudes)\n", "\n", "\n", - "ampform_intensity_expr = unfold_intensity(ampform_model)\n", - "dpd_intensity_expr = unfold_intensity(dpd_model)" + "ampform_intensity_expr = unfold_intensity(AMPFORM_MODEL)\n", + "dpd_intensity_expr = unfold_intensity(DPD_MODEL)" ] }, { @@ -548,11 +458,11 @@ "source": [ "ampform_func = perform_cached_lambdify(\n", " ampform_intensity_expr,\n", - " parameters=ampform_model.parameter_defaults,\n", + " parameters=AMPFORM_MODEL.parameter_defaults,\n", ")\n", "dpd_func = perform_cached_lambdify(\n", " dpd_intensity_expr,\n", - " parameters=dpd_model.parameter_defaults,\n", + " parameters=DPD_MODEL.parameter_defaults,\n", ")" ] }, @@ -626,8 +536,8 @@ "outputs": [], "source": [ "def create_sliders() -> dict[str, ToggleButtons]:\n", - " all_parameters = dict(ampform_model.parameter_defaults.items())\n", - " all_parameters.update(dpd_model.parameter_defaults)\n", + " all_parameters = dict(AMPFORM_MODEL.parameter_defaults.items())\n", + " all_parameters.update(DPD_MODEL.parameter_defaults)\n", " sliders = {}\n", " for symbol, value in all_parameters.items():\n", " value = \"+1\"\n", @@ -712,6 +622,7 @@ "code_prompt_show": "Generate comparison widget" }, "tags": [ + "full-width", "hide-input", "scroll-input" ] @@ -728,17 +639,18 @@ " (ax_s1, ax_s2, ax_s3),\n", " (ax_t1, ax_t2, ax_t3),\n", ") = axes\n", - "final_state = ampform_model.reaction_info.final_state\n", "for ax in axes[:, 0].flatten():\n", " ax.set_ylabel(\"Intensity (a.u.)\")\n", "for ax in axes[:, 1:].flatten():\n", " ax.set_yticks([])\n", - "ax_s1.set_xlabel(f\"$m({FINAL_STATE[1].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_s2.set_xlabel(f\"$m({FINAL_STATE[0].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_s3.set_xlabel(f\"$m({FINAL_STATE[0].latex}, {FINAL_STATE[1].latex})$\")\n", - "ax_t1.set_xlabel(Rf\"$\\theta({FINAL_STATE[1].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_t2.set_xlabel(Rf\"$\\theta({FINAL_STATE[0].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_t3.set_xlabel(Rf\"$\\theta({FINAL_STATE[0].latex}, {FINAL_STATE[1].latex})$\")\n", + "\n", + "final_state = DECAY.final_state\n", + "ax_s1.set_xlabel(f\"$m({final_state[2].latex}, {final_state[3].latex})$\")\n", + "ax_s2.set_xlabel(f\"$m({final_state[1].latex}, {final_state[3].latex})$\")\n", + "ax_s3.set_xlabel(f\"$m({final_state[1].latex}, {final_state[2].latex})$\")\n", + "ax_t1.set_xlabel(Rf\"$\\theta({final_state[2].latex}, {final_state[3].latex})$\")\n", + "ax_t2.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[3].latex})$\")\n", + "ax_t3.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[2].latex})$\")\n", "fig.suptitle(f'Selected resonances: ${\", \".join(resonance_selector.value)}$')\n", "fig.tight_layout()\n", "\n", @@ -935,13 +847,13 @@ " ampform_expr = _simplify(\n", " ampform_intensity_expr,\n", " ampform_pars,\n", - " ampform_model.kinematic_variables,\n", + " AMPFORM_MODEL.kinematic_variables,\n", " selected_resonances,\n", " )\n", " dpd_expr = _simplify(\n", " dpd_intensity_expr,\n", " dpd_pars,\n", - " dpd_model.variables,\n", + " DPD_MODEL.variables,\n", " selected_resonances,\n", " )\n", " else:\n", @@ -968,26 +880,6 @@ ")\n", "display(output, ui)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if NO_TQDM:\n", - " filename = \"jpsi2phipipi-comparison.svg\"\n", - " plt.savefig(filename)\n", - " display(SVG(filename))" - ] } ], "metadata": { diff --git a/docs/comparison/jpsi2pipipi.ipynb b/docs/comparison/jpsi2pipipi.ipynb index 6c6ff454..62ffd5ac 100644 --- a/docs/comparison/jpsi2pipipi.ipynb +++ b/docs/comparison/jpsi2pipipi.ipynb @@ -2,7 +2,9 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "# J/ψ → π⁰ π⁺ π⁻\n", "\n", @@ -31,12 +33,11 @@ "source": [ "from __future__ import annotations\n", "\n", - "import itertools\n", "import logging\n", "import os\n", "import warnings\n", "from textwrap import dedent\n", - "from typing import TYPE_CHECKING, Iterable\n", + "from typing import TYPE_CHECKING\n", "\n", "import ampform\n", "import graphviz\n", @@ -48,7 +49,7 @@ "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from ampform.kinematics.lorentz import FourMomentumSymbol, InvariantMass\n", "from ampform.sympy import perform_cached_doit\n", - "from IPython.display import SVG, Latex, Markdown, clear_output, display\n", + "from IPython.display import Latex, Markdown, clear_output, display\n", "from ipywidgets import (\n", " Accordion,\n", " Checkbox,\n", @@ -66,7 +67,8 @@ "from tensorwaves.data.transform import SympyDataTransformer\n", "\n", "from ampform_dpd import DalitzPlotDecompositionBuilder\n", - "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecay, ThreeBodyDecayChain\n", + "from ampform_dpd.adapter.qrules import to_three_body_decay\n", + "from ampform_dpd.decay import Particle\n", "from ampform_dpd.io import (\n", " as_markdown_table,\n", " aslatex,\n", @@ -74,7 +76,6 @@ " perform_cached_lambdify,\n", " simplify_latex_rendering,\n", ")\n", - "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n", "\n", "if TYPE_CHECKING:\n", " from ampform.helicity import HelicityModel\n", @@ -105,11 +106,8 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "mystnb": { - "code_prompt_show": "Define initial and final state" + "code_prompt_show": "Generate transitions" }, "tags": [ "hide-input" @@ -117,26 +115,16 @@ }, "outputs": [], "source": [ - "PDG = qrules.load_pdg()\n", - "PARTICLE_DB = {\n", - " p.name: Particle(\n", - " name=p.name,\n", - " latex=p.latex,\n", - " spin=p.spin,\n", - " parity=int(p.parity),\n", - " mass=p.mass,\n", - " width=p.width,\n", - " )\n", - " for p in PDG\n", - " if p.parity is not None\n", - "}\n", - "INITIAL_STATE = PARTICLE_DB[\"J/psi(1S)\"]\n", - "FS_zero = PARTICLE_DB[\"pi0\"]\n", - "FS_neg = PARTICLE_DB[\"pi-\"]\n", - "FS_pos = PARTICLE_DB[\"pi+\"]\n", - "PARTICLE_TO_ID = {INITIAL_STATE: 0, FS_zero: 1, FS_neg: 2, FS_pos: 3}\n", - "_, *FINAL_STATE = PARTICLE_TO_ID\n", - "Markdown(as_markdown_table(list(PARTICLE_TO_ID)))" + "REACTION = qrules.generate_transitions(\n", + " initial_state=\"J/psi(1S)\",\n", + " final_state=[\"pi0\", \"pi-\", \"pi+\"],\n", + " allowed_intermediate_particles=[\"a(0)(980)\", \"rho(770)\"],\n", + " mass_conservation_factor=0,\n", + " formalism=\"helicity\",\n", + ")\n", + "REACTION123 = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", + "graphviz.Source(dot)" ] }, { @@ -146,25 +134,14 @@ "jupyter": { "source_hidden": true }, - "mystnb": { - "code_prompt_show": "Generate transitions" - }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ - "REACTION = qrules.generate_transitions(\n", - " initial_state=INITIAL_STATE.name,\n", - " final_state=[p.name for p in FINAL_STATE],\n", - " allowed_intermediate_particles=[\"a(0)(980)\", \"rho(770)\"],\n", - " mass_conservation_factor=0,\n", - " formalism=\"helicity\",\n", - ")\n", - "REACTION123 = relabel_edge_ids(REACTION)\n", - "dot = qrules.io.asdot(REACTION123, collapse_graphs=True)\n", - "graphviz.Source(dot)" + "DECAY = to_three_body_decay(REACTION123.transitions, min_ls=True)\n", + "Markdown(as_markdown_table([DECAY.initial_state, *DECAY.final_state.values()]))" ] }, { @@ -180,12 +157,11 @@ }, "outputs": [], "source": [ - "qrules_resonances = sorted(\n", - " REACTION.get_intermediate_particles(),\n", - " key=lambda p: (p.charge, p.mass, p.name),\n", + "resonances = sorted(\n", + " {t.resonance for t in DECAY.chains},\n", + " key=lambda p: (p.name[0], p.mass),\n", ")\n", - "resonance_names = [p.name for p in qrules_resonances]\n", - "resonances = [PARTICLE_DB[name] for name in resonance_names]\n", + "resonance_names = [p.name for p in resonances]\n", "Markdown(as_markdown_table(resonances))" ] }, @@ -197,84 +173,11 @@ "source_hidden": true }, "tags": [ - "hide-input", - "scroll-input" + "hide-input" ] }, "outputs": [], "source": [ - "def load_three_body_decay(\n", - " resonance_names: Iterable[str],\n", - " particle_definitions: dict[str, Particle],\n", - " min_ls: bool = True,\n", - ") -> ThreeBodyDecay:\n", - " _resonances = [particle_definitions[name] for name in resonance_names]\n", - " chains: list[ThreeBodyDecayChain] = []\n", - " for res in _resonances:\n", - " chains.extend(_create_isobar(res, min_ls))\n", - " return ThreeBodyDecay(\n", - " states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n", - " chains=tuple(chains),\n", - " )\n", - "\n", - "\n", - "def _create_isobar(resonance: Particle, min_ls: bool) -> list[ThreeBodyDecayChain]:\n", - " if resonance.name.endswith(\"-\"):\n", - " child1, child2, spectator = FS_zero, FS_neg, FS_pos\n", - " elif resonance.name.endswith(\"+\"):\n", - " child1, child2, spectator = FS_pos, FS_zero, FS_neg\n", - " else:\n", - " child1, child2, spectator = FS_neg, FS_pos, FS_zero\n", - " prod_ls_couplings = _generate_ls(\n", - " INITIAL_STATE, resonance, spectator, conserve_parity=False\n", - " )\n", - " dec_ls_couplings = _generate_ls(resonance, child1, child2, conserve_parity=True)\n", - " if min_ls:\n", - " decay = IsobarNode(\n", - " parent=INITIAL_STATE,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=min(dec_ls_couplings),\n", - " ),\n", - " child2=spectator,\n", - " interaction=min(prod_ls_couplings),\n", - " )\n", - " return [ThreeBodyDecayChain(decay)]\n", - " chains = []\n", - " for dec_ls, prod_ls in itertools.product(dec_ls_couplings, prod_ls_couplings):\n", - " decay = IsobarNode(\n", - " parent=INITIAL_STATE,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=dec_ls,\n", - " ),\n", - " child2=spectator,\n", - " interaction=prod_ls,\n", - " )\n", - " chains.append(ThreeBodyDecayChain(decay))\n", - " return chains\n", - "\n", - "\n", - "def _generate_ls(\n", - " parent: Particle, child1: Particle, child2: Particle, conserve_parity: bool\n", - ") -> list[tuple[int, sp.Rational]]:\n", - " ls = generate_ls_couplings(parent.spin, child1.spin, child2.spin)\n", - " if conserve_parity:\n", - " return filter_parity_violating_ls(\n", - " ls, parent.parity, child1.parity, child2.parity\n", - " )\n", - " return ls\n", - "\n", - "\n", - "DECAY = load_three_body_decay(\n", - " resonance_names,\n", - " particle_definitions=PARTICLE_DB,\n", - " min_ls=True,\n", - ")\n", "Latex(aslatex(DECAY, with_jp=True))" ] }, @@ -312,8 +215,9 @@ "outputs": [], "source": [ "model_builder = DalitzPlotDecompositionBuilder(DECAY, min_ls=True)\n", - "dpd_model = model_builder.formulate(reference_subsystem=1)\n", - "dpd_model.intensity" + "DPD_MODEL = model_builder.formulate(reference_subsystem=1)\n", + "del model_builder\n", + "DPD_MODEL.intensity.cleanup()" ] }, { @@ -330,12 +234,14 @@ }, "outputs": [], "source": [ - "Latex(aslatex(dpd_model.amplitudes))" + "Latex(aslatex(DPD_MODEL.amplitudes))" ] }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "There is an isobar Wigner-$d$ function, which takes the following helicity angles as argument:" ] @@ -353,7 +259,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(dpd_model.variables))" + "Latex(aslatex(DPD_MODEL.variables))" ] }, { @@ -370,10 +276,11 @@ "outputs": [], "source": [ "masses = {\n", - " sp.Symbol(f\"m{i}\", nonnegative=True): round(p.mass, 7)\n", - " for p, i in PARTICLE_TO_ID.items()\n", + " symbol: sp.sympify(value)\n", + " for symbol, value in DPD_MODEL.parameter_defaults.items()\n", + " if symbol.name.startswith(\"m\")\n", + " if len(symbol.name) == 2\n", "}\n", - "dpd_model.parameter_defaults.update(masses)\n", "Latex(aslatex(masses))" ] }, @@ -402,8 +309,10 @@ "source": [ "model_builder = ampform.get_builder(REACTION)\n", "model_builder.use_helicity_couplings = False\n", - "ampform_model = model_builder.formulate()\n", - "ampform_model.intensity" + "model_builder.config.scalar_initial_state_mass = True\n", + "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", + "AMPFORM_MODEL = model_builder.formulate()\n", + "AMPFORM_MODEL.intensity" ] }, { @@ -420,7 +329,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(ampform_model.amplitudes))" + "Latex(aslatex(AMPFORM_MODEL.amplitudes))" ] }, { @@ -436,7 +345,7 @@ }, "outputs": [], "source": [ - "Latex(aslatex(ampform_model.kinematic_variables))" + "Latex(aslatex(AMPFORM_MODEL.kinematic_variables))" ] }, { @@ -453,6 +362,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "mystnb": { "code_prompt_show": "Formulate kinematic variables in terms of four-momenta" }, @@ -465,6 +377,7 @@ "p1, p2, p3 = tuple(FourMomentumSymbol(f\"p{i}\", shape=[]) for i in (0, 1, 2))\n", "s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", "mass_definitions = {\n", + " **masses,\n", " s1: InvariantMass(p2 + p3) ** 2,\n", " s2: InvariantMass(p1 + p3) ** 2,\n", " s3: InvariantMass(p1 + p2) ** 2,\n", @@ -473,16 +386,13 @@ " sp.Symbol(\"m_12\", nonnegative=True): InvariantMass(p2 + p3),\n", "}\n", "dpd_variables = {\n", - " sp.Symbol(f\"m{i}\", nonnegative=True): sp.Float(p.mass)\n", - " for i, p in enumerate(PARTICLE_TO_ID)\n", + " symbol: expr.doit().xreplace(DPD_MODEL.variables).xreplace(mass_definitions)\n", + " for symbol, expr in DPD_MODEL.variables.items()\n", "}\n", - "for symbol, expr in dpd_model.variables.items():\n", - " expr = expr.doit().xreplace(mass_definitions).xreplace(dpd_variables)\n", - " dpd_variables[symbol] = expr\n", "dpd_transformer = SympyDataTransformer.from_sympy(dpd_variables, backend=\"jax\")\n", "\n", "ampform_transformer = SympyDataTransformer.from_sympy(\n", - " ampform_model.kinematic_variables, backend=\"jax\"\n", + " AMPFORM_MODEL.kinematic_variables, backend=\"jax\"\n", ")" ] }, @@ -503,7 +413,7 @@ " return phsp_generator.generate(size, rng)\n", "\n", "\n", - "phsp = generate_phase_space(ampform_model.reaction_info, size=100_000)\n", + "phsp = generate_phase_space(AMPFORM_MODEL.reaction_info, size=100_000)\n", "ampform_phsp = ampform_transformer(phsp)\n", "dpd_phsp = dpd_transformer(phsp)" ] @@ -534,8 +444,8 @@ " return unfolded_intensity.xreplace(unfolded_amplitudes)\n", "\n", "\n", - "ampform_intensity_expr = unfold_intensity(ampform_model)\n", - "dpd_intensity_expr = unfold_intensity(dpd_model)" + "ampform_intensity_expr = unfold_intensity(AMPFORM_MODEL)\n", + "dpd_intensity_expr = unfold_intensity(DPD_MODEL)" ] }, { @@ -548,11 +458,11 @@ "source": [ "ampform_func = perform_cached_lambdify(\n", " ampform_intensity_expr,\n", - " parameters=ampform_model.parameter_defaults,\n", + " parameters=AMPFORM_MODEL.parameter_defaults,\n", ")\n", "dpd_func = perform_cached_lambdify(\n", " dpd_intensity_expr,\n", - " parameters=dpd_model.parameter_defaults,\n", + " parameters=DPD_MODEL.parameter_defaults,\n", ")" ] }, @@ -626,8 +536,8 @@ "outputs": [], "source": [ "def create_sliders() -> dict[str, ToggleButtons]:\n", - " all_parameters = dict(ampform_model.parameter_defaults.items())\n", - " all_parameters.update(dpd_model.parameter_defaults)\n", + " all_parameters = dict(AMPFORM_MODEL.parameter_defaults.items())\n", + " all_parameters.update(DPD_MODEL.parameter_defaults)\n", " sliders = {}\n", " for symbol, value in all_parameters.items():\n", " value = \"+1\"\n", @@ -712,6 +622,7 @@ "code_prompt_show": "Generate comparison widget" }, "tags": [ + "full-width", "hide-input", "scroll-input" ] @@ -728,17 +639,18 @@ " (ax_s1, ax_s2, ax_s3),\n", " (ax_t1, ax_t2, ax_t3),\n", ") = axes\n", - "final_state = ampform_model.reaction_info.final_state\n", "for ax in axes[:, 0].flatten():\n", " ax.set_ylabel(\"Intensity (a.u.)\")\n", "for ax in axes[:, 1:].flatten():\n", " ax.set_yticks([])\n", - "ax_s1.set_xlabel(f\"$m({FINAL_STATE[1].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_s2.set_xlabel(f\"$m({FINAL_STATE[0].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_s3.set_xlabel(f\"$m({FINAL_STATE[0].latex}, {FINAL_STATE[1].latex})$\")\n", - "ax_t1.set_xlabel(Rf\"$\\theta({FINAL_STATE[1].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_t2.set_xlabel(Rf\"$\\theta({FINAL_STATE[0].latex}, {FINAL_STATE[2].latex})$\")\n", - "ax_t3.set_xlabel(Rf\"$\\theta({FINAL_STATE[0].latex}, {FINAL_STATE[1].latex})$\")\n", + "\n", + "final_state = DECAY.final_state\n", + "ax_s1.set_xlabel(f\"$m({final_state[2].latex}, {final_state[3].latex})$\")\n", + "ax_s2.set_xlabel(f\"$m({final_state[1].latex}, {final_state[3].latex})$\")\n", + "ax_s3.set_xlabel(f\"$m({final_state[1].latex}, {final_state[2].latex})$\")\n", + "ax_t1.set_xlabel(Rf\"$\\theta({final_state[2].latex}, {final_state[3].latex})$\")\n", + "ax_t2.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[3].latex})$\")\n", + "ax_t3.set_xlabel(Rf\"$\\theta({final_state[1].latex}, {final_state[2].latex})$\")\n", "fig.suptitle(f'Selected resonances: ${\", \".join(resonance_selector.value)}$')\n", "fig.tight_layout()\n", "\n", @@ -933,13 +845,13 @@ " ampform_expr = _simplify(\n", " ampform_intensity_expr,\n", " ampform_pars,\n", - " ampform_model.kinematic_variables,\n", + " AMPFORM_MODEL.kinematic_variables,\n", " selected_resonances,\n", " )\n", " dpd_expr = _simplify(\n", " dpd_intensity_expr,\n", " dpd_pars,\n", - " dpd_model.variables,\n", + " DPD_MODEL.variables,\n", " selected_resonances,\n", " )\n", " else:\n", @@ -966,26 +878,6 @@ ")\n", "display(output, ui)" ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if NO_TQDM:\n", - " filename = \"jpsi2pipipi-comparison.svg\"\n", - " plt.savefig(filename)\n", - " display(SVG(filename))" - ] } ], "metadata": { diff --git a/docs/conf.py b/docs/conf.py index 47dabb5a..9e88c084 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -36,18 +36,21 @@ add_module_names = False api_github_repo = f"{ORGANIZATION}/{REPO_NAME}" api_target_substitutions: dict[str, str | tuple[str, str]] = { - "Literal[(-1, 1)]": "typing.Literal", + "FrozenTransition": "qrules.topology.FrozenTransition", "Literal[-1, 1]": "typing.Literal", + "Literal[(-1, 1)]": "typing.Literal", "OuterStates": ("obj", "ampform_dpd.decay.OuterStates"), "ParameterValue": ("obj", "tensorwaves.interface.ParameterValue"), "ParametrizedBackendFunction": "tensorwaves.function.ParametrizedBackendFunction", "PoolSum": "ampform.sympy.PoolSum", "PositionalArgumentFunction": "tensorwaves.function.PositionalArgumentFunction", + "qrules.topology.EdgeType": "typing.TypeVar", + "qrules.topology.NodeType": "typing.TypeVar", + "sp.acos": "sympy.functions.elementary.trigonometric.acos", "sp.Expr": "sympy.core.expr.Expr", "sp.Indexed": "sympy.tensor.indexed.Indexed", "sp.Rational": "sympy.core.numbers.Rational", "sp.Symbol": "sympy.core.symbol.Symbol", - "sp.acos": "sympy.functions.elementary.trigonometric.acos", } api_target_types: dict[str, str] = {} author = "Common Partial Wave Analysis" diff --git a/docs/jpsi2ksp.ipynb b/docs/jpsi2ksp.ipynb index 7d10d2f3..9068f7ed 100644 --- a/docs/jpsi2ksp.ipynb +++ b/docs/jpsi2ksp.ipynb @@ -2,7 +2,9 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "# J/ψ → K⁰ Σ⁺ p̅\n", "\n", @@ -26,11 +28,10 @@ "source": [ "from __future__ import annotations\n", "\n", - "import itertools\n", "import logging\n", "import os\n", "import warnings\n", - "from typing import TYPE_CHECKING, Any, Iterable\n", + "from typing import TYPE_CHECKING, Any\n", "\n", "import graphviz\n", "import jax.numpy as jnp\n", @@ -45,15 +46,15 @@ "from tensorwaves.data.transform import SympyDataTransformer\n", "from tqdm.auto import tqdm\n", "\n", - "from ampform_dpd import DalitzPlotDecompositionBuilder, _get_particle\n", - "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecay, ThreeBodyDecayChain\n", + "from ampform_dpd import DalitzPlotDecompositionBuilder, get_particle\n", + "from ampform_dpd.adapter.qrules import to_three_body_decay\n", + "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecayChain\n", "from ampform_dpd.io import (\n", " as_markdown_table,\n", " aslatex,\n", " perform_cached_lambdify,\n", " simplify_latex_rendering,\n", ")\n", - "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n", "\n", "simplify_latex_rendering()\n", "logging.getLogger(\"absl\").setLevel(logging.ERROR) # mute JAX\n", @@ -97,7 +98,7 @@ " initial_state=\"J/psi(1S)\",\n", " final_state=[\"K0\", \"Sigma+\", \"p~\"],\n", " allowed_interaction_types=\"strong\",\n", - " formalism=\"helicity\", # AmpForm-DPD already generates LS couplings\n", + " formalism=\"canonical-helicity\",\n", " mass_conservation_factor=0.05,\n", ")\n", "REACTION = relabel_edge_ids(REACTION)\n", @@ -109,31 +110,17 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ - "PDG = qrules.load_pdg()\n", - "PARTICLE_DB = {\n", - " p.name: Particle(\n", - " name=p.name,\n", - " latex=p.latex,\n", - " spin=p.spin,\n", - " parity=int(p.parity),\n", - " mass=p.mass,\n", - " width=p.width,\n", - " )\n", - " for p in PDG\n", - " if p.parity is not None\n", - "}\n", - "Jpsi = PARTICLE_DB[\"J/psi(1S)\"]\n", - "K = PARTICLE_DB[\"K0\"]\n", - "Σ = PARTICLE_DB[\"Sigma+\"]\n", - "pbar = PARTICLE_DB[\"p~\"]\n", - "PARTICLE_TO_ID = {Jpsi: 0, K: 1, Σ: 2, pbar: 3}\n", - "Markdown(as_markdown_table(list(PARTICLE_TO_ID)))" + "DECAY = to_three_body_decay(REACTION.transitions, min_ls=True)\n", + "Markdown(as_markdown_table([DECAY.initial_state, *DECAY.final_state.values()]))" ] }, { @@ -149,17 +136,11 @@ }, "outputs": [], "source": [ - "resonance_names = [\n", - " \"Sigma(1660)~-\",\n", - " \"Sigma(1670)~-\",\n", - " \"Sigma(1750)~-\",\n", - " \"Sigma(1775)~-\",\n", - " \"Sigma(1910)~-\",\n", - " \"N(1700)+\",\n", - " \"N(1710)+\",\n", - " \"N(1720)+\",\n", - "]\n", - "resonances = sorted(PARTICLE_DB[name] for name in resonance_names)\n", + "resonances = sorted(\n", + " {t.resonance for t in DECAY.chains},\n", + " key=lambda p: (p.name[0], p.mass),\n", + ")\n", + "resonance_names = [p.name for p in resonances]\n", "Markdown(as_markdown_table(resonances))" ] }, @@ -167,87 +148,12 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ - "def load_three_body_decay(\n", - " resonance_names: Iterable[str],\n", - " particle_definitions: dict[str, Particle],\n", - " min_ls: bool = True,\n", - ") -> ThreeBodyDecay:\n", - " resonances = [particle_definitions[name] for name in resonance_names]\n", - " chains: list[ThreeBodyDecayChain] = []\n", - " for res in resonances:\n", - " chains.extend(_create_isobar(res, min_ls))\n", - " return ThreeBodyDecay(\n", - " states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n", - " chains=tuple(sorted(chains)),\n", - " )\n", - "\n", - "\n", - "def _create_isobar(resonance: Particle, min_ls: bool) -> list[ThreeBodyDecayChain]:\n", - " if resonance.name.startswith(\"Sigma\"):\n", - " child1, child2, spectator = pbar, K, Σ\n", - " elif resonance.name.startswith(\"N\"):\n", - " child1, child2, spectator = K, Σ, pbar\n", - " elif resonance.name.startswith(\"K\"):\n", - " child1, child2, spectator = Σ, pbar, K\n", - " else:\n", - " raise NotImplementedError\n", - " prod_ls_couplings = _generate_ls(Jpsi, resonance, spectator, conserve_parity=True)\n", - " dec_ls_couplings = _generate_ls(resonance, child1, child2, conserve_parity=True)\n", - " if min_ls:\n", - " decay = IsobarNode(\n", - " parent=Jpsi,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=min(dec_ls_couplings),\n", - " ),\n", - " child2=spectator,\n", - " interaction=min(prod_ls_couplings),\n", - " )\n", - " return [ThreeBodyDecayChain(decay)]\n", - " chains = []\n", - " for dec_ls, prod_ls in itertools.product(dec_ls_couplings, prod_ls_couplings):\n", - " decay = IsobarNode(\n", - " parent=Jpsi,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=dec_ls,\n", - " ),\n", - " child2=spectator,\n", - " interaction=prod_ls,\n", - " )\n", - " chains.append(ThreeBodyDecayChain(decay))\n", - " return chains\n", - "\n", - "\n", - "def _generate_ls(\n", - " parent: Particle, child1: Particle, child2: Particle, conserve_parity: bool\n", - ") -> list[tuple[int, sp.Rational]]:\n", - " ls = generate_ls_couplings(parent.spin, child1.spin, child2.spin)\n", - " if conserve_parity:\n", - " return filter_parity_violating_ls(\n", - " ls, parent.parity, child1.parity, child2.parity\n", - " )\n", - " return ls\n", - "\n", - "\n", - "DECAY = load_three_body_decay(\n", - " resonance_names,\n", - " particle_definitions=PARTICLE_DB,\n", - " min_ls=True,\n", - ")\n", "Latex(aslatex(DECAY, with_jp=True))" ] }, @@ -347,7 +253,7 @@ " angular_momentum: Any\n", " meson_radius: Any\n", "\n", - " _latex_repr_ = R\"\\mathcal{{F}}_{{{angular_momentum}}}\\left({s}, {m1}, {m1}\\right)\"\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", @@ -391,14 +297,16 @@ " production_node = decay_chain.decay\n", " assert isinstance(production_node.child1, IsobarNode), \"Not a 3-body isobar decay\"\n", " decay_node = production_node.child1\n", - "\n", " s = _get_mandelstam_s(decay_chain)\n", " parameter_defaults = {}\n", " production_ff, new_pars = _create_form_factor(s, production_node)\n", + " _is_messed_up(new_pars)\n", " parameter_defaults.update(new_pars)\n", " decay_ff, new_pars = _create_form_factor(s, decay_node)\n", + " _is_messed_up(new_pars)\n", " parameter_defaults.update(new_pars)\n", " breit_wigner, new_pars = _create_breit_wigner(s, decay_node)\n", + " _is_messed_up(new_pars)\n", " parameter_defaults.update(new_pars)\n", " return (\n", " production_ff * decay_ff * breit_wigner,\n", @@ -406,11 +314,20 @@ " )\n", "\n", "\n", + "def _is_messed_up(parameter_defaults):\n", + " for key in parameter_defaults:\n", + " if key.name == \"m0\" and not key.is_nonnegative:\n", + " raise ValueError\n", + "\n", + "\n", "def _create_form_factor(\n", " s: sp.Symbol, isobar: IsobarNode\n", ") -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", " assert isobar.interaction is not None, \"Need LS-couplings\"\n", - " inv_mass = _create_mass_symbol(isobar.parent)\n", + " if isobar.parent.name == \"J/psi(1S)\":\n", + " inv_mass = sp.Symbol(\"m0\", nonnegative=True)\n", + " else:\n", + " inv_mass = _get_mandelstam_s(isobar)\n", " outgoing_state_mass1 = _create_mass_symbol(isobar.child1)\n", " outgoing_state_mass2 = _create_mass_symbol(isobar.child2)\n", " meson_radius = _create_meson_radius_symbol(isobar.parent)\n", @@ -423,7 +340,11 @@ " )\n", " parameter_defaults = {\n", " meson_radius: 1,\n", + " outgoing_state_mass1: get_particle(isobar.child1).mass,\n", + " outgoing_state_mass2: get_particle(isobar.child2).mass,\n", " }\n", + " if not inv_mass.name.startswith(\"s\"):\n", + " parameter_defaults[inv_mass] = get_particle(isobar).mass\n", " return form_factor, parameter_defaults\n", "\n", "\n", @@ -456,31 +377,35 @@ "\n", "\n", "def _create_meson_radius_symbol(isobar: IsobarNode) -> sp.Symbol:\n", - " if _get_particle(isobar) is Jpsi:\n", + " if get_particle(isobar).name == \"J/psi(1S)\":\n", " return sp.Symbol(R\"R_{J/\\psi}\")\n", " return sp.Symbol(R\"R_\\mathrm{res}\")\n", "\n", "\n", "def _create_mass_symbol(particle: IsobarNode | Particle) -> sp.Symbol:\n", - " particle = _get_particle(particle)\n", - " state_id = PARTICLE_TO_ID.get(particle)\n", - " if state_id is not None:\n", - " return sp.Symbol(f\"m{state_id}\", nonnegative=True)\n", + " particle = get_particle(particle)\n", " return sp.Symbol(f\"m_{{{particle.latex}}}\", nonnegative=True)\n", "\n", "\n", - "def _get_mandelstam_s(decay: ThreeBodyDecayChain) -> sp.Symbol:\n", + "def _get_mandelstam_s(decay: ThreeBodyDecayChain | IsobarNode) -> sp.Symbol:\n", " s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", - " m1, m2, m3 = map(_create_mass_symbol, [K, Σ, pbar])\n", - " decay_masses = {_create_mass_symbol(p) for p in decay.decay_products}\n", - " if decay_masses == {m2, m3}:\n", + " decay_products = {p.name for p in _get_decay_products(decay)}\n", + " if decay_products == {\"Sigma+\", \"p~\"}:\n", " return s1\n", - " if decay_masses == {m1, m3}:\n", + " if decay_products == {\"K0\", \"p~\"}:\n", " return s2\n", - " if decay_masses == {m1, m2}:\n", + " if decay_products == {\"K0\", \"Sigma+\"}:\n", " return s3\n", - " msg = f\"Cannot find Mandelstam variable for {''.join(decay_masses)}\"\n", - " raise NotImplementedError(msg)" + " msg = f\"Cannot find Mandelstam variable for {', '.join(decay_products)}\"\n", + " raise NotImplementedError(msg)\n", + "\n", + "\n", + "def _get_decay_products(\n", + " decay: ThreeBodyDecayChain | IsobarNode,\n", + ") -> tuple[Particle, Particle]:\n", + " if isinstance(decay, ThreeBodyDecayChain):\n", + " return decay.decay_products\n", + " return decay.children" ] }, { @@ -539,29 +464,6 @@ "Latex(aslatex(model.variables))" ] }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "masses = {\n", - " _create_mass_symbol(Jpsi): Jpsi.mass,\n", - " _create_mass_symbol(K): K.mass,\n", - " _create_mass_symbol(Σ): Σ.mass,\n", - " _create_mass_symbol(pbar): pbar.mass,\n", - "}\n", - "model.parameter_defaults.update(masses)\n", - "Latex(aslatex(masses))" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -573,9 +475,6 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [ "hide-input", "full-width", @@ -697,6 +596,28 @@ "First, the data transformer defined above expects values for the masses. We have already defined these values above, but we need to convert them from {mod}`sympy` objects to numerical data:" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "masses = {\n", + " symbol: value\n", + " for symbol, value in model.parameter_defaults.items()\n", + " if str(symbol).startswith(\"m\")\n", + " if len(str(symbol)) == 2\n", + "}\n", + "Latex(aslatex(masses))" + ] + }, { "cell_type": "code", "execution_count": null, @@ -749,7 +670,7 @@ "outputs": [], "source": [ "s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", - "m0, m1, m2, m3 = sorted(masses, key=str)\n", + "m0, m1, m2, m3 = sp.symbols(\"m:4\", nonnegative=True)\n", "s1_expr = compute_third_mandelstam(s3, s2, m0, m1, m2, m3)\n", "Latex(aslatex({s1: s1_expr}))" ] @@ -933,14 +854,14 @@ "i1, i2 = 0, 0\n", "for chain in tqdm(model.decay.chains, disable=NO_TQDM):\n", " resonance = chain.resonance\n", - " decay_product = set(chain.decay_products)\n", - " if decay_product == {K, Σ}:\n", + " decay_product = {p.name for p in chain.decay_products}\n", + " if decay_product == {\"K0\", \"Sigma+\"}:\n", " ax = ax1\n", " i1 += 1\n", " i = i1\n", " projection_axis = 0\n", " x_data = x\n", - " elif decay_product == {K, pbar}:\n", + " elif decay_product == {\"K0\", \"p~\"}:\n", " ax = ax2\n", " i2 += 1\n", " i = i2\n", diff --git a/docs/lc2pkpi.ipynb b/docs/lc2pkpi.ipynb index b3a41cbf..e2c1721f 100644 --- a/docs/lc2pkpi.ipynb +++ b/docs/lc2pkpi.ipynb @@ -26,20 +26,22 @@ "source": [ "from __future__ import annotations\n", "\n", - "import itertools\n", - "from typing import Iterable\n", + "import warnings\n", "\n", + "import graphviz\n", "import qrules\n", "import sympy as sp\n", + "from ampform.helicity.align.dpd import relabel_edge_ids\n", "from IPython.display import Latex, Markdown\n", "\n", - "from ampform_dpd import DalitzPlotDecompositionBuilder\n", - "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecay, ThreeBodyDecayChain\n", + "from ampform_dpd import DalitzPlotDecompositionBuilder, get_particle\n", + "from ampform_dpd.adapter.qrules import load_particles, to_three_body_decay\n", + "from ampform_dpd.decay import IsobarNode, Particle, ThreeBodyDecayChain\n", "from ampform_dpd.dynamics import BreitWignerMinL\n", "from ampform_dpd.io import as_markdown_table, aslatex, simplify_latex_rendering\n", - "from ampform_dpd.spin import filter_parity_violating_ls, generate_ls_couplings\n", "\n", - "simplify_latex_rendering()" + "simplify_latex_rendering()\n", + "warnings.simplefilter(\"ignore\")" ] }, { @@ -62,33 +64,35 @@ }, "outputs": [], "source": [ - "PDG = qrules.load_pdg()\n", - "PARTICLE_DB = {\n", - " p.name: Particle(\n", - " name=p.name,\n", - " latex=p.latex,\n", - " spin=p.spin,\n", - " parity=int(p.parity),\n", - " mass=p.mass,\n", - " width=p.width,\n", - " )\n", - " for p in PDG\n", - " if p.parity is not None\n", - "}\n", - "PARTICLE_DB[\"Lambda(2000)\"] = Particle(\n", - " name=\"Lambda(2000)\",\n", - " latex=R\"\\Lambda(2000)\",\n", - " spin=0.5,\n", - " parity=-1,\n", - " mass=2.0,\n", - " width=(0.020 + 0.400) / 2,\n", + "PARTICLES = load_particles()\n", + "for name in [\n", + " \"K*(1410)~0\",\n", + " \"K(2)*(1430)~0\",\n", + " \"K*(1680)~0\",\n", + " \"Delta(1620)++\",\n", + " \"Delta(1900)++\",\n", + " \"Delta(1910)++\",\n", + " \"Delta(1920)++\",\n", + " \"Lambda(1800)\",\n", + " \"Lambda(1810)\",\n", + " \"Lambda(1890)\",\n", + "]:\n", + " PARTICLES.remove(PARTICLES[name])\n", + "STM = qrules.StateTransitionManager(\n", + " initial_state=[\"Lambda(c)+\"],\n", + " final_state=[\"p\", \"K-\", \"pi+\"],\n", + " mass_conservation_factor=3,\n", + " allowed_intermediate_particles=[\"K\", \"Delta\", \"Lambda\"],\n", + " particle_db=PARTICLES,\n", + " max_angular_momentum=2,\n", + " formalism=\"canonical-helicity\",\n", ")\n", - "Λc = PARTICLE_DB[\"Lambda(c)+\"]\n", - "p = PARTICLE_DB[\"p\"]\n", - "π = PARTICLE_DB[\"pi+\"]\n", - "K = PARTICLE_DB[\"K-\"]\n", - "PARTICLE_TO_ID = {Λc: 0, p: 1, π: 2, K: 3}\n", - "Markdown(as_markdown_table(list(PARTICLE_TO_ID)))" + "STM.set_allowed_interaction_types([qrules.InteractionType.STRONG], node_id=1)\n", + "problem_sets = STM.create_problem_sets()\n", + "REACTION = STM.find_solutions(problem_sets)\n", + "REACTION = relabel_edge_ids(REACTION)\n", + "dot = qrules.io.asdot(REACTION, collapse_graphs=True)\n", + "graphviz.Source(dot)" ] }, { @@ -104,22 +108,8 @@ }, "outputs": [], "source": [ - "resonance_names = [\n", - " \"Lambda(1405)\",\n", - " \"Lambda(1520)\",\n", - " \"Lambda(1600)\",\n", - " \"Lambda(1670)\",\n", - " \"Lambda(1690)\",\n", - " \"Lambda(2000)\",\n", - " \"Delta(1232)+\",\n", - " \"Delta(1600)+\",\n", - " \"Delta(1700)+\",\n", - " \"K(0)*(700)+\",\n", - " \"K*(892)0\",\n", - " \"K(2)*(1430)0\",\n", - "]\n", - "resonances = [PARTICLE_DB[name] for name in resonance_names]\n", - "Markdown(as_markdown_table(resonances))" + "DECAY = to_three_body_decay(REACTION.transitions, min_ls=True)\n", + "Markdown(as_markdown_table([DECAY.initial_state, *DECAY.final_state.values()]))" ] }, { @@ -135,78 +125,24 @@ }, "outputs": [], "source": [ - "def load_three_body_decay(\n", - " resonance_names: Iterable[str],\n", - " particle_definitions: dict[str, Particle],\n", - " min_ls: bool = True,\n", - ") -> ThreeBodyDecay:\n", - " resonances = [particle_definitions[name] for name in resonance_names]\n", - " chains: list[ThreeBodyDecayChain] = []\n", - " for res in resonances:\n", - " chains.extend(_create_isobar(res, min_ls))\n", - " return ThreeBodyDecay(\n", - " states={state_id: particle for particle, state_id in PARTICLE_TO_ID.items()},\n", - " chains=tuple(chains),\n", - " )\n", - "\n", - "\n", - "def _create_isobar(resonance: Particle, min_ls: bool) -> list[ThreeBodyDecayChain]:\n", - " if resonance.name.startswith(\"K\"):\n", - " child1, child2, spectator = π, K, p\n", - " elif resonance.name.startswith(\"L\"):\n", - " child1, child2, spectator = K, p, π\n", - " elif resonance.name.startswith(\"D\"):\n", - " child1, child2, spectator = p, π, K\n", - " else:\n", - " raise NotImplementedError\n", - " prod_ls_couplings = _generate_ls(Λc, resonance, spectator, conserve_parity=False)\n", - " dec_ls_couplings = _generate_ls(resonance, child1, child2, conserve_parity=True)\n", - " if min_ls:\n", - " decay = IsobarNode(\n", - " parent=Λc,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=min(dec_ls_couplings),\n", - " ),\n", - " child2=spectator,\n", - " interaction=min(prod_ls_couplings),\n", - " )\n", - " return [ThreeBodyDecayChain(decay)]\n", - " chains = []\n", - " for dec_ls, prod_ls in itertools.product(dec_ls_couplings, prod_ls_couplings):\n", - " decay = IsobarNode(\n", - " parent=Λc,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=child1,\n", - " child2=child2,\n", - " interaction=dec_ls,\n", - " ),\n", - " child2=spectator,\n", - " interaction=prod_ls,\n", - " )\n", - " chains.append(ThreeBodyDecayChain(decay))\n", - " return chains\n", - "\n", - "\n", - "def _generate_ls(\n", - " parent: Particle, child1: Particle, child2: Particle, conserve_parity: bool\n", - ") -> list[tuple[int, sp.Rational]]:\n", - " ls = generate_ls_couplings(parent.spin, child1.spin, child2.spin)\n", - " if conserve_parity:\n", - " return filter_parity_violating_ls(\n", - " ls, parent.parity, child1.parity, child2.parity\n", - " )\n", - " return ls\n", - "\n", - "\n", - "DECAY = load_three_body_decay(\n", - " resonance_names,\n", - " particle_definitions=PARTICLE_DB,\n", - " min_ls=True,\n", + "resonances = sorted(\n", + " {t.resonance for t in DECAY.chains},\n", + " key=lambda p: (p.name[0], p.mass),\n", ")\n", + "resonance_names = [p.name for p in resonances]\n", + "Markdown(as_markdown_table(resonances))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ "Latex(aslatex(DECAY, with_jp=True))" ] }, @@ -256,13 +192,15 @@ " decay_chain: ThreeBodyDecayChain,\n", ") -> tuple[BreitWignerMinL, dict[sp.Symbol, float]]:\n", " s = _get_mandelstam_s(decay_chain)\n", - " child1_mass, child2_mass = map(_to_mass_symbol, decay_chain.decay_products)\n", + " child1_mass, child2_mass = map(_create_mass_symbol, decay_chain.decay_products)\n", " l_dec = sp.Rational(decay_chain.outgoing_ls.L)\n", " l_prod = sp.Rational(decay_chain.incoming_ls.L)\n", - " parent_mass = sp.Symbol(f\"m_{{{decay_chain.parent.latex}}}\")\n", - " spectator_mass = sp.Symbol(f\"m_{{{decay_chain.spectator.latex}}}\")\n", - " resonance_mass = sp.Symbol(f\"m_{{{decay_chain.resonance.latex}}}\")\n", - " resonance_width = sp.Symbol(Rf\"\\Gamma_{{{decay_chain.resonance.latex}}}\")\n", + " parent_mass = sp.Symbol(f\"m_{{{decay_chain.parent.latex}}}\", nonnegative=True)\n", + " spectator_mass = sp.Symbol(f\"m_{{{decay_chain.spectator.latex}}}\", nonnegative=True)\n", + " resonance_mass = sp.Symbol(f\"m_{{{decay_chain.resonance.latex}}}\", nonnegative=True)\n", + " resonance_width = sp.Symbol(\n", + " Rf\"\\Gamma_{{{decay_chain.resonance.latex}}}\", nonnegative=True\n", + " )\n", " R_dec = sp.Symbol(R\"R_\\mathrm{res}\")\n", " R_prod = sp.Symbol(R\"R_{\\Lambda_c}\")\n", " parameter_defaults = {\n", @@ -292,25 +230,28 @@ " return dynamics, parameter_defaults\n", "\n", "\n", - "def _get_mandelstam_s(decay: ThreeBodyDecayChain) -> sp.Symbol:\n", + "def _create_mass_symbol(particle: IsobarNode | Particle) -> sp.Symbol:\n", + " particle = get_particle(particle)\n", + " return sp.Symbol(f\"m_{{{particle.latex}}}\", nonnegative=True)\n", + "\n", + "\n", + "def _get_mandelstam_s(decay: ThreeBodyDecayChain | IsobarNode) -> sp.Symbol:\n", " s1, s2, s3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", - " m1, m2, m3 = map(_to_mass_symbol, [p, π, K])\n", - " decay_masses = {_to_mass_symbol(p) for p in decay.decay_products}\n", - " if decay_masses == {m2, m3}:\n", + " decay_products = {p.name for p in _get_decay_products(decay)}\n", + " if decay_products == {\"p\", \"pi+\"}:\n", " return s1\n", - " if decay_masses == {m1, m3}:\n", + " if decay_products == {\"pi+\", \"K-\"}:\n", " return s2\n", - " if decay_masses == {m1, m2}:\n", + " if decay_products == {\"K-\", \"p\"}:\n", " return s3\n", - " msg = f\"Cannot find Mandelstam variable for {''.join(decay_masses)}\"\n", + " msg = f\"Cannot find Mandelstam variable for {', '.join(decay_products)}\"\n", " raise NotImplementedError(msg)\n", "\n", "\n", - "def _to_mass_symbol(particle: Particle) -> sp.Symbol:\n", - " state_id = PARTICLE_TO_ID.get(particle)\n", - " if state_id is not None:\n", - " return sp.Symbol(f\"m{state_id}\", nonnegative=True)\n", - " return sp.Symbol(f\"m_{{{particle.latex}}}\", nonnegative=True)" + "def _get_decay_products(decay: ThreeBodyDecayChain | IsobarNode) -> list[Particle]:\n", + " if isinstance(decay, ThreeBodyDecayChain):\n", + " return decay.decay_products\n", + " return decay.children" ] }, { @@ -371,7 +312,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.17" + "version": "3.9.19" } }, "nbformat": 4,