From bb13735c814ec71e501b92dc564c1cb1ca41d163 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:28 +0200 Subject: [PATCH 01/42] MAINT: import `Lc2ppiK.json` from RUB-EP1/amplitude-serialization@334de9f https://raw.githubusercontent.com/RUB-EP1/amplitude-serialization/334de9f/Lc2ppiK.json --- .cspell.json | 1 + docs/.gitignore | 1 + docs/Lc2ppiK.json | 978 ++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 980 insertions(+) create mode 100644 docs/Lc2ppiK.json diff --git a/.cspell.json b/.cspell.json index ddaae994..b9bd5089 100644 --- a/.cspell.json +++ b/.cspell.json @@ -58,6 +58,7 @@ "eqref", "figsize", "fontsize", + "formfactor", "gcov", "graphviz", "ipykernel", diff --git a/docs/.gitignore b/docs/.gitignore index 028cb435..53f8e968 100644 --- a/docs/.gitignore +++ b/docs/.gitignore @@ -2,3 +2,4 @@ *.json *.zip /api/ +!*.json diff --git a/docs/Lc2ppiK.json b/docs/Lc2ppiK.json new file mode 100644 index 00000000..d4d79b32 --- /dev/null +++ b/docs/Lc2ppiK.json @@ -0,0 +1,978 @@ +{ + "distributions": [ + { + "name": "my_model_for_reaction_intensity", + "type": "hadronic_cross_section_unpolarized_dist", + "decay_description": { + "kinematics": { + "initial_state": { + "index": 0, + "name": "Lc", + "spin": "1/2", + "mass": 2.28646 + }, + "final_state": [ + { + "index": 1, + "name": "p", + "spin": "1/2", + "mass": 0.938272046 + }, + { + "index": 2, + "name": "pi", + "spin": "0", + "mass": 0.13957018 + }, + { + "index": 3, + "name": "K", + "spin": "0", + "mass": 0.493677 + } + ] + }, + "reference_topology": [[3, 1], 2], + "chains": [ + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L1405_Flatte" + } + ], + "weight": "7.38649400481717 + 1.971018433257411i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "+", + "node": [3, 1], + "formfactor": "" + } + ], + "topology": [[3, 1], 2], + "name": "L1405" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L1405_Flatte" + } + ], + "weight": "-3.2332358574805515 + 2.2557724553615772i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "+", + "node": [3, 1], + "formfactor": "" + } + ], + "topology": [[3, 1], 2], + "name": "L1405" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [3, 1], + "parametrization": "L1520_BW" + } + ], + "weight": "0.146999 + 0.022162i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "-", + "node": [3, 1], + "formfactor": "BlattWeisskopf_resonance_l2" + } + ], + "topology": [[3, 1], 2], + "name": "L1520" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [3, 1], + "parametrization": "L1520_BW" + } + ], + "weight": "-0.0803435 + 0.7494165i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "-", + "node": [3, 1], + "formfactor": "BlattWeisskopf_resonance_l2" + } + ], + "topology": [[3, 1], 2], + "name": "L1520" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L1600_BW" + } + ], + "weight": "4.929406127531439 - 0.5956915012088891i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "-", + "node": [3, 1], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[3, 1], 2], + "name": "L1600" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L1600_BW" + } + ], + "weight": "-3.4228557332438796 - 2.179858885546952i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "-", + "node": [3, 1], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[3, 1], 2], + "name": "L1600" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L1670_BW" + } + ], + "weight": "-0.24012285628923374 - 0.10230279488850731i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "+", + "node": [3, 1], + "formfactor": "" + } + ], + "topology": [[3, 1], 2], + "name": "L1670" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L1670_BW" + } + ], + "weight": "-0.40374241570833247 + 0.7154739757283278i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "+", + "node": [3, 1], + "formfactor": "" + } + ], + "topology": [[3, 1], 2], + "name": "L1670" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [3, 1], + "parametrization": "L1690_BW" + } + ], + "weight": "-0.192886 - 0.0551175i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "-", + "node": [3, 1], + "formfactor": "BlattWeisskopf_resonance_l2" + } + ], + "topology": [[3, 1], 2], + "name": "L1690" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [3, 1], + "parametrization": "L1690_BW" + } + ], + "weight": "-1.365296 - 0.1768065i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "-", + "node": [3, 1], + "formfactor": "BlattWeisskopf_resonance_l2" + } + ], + "topology": [[3, 1], 2], + "name": "L1690" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L2000_BW" + } + ], + "weight": "-3.0661953154540726 - 2.684313105886122i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "+", + "node": [3, 1], + "formfactor": "" + } + ], + "topology": [[3, 1], 2], + "name": "L2000" + }, + { + "propagators": [ + { + "spin": "1/2", + "node": [3, 1], + "parametrization": "L2000_BW" + } + ], + "weight": "-5.667359734940468 - 5.38391527459506i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[3, 1], 2], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "1/2"], + "parity_factor": "+", + "node": [3, 1], + "formfactor": "" + } + ], + "topology": [[3, 1], 2], + "name": "L2000" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [1, 2], + "parametrization": "D(1232)_BW" + } + ], + "weight": "-3.3890955 + 1.5259025i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[1, 2], 3], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["1/2", "0"], + "parity_factor": "+", + "node": [1, 2], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[1, 2], 3], + "name": "D(1232)" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [1, 2], + "parametrization": "D(1232)_BW" + } + ], + "weight": "-6.4935965 + 2.264168i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[1, 2], 3], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["1/2", "0"], + "parity_factor": "+", + "node": [1, 2], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[1, 2], 3], + "name": "D(1232)" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [1, 2], + "parametrization": "D(1600)_BW" + } + ], + "weight": "5.7007925 - 1.5627555i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[1, 2], 3], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["1/2", "0"], + "parity_factor": "+", + "node": [1, 2], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[1, 2], 3], + "name": "D(1600)" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [1, 2], + "parametrization": "D(1600)_BW" + } + ], + "weight": "3.3646055 - 0.5011915i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[1, 2], 3], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["1/2", "0"], + "parity_factor": "+", + "node": [1, 2], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[1, 2], 3], + "name": "D(1600)" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [1, 2], + "parametrization": "D(1700)_BW" + } + ], + "weight": "-5.18914 - 0.717436i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1/2", "0"], + "node": [[1, 2], 3], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["1/2", "0"], + "parity_factor": "-", + "node": [1, 2], + "formfactor": "BlattWeisskopf_resonance_l2" + } + ], + "topology": [[1, 2], 3], + "name": "D(1700)" + }, + { + "propagators": [ + { + "spin": "3/2", + "node": [1, 2], + "parametrization": "D(1700)_BW" + } + ], + "weight": "-6.437051 - 1.052785i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1/2", "0"], + "node": [[1, 2], 3], + "formfactor": "BlattWeisskopf_b_decay_l1" + }, + { + "type": "parity", + "helicities": ["1/2", "0"], + "parity_factor": "-", + "node": [1, 2], + "formfactor": "BlattWeisskopf_resonance_l2" + } + ], + "topology": [[1, 2], 3], + "name": "D(1700)" + }, + { + "propagators": [ + { + "spin": "0", + "node": [2, 3], + "parametrization": "K700_BuggBW" + } + ], + "weight": "0.068908 + 2.521444i", + "vertices": [ + { + "type": "helicity", + "helicities": ["0", "1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "" + } + ], + "topology": [[2, 3], 1], + "name": "K(700)" + }, + { + "propagators": [ + { + "spin": "0", + "node": [2, 3], + "parametrization": "K700_BuggBW" + } + ], + "weight": "-2.68563 + 0.03849i", + "vertices": [ + { + "type": "helicity", + "helicities": ["0", "-1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "" + } + ], + "topology": [[2, 3], 1], + "name": "K(700)" + }, + { + "propagators": [ + { + "spin": "1", + "node": [2, 3], + "parametrization": "K892_BW" + } + ], + "weight": "0.6885560139393164 - 0.5922539890384868i", + "vertices": [ + { + "type": "helicity", + "helicities": ["-1", "-1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[2, 3], 1], + "name": "K(892)" + }, + { + "propagators": [ + { + "spin": "1", + "node": [2, 3], + "parametrization": "K892_BW" + } + ], + "weight": "-0.4198173614898905 - 2.398905956940163i", + "vertices": [ + { + "type": "helicity", + "helicities": ["0", "1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[2, 3], 1], + "name": "K(892)" + }, + { + "propagators": [ + { + "spin": "1", + "node": [2, 3], + "parametrization": "K892_BW" + } + ], + "weight": "0.5773502691896258 + 0.0i", + "vertices": [ + { + "type": "helicity", + "helicities": ["0", "-1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[2, 3], 1], + "name": "K(892)" + }, + { + "propagators": [ + { + "spin": "1", + "node": [2, 3], + "parametrization": "K892_BW" + } + ], + "weight": "-1.8137146937446735 - 1.9014511500518056i", + "vertices": [ + { + "type": "helicity", + "helicities": ["1", "1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "BlattWeisskopf_resonance_l1" + } + ], + "topology": [[2, 3], 1], + "name": "K(892)" + }, + { + "propagators": [ + { + "spin": "0", + "node": [2, 3], + "parametrization": "K1430_BuggBW" + } + ], + "weight": "-6.71516 + 10.479411i", + "vertices": [ + { + "type": "helicity", + "helicities": ["0", "1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "" + } + ], + "topology": [[2, 3], 1], + "name": "K(1430)" + }, + { + "propagators": [ + { + "spin": "0", + "node": [2, 3], + "parametrization": "K1430_BuggBW" + } + ], + "weight": "0.219754 + 8.741196i", + "vertices": [ + { + "type": "helicity", + "helicities": ["0", "-1/2"], + "node": [[2, 3], 1], + "formfactor": "" + }, + { + "type": "parity", + "helicities": ["0", "0"], + "parity_factor": "+", + "node": [2, 3], + "formfactor": "" + } + ], + "topology": [[2, 3], 1], + "name": "K(1430)" + } + ], + "appendix": {} + }, + "variables": [ + { + "node": [1, 2], + "mass_angles": ["m_12", "cos_theta_12", "phi_12"] + }, + { + "node": [[1, 2], 3], + "mass_angles": ["m_12_3", "cos_theta_12_3", "phi_12_3"] + } + ], + "parameters": [] + } + ], + "functions": [ + { + "name": "L1405_Flatte", + "type": "BreitWigner", + "mass": 1.4051, + "channels": [ + { + "gsq": 0.23395150538434703, + "ma": 0.938272046, + "mb": 0.493677, + "l": 0, + "d": 0 + }, + { + "gsq": 0.23395150538434703, + "ma": 1.18937, + "mb": 0.13957018, + "l": 0, + "d": 0 + } + ] + }, + { + "name": "L1690_BW", + "l": 2, + "mb": 0.938272046, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.69, + "ma": 0.493677, + "width": 0.07 + }, + { + "name": "D(1232)_BW", + "l": 1, + "mb": 0.13957018, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.232, + "ma": 0.938272046, + "width": 0.117 + }, + { + "name": "L1520_BW", + "l": 2, + "mb": 0.938272046, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.518467, + "ma": 0.493677, + "width": 0.015195 + }, + { + "name": "L1600_BW", + "l": 1, + "mb": 0.938272046, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.63, + "ma": 0.493677, + "width": 0.25 + }, + { + "name": "L2000_BW", + "l": 0, + "mb": 0.938272046, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.98819, + "ma": 0.493677, + "width": 0.17926 + }, + { + "name": "D(1600)_BW", + "l": 1, + "mb": 0.13957018, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.64, + "ma": 0.938272046, + "width": 0.3 + }, + { + "name": "D(1700)_BW", + "l": 2, + "mb": 0.13957018, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.69, + "ma": 0.938272046, + "width": 0.38 + }, + { + "name": "K892_BW", + "l": 1, + "mb": 0.493677, + "type": "BreitWigner", + "d": 1.5, + "mass": 0.8955, + "ma": 0.13957018, + "width": 0.047299999999999995 + }, + { + "name": "K700_BuggBW", + "slope": 0.94106, + "type": "BreitWignerWidthExp", + "mass": 0.824, + "width": 0.478 + }, + { + "name": "K1430_BuggBW", + "slope": 0.020981, + "type": "BreitWignerWidthExp", + "mass": 1.375, + "width": 0.19 + }, + { + "name": "L1670_BW", + "l": 0, + "mb": 0.938272046, + "type": "BreitWigner", + "d": 1.5, + "mass": 1.67, + "ma": 0.493677, + "width": 0.03 + }, + { + "name": "BlattWeisskopf_resonance_l1", + "type": "BlattWeisskopf", + "radius": 1.5, + "l": 1 + }, + { + "name": "BlattWeisskopf_resonance_l2", + "type": "BlattWeisskopf", + "radius": 1.5, + "l": 2 + }, + { + "name": "BlattWeisskopf_b_decay_l1", + "type": "BlattWeisskopf", + "radius": 5.0, + "l": 1 + } + ], + "domains": [ + { + "name": "default", + "type": "product_domain", + "axes": [ + { + "name": "cos_theta_12", + "min": -1.0, + "max": 1.0 + }, + { + "name": "phi_12", + "min": -3.14, + "max": 3.14 + }, + { + "name": "m_12", + "min": 1.0, + "max": 5.5 + }, + { + "name": "cos_theta_12_3", + "min": -1.0, + "max": 1.0 + }, + { + "name": "phi_12_3", + "min": -3.14, + "max": 3.14 + }, + { + "name": "m_12_3", + "min": 1.0, + "max": 5.5 + } + ] + } + ], + "misc": { + "amplitude_model_checksums": [ + { + "point": "validation_point", + "distribution": "my_model_for_reaction_intensity", + "value": 1.3 + } + ] + }, + "parameter_points": [ + { + "name": "validation_point", + "parameters": [ + { + "name": "cos_theta_12", + "value": 0.23678498926410071 + }, + { + "name": "phi_12", + "value": 0.0 + }, + { + "name": "m_12", + "value": 1.7058127512193626 + }, + { + "name": "cos_theta_12_3", + "value": 0.0 + }, + { + "name": "phi_12_3", + "value": 0.0 + }, + { + "name": "m_12_3", + "value": 2.28646 + } + ] + } + ] +} From d4f456f50cc1ac2f9deae16acd20fd2faa6c8cbd Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:32 +0200 Subject: [PATCH 02/42] DOC: add skeleton serialization notebook --- .pre-commit-config.yaml | 1 + docs/index.md | 1 + docs/serialization.ipynb | 74 ++++++++++++++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+) create mode 100644 docs/serialization.ipynb diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 054ed876..5e9fa5c0 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,6 +21,7 @@ repos: - --extra-keys - | cell.attachments + cell.id cell.metadata.code_folding cell.metadata.editable cell.metadata.id diff --git a/docs/index.md b/docs/index.md index d46e898b..17defe25 100644 --- a/docs/index.md +++ b/docs/index.md @@ -171,6 +171,7 @@ maxdepth: 1 lc2pkpi jpsi2ksp xib2pkk +serialization ``` ```{toctree} diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb new file mode 100644 index 00000000..bc6f5e40 --- /dev/null +++ b/docs/serialization.ipynb @@ -0,0 +1,74 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Model serialization" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Import model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "import json\n", + "\n", + "with open(\"Lc2ppiK.json\") as stream:\n", + " model_definition = json.load(stream)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "remove-cell" + ] + }, + "outputs": [], + "source": [ + "from IPython.display import JSON\n", + "\n", + "JSON(model_definition)" + ] + } + ], + "metadata": { + "colab": { + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.19" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From 02c2befaa97ab63550a7634356cf04dbc889dac6 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:33 +0200 Subject: [PATCH 03/42] DOC: load kinematics section --- docs/serialization.ipynb | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index bc6f5e40..cec8f8c4 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -25,7 +25,7 @@ "import json\n", "\n", "with open(\"Lc2ppiK.json\") as stream:\n", - " model_definition = json.load(stream)" + " model_def = json.load(stream)" ] }, { @@ -43,7 +43,27 @@ "source": [ "from IPython.display import JSON\n", "\n", - "JSON(model_definition)" + "JSON(model_def)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "kinematics_def = model_def[\"distributions\"][0][\"decay_description\"][\"kinematics\"]\n", + "kinematics_def" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "initial_state_def = kinematics_def[\"initial_state\"]\n", + "final_state_def = kinematics_def[\"final_state\"]" ] } ], From 4b38ac911e0dfcd1ce57a3715cf6187ad81e3f36 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:34 +0200 Subject: [PATCH 04/42] ENH: support undefined `Particle.parity` --- src/ampform_dpd/decay.py | 2 +- src/ampform_dpd/io.py | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/decay.py b/src/ampform_dpd/decay.py index d545a5e1..813097ce 100644 --- a/src/ampform_dpd/decay.py +++ b/src/ampform_dpd/decay.py @@ -30,7 +30,7 @@ class Particle: name: str latex: str spin: sp.Rational = field(converter=to_rational, validator=assert_spin_value) - parity: Literal[-1, 1] + parity: Literal[-1, 1] | None mass: float width: float diff --git a/src/ampform_dpd/io.py b/src/ampform_dpd/io.py index 038d0f4d..b34aa60f 100644 --- a/src/ampform_dpd/io.py +++ b/src/ampform_dpd/io.py @@ -87,11 +87,13 @@ def _(obj: Particle, with_jp: bool = False, only_jp: bool = False, **kwargs) -> def _render_jp(particle: Particle) -> str: - parity = "-" if particle.parity < 0 else "+" if particle.spin.denominator == 1: spin = sp.latex(particle.spin) else: spin = Rf"\frac{{{particle.spin.numerator}}}{{{particle.spin.denominator}}}" + if particle.parity is None: + return f"J={spin}" + parity = "-" if particle.parity < 0 else "+" return f"{spin}^{parity}" From 6f5502205a6f4c1d357985d2ba315157b61c6563 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:35 +0200 Subject: [PATCH 05/42] ENH: hide LS columns if undefined --- src/ampform_dpd/io.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/io.py b/src/ampform_dpd/io.py index b34aa60f..f4c517e1 100644 --- a/src/ampform_dpd/io.py +++ b/src/ampform_dpd/io.py @@ -178,9 +178,11 @@ def _as_decay_markdown_table(decay_chains: Sequence[ThreeBodyDecayChain]) -> str R"$J^P$", R"mass (MeV)", R"width (MeV)", - R"$L_\mathrm{dec}^\mathrm{min}$", - R"$L_\mathrm{prod}^\mathrm{min}$", ] + if any(c.outgoing_ls is not None for c in decay_chains): + column_names.append(R"$L_\mathrm{dec}^\mathrm{min}$") + if any(c.incoming_ls is not None for c in decay_chains): + column_names.append(R"$L_\mathrm{prod}^\mathrm{min}$") src = _create_markdown_table_header(column_names) for chain in decay_chains: child1, child2 = map(aslatex, chain.decay_products) From 0be4c7328c77b47c002961b537a2e3d1e3d3445a Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:35 +0200 Subject: [PATCH 06/42] DOC: define `to_decay()` function for deserialization --- .cspell.json | 1 + docs/serialization.ipynb | 188 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 182 insertions(+), 7 deletions(-) diff --git a/.cspell.json b/.cspell.json index b9bd5089..6639a411 100644 --- a/.cspell.json +++ b/.cspell.json @@ -85,6 +85,7 @@ "nbformat", "ncols", "ndarray", + "noqa", "noreply", "nrows", "pandoc", diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index cec8f8c4..d823e1d2 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -9,11 +9,26 @@ }, { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "## Import model" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "from __future__ import annotations" + ] + }, { "cell_type": "code", "execution_count": null, @@ -46,24 +61,183 @@ "JSON(model_def)" ] }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Construct `ThreeBodyDecay`" + ] + }, { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-cell", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "from typing import Callable\n", + "from warnings import warn\n", + "\n", + "from ampform_dpd.decay import (\n", + " FinalStateID,\n", + " IsobarNode,\n", + " Particle,\n", + " State,\n", + " ThreeBodyDecay,\n", + " ThreeBodyDecayChain,\n", + ")\n", + "\n", + "\n", + "def to_decay(\n", + " definition: dict, to_latex: Callable[[str], str] | None = None\n", + ") -> ThreeBodyDecay:\n", + " distributions_def = model_def[\"distributions\"]\n", + " n_distributions = len(distributions_def)\n", + " if n_distributions == 0:\n", + " msg = \"The serialized model does not have any distributions\"\n", + " raise ValueError(msg)\n", + " if n_distributions > 1:\n", + " msg = f\"There are {n_distributions} distributions, but expecting one only\"\n", + " warn(msg, category=UserWarning)\n", + " decay_description = distributions_def[0][\"decay_description\"]\n", + " initial_state, final_state = _get_states(decay_description[\"kinematics\"])\n", + " return ThreeBodyDecay(\n", + " states={\n", + " initial_state.index: initial_state,\n", + " **final_state,\n", + " },\n", + " chains=sorted({\n", + " to_decay_chain(chain, initial_state, final_state, to_latex)\n", + " for chain in decay_description[\"chains\"]\n", + " }),\n", + " )\n", + "\n", + "\n", + "def to_decay_chain(\n", + " chain_def: dict,\n", + " initial_state: State,\n", + " final_state: dict[FinalStateID, State],\n", + " to_latex: Callable[[str], str] | None = None,\n", + ") -> ThreeBodyDecayChain:\n", + " vertices = chain_def[\"vertices\"]\n", + " if to_latex is None:\n", + " to_latex = lambda x: x # noqa:E731\n", + " resonance = Particle(\n", + " name=chain_def[\"name\"],\n", + " latex=name_to_latex(chain_def[\"name\"]),\n", + " spin=chain_def[\"propagators\"][0][\"spin\"],\n", + " mass=0,\n", + " width=0,\n", + " parity=None,\n", + " )\n", + " return ThreeBodyDecayChain(\n", + " decay=IsobarNode(\n", + " parent=initial_state,\n", + " child1=IsobarNode(\n", + " parent=resonance,\n", + " child1=final_state[vertices[1][\"node\"][0]],\n", + " child2=final_state[vertices[1][\"node\"][1]],\n", + " ),\n", + " child2=final_state[vertices[0][\"node\"][1]],\n", + " )\n", + " )\n", + "\n", + "\n", + "def _get_states(kinematics: dict) -> tuple[State, dict[FinalStateID, State]]:\n", + " initial_state_def = kinematics[\"initial_state\"]\n", + " final_state_def = kinematics[\"final_state\"]\n", + " initial_state = dict_to_particle(initial_state_def)\n", + " final_state = {p[\"index\"]: dict_to_particle(p) for p in final_state_def}\n", + " return initial_state, final_state\n", + "\n", + "\n", + "def dict_to_particle(dct: dict) -> State:\n", + " return State(\n", + " name=dct[\"name\"],\n", + " latex=name_to_latex(dct[\"name\"]),\n", + " mass=dct[\"mass\"],\n", + " width=0,\n", + " spin=dct[\"spin\"],\n", + " parity=None,\n", + " index=dct[\"index\"],\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, "outputs": [], "source": [ - "kinematics_def = model_def[\"distributions\"][0][\"decay_description\"][\"kinematics\"]\n", - "kinematics_def" + "def name_to_latex(name: str) -> str:\n", + " latex = {\n", + " \"Lc\": R\"\\Lambda_c^+\",\n", + " \"pi\": R\"\\pi^+\",\n", + " \"K\": \"K^-\",\n", + " \"p\": \"p\",\n", + " }.get(name)\n", + " if latex is not None:\n", + " return latex\n", + " mass_str = name[1:].strip(\"(\").strip(\")\")\n", + " subsystem_letter = name[0]\n", + " subsystem = {\"D\": \"D\", \"K\": \"K\", \"L\": R\"\\Lambda\"}.get(subsystem_letter)\n", + " if subsystem is None:\n", + " return name\n", + " return f\"{subsystem}({mass_str})\"\n", + "\n", + "\n", + "decay = to_decay(model_def, to_latex=name_to_latex)" ] }, { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, "outputs": [], "source": [ - "initial_state_def = kinematics_def[\"initial_state\"]\n", - "final_state_def = kinematics_def[\"final_state\"]" + "from IPython.display import Math\n", + "\n", + "from ampform_dpd.io import aslatex\n", + "\n", + "Math(aslatex(decay))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "from IPython.display import Markdown\n", + "\n", + "from ampform_dpd.io import as_markdown_table\n", + "\n", + "Markdown(as_markdown_table(decay))" ] } ], From 4a871e3d906e95ba65001e7cbc7c5c21f83db520 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:36 +0200 Subject: [PATCH 07/42] MAINT: update model definition from RUB-EP1/amplitude-serialization@8c06a81 --- docs/Lc2ppiK.json | 82 +++++++++++++++++++++++------------------------ 1 file changed, 41 insertions(+), 41 deletions(-) diff --git a/docs/Lc2ppiK.json b/docs/Lc2ppiK.json index d4d79b32..45d74168 100644 --- a/docs/Lc2ppiK.json +++ b/docs/Lc2ppiK.json @@ -2,7 +2,7 @@ "distributions": [ { "name": "my_model_for_reaction_intensity", - "type": "hadronic_cross_section_unpolarized_dist", + "type": "HadronicUnpolarizedIntensity", "decay_description": { "kinematics": { "initial_state": { @@ -363,7 +363,7 @@ { "spin": "3/2", "node": [1, 2], - "parametrization": "D(1232)_BW" + "parametrization": "D1232_BW" } ], "weight": "-3.3890955 + 1.5259025i", @@ -383,14 +383,14 @@ } ], "topology": [[1, 2], 3], - "name": "D(1232)" + "name": "D1232" }, { "propagators": [ { "spin": "3/2", "node": [1, 2], - "parametrization": "D(1232)_BW" + "parametrization": "D1232_BW" } ], "weight": "-6.4935965 + 2.264168i", @@ -410,14 +410,14 @@ } ], "topology": [[1, 2], 3], - "name": "D(1232)" + "name": "D1232" }, { "propagators": [ { "spin": "3/2", "node": [1, 2], - "parametrization": "D(1600)_BW" + "parametrization": "D1600_BW" } ], "weight": "5.7007925 - 1.5627555i", @@ -437,14 +437,14 @@ } ], "topology": [[1, 2], 3], - "name": "D(1600)" + "name": "D1600" }, { "propagators": [ { "spin": "3/2", "node": [1, 2], - "parametrization": "D(1600)_BW" + "parametrization": "D1600_BW" } ], "weight": "3.3646055 - 0.5011915i", @@ -464,14 +464,14 @@ } ], "topology": [[1, 2], 3], - "name": "D(1600)" + "name": "D1600" }, { "propagators": [ { "spin": "3/2", "node": [1, 2], - "parametrization": "D(1700)_BW" + "parametrization": "D1700_BW" } ], "weight": "-5.18914 - 0.717436i", @@ -491,14 +491,14 @@ } ], "topology": [[1, 2], 3], - "name": "D(1700)" + "name": "D1700" }, { "propagators": [ { "spin": "3/2", "node": [1, 2], - "parametrization": "D(1700)_BW" + "parametrization": "D1700_BW" } ], "weight": "-6.437051 - 1.052785i", @@ -518,7 +518,7 @@ } ], "topology": [[1, 2], 3], - "name": "D(1700)" + "name": "D1700" }, { "propagators": [ @@ -545,7 +545,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(700)" + "name": "K700" }, { "propagators": [ @@ -572,7 +572,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(700)" + "name": "K700" }, { "propagators": [ @@ -599,7 +599,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(892)" + "name": "K892" }, { "propagators": [ @@ -626,7 +626,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(892)" + "name": "K892" }, { "propagators": [ @@ -653,7 +653,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(892)" + "name": "K892" }, { "propagators": [ @@ -680,7 +680,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(892)" + "name": "K892" }, { "propagators": [ @@ -707,7 +707,7 @@ } ], "topology": [[2, 3], 1], - "name": "K(1430)" + "name": "K1430" }, { "propagators": [ @@ -734,19 +734,19 @@ } ], "topology": [[2, 3], 1], - "name": "K(1430)" + "name": "K1430" } ], "appendix": {} }, "variables": [ { - "node": [1, 2], - "mass_angles": ["m_12", "cos_theta_12", "phi_12"] + "node": [3, 1], + "mass_angles": ["m_31", "phi_31", "cos_theta_31"] }, { - "node": [[1, 2], 3], - "mass_angles": ["m_12_3", "cos_theta_12_3", "phi_12_3"] + "node": [[3, 1], 2], + "mass_angles": ["m_31_2", "phi_31_2", "cos_theta_31_2"] } ], "parameters": [] @@ -755,7 +755,7 @@ "functions": [ { "name": "L1405_Flatte", - "type": "BreitWigner", + "type": "MultichannelBreitWigner", "mass": 1.4051, "channels": [ { @@ -785,7 +785,7 @@ "width": 0.07 }, { - "name": "D(1232)_BW", + "name": "D1232_BW", "l": 1, "mb": 0.13957018, "type": "BreitWigner", @@ -825,7 +825,7 @@ "width": 0.17926 }, { - "name": "D(1600)_BW", + "name": "D1600_BW", "l": 1, "mb": 0.13957018, "type": "BreitWigner", @@ -835,7 +835,7 @@ "width": 0.3 }, { - "name": "D(1700)_BW", + "name": "D1700_BW", "l": 2, "mb": 0.13957018, "type": "BreitWigner", @@ -857,14 +857,14 @@ { "name": "K700_BuggBW", "slope": 0.94106, - "type": "BreitWignerWidthExp", + "type": "BreitWignerWidthExpLikeBugg", "mass": 0.824, "width": 0.478 }, { "name": "K1430_BuggBW", "slope": 0.020981, - "type": "BreitWignerWidthExp", + "type": "BreitWignerWidthExpLikeBugg", "mass": 1.375, "width": 0.19 }, @@ -938,9 +938,9 @@ "misc": { "amplitude_model_checksums": [ { - "point": "validation_point", + "name": "validation_point", "distribution": "my_model_for_reaction_intensity", - "value": 1.3 + "value": 9345.853380852355 } ] }, @@ -949,27 +949,27 @@ "name": "validation_point", "parameters": [ { - "name": "cos_theta_12", - "value": 0.23678498926410071 + "name": "cos_theta_31", + "value": -0.2309352648098208 }, { - "name": "phi_12", + "name": "phi_31", "value": 0.0 }, { - "name": "m_12", - "value": 1.7058127512193626 + "name": "m_31", + "value": 1.9101377207489973 }, { - "name": "cos_theta_12_3", + "name": "cos_theta_31_2", "value": 0.0 }, { - "name": "phi_12_3", + "name": "phi_31_2", "value": 0.0 }, { - "name": "m_12_3", + "name": "m_31_2", "value": 2.28646 } ] From 08e27bcf33b647d0def9bd4453a4b35f56235689 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:37 +0200 Subject: [PATCH 08/42] FEAT: define `DefinedExpression` struct --- src/ampform_dpd/__init__.py | 28 ++++++++++++++++++++++++++-- 1 file changed, 26 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/__init__.py b/src/ampform_dpd/__init__.py index 66363c8c..df691389 100644 --- a/src/ampform_dpd/__init__.py +++ b/src/ampform_dpd/__init__.py @@ -2,15 +2,16 @@ from __future__ import annotations +from collections import abc from functools import lru_cache from itertools import product -from typing import Literal, Protocol +from typing import Any, Literal, Protocol from warnings import warn import sympy as sp from ampform.kinematics.phasespace import compute_third_mandelstam from ampform.sympy import PoolSum -from attrs import field, frozen +from attrs import define, field, frozen from sympy.core.symbol import Str from sympy.physics.quantum.spin import CG, WignerD from sympy.physics.quantum.spin import Rotation as Wigner @@ -396,6 +397,29 @@ def __call__( ) -> tuple[sp.Expr, dict[sp.Symbol, float | complex]]: ... +@define +class DefinedExpression: + expression: sp.Expr = sp.S.One + definitions: dict[sp.Symbol, complex | float] = field(factory=dict) + + def __mul__(self, other: Any) -> DefinedExpression: + if isinstance(other, DefinedExpression): + return DefinedExpression( + expression=self.expression * other.expression, + definitions={**self.definitions, **other.definitions}, + ) + if isinstance(other, abc.Sequence) and len(other) == 2: # noqa: PLR2004 + expression, definitions = other + return DefinedExpression( + expression=self.expression * expression, + definitions={**self.definitions, **definitions}, + ) + return DefinedExpression( + expression=self.expression * other, + definitions=self.definitions, + ) + + def formulate_non_resonant( decay_chain: ThreeBodyDecayChain, ) -> tuple[sp.Expr, dict[sp.Symbol, float | complex]]: From 93cb5cd934ec6783bd5027b3f5eab37ae77780f1 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:38 +0200 Subject: [PATCH 09/42] BREAK: remove `ampform_dpd.simplify_latex_rendering()` alias --- src/ampform_dpd/__init__.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/ampform_dpd/__init__.py b/src/ampform_dpd/__init__.py index df691389..0f28efcd 100644 --- a/src/ampform_dpd/__init__.py +++ b/src/ampform_dpd/__init__.py @@ -29,9 +29,6 @@ get_decay_product_ids, to_particle, ) -from ampform_dpd.io import ( - simplify_latex_rendering, # noqa: F401 # pyright:ignore[reportUnusedImport] -) from ampform_dpd.spin import create_spin_range From 6a7e31d634acce46b3b2d2b83c39b12e085cde5d Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:39 +0200 Subject: [PATCH 10/42] DOC: build amplitudes with dynamics --- .cspell.json | 5 + docs/serialization.ipynb | 1232 ++++++++++++++++++++++++++++++++++++-- pyproject.toml | 3 + 3 files changed, 1187 insertions(+), 53 deletions(-) diff --git a/.cspell.json b/.cspell.json index 6639a411..be509934 100644 --- a/.cspell.json +++ b/.cspell.json @@ -38,6 +38,7 @@ "PyPI", "absl", "arange", + "asdict", "asdot", "aslatex", "autoscale", @@ -51,6 +52,7 @@ "colorbar", "commitlint", "concat", + "difflib", "docnb", "elif", "endswith", @@ -61,6 +63,7 @@ "formfactor", "gcov", "graphviz", + "imag", "ipykernel", "ipynb", "ipython", @@ -105,6 +108,7 @@ "savefig", "sdist", "seealso", + "setattr", "setuptools", "sharey", "simplefilter", @@ -158,6 +162,7 @@ "pytest", "PYTHONHASHSEED", "QRules", + "recoupling", "sympify", "SymPy", "TensorWaves", diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index d823e1d2..9631ca89 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -20,13 +20,51 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "tags": [ - "remove-input" + "hide-cell" ] }, "outputs": [], "source": [ - "from __future__ import annotations" + "from __future__ import annotations\n", + "\n", + "import json\n", + "from collections import abc\n", + "from difflib import get_close_matches\n", + "from typing import Any, Callable, Sequence, Union\n", + "from warnings import warn\n", + "\n", + "import sympy as sp\n", + "from ampform.dynamics import BlattWeisskopfSquared\n", + "from ampform.kinematics.phasespace import Kallen\n", + "from ampform.sympy import argument, unevaluated\n", + "from attrs import asdict, frozen\n", + "from IPython.display import JSON, Markdown, Math\n", + "from sympy.core.symbol import Str\n", + "from sympy.physics.quantum.spin import Rotation as Wigner\n", + "from sympy.printing.latex import LatexPrinter\n", + "\n", + "from ampform_dpd import (\n", + " DefinedExpression,\n", + " _create_coupling_symbol, # noqa:PLC2701\n", + ")\n", + "from ampform_dpd.angles import formulate_scattering_angle\n", + "from ampform_dpd.decay import (\n", + " FinalStateID,\n", + " IsobarNode,\n", + " Particle,\n", + " State,\n", + " StateID,\n", + " ThreeBodyDecay,\n", + " ThreeBodyDecayChain,\n", + ")\n", + "from ampform_dpd.dynamics import BuggBreitWigner, EnergyDependentWidth, FormFactor, P\n", + "from ampform_dpd.io import as_markdown_table, aslatex, simplify_latex_rendering\n", + "\n", + "simplify_latex_rendering()" ] }, { @@ -37,33 +75,27 @@ }, "outputs": [], "source": [ - "import json\n", - "\n", "with open(\"Lc2ppiK.json\") as stream:\n", - " model_def = json.load(stream)" + " MODEL_DEFINITION = json.load(stream)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [ "remove-cell" ] }, "outputs": [], "source": [ - "from IPython.display import JSON\n", - "\n", - "JSON(model_def)" + "JSON(MODEL_DEFINITION)" ] }, { "cell_type": "markdown", "metadata": { + "jp-MarkdownHeadingCollapsed": true, "tags": [] }, "source": [ @@ -84,42 +116,35 @@ }, "outputs": [], "source": [ - "from typing import Callable\n", - "from warnings import warn\n", + "def to_decay(\n", + " model: dict, to_latex: Callable[[str], str] | None = None\n", + ") -> ThreeBodyDecay:\n", + " initial_state = _get_initial_state(model)\n", + " final_state = _get_final_state(model)\n", + " return ThreeBodyDecay(\n", + " states=_get_states(model),\n", + " chains=sorted({\n", + " to_decay_chain(chain, initial_state, final_state, to_latex)\n", + " for chain in _get_decay_chains(model)\n", + " }),\n", + " )\n", "\n", - "from ampform_dpd.decay import (\n", - " FinalStateID,\n", - " IsobarNode,\n", - " Particle,\n", - " State,\n", - " ThreeBodyDecay,\n", - " ThreeBodyDecayChain,\n", - ")\n", "\n", + "def _get_decay_chains(model: dict) -> list[dict]:\n", + " distribution_def = _get_distribution_def(model)\n", + " return distribution_def[\"decay_description\"][\"chains\"]\n", "\n", - "def to_decay(\n", - " definition: dict, to_latex: Callable[[str], str] | None = None\n", - ") -> ThreeBodyDecay:\n", - " distributions_def = model_def[\"distributions\"]\n", - " n_distributions = len(distributions_def)\n", + "\n", + "def _get_distribution_def(model: dict) -> dict:\n", + " distribution_defs = MODEL_DEFINITION[\"distributions\"]\n", + " n_distributions = len(distribution_defs)\n", " if n_distributions == 0:\n", " msg = \"The serialized model does not have any distributions\"\n", " raise ValueError(msg)\n", " if n_distributions > 1:\n", " msg = f\"There are {n_distributions} distributions, but expecting one only\"\n", " warn(msg, category=UserWarning)\n", - " decay_description = distributions_def[0][\"decay_description\"]\n", - " initial_state, final_state = _get_states(decay_description[\"kinematics\"])\n", - " return ThreeBodyDecay(\n", - " states={\n", - " initial_state.index: initial_state,\n", - " **final_state,\n", - " },\n", - " chains=sorted({\n", - " to_decay_chain(chain, initial_state, final_state, to_latex)\n", - " for chain in decay_description[\"chains\"]\n", - " }),\n", - " )\n", + " return distribution_defs[0]\n", "\n", "\n", "def to_decay_chain(\n", @@ -152,12 +177,25 @@ " )\n", "\n", "\n", - "def _get_states(kinematics: dict) -> tuple[State, dict[FinalStateID, State]]:\n", - " initial_state_def = kinematics[\"initial_state\"]\n", + "def _get_states(model: dict) -> dict[StateID, State]:\n", + " initial_state = _get_initial_state(model)\n", + " final_state = _get_final_state(model)\n", + " return {initial_state.index: initial_state, **final_state}\n", + "\n", + "\n", + "def _get_initial_state(model: dict) -> dict[StateID, State]:\n", + " distribution_def = _get_distribution_def(model)\n", + " decay_description = distribution_def[\"decay_description\"]\n", + " kinematics = decay_description[\"kinematics\"]\n", + " return dict_to_particle(kinematics[\"initial_state\"])\n", + "\n", + "\n", + "def _get_final_state(model: dict) -> dict[StateID, State]:\n", + " distribution_def = _get_distribution_def(model)\n", + " decay_description = distribution_def[\"decay_description\"]\n", + " kinematics = decay_description[\"kinematics\"]\n", " final_state_def = kinematics[\"final_state\"]\n", - " initial_state = dict_to_particle(initial_state_def)\n", - " final_state = {p[\"index\"]: dict_to_particle(p) for p in final_state_def}\n", - " return initial_state, final_state\n", + " return {p[\"index\"]: dict_to_particle(p) for p in final_state_def}\n", "\n", "\n", "def dict_to_particle(dct: dict) -> State:\n", @@ -176,6 +214,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "tags": [] }, "outputs": [], @@ -197,27 +238,242 @@ " return f\"{subsystem}({mass_str})\"\n", "\n", "\n", - "decay = to_decay(model_def, to_latex=name_to_latex)" + "DECAY = to_decay(MODEL_DEFINITION, to_latex=name_to_latex)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "Math(aslatex(DECAY))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ - "from IPython.display import Math\n", + "Markdown(as_markdown_table(DECAY))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Dynamics" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{seealso} [RUB-EP1/amplitude-serialization#22](https://github.com/RUB-EP1/amplitude-serialization/issues/22)\n", + ":::" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, + "source": [ + "### Function look-up mechanism" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "function_defs = MODEL_DEFINITION[\"functions\"]\n", + "{f[\"type\"] for f in function_defs}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "CHAIN_DEFINITIONS = _get_decay_chains(MODEL_DEFINITION)\n", + "CHAIN_DEFINITIONS[2]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def get_function_definition(function_name: str, model: dict) -> dict:\n", + " function_definitions = model[\"functions\"]\n", + " for function_def in function_definitions:\n", + " if function_def[\"name\"] == function_name:\n", + " return function_def\n", + " existing_names = {f[\"name\"] for f in function_definitions}\n", + " msg = f\"Could not find function with name {function_name!r}.\"\n", + " candidates = get_close_matches(function_name, existing_names)\n", + " if candidates:\n", + " msg += f\" Did you mean any of these? {', '.join(sorted(candidates))}\"\n", + " raise KeyError(msg)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "raises-exception" + ] + }, + "outputs": [], + "source": [ + "get_function_definition(\"BlattWeiskopf\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "### Vertices" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Blatt-Weisskopf form factor" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "get_function_definition(\"BlattWeisskopf_resonance_l1\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "get_function_definition(\"BlattWeisskopf_resonance_l2\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _node_to_mass(node_item: int | Sequence[int]) -> sp.Symbol:\n", + " if isinstance(node_item, int):\n", + " return sp.Symbol(f\"m{node_item}\", nonnegative=True)\n", + " if (\n", + " isinstance(node_item, abc.Sequence)\n", + " and all(isinstance(i, int) for i in node_item)\n", + " and len(node_item) == 2\n", + " ):\n", + " ij = \"\".join(str(i) for i in node_item)\n", + " return sp.Symbol(f\"s{ij}\", nonnegative=True)\n", + " msg = f\"Cannot create mass symbol for node {node_item}\"\n", + " raise NotImplementedError(msg)\n", + "\n", + "\n", + "sp.Tuple(_node_to_mass([2, 3]), _node_to_mass(1))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _node_to_mandelstam(node: Topology) -> sp.Symbol:\n", + " if all(isinstance(i, int) for i in node):\n", + " return _node_to_mass(node)\n", + " return _node_to_mass(0)\n", + "\n", + "\n", + "sp.Tuple(_node_to_mandelstam([2, 3]), _node_to_mandelstam([[2, 3], 1]))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "def formulate_form_factor(vertex: dict, model: dict) -> DefinedExpression:\n", + " function_name = vertex.get(\"formfactor\")\n", + " if not function_name:\n", + " return DefinedExpression()\n", + " function_definition = get_function_definition(function_name, model)\n", + " function_type = function_definition[\"type\"]\n", + " if function_type == \"BlattWeisskopf\":\n", + " node = vertex[\"node\"]\n", + " s = _node_to_mandelstam(node)\n", + " m1, m2 = (_node_to_mass(i) for i in node)\n", + " if all(isinstance(i, int) for i in node):\n", + " meson_radius = sp.Symbol(R\"R_\\mathrm{res}\")\n", + " else:\n", + " initial_state = _get_initial_state(model)\n", + " meson_radius = sp.Symbol(f\"R_{{{initial_state.latex}}}\")\n", + " angular_momentum = int(function_definition[\"l\"])\n", + " return DefinedExpression(\n", + " expression=FormFactor(s, m1, m2, angular_momentum, meson_radius),\n", + " definitions={\n", + " meson_radius: function_definition[\"radius\"],\n", + " },\n", + " )\n", + " msg = f\"No form factor implementation for {function_name!r}\"\n", + " raise NotImplementedError(msg)\n", "\n", - "from ampform_dpd.io import aslatex\n", "\n", - "Math(aslatex(decay))" + "CHAIN_2 = CHAIN_DEFINITIONS[2]\n", + "Math(aslatex(formulate_form_factor(CHAIN_2[\"vertices\"][0], MODEL_DEFINITION)))" ] }, { @@ -233,11 +489,881 @@ }, "outputs": [], "source": [ - "from IPython.display import Markdown\n", + "s, m1, m2, d, L, z = sp.symbols(\"s m1 m2 d L z\")\n", + "exprs = [\n", + " FormFactor(s, m1, m2, L, d),\n", + " BlattWeisskopfSquared(z, L),\n", + "]\n", + "Math(aslatex({e: e.evaluate() for e in exprs}))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Propagators" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Breit-Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated\n", + "class BreitWigner(sp.Expr):\n", + " s: Any\n", + " mass: Any\n", + " width: Any\n", + " m1: Any = 0\n", + " m2: Any = 0\n", + " angular_momentum: Any = 0\n", + " meson_radius: Any = 1\n", + " _latex_repr_ = R\"\\mathcal{{R}}^\\mathrm{{BW}}_{{{angular_momentum}}}\\left({s}\\right)\" # noqa:RUF027\n", + "\n", + " def evaluate(self):\n", + " width = self.energy_dependent_width()\n", + " return BreitWigner(self.s, self.mass, width)\n", "\n", - "from ampform_dpd.io import as_markdown_table\n", + " def energy_dependent_width(self) -> sp.Expr:\n", + " s, m0, Γ0, m1, m2, L, d = self.args\n", + " if L == 0 and m1 == 0 and m2 == 0:\n", + " return Γ0\n", + " return EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d)\n", "\n", - "Markdown(as_markdown_table(decay))" + " def _latex_repr_(self, printer: LatexPrinter, *args) -> str:\n", + " s = printer._print(self.s)\n", + " function_symbol = R\"\\mathcal{R}^\\mathrm{BW}\"\n", + " arg = Rf\"\\left({s}\\right)\"\n", + " if self.angular_momentum == 0:\n", + " if self.m1 == 0 and self.m2 == 0:\n", + " width = printer._print(self.energy_dependent_width())\n", + " arg = Rf\"\\left({s}; {width}\\right)\"\n", + " return Rf\"{function_symbol}{arg}\"\n", + " L = printer._print(self.angular_momentum)\n", + " return Rf\"{function_symbol}_{{{L}}}{arg}\"\n", + "\n", + "\n", + "x, y, z = sp.symbols(\"x:z\")\n", + "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", + "exprs = [\n", + " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", + " BreitWigner(s, m0, Γ0),\n", + " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", + " FormFactor(s, m1, m2, L, d),\n", + " P(s, m1, m2),\n", + " Kallen(x, y, z),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CHAIN_DEFINITIONS[20]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "get_function_definition(\"K892_BW\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def _formulate_breit_wigner(\n", + " propagator: dict, function_definition: dict, resonance: str, **kwargs\n", + ") -> DefinedExpression:\n", + " node = propagator[\"node\"]\n", + " i, j = node\n", + " s = _node_to_mandelstam(node)\n", + " mass = sp.Symbol(f\"m_{{{resonance}}}\")\n", + " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\")\n", + " m1 = _node_to_mass(i)\n", + " m2 = _node_to_mass(j)\n", + " angular_momentum = int(function_definition[\"l\"])\n", + " d = sp.Symbol(R\"R_\\mathrm{res}\")\n", + " return DefinedExpression(\n", + " expression=BreitWigner(s, mass, width, m1, m2, angular_momentum, d),\n", + " definitions={\n", + " mass: function_definition[\"mass\"],\n", + " width: function_definition[\"width\"],\n", + " m1: function_definition[\"ma\"],\n", + " m2: function_definition[\"mb\"],\n", + " d: function_definition[\"d\"],\n", + " },\n", + " )\n", + "\n", + "\n", + "CHAIN_20 = CHAIN_DEFINITIONS[20]\n", + "K892_BW = _formulate_breit_wigner(\n", + " propagator=CHAIN_20[\"propagators\"][0],\n", + " function_definition=get_function_definition(\"K892_BW\", MODEL_DEFINITION),\n", + " resonance=name_to_latex(CHAIN_20[\"name\"]),\n", + ")\n", + "Math(aslatex(K892_BW))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Multi-channel Breit-Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@unevaluated\n", + "class MultichannelBreitWigner(sp.Expr):\n", + " s: Any\n", + " mass: Any\n", + " channels: list[ChannelArguments] = argument(sympify=False)\n", + "\n", + " def evaluate(self):\n", + " s = self.s\n", + " m0 = self.mass\n", + " width = sum(channel.formulate_width(s, m0) for channel in self.channels)\n", + " return BreitWigner(s, m0, width)\n", + "\n", + " def _latex_repr_(self, printer: LatexPrinter, *args) -> str:\n", + " latex = R\"\\mathcal{R}^\\mathrm{BW}_\\mathrm{multi}\\left(\"\n", + " latex += printer._print(self.s) + \"; \"\n", + " latex += \", \".join(printer._print(channel.width) for channel in self.channels)\n", + " latex += R\"\\right)\"\n", + " return latex\n", + "\n", + "\n", + "@frozen\n", + "class ChannelArguments:\n", + " width: Any\n", + " m1: Any = 0\n", + " m2: Any = 0\n", + " angular_momentum: Any = 0\n", + " meson_radius: Any = 1\n", + "\n", + " def __attrs_post_init__(self) -> None:\n", + " for name, value in asdict(self).items():\n", + " object.__setattr__(self, name, sp.sympify(value))\n", + "\n", + " def formulate_width(self, s: Any, m0: Any) -> sp.Expr:\n", + " Γ0 = self.width\n", + " m1 = self.m1\n", + " m2 = self.m2\n", + " L = self.angular_momentum\n", + " R = self.meson_radius\n", + " ff = FormFactor(s, m1, m2, L, R) ** 2\n", + " return Γ0 * m0 / sp.sqrt(s) * ff\n", + "\n", + "\n", + "x, y, z = sp.symbols(\"x:z\")\n", + "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", + "channels = [\n", + " ChannelArguments(\n", + " sp.Symbol(f\"Gamma{i}\"),\n", + " sp.Symbol(f\"m_{{a,{i}}}\"),\n", + " sp.Symbol(f\"m_{{b,{i}}}\"),\n", + " sp.Symbol(f\"L{i}\"),\n", + " d,\n", + " )\n", + " for i in [1, 2]\n", + "]\n", + "exprs = [\n", + " MultichannelBreitWigner(s, m0, channels),\n", + " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", + " BreitWigner(s, m0, Γ0),\n", + " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", + " FormFactor(s, m1, m2, L, d),\n", + " P(s, m1, m2),\n", + " Kallen(x, y, z),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CHAIN_DEFINITIONS[0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "get_function_definition(\"L1405_Flatte\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "def _formulate_multichannel_breit_wigner(\n", + " propagator: dict, function_definition: dict, resonance: str, **kwargs\n", + ") -> DefinedExpression:\n", + " channel_definitions = function_definition[\"channels\"]\n", + " if len(channel_definitions) < 2:\n", + " msg = \"Need at least two channels for a multi-channel Breit-Wigner\"\n", + " raise NotImplementedError(msg)\n", + " node = propagator[\"node\"]\n", + " i, j = node\n", + " s = _node_to_mandelstam(node)\n", + " resonance_latex = name_to_latex(resonance)\n", + " mass = sp.Symbol(f\"m_{{{resonance}}}\")\n", + " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\")\n", + " m1 = _node_to_mass(i)\n", + " m2 = _node_to_mass(j)\n", + " angular_momentum = int(channel_definitions[0][\"l\"])\n", + " d = sp.Symbol(f\"R_{{{resonance_latex}}}\")\n", + " channels = [ChannelArguments(width, m1, m2, angular_momentum, d)]\n", + " parameter_defaults = {\n", + " mass: function_definition[\"mass\"],\n", + " width: channel_definitions[0][\"gsq\"],\n", + " m1: channel_definitions[0][\"ma\"],\n", + " m2: channel_definitions[0][\"mb\"],\n", + " d: channel_definitions[0][\"d\"],\n", + " }\n", + " for i, channel_definition in enumerate(channel_definitions[1:], 2):\n", + " Γi = sp.Symbol(\n", + " Rf\"\\Gamma_{{{resonance_latex}}}^\\text{{ch. {i}}}\", nonnegative=True\n", + " )\n", + " mi1 = sp.Symbol(f\"m_{{a,{i}}}\", nonnegative=True)\n", + " mi2 = sp.Symbol(f\"m_{{b,{i}}}\", nonnegative=True)\n", + " angular_momentum = int(channel_definition[\"l\"])\n", + " channels.append(ChannelArguments(Γi, mi1, mi2, angular_momentum, d))\n", + " parameter_defaults.update({\n", + " mi1: channel_definition[\"ma\"],\n", + " mi2: channel_definition[\"mb\"],\n", + " })\n", + " return DefinedExpression(\n", + " expression=MultichannelBreitWigner(s, mass, channels),\n", + " definitions=parameter_defaults,\n", + " )\n", + "\n", + "\n", + "CHAIN_0 = CHAIN_DEFINITIONS[0]\n", + "L1405_Flatte = _formulate_multichannel_breit_wigner(\n", + " propagator=CHAIN_0[\"propagators\"][0],\n", + " function_definition=get_function_definition(\"L1405_Flatte\", MODEL_DEFINITION),\n", + " resonance=name_to_latex(CHAIN_0[\"name\"]),\n", + ")\n", + "Math(aslatex(L1405_Flatte))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Breit-Wigner with exponential" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "s, m0, Γ0, m1, m2, γ = sp.symbols(\"s m0 Gamma0 m1 m2 gamma\")\n", + "expr = BuggBreitWigner(s, m0, Γ0, m1, m2, γ)\n", + "Math(aslatex({expr: expr.doit(deep=False)}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CHAIN_DEFINITIONS[18]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "get_function_definition(\"K700_BuggBW\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _formulate_bugg_breit_wigner(\n", + " propagator: dict, function_definition: dict, resonance: str, model: dict, **kwargs\n", + ") -> DefinedExpression:\n", + " node = propagator[\"node\"]\n", + " i, j = node\n", + " s = _node_to_mandelstam(node)\n", + " mass = sp.Symbol(f\"m_{{{resonance}}}\", nonnegative=True)\n", + " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\", nonnegative=True)\n", + " γ = sp.Symbol(Rf\"\\gamma_{{{resonance}}}\")\n", + " m1 = _node_to_mass(i)\n", + " m2 = _node_to_mass(j)\n", + " final_state = _get_final_state(model)\n", + " return DefinedExpression(\n", + " expression=BuggBreitWigner(s, mass, width, m1, m2, γ),\n", + " definitions={\n", + " m0: function_definition[\"mass\"],\n", + " Γ0: function_definition[\"width\"],\n", + " m1: final_state[i].mass,\n", + " m2: final_state[j].mass,\n", + " γ: function_definition[\"slope\"],\n", + " },\n", + " )\n", + "\n", + "\n", + "CHAIN_18 = CHAIN_DEFINITIONS[18]\n", + "K700_BuggBW = _formulate_bugg_breit_wigner(\n", + " propagator=CHAIN_0[\"propagators\"][0],\n", + " function_definition=get_function_definition(\"K700_BuggBW\", MODEL_DEFINITION),\n", + " resonance=name_to_latex(CHAIN_0[\"name\"]),\n", + " model=MODEL_DEFINITION,\n", + ")\n", + "Math(aslatex(K700_BuggBW))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### General propagator dynamics builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def formulate_dynamics(chain_definition: dict, model: dict) -> DefinedExpression:\n", + " expr = DefinedExpression()\n", + " for propagator in chain_definition[\"propagators\"]:\n", + " parametrization = propagator[\"parametrization\"]\n", + " function_definition = get_function_definition(parametrization, model)\n", + " function_type = function_definition[\"type\"]\n", + " if function_type == \"BreitWigner\":\n", + " dynamics_builder = _formulate_breit_wigner\n", + " elif function_type == \"MultichannelBreitWigner\":\n", + " dynamics_builder = _formulate_multichannel_breit_wigner\n", + " elif function_type == \"BreitWignerWidthExpLikeBugg\":\n", + " dynamics_builder = _formulate_bugg_breit_wigner\n", + " else:\n", + " msg = f\"No dynamics implementation for function type {function_type!r}\"\n", + " raise NotImplementedError(msg)\n", + " expr *= dynamics_builder(\n", + " propagator,\n", + " function_definition=get_function_definition(parametrization, model),\n", + " resonance=name_to_latex(chain_definition[\"name\"]),\n", + " model=model,\n", + " )\n", + " return expr" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "Math(aslatex(formulate_dynamics(CHAIN_DEFINITIONS[0], MODEL_DEFINITION)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "Math(aslatex(formulate_dynamics(CHAIN_DEFINITIONS[18], MODEL_DEFINITION)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "Math(aslatex(formulate_dynamics(CHAIN_DEFINITIONS[20], MODEL_DEFINITION)))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Construct `AmplitudeModel`" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Unpolarized intensity" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [] + }, + "outputs": [], + "source": [ + "Topology = list[Union[int, \"Topology\"]]\n", + "\n", + "\n", + "def get_reference_subsystem(model: dict) -> FinalStateID:\n", + " topology = get_reference_topology(model)\n", + " return get_spectator_id(topology)\n", + "\n", + "\n", + "def get_spectator_id(topology: Topology) -> FinalStateID:\n", + " spectator_candidates = {i for i in topology if isinstance(i, int)}\n", + " if len(spectator_candidates) != 1:\n", + " msg = f\"Reference topology {topology} seems not to be a three-body decay\"\n", + " raise ValueError(msg)\n", + " return next(iter(spectator_candidates))\n", + "\n", + "\n", + "def get_reference_topology(model: dict) -> Topology:\n", + " distribution_def = _get_distribution_def(model)\n", + " return distribution_def[\"decay_description\"][\"reference_topology\"]\n", + "\n", + "\n", + "get_reference_subsystem(MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [] + }, + "outputs": [], + "source": [ + "def get_existing_subsystem_ids(model: dict) -> list[FinalStateID]:\n", + " distribution_def = _get_distribution_def(model)\n", + " chain_defs = distribution_def[\"decay_description\"][\"chains\"]\n", + " subsystem_ids = {get_spectator_id(c[\"topology\"]) for c in chain_defs}\n", + " return sorted(subsystem_ids)\n", + "\n", + "\n", + "get_existing_subsystem_ids(MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [] + }, + "outputs": [], + "source": [ + "from ampform.sympy import PoolSum\n", + "\n", + "from ampform_dpd import (\n", + " _AlignmentWignerGenerator, # noqa:PLC2701,RUF100\n", + " _generate_amplitude_index_bases, # noqa:PLC2701,RUF100\n", + ")\n", + "from ampform_dpd.spin import create_spin_range\n", + "\n", + "\n", + "def formulate_aligned_amplitude(\n", + " model: dict,\n", + " λ1: sp.Rational | sp.Symbol,\n", + " λ2: sp.Rational | sp.Symbol,\n", + " λ3: sp.Rational | sp.Symbol,\n", + ") -> tuple[PoolSum, dict[sp.Symbol, sp.Expr]]:\n", + " reference_subsystem = get_reference_subsystem(model)\n", + " wigner_generator = _AlignmentWignerGenerator(reference_subsystem)\n", + " λ0 = sp.Symbol(R\"\\lambda_0\", rational=True)\n", + " _λ0, _λ1, _λ2, _λ3 = sp.symbols(R\"\\lambda_(0:4)^{\\prime}\", rational=True)\n", + " distribution_def = _get_distribution_def(model)\n", + " decay_description = distribution_def[\"decay_description\"]\n", + " states = _get_states(decay_description[\"kinematics\"])\n", + " j0, j1, j2, j3 = (states[i].spin for i in sorted(states))\n", + " A = _generate_amplitude_index_bases()\n", + " amp_expr = PoolSum(\n", + " sum(\n", + " A[k][_λ0, _λ1, _λ2, _λ3]\n", + " * wigner_generator(j0, λ0, _λ0, rotated_state=0, aligned_subsystem=k)\n", + " * wigner_generator(j1, _λ1, λ1, rotated_state=1, aligned_subsystem=k)\n", + " * wigner_generator(j2, _λ2, λ2, rotated_state=2, aligned_subsystem=k)\n", + " * wigner_generator(j3, _λ3, λ3, rotated_state=3, aligned_subsystem=k)\n", + " for k in get_existing_subsystem_ids(model)\n", + " ),\n", + " (_λ0, create_spin_range(j0)),\n", + " (_λ1, create_spin_range(j1)),\n", + " (_λ2, create_spin_range(j2)),\n", + " (_λ3, create_spin_range(j3)),\n", + " )\n", + " return amp_expr, wigner_generator.angle_definitions\n", + "\n", + "\n", + "λ1, λ2, λ3 = sp.symbols(R\"\\lambda_(1:4)\", rational=True)\n", + "amplitude_expr, _ = formulate_aligned_amplitude(MODEL_DEFINITION, λ1, λ2, λ3)\n", + "amplitude_expr.cleanup()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Amplitude for the decay chain" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CHAIN_DEFINITIONS[0]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _formulate_recoupling(\n", + " chain_definition: dict, vertex_idx: int\n", + ") -> tuple[sp.Indexed, sp.Number]:\n", + " vertex_definitions = chain_definition[\"vertices\"]\n", + " if len(vertex_definitions) != 2:\n", + " msg = f\"Not a three-body decay: there are {len(vertex_definitions)} vertices\"\n", + " raise ValueError(msg)\n", + " if vertex_idx not in {0, 1}:\n", + " msg = f\"Vertex index out of range. Can either be 0 or 1, not {vertex_idx}.\"\n", + " raise ValueError(msg)\n", + " vertex = chain_definition[\"vertices\"][vertex_idx]\n", + " vertex_type = vertex[\"type\"]\n", + " if vertex_type in {\"helicity\", \"parity\"}:\n", + " helicity_coupling = True\n", + " interaction = None\n", + " elif vertex_type == \"ls\":\n", + " msg = \"No implementation yet for LS-coupling models\"\n", + " raise NotImplementedError(msg)\n", + " else:\n", + " msg = f\"No implementation for vertex of type {vertex_type!r}\"\n", + " raise NotImplementedError(msg)\n", + " symbol = _create_coupling_symbol(\n", + " helicity_coupling,\n", + " resonance=Str(name_to_latex(chain_definition[\"name\"])),\n", + " helicities=tuple(sp.Rational(v) for v in vertex[\"helicities\"]),\n", + " interaction=interaction,\n", + " typ=\"production\" if vertex_idx == 0 else \"decay\",\n", + " )\n", + " value = sp.S.One\n", + " _warn_once(\"No implementation yet for value of the helicity recoupling factor\")\n", + " return symbol, value\n", + "\n", + "\n", + "def _warn_once(msg: str) -> None:\n", + " warn(msg, category=UserWarning)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "scroll-output" + ] + }, + "outputs": [], + "source": [ + "Math(aslatex(dict(_formulate_recoupling(CHAIN_DEFINITIONS[0], i) for i in range(2))))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _get_final_state_helicities(\n", + " chain_definition: dict,\n", + ") -> dict[FinalStateID, sp.Rational]:\n", + " vertices = chain_definition[\"vertices\"]\n", + " helicities: dict[FinalStateID, sp.Rational] = {}\n", + " for v in vertices:\n", + " for helicity, node in zip(v[\"helicities\"], v[\"node\"]):\n", + " if not isinstance(node, int):\n", + " continue\n", + " helicities[node] = sp.Rational(helicity)\n", + " return {i: helicities[i] for i in sorted(helicities)}\n", + "\n", + "\n", + "_get_final_state_helicities(CHAIN_DEFINITIONS[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _get_resonance_helicity(\n", + " chain_definition: dict,\n", + ") -> tuple[tuple[FinalStateID, FinalStateID], sp.Rational]:\n", + " vertices = chain_definition[\"vertices\"]\n", + " for vertex in vertices:\n", + " node = vertex[\"node\"]\n", + " if all(isinstance(i, int) for i in node):\n", + " continue\n", + " helicities = vertex[\"helicities\"]\n", + " for helicity, sub_node in zip(helicities, node):\n", + " if isinstance(sub_node, abc.Sequence) and len(sub_node) == 2:\n", + " return tuple(sub_node), sp.Rational(helicity)\n", + " msg = \"Could not find a resonance node\"\n", + " raise ValueError(msg)\n", + "\n", + "\n", + "_get_resonance_helicity(CHAIN_DEFINITIONS[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _get_spectator_helicity(chain_definition: dict) -> tuple[int, sp.Rational]:\n", + " vertices = chain_definition[\"vertices\"]\n", + " for vertex in vertices:\n", + " node = vertex[\"node\"]\n", + " if all(isinstance(i, int) for i in node):\n", + " continue\n", + " helicities = vertex[\"helicities\"]\n", + " for helicity, state_id in zip(helicities, node):\n", + " if isinstance(state_id, int):\n", + " return state_id, sp.Rational(helicity)\n", + " msg = \"Could not find vertex that contains a spectator\"\n", + " raise ValueError(msg)\n", + "\n", + "\n", + "_get_spectator_helicity(CHAIN_DEFINITIONS[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _get_decay_product_helicities(\n", + " chain_definition: dict,\n", + ") -> tuple[tuple[int, sp.Rational], tuple[int, sp.Rational]]:\n", + " vertices = chain_definition[\"vertices\"]\n", + " for v in vertices:\n", + " node = v[\"node\"]\n", + " if all(isinstance(i, int) for i in node):\n", + " helicities = v[\"helicities\"]\n", + " return tuple((i, sp.Rational(λ)) for i, λ in zip(node, helicities))\n", + " msg = \"Could not fine a helicity for any resonance node\"\n", + " raise ValueError(msg)\n", + "\n", + "\n", + "_get_decay_product_helicities(CHAIN_DEFINITIONS[0])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def _get_weight(chain_definition: dict) -> tuple[sp.Symbol, complex | float]:\n", + " value: complex | float\n", + " value = complex(str(chain_definition[\"weight\"]).replace(\" \", \"\").replace(\"i\", \"j\"))\n", + " if not value.imag:\n", + " value = value.real\n", + " resonance_latex = name_to_latex(chain_definition[\"name\"])\n", + " _, resonance_helicity = _get_resonance_helicity(chain_definition)\n", + " c = sp.IndexedBase(f\"c^{{{resonance_latex}[{resonance_helicity}]}}\")\n", + " λ1, λ2, λ3 = _get_final_state_helicities(chain_definition).values()\n", + " symbol = c[λ1, λ2, λ3]\n", + " return symbol, value\n", + "\n", + "\n", + "Math(aslatex(dict([_get_weight(CHAIN_DEFINITIONS[0])])))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def formulate_chain_amplitude(model: dict, chain_idx: int) -> sp.Expr:\n", + " chain_defs = _get_decay_chains(model)\n", + " chain_definition = chain_defs[chain_idx]\n", + " definitions = {}\n", + " # -----------------------\n", + " dynamics = formulate_dynamics(chain_definition, model)\n", + " for vertex in chain_definition[\"vertices\"]:\n", + " dynamics *= formulate_form_factor(vertex, model)\n", + " definitions.update(dynamics.definitions)\n", + " # -----------------------\n", + " weight, weight_val = _get_weight(chain_definition)\n", + " h_prod, h_prod_val = _formulate_recoupling(chain_definition, vertex_idx=0)\n", + " h_dec, h_dec_val = _formulate_recoupling(chain_definition, vertex_idx=1)\n", + " definitions[weight] = weight_val\n", + " definitions[h_prod] = h_prod_val\n", + " definitions[h_dec] = h_dec_val\n", + " # -----------------------\n", + " (i, λi), (j, λj) = _get_decay_product_helicities(chain_definition)\n", + " θij, θij_expr = formulate_scattering_angle(i, j)\n", + " definitions[θij] = θij_expr\n", + " jR = sp.Rational(chain_definition[\"propagators\"][0][\"spin\"])\n", + " _, λR = _get_resonance_helicity(chain_definition)\n", + " # -----------------------\n", + " A = _generate_amplitude_index_bases()\n", + " subsystem_id = get_spectator_id(chain_definition[\"topology\"])\n", + " k, λk = _get_spectator_helicity(chain_definition)\n", + " λ = {i: λi, j: λj, k: λk}\n", + " amplitude_symbol = A[subsystem_id][λ[1], λ[2], λ[3]]\n", + " definitions[amplitude_symbol] = (\n", + " weight * h_prod * h_dec * Wigner.d(jR, λR, λi - λj, θij) * dynamics.expression\n", + " )\n", + " return definitions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "definitions = formulate_chain_amplitude(MODEL_DEFINITION, chain_idx=0)\n", + "Math(aslatex(definitions))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Full amplitude model" ] } ], diff --git a/pyproject.toml b/pyproject.toml index afaff4e9..49ba4334 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -338,6 +338,9 @@ ignore-names = [ [tool.ruff.lint.pydocstyle] convention = "google" +[tool.ruff.lint.pylint] +allow-dunder-method-names = ["_latex_repr_"] + [tool.tomlsort] all = false ignore_case = true From c5890bc718273c64189f1b8202bf5ce72c286d28 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:39 +0200 Subject: [PATCH 11/42] ENH: implement `aslatex()` for `DefinedExpression` --- src/ampform_dpd/io.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/src/ampform_dpd/io.py b/src/ampform_dpd/io.py index f4c517e1..b843ca25 100644 --- a/src/ampform_dpd/io.py +++ b/src/ampform_dpd/io.py @@ -30,6 +30,7 @@ from ampform.io import aslatex from tensorwaves.function.sympy import create_function, create_parametrized_function +from ampform_dpd import DefinedExpression from ampform_dpd._cache import get_readable_hash, get_system_cache_directory from ampform_dpd.decay import ( IsobarNode, @@ -86,6 +87,21 @@ def _(obj: Particle, with_jp: bool = False, only_jp: bool = False, **kwargs) -> return obj.latex +@aslatex.register(DefinedExpression) +def _(obj: DefinedExpression, **kwargs) -> str: + latex = R"\begin{array}{rcl}" + "\n" + expr = obj.expression + unfolded = expr.doit(deep=False) + if expr == unfolded: + latex += Rf" {aslatex(obj.expression, **kwargs)} \\" + "\n" + else: + latex += Rf" {aslatex(expr)} &=& {aslatex(unfolded)} \\" + "\n" + for lhs, rhs in obj.definitions.items(): + latex += Rf" {aslatex(lhs)} &=& {aslatex(rhs)} \\" + "\n" + latex += R"\end{array}" + return latex + + def _render_jp(particle: Particle) -> str: if particle.spin.denominator == 1: spin = sp.latex(particle.spin) From ce0aa753fab74faef186042e8893e3988e9303ac Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:40 +0200 Subject: [PATCH 12/42] FEAT: implement model serialization (validation fails) --- .cspell.json | 2 + docs/serialization.ipynb | 713 ++++++++++++++++++++++++++++++++++----- pyproject.toml | 8 + 3 files changed, 647 insertions(+), 76 deletions(-) diff --git a/.cspell.json b/.cspell.json index be509934..219f58e3 100644 --- a/.cspell.json +++ b/.cspell.json @@ -61,6 +61,7 @@ "figsize", "fontsize", "formfactor", + "funcs", "gcov", "graphviz", "imag", @@ -163,6 +164,7 @@ "PYTHONHASHSEED", "QRules", "recoupling", + "recouplings", "sympify", "SymPy", "TensorWaves", diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 9631ca89..b23ff5cb 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -24,7 +24,8 @@ "source_hidden": true }, "tags": [ - "hide-cell" + "hide-input", + "scroll-input" ] }, "outputs": [], @@ -34,22 +35,32 @@ "import json\n", "from collections import abc\n", "from difflib import get_close_matches\n", - "from typing import Any, Callable, Sequence, Union\n", + "from itertools import product\n", + "from typing import Any, Callable, Literal, Sequence, Union\n", "from warnings import warn\n", "\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", "import sympy as sp\n", "from ampform.dynamics import BlattWeisskopfSquared\n", "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import argument, unevaluated\n", + "from ampform.sympy import PoolSum, argument, perform_cached_doit, unevaluated\n", "from attrs import asdict, frozen\n", "from IPython.display import JSON, Markdown, Math\n", - "from sympy.core.symbol import Str\n", + "from sympy.functions.special.tensor_functions import KroneckerDelta as δ\n", + "from sympy.physics.quantum.cg import CG\n", "from sympy.physics.quantum.spin import Rotation as Wigner\n", "from sympy.printing.latex import LatexPrinter\n", + "from tqdm.auto import tqdm\n", "\n", "from ampform_dpd import (\n", + " AmplitudeModel,\n", " DefinedExpression,\n", - " _create_coupling_symbol, # noqa:PLC2701\n", + " _AlignmentWignerGenerator,\n", + " _generate_amplitude_index_bases,\n", + " create_mass_symbol_mapping,\n", + " formulate_invariants,\n", ")\n", "from ampform_dpd.angles import formulate_scattering_angle\n", "from ampform_dpd.decay import (\n", @@ -62,7 +73,13 @@ " ThreeBodyDecayChain,\n", ")\n", "from ampform_dpd.dynamics import BuggBreitWigner, EnergyDependentWidth, FormFactor, P\n", - "from ampform_dpd.io import as_markdown_table, aslatex, simplify_latex_rendering\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 create_spin_range\n", "\n", "simplify_latex_rendering()" ] @@ -410,8 +427,8 @@ " and all(isinstance(i, int) for i in node_item)\n", " and len(node_item) == 2\n", " ):\n", - " ij = \"\".join(str(i) for i in node_item)\n", - " return sp.Symbol(f\"s{ij}\", nonnegative=True)\n", + " k, *_ = {1, 2, 3} - set(node_item)\n", + " return sp.Symbol(f\"sigma{k}\", nonnegative=True)\n", " msg = f\"Cannot create mass symbol for node {node_item}\"\n", " raise NotImplementedError(msg)\n", "\n", @@ -536,11 +553,13 @@ " m2: Any = 0\n", " angular_momentum: Any = 0\n", " meson_radius: Any = 1\n", - " _latex_repr_ = R\"\\mathcal{{R}}^\\mathrm{{BW}}_{{{angular_momentum}}}\\left({s}\\right)\" # noqa:RUF027\n", "\n", " def evaluate(self):\n", " width = self.energy_dependent_width()\n", - " return BreitWigner(self.s, self.mass, width)\n", + " expr = SimpleBreitWigner(self.s, self.mass, width)\n", + " if self.angular_momentum == 0 and self.m1 == 0 and self.m2 == 0:\n", + " return expr.evaluate()\n", + " return expr\n", "\n", " def energy_dependent_width(self) -> sp.Expr:\n", " s, m0, Γ0, m1, m2, L, d = self.args\n", @@ -551,21 +570,32 @@ " def _latex_repr_(self, printer: LatexPrinter, *args) -> str:\n", " s = printer._print(self.s)\n", " function_symbol = R\"\\mathcal{R}^\\mathrm{BW}\"\n", - " arg = Rf\"\\left({s}\\right)\"\n", - " if self.angular_momentum == 0:\n", - " if self.m1 == 0 and self.m2 == 0:\n", - " width = printer._print(self.energy_dependent_width())\n", - " arg = Rf\"\\left({s}; {width}\\right)\"\n", - " return Rf\"{function_symbol}{arg}\"\n", + " mass = printer._print(self.mass)\n", + " width = printer._print(self.width)\n", + " arg = Rf\"\\left({s}; {mass}, {width}\\right)\"\n", " L = printer._print(self.angular_momentum)\n", + " if isinstance(self.angular_momentum, sp.Integer):\n", + " return Rf\"{function_symbol}_{{L={L}}}{arg}\"\n", " return Rf\"{function_symbol}_{{{L}}}{arg}\"\n", "\n", "\n", + "@unevaluated\n", + "class SimpleBreitWigner(sp.Expr):\n", + " s: Any\n", + " mass: Any\n", + " width: Any\n", + " _latex_repr_ = R\"\\mathcal{{R}}^\\mathrm{{BW}}\\left({s}; {mass}, {width}\\right)\"\n", + "\n", + " def evaluate(self):\n", + " s, m0, Γ0 = self.args\n", + " return 1 / (m0**2 - s - m0 * Γ0 * 1j)\n", + "\n", + "\n", "x, y, z = sp.symbols(\"x:z\")\n", "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", "exprs = [\n", " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", - " BreitWigner(s, m0, Γ0),\n", + " SimpleBreitWigner(s, m0, Γ0),\n", " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", " FormFactor(s, m1, m2, L, d),\n", " P(s, m1, m2),\n", @@ -876,8 +906,8 @@ " return DefinedExpression(\n", " expression=BuggBreitWigner(s, mass, width, m1, m2, γ),\n", " definitions={\n", - " m0: function_definition[\"mass\"],\n", - " Γ0: function_definition[\"width\"],\n", + " mass: function_definition[\"mass\"],\n", + " width: function_definition[\"width\"],\n", " m1: final_state[i].mass,\n", " m2: final_state[j].mass,\n", " γ: function_definition[\"slope\"],\n", @@ -1053,25 +1083,16 @@ }, "outputs": [], "source": [ - "from ampform.sympy import PoolSum\n", - "\n", - "from ampform_dpd import (\n", - " _AlignmentWignerGenerator, # noqa:PLC2701,RUF100\n", - " _generate_amplitude_index_bases, # noqa:PLC2701,RUF100\n", - ")\n", - "from ampform_dpd.spin import create_spin_range\n", - "\n", - "\n", "def formulate_aligned_amplitude(\n", " model: dict,\n", + " λ0: sp.Rational | sp.Symbol,\n", " λ1: sp.Rational | sp.Symbol,\n", " λ2: sp.Rational | sp.Symbol,\n", " λ3: sp.Rational | sp.Symbol,\n", ") -> tuple[PoolSum, dict[sp.Symbol, sp.Expr]]:\n", " reference_subsystem = get_reference_subsystem(model)\n", " wigner_generator = _AlignmentWignerGenerator(reference_subsystem)\n", - " λ0 = sp.Symbol(R\"\\lambda_0\", rational=True)\n", - " _λ0, _λ1, _λ2, _λ3 = sp.symbols(R\"\\lambda_(0:4)^{\\prime}\", rational=True)\n", + " _λ0, _λ1, _λ2, _λ3 = sp.symbols(R\"\\lambda_(:4)^{\\prime}\", rational=True)\n", " distribution_def = _get_distribution_def(model)\n", " decay_description = distribution_def[\"decay_description\"]\n", " states = _get_states(decay_description[\"kinematics\"])\n", @@ -1094,9 +1115,9 @@ " return amp_expr, wigner_generator.angle_definitions\n", "\n", "\n", - "λ1, λ2, λ3 = sp.symbols(R\"\\lambda_(1:4)\", rational=True)\n", - "amplitude_expr, _ = formulate_aligned_amplitude(MODEL_DEFINITION, λ1, λ2, λ3)\n", - "amplitude_expr.cleanup()" + "λ0, λ1, λ2, λ3 = sp.symbols(\"lambda(:4)\", rational=True)\n", + "amplitude_expr, _ = formulate_aligned_amplitude(MODEL_DEFINITION, λ0, λ1, λ2, λ3)\n", + "amplitude_expr" ] }, { @@ -1106,6 +1127,120 @@ "### Amplitude for the decay chain" ] }, + { + "cell_type": "markdown", + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, + "source": [ + "#### Helicity recouplings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "scroll-input" + ] + }, + "outputs": [], + "source": [ + "@unevaluated\n", + "class HelicityRecoupling(sp.Expr):\n", + " λa: sp.Rational | sp.Symbol\n", + " λb: sp.Rational | sp.Symbol\n", + " λa0: sp.Rational | sp.Symbol\n", + " λb0: sp.Rational | sp.Symbol\n", + " _latex_repr_ = R\"\\mathcal{{H}}^\\text{{helicity}}\\left({λa},{λb}|{λa0},{λb0}\\right)\"\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " λa, λb, λa0, λb0 = self.args\n", + " return δ(λa, λa0) * δ(λb, λb0)\n", + "\n", + "\n", + "@unevaluated\n", + "class ParityRecoupling(sp.Expr):\n", + " λa: sp.Rational | sp.Symbol\n", + " λb: sp.Rational | sp.Symbol\n", + " λa0: sp.Rational | sp.Symbol\n", + " λb0: sp.Rational | sp.Symbol\n", + " f: sp.Integer | sp.Symbol\n", + " _latex_repr_ = (\n", + " R\"\\mathcal{{H}}^\\text{{parity}}\\left({λa},{λb}|{λa0},{λb0},{f}\\right)\"\n", + " )\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " λa, λb, λa0, λb0, f = self.args\n", + " return δ(λa, λa0) * δ(λb, λb0) + f * δ(λa, -λa0) * δ(λb, -λb0)\n", + "\n", + "\n", + "@unevaluated\n", + "class LSRecoupling(sp.Expr):\n", + " λa: sp.Rational | sp.Symbol\n", + " λb: sp.Rational | sp.Symbol\n", + " l: sp.Integer | sp.Symbol\n", + " s: sp.Rational | sp.Symbol\n", + " ja: sp.Rational | sp.Symbol\n", + " jb: sp.Rational | sp.Symbol\n", + " j: sp.Rational | sp.Symbol\n", + " _latex_repr_ = (\n", + " R\"\\mathcal{{H}}^\\text{{parity}}\\left({λa},{λb}|{l},{s},{ja},{jb},{j}\\right)\"\n", + " )\n", + "\n", + " def evaluate(self) -> sp.Expr:\n", + " λa, λb, l, s, ja, jb, j = self.args\n", + " return (\n", + " sp.sqrt((2 * l + 1) / (2 * j + 1))\n", + " * CG(ja, λa, jb, -λb, s, λa - λb)\n", + " * CG(l, 0, s, λa - λb, j, λa - λb)\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "λa = sp.Symbol(R\"\\lambda_a\", rational=True)\n", + "λb = sp.Symbol(R\"\\lambda_b\", rational=True)\n", + "λa0 = sp.Symbol(R\"\\lambda_a^0\", rational=True)\n", + "λb0 = sp.Symbol(R\"\\lambda_b^0\", rational=True)\n", + "f = sp.Symbol(\"f\", integer=True)\n", + "l = sp.Symbol(\"l\", integer=True, nonnegative=True)\n", + "s = sp.Symbol(\"s\", nonnegative=True, rational=True)\n", + "ja = sp.Symbol(\"j_a\", nonnegative=True, rational=True)\n", + "jb = sp.Symbol(\"j_b\", nonnegative=True, rational=True)\n", + "j = sp.Symbol(\"j\", nonnegative=True, rational=True)\n", + "exprs = [\n", + " HelicityRecoupling(λa, λb, λa0, λb0),\n", + " ParityRecoupling(λa, λb, λa0, λb0, f),\n", + " LSRecoupling(λa, λb, l, s, ja, jb, j),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, + "source": [ + "#### Recoupling deserialization" + ] + }, { "cell_type": "code", "execution_count": null, @@ -1125,9 +1260,8 @@ }, "outputs": [], "source": [ - "def _formulate_recoupling(\n", - " chain_definition: dict, vertex_idx: int\n", - ") -> tuple[sp.Indexed, sp.Number]:\n", + "def _formulate_recoupling(model: dict, chain_idx: int, vertex_idx: int) -> sp.Expr:\n", + " chain_definition = _get_decay_chains(model)[chain_idx]\n", " vertex_definitions = chain_definition[\"vertices\"]\n", " if len(vertex_definitions) != 2:\n", " msg = f\"Not a three-body decay: there are {len(vertex_definitions)} vertices\"\n", @@ -1137,42 +1271,104 @@ " raise ValueError(msg)\n", " vertex = chain_definition[\"vertices\"][vertex_idx]\n", " vertex_type = vertex[\"type\"]\n", + " node = vertex[\"node\"]\n", + " λa, λb = map(_get_helicity_symbol, node)\n", " if vertex_type in {\"helicity\", \"parity\"}:\n", - " helicity_coupling = True\n", - " interaction = None\n", - " elif vertex_type == \"ls\":\n", - " msg = \"No implementation yet for LS-coupling models\"\n", - " raise NotImplementedError(msg)\n", - " else:\n", - " msg = f\"No implementation for vertex of type {vertex_type!r}\"\n", - " raise NotImplementedError(msg)\n", - " symbol = _create_coupling_symbol(\n", - " helicity_coupling,\n", - " resonance=Str(name_to_latex(chain_definition[\"name\"])),\n", - " helicities=tuple(sp.Rational(v) for v in vertex[\"helicities\"]),\n", - " interaction=interaction,\n", - " typ=\"production\" if vertex_idx == 0 else \"decay\",\n", - " )\n", - " value = sp.S.One\n", - " _warn_once(\"No implementation yet for value of the helicity recoupling factor\")\n", - " return symbol, value\n", + " λa0, λb0 = (sp.Rational(v) for v in vertex[\"helicities\"])\n", + " if vertex_type == \"parity\":\n", + " f = _sign_to_value(vertex.get(\"parity_factor\", \"+\"))\n", + " return ParityRecoupling(λa, λb, λa0, λb0, f)\n", + " return HelicityRecoupling(λa, λb, λa0, λb0)\n", + " if vertex_type == \"ls\":\n", + " l = int(vertex[\"l\"])\n", + " s = sp.Rational(vertex[\"s\"])\n", + " ja, jb = _get_child_spins(model, chain_idx, vertex_idx)\n", + " j = _get_parent_spin(vertex)\n", + " return LSRecoupling(λa, λb, l, s, ja, jb, j)\n", + " msg = f\"No implementation for vertex of type {vertex_type!r}\"\n", + " raise NotImplementedError(msg)\n", + "\n", + "\n", + "def _sign_to_value(sign: Literal[\"\", \"-\", \"+\"]) -> Literal[0, -1, 1]:\n", + " stripped_sign = sign.strip()\n", + " if stripped_sign == \"-\":\n", + " return -1\n", + " if not stripped_sign:\n", + " return 0\n", + " if stripped_sign == \"+\":\n", + " return +1\n", + " msg = f\"Cannot convert {sign!r} to value\"\n", + " raise NotImplementedError(msg)\n", "\n", "\n", - "def _warn_once(msg: str) -> None:\n", - " warn(msg, category=UserWarning)" + "def _get_parent_spin(\n", + " model: dict, chain_idx: int, vertex_idx: int\n", + ") -> tuple[sp.Rational, sp.Rational]:\n", + " chain_definition = _get_decay_chains(model)[chain_idx]\n", + " vertex = chain_definition[\"vertices\"]\n", + " node = vertex[\"node\"]\n", + " initial_state = _get_initial_state(model)\n", + " if all(isinstance(i, int) for i in node):\n", + " return __get_propagator_spin(vertex)\n", + " return initial_state.spin\n", + "\n", + "\n", + "def _get_child_spins(\n", + " model: dict, chain_idx: int, vertex_idx: int\n", + ") -> tuple[sp.Rational, sp.Rational]:\n", + " chain_definition = _get_decay_chains(model)[chain_idx]\n", + " vertex = chain_definition[\"vertices\"]\n", + " node = vertex[\"node\"]\n", + " final_state = _get_final_state(model)\n", + " spins = []\n", + " for node_item in node:\n", + " if isinstance(node_item, int):\n", + " spins.append(sp.Rational(final_state[node_item]))\n", + " else:\n", + " spins.append(__get_propagator_spin(vertex[\"propagators\"]))\n", + " return tuple(spins)\n", + "\n", + "\n", + "def __get_propagator_spin(vertex: dict) -> sp.Rational:\n", + " propagators = vertex[\"propagators\"]\n", + " if len(propagators) != 1:\n", + " msg = f\"There are {len(propagators)} propagators, not a three-body decay\"\n", + " raise ValueError(msg)\n", + " return sp.Rational(propagators[0][\"spin\"])\n", + "\n", + "\n", + "def _get_helicity_symbol(node: int | Topology) -> sp.Symbol:\n", + " if isinstance(node, int):\n", + " return sp.Symbol(f\"lambda{node}\", rational=True)\n", + " return sp.Symbol(R\"\\lambda_R\", rational=True)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "tags": [ "scroll-output" ] }, "outputs": [], "source": [ - "Math(aslatex(dict(_formulate_recoupling(CHAIN_DEFINITIONS[0], i) for i in range(2))))" + "recouplings = [\n", + " _formulate_recoupling(MODEL_DEFINITION, chain_idx=0, vertex_idx=i) for i in range(2)\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in recouplings}))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "jp-MarkdownHeadingCollapsed": true + }, + "source": [ + "#### Chain amplitudes" ] }, { @@ -1312,41 +1508,55 @@ { "cell_type": "code", "execution_count": null, - "metadata": {}, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, "outputs": [], "source": [ - "def formulate_chain_amplitude(model: dict, chain_idx: int) -> sp.Expr:\n", + "def formulate_chain_amplitude(\n", + " λ0: sp.Rational,\n", + " λ1: sp.Rational,\n", + " λ2: sp.Rational,\n", + " λ3: sp.Rational,\n", + " model: dict,\n", + " chain_idx: int,\n", + ") -> sp.Expr:\n", " chain_defs = _get_decay_chains(model)\n", " chain_definition = chain_defs[chain_idx]\n", - " definitions = {}\n", " # -----------------------\n", " dynamics = formulate_dynamics(chain_definition, model)\n", " for vertex in chain_definition[\"vertices\"]:\n", " dynamics *= formulate_form_factor(vertex, model)\n", - " definitions.update(dynamics.definitions)\n", " # -----------------------\n", " weight, weight_val = _get_weight(chain_definition)\n", - " h_prod, h_prod_val = _formulate_recoupling(chain_definition, vertex_idx=0)\n", - " h_dec, h_dec_val = _formulate_recoupling(chain_definition, vertex_idx=1)\n", - " definitions[weight] = weight_val\n", - " definitions[h_prod] = h_prod_val\n", - " definitions[h_dec] = h_dec_val\n", " # -----------------------\n", - " (i, λi), (j, λj) = _get_decay_product_helicities(chain_definition)\n", + " (i, λi_val), (j, λj_val) = _get_decay_product_helicities(chain_definition)\n", " θij, θij_expr = formulate_scattering_angle(i, j)\n", - " definitions[θij] = θij_expr\n", " jR = sp.Rational(chain_definition[\"propagators\"][0][\"spin\"])\n", - " _, λR = _get_resonance_helicity(chain_definition)\n", + " R_node, λR_val = _get_resonance_helicity(chain_definition)\n", + " λR = _get_helicity_symbol(R_node)\n", " # -----------------------\n", " A = _generate_amplitude_index_bases()\n", " subsystem_id = get_spectator_id(chain_definition[\"topology\"])\n", - " k, λk = _get_spectator_helicity(chain_definition)\n", - " λ = {i: λi, j: λj, k: λk}\n", - " amplitude_symbol = A[subsystem_id][λ[1], λ[2], λ[3]]\n", - " definitions[amplitude_symbol] = (\n", - " weight * h_prod * h_dec * Wigner.d(jR, λR, λi - λj, θij) * dynamics.expression\n", + " h_prod = _formulate_recoupling(chain_definition, chain_idx, vertex_idx=0)\n", + " h_dec = _formulate_recoupling(chain_definition, chain_idx, vertex_idx=1)\n", + " amplitude_expression = (\n", + " weight\n", + " * h_prod\n", + " * h_dec\n", + " * Wigner.d(jR, λR, λi_val - λj_val, θij)\n", + " * dynamics.expression\n", " )\n", - " return definitions" + " amplitude_expression = amplitude_expression.subs({λR: λR_val})\n", + " amplitude_symbol = A[subsystem_id][λ0, λ1, λ2, λ3]\n", + " return {\n", + " amplitude_symbol: amplitude_expression,\n", + " weight: weight_val,\n", + " **dynamics.definitions,\n", + " θij: θij_expr,\n", + " }" ] }, { @@ -1355,7 +1565,7 @@ "metadata": {}, "outputs": [], "source": [ - "definitions = formulate_chain_amplitude(MODEL_DEFINITION, chain_idx=0)\n", + "definitions = formulate_chain_amplitude(λ0, λ1, λ2, λ3, MODEL_DEFINITION, chain_idx=0)\n", "Math(aslatex(definitions))" ] }, @@ -1365,6 +1575,357 @@ "source": [ "### Full amplitude model" ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + } + }, + "outputs": [], + "source": [ + "def formulate(model: dict, cleanup_summations: bool = False) -> AmplitudeModel:\n", + " states = _get_states(model)\n", + " helicity_symbols = sp.symbols(\"lambda(:4)\", rational=True)\n", + " allowed_helicities = {\n", + " symbol: create_spin_range(states[i].spin)\n", + " for i, symbol in enumerate(helicity_symbols)\n", + " }\n", + " amplitude_definitions = {}\n", + " angle_definitions = {}\n", + " parameter_defaults = {}\n", + " n_chains = len(_get_decay_chains(model))\n", + " helicity_values: tuple[sp.Rational, sp.Rational, sp.Rational, sp.Rational]\n", + " for helicity_values in product(*allowed_helicities.values()):\n", + " for chain_idx in range(n_chains):\n", + " amp_defs = formulate_chain_amplitude(*helicity_values, model, chain_idx)\n", + " (amp_symbol, amp_expr), *parameters, (θij, θij_expr) = amp_defs.items()\n", + " helicity_substitutions = dict(zip(helicity_symbols, helicity_values))\n", + " amplitude_definitions[amp_symbol] = amp_expr.subs(helicity_substitutions)\n", + " angle_definitions[θij] = θij_expr\n", + " parameter_defaults.update(dict(parameters))\n", + " aligned_amp, zeta_defs = formulate_aligned_amplitude(model, *helicity_symbols)\n", + " angle_definitions.update(zeta_defs)\n", + " decay = to_decay(model)\n", + " masses = create_mass_symbol_mapping(decay)\n", + " parameter_defaults.update(masses)\n", + " if cleanup_summations:\n", + " aligned_amp = aligned_amp.cleanup() # type:ignore[assignment]\n", + " intensity = PoolSum(\n", + " sp.Abs(aligned_amp) ** 2,\n", + " *allowed_helicities.items(),\n", + " )\n", + " if cleanup_summations:\n", + " intensity = intensity.cleanup() # type:ignore[assignment]\n", + " return AmplitudeModel(\n", + " decay=decay,\n", + " intensity=intensity,\n", + " amplitudes=amplitude_definitions,\n", + " variables=angle_definitions,\n", + " parameter_defaults=parameter_defaults,\n", + " masses=masses,\n", + " invariants=formulate_invariants(decay),\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Symbolic result" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "MODEL = formulate(MODEL_DEFINITION, cleanup_summations=True)\n", + "MODEL.intensity" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "Math(aslatex(MODEL.amplitudes))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "Math(aslatex(MODEL.variables))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "Math(aslatex({**MODEL.invariants, **MODEL.masses}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "intensity_expr = MODEL.full_expression.xreplace(MODEL.variables)\n", + "intensity_expr = intensity_expr.xreplace(MODEL.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "free_symbols = intensity_expr.free_symbols\n", + "assert len(free_symbols) == 3\n", + "assert str(sorted(free_symbols, key=str)) == \"[sigma1, sigma2, sigma3]\"" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Numeric results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [] + }, + "outputs": [], + "source": [ + "intensity_funcs = {}\n", + "for s, s_expr in tqdm(MODEL.invariants.items()):\n", + " k = int(str(s)[-1])\n", + " s_expr = s_expr.xreplace(MODEL.masses).doit()\n", + " expr = perform_cached_doit(intensity_expr.xreplace({s: s_expr}))\n", + " func = perform_cached_lambdify(expr, backend=\"jax\")\n", + " assert len(func.argument_order) == 2\n", + " intensity_funcs[k] = func" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Validation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "checksums = {\n", + " misc_key: {checksum[\"name\"]: checksum[\"value\"] for checksum in misc_value}\n", + " for misc_key, misc_value in MODEL_DEFINITION[\"misc\"].items()\n", + " if \"checksum\" in misc_key\n", + "}\n", + "checksums" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "checksum_points = {\n", + " point[\"name\"]: {par[\"name\"]: par[\"value\"] for par in point[\"parameters\"]}\n", + " for point in MODEL_DEFINITION[\"parameter_points\"]\n", + "}\n", + "checksum_points" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "array = []\n", + "for distribution_name, checksum in checksums.items():\n", + " for point_name, expected in checksum.items():\n", + " parameters = checksum_points[point_name]\n", + " s1 = parameters[\"m_31_2\"] ** 2\n", + " s2 = parameters[\"m_31\"] ** 2\n", + " computed = intensity_funcs[3]({\"sigma1\": s1, \"sigma2\": s2})\n", + " status = \"🟢\" if computed == expected else \"🔴\"\n", + " array.append((distribution_name, point_name, computed, expected, status))\n", + "pd.DataFrame(array, columns=[\"Distribution\", \"Point\", \"Computed\", \"Expected\", \"Status\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Dalitz plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "i, j = (2, 1)\n", + "k, *_ = {1, 2, 3} - {i, j}\n", + "σk, σk_expr = list(MODEL.invariants.items())[k - 1]\n", + "Math(aslatex({σk: σk_expr}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Define meshgrid for Dalitz plot" + }, + "tags": [ + "hide-input", + "remove-output" + ] + }, + "outputs": [], + "source": [ + "resolution = 1_000\n", + "m = sorted(MODEL.masses, key=str)\n", + "x_min = float(((m[j] + m[k]) ** 2).xreplace(MODEL.masses))\n", + "x_max = float(((m[0] - m[i]) ** 2).xreplace(MODEL.masses))\n", + "y_min = float(((m[i] + m[k]) ** 2).xreplace(MODEL.masses))\n", + "y_max = float(((m[0] - m[j]) ** 2).xreplace(MODEL.masses))\n", + "x_diff = x_max - x_min\n", + "y_diff = y_max - y_min\n", + "x_min -= 0.05 * x_diff\n", + "x_max += 0.05 * x_diff\n", + "y_min -= 0.05 * y_diff\n", + "y_max += 0.05 * y_diff\n", + "X, Y = jnp.meshgrid(\n", + " jnp.linspace(x_min, x_max, num=resolution),\n", + " jnp.linspace(y_min, y_max, num=resolution),\n", + ")\n", + "dalitz_data = {\n", + " f\"sigma{i}\": X,\n", + " f\"sigma{j}\": Y,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "mystnb": { + "code_prompt_show": "Prepare parametrized numerical function" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "intensities = intensity_funcs[k](dalitz_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "assert not jnp.all(jnp.isnan(intensities)), \"All intensities are NaN\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def get_decay_products(subsystem_id: int) -> tuple[State, State]:\n", + " return tuple(s for s in DECAY.final_state.values() if s.index != subsystem_id)\n", + "\n", + "\n", + "plt.rc(\"font\", size=18)\n", + "I_tot = jnp.nansum(intensities)\n", + "normalized_intensities = intensities / I_tot\n", + "\n", + "fig, ax = plt.subplots(figsize=(14, 10))\n", + "mesh = ax.pcolormesh(X, Y, normalized_intensities)\n", + "ax.set_aspect(\"equal\")\n", + "c_bar = plt.colorbar(mesh, ax=ax, pad=0.01)\n", + "c_bar.ax.set_ylabel(\"Normalized intensity (a.u.)\")\n", + "sigma_labels = {\n", + " i: Rf\"$\\sigma_{i} = M^2\\left({' '.join(p.latex for p in get_decay_products(i))}\\right)$\"\n", + " for i in (1, 2, 3)\n", + "}\n", + "ax.set_xlabel(sigma_labels[i])\n", + "ax.set_ylabel(sigma_labels[j])\n", + "plt.show()" + ] } ], "metadata": { diff --git a/pyproject.toml b/pyproject.toml index 49ba4334..1f9dd944 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,6 +54,7 @@ doc = [ "ipywidgets", "matplotlib", "myst-nb >=0.14", # nb_render_markdown_format for Markdown tables + "pandas", "sphinx-api-relink >=0.0.4", "sphinx-book-theme", "sphinx-codeautolink[ipython]", @@ -321,6 +322,13 @@ ignore-names = [ ] "docs/conf.py" = ["D100"] "jpsi2ksp.ipynb" = ["PLC2701"] +"serialization.ipynb" = [ + "E741", + "N813", + "PLC2403", + "PLC2701", + "RUF100", +] "setup.py" = ["D100"] "src/ampform_dpd/io.py" = ["S403"] "tests/*" = [ From 7ac413657517769648c9356d631c8174e4c7bd95 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:41 +0200 Subject: [PATCH 13/42] ENH: render L in energy-dependent width --- src/ampform_dpd/dynamics/__init__.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/dynamics/__init__.py b/src/ampform_dpd/dynamics/__init__.py index 4ccc83a0..e4415f73 100644 --- a/src/ampform_dpd/dynamics/__init__.py +++ b/src/ampform_dpd/dynamics/__init__.py @@ -2,13 +2,16 @@ from __future__ import annotations -from typing import Any +from typing import TYPE_CHECKING, Any import sympy as sp from ampform.dynamics import formulate_form_factor from ampform.kinematics.phasespace import Kallen from ampform.sympy import unevaluated +if TYPE_CHECKING: + from sympy.printing.latex import LatexPrinter + @unevaluated class RelativisticBreitWigner(sp.Expr): @@ -148,7 +151,6 @@ class EnergyDependentWidth(sp.Expr): m2: Any L: Any R: Any - _latex_repr_ = R"\Gamma\left({s}\right)" def evaluate(self): s, m0, Γ0, m1, m2, L, R = self.args @@ -164,6 +166,15 @@ def evaluate(self): evaluate=False, ) + def _latex_repr_(self, printer: LatexPrinter) -> str: + s = printer._print(self.s) # pyright:ignore[reportPrivateUsage] + if self.L == 0: + if self.m1 == 0 and self.m2 == 0: + return printer._print(self.Γ0) # pyright:ignore[reportPrivateUsage] + return Rf"\Gamma\left({s}\right)" + L = printer._print(self.L) # pyright:ignore[reportPrivateUsage] + return Rf"\Gamma_{{{L}}}\left({s}\right)" + @unevaluated class BlattWeisskopf(sp.Expr): From 526e9ebc7370a095813ad71f42b191035a414209 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:42 +0200 Subject: [PATCH 14/42] MAINT: convert `io` to subpackage --- pyproject.toml | 2 +- src/ampform_dpd/{io.py => io/__init__.py} | 0 2 files changed, 1 insertion(+), 1 deletion(-) rename src/ampform_dpd/{io.py => io/__init__.py} (100%) diff --git a/pyproject.toml b/pyproject.toml index 1f9dd944..2d2f4136 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -330,7 +330,7 @@ ignore-names = [ "RUF100", ] "setup.py" = ["D100"] -"src/ampform_dpd/io.py" = ["S403"] +"src/ampform_dpd/io/__init__.py" = ["S403"] "tests/*" = [ "D", "INP001", diff --git a/src/ampform_dpd/io.py b/src/ampform_dpd/io/__init__.py similarity index 100% rename from src/ampform_dpd/io.py rename to src/ampform_dpd/io/__init__.py From 2146562b6e37c3b11827343e43c051951d7c8b97 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:43 +0200 Subject: [PATCH 15/42] FEAT: implement `serialization.decay` module --- docs/conf.py | 3 + docs/serialization.ipynb | 1820 +----------------- pyproject.toml | 2 + src/ampform_dpd/io/serialization/__init__.py | 0 src/ampform_dpd/io/serialization/decay.py | 154 ++ src/ampform_dpd/io/serialization/format.py | 140 ++ tests/io_serialization/__init__.py | 0 tests/io_serialization/conftest.py | 18 + tests/io_serialization/test_decay.py | 45 + 9 files changed, 374 insertions(+), 1808 deletions(-) create mode 100644 src/ampform_dpd/io/serialization/__init__.py create mode 100644 src/ampform_dpd/io/serialization/decay.py create mode 100644 src/ampform_dpd/io/serialization/format.py create mode 100644 tests/io_serialization/__init__.py create mode 100644 tests/io_serialization/conftest.py create mode 100644 tests/io_serialization/test_decay.py diff --git a/docs/conf.py b/docs/conf.py index 009e4d43..41652019 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -44,6 +44,7 @@ "InitialStateID": ("obj", "ampform_dpd.decay.InitialStateID"), "Literal[-1, 1]": "typing.Literal", "Literal[(-1, 1)]": "typing.Literal", + "Node": ("obj", "ampform_dpd.io.serialization.format.Node"), "ParameterValue": ("obj", "tensorwaves.interface.ParameterValue"), "ParametrizedBackendFunction": "tensorwaves.function.ParametrizedBackendFunction", "PoolSum": "ampform.sympy.PoolSum", @@ -57,6 +58,8 @@ "sp.Symbol": "sympy.core.symbol.Symbol", "StateID": ("obj", "ampform_dpd.decay.StateID"), "StateIDTemplate": ("obj", "ampform_dpd.decay.StateID"), + "Topology": ("obj", "ampform_dpd.io.serialization.format.Topology"), + "typing_extensions.Required": ("obj", "typing.Required"), } api_target_types: dict[str, str] = {} author = "Common Partial Wave Analysis" diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index b23ff5cb..415a9090 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -2,7 +2,9 @@ "cells": [ { "cell_type": "markdown", - "metadata": {}, + "metadata": { + "tags": [] + }, "source": [ "# Model serialization" ] @@ -20,9 +22,6 @@ "cell_type": "code", "execution_count": null, "metadata": { - "jupyter": { - "source_hidden": true - }, "tags": [ "hide-input", "scroll-input" @@ -33,55 +32,11 @@ "from __future__ import annotations\n", "\n", "import json\n", - "from collections import abc\n", - "from difflib import get_close_matches\n", - "from itertools import product\n", - "from typing import Any, Callable, Literal, Sequence, Union\n", - "from warnings import warn\n", "\n", - "import jax.numpy as jnp\n", - "import matplotlib.pyplot as plt\n", - "import pandas as pd\n", - "import sympy as sp\n", - "from ampform.dynamics import BlattWeisskopfSquared\n", - "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import PoolSum, argument, perform_cached_doit, unevaluated\n", - "from attrs import asdict, frozen\n", "from IPython.display import JSON, Markdown, Math\n", - "from sympy.functions.special.tensor_functions import KroneckerDelta as δ\n", - "from sympy.physics.quantum.cg import CG\n", - "from sympy.physics.quantum.spin import Rotation as Wigner\n", - "from sympy.printing.latex import LatexPrinter\n", - "from tqdm.auto import tqdm\n", "\n", - "from ampform_dpd import (\n", - " AmplitudeModel,\n", - " DefinedExpression,\n", - " _AlignmentWignerGenerator,\n", - " _generate_amplitude_index_bases,\n", - " create_mass_symbol_mapping,\n", - " formulate_invariants,\n", - ")\n", - "from ampform_dpd.angles import formulate_scattering_angle\n", - "from ampform_dpd.decay import (\n", - " FinalStateID,\n", - " IsobarNode,\n", - " Particle,\n", - " State,\n", - " StateID,\n", - " ThreeBodyDecay,\n", - " ThreeBodyDecayChain,\n", - ")\n", - "from ampform_dpd.dynamics import BuggBreitWigner, EnergyDependentWidth, FormFactor, P\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 create_spin_range\n", - "\n", - "simplify_latex_rendering()" + "from ampform_dpd.io import as_markdown_table, aslatex\n", + "from ampform_dpd.io.serialization.decay import to_decay" ] }, { @@ -112,7 +67,6 @@ { "cell_type": "markdown", "metadata": { - "jp-MarkdownHeadingCollapsed": true, "tags": [] }, "source": [ @@ -126,117 +80,14 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Name-to-LaTeX converter" + }, "tags": [ - "hide-cell", - "scroll-input" + "hide-input" ] }, "outputs": [], - "source": [ - "def to_decay(\n", - " model: dict, to_latex: Callable[[str], str] | None = None\n", - ") -> ThreeBodyDecay:\n", - " initial_state = _get_initial_state(model)\n", - " final_state = _get_final_state(model)\n", - " return ThreeBodyDecay(\n", - " states=_get_states(model),\n", - " chains=sorted({\n", - " to_decay_chain(chain, initial_state, final_state, to_latex)\n", - " for chain in _get_decay_chains(model)\n", - " }),\n", - " )\n", - "\n", - "\n", - "def _get_decay_chains(model: dict) -> list[dict]:\n", - " distribution_def = _get_distribution_def(model)\n", - " return distribution_def[\"decay_description\"][\"chains\"]\n", - "\n", - "\n", - "def _get_distribution_def(model: dict) -> dict:\n", - " distribution_defs = MODEL_DEFINITION[\"distributions\"]\n", - " n_distributions = len(distribution_defs)\n", - " if n_distributions == 0:\n", - " msg = \"The serialized model does not have any distributions\"\n", - " raise ValueError(msg)\n", - " if n_distributions > 1:\n", - " msg = f\"There are {n_distributions} distributions, but expecting one only\"\n", - " warn(msg, category=UserWarning)\n", - " return distribution_defs[0]\n", - "\n", - "\n", - "def to_decay_chain(\n", - " chain_def: dict,\n", - " initial_state: State,\n", - " final_state: dict[FinalStateID, State],\n", - " to_latex: Callable[[str], str] | None = None,\n", - ") -> ThreeBodyDecayChain:\n", - " vertices = chain_def[\"vertices\"]\n", - " if to_latex is None:\n", - " to_latex = lambda x: x # noqa:E731\n", - " resonance = Particle(\n", - " name=chain_def[\"name\"],\n", - " latex=name_to_latex(chain_def[\"name\"]),\n", - " spin=chain_def[\"propagators\"][0][\"spin\"],\n", - " mass=0,\n", - " width=0,\n", - " parity=None,\n", - " )\n", - " return ThreeBodyDecayChain(\n", - " decay=IsobarNode(\n", - " parent=initial_state,\n", - " child1=IsobarNode(\n", - " parent=resonance,\n", - " child1=final_state[vertices[1][\"node\"][0]],\n", - " child2=final_state[vertices[1][\"node\"][1]],\n", - " ),\n", - " child2=final_state[vertices[0][\"node\"][1]],\n", - " )\n", - " )\n", - "\n", - "\n", - "def _get_states(model: dict) -> dict[StateID, State]:\n", - " initial_state = _get_initial_state(model)\n", - " final_state = _get_final_state(model)\n", - " return {initial_state.index: initial_state, **final_state}\n", - "\n", - "\n", - "def _get_initial_state(model: dict) -> dict[StateID, State]:\n", - " distribution_def = _get_distribution_def(model)\n", - " decay_description = distribution_def[\"decay_description\"]\n", - " kinematics = decay_description[\"kinematics\"]\n", - " return dict_to_particle(kinematics[\"initial_state\"])\n", - "\n", - "\n", - "def _get_final_state(model: dict) -> dict[StateID, State]:\n", - " distribution_def = _get_distribution_def(model)\n", - " decay_description = distribution_def[\"decay_description\"]\n", - " kinematics = decay_description[\"kinematics\"]\n", - " final_state_def = kinematics[\"final_state\"]\n", - " return {p[\"index\"]: dict_to_particle(p) for p in final_state_def}\n", - "\n", - "\n", - "def dict_to_particle(dct: dict) -> State:\n", - " return State(\n", - " name=dct[\"name\"],\n", - " latex=name_to_latex(dct[\"name\"]),\n", - " mass=dct[\"mass\"],\n", - " width=0,\n", - " spin=dct[\"spin\"],\n", - " parity=None,\n", - " index=dct[\"index\"],\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [] - }, - "outputs": [], "source": [ "def name_to_latex(name: str) -> str:\n", " latex = {\n", @@ -252,22 +103,18 @@ " subsystem = {\"D\": \"D\", \"K\": \"K\", \"L\": R\"\\Lambda\"}.get(subsystem_letter)\n", " if subsystem is None:\n", " return name\n", - " return f\"{subsystem}({mass_str})\"\n", - "\n", - "\n", - "DECAY = to_decay(MODEL_DEFINITION, to_latex=name_to_latex)" + " return f\"{subsystem}({mass_str})\"" ] }, { "cell_type": "code", "execution_count": null, "metadata": { - "tags": [ - "hide-input" - ] + "tags": [] }, "outputs": [], "source": [ + "DECAY = to_decay(MODEL_DEFINITION, to_latex=name_to_latex)\n", "Math(aslatex(DECAY))" ] }, @@ -283,1649 +130,6 @@ "source": [ "Markdown(as_markdown_table(DECAY))" ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "## Dynamics" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - ":::{seealso} [RUB-EP1/amplitude-serialization#22](https://github.com/RUB-EP1/amplitude-serialization/issues/22)\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Function look-up mechanism" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "function_defs = MODEL_DEFINITION[\"functions\"]\n", - "{f[\"type\"] for f in function_defs}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "CHAIN_DEFINITIONS = _get_decay_chains(MODEL_DEFINITION)\n", - "CHAIN_DEFINITIONS[2]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "def get_function_definition(function_name: str, model: dict) -> dict:\n", - " function_definitions = model[\"functions\"]\n", - " for function_def in function_definitions:\n", - " if function_def[\"name\"] == function_name:\n", - " return function_def\n", - " existing_names = {f[\"name\"] for f in function_definitions}\n", - " msg = f\"Could not find function with name {function_name!r}.\"\n", - " candidates = get_close_matches(function_name, existing_names)\n", - " if candidates:\n", - " msg += f\" Did you mean any of these? {', '.join(sorted(candidates))}\"\n", - " raise KeyError(msg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "get_function_definition(\"BlattWeiskopf\", MODEL_DEFINITION)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "### Vertices" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "#### Blatt-Weisskopf form factor" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "get_function_definition(\"BlattWeisskopf_resonance_l1\", MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "get_function_definition(\"BlattWeisskopf_resonance_l2\", MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _node_to_mass(node_item: int | Sequence[int]) -> sp.Symbol:\n", - " if isinstance(node_item, int):\n", - " return sp.Symbol(f\"m{node_item}\", nonnegative=True)\n", - " if (\n", - " isinstance(node_item, abc.Sequence)\n", - " and all(isinstance(i, int) for i in node_item)\n", - " and len(node_item) == 2\n", - " ):\n", - " k, *_ = {1, 2, 3} - set(node_item)\n", - " return sp.Symbol(f\"sigma{k}\", nonnegative=True)\n", - " msg = f\"Cannot create mass symbol for node {node_item}\"\n", - " raise NotImplementedError(msg)\n", - "\n", - "\n", - "sp.Tuple(_node_to_mass([2, 3]), _node_to_mass(1))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _node_to_mandelstam(node: Topology) -> sp.Symbol:\n", - " if all(isinstance(i, int) for i in node):\n", - " return _node_to_mass(node)\n", - " return _node_to_mass(0)\n", - "\n", - "\n", - "sp.Tuple(_node_to_mandelstam([2, 3]), _node_to_mandelstam([[2, 3], 1]))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "def formulate_form_factor(vertex: dict, model: dict) -> DefinedExpression:\n", - " function_name = vertex.get(\"formfactor\")\n", - " if not function_name:\n", - " return DefinedExpression()\n", - " function_definition = get_function_definition(function_name, model)\n", - " function_type = function_definition[\"type\"]\n", - " if function_type == \"BlattWeisskopf\":\n", - " node = vertex[\"node\"]\n", - " s = _node_to_mandelstam(node)\n", - " m1, m2 = (_node_to_mass(i) for i in node)\n", - " if all(isinstance(i, int) for i in node):\n", - " meson_radius = sp.Symbol(R\"R_\\mathrm{res}\")\n", - " else:\n", - " initial_state = _get_initial_state(model)\n", - " meson_radius = sp.Symbol(f\"R_{{{initial_state.latex}}}\")\n", - " angular_momentum = int(function_definition[\"l\"])\n", - " return DefinedExpression(\n", - " expression=FormFactor(s, m1, m2, angular_momentum, meson_radius),\n", - " definitions={\n", - " meson_radius: function_definition[\"radius\"],\n", - " },\n", - " )\n", - " msg = f\"No form factor implementation for {function_name!r}\"\n", - " raise NotImplementedError(msg)\n", - "\n", - "\n", - "CHAIN_2 = CHAIN_DEFINITIONS[2]\n", - "Math(aslatex(formulate_form_factor(CHAIN_2[\"vertices\"][0], MODEL_DEFINITION)))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s, m1, m2, d, L, z = sp.symbols(\"s m1 m2 d L z\")\n", - "exprs = [\n", - " FormFactor(s, m1, m2, L, d),\n", - " BlattWeisskopfSquared(z, L),\n", - "]\n", - "Math(aslatex({e: e.evaluate() for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Propagators" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "#### Breit-Wigner" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class BreitWigner(sp.Expr):\n", - " s: Any\n", - " mass: Any\n", - " width: Any\n", - " m1: Any = 0\n", - " m2: Any = 0\n", - " angular_momentum: Any = 0\n", - " meson_radius: Any = 1\n", - "\n", - " def evaluate(self):\n", - " width = self.energy_dependent_width()\n", - " expr = SimpleBreitWigner(self.s, self.mass, width)\n", - " if self.angular_momentum == 0 and self.m1 == 0 and self.m2 == 0:\n", - " return expr.evaluate()\n", - " return expr\n", - "\n", - " def energy_dependent_width(self) -> sp.Expr:\n", - " s, m0, Γ0, m1, m2, L, d = self.args\n", - " if L == 0 and m1 == 0 and m2 == 0:\n", - " return Γ0\n", - " return EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d)\n", - "\n", - " def _latex_repr_(self, printer: LatexPrinter, *args) -> str:\n", - " s = printer._print(self.s)\n", - " function_symbol = R\"\\mathcal{R}^\\mathrm{BW}\"\n", - " mass = printer._print(self.mass)\n", - " width = printer._print(self.width)\n", - " arg = Rf\"\\left({s}; {mass}, {width}\\right)\"\n", - " L = printer._print(self.angular_momentum)\n", - " if isinstance(self.angular_momentum, sp.Integer):\n", - " return Rf\"{function_symbol}_{{L={L}}}{arg}\"\n", - " return Rf\"{function_symbol}_{{{L}}}{arg}\"\n", - "\n", - "\n", - "@unevaluated\n", - "class SimpleBreitWigner(sp.Expr):\n", - " s: Any\n", - " mass: Any\n", - " width: Any\n", - " _latex_repr_ = R\"\\mathcal{{R}}^\\mathrm{{BW}}\\left({s}; {mass}, {width}\\right)\"\n", - "\n", - " def evaluate(self):\n", - " s, m0, Γ0 = self.args\n", - " return 1 / (m0**2 - s - m0 * Γ0 * 1j)\n", - "\n", - "\n", - "x, y, z = sp.symbols(\"x:z\")\n", - "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", - "exprs = [\n", - " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", - " SimpleBreitWigner(s, m0, Γ0),\n", - " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", - " FormFactor(s, m1, m2, L, d),\n", - " P(s, m1, m2),\n", - " Kallen(x, y, z),\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "CHAIN_DEFINITIONS[20]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "get_function_definition(\"K892_BW\", MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def _formulate_breit_wigner(\n", - " propagator: dict, function_definition: dict, resonance: str, **kwargs\n", - ") -> DefinedExpression:\n", - " node = propagator[\"node\"]\n", - " i, j = node\n", - " s = _node_to_mandelstam(node)\n", - " mass = sp.Symbol(f\"m_{{{resonance}}}\")\n", - " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\")\n", - " m1 = _node_to_mass(i)\n", - " m2 = _node_to_mass(j)\n", - " angular_momentum = int(function_definition[\"l\"])\n", - " d = sp.Symbol(R\"R_\\mathrm{res}\")\n", - " return DefinedExpression(\n", - " expression=BreitWigner(s, mass, width, m1, m2, angular_momentum, d),\n", - " definitions={\n", - " mass: function_definition[\"mass\"],\n", - " width: function_definition[\"width\"],\n", - " m1: function_definition[\"ma\"],\n", - " m2: function_definition[\"mb\"],\n", - " d: function_definition[\"d\"],\n", - " },\n", - " )\n", - "\n", - "\n", - "CHAIN_20 = CHAIN_DEFINITIONS[20]\n", - "K892_BW = _formulate_breit_wigner(\n", - " propagator=CHAIN_20[\"propagators\"][0],\n", - " function_definition=get_function_definition(\"K892_BW\", MODEL_DEFINITION),\n", - " resonance=name_to_latex(CHAIN_20[\"name\"]),\n", - ")\n", - "Math(aslatex(K892_BW))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "#### Multi-channel Breit-Wigner" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "@unevaluated\n", - "class MultichannelBreitWigner(sp.Expr):\n", - " s: Any\n", - " mass: Any\n", - " channels: list[ChannelArguments] = argument(sympify=False)\n", - "\n", - " def evaluate(self):\n", - " s = self.s\n", - " m0 = self.mass\n", - " width = sum(channel.formulate_width(s, m0) for channel in self.channels)\n", - " return BreitWigner(s, m0, width)\n", - "\n", - " def _latex_repr_(self, printer: LatexPrinter, *args) -> str:\n", - " latex = R\"\\mathcal{R}^\\mathrm{BW}_\\mathrm{multi}\\left(\"\n", - " latex += printer._print(self.s) + \"; \"\n", - " latex += \", \".join(printer._print(channel.width) for channel in self.channels)\n", - " latex += R\"\\right)\"\n", - " return latex\n", - "\n", - "\n", - "@frozen\n", - "class ChannelArguments:\n", - " width: Any\n", - " m1: Any = 0\n", - " m2: Any = 0\n", - " angular_momentum: Any = 0\n", - " meson_radius: Any = 1\n", - "\n", - " def __attrs_post_init__(self) -> None:\n", - " for name, value in asdict(self).items():\n", - " object.__setattr__(self, name, sp.sympify(value))\n", - "\n", - " def formulate_width(self, s: Any, m0: Any) -> sp.Expr:\n", - " Γ0 = self.width\n", - " m1 = self.m1\n", - " m2 = self.m2\n", - " L = self.angular_momentum\n", - " R = self.meson_radius\n", - " ff = FormFactor(s, m1, m2, L, R) ** 2\n", - " return Γ0 * m0 / sp.sqrt(s) * ff\n", - "\n", - "\n", - "x, y, z = sp.symbols(\"x:z\")\n", - "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", - "channels = [\n", - " ChannelArguments(\n", - " sp.Symbol(f\"Gamma{i}\"),\n", - " sp.Symbol(f\"m_{{a,{i}}}\"),\n", - " sp.Symbol(f\"m_{{b,{i}}}\"),\n", - " sp.Symbol(f\"L{i}\"),\n", - " d,\n", - " )\n", - " for i in [1, 2]\n", - "]\n", - "exprs = [\n", - " MultichannelBreitWigner(s, m0, channels),\n", - " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", - " BreitWigner(s, m0, Γ0),\n", - " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", - " FormFactor(s, m1, m2, L, d),\n", - " P(s, m1, m2),\n", - " Kallen(x, y, z),\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "CHAIN_DEFINITIONS[0]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "get_function_definition(\"L1405_Flatte\", MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def _formulate_multichannel_breit_wigner(\n", - " propagator: dict, function_definition: dict, resonance: str, **kwargs\n", - ") -> DefinedExpression:\n", - " channel_definitions = function_definition[\"channels\"]\n", - " if len(channel_definitions) < 2:\n", - " msg = \"Need at least two channels for a multi-channel Breit-Wigner\"\n", - " raise NotImplementedError(msg)\n", - " node = propagator[\"node\"]\n", - " i, j = node\n", - " s = _node_to_mandelstam(node)\n", - " resonance_latex = name_to_latex(resonance)\n", - " mass = sp.Symbol(f\"m_{{{resonance}}}\")\n", - " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\")\n", - " m1 = _node_to_mass(i)\n", - " m2 = _node_to_mass(j)\n", - " angular_momentum = int(channel_definitions[0][\"l\"])\n", - " d = sp.Symbol(f\"R_{{{resonance_latex}}}\")\n", - " channels = [ChannelArguments(width, m1, m2, angular_momentum, d)]\n", - " parameter_defaults = {\n", - " mass: function_definition[\"mass\"],\n", - " width: channel_definitions[0][\"gsq\"],\n", - " m1: channel_definitions[0][\"ma\"],\n", - " m2: channel_definitions[0][\"mb\"],\n", - " d: channel_definitions[0][\"d\"],\n", - " }\n", - " for i, channel_definition in enumerate(channel_definitions[1:], 2):\n", - " Γi = sp.Symbol(\n", - " Rf\"\\Gamma_{{{resonance_latex}}}^\\text{{ch. {i}}}\", nonnegative=True\n", - " )\n", - " mi1 = sp.Symbol(f\"m_{{a,{i}}}\", nonnegative=True)\n", - " mi2 = sp.Symbol(f\"m_{{b,{i}}}\", nonnegative=True)\n", - " angular_momentum = int(channel_definition[\"l\"])\n", - " channels.append(ChannelArguments(Γi, mi1, mi2, angular_momentum, d))\n", - " parameter_defaults.update({\n", - " mi1: channel_definition[\"ma\"],\n", - " mi2: channel_definition[\"mb\"],\n", - " })\n", - " return DefinedExpression(\n", - " expression=MultichannelBreitWigner(s, mass, channels),\n", - " definitions=parameter_defaults,\n", - " )\n", - "\n", - "\n", - "CHAIN_0 = CHAIN_DEFINITIONS[0]\n", - "L1405_Flatte = _formulate_multichannel_breit_wigner(\n", - " propagator=CHAIN_0[\"propagators\"][0],\n", - " function_definition=get_function_definition(\"L1405_Flatte\", MODEL_DEFINITION),\n", - " resonance=name_to_latex(CHAIN_0[\"name\"]),\n", - ")\n", - "Math(aslatex(L1405_Flatte))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "#### Breit-Wigner with exponential" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s, m0, Γ0, m1, m2, γ = sp.symbols(\"s m0 Gamma0 m1 m2 gamma\")\n", - "expr = BuggBreitWigner(s, m0, Γ0, m1, m2, γ)\n", - "Math(aslatex({expr: expr.doit(deep=False)}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "CHAIN_DEFINITIONS[18]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "get_function_definition(\"K700_BuggBW\", MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _formulate_bugg_breit_wigner(\n", - " propagator: dict, function_definition: dict, resonance: str, model: dict, **kwargs\n", - ") -> DefinedExpression:\n", - " node = propagator[\"node\"]\n", - " i, j = node\n", - " s = _node_to_mandelstam(node)\n", - " mass = sp.Symbol(f\"m_{{{resonance}}}\", nonnegative=True)\n", - " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\", nonnegative=True)\n", - " γ = sp.Symbol(Rf\"\\gamma_{{{resonance}}}\")\n", - " m1 = _node_to_mass(i)\n", - " m2 = _node_to_mass(j)\n", - " final_state = _get_final_state(model)\n", - " return DefinedExpression(\n", - " expression=BuggBreitWigner(s, mass, width, m1, m2, γ),\n", - " definitions={\n", - " mass: function_definition[\"mass\"],\n", - " width: function_definition[\"width\"],\n", - " m1: final_state[i].mass,\n", - " m2: final_state[j].mass,\n", - " γ: function_definition[\"slope\"],\n", - " },\n", - " )\n", - "\n", - "\n", - "CHAIN_18 = CHAIN_DEFINITIONS[18]\n", - "K700_BuggBW = _formulate_bugg_breit_wigner(\n", - " propagator=CHAIN_0[\"propagators\"][0],\n", - " function_definition=get_function_definition(\"K700_BuggBW\", MODEL_DEFINITION),\n", - " resonance=name_to_latex(CHAIN_0[\"name\"]),\n", - " model=MODEL_DEFINITION,\n", - ")\n", - "Math(aslatex(K700_BuggBW))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### General propagator dynamics builder" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def formulate_dynamics(chain_definition: dict, model: dict) -> DefinedExpression:\n", - " expr = DefinedExpression()\n", - " for propagator in chain_definition[\"propagators\"]:\n", - " parametrization = propagator[\"parametrization\"]\n", - " function_definition = get_function_definition(parametrization, model)\n", - " function_type = function_definition[\"type\"]\n", - " if function_type == \"BreitWigner\":\n", - " dynamics_builder = _formulate_breit_wigner\n", - " elif function_type == \"MultichannelBreitWigner\":\n", - " dynamics_builder = _formulate_multichannel_breit_wigner\n", - " elif function_type == \"BreitWignerWidthExpLikeBugg\":\n", - " dynamics_builder = _formulate_bugg_breit_wigner\n", - " else:\n", - " msg = f\"No dynamics implementation for function type {function_type!r}\"\n", - " raise NotImplementedError(msg)\n", - " expr *= dynamics_builder(\n", - " propagator,\n", - " function_definition=get_function_definition(parametrization, model),\n", - " resonance=name_to_latex(chain_definition[\"name\"]),\n", - " model=model,\n", - " )\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "Math(aslatex(formulate_dynamics(CHAIN_DEFINITIONS[0], MODEL_DEFINITION)))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "Math(aslatex(formulate_dynamics(CHAIN_DEFINITIONS[18], MODEL_DEFINITION)))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "Math(aslatex(formulate_dynamics(CHAIN_DEFINITIONS[20], MODEL_DEFINITION)))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "## Construct `AmplitudeModel`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Unpolarized intensity" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [] - }, - "outputs": [], - "source": [ - "Topology = list[Union[int, \"Topology\"]]\n", - "\n", - "\n", - "def get_reference_subsystem(model: dict) -> FinalStateID:\n", - " topology = get_reference_topology(model)\n", - " return get_spectator_id(topology)\n", - "\n", - "\n", - "def get_spectator_id(topology: Topology) -> FinalStateID:\n", - " spectator_candidates = {i for i in topology if isinstance(i, int)}\n", - " if len(spectator_candidates) != 1:\n", - " msg = f\"Reference topology {topology} seems not to be a three-body decay\"\n", - " raise ValueError(msg)\n", - " return next(iter(spectator_candidates))\n", - "\n", - "\n", - "def get_reference_topology(model: dict) -> Topology:\n", - " distribution_def = _get_distribution_def(model)\n", - " return distribution_def[\"decay_description\"][\"reference_topology\"]\n", - "\n", - "\n", - "get_reference_subsystem(MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [] - }, - "outputs": [], - "source": [ - "def get_existing_subsystem_ids(model: dict) -> list[FinalStateID]:\n", - " distribution_def = _get_distribution_def(model)\n", - " chain_defs = distribution_def[\"decay_description\"][\"chains\"]\n", - " subsystem_ids = {get_spectator_id(c[\"topology\"]) for c in chain_defs}\n", - " return sorted(subsystem_ids)\n", - "\n", - "\n", - "get_existing_subsystem_ids(MODEL_DEFINITION)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [] - }, - "outputs": [], - "source": [ - "def formulate_aligned_amplitude(\n", - " model: dict,\n", - " λ0: sp.Rational | sp.Symbol,\n", - " λ1: sp.Rational | sp.Symbol,\n", - " λ2: sp.Rational | sp.Symbol,\n", - " λ3: sp.Rational | sp.Symbol,\n", - ") -> tuple[PoolSum, dict[sp.Symbol, sp.Expr]]:\n", - " reference_subsystem = get_reference_subsystem(model)\n", - " wigner_generator = _AlignmentWignerGenerator(reference_subsystem)\n", - " _λ0, _λ1, _λ2, _λ3 = sp.symbols(R\"\\lambda_(:4)^{\\prime}\", rational=True)\n", - " distribution_def = _get_distribution_def(model)\n", - " decay_description = distribution_def[\"decay_description\"]\n", - " states = _get_states(decay_description[\"kinematics\"])\n", - " j0, j1, j2, j3 = (states[i].spin for i in sorted(states))\n", - " A = _generate_amplitude_index_bases()\n", - " amp_expr = PoolSum(\n", - " sum(\n", - " A[k][_λ0, _λ1, _λ2, _λ3]\n", - " * wigner_generator(j0, λ0, _λ0, rotated_state=0, aligned_subsystem=k)\n", - " * wigner_generator(j1, _λ1, λ1, rotated_state=1, aligned_subsystem=k)\n", - " * wigner_generator(j2, _λ2, λ2, rotated_state=2, aligned_subsystem=k)\n", - " * wigner_generator(j3, _λ3, λ3, rotated_state=3, aligned_subsystem=k)\n", - " for k in get_existing_subsystem_ids(model)\n", - " ),\n", - " (_λ0, create_spin_range(j0)),\n", - " (_λ1, create_spin_range(j1)),\n", - " (_λ2, create_spin_range(j2)),\n", - " (_λ3, create_spin_range(j3)),\n", - " )\n", - " return amp_expr, wigner_generator.angle_definitions\n", - "\n", - "\n", - "λ0, λ1, λ2, λ3 = sp.symbols(\"lambda(:4)\", rational=True)\n", - "amplitude_expr, _ = formulate_aligned_amplitude(MODEL_DEFINITION, λ0, λ1, λ2, λ3)\n", - "amplitude_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Amplitude for the decay chain" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Helicity recouplings" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class HelicityRecoupling(sp.Expr):\n", - " λa: sp.Rational | sp.Symbol\n", - " λb: sp.Rational | sp.Symbol\n", - " λa0: sp.Rational | sp.Symbol\n", - " λb0: sp.Rational | sp.Symbol\n", - " _latex_repr_ = R\"\\mathcal{{H}}^\\text{{helicity}}\\left({λa},{λb}|{λa0},{λb0}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " λa, λb, λa0, λb0 = self.args\n", - " return δ(λa, λa0) * δ(λb, λb0)\n", - "\n", - "\n", - "@unevaluated\n", - "class ParityRecoupling(sp.Expr):\n", - " λa: sp.Rational | sp.Symbol\n", - " λb: sp.Rational | sp.Symbol\n", - " λa0: sp.Rational | sp.Symbol\n", - " λb0: sp.Rational | sp.Symbol\n", - " f: sp.Integer | sp.Symbol\n", - " _latex_repr_ = (\n", - " R\"\\mathcal{{H}}^\\text{{parity}}\\left({λa},{λb}|{λa0},{λb0},{f}\\right)\"\n", - " )\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " λa, λb, λa0, λb0, f = self.args\n", - " return δ(λa, λa0) * δ(λb, λb0) + f * δ(λa, -λa0) * δ(λb, -λb0)\n", - "\n", - "\n", - "@unevaluated\n", - "class LSRecoupling(sp.Expr):\n", - " λa: sp.Rational | sp.Symbol\n", - " λb: sp.Rational | sp.Symbol\n", - " l: sp.Integer | sp.Symbol\n", - " s: sp.Rational | sp.Symbol\n", - " ja: sp.Rational | sp.Symbol\n", - " jb: sp.Rational | sp.Symbol\n", - " j: sp.Rational | sp.Symbol\n", - " _latex_repr_ = (\n", - " R\"\\mathcal{{H}}^\\text{{parity}}\\left({λa},{λb}|{l},{s},{ja},{jb},{j}\\right)\"\n", - " )\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " λa, λb, l, s, ja, jb, j = self.args\n", - " return (\n", - " sp.sqrt((2 * l + 1) / (2 * j + 1))\n", - " * CG(ja, λa, jb, -λb, s, λa - λb)\n", - " * CG(l, 0, s, λa - λb, j, λa - λb)\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "λa = sp.Symbol(R\"\\lambda_a\", rational=True)\n", - "λb = sp.Symbol(R\"\\lambda_b\", rational=True)\n", - "λa0 = sp.Symbol(R\"\\lambda_a^0\", rational=True)\n", - "λb0 = sp.Symbol(R\"\\lambda_b^0\", rational=True)\n", - "f = sp.Symbol(\"f\", integer=True)\n", - "l = sp.Symbol(\"l\", integer=True, nonnegative=True)\n", - "s = sp.Symbol(\"s\", nonnegative=True, rational=True)\n", - "ja = sp.Symbol(\"j_a\", nonnegative=True, rational=True)\n", - "jb = sp.Symbol(\"j_b\", nonnegative=True, rational=True)\n", - "j = sp.Symbol(\"j\", nonnegative=True, rational=True)\n", - "exprs = [\n", - " HelicityRecoupling(λa, λb, λa0, λb0),\n", - " ParityRecoupling(λa, λb, λa0, λb0, f),\n", - " LSRecoupling(λa, λb, l, s, ja, jb, j),\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Recoupling deserialization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "CHAIN_DEFINITIONS[0]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _formulate_recoupling(model: dict, chain_idx: int, vertex_idx: int) -> sp.Expr:\n", - " chain_definition = _get_decay_chains(model)[chain_idx]\n", - " vertex_definitions = chain_definition[\"vertices\"]\n", - " if len(vertex_definitions) != 2:\n", - " msg = f\"Not a three-body decay: there are {len(vertex_definitions)} vertices\"\n", - " raise ValueError(msg)\n", - " if vertex_idx not in {0, 1}:\n", - " msg = f\"Vertex index out of range. Can either be 0 or 1, not {vertex_idx}.\"\n", - " raise ValueError(msg)\n", - " vertex = chain_definition[\"vertices\"][vertex_idx]\n", - " vertex_type = vertex[\"type\"]\n", - " node = vertex[\"node\"]\n", - " λa, λb = map(_get_helicity_symbol, node)\n", - " if vertex_type in {\"helicity\", \"parity\"}:\n", - " λa0, λb0 = (sp.Rational(v) for v in vertex[\"helicities\"])\n", - " if vertex_type == \"parity\":\n", - " f = _sign_to_value(vertex.get(\"parity_factor\", \"+\"))\n", - " return ParityRecoupling(λa, λb, λa0, λb0, f)\n", - " return HelicityRecoupling(λa, λb, λa0, λb0)\n", - " if vertex_type == \"ls\":\n", - " l = int(vertex[\"l\"])\n", - " s = sp.Rational(vertex[\"s\"])\n", - " ja, jb = _get_child_spins(model, chain_idx, vertex_idx)\n", - " j = _get_parent_spin(vertex)\n", - " return LSRecoupling(λa, λb, l, s, ja, jb, j)\n", - " msg = f\"No implementation for vertex of type {vertex_type!r}\"\n", - " raise NotImplementedError(msg)\n", - "\n", - "\n", - "def _sign_to_value(sign: Literal[\"\", \"-\", \"+\"]) -> Literal[0, -1, 1]:\n", - " stripped_sign = sign.strip()\n", - " if stripped_sign == \"-\":\n", - " return -1\n", - " if not stripped_sign:\n", - " return 0\n", - " if stripped_sign == \"+\":\n", - " return +1\n", - " msg = f\"Cannot convert {sign!r} to value\"\n", - " raise NotImplementedError(msg)\n", - "\n", - "\n", - "def _get_parent_spin(\n", - " model: dict, chain_idx: int, vertex_idx: int\n", - ") -> tuple[sp.Rational, sp.Rational]:\n", - " chain_definition = _get_decay_chains(model)[chain_idx]\n", - " vertex = chain_definition[\"vertices\"]\n", - " node = vertex[\"node\"]\n", - " initial_state = _get_initial_state(model)\n", - " if all(isinstance(i, int) for i in node):\n", - " return __get_propagator_spin(vertex)\n", - " return initial_state.spin\n", - "\n", - "\n", - "def _get_child_spins(\n", - " model: dict, chain_idx: int, vertex_idx: int\n", - ") -> tuple[sp.Rational, sp.Rational]:\n", - " chain_definition = _get_decay_chains(model)[chain_idx]\n", - " vertex = chain_definition[\"vertices\"]\n", - " node = vertex[\"node\"]\n", - " final_state = _get_final_state(model)\n", - " spins = []\n", - " for node_item in node:\n", - " if isinstance(node_item, int):\n", - " spins.append(sp.Rational(final_state[node_item]))\n", - " else:\n", - " spins.append(__get_propagator_spin(vertex[\"propagators\"]))\n", - " return tuple(spins)\n", - "\n", - "\n", - "def __get_propagator_spin(vertex: dict) -> sp.Rational:\n", - " propagators = vertex[\"propagators\"]\n", - " if len(propagators) != 1:\n", - " msg = f\"There are {len(propagators)} propagators, not a three-body decay\"\n", - " raise ValueError(msg)\n", - " return sp.Rational(propagators[0][\"spin\"])\n", - "\n", - "\n", - "def _get_helicity_symbol(node: int | Topology) -> sp.Symbol:\n", - " if isinstance(node, int):\n", - " return sp.Symbol(f\"lambda{node}\", rational=True)\n", - " return sp.Symbol(R\"\\lambda_R\", rational=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-output" - ] - }, - "outputs": [], - "source": [ - "recouplings = [\n", - " _formulate_recoupling(MODEL_DEFINITION, chain_idx=0, vertex_idx=i) for i in range(2)\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in recouplings}))" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Chain amplitudes" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _get_final_state_helicities(\n", - " chain_definition: dict,\n", - ") -> dict[FinalStateID, sp.Rational]:\n", - " vertices = chain_definition[\"vertices\"]\n", - " helicities: dict[FinalStateID, sp.Rational] = {}\n", - " for v in vertices:\n", - " for helicity, node in zip(v[\"helicities\"], v[\"node\"]):\n", - " if not isinstance(node, int):\n", - " continue\n", - " helicities[node] = sp.Rational(helicity)\n", - " return {i: helicities[i] for i in sorted(helicities)}\n", - "\n", - "\n", - "_get_final_state_helicities(CHAIN_DEFINITIONS[0])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _get_resonance_helicity(\n", - " chain_definition: dict,\n", - ") -> tuple[tuple[FinalStateID, FinalStateID], sp.Rational]:\n", - " vertices = chain_definition[\"vertices\"]\n", - " for vertex in vertices:\n", - " node = vertex[\"node\"]\n", - " if all(isinstance(i, int) for i in node):\n", - " continue\n", - " helicities = vertex[\"helicities\"]\n", - " for helicity, sub_node in zip(helicities, node):\n", - " if isinstance(sub_node, abc.Sequence) and len(sub_node) == 2:\n", - " return tuple(sub_node), sp.Rational(helicity)\n", - " msg = \"Could not find a resonance node\"\n", - " raise ValueError(msg)\n", - "\n", - "\n", - "_get_resonance_helicity(CHAIN_DEFINITIONS[0])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _get_spectator_helicity(chain_definition: dict) -> tuple[int, sp.Rational]:\n", - " vertices = chain_definition[\"vertices\"]\n", - " for vertex in vertices:\n", - " node = vertex[\"node\"]\n", - " if all(isinstance(i, int) for i in node):\n", - " continue\n", - " helicities = vertex[\"helicities\"]\n", - " for helicity, state_id in zip(helicities, node):\n", - " if isinstance(state_id, int):\n", - " return state_id, sp.Rational(helicity)\n", - " msg = \"Could not find vertex that contains a spectator\"\n", - " raise ValueError(msg)\n", - "\n", - "\n", - "_get_spectator_helicity(CHAIN_DEFINITIONS[0])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _get_decay_product_helicities(\n", - " chain_definition: dict,\n", - ") -> tuple[tuple[int, sp.Rational], tuple[int, sp.Rational]]:\n", - " vertices = chain_definition[\"vertices\"]\n", - " for v in vertices:\n", - " node = v[\"node\"]\n", - " if all(isinstance(i, int) for i in node):\n", - " helicities = v[\"helicities\"]\n", - " return tuple((i, sp.Rational(λ)) for i, λ in zip(node, helicities))\n", - " msg = \"Could not fine a helicity for any resonance node\"\n", - " raise ValueError(msg)\n", - "\n", - "\n", - "_get_decay_product_helicities(CHAIN_DEFINITIONS[0])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def _get_weight(chain_definition: dict) -> tuple[sp.Symbol, complex | float]:\n", - " value: complex | float\n", - " value = complex(str(chain_definition[\"weight\"]).replace(\" \", \"\").replace(\"i\", \"j\"))\n", - " if not value.imag:\n", - " value = value.real\n", - " resonance_latex = name_to_latex(chain_definition[\"name\"])\n", - " _, resonance_helicity = _get_resonance_helicity(chain_definition)\n", - " c = sp.IndexedBase(f\"c^{{{resonance_latex}[{resonance_helicity}]}}\")\n", - " λ1, λ2, λ3 = _get_final_state_helicities(chain_definition).values()\n", - " symbol = c[λ1, λ2, λ3]\n", - " return symbol, value\n", - "\n", - "\n", - "Math(aslatex(dict([_get_weight(CHAIN_DEFINITIONS[0])])))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def formulate_chain_amplitude(\n", - " λ0: sp.Rational,\n", - " λ1: sp.Rational,\n", - " λ2: sp.Rational,\n", - " λ3: sp.Rational,\n", - " model: dict,\n", - " chain_idx: int,\n", - ") -> sp.Expr:\n", - " chain_defs = _get_decay_chains(model)\n", - " chain_definition = chain_defs[chain_idx]\n", - " # -----------------------\n", - " dynamics = formulate_dynamics(chain_definition, model)\n", - " for vertex in chain_definition[\"vertices\"]:\n", - " dynamics *= formulate_form_factor(vertex, model)\n", - " # -----------------------\n", - " weight, weight_val = _get_weight(chain_definition)\n", - " # -----------------------\n", - " (i, λi_val), (j, λj_val) = _get_decay_product_helicities(chain_definition)\n", - " θij, θij_expr = formulate_scattering_angle(i, j)\n", - " jR = sp.Rational(chain_definition[\"propagators\"][0][\"spin\"])\n", - " R_node, λR_val = _get_resonance_helicity(chain_definition)\n", - " λR = _get_helicity_symbol(R_node)\n", - " # -----------------------\n", - " A = _generate_amplitude_index_bases()\n", - " subsystem_id = get_spectator_id(chain_definition[\"topology\"])\n", - " h_prod = _formulate_recoupling(chain_definition, chain_idx, vertex_idx=0)\n", - " h_dec = _formulate_recoupling(chain_definition, chain_idx, vertex_idx=1)\n", - " amplitude_expression = (\n", - " weight\n", - " * h_prod\n", - " * h_dec\n", - " * Wigner.d(jR, λR, λi_val - λj_val, θij)\n", - " * dynamics.expression\n", - " )\n", - " amplitude_expression = amplitude_expression.subs({λR: λR_val})\n", - " amplitude_symbol = A[subsystem_id][λ0, λ1, λ2, λ3]\n", - " return {\n", - " amplitude_symbol: amplitude_expression,\n", - " weight: weight_val,\n", - " **dynamics.definitions,\n", - " θij: θij_expr,\n", - " }" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "definitions = formulate_chain_amplitude(λ0, λ1, λ2, λ3, MODEL_DEFINITION, chain_idx=0)\n", - "Math(aslatex(definitions))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Full amplitude model" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def formulate(model: dict, cleanup_summations: bool = False) -> AmplitudeModel:\n", - " states = _get_states(model)\n", - " helicity_symbols = sp.symbols(\"lambda(:4)\", rational=True)\n", - " allowed_helicities = {\n", - " symbol: create_spin_range(states[i].spin)\n", - " for i, symbol in enumerate(helicity_symbols)\n", - " }\n", - " amplitude_definitions = {}\n", - " angle_definitions = {}\n", - " parameter_defaults = {}\n", - " n_chains = len(_get_decay_chains(model))\n", - " helicity_values: tuple[sp.Rational, sp.Rational, sp.Rational, sp.Rational]\n", - " for helicity_values in product(*allowed_helicities.values()):\n", - " for chain_idx in range(n_chains):\n", - " amp_defs = formulate_chain_amplitude(*helicity_values, model, chain_idx)\n", - " (amp_symbol, amp_expr), *parameters, (θij, θij_expr) = amp_defs.items()\n", - " helicity_substitutions = dict(zip(helicity_symbols, helicity_values))\n", - " amplitude_definitions[amp_symbol] = amp_expr.subs(helicity_substitutions)\n", - " angle_definitions[θij] = θij_expr\n", - " parameter_defaults.update(dict(parameters))\n", - " aligned_amp, zeta_defs = formulate_aligned_amplitude(model, *helicity_symbols)\n", - " angle_definitions.update(zeta_defs)\n", - " decay = to_decay(model)\n", - " masses = create_mass_symbol_mapping(decay)\n", - " parameter_defaults.update(masses)\n", - " if cleanup_summations:\n", - " aligned_amp = aligned_amp.cleanup() # type:ignore[assignment]\n", - " intensity = PoolSum(\n", - " sp.Abs(aligned_amp) ** 2,\n", - " *allowed_helicities.items(),\n", - " )\n", - " if cleanup_summations:\n", - " intensity = intensity.cleanup() # type:ignore[assignment]\n", - " return AmplitudeModel(\n", - " decay=decay,\n", - " intensity=intensity,\n", - " amplitudes=amplitude_definitions,\n", - " variables=angle_definitions,\n", - " parameter_defaults=parameter_defaults,\n", - " masses=masses,\n", - " invariants=formulate_invariants(decay),\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "## Symbolic result" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MODEL = formulate(MODEL_DEFINITION, cleanup_summations=True)\n", - "MODEL.intensity" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "Math(aslatex(MODEL.amplitudes))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "Math(aslatex(MODEL.variables))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "Math(aslatex({**MODEL.invariants, **MODEL.masses}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensity_expr = MODEL.full_expression.xreplace(MODEL.variables)\n", - "intensity_expr = intensity_expr.xreplace(MODEL.parameter_defaults)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "free_symbols = intensity_expr.free_symbols\n", - "assert len(free_symbols) == 3\n", - "assert str(sorted(free_symbols, key=str)) == \"[sigma1, sigma2, sigma3]\"" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "## Numeric results" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [] - }, - "outputs": [], - "source": [ - "intensity_funcs = {}\n", - "for s, s_expr in tqdm(MODEL.invariants.items()):\n", - " k = int(str(s)[-1])\n", - " s_expr = s_expr.xreplace(MODEL.masses).doit()\n", - " expr = perform_cached_doit(intensity_expr.xreplace({s: s_expr}))\n", - " func = perform_cached_lambdify(expr, backend=\"jax\")\n", - " assert len(func.argument_order) == 2\n", - " intensity_funcs[k] = func" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "### Validation" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "checksums = {\n", - " misc_key: {checksum[\"name\"]: checksum[\"value\"] for checksum in misc_value}\n", - " for misc_key, misc_value in MODEL_DEFINITION[\"misc\"].items()\n", - " if \"checksum\" in misc_key\n", - "}\n", - "checksums" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [] - }, - "outputs": [], - "source": [ - "checksum_points = {\n", - " point[\"name\"]: {par[\"name\"]: par[\"value\"] for par in point[\"parameters\"]}\n", - " for point in MODEL_DEFINITION[\"parameter_points\"]\n", - "}\n", - "checksum_points" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "array = []\n", - "for distribution_name, checksum in checksums.items():\n", - " for point_name, expected in checksum.items():\n", - " parameters = checksum_points[point_name]\n", - " s1 = parameters[\"m_31_2\"] ** 2\n", - " s2 = parameters[\"m_31\"] ** 2\n", - " computed = intensity_funcs[3]({\"sigma1\": s1, \"sigma2\": s2})\n", - " status = \"🟢\" if computed == expected else \"🔴\"\n", - " array.append((distribution_name, point_name, computed, expected, status))\n", - "pd.DataFrame(array, columns=[\"Distribution\", \"Point\", \"Computed\", \"Expected\", \"Status\"])" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [] - }, - "source": [ - "### Dalitz plot" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "i, j = (2, 1)\n", - "k, *_ = {1, 2, 3} - {i, j}\n", - "σk, σk_expr = list(MODEL.invariants.items())[k - 1]\n", - "Math(aslatex({σk: σk_expr}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define meshgrid for Dalitz plot" - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "resolution = 1_000\n", - "m = sorted(MODEL.masses, key=str)\n", - "x_min = float(((m[j] + m[k]) ** 2).xreplace(MODEL.masses))\n", - "x_max = float(((m[0] - m[i]) ** 2).xreplace(MODEL.masses))\n", - "y_min = float(((m[i] + m[k]) ** 2).xreplace(MODEL.masses))\n", - "y_max = float(((m[0] - m[j]) ** 2).xreplace(MODEL.masses))\n", - "x_diff = x_max - x_min\n", - "y_diff = y_max - y_min\n", - "x_min -= 0.05 * x_diff\n", - "x_max += 0.05 * x_diff\n", - "y_min -= 0.05 * y_diff\n", - "y_max += 0.05 * y_diff\n", - "X, Y = jnp.meshgrid(\n", - " jnp.linspace(x_min, x_max, num=resolution),\n", - " jnp.linspace(y_min, y_max, num=resolution),\n", - ")\n", - "dalitz_data = {\n", - " f\"sigma{i}\": X,\n", - " f\"sigma{j}\": Y,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Prepare parametrized numerical function" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "intensities = intensity_funcs[k](dalitz_data)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "assert not jnp.all(jnp.isnan(intensities)), \"All intensities are NaN\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def get_decay_products(subsystem_id: int) -> tuple[State, State]:\n", - " return tuple(s for s in DECAY.final_state.values() if s.index != subsystem_id)\n", - "\n", - "\n", - "plt.rc(\"font\", size=18)\n", - "I_tot = jnp.nansum(intensities)\n", - "normalized_intensities = intensities / I_tot\n", - "\n", - "fig, ax = plt.subplots(figsize=(14, 10))\n", - "mesh = ax.pcolormesh(X, Y, normalized_intensities)\n", - "ax.set_aspect(\"equal\")\n", - "c_bar = plt.colorbar(mesh, ax=ax, pad=0.01)\n", - "c_bar.ax.set_ylabel(\"Normalized intensity (a.u.)\")\n", - "sigma_labels = {\n", - " i: Rf\"$\\sigma_{i} = M^2\\left({' '.join(p.latex for p in get_decay_products(i))}\\right)$\"\n", - " for i in (1, 2, 3)\n", - "}\n", - "ax.set_xlabel(sigma_labels[i])\n", - "ax.set_ylabel(sigma_labels[j])\n", - "plt.show()" - ] } ], "metadata": { diff --git a/pyproject.toml b/pyproject.toml index 2d2f4136..7f663b5a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -32,6 +32,7 @@ dependencies = [ "qrules >=0.10.0", "sympy >=1.10", # module sympy.printing.numpy and array expressions with shape kwarg "tensorwaves[jax]", + 'typing-extensions; python_version <"3.11.0"', ] description = "Symbolic expressions for Dalitz-Plot Decomposition" dynamic = ["version"] @@ -331,6 +332,7 @@ ignore-names = [ ] "setup.py" = ["D100"] "src/ampform_dpd/io/__init__.py" = ["S403"] +"src/ampform_dpd/io/serialization/format.py" = ["E741"] "tests/*" = [ "D", "INP001", diff --git a/src/ampform_dpd/io/serialization/__init__.py b/src/ampform_dpd/io/serialization/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/ampform_dpd/io/serialization/decay.py b/src/ampform_dpd/io/serialization/decay.py new file mode 100644 index 00000000..20373c5b --- /dev/null +++ b/src/ampform_dpd/io/serialization/decay.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING, Callable, Iterable +from warnings import warn + +import sympy as sp + +from ampform_dpd.decay import ( + FinalStateID, + IsobarNode, + Particle, + State, + StateID, + ThreeBodyDecay, + ThreeBodyDecayChain, +) + +if TYPE_CHECKING: + from ampform_dpd.io.serialization.format import ( + DecayChain, + Distribution, + ModelDefinition, + StateDefinition, + Vertex, + ) + + +def to_decay( + model: ModelDefinition, to_latex: Callable[[str], str] | None = None +) -> ThreeBodyDecay: + initial_state = get_initial_state(model, to_latex) + final_state = get_final_state(model, to_latex) + return ThreeBodyDecay( + states=_get_states(model, to_latex), + chains=sorted({ + to_decay_chain(chain, initial_state, final_state, to_latex) + for chain in _get_decay_chains(model) + }), + ) + + +def to_decay_chain( + chain_definition: DecayChain, + initial_state: State, + final_state: dict[FinalStateID, State], + to_latex: Callable[[str], str] | None = None, +) -> ThreeBodyDecayChain: + vertices = chain_definition["vertices"] + if to_latex is None: + to_latex = lambda x: x # noqa:E731 + resonance = Particle( + name=chain_definition["name"], + latex=to_latex(chain_definition["name"]), + spin=sp.Rational(chain_definition["propagators"][0]["spin"]), + mass=0, + width=0, + parity=None, + ) + child1_id, child2_id = __find_decay_product_ids(vertices) + spectator_id = __find_spectator_id(vertices) + return ThreeBodyDecayChain( + decay=IsobarNode( + parent=initial_state, + child1=IsobarNode( + parent=resonance, + child1=final_state[child1_id], + child2=final_state[child2_id], + ), + child2=final_state[spectator_id], + ) + ) + + +def __find_spectator_id(vertices: Iterable[Vertex]) -> FinalStateID: + for vertex in vertices: + node = vertex["node"] + if any(not isinstance(i, int) for i in node): + for state_id in node: + if isinstance(state_id, int): + return state_id + msg = "Could not find a production node with a spectator" + raise ValueError(msg) + + +def __find_decay_product_ids( + vertices: Iterable[Vertex], +) -> tuple[FinalStateID, FinalStateID]: + for vertex in vertices: + node = vertex["node"] + if all(isinstance(i, int) for i in node) and len(node) == 2: # noqa: PLR2004 + return tuple(node) # type:ignore[return-value] + msg = "Could not find a node that has two final state items (decay node)" + raise ValueError(msg) + + +def _get_decay_chains(model: ModelDefinition) -> list[DecayChain]: + distribution_def = _get_distribution_def(model) + return distribution_def["decay_description"]["chains"] + + +def _get_distribution_def(model: ModelDefinition) -> Distribution: + distribution_defs = model["distributions"] + n_distributions = len(distribution_defs) + if n_distributions == 0: + msg = "The serialized model does not have any distributions" + raise ValueError(msg) + if n_distributions > 1: + msg = f"There are {n_distributions} distributions, but expecting one only" + warn(msg, category=UserWarning) + return distribution_defs[0] + + +def _get_states( + model: ModelDefinition, to_latex: Callable[[str], str] | None = None +) -> dict[StateID, State]: + initial_state = get_initial_state(model, to_latex) + final_state = get_final_state(model, to_latex) + return {initial_state.index: initial_state, **final_state} # type:ignore[dict-item] + + +def get_initial_state( + model: ModelDefinition, to_latex: Callable[[str], str] | None = None +) -> State: + distribution_def = _get_distribution_def(model) + decay_description = distribution_def["decay_description"] + kinematics = decay_description["kinematics"] + return _to_particle(kinematics["initial_state"], to_latex) + + +def get_final_state( + model: ModelDefinition, to_latex: Callable[[str], str] | None = None +) -> dict[FinalStateID, State]: + distribution_def = _get_distribution_def(model) + decay_description = distribution_def["decay_description"] + kinematics = decay_description["kinematics"] + final_state_def = kinematics["final_state"] + final_state = [_to_particle(p, to_latex) for p in final_state_def] + return {p.index: p for p in final_state} + + +def _to_particle( + definition: StateDefinition, to_latex: Callable[[str], str] | None = None +) -> State: + if to_latex is None: + to_latex = lambda x: x # noqa:E731 + return State( + name=definition["name"], + latex=to_latex(definition["name"]), + mass=definition["mass"], + width=0, + spin=sp.Rational(definition["spin"]), + parity=None, + index=definition["index"], + ) diff --git a/src/ampform_dpd/io/serialization/format.py b/src/ampform_dpd/io/serialization/format.py new file mode 100644 index 00000000..d0987910 --- /dev/null +++ b/src/ampform_dpd/io/serialization/format.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import difflib +import sys +from typing import TYPE_CHECKING, Literal, TypedDict, Union + +from ampform_dpd.decay import FinalStateID + +if TYPE_CHECKING: + from ampform_dpd.decay import StateID + +if sys.version_info >= (3, 11): + from typing import Required +else: + from typing_extensions import Required + + +class ModelDefinition(TypedDict): + distributions: Required[list] + functions: Required[list[FunctionDefinition]] + domains: Required[list] + misc: dict + parameter_points: list + + +class Distribution(TypedDict): + name: Required[str] + type: Required[str] + decay_description: Required[DecayDescription] + variables: Required[list] + parameters: list + + +class DecayDescription(TypedDict): + kinematics: Required[KinematicsDefinition] + reference_topology: Required[Topology] + chains: Required[list[DecayChain]] + appendix: dict + + +Topology = list[Union[FinalStateID, "Topology"]] +"""Topology definition as a list of final state IDs.""" + + +class DecayChain(TypedDict): + name: Required[str] + propagators: Required[list[Propagator]] + vertices: Required[list[Vertex]] + topology: Required[Topology] + kinematics: Required[KinematicsDefinition] + weight: Required[str] + + +class Propagator(TypedDict): + spin: str + node: Node + parametrization: str + + +class Vertex(TypedDict): + type: Required[Literal["helicity", "ls", "parity"]] + node: Required[Node] + formfactor: str + + +class HelicityVertex(Vertex): + helicities: Required[tuple[str, str]] + + +class ParityVertex(HelicityVertex): + parity_factor: ParityFactor + + +class LSVertex(Vertex): + l: str + s: str + + +Node = tuple[Union[FinalStateID, "Node"], Union[FinalStateID, "Node"]] +"""Node definition within a `.Topology`.""" +ParityFactor = Literal["+", "-", ""] + + +class KinematicsDefinition(TypedDict): + initial_state: Required[StateDefinition] + final_state: Required[list[StateDefinition]] + + +class StateDefinition(TypedDict): + index: StateID + name: str + spin: str + mass: float + + +class FunctionDefinition(TypedDict): + name: Required[str] + type: Required[str] + + +class BlattWeisskopfDefinition(FunctionDefinition): + l: int + radius: float + + +class BreitWignerDefinition(FunctionDefinition): + mass: float + width: float + ma: float + mb: float + l: float + d: float + + +class MultichannelBreitWignerDefinition(FunctionDefinition): + mass: float + channels: list[ChannelParameters] + + +class ChannelParameters(TypedDict): + gsq: float + ma: float + mb: float + l: float + d: float + + +def get_function_definition( + function_name: str, model: ModelDefinition +) -> FunctionDefinition: + function_definitions = model["functions"] + for function_def in function_definitions: + if function_def["name"] == function_name: + return function_def + existing_names = {f["name"] for f in function_definitions} + msg = f"Could not find function with name {function_name!r}." + candidates = difflib.get_close_matches(function_name, existing_names) + if candidates: + msg += f" Did you mean any of these? {', '.join(sorted(candidates))}" + raise KeyError(msg) diff --git a/tests/io_serialization/__init__.py b/tests/io_serialization/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/io_serialization/conftest.py b/tests/io_serialization/conftest.py new file mode 100644 index 00000000..7f4c08c9 --- /dev/null +++ b/tests/io_serialization/conftest.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +import json +from pathlib import Path +from typing import TYPE_CHECKING + +import pytest + +if TYPE_CHECKING: + from ampform_dpd.io.serialization.format import ModelDefinition + + +@pytest.fixture(scope="session") +def model_definition() -> ModelDefinition: + this_directory = Path(__file__).parent + doc_directory = this_directory.parent.parent / "docs" + with open(doc_directory / "Lc2ppiK.json") as stream: + return json.load(stream) diff --git a/tests/io_serialization/test_decay.py b/tests/io_serialization/test_decay.py new file mode 100644 index 00000000..2974f955 --- /dev/null +++ b/tests/io_serialization/test_decay.py @@ -0,0 +1,45 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import sympy as sp + +from ampform_dpd.io.serialization.decay import ( + _get_decay_chains, # pyright:ignore[reportPrivateUsage] + get_final_state, + get_initial_state, + to_decay_chain, +) + +if TYPE_CHECKING: + from ampform_dpd.io.serialization.format import ModelDefinition + + +def test_get_final_state(model_definition: ModelDefinition): + final_state = get_final_state(model_definition) + assert len(final_state) == 3 + assert set(final_state) == {1, 2, 3} + assert {p.name for p in final_state.values()} == {"p", "pi", "K"} + + +def test_get_initial_state(model_definition: ModelDefinition): + initial_state = get_initial_state(model_definition) + assert initial_state.index == 0 + assert initial_state.name == "Lc" + assert initial_state.latex == "Lc" + assert initial_state.spin is sp.Rational(1 / 2) + assert initial_state.mass == 2.28646 + assert initial_state.parity is None + + +def test_to_decay_chain(model_definition: ModelDefinition): + chain_definitions = _get_decay_chains(model_definition) + chain = to_decay_chain( + chain_definitions[0], + initial_state=get_initial_state(model_definition), + final_state=get_final_state(model_definition), + ) + assert chain.resonance.name == "L1405" + assert chain.resonance.spin is sp.Rational(1 / 2) + assert tuple(p.name for p in chain.decay_products) == ("K", "p") + assert chain.spectator.name == "pi" From 497193f5591bdb21f2933a7e47a7449f809c6b4a Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:44 +0200 Subject: [PATCH 16/42] FEAT: implement standard Breit-Wigner expressions --- src/ampform_dpd/dynamics/__init__.py | 93 +++++++++++++++++++++++++++- 1 file changed, 92 insertions(+), 1 deletion(-) diff --git a/src/ampform_dpd/dynamics/__init__.py b/src/ampform_dpd/dynamics/__init__.py index e4415f73..b622d315 100644 --- a/src/ampform_dpd/dynamics/__init__.py +++ b/src/ampform_dpd/dynamics/__init__.py @@ -1,5 +1,6 @@ """Functions for dynamics lineshapes and kinematics.""" +# pyright:reportPrivateUsage=false from __future__ import annotations from typing import TYPE_CHECKING, Any @@ -7,7 +8,8 @@ import sympy as sp from ampform.dynamics import formulate_form_factor from ampform.kinematics.phasespace import Kallen -from ampform.sympy import unevaluated +from ampform.sympy import argument, unevaluated +from attrs import asdict, frozen if TYPE_CHECKING: from sympy.printing.latex import LatexPrinter @@ -213,3 +215,92 @@ def evaluate(self): angular_momentum=angular_momentum, meson_radius=meson_radius, ) + + +@unevaluated +class MultichannelBreitWigner(sp.Expr): + s: Any + mass: Any + channels: list[ChannelArguments] = argument(sympify=False) + + def evaluate(self): + s = self.s + m0 = self.mass + width = sum(channel.formulate_width(s, m0) for channel in self.channels) + return BreitWigner(s, m0, width) + + def _latex_repr_(self, printer: LatexPrinter, *args) -> str: + latex = R"\mathcal{R}^\mathrm{BW}_\mathrm{multi}\left(" + latex += printer._print(self.s) + "; " + latex += ", ".join(printer._print(channel.width) for channel in self.channels) + latex += R"\right)" + return latex + + +@frozen +class ChannelArguments: + width: Any + m1: Any = 0 + m2: Any = 0 + angular_momentum: Any = 0 + meson_radius: Any = 1 + + def __attrs_post_init__(self) -> None: + for name, value in asdict(self).items(): + object.__setattr__(self, name, sp.sympify(value)) + + def formulate_width(self, s: Any, m0: Any) -> sp.Expr: + Γ0 = self.width + m1 = self.m1 + m2 = self.m2 + L = self.angular_momentum + R = self.meson_radius + ff = FormFactor(s, m1, m2, L, R) ** 2 + return Γ0 * m0 / sp.sqrt(s) * ff + + +@unevaluated +class BreitWigner(sp.Expr): + s: Any + mass: Any + width: Any + m1: Any = 0 + m2: Any = 0 + angular_momentum: Any = 0 + meson_radius: Any = 1 + + def evaluate(self): + width = self.energy_dependent_width() + expr = SimpleBreitWigner(self.s, self.mass, width) + if self.angular_momentum == 0 and self.m1 == 0 and self.m2 == 0: + return expr.evaluate() + return expr + + def energy_dependent_width(self) -> sp.Expr: + s, m0, Γ0, m1, m2, L, d = self.args + if L == 0 and m1 == 0 and m2 == 0: + return Γ0 # type:ignore[return-value] + return EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d) + + def _latex_repr_(self, printer: LatexPrinter, *args) -> str: + s = printer._print(self.s) + function_symbol = R"\mathcal{R}^\mathrm{BW}" + mass = printer._print(self.mass) + width = printer._print(self.width) + arg = Rf"\left({s}; {mass}, {width}\right)" + L = printer._print(self.angular_momentum) + if isinstance(self.angular_momentum, sp.Integer): + return Rf"{function_symbol}_{{L={L}}}{arg}" + return Rf"{function_symbol}_{{{L}}}{arg}" + + +@unevaluated +class SimpleBreitWigner(sp.Expr): + s: Any + mass: Any + width: Any + _latex_repr_ = R"\mathcal{{R}}^\mathrm{{BW}}\left({s}; {mass}, {width}\right)" + + def evaluate(self): + s, m0, Γ0 = self.args + return 1 / (m0**2 - s - m0 * Γ0 * 1j) From d4449e0a1d23e58a7016875332fa7670ca0fde8b Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:46 +0200 Subject: [PATCH 17/42] DX: add tests for `get_function_definition` --- tests/io_serialization/test_format.py | 34 +++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) create mode 100644 tests/io_serialization/test_format.py diff --git a/tests/io_serialization/test_format.py b/tests/io_serialization/test_format.py new file mode 100644 index 00000000..f51fceac --- /dev/null +++ b/tests/io_serialization/test_format.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import pytest + +from ampform_dpd.io.serialization.format import ModelDefinition, get_function_definition + + +def test_get_function_definition_blatt_weisskopf(model_definition: ModelDefinition): + func_def = get_function_definition("BlattWeisskopf_resonance_l1", model_definition) + assert func_def == { + "name": "BlattWeisskopf_resonance_l1", + "type": "BlattWeisskopf", + "radius": 1.5, + "l": 1, + } + func_def = get_function_definition("BlattWeisskopf_resonance_l2", model_definition) + assert func_def == { + "name": "BlattWeisskopf_resonance_l2", + "type": "BlattWeisskopf", + "radius": 1.5, + "l": 2, + } + + +def test_get_function_definition_missing(model_definition: ModelDefinition): + # cspell:ignore Weiskopf + with pytest.raises( + KeyError, + match=( + r" Did you mean any of these\?" + " BlattWeisskopf_b_decay_l1, BlattWeisskopf_resonance_l1, BlattWeisskopf_resonance_l2" + ), + ): + get_function_definition("BlattWeiskopf", model_definition) From 46803a8307331eb2d5ea02a345891d17b0245058 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:47 +0200 Subject: [PATCH 18/42] FEAT: implement `dynamics` module --- docs/conf.py | 1 + docs/serialization.ipynb | 398 ++++++++++++++++++- src/ampform_dpd/io/serialization/dynamics.py | 203 ++++++++++ tests/io_serialization/test_dynamics.py | 0 4 files changed, 598 insertions(+), 4 deletions(-) create mode 100644 src/ampform_dpd/io/serialization/dynamics.py create mode 100644 tests/io_serialization/test_dynamics.py diff --git a/docs/conf.py b/docs/conf.py index 41652019..6cefc719 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -37,6 +37,7 @@ api_github_repo = f"{ORGANIZATION}/{REPO_NAME}" api_target_substitutions: dict[str, str | tuple[str, str]] = { "ampform_dpd.decay.StateIDTemplate": ("obj", "ampform_dpd.decay.StateID"), + "ampform_dpd.io.serialization.dynamics.T": "typing.TypeVar", "DecayNode": ("obj", "ampform_dpd.decay.DecayNode"), "FinalState": ("obj", "ampform_dpd.decay.FinalState"), "FinalStateID": ("obj", "ampform_dpd.decay.FinalStateID"), diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 415a9090..0d6283bc 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -22,6 +22,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "tags": [ "hide-input", "scroll-input" @@ -33,10 +36,42 @@ "\n", "import json\n", "\n", - "from IPython.display import JSON, Markdown, Math\n", + "import sympy as sp\n", + "from ampform.dynamics import BlattWeisskopfSquared\n", + "from ampform.dynamics.phasespace import BreakupMomentumSquared\n", + "from ampform.kinematics.phasespace import Kallen\n", + "from IPython.display import JSON, Markdown, Math, display\n", "\n", + "from ampform_dpd import DefinedExpression\n", + "from ampform_dpd.dynamics import (\n", + " BreitWigner,\n", + " BuggBreitWigner,\n", + " ChannelArguments,\n", + " EnergyDependentWidth,\n", + " FormFactor,\n", + " MultichannelBreitWigner,\n", + " P,\n", + " SimpleBreitWigner,\n", + ")\n", "from ampform_dpd.io import as_markdown_table, aslatex\n", - "from ampform_dpd.io.serialization.decay import to_decay" + "from ampform_dpd.io.serialization.decay import (\n", + " _get_decay_chains,\n", + " get_final_state,\n", + " to_decay,\n", + ")\n", + "from ampform_dpd.io.serialization.dynamics import (\n", + " formulate_breit_wigner,\n", + " formulate_dynamics,\n", + " formulate_form_factor,\n", + " formulate_multichannel_breit_wigner,\n", + " to_mandelstam_symbol,\n", + " to_mass_symbol,\n", + ")\n", + "from ampform_dpd.io.serialization.format import (\n", + " ModelDefinition,\n", + " Propagator,\n", + " get_function_definition,\n", + ")" ] }, { @@ -67,6 +102,7 @@ { "cell_type": "markdown", "metadata": { + "jp-MarkdownHeadingCollapsed": true, "tags": [] }, "source": [ @@ -89,7 +125,7 @@ }, "outputs": [], "source": [ - "def name_to_latex(name: str) -> str:\n", + "def to_latex(name: str) -> str:\n", " latex = {\n", " \"Lc\": R\"\\Lambda_c^+\",\n", " \"pi\": R\"\\pi^+\",\n", @@ -114,7 +150,7 @@ }, "outputs": [], "source": [ - "DECAY = to_decay(MODEL_DEFINITION, to_latex=name_to_latex)\n", + "DECAY = to_decay(MODEL_DEFINITION, to_latex=to_latex)\n", "Math(aslatex(DECAY))" ] }, @@ -130,6 +166,360 @@ "source": [ "Markdown(as_markdown_table(DECAY))" ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Dynamics" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{seealso} [RUB-EP1/amplitude-serialization#22](https://github.com/RUB-EP1/amplitude-serialization/issues/22)\n", + ":::" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "CHAIN_DEFS = _get_decay_chains(MODEL_DEFINITION)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Vertices" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Blatt-Weisskopf form factor" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "z = sp.Symbol(\"z\", nonnegative=True)\n", + "s, m1, m2, L, d = sp.symbols(\"s m1 m2 L R\")\n", + "exprs = [\n", + " FormFactor(s, m1, m2, L, d),\n", + " BlattWeisskopfSquared(z, L),\n", + " BreakupMomentumSquared(s, m1, m2),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "ff_L1520 = formulate_form_factor(\n", + " vertex=CHAIN_DEFS[2][\"vertices\"][0],\n", + " model=MODEL_DEFINITION,\n", + ")\n", + "Math(aslatex(ff_L1520))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Propagators" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Breit-Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "x, y, z = sp.symbols(\"x:z\")\n", + "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", + "exprs = [\n", + " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", + " SimpleBreitWigner(s, m0, Γ0),\n", + " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", + " FormFactor(s, m1, m2, L, d),\n", + " P(s, m1, m2),\n", + " Kallen(x, y, z),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "K892_BW = formulate_breit_wigner(\n", + " propagator=CHAIN_DEFS[20][\"propagators\"][0],\n", + " resonance=to_latex(CHAIN_DEFS[20][\"name\"]),\n", + " model=MODEL_DEFINITION,\n", + ")\n", + "Math(aslatex(K892_BW))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Multi-channel Breit-Wigner" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "x, y, z = sp.symbols(\"x:z\")\n", + "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", + "channels = [\n", + " ChannelArguments(\n", + " sp.Symbol(f\"Gamma{i}\"),\n", + " sp.Symbol(f\"m_{{a,{i}}}\"),\n", + " sp.Symbol(f\"m_{{b,{i}}}\"),\n", + " sp.Symbol(f\"L{i}\"),\n", + " d,\n", + " )\n", + " for i in [1, 2]\n", + "]\n", + "exprs = [\n", + " MultichannelBreitWigner(s, m0, channels),\n", + " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", + " BreitWigner(s, m0, Γ0),\n", + " EnergyDependentWidth(s, m0, Γ0, m1, m2, L, d),\n", + " FormFactor(s, m1, m2, L, d),\n", + " P(s, m1, m2),\n", + " Kallen(x, y, z),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "full-width" + ] + }, + "outputs": [], + "source": [ + "L1405_Flatte = formulate_multichannel_breit_wigner(\n", + " propagator=CHAIN_DEFS[0][\"propagators\"][0],\n", + " resonance=to_latex(CHAIN_DEFS[0][\"name\"]),\n", + " model=MODEL_DEFINITION,\n", + ")\n", + "Math(aslatex(L1405_Flatte))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "#### Breit-Wigner with exponential" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The model contains one lineshape function that is not standard, so we have to implement a custom propagator dynamics builder for this." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "s, m0, Γ0, m1, m2, γ = sp.symbols(\"s m0 Gamma0 m1 m2 gamma\")\n", + "expr = BuggBreitWigner(s, m0, Γ0, m1, m2, γ)\n", + "Math(aslatex({expr: expr.doit(deep=False)}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "CHAIN_DEFS[18]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "get_function_definition(\"K700_BuggBW\", MODEL_DEFINITION)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def formulate_bugg_breit_wigner(\n", + " propagator: Propagator, resonance: str, model: ModelDefinition\n", + ") -> DefinedExpression:\n", + " function_definition = get_function_definition(propagator[\"parametrization\"], model)\n", + " node = propagator[\"node\"]\n", + " i, j = node\n", + " s = to_mandelstam_symbol(node)\n", + " mass = sp.Symbol(f\"m_{{{resonance}}}\", nonnegative=True)\n", + " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\", nonnegative=True)\n", + " γ = sp.Symbol(Rf\"\\gamma_{{{resonance}}}\")\n", + " m1 = to_mass_symbol(i)\n", + " m2 = to_mass_symbol(j)\n", + " final_state = get_final_state(model)\n", + " return DefinedExpression(\n", + " expression=BuggBreitWigner(s, mass, width, m1, m2, γ),\n", + " definitions={\n", + " mass: function_definition[\"mass\"],\n", + " width: function_definition[\"width\"],\n", + " m1: final_state[i].mass,\n", + " m2: final_state[j].mass,\n", + " γ: function_definition[\"slope\"],\n", + " },\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "CHAIN_18 = CHAIN_DEFS[18]\n", + "K700_BuggBW = formulate_bugg_breit_wigner(\n", + " propagator=CHAIN_18[\"propagators\"][0],\n", + " resonance=to_latex(CHAIN_18[\"name\"]),\n", + " model=MODEL_DEFINITION,\n", + ")\n", + "Math(aslatex(K700_BuggBW))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### General propagator dynamics builder" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "DYNAMICS_BUILDERS = {\n", + " \"BreitWignerWidthExpLikeBugg\": formulate_bugg_breit_wigner,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "exprs = [\n", + " formulate_dynamics(CHAIN_DEFS[0], MODEL_DEFINITION, to_latex, DYNAMICS_BUILDERS),\n", + " formulate_dynamics(CHAIN_DEFS[18], MODEL_DEFINITION, to_latex, DYNAMICS_BUILDERS),\n", + " formulate_dynamics(CHAIN_DEFS[20], MODEL_DEFINITION, to_latex, DYNAMICS_BUILDERS),\n", + "]\n", + "for expr in exprs:\n", + " display(Math(aslatex(expr)))" + ] } ], "metadata": { diff --git a/src/ampform_dpd/io/serialization/dynamics.py b/src/ampform_dpd/io/serialization/dynamics.py new file mode 100644 index 00000000..31b49a79 --- /dev/null +++ b/src/ampform_dpd/io/serialization/dynamics.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +from collections import abc +from typing import TYPE_CHECKING, Callable, Protocol, TypeVar, cast + +import sympy as sp + +from ampform_dpd import DefinedExpression +from ampform_dpd.dynamics import ( + BreitWigner, + ChannelArguments, + FormFactor, + MultichannelBreitWigner, +) +from ampform_dpd.io.serialization.decay import get_initial_state +from ampform_dpd.io.serialization.format import ( + BlattWeisskopfDefinition, + BreitWignerDefinition, + DecayChain, + ModelDefinition, + MultichannelBreitWignerDefinition, + Propagator, + Vertex, + get_function_definition, +) + +if TYPE_CHECKING: + from ampform_dpd.io.serialization.format import Node + +T = TypeVar("T") + + +def identity_function(x: T) -> T: + return x + + +class PropagatorDynamicsBuilder(Protocol): + def __call__( + self, + propagator: Propagator, + resonance: str, + model: ModelDefinition, + ) -> DefinedExpression: ... + + +def formulate_dynamics( + chain_definition: DecayChain, + model: ModelDefinition, + to_latex: Callable[[str], str] = identity_function, + additional_definitions: dict[str, PropagatorDynamicsBuilder] | None = None, +) -> DefinedExpression: + definitions: dict[str, PropagatorDynamicsBuilder] = { + "BreitWigner": formulate_breit_wigner, + "MultichannelBreitWigner": formulate_multichannel_breit_wigner, + } + if additional_definitions is not None: + definitions.update(additional_definitions) + expr = DefinedExpression() + for propagator in chain_definition["propagators"]: + parametrization = propagator["parametrization"] + function_definition = get_function_definition(parametrization, model) + function_type = function_definition["type"] + dynamics_builder = definitions.get(function_type) + if dynamics_builder is None: + msg = f"No dynamics implementation for function type {function_type!r}" + raise NotImplementedError(msg) + expr *= dynamics_builder( + propagator, + resonance=to_latex(chain_definition["name"]), + model=model, + ) + return expr + + +def formulate_form_factor(vertex: Vertex, model: ModelDefinition) -> DefinedExpression: + function_name = vertex.get("formfactor") + if not function_name: + return DefinedExpression() + function_definition = get_function_definition(function_name, model) + function_definition = cast(BlattWeisskopfDefinition, function_definition) + function_type = function_definition["type"] + if function_type == "BlattWeisskopf": + node = vertex["node"] + s = to_mandelstam_symbol(node) + m1, m2 = (to_mass_symbol(i) for i in node) + if all(isinstance(i, int) for i in node): + meson_radius = sp.Symbol(R"R_\mathrm{res}") + else: + initial_state = get_initial_state(model) + meson_radius = sp.Symbol(f"R_{{{initial_state.latex}}}") + angular_momentum = int(function_definition["l"]) + return DefinedExpression( + expression=FormFactor(s, m1, m2, angular_momentum, meson_radius), + definitions={ + meson_radius: function_definition["radius"], + }, + ) + msg = f"No form factor implementation for {function_name!r}" + raise NotImplementedError(msg) + + +def formulate_breit_wigner( + propagator: Propagator, resonance: str, model: ModelDefinition +) -> DefinedExpression: + function_definition = get_function_definition(propagator["parametrization"], model) + function_definition = cast(BreitWignerDefinition, function_definition) + node = propagator["node"] + i, j = node + s = to_mandelstam_symbol(node) + mass = sp.Symbol(f"m_{{{resonance}}}") + width = sp.Symbol(Rf"\Gamma_{{{resonance}}}") + m1 = to_mass_symbol(i) + m2 = to_mass_symbol(j) + angular_momentum = int(function_definition["l"]) + d = sp.Symbol(R"R_\mathrm{res}") + return DefinedExpression( + expression=BreitWigner(s, mass, width, m1, m2, angular_momentum, d), + definitions={ + mass: function_definition["mass"], + width: function_definition["width"], + m1: function_definition["ma"], + m2: function_definition["mb"], + d: function_definition["d"], + }, + ) + + +def formulate_multichannel_breit_wigner( # noqa: PLR0914 + propagator: Propagator, resonance: str, model: ModelDefinition +) -> DefinedExpression: + function_definition = get_function_definition(propagator["parametrization"], model) + function_definition = cast(MultichannelBreitWignerDefinition, function_definition) + channel_definitions = function_definition["channels"] + if len(channel_definitions) < 2: # noqa: PLR2004 + msg = "Need at least two channels for a multi-channel Breit-Wigner" + raise NotImplementedError(msg) + node = propagator["node"] + i, j = node + s = to_mandelstam_symbol(node) + mass = sp.Symbol(f"m_{{{resonance}}}") + width = sp.Symbol(Rf"\Gamma_{{{resonance}}}") + m1 = to_mass_symbol(i) + m2 = to_mass_symbol(j) + angular_momentum = int(channel_definitions[0]["l"]) + d = sp.Symbol(f"R_{{{resonance}}}") + channels = [ChannelArguments(width, m1, m2, angular_momentum, d)] + parameter_defaults: dict[sp.Symbol, complex | float] = { + mass: function_definition["mass"], + width: channel_definitions[0]["gsq"], + m1: channel_definitions[0]["ma"], + m2: channel_definitions[0]["mb"], + d: channel_definitions[0]["d"], + } + for channel_idx, channel_definition in enumerate(channel_definitions[1:], 2): + Γi = sp.Symbol( + Rf"\Gamma_{{{resonance}}}^\text{{ch. {channel_idx}}}", nonnegative=True + ) + mi1 = sp.Symbol(f"m_{{a,{channel_idx}}}", nonnegative=True) + mi2 = sp.Symbol(f"m_{{b,{channel_idx}}}", nonnegative=True) + angular_momentum = int(channel_definition["l"]) + channels.append(ChannelArguments(Γi, mi1, mi2, angular_momentum, d)) + parameter_defaults.update({ + mi1: channel_definition["ma"], + mi2: channel_definition["mb"], + }) + return DefinedExpression( + expression=MultichannelBreitWigner(s, mass, channels), + definitions=parameter_defaults, + ) + + +def to_mandelstam_symbol(node: Node) -> sp.Symbol: + """Create a Mandelstam symbol for a node. + + >>> to_mandelstam_symbol([3, 2]) + sigma1 + >>> to_mandelstam_symbol([1, [2, 3]]) + m0 + """ + if all(isinstance(i, int) for i in node): + return to_mass_symbol(node) + return to_mass_symbol(0) + + +def to_mass_symbol(node_item: int | Node) -> sp.Symbol: + """Create a mass symbol for a node. + + >>> to_mass_symbol(1) + m1 + >>> to_mass_symbol((1, 2)) + sigma3 + """ + if isinstance(node_item, int): + return sp.Symbol(f"m{node_item}", nonnegative=True) + if ( + isinstance(node_item, abc.Sequence) + and all(isinstance(i, int) for i in node_item) + and len(node_item) == 2 # noqa: PLR2004 + ): + k, *_ = {1, 2, 3} - set(node_item) # type:ignore[arg-type] + return sp.Symbol(f"sigma{k}", nonnegative=True) + msg = f"Cannot create mass symbol for node {node_item}" + raise NotImplementedError(msg) diff --git a/tests/io_serialization/test_dynamics.py b/tests/io_serialization/test_dynamics.py new file mode 100644 index 00000000..e69de29b From 55d54abc7fa4a2d8b2abc7120dd728d798d0322b Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:48 +0200 Subject: [PATCH 19/42] BREAK: improve organization `format` and `decay` module --- docs/serialization.ipynb | 9 ++---- src/ampform_dpd/io/serialization/decay.py | 30 ++++--------------- src/ampform_dpd/io/serialization/format.py | 18 ++++++++++++ tests/io_serialization/test_decay.py | 34 ++++++++++++++++++++-- 4 files changed, 59 insertions(+), 32 deletions(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 0d6283bc..e30a1ccf 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -54,11 +54,7 @@ " SimpleBreitWigner,\n", ")\n", "from ampform_dpd.io import as_markdown_table, aslatex\n", - "from ampform_dpd.io.serialization.decay import (\n", - " _get_decay_chains,\n", - " get_final_state,\n", - " to_decay,\n", - ")\n", + "from ampform_dpd.io.serialization.decay import get_final_state, to_decay\n", "from ampform_dpd.io.serialization.dynamics import (\n", " formulate_breit_wigner,\n", " formulate_dynamics,\n", @@ -70,6 +66,7 @@ "from ampform_dpd.io.serialization.format import (\n", " ModelDefinition,\n", " Propagator,\n", + " get_decay_chains,\n", " get_function_definition,\n", ")" ] @@ -194,7 +191,7 @@ }, "outputs": [], "source": [ - "CHAIN_DEFS = _get_decay_chains(MODEL_DEFINITION)" + "CHAIN_DEFS = get_decay_chains(MODEL_DEFINITION)" ] }, { diff --git a/src/ampform_dpd/io/serialization/decay.py b/src/ampform_dpd/io/serialization/decay.py index 20373c5b..97102b0c 100644 --- a/src/ampform_dpd/io/serialization/decay.py +++ b/src/ampform_dpd/io/serialization/decay.py @@ -1,7 +1,6 @@ from __future__ import annotations from typing import TYPE_CHECKING, Callable, Iterable -from warnings import warn import sympy as sp @@ -14,11 +13,11 @@ ThreeBodyDecay, ThreeBodyDecayChain, ) +from ampform_dpd.io.serialization.format import get_decay_chains, get_distribution_def if TYPE_CHECKING: from ampform_dpd.io.serialization.format import ( DecayChain, - Distribution, ModelDefinition, StateDefinition, Vertex, @@ -31,10 +30,10 @@ def to_decay( initial_state = get_initial_state(model, to_latex) final_state = get_final_state(model, to_latex) return ThreeBodyDecay( - states=_get_states(model, to_latex), + states=get_states(model, to_latex), chains=sorted({ to_decay_chain(chain, initial_state, final_state, to_latex) - for chain in _get_decay_chains(model) + for chain in get_decay_chains(model) }), ) @@ -93,24 +92,7 @@ def __find_decay_product_ids( raise ValueError(msg) -def _get_decay_chains(model: ModelDefinition) -> list[DecayChain]: - distribution_def = _get_distribution_def(model) - return distribution_def["decay_description"]["chains"] - - -def _get_distribution_def(model: ModelDefinition) -> Distribution: - distribution_defs = model["distributions"] - n_distributions = len(distribution_defs) - if n_distributions == 0: - msg = "The serialized model does not have any distributions" - raise ValueError(msg) - if n_distributions > 1: - msg = f"There are {n_distributions} distributions, but expecting one only" - warn(msg, category=UserWarning) - return distribution_defs[0] - - -def _get_states( +def get_states( model: ModelDefinition, to_latex: Callable[[str], str] | None = None ) -> dict[StateID, State]: initial_state = get_initial_state(model, to_latex) @@ -121,7 +103,7 @@ def _get_states( def get_initial_state( model: ModelDefinition, to_latex: Callable[[str], str] | None = None ) -> State: - distribution_def = _get_distribution_def(model) + distribution_def = get_distribution_def(model) decay_description = distribution_def["decay_description"] kinematics = decay_description["kinematics"] return _to_particle(kinematics["initial_state"], to_latex) @@ -130,7 +112,7 @@ def get_initial_state( def get_final_state( model: ModelDefinition, to_latex: Callable[[str], str] | None = None ) -> dict[FinalStateID, State]: - distribution_def = _get_distribution_def(model) + distribution_def = get_distribution_def(model) decay_description = distribution_def["decay_description"] kinematics = decay_description["kinematics"] final_state_def = kinematics["final_state"] diff --git a/src/ampform_dpd/io/serialization/format.py b/src/ampform_dpd/io/serialization/format.py index d0987910..d3fdc571 100644 --- a/src/ampform_dpd/io/serialization/format.py +++ b/src/ampform_dpd/io/serialization/format.py @@ -3,6 +3,7 @@ import difflib import sys from typing import TYPE_CHECKING, Literal, TypedDict, Union +from warnings import warn from ampform_dpd.decay import FinalStateID @@ -125,6 +126,23 @@ class ChannelParameters(TypedDict): d: float +def get_decay_chains(model: ModelDefinition) -> list[DecayChain]: + distribution_def = get_distribution_def(model) + return distribution_def["decay_description"]["chains"] + + +def get_distribution_def(model: ModelDefinition) -> Distribution: + distribution_defs = model["distributions"] + n_distributions = len(distribution_defs) + if n_distributions == 0: + msg = "The serialized model does not have any distributions" + raise ValueError(msg) + if n_distributions > 1: + msg = f"There are {n_distributions} distributions, but expecting one only" + warn(msg, category=UserWarning) + return distribution_defs[0] + + def get_function_definition( function_name: str, model: ModelDefinition ) -> FunctionDefinition: diff --git a/tests/io_serialization/test_decay.py b/tests/io_serialization/test_decay.py index 2974f955..ceb24832 100644 --- a/tests/io_serialization/test_decay.py +++ b/tests/io_serialization/test_decay.py @@ -5,11 +5,14 @@ import sympy as sp from ampform_dpd.io.serialization.decay import ( - _get_decay_chains, # pyright:ignore[reportPrivateUsage] + get_decay_chains, # pyright:ignore[reportPrivateUsage] get_final_state, get_initial_state, + get_states, + to_decay, to_decay_chain, ) +from ampform_dpd.io.serialization.format import ModelDefinition if TYPE_CHECKING: from ampform_dpd.io.serialization.format import ModelDefinition @@ -32,8 +35,35 @@ def test_get_initial_state(model_definition: ModelDefinition): assert initial_state.parity is None +def test_get_states(model_definition: ModelDefinition): + states = get_states(model_definition) + assert len(states) == 4 + assert {i: s.name for i, s in states.items()} == {0: "Lc", 1: "p", 2: "pi", 3: "K"} + + +def test_to_decay(model_definition: ModelDefinition): + decay = to_decay(model_definition) + assert decay.initial_state.name == "Lc" + assert {p.name for p in decay.final_state.values()} == {"p", "pi", "K"} + assert len(decay.chains) == 12 + assert [c.resonance.name for c in decay.chains] == [ + "D1232", + "D1600", + "D1700", + "K1430", + "K700", + "K892", + "L1405", + "L1520", + "L1600", + "L1670", + "L1690", + "L2000", + ] + + def test_to_decay_chain(model_definition: ModelDefinition): - chain_definitions = _get_decay_chains(model_definition) + chain_definitions = get_decay_chains(model_definition) chain = to_decay_chain( chain_definitions[0], initial_state=get_initial_state(model_definition), From 69270576837d0561e2fdbcd6d27e55044e88ef8a Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:50 +0200 Subject: [PATCH 20/42] FEAT: implement `amplitude` module --- docs/serialization.ipynb | 481 +++++++++++++++++- pyproject.toml | 1 + src/ampform_dpd/io/serialization/amplitude.py | 398 +++++++++++++++ src/ampform_dpd/io/serialization/decay.py | 19 +- src/ampform_dpd/io/serialization/format.py | 5 + tests/io_serialization/test_amplitude.py | 63 +++ 6 files changed, 961 insertions(+), 6 deletions(-) create mode 100644 src/ampform_dpd/io/serialization/amplitude.py create mode 100644 tests/io_serialization/test_amplitude.py diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index e30a1ccf..6450291b 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -36,13 +36,19 @@ "\n", "import json\n", "\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", "import sympy as sp\n", "from ampform.dynamics import BlattWeisskopfSquared\n", "from ampform.dynamics.phasespace import BreakupMomentumSquared\n", "from ampform.kinematics.phasespace import Kallen\n", - "from IPython.display import JSON, Markdown, Math, display\n", + "from ampform.sympy import perform_cached_doit\n", + "from IPython.display import JSON, Markdown, Math\n", + "from tqdm.auto import tqdm\n", "\n", "from ampform_dpd import DefinedExpression\n", + "from ampform_dpd.decay import FinalStateID, State, ThreeBodyDecay\n", "from ampform_dpd.dynamics import (\n", " BreitWigner,\n", " BuggBreitWigner,\n", @@ -53,7 +59,21 @@ " P,\n", " SimpleBreitWigner,\n", ")\n", - "from ampform_dpd.io import as_markdown_table, aslatex\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.io.serialization.amplitude import (\n", + " HelicityRecoupling,\n", + " LSRecoupling,\n", + " ParityRecoupling,\n", + " formulate,\n", + " formulate_aligned_amplitude,\n", + " formulate_chain_amplitude,\n", + " formulate_recoupling,\n", + ")\n", "from ampform_dpd.io.serialization.decay import get_final_state, to_decay\n", "from ampform_dpd.io.serialization.dynamics import (\n", " formulate_breit_wigner,\n", @@ -68,7 +88,9 @@ " Propagator,\n", " get_decay_chains,\n", " get_function_definition,\n", - ")" + ")\n", + "\n", + "simplify_latex_rendering()" ] }, { @@ -514,8 +536,457 @@ " formulate_dynamics(CHAIN_DEFS[18], MODEL_DEFINITION, to_latex, DYNAMICS_BUILDERS),\n", " formulate_dynamics(CHAIN_DEFS[20], MODEL_DEFINITION, to_latex, DYNAMICS_BUILDERS),\n", "]\n", - "for expr in exprs:\n", - " display(Math(aslatex(expr)))" + "Math(aslatex(exprs))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Construct `AmplitudeModel`" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Unpolarized intensity" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "full-width" + ] + }, + "outputs": [], + "source": [ + "λ0, λ1, λ2, λ3 = sp.symbols(\"lambda(:4)\", rational=True)\n", + "amplitude_expr, _ = formulate_aligned_amplitude(MODEL_DEFINITION, λ0, λ1, λ2, λ3)\n", + "amplitude_expr.cleanup()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Amplitude for the decay chain" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Helicity recouplings" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "λa = sp.Symbol(R\"\\lambda_a\", rational=True)\n", + "λb = sp.Symbol(R\"\\lambda_b\", rational=True)\n", + "λa0 = sp.Symbol(R\"\\lambda_a^0\", rational=True)\n", + "λb0 = sp.Symbol(R\"\\lambda_b^0\", rational=True)\n", + "f = sp.Symbol(\"f\", integer=True)\n", + "l = sp.Symbol(\"l\", integer=True, nonnegative=True)\n", + "s = sp.Symbol(\"s\", nonnegative=True, rational=True)\n", + "ja = sp.Symbol(\"j_a\", nonnegative=True, rational=True)\n", + "jb = sp.Symbol(\"j_b\", nonnegative=True, rational=True)\n", + "j = sp.Symbol(\"j\", nonnegative=True, rational=True)\n", + "exprs = [\n", + " HelicityRecoupling(λa, λb, λa0, λb0),\n", + " ParityRecoupling(λa, λb, λa0, λb0, f),\n", + " LSRecoupling(λa, λb, l, s, ja, jb, j),\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Recoupling deserialization" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "recouplings = [\n", + " formulate_recoupling(MODEL_DEFINITION, chain_idx=0, vertex_idx=i) for i in range(2)\n", + "]\n", + "Math(aslatex({e: e.doit(deep=False) for e in recouplings}))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Chain amplitudes" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "full-width" + ] + }, + "outputs": [], + "source": [ + "definitions = formulate_chain_amplitude(λ0, λ1, λ2, λ3, MODEL_DEFINITION, chain_idx=0)\n", + "Math(aslatex(definitions))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Full amplitude model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "full-width" + ] + }, + "outputs": [], + "source": [ + "MODEL = formulate(\n", + " MODEL_DEFINITION,\n", + " additional_builders=DYNAMICS_BUILDERS,\n", + " cleanup_summations=True,\n", + " to_latex=to_latex,\n", + ")\n", + "MODEL.intensity" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input", + "full-width" + ] + }, + "outputs": [], + "source": [ + "Math(aslatex(MODEL.amplitudes))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "Math(aslatex(MODEL.variables))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "Math(aslatex({**MODEL.invariants, **MODEL.masses}))" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "## Numeric results" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "intensity_expr = MODEL.full_expression.xreplace(MODEL.variables)\n", + "intensity_expr = intensity_expr.xreplace(MODEL.parameter_defaults)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "free_symbols = intensity_expr.free_symbols\n", + "assert len(free_symbols) == 3\n", + "assert str(sorted(free_symbols, key=str)) == \"[sigma1, sigma2, sigma3]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "mystnb": { + "code_prompt_show": "Lambdify to numeric function" + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "intensity_funcs = {}\n", + "for s, s_expr in tqdm(MODEL.invariants.items()):\n", + " k = int(str(s)[-1])\n", + " s_expr = s_expr.xreplace(MODEL.masses).doit()\n", + " expr = perform_cached_doit(intensity_expr.xreplace({s: s_expr}))\n", + " func = perform_cached_lambdify(expr, backend=\"jax\")\n", + " assert len(func.argument_order) == 2\n", + " intensity_funcs[k] = func" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Validation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "checksums = {\n", + " misc_key: {checksum[\"name\"]: checksum[\"value\"] for checksum in misc_value}\n", + " for misc_key, misc_value in MODEL_DEFINITION[\"misc\"].items()\n", + " if \"checksum\" in misc_key\n", + "}\n", + "checksums" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [] + }, + "outputs": [], + "source": [ + "checksum_points = {\n", + " point[\"name\"]: {par[\"name\"]: par[\"value\"] for par in point[\"parameters\"]}\n", + " for point in MODEL_DEFINITION[\"parameter_points\"]\n", + "}\n", + "checksum_points" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "array = []\n", + "for distribution_name, checksum in checksums.items():\n", + " for point_name, expected in checksum.items():\n", + " parameters = checksum_points[point_name]\n", + " s1 = parameters[\"m_31_2\"] ** 2\n", + " s2 = parameters[\"m_31\"] ** 2\n", + " computed = intensity_funcs[3]({\"sigma1\": s1, \"sigma2\": s2})\n", + " status = \"🟢\" if computed == expected else \"🔴\"\n", + " array.append((distribution_name, point_name, computed, expected, status))\n", + "pd.DataFrame(array, columns=[\"Distribution\", \"Point\", \"Computed\", \"Expected\", \"Status\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + "### Dalitz plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "i, j = (2, 1)\n", + "k, *_ = {1, 2, 3} - {i, j}\n", + "σk, σk_expr = list(MODEL.invariants.items())[k - 1]\n", + "Math(aslatex({σk: σk_expr}))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "mystnb": { + "code_prompt_show": "Define meshgrid for Dalitz plot" + }, + "tags": [ + "hide-input", + "remove-output" + ] + }, + "outputs": [], + "source": [ + "resolution = 1_000\n", + "m = sorted(MODEL.masses, key=str)\n", + "x_min = float(((m[j] + m[k]) ** 2).xreplace(MODEL.masses))\n", + "x_max = float(((m[0] - m[i]) ** 2).xreplace(MODEL.masses))\n", + "y_min = float(((m[i] + m[k]) ** 2).xreplace(MODEL.masses))\n", + "y_max = float(((m[0] - m[j]) ** 2).xreplace(MODEL.masses))\n", + "x_diff = x_max - x_min\n", + "y_diff = y_max - y_min\n", + "x_min -= 0.05 * x_diff\n", + "x_max += 0.05 * x_diff\n", + "y_min -= 0.05 * y_diff\n", + "y_max += 0.05 * y_diff\n", + "X, Y = jnp.meshgrid(\n", + " jnp.linspace(x_min, x_max, num=resolution),\n", + " jnp.linspace(y_min, y_max, num=resolution),\n", + ")\n", + "dalitz_data = {\n", + " f\"sigma{i}\": X,\n", + " f\"sigma{j}\": Y,\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "mystnb": { + "code_prompt_show": "Prepare parametrized numerical function" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "intensities = intensity_funcs[k](dalitz_data)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "remove-input" + ] + }, + "outputs": [], + "source": [ + "assert not jnp.all(jnp.isnan(intensities)), \"All intensities are NaN\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "jupyter": { + "source_hidden": true + }, + "tags": [ + "hide-input" + ] + }, + "outputs": [], + "source": [ + "def get_decay_products(\n", + " decay: ThreeBodyDecay, subsystem_id: FinalStateID\n", + ") -> tuple[State, State]:\n", + " if subsystem_id not in decay.final_state:\n", + " msg = f\"Subsystem ID {subsystem_id} is not a valid final state ID\"\n", + " raise ValueError(msg)\n", + " return tuple(s for s in decay.final_state.values() if s.index != subsystem_id)\n", + "\n", + "\n", + "plt.rc(\"font\", size=18)\n", + "I_tot = jnp.nansum(intensities)\n", + "normalized_intensities = intensities / I_tot\n", + "\n", + "fig, ax = plt.subplots(figsize=(14, 10))\n", + "mesh = ax.pcolormesh(X, Y, normalized_intensities)\n", + "ax.set_aspect(\"equal\")\n", + "c_bar = plt.colorbar(mesh, ax=ax, pad=0.01)\n", + "c_bar.ax.set_ylabel(\"Normalized intensity (a.u.)\")\n", + "sigma_labels = {\n", + " i: Rf\"$\\sigma_{i} = M^2\\left({' '.join(p.latex for p in get_decay_products(DECAY, i))}\\right)$\"\n", + " for i in (1, 2, 3)\n", + "}\n", + "ax.set_xlabel(sigma_labels[i])\n", + "ax.set_ylabel(sigma_labels[j])\n", + "plt.show()" ] } ], diff --git a/pyproject.toml b/pyproject.toml index 7f663b5a..5ddcdc7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -332,6 +332,7 @@ ignore-names = [ ] "setup.py" = ["D100"] "src/ampform_dpd/io/__init__.py" = ["S403"] +"src/ampform_dpd/io/serialization/amplitude.py" = ["E741"] "src/ampform_dpd/io/serialization/format.py" = ["E741"] "tests/*" = [ "D", diff --git a/src/ampform_dpd/io/serialization/amplitude.py b/src/ampform_dpd/io/serialization/amplitude.py new file mode 100644 index 00000000..29b13624 --- /dev/null +++ b/src/ampform_dpd/io/serialization/amplitude.py @@ -0,0 +1,398 @@ +from __future__ import annotations + +from collections import abc +from itertools import product +from typing import TYPE_CHECKING, Any, Callable, Literal, cast + +import sympy as sp +from ampform.sympy import PoolSum, unevaluated +from sympy.functions.special.tensor_functions import ( + KroneckerDelta as δ, # noqa: N813, PLC2403 +) +from sympy.physics.quantum.cg import CG +from sympy.physics.quantum.spin import Rotation as Wigner + +from ampform_dpd import ( + AmplitudeModel, # pyright:ignore[reportPrivateUsage] + _AlignmentWignerGenerator, # pyright:ignore[reportPrivateUsage] + _generate_amplitude_index_bases, # pyright:ignore[reportPrivateUsage] + create_mass_symbol_mapping, + formulate_invariants, # pyright:ignore[reportPrivateUsage] +) +from ampform_dpd.angles import formulate_scattering_angle +from ampform_dpd.io.serialization.decay import ( + get_final_state, + get_initial_state, + get_spectator_id, + get_states, + to_decay, +) +from ampform_dpd.io.serialization.dynamics import ( + PropagatorDynamicsBuilder, + formulate_dynamics, + formulate_form_factor, + identity_function, +) +from ampform_dpd.io.serialization.format import ( + DecayChain, + HelicityVertex, + LSVertex, + Node, + ParityFactor, + ParityVertex, + get_decay_chains, + get_distribution_def, + get_reference_topology, +) +from ampform_dpd.spin import create_spin_range + +if TYPE_CHECKING: + from ampform_dpd.decay import FinalStateID + from ampform_dpd.io.serialization.format import ModelDefinition + + +def formulate( # noqa: PLR0914 + model: ModelDefinition, + cleanup_summations: bool = False, + to_latex: Callable[[str], str] = identity_function, + additional_builders: dict[str, PropagatorDynamicsBuilder] | None = None, +) -> AmplitudeModel: + states = get_states(model) + helicity_symbols = sp.symbols("lambda(:4)", rational=True) + allowed_helicities = { + symbol: create_spin_range(states[i].spin) # type:ignore[index] + for i, symbol in enumerate(helicity_symbols) + } + amplitude_definitions = {} + angle_definitions = {} + parameter_defaults = {} + n_chains = len(get_decay_chains(model)) + helicity_values: tuple[sp.Rational, sp.Rational, sp.Rational, sp.Rational] + for helicity_values in product(*allowed_helicities.values()): # type:ignore[assignment] + for chain_idx in range(n_chains): + amp_defs = formulate_chain_amplitude( + *helicity_values, model, chain_idx, to_latex, additional_builders + ) + (amp_symbol, amp_expr), *parameters, (θij, θij_expr) = amp_defs.items() + if not isinstance(amp_expr, sp.Expr): + msg = f"Expected an expression, got {amp_expr!r}" + raise TypeError(msg) + helicity_substitutions = dict(zip(helicity_symbols, helicity_values)) + amplitude_definitions[amp_symbol] = amp_expr.subs(helicity_substitutions) + angle_definitions[θij] = θij_expr + parameter_defaults.update(dict(parameters)) + aligned_amp, zeta_defs = formulate_aligned_amplitude(model, *helicity_symbols) + angle_definitions.update(zeta_defs) + decay = to_decay(model) + masses = create_mass_symbol_mapping(decay) + parameter_defaults.update(masses) + if cleanup_summations: + aligned_amp = aligned_amp.cleanup() # type:ignore[assignment] + intensity = PoolSum( + sp.Abs(aligned_amp) ** 2, + *allowed_helicities.items(), + ) + if cleanup_summations: + intensity = intensity.cleanup() # type:ignore[assignment] + return AmplitudeModel( + decay=decay, + intensity=intensity, + amplitudes=amplitude_definitions, # type:ignore[arg-type] + variables=angle_definitions, # type:ignore[arg-type] + parameter_defaults=parameter_defaults, # type:ignore[arg-type] + masses=masses, + invariants=formulate_invariants(decay), + ) + + +def formulate_chain_amplitude( # noqa: PLR0914, PLR0917 + λ0: sp.Rational, + λ1: sp.Rational, + λ2: sp.Rational, + λ3: sp.Rational, + model: ModelDefinition, + chain_idx: int, + to_latex: Callable[[str], str] = identity_function, + additional_builders: dict[str, PropagatorDynamicsBuilder] | None = None, +) -> dict[sp.Symbol, complex | float | sp.Expr]: + chain_defs = get_decay_chains(model) + chain_definition = chain_defs[chain_idx] + # ----------------------- + dynamics = formulate_dynamics( + chain_definition, model, to_latex, additional_builders + ) + for vertex in chain_definition["vertices"]: + dynamics *= formulate_form_factor(vertex, model) + # ----------------------- + weight, weight_val = _get_weight(chain_definition) + # ----------------------- + (i, λi_val), (j, λj_val) = _get_decay_product_helicities(chain_definition) + θij, θij_expr = formulate_scattering_angle(i, j) + jR = sp.Rational(chain_definition["propagators"][0]["spin"]) # noqa: N806 + R_node, λR_val = _get_resonance_helicity(chain_definition) # noqa: N806 + λR = _get_helicity_symbol(R_node) + # ----------------------- + A = _generate_amplitude_index_bases() + subsystem_id = get_spectator_id(chain_definition["topology"]) + h_prod = formulate_recoupling(model, chain_idx, vertex_idx=0) + h_dec = formulate_recoupling(model, chain_idx, vertex_idx=1) + amplitude_expression = ( + weight + * h_prod + * h_dec + * Wigner.d(jR, λR, λi_val - λj_val, θij) + * dynamics.expression + ) + amplitude_expression = amplitude_expression.subs({λR: λR_val}) + amplitude_symbol = A[subsystem_id][λ0, λ1, λ2, λ3] + return { + amplitude_symbol: amplitude_expression, + weight: weight_val, + **dynamics.definitions, + θij: θij_expr, + } + + +def _get_decay_product_helicities( + chain_definition: DecayChain, +) -> tuple[tuple[int, sp.Rational], tuple[int, sp.Rational]]: + vertices = chain_definition["vertices"] + for vertex in vertices: + node = vertex["node"] + if all(isinstance(i, int) for i in node): + helicities = vertex.get("helicities") + if helicities is None: + msg = "Vertex does not contain helicities. Is it an LS vertex?" + raise ValueError(msg, vertex) + return tuple((i, sp.Rational(λ)) for i, λ in zip(node, helicities)) # type:ignore[assignment,call-overload,return-value] + msg = "Could not fine a helicity for any resonance node" + raise ValueError(msg) + + +def formulate_aligned_amplitude( + model: ModelDefinition, + λ0: sp.Rational | sp.Symbol, + λ1: sp.Rational | sp.Symbol, + λ2: sp.Rational | sp.Symbol, + λ3: sp.Rational | sp.Symbol, +) -> tuple[PoolSum, dict[sp.Symbol, sp.Expr]]: + reference_topology = get_reference_topology(model) + reference_subsystem = get_spectator_id(reference_topology) + wigner_generator = _AlignmentWignerGenerator(reference_subsystem) + _λ0, _λ1, _λ2, _λ3 = sp.symbols(R"\lambda_(:4)^{\prime}", rational=True) + states = get_states(model) + j0, j1, j2, j3 = (states[i].spin for i in sorted(states)) + A = _generate_amplitude_index_bases() + amp_expr = PoolSum( + sum( + A[k][_λ0, _λ1, _λ2, _λ3] + * wigner_generator(j0, λ0, _λ0, rotated_state=0, aligned_subsystem=k) + * wigner_generator(j1, _λ1, λ1, rotated_state=1, aligned_subsystem=k) + * wigner_generator(j2, _λ2, λ2, rotated_state=2, aligned_subsystem=k) + * wigner_generator(j3, _λ3, λ3, rotated_state=3, aligned_subsystem=k) + for k in get_existing_subsystem_ids(model) + ), + (_λ0, create_spin_range(j0)), + (_λ1, create_spin_range(j1)), + (_λ2, create_spin_range(j2)), + (_λ3, create_spin_range(j3)), + ) + return amp_expr, wigner_generator.angle_definitions # type:ignore[return-value] + + +def _get_weight( + chain_definition: DecayChain, to_latex: Callable[[str], str] = identity_function +) -> tuple[sp.Symbol, complex | float]: + value: complex | float + value = complex(str(chain_definition["weight"]).replace(" ", "").replace("i", "j")) + if not value.imag: + value = value.real + resonance_latex = to_latex(chain_definition["name"]) + _, resonance_helicity = _get_resonance_helicity(chain_definition) + c = sp.IndexedBase(f"c^{{{resonance_latex}[{resonance_helicity}]}}") + λ1, λ2, λ3 = _get_final_state_helicities(chain_definition).values() + symbol = c[λ1, λ2, λ3] + return symbol, value + + +def _get_resonance_helicity( + chain_definition: DecayChain, +) -> tuple[tuple[FinalStateID, FinalStateID], sp.Rational]: + vertices = chain_definition["vertices"] + for vertex in vertices: + node = vertex["node"] + if all(isinstance(i, int) for i in node): + continue + vertex = cast(HelicityVertex, vertex) + helicities = vertex.get("helicities") + if helicities is None: # pyright:ignore[reportUnnecessaryComparison] + msg = "Vertex does not contain helicities. Is it an LS vertex?" + raise ValueError(msg, vertex) + for helicity, sub_node in zip(helicities, node): + if isinstance(sub_node, abc.Sequence) and len(sub_node) == 2: # noqa: PLR2004 + return tuple(sub_node), sp.Rational(helicity) # type:ignore[return-value] + msg = "Could not find a resonance node" + raise ValueError(msg) + + +def _get_final_state_helicities( + chain_definition: DecayChain, +) -> dict[FinalStateID, sp.Rational]: + vertices = chain_definition["vertices"] + collected_helicities: dict[FinalStateID, sp.Rational] = {} + for vertex in vertices: + vertex = cast(HelicityVertex, vertex) + helicities = vertex.get("helicities") + if helicities is None: # pyright:ignore[reportUnnecessaryComparison] + msg = "Vertex does not contain helicities. Is it an LS vertex?" + raise ValueError(msg, vertex) # type:ignore[index] + for helicity, node in zip(helicities, vertex["node"]): + if not isinstance(node, int): + continue + collected_helicities[node] = sp.Rational(helicity) + return {i: collected_helicities[i] for i in sorted(collected_helicities)} + + +def formulate_recoupling( # noqa: PLR0914 + model: ModelDefinition, chain_idx: int, vertex_idx: int +) -> sp.Expr: + chain_definition = get_decay_chains(model)[chain_idx] + vertex_definitions = chain_definition["vertices"] + if len(vertex_definitions) != 2: # noqa: PLR2004 + msg = f"Not a three-body decay: there are {len(vertex_definitions)} vertices" + raise ValueError(msg) + if vertex_idx not in {0, 1}: + msg = f"Vertex index out of range. Can either be 0 or 1, not {vertex_idx}." + raise ValueError(msg) + vertex = chain_definition["vertices"][vertex_idx] + vertex_type = vertex["type"] + node = vertex["node"] + λa, λb = map(_get_helicity_symbol, node) + if vertex_type in {"helicity", "parity"}: + vertex = cast(HelicityVertex, vertex) + λa0, λb0 = (sp.Rational(v) for v in vertex["helicities"]) + if vertex_type == "parity": + vertex = cast(ParityVertex, vertex) + f = _sign_to_value(vertex.get("parity_factor", "+")) + return ParityRecoupling(λa, λb, λa0, λb0, f) + return HelicityRecoupling(λa, λb, λa0, λb0) + if vertex_type == "ls": + vertex = cast(LSVertex, vertex) + l = int(vertex["l"]) + s = sp.Rational(vertex["s"]) + ja, jb = _get_child_spins(model, chain_idx, vertex_idx) + j = _get_parent_spin(model, chain_idx, vertex_idx) + return LSRecoupling(λa, λb, l, s, ja, jb, j) + msg = f"No implementation for vertex of type {vertex_type!r}" + raise NotImplementedError(msg) + + +def _sign_to_value(sign: ParityFactor) -> Literal[0, -1, 1]: + stripped_sign = sign.strip() + if stripped_sign == "-": + return -1 + if not stripped_sign: + return 0 + if stripped_sign == "+": + return +1 + msg = f"Cannot convert {sign!r} to value" + raise NotImplementedError(msg) + + +def _get_parent_spin( + model: ModelDefinition, chain_idx: int, vertex_idx: int +) -> sp.Rational: + chain_definition = get_decay_chains(model)[chain_idx] + vertex = chain_definition["vertices"][vertex_idx] + if all(isinstance(i, int) for i in vertex["node"]): + return __get_propagator_spin(chain_definition) + initial_state = get_initial_state(model) + return initial_state.spin + + +def _get_child_spins( + model: ModelDefinition, chain_idx: int, vertex_idx: int +) -> tuple[sp.Rational, sp.Rational]: + chain_definition = get_decay_chains(model)[chain_idx] + vertex = chain_definition["vertices"][vertex_idx] + node = vertex["node"] + final_state = get_final_state(model) + spins = [] + for node_item in node: + if isinstance(node_item, int): + spins.append(sp.Rational(final_state[node_item])) + else: + spins.append(__get_propagator_spin(chain_definition)) + return tuple(spins) + + +def __get_propagator_spin(chain_definition: DecayChain) -> sp.Rational: + propagators = chain_definition["propagators"] + if len(propagators) != 1: + msg = f"There are {len(propagators)} propagators, not a three-body decay" + raise ValueError(msg) + return sp.Rational(propagators[0]["spin"]) + + +def _get_helicity_symbol(node: int | Node) -> sp.Symbol: + if isinstance(node, int): + return sp.Symbol(f"lambda{node}", rational=True) + return sp.Symbol(R"\lambda_R", rational=True) + + +def get_existing_subsystem_ids(model: ModelDefinition) -> list[FinalStateID]: + distribution_def = get_distribution_def(model) + chain_defs = distribution_def["decay_description"]["chains"] + subsystem_ids = {get_spectator_id(c["topology"]) for c in chain_defs} + return sorted(subsystem_ids) + + +@unevaluated +class HelicityRecoupling(sp.Expr): + λa: sp.Rational | sp.Symbol + λb: sp.Rational | sp.Symbol + λa0: sp.Rational | sp.Symbol + λb0: sp.Rational | sp.Symbol + _latex_repr_ = R"\mathcal{{H}}^\text{{helicity}}\left({λa},{λb}|{λa0},{λb0}\right)" + + def evaluate(self) -> sp.Expr: + λa, λb, λa0, λb0 = self.args + return δ(λa, λa0) * δ(λb, λb0) + + +@unevaluated +class ParityRecoupling(sp.Expr): + λa: Any + λb: Any + λa0: Any + λb0: Any + f: Any + _latex_repr_ = ( + R"\mathcal{{H}}^\text{{parity}}\left({λa},{λb}|{λa0},{λb0},{f}\right)" + ) + + def evaluate(self) -> sp.Expr: + λa, λb, λa0, λb0, f = self.args + return δ(λa, λa0) * δ(λb, λb0) + f * δ(λa, -λa0) * δ(λb, -λb0) # type:ignore[operator] + + +@unevaluated +class LSRecoupling(sp.Expr): + λa: Any + λb: Any + l: Any + s: Any + ja: Any + jb: Any + j: Any + _latex_repr_ = ( + R"\mathcal{{H}}^\text{{parity}}\left({λa},{λb}|{l},{s},{ja},{jb},{j}\right)" + ) + + def evaluate(self) -> sp.Expr: + λa, λb, l, s, ja, jb, j = self.args + return ( + sp.sqrt((2 * l + 1) / (2 * j + 1)) # type:ignore[operator] + * CG(ja, λa, jb, -λb, s, λa - λb) # type:ignore[operator] + * CG(l, 0, s, λa - λb, j, λa - λb) # type:ignore[operator] + ) diff --git a/src/ampform_dpd/io/serialization/decay.py b/src/ampform_dpd/io/serialization/decay.py index 97102b0c..ae2bd893 100644 --- a/src/ampform_dpd/io/serialization/decay.py +++ b/src/ampform_dpd/io/serialization/decay.py @@ -13,7 +13,11 @@ ThreeBodyDecay, ThreeBodyDecayChain, ) -from ampform_dpd.io.serialization.format import get_decay_chains, get_distribution_def +from ampform_dpd.io.serialization.format import ( + Topology, + get_decay_chains, + get_distribution_def, +) if TYPE_CHECKING: from ampform_dpd.io.serialization.format import ( @@ -134,3 +138,16 @@ def _to_particle( parity=None, index=definition["index"], ) + + +def get_spectator_id(topology: Topology) -> FinalStateID: + """Get the spectator ID from a reference topology. + + >>> get_spectator_id([1, [2, 3]]) + 1 + """ + spectator_candidates = {i for i in topology if isinstance(i, int)} + if len(spectator_candidates) != 1: + msg = f"Reference topology {topology} seems not to be a three-body decay" + raise ValueError(msg) + return next(iter(spectator_candidates)) diff --git a/src/ampform_dpd/io/serialization/format.py b/src/ampform_dpd/io/serialization/format.py index d3fdc571..31504f3a 100644 --- a/src/ampform_dpd/io/serialization/format.py +++ b/src/ampform_dpd/io/serialization/format.py @@ -156,3 +156,8 @@ def get_function_definition( if candidates: msg += f" Did you mean any of these? {', '.join(sorted(candidates))}" raise KeyError(msg) + + +def get_reference_topology(model: ModelDefinition) -> Topology: + distribution_def = get_distribution_def(model) + return distribution_def["decay_description"]["reference_topology"] diff --git a/tests/io_serialization/test_amplitude.py b/tests/io_serialization/test_amplitude.py new file mode 100644 index 00000000..3139084c --- /dev/null +++ b/tests/io_serialization/test_amplitude.py @@ -0,0 +1,63 @@ +# pyright:reportPrivateUsage=false +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +import sympy as sp + +from ampform_dpd.io.serialization.amplitude import ( + _get_decay_product_helicities, + _get_final_state_helicities, + _get_resonance_helicity, + _get_weight, + get_existing_subsystem_ids, +) +from ampform_dpd.io.serialization.format import get_decay_chains + +if TYPE_CHECKING: + from ampform_dpd.io.serialization.format import ModelDefinition + + +def test_get_decay_product_helicities(model_definition: ModelDefinition): + chain_defs = get_decay_chains(model_definition) + half = sp.Rational(1 / 2) + assert _get_decay_product_helicities(chain_defs[0]) == ((3, 0), (1, +half)) + assert _get_decay_product_helicities(chain_defs[15]) == ((1, +half), (2, 0)) + assert _get_decay_product_helicities(chain_defs[-1]) == ((2, 0), (3, 0)) + + +def test_get_existing_subsystem_ids(model_definition: ModelDefinition): + assert get_existing_subsystem_ids(model_definition) == [1, 2, 3] + + +@pytest.mark.parametrize("chain_id", range(26)) +def test_get_final_state_helicities(model_definition: ModelDefinition, chain_id: int): + chain_defs = get_decay_chains(model_definition) + assert len(chain_defs) == 26 + if chain_id in {19, 20, 22, 25}: + λp = -sp.Rational(1 / 2) + else: + λp = +sp.Rational(1 / 2) + assert _get_final_state_helicities(chain_defs[chain_id]) == {1: λp, 2: 0, 3: 0} + + +def test_get_resonance_helicity(model_definition: ModelDefinition): + chain_defs = get_decay_chains(model_definition) + half = sp.Rational(1 / 2) + node, helicity = _get_resonance_helicity(chain_defs[0]) + assert node == (3, 1) + assert helicity == +half + node, helicity = _get_resonance_helicity(chain_defs[1]) + assert node == (3, 1) + assert helicity == -half + node, helicity = _get_resonance_helicity(chain_defs[-1]) + assert node == (2, 3) + assert helicity == 0 + + +def test_get_weight(model_definition: ModelDefinition): + chain_defs = get_decay_chains(model_definition) + symbol, value = _get_weight(chain_defs[0]) + assert symbol.name == "c^{L1405[1/2]}[1/2, 0, 0]" + assert value == 7.38649400481717 + 1.971018433257411j From 07b32608f976c4b3be6cb32cccd26ead9c34236e Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Mon, 20 May 2024 20:18:51 +0200 Subject: [PATCH 21/42] DOC: remove markdown table for decay --- docs/serialization.ipynb | 24 +++--------------------- 1 file changed, 3 insertions(+), 21 deletions(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 6450291b..766aad89 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -44,7 +44,7 @@ "from ampform.dynamics.phasespace import BreakupMomentumSquared\n", "from ampform.kinematics.phasespace import Kallen\n", "from ampform.sympy import perform_cached_doit\n", - "from IPython.display import JSON, Markdown, Math\n", + "from IPython.display import JSON, Math\n", "from tqdm.auto import tqdm\n", "\n", "from ampform_dpd import DefinedExpression\n", @@ -59,12 +59,7 @@ " P,\n", " SimpleBreitWigner,\n", ")\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.io import aslatex, perform_cached_lambdify, simplify_latex_rendering\n", "from ampform_dpd.io.serialization.amplitude import (\n", " HelicityRecoupling,\n", " LSRecoupling,\n", @@ -170,20 +165,7 @@ "outputs": [], "source": [ "DECAY = to_decay(MODEL_DEFINITION, to_latex=to_latex)\n", - "Math(aslatex(DECAY))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "Markdown(as_markdown_table(DECAY))" + "Math(aslatex(DECAY, with_jp=True))" ] }, { From a8cf5a0665318f0d9bc870de89a4201870604f4d Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 20:10:15 +0200 Subject: [PATCH 22/42] DOC: split lines in amplitude rendering --- docs/serialization.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 766aad89..8394caeb 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -681,7 +681,7 @@ }, "outputs": [], "source": [ - "Math(aslatex(MODEL.amplitudes))" + "Math(aslatex(MODEL.amplitudes, terms_per_line=1))" ] }, { From 00a0ac17a7be1fd282a922551bde68c3b4b1dde2 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 16:09:49 +0200 Subject: [PATCH 23/42] FIX: make `MultichannelBreitWigner` immutable --- docs/serialization.ipynb | 4 ++-- src/ampform_dpd/dynamics/__init__.py | 2 +- src/ampform_dpd/io/serialization/dynamics.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 8394caeb..a417b3b2 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -337,7 +337,7 @@ "source": [ "x, y, z = sp.symbols(\"x:z\")\n", "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", - "channels = [\n", + "channels = tuple(\n", " ChannelArguments(\n", " sp.Symbol(f\"Gamma{i}\"),\n", " sp.Symbol(f\"m_{{a,{i}}}\"),\n", @@ -346,7 +346,7 @@ " d,\n", " )\n", " for i in [1, 2]\n", - "]\n", + ")\n", "exprs = [\n", " MultichannelBreitWigner(s, m0, channels),\n", " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", diff --git a/src/ampform_dpd/dynamics/__init__.py b/src/ampform_dpd/dynamics/__init__.py index b622d315..056f3a71 100644 --- a/src/ampform_dpd/dynamics/__init__.py +++ b/src/ampform_dpd/dynamics/__init__.py @@ -221,7 +221,7 @@ def evaluate(self): class MultichannelBreitWigner(sp.Expr): s: Any mass: Any - channels: list[ChannelArguments] = argument(sympify=False) + channels: tuple[ChannelArguments, ...] = argument(sympify=False) def evaluate(self): s = self.s diff --git a/src/ampform_dpd/io/serialization/dynamics.py b/src/ampform_dpd/io/serialization/dynamics.py index 31b49a79..09941c62 100644 --- a/src/ampform_dpd/io/serialization/dynamics.py +++ b/src/ampform_dpd/io/serialization/dynamics.py @@ -164,7 +164,7 @@ def formulate_multichannel_breit_wigner( # noqa: PLR0914 mi2: channel_definition["mb"], }) return DefinedExpression( - expression=MultichannelBreitWigner(s, mass, channels), + expression=MultichannelBreitWigner(s, mass, tuple(channels)), definitions=parameter_defaults, ) From 88b22e49e487c175b3c8b496e162491704dfa58f Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 16:11:34 +0200 Subject: [PATCH 24/42] FIX: write latex resonance names in weights --- src/ampform_dpd/io/serialization/amplitude.py | 2 +- tests/io_serialization/test_amplitude.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/io/serialization/amplitude.py b/src/ampform_dpd/io/serialization/amplitude.py index 29b13624..5c346ae4 100644 --- a/src/ampform_dpd/io/serialization/amplitude.py +++ b/src/ampform_dpd/io/serialization/amplitude.py @@ -124,7 +124,7 @@ def formulate_chain_amplitude( # noqa: PLR0914, PLR0917 for vertex in chain_definition["vertices"]: dynamics *= formulate_form_factor(vertex, model) # ----------------------- - weight, weight_val = _get_weight(chain_definition) + weight, weight_val = _get_weight(chain_definition, to_latex) # ----------------------- (i, λi_val), (j, λj_val) = _get_decay_product_helicities(chain_definition) θij, θij_expr = formulate_scattering_angle(i, j) diff --git a/tests/io_serialization/test_amplitude.py b/tests/io_serialization/test_amplitude.py index 3139084c..f8ed8c31 100644 --- a/tests/io_serialization/test_amplitude.py +++ b/tests/io_serialization/test_amplitude.py @@ -59,5 +59,5 @@ def test_get_resonance_helicity(model_definition: ModelDefinition): def test_get_weight(model_definition: ModelDefinition): chain_defs = get_decay_chains(model_definition) symbol, value = _get_weight(chain_defs[0]) - assert symbol.name == "c^{L1405[1/2]}[1/2, 0, 0]" + assert symbol.name == R"c^{L1405[1/2]}_{\frac{1}{2}, 0, 0}" assert value == 7.38649400481717 + 1.971018433257411j From 74ff70d5c77007fe9fdeb86766a6e7a445e6cf1b Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 16:23:50 +0200 Subject: [PATCH 25/42] FIX: improve symbol assumptions --- docs/jpsi2ksp.ipynb | 2 +- docs/lc2pkpi.ipynb | 10 +++++----- docs/serialization.ipynb | 18 +++++++++--------- docs/xib2pkk.ipynb | 2 +- src/ampform_dpd/dynamics/builder.py | 4 ++-- src/ampform_dpd/io/serialization/dynamics.py | 16 ++++++++-------- 6 files changed, 26 insertions(+), 26 deletions(-) diff --git a/docs/jpsi2ksp.ipynb b/docs/jpsi2ksp.ipynb index 984f99e7..20a91d0f 100644 --- a/docs/jpsi2ksp.ipynb +++ b/docs/jpsi2ksp.ipynb @@ -247,7 +247,7 @@ }, "outputs": [], "source": [ - "s, m0, w0, m1, m2, L, R, z = sp.symbols(\"s m0 Gamma0 m1 m2 L R z\")\n", + "s, m0, w0, m1, m2, L, R, z = sp.symbols(\"s m0 Gamma0 m1 m2 L R z\", nonnegative=True)\n", "exprs = [\n", " RelativisticBreitWigner(s, m0, w0, m1, m2, L, R),\n", " EnergyDependentWidth(s, m0, w0, m1, m2, L, R),\n", diff --git a/docs/lc2pkpi.ipynb b/docs/lc2pkpi.ipynb index 39e098ed..538efef0 100644 --- a/docs/lc2pkpi.ipynb +++ b/docs/lc2pkpi.ipynb @@ -169,9 +169,9 @@ "outputs": [], "source": [ "s, m0, Γ0, m1, m2 = sp.symbols(\"s m0 Gamma0 m1 m2\", nonnegative=True)\n", - "m_top, m_spec = sp.symbols(R\"m_\\mathrm{top} m_\\mathrm{spectator}\")\n", - "R_dec, R_prod = sp.symbols(R\"R_\\mathrm{res} R_{\\Lambda_c}\")\n", - "l_Λc, l_R = sp.symbols(R\"l_{\\Lambda_c} l_R\", integer=True, positive=True)\n", + "m_top, m_spec = sp.symbols(R\"m_\\mathrm{top} m_\\mathrm{spectator}\", nonnegative=True)\n", + "R_dec, R_prod = sp.symbols(R\"R_\\mathrm{res} R_{\\Lambda_c}\", nonnegative=True)\n", + "l_Λc, l_R = sp.symbols(R\"l_{\\Lambda_c} l_R\", integer=True, nonnegative=True)\n", "bw = BreitWignerMinL(s, m_top, m_spec, m0, Γ0, m1, m2, l_R, l_Λc, R_dec, R_prod)\n", "Latex(aslatex({bw: bw.doit(deep=False)}))" ] @@ -205,8 +205,8 @@ " 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", + " R_dec = sp.Symbol(R\"R_\\mathrm{res}\", nonnegative=True)\n", + " R_prod = sp.Symbol(R\"R_{\\Lambda_c}\", nonnegative=True)\n", " parameter_defaults = {\n", " parent_mass: decay_chain.parent.mass,\n", " spectator_mass: decay_chain.spectator.mass,\n", diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index a417b3b2..3f0111ac 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -229,7 +229,7 @@ "outputs": [], "source": [ "z = sp.Symbol(\"z\", nonnegative=True)\n", - "s, m1, m2, L, d = sp.symbols(\"s m1 m2 L R\")\n", + "s, m1, m2, L, d = sp.symbols(\"s m1 m2 L R\", nonnegative=True)\n", "exprs = [\n", " FormFactor(s, m1, m2, L, d),\n", " BlattWeisskopfSquared(z, L),\n", @@ -285,7 +285,7 @@ "outputs": [], "source": [ "x, y, z = sp.symbols(\"x:z\")\n", - "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", + "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\", nonnegative=True)\n", "exprs = [\n", " BreitWigner(s, m0, Γ0, m1, m2, L, d),\n", " SimpleBreitWigner(s, m0, Γ0),\n", @@ -336,13 +336,13 @@ "outputs": [], "source": [ "x, y, z = sp.symbols(\"x:z\")\n", - "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\")\n", + "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\", nonnegative=True)\n", "channels = tuple(\n", " ChannelArguments(\n", - " sp.Symbol(f\"Gamma{i}\"),\n", - " sp.Symbol(f\"m_{{a,{i}}}\"),\n", - " sp.Symbol(f\"m_{{b,{i}}}\"),\n", - " sp.Symbol(f\"L{i}\"),\n", + " sp.Symbol(f\"Gamma{i}\", nonnegative=True),\n", + " sp.Symbol(f\"m_{{a,{i}}}\", nonnegative=True),\n", + " sp.Symbol(f\"m_{{b,{i}}}\", nonnegative=True),\n", + " sp.Symbol(f\"L{i}\", integer=True, nonnegative=True),\n", " d,\n", " )\n", " for i in [1, 2]\n", @@ -406,7 +406,7 @@ }, "outputs": [], "source": [ - "s, m0, Γ0, m1, m2, γ = sp.symbols(\"s m0 Gamma0 m1 m2 gamma\")\n", + "s, m0, Γ0, m1, m2, γ = sp.symbols(\"s m0 Gamma0 m1 m2 gamma\", nonnegative=True)\n", "expr = BuggBreitWigner(s, m0, Γ0, m1, m2, γ)\n", "Math(aslatex({expr: expr.doit(deep=False)}))" ] @@ -446,7 +446,7 @@ " s = to_mandelstam_symbol(node)\n", " mass = sp.Symbol(f\"m_{{{resonance}}}\", nonnegative=True)\n", " width = sp.Symbol(Rf\"\\Gamma_{{{resonance}}}\", nonnegative=True)\n", - " γ = sp.Symbol(Rf\"\\gamma_{{{resonance}}}\")\n", + " γ = sp.Symbol(Rf\"\\gamma_{{{resonance}}}\", nonnegative=True)\n", " m1 = to_mass_symbol(i)\n", " m2 = to_mass_symbol(j)\n", " final_state = get_final_state(model)\n", diff --git a/docs/xib2pkk.ipynb b/docs/xib2pkk.ipynb index e25629ed..dfa1d576 100644 --- a/docs/xib2pkk.ipynb +++ b/docs/xib2pkk.ipynb @@ -273,7 +273,7 @@ }, "outputs": [], "source": [ - "s, m0, w0, m1, m2, L, R, z = sp.symbols(\"s m0 Gamma0 m1 m2 L R z\")\n", + "s, m0, w0, m1, m2, L, R, z = sp.symbols(\"s m0 Gamma0 m1 m2 L R z\", nonnegative=True)\n", "exprs = [\n", " RelativisticBreitWigner(s, m0, w0, m1, m2, L, R),\n", " EnergyDependentWidth(s, m0, w0, m1, m2, L, R),\n", diff --git a/src/ampform_dpd/dynamics/builder.py b/src/ampform_dpd/dynamics/builder.py index 15c96c76..eeba3325 100644 --- a/src/ampform_dpd/dynamics/builder.py +++ b/src/ampform_dpd/dynamics/builder.py @@ -102,8 +102,8 @@ def _get_angular_momentum(isobar: IsobarNode) -> int: def _create_meson_radius_symbol(isobar: IsobarNode) -> sp.Symbol: if isinstance(isobar.parent, State): - return sp.Symbol(Rf"R_{{{isobar.parent.latex}}}") - return sp.Symbol(R"R_\mathrm{res}") + return sp.Symbol(Rf"R_{{{isobar.parent.latex}}}", nonnegative=True) + return sp.Symbol(R"R_\mathrm{res}", nonnegative=True) def create_mass_symbol(particle: IsobarNode | Particle) -> sp.Symbol: diff --git a/src/ampform_dpd/io/serialization/dynamics.py b/src/ampform_dpd/io/serialization/dynamics.py index 09941c62..4a311eff 100644 --- a/src/ampform_dpd/io/serialization/dynamics.py +++ b/src/ampform_dpd/io/serialization/dynamics.py @@ -84,10 +84,10 @@ def formulate_form_factor(vertex: Vertex, model: ModelDefinition) -> DefinedExpr s = to_mandelstam_symbol(node) m1, m2 = (to_mass_symbol(i) for i in node) if all(isinstance(i, int) for i in node): - meson_radius = sp.Symbol(R"R_\mathrm{res}") + meson_radius = sp.Symbol(R"R_\mathrm{res}", nonnegative=True) else: initial_state = get_initial_state(model) - meson_radius = sp.Symbol(f"R_{{{initial_state.latex}}}") + meson_radius = sp.Symbol(f"R_{{{initial_state.latex}}}", nonnegative=True) angular_momentum = int(function_definition["l"]) return DefinedExpression( expression=FormFactor(s, m1, m2, angular_momentum, meson_radius), @@ -107,12 +107,12 @@ def formulate_breit_wigner( node = propagator["node"] i, j = node s = to_mandelstam_symbol(node) - mass = sp.Symbol(f"m_{{{resonance}}}") - width = sp.Symbol(Rf"\Gamma_{{{resonance}}}") + mass = sp.Symbol(f"m_{{{resonance}}}", nonnegative=True) + width = sp.Symbol(Rf"\Gamma_{{{resonance}}}", nonnegative=True) m1 = to_mass_symbol(i) m2 = to_mass_symbol(j) angular_momentum = int(function_definition["l"]) - d = sp.Symbol(R"R_\mathrm{res}") + d = sp.Symbol(R"R_\mathrm{res}", nonnegative=True) return DefinedExpression( expression=BreitWigner(s, mass, width, m1, m2, angular_momentum, d), definitions={ @@ -137,12 +137,12 @@ def formulate_multichannel_breit_wigner( # noqa: PLR0914 node = propagator["node"] i, j = node s = to_mandelstam_symbol(node) - mass = sp.Symbol(f"m_{{{resonance}}}") - width = sp.Symbol(Rf"\Gamma_{{{resonance}}}") + mass = sp.Symbol(f"m_{{{resonance}}}", nonnegative=True) + width = sp.Symbol(Rf"\Gamma_{{{resonance}}}", nonnegative=True) m1 = to_mass_symbol(i) m2 = to_mass_symbol(j) angular_momentum = int(channel_definitions[0]["l"]) - d = sp.Symbol(f"R_{{{resonance}}}") + d = sp.Symbol(f"R_{{{resonance}}}", nonnegative=True) channels = [ChannelArguments(width, m1, m2, angular_momentum, d)] parameter_defaults: dict[sp.Symbol, complex | float] = { mass: function_definition["mass"], From 641c80985ca46dd2fa3dcdfc668733eb98588a8c Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 16:28:13 +0200 Subject: [PATCH 26/42] FIX: insert widths into definitions --- src/ampform_dpd/io/serialization/dynamics.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/ampform_dpd/io/serialization/dynamics.py b/src/ampform_dpd/io/serialization/dynamics.py index 4a311eff..8b9dbd37 100644 --- a/src/ampform_dpd/io/serialization/dynamics.py +++ b/src/ampform_dpd/io/serialization/dynamics.py @@ -162,6 +162,7 @@ def formulate_multichannel_breit_wigner( # noqa: PLR0914 parameter_defaults.update({ mi1: channel_definition["ma"], mi2: channel_definition["mb"], + Γi: channel_definition["gsq"], }) return DefinedExpression( expression=MultichannelBreitWigner(s, mass, tuple(channels)), From 73221d94e995707e1ab55b41767a6efdfdfd1473 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 18:13:12 +0200 Subject: [PATCH 27/42] BREAK: make `ChannelArguments` an expression class --- docs/serialization.ipynb | 12 +++++---- src/ampform_dpd/dynamics/__init__.py | 28 ++++++++------------ src/ampform_dpd/io/serialization/dynamics.py | 4 +-- 3 files changed, 20 insertions(+), 24 deletions(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 3f0111ac..e549b1a9 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -339,11 +339,13 @@ "s, m0, Γ0, m1, m2, L, d = sp.symbols(\"s m0 Gamma0 m1 m2 L R\", nonnegative=True)\n", "channels = tuple(\n", " ChannelArguments(\n", - " sp.Symbol(f\"Gamma{i}\", nonnegative=True),\n", - " sp.Symbol(f\"m_{{a,{i}}}\", nonnegative=True),\n", - " sp.Symbol(f\"m_{{b,{i}}}\", nonnegative=True),\n", - " sp.Symbol(f\"L{i}\", integer=True, nonnegative=True),\n", - " d,\n", + " s,\n", + " m0,\n", + " width=sp.Symbol(f\"Gamma{i}\", nonnegative=True),\n", + " m1=sp.Symbol(f\"m_{{a,{i}}}\", nonnegative=True),\n", + " m2=sp.Symbol(f\"m_{{b,{i}}}\", nonnegative=True),\n", + " angular_momentum=sp.Symbol(f\"L{i}\", integer=True, nonnegative=True),\n", + " meson_radius=d,\n", " )\n", " for i in [1, 2]\n", ")\n", diff --git a/src/ampform_dpd/dynamics/__init__.py b/src/ampform_dpd/dynamics/__init__.py index 056f3a71..92d5b706 100644 --- a/src/ampform_dpd/dynamics/__init__.py +++ b/src/ampform_dpd/dynamics/__init__.py @@ -8,8 +8,7 @@ import sympy as sp from ampform.dynamics import formulate_form_factor from ampform.kinematics.phasespace import Kallen -from ampform.sympy import argument, unevaluated -from attrs import asdict, frozen +from ampform.sympy import unevaluated if TYPE_CHECKING: from sympy.printing.latex import LatexPrinter @@ -221,12 +220,12 @@ def evaluate(self): class MultichannelBreitWigner(sp.Expr): s: Any mass: Any - channels: tuple[ChannelArguments, ...] = argument(sympify=False) + channels: tuple[ChannelArguments, ...] def evaluate(self): s = self.s m0 = self.mass - width = sum(channel.formulate_width(s, m0) for channel in self.channels) + width = sum(channel.evaluate() for channel in self.channels) return BreitWigner(s, m0, width) def _latex_repr_(self, printer: LatexPrinter, *args) -> str: @@ -237,26 +236,21 @@ def _latex_repr_(self, printer: LatexPrinter, *args) -> str: return latex -@frozen -class ChannelArguments: +@unevaluated +class ChannelArguments(sp.Expr): + s: Any + m0: Any width: Any m1: Any = 0 m2: Any = 0 angular_momentum: Any = 0 meson_radius: Any = 1 + _latex_repr_ = R"\Gamma^\text{channel}\left({{s}}, {{m0}}, {{width}}\right)" - def __attrs_post_init__(self) -> None: - for name, value in asdict(self).items(): - object.__setattr__(self, name, sp.sympify(value)) - - def formulate_width(self, s: Any, m0: Any) -> sp.Expr: - Γ0 = self.width - m1 = self.m1 - m2 = self.m2 - L = self.angular_momentum - R = self.meson_radius + def evaluate(self) -> sp.Expr: + s, m0, Γ0, m1, m2, L, R = self.args ff = FormFactor(s, m1, m2, L, R) ** 2 - return Γ0 * m0 / sp.sqrt(s) * ff + return Γ0 * m0 / sp.sqrt(s) * ff # type:ignore[operator] @unevaluated diff --git a/src/ampform_dpd/io/serialization/dynamics.py b/src/ampform_dpd/io/serialization/dynamics.py index 8b9dbd37..aec2f369 100644 --- a/src/ampform_dpd/io/serialization/dynamics.py +++ b/src/ampform_dpd/io/serialization/dynamics.py @@ -143,7 +143,7 @@ def formulate_multichannel_breit_wigner( # noqa: PLR0914 m2 = to_mass_symbol(j) angular_momentum = int(channel_definitions[0]["l"]) d = sp.Symbol(f"R_{{{resonance}}}", nonnegative=True) - channels = [ChannelArguments(width, m1, m2, angular_momentum, d)] + channels = [ChannelArguments(s, mass, width, m1, m2, angular_momentum, d)] parameter_defaults: dict[sp.Symbol, complex | float] = { mass: function_definition["mass"], width: channel_definitions[0]["gsq"], @@ -158,7 +158,7 @@ def formulate_multichannel_breit_wigner( # noqa: PLR0914 mi1 = sp.Symbol(f"m_{{a,{channel_idx}}}", nonnegative=True) mi2 = sp.Symbol(f"m_{{b,{channel_idx}}}", nonnegative=True) angular_momentum = int(channel_definition["l"]) - channels.append(ChannelArguments(Γi, mi1, mi2, angular_momentum, d)) + channels.append(ChannelArguments(s, mass, Γi, mi1, mi2, angular_momentum, d)) parameter_defaults.update({ mi1: channel_definition["ma"], mi2: channel_definition["mb"], From b297eb1b7a84c78fa422c851055237b65711a044 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 19:41:56 +0200 Subject: [PATCH 28/42] BEHAVIOR: write weight as `Symbol` not `Indexed` --- src/ampform_dpd/io/serialization/amplitude.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/ampform_dpd/io/serialization/amplitude.py b/src/ampform_dpd/io/serialization/amplitude.py index 5c346ae4..e37a56e4 100644 --- a/src/ampform_dpd/io/serialization/amplitude.py +++ b/src/ampform_dpd/io/serialization/amplitude.py @@ -209,9 +209,9 @@ def _get_weight( value = value.real resonance_latex = to_latex(chain_definition["name"]) _, resonance_helicity = _get_resonance_helicity(chain_definition) - c = sp.IndexedBase(f"c^{{{resonance_latex}[{resonance_helicity}]}}") - λ1, λ2, λ3 = _get_final_state_helicities(chain_definition).values() - symbol = c[λ1, λ2, λ3] + helicities = _get_final_state_helicities(chain_definition).values() + subscript = ", ".join(sp.latex(λ) for λ in helicities) + symbol = sp.Symbol(f"c^{{{resonance_latex}[{resonance_helicity}]}}_{{{subscript}}}") return symbol, value From 3effb9b765622a97fd4624faf877cc0e2e48012b Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 20:00:35 +0200 Subject: [PATCH 29/42] ENH: show reason for `assert` failure --- docs/serialization.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index e549b1a9..dab3a7aa 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -770,7 +770,7 @@ " s_expr = s_expr.xreplace(MODEL.masses).doit()\n", " expr = perform_cached_doit(intensity_expr.xreplace({s: s_expr}))\n", " func = perform_cached_lambdify(expr, backend=\"jax\")\n", - " assert len(func.argument_order) == 2\n", + " assert len(func.argument_order) == 2, func.argument_order\n", " intensity_funcs[k] = func" ] }, From c13280dd3c0b1590b47872b2279da948f8aa0c2a Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 16:27:10 +0200 Subject: [PATCH 30/42] FIX: avoid overwriting chain amplitudes --- docs/serialization.ipynb | 3 ++- src/ampform_dpd/io/serialization/amplitude.py | 6 ++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index dab3a7aa..bffeb873 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -678,7 +678,8 @@ "metadata": { "tags": [ "hide-input", - "full-width" + "full-width", + "scroll-output" ] }, "outputs": [], diff --git a/src/ampform_dpd/io/serialization/amplitude.py b/src/ampform_dpd/io/serialization/amplitude.py index e37a56e4..8aacdb82 100644 --- a/src/ampform_dpd/io/serialization/amplitude.py +++ b/src/ampform_dpd/io/serialization/amplitude.py @@ -63,7 +63,7 @@ def formulate( # noqa: PLR0914 symbol: create_spin_range(states[i].spin) # type:ignore[index] for i, symbol in enumerate(helicity_symbols) } - amplitude_definitions = {} + amplitude_definitions = {} # type:ignore[var-annotated] angle_definitions = {} parameter_defaults = {} n_chains = len(get_decay_chains(model)) @@ -78,7 +78,9 @@ def formulate( # noqa: PLR0914 msg = f"Expected an expression, got {amp_expr!r}" raise TypeError(msg) helicity_substitutions = dict(zip(helicity_symbols, helicity_values)) - amplitude_definitions[amp_symbol] = amp_expr.subs(helicity_substitutions) + existing_amplitude = amplitude_definitions.get(amp_symbol, sp.Integer(0)) + existing_amplitude += amp_expr.subs(helicity_substitutions) + amplitude_definitions[amp_symbol] = existing_amplitude angle_definitions[θij] = θij_expr parameter_defaults.update(dict(parameters)) aligned_amp, zeta_defs = formulate_aligned_amplitude(model, *helicity_symbols) From ad6c3646f0f4087c6914f3d45736d956d4c75ba2 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 21:53:18 +0200 Subject: [PATCH 31/42] FIX: use `typing.List` in type alias --- src/ampform_dpd/io/serialization/format.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/io/serialization/format.py b/src/ampform_dpd/io/serialization/format.py index 31504f3a..d2c8f424 100644 --- a/src/ampform_dpd/io/serialization/format.py +++ b/src/ampform_dpd/io/serialization/format.py @@ -2,7 +2,7 @@ import difflib import sys -from typing import TYPE_CHECKING, Literal, TypedDict, Union +from typing import TYPE_CHECKING, List, Literal, TypedDict, Union from warnings import warn from ampform_dpd.decay import FinalStateID @@ -39,7 +39,7 @@ class DecayDescription(TypedDict): appendix: dict -Topology = list[Union[FinalStateID, "Topology"]] +Topology = List[Union[FinalStateID, "Topology"]] """Topology definition as a list of final state IDs.""" From e74cf7f7f49dd97a8542934466ac752f72f8d4b5 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 21:56:24 +0200 Subject: [PATCH 32/42] FIX: set correct builder signature --- src/ampform_dpd/dynamics/builder.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/ampform_dpd/dynamics/builder.py b/src/ampform_dpd/dynamics/builder.py index eeba3325..a6a1fa46 100644 --- a/src/ampform_dpd/dynamics/builder.py +++ b/src/ampform_dpd/dynamics/builder.py @@ -23,7 +23,7 @@ def formulate_breit_wigner_with_form_factor( decay: ThreeBodyDecayChain, -) -> tuple[sp.Expr, dict[sp.Symbol, float]]: +) -> tuple[sp.Expr, dict[sp.Symbol, complex | float]]: decay_node = decay.decay_node s = get_mandelstam_s(decay_node) parameter_defaults = {} @@ -41,7 +41,7 @@ def formulate_breit_wigner_with_form_factor( def _create_form_factor( s: sp.Symbol, isobar: IsobarNode -) -> tuple[sp.Expr, dict[sp.Symbol, float]]: +) -> tuple[sp.Expr, dict[sp.Symbol, complex | float]]: if isinstance(isobar.parent, State): inv_mass = sp.Symbol("m0", nonnegative=True) else: @@ -56,7 +56,7 @@ def _create_form_factor( angular_momentum=_get_angular_momentum(isobar), meson_radius=meson_radius, ) - parameter_defaults = { + parameter_defaults: dict[sp.Symbol, complex | float] = { meson_radius: 1, outgoing_state_mass1: to_particle(isobar.child1).mass, outgoing_state_mass2: to_particle(isobar.child2).mass, @@ -68,7 +68,7 @@ def _create_form_factor( def _create_breit_wigner( s: sp.Symbol, isobar: DecayNode -) -> tuple[sp.Expr, dict[sp.Symbol, float]]: +) -> tuple[sp.Expr, dict[sp.Symbol, complex | float]]: outgoing_state_mass1 = create_mass_symbol(isobar.child1) outgoing_state_mass2 = create_mass_symbol(isobar.child2) angular_momentum = _get_angular_momentum(isobar) @@ -85,7 +85,7 @@ def _create_breit_wigner( angular_momentum=angular_momentum, meson_radius=meson_radius, ) - parameter_defaults = { + parameter_defaults: dict[sp.Symbol, complex | float] = { res_mass: isobar.parent.mass, res_width: isobar.parent.width, meson_radius: 1, From 3bb16101ffdb5d066a657c523a83a9e2d0d55328 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 22:42:20 +0200 Subject: [PATCH 33/42] FIX: downgrade to AmpForm v0.15.1 --- .constraints/py3.10.txt | 2 +- .constraints/py3.11.txt | 2 +- .constraints/py3.12.txt | 2 +- .constraints/py3.8.txt | 2 +- .constraints/py3.9.txt | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.constraints/py3.10.txt b/.constraints/py3.10.txt index 45d91c23..d1a84608 100644 --- a/.constraints/py3.10.txt +++ b/.constraints/py3.10.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.3 +ampform==0.15.1 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.11.txt b/.constraints/py3.11.txt index 856bdcbb..bc21fcce 100644 --- a/.constraints/py3.11.txt +++ b/.constraints/py3.11.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.3 +ampform==0.15.1 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.12.txt b/.constraints/py3.12.txt index c5d58763..b7c7a7fc 100644 --- a/.constraints/py3.12.txt +++ b/.constraints/py3.12.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.3 +ampform==0.15.1 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.8.txt b/.constraints/py3.8.txt index dc62b24a..9b9d34c4 100644 --- a/.constraints/py3.8.txt +++ b/.constraints/py3.8.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.4 alabaster==0.7.13 -ampform==0.15.3 +ampform==0.15.1 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.9.txt b/.constraints/py3.9.txt index d41bd3d3..7e3b5e2e 100644 --- a/.constraints/py3.9.txt +++ b/.constraints/py3.9.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.3 +ampform==0.15.1 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 From d72ccedc24ba3e92dc1d932d67be47d354a06124 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Tue, 21 May 2024 22:57:45 +0200 Subject: [PATCH 34/42] FIX: use `Tuple` in type alias definition https://github.com/ComPWA/ampform-dpd/actions/runs/9181214979/job/25247435241?pr=132 --- src/ampform_dpd/io/serialization/format.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/io/serialization/format.py b/src/ampform_dpd/io/serialization/format.py index d2c8f424..a09f779b 100644 --- a/src/ampform_dpd/io/serialization/format.py +++ b/src/ampform_dpd/io/serialization/format.py @@ -2,7 +2,7 @@ import difflib import sys -from typing import TYPE_CHECKING, List, Literal, TypedDict, Union +from typing import TYPE_CHECKING, List, Literal, Tuple, TypedDict, Union from warnings import warn from ampform_dpd.decay import FinalStateID @@ -77,7 +77,7 @@ class LSVertex(Vertex): s: str -Node = tuple[Union[FinalStateID, "Node"], Union[FinalStateID, "Node"]] +Node = Tuple[Union[FinalStateID, "Node"], Union[FinalStateID, "Node"]] """Node definition within a `.Topology`.""" ParityFactor = Literal["+", "-", ""] From 2c29ea07f61e2d7b48d0c2394211e8fe05077b83 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Wed, 22 May 2024 15:28:09 +0200 Subject: [PATCH 35/42] ENH: upgrade to AmpForm v0.15.4 https://github.com/ComPWA/ampform/releases/0.15.4 --- .constraints/py3.10.txt | 2 +- .constraints/py3.11.txt | 2 +- .constraints/py3.12.txt | 2 +- .constraints/py3.8.txt | 2 +- .constraints/py3.9.txt | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.constraints/py3.10.txt b/.constraints/py3.10.txt index d1a84608..febeb41c 100644 --- a/.constraints/py3.10.txt +++ b/.constraints/py3.10.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.1 +ampform==0.15.4 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.11.txt b/.constraints/py3.11.txt index bc21fcce..c9ee0580 100644 --- a/.constraints/py3.11.txt +++ b/.constraints/py3.11.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.1 +ampform==0.15.4 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.12.txt b/.constraints/py3.12.txt index b7c7a7fc..fe33a291 100644 --- a/.constraints/py3.12.txt +++ b/.constraints/py3.12.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.1 +ampform==0.15.4 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.8.txt b/.constraints/py3.8.txt index 9b9d34c4..ad9b0108 100644 --- a/.constraints/py3.8.txt +++ b/.constraints/py3.8.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.4 alabaster==0.7.13 -ampform==0.15.1 +ampform==0.15.4 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 diff --git a/.constraints/py3.9.txt b/.constraints/py3.9.txt index 7e3b5e2e..ee215940 100644 --- a/.constraints/py3.9.txt +++ b/.constraints/py3.9.txt @@ -3,7 +3,7 @@ absl-py==2.1.0 accessible-pygments==0.0.5 alabaster==0.7.16 -ampform==0.15.1 +ampform==0.15.4 anyio==4.3.0 argon2-cffi==23.1.0 argon2-cffi-bindings==21.2.0 From 139505c3724b59ab9d2f515c478fdd079cad5f19 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 08:03:09 +0200 Subject: [PATCH 36/42] MAINT: remove empty `test_dynamics.py` --- tests/io_serialization/test_dynamics.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 tests/io_serialization/test_dynamics.py diff --git a/tests/io_serialization/test_dynamics.py b/tests/io_serialization/test_dynamics.py deleted file mode 100644 index e69de29b..00000000 From 0c397e19139c770e1f1525e7106cefb23ab8312c Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 08:40:22 +0200 Subject: [PATCH 37/42] DOC: hide caching warnings --- docs/serialization.ipynb | 2 ++ 1 file changed, 2 insertions(+) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index bffeb873..d8da66cf 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -35,6 +35,7 @@ "from __future__ import annotations\n", "\n", "import json\n", + "import logging\n", "\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", @@ -85,6 +86,7 @@ " get_function_definition,\n", ")\n", "\n", + "logging.getLogger(\"ampform.sympy\").setLevel(logging.ERROR)\n", "simplify_latex_rendering()" ] }, From 507066a1c2c0961092e2b7b53fd4f269253355d9 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 08:40:56 +0200 Subject: [PATCH 38/42] DX: limit number of rendered amplitudes in notebook --- docs/serialization.ipynb | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index d8da66cf..5a51a241 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -36,6 +36,7 @@ "\n", "import json\n", "import logging\n", + "import os\n", "\n", "import jax.numpy as jnp\n", "import matplotlib.pyplot as plt\n", @@ -686,7 +687,13 @@ }, "outputs": [], "source": [ - "Math(aslatex(MODEL.amplitudes, terms_per_line=1))" + "if \"EXECUTE_NB\" in os.environ:\n", + " selected_amplitudes = MODEL.amplitudes\n", + "else:\n", + " selected_amplitudes = {\n", + " k: v for i, (k, v) in enumerate(MODEL.amplitudes.items()) if i < 2\n", + " }\n", + "Math(aslatex(selected_amplitudes, terms_per_line=1))" ] }, { From 83b6983489e754c291bba561da0b65ded89debed Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 08:43:18 +0200 Subject: [PATCH 39/42] DOC: add preview warning for `ampform_dpd.io.serialization` --- docs/serialization.ipynb | 31 +++++++++++++++++++- src/ampform_dpd/io/serialization/__init__.py | 5 ++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 5a51a241..1898190d 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -9,6 +9,17 @@ "# Model serialization" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "This notebooks illustrates the use of the {mod}`ampform_dpd.io.serialization` module, which can be used to build amplitude models from the [amplitude-serialization](https://rub-ep1.github.io/amplitude-serialization) initiative.\n", + "\n", + ":::{warning}\n", + "The {mode}`ampform_dpd.io.serialization` module is a preview feature.\n", + ":::" + ] + }, { "cell_type": "markdown", "metadata": { @@ -679,6 +690,9 @@ "cell_type": "code", "execution_count": null, "metadata": { + "jupyter": { + "source_hidden": true + }, "tags": [ "hide-input", "full-width", @@ -793,6 +807,17 @@ "### Validation" ] }, + { + "cell_type": "markdown", + "metadata": { + "tags": [] + }, + "source": [ + ":::{error}\n", + "The following serves as a numerical check on whether the amplitude model has been deserialized correctly. For now, this is not the case, see [ComPWA/ampform-dpd#133](https://github.com/ComPWA/ampform-dpd/issues/133) for updates.\n", + ":::" + ] + }, { "cell_type": "code", "execution_count": null, @@ -950,8 +975,12 @@ "jupyter": { "source_hidden": true }, + "mystnb": { + "code_prompt_show": "Dalitz plot is not yet correct" + }, "tags": [ - "hide-input" + "hide-input", + "hide-output" ] }, "outputs": [], diff --git a/src/ampform_dpd/io/serialization/__init__.py b/src/ampform_dpd/io/serialization/__init__.py index e69de29b..297f29bb 100644 --- a/src/ampform_dpd/io/serialization/__init__.py +++ b/src/ampform_dpd/io/serialization/__init__.py @@ -0,0 +1,5 @@ +"""Amplitude model serialization + +See :doc:`/usage/serialization` for more information and +`https://rub-ep1.github.io/amplitude-serialization/`_ for more information. +""" From 8c2f05d5ba657bf615806bf5e90b9c25c1c7744a Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 09:15:21 +0200 Subject: [PATCH 40/42] MAINT: import `Lc2ppiK.json` from RUB-EP1/amplitude-serialization@678cdff https://raw.githubusercontent.com/RUB-EP1/amplitude-serialization/678cdff/Lc2ppiK.json --- .cspell.json | 1 + docs/Lc2ppiK.json | 4 ++-- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/.cspell.json b/.cspell.json index 219f58e3..15846a25 100644 --- a/.cspell.json +++ b/.cspell.json @@ -52,6 +52,7 @@ "colorbar", "commitlint", "concat", + "costheta", "difflib", "docnb", "elif", diff --git a/docs/Lc2ppiK.json b/docs/Lc2ppiK.json index 45d74168..01479b23 100644 --- a/docs/Lc2ppiK.json +++ b/docs/Lc2ppiK.json @@ -742,11 +742,11 @@ "variables": [ { "node": [3, 1], - "mass_angles": ["m_31", "phi_31", "cos_theta_31"] + "mass_phi_costheta": ["m_31", "phi_31", "cos_theta_31"] }, { "node": [[3, 1], 2], - "mass_angles": ["m_31_2", "phi_31_2", "cos_theta_31_2"] + "mass_phi_costheta": ["m_31_2", "phi_31_2", "cos_theta_31_2"] } ], "parameters": [] From 6f2d9878ad9296bc459b7506727bd07bfb634937 Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 09:18:05 +0200 Subject: [PATCH 41/42] FIX: update link to `serialization` notebook --- src/ampform_dpd/io/serialization/__init__.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/ampform_dpd/io/serialization/__init__.py b/src/ampform_dpd/io/serialization/__init__.py index 297f29bb..c9530d7c 100644 --- a/src/ampform_dpd/io/serialization/__init__.py +++ b/src/ampform_dpd/io/serialization/__init__.py @@ -1,5 +1,9 @@ """Amplitude model serialization -See :doc:`/usage/serialization` for more information and -`https://rub-ep1.github.io/amplitude-serialization/`_ for more information. +See :doc:`/serialization` for more information and +https://rub-ep1.github.io/amplitude-serialization for more information. + +.. warning:: + This module is in preview, see https://github.com/ComPWA/ampform-dpd/issues/133 for + updates. """ From 9d3c873a0755c28574c9c8f497c02aaebdccbe0b Mon Sep 17 00:00:00 2001 From: Remco de Boer <29308176+redeboer@users.noreply.github.com> Date: Sun, 26 May 2024 10:45:37 +0200 Subject: [PATCH 42/42] FIX: fix `mode` -> `mod` typo --- docs/serialization.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/serialization.ipynb b/docs/serialization.ipynb index 1898190d..74e9dd2b 100644 --- a/docs/serialization.ipynb +++ b/docs/serialization.ipynb @@ -16,7 +16,7 @@ "This notebooks illustrates the use of the {mod}`ampform_dpd.io.serialization` module, which can be used to build amplitude models from the [amplitude-serialization](https://rub-ep1.github.io/amplitude-serialization) initiative.\n", "\n", ":::{warning}\n", - "The {mode}`ampform_dpd.io.serialization` module is a preview feature.\n", + "The {mod}`ampform_dpd.io.serialization` module is a preview feature.\n", ":::" ] },