{ "cells": [ { "cell_type": "markdown", "metadata": { "tags": [] }, "source": [ "# Cross-check with LHCb data\n", "\n", "```{autolink-concat}\n", "```" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Import Python libraries" }, "tags": [ "hide-cell", "scroll-input" ] }, "outputs": [], "source": [ "from __future__ import annotations\n", "\n", "import json\n", "import logging\n", "import os\n", "from functools import lru_cache\n", "from textwrap import dedent\n", "\n", "import numpy as np\n", "import sympy as sp\n", "from IPython.display import Markdown, Math\n", "from tqdm.auto import tqdm\n", "\n", "from polarimetry.amplitude import AmplitudeModel, simplify_latex_rendering\n", "from polarimetry.data import create_data_transformer\n", "from polarimetry.io import (\n", " as_latex,\n", " display_latex,\n", " mute_jax_warnings,\n", " perform_cached_doit,\n", " perform_cached_lambdify,\n", ")\n", "from polarimetry.lhcb import (\n", " get_conversion_factor,\n", " get_conversion_factor_ls,\n", " load_model,\n", " load_model_builder,\n", " parameter_key_to_symbol,\n", ")\n", "from polarimetry.lhcb.particle import load_particles\n", "\n", "\n", "@lru_cache(maxsize=None)\n", "def load_model_cached(model_id: int | str) -> AmplitudeModel:\n", " return load_model(MODEL_FILE, PARTICLES, model_id)\n", "\n", "\n", "mute_jax_warnings()\n", "simplify_latex_rendering()\n", "NO_TQDM = \"EXECUTE_NB\" in os.environ\n", "if NO_TQDM:\n", " logging.getLogger().setLevel(logging.ERROR)\n", "\n", "MODEL_FILE = \"../data/model-definitions.yaml\"\n", "PARTICLES = load_particles(\"../data/particle-definitions.yaml\")\n", "DEFAULT_MODEL = load_model_cached(model_id=0)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [] }, "outputs": [], "source": [ "with open(\"../data/crosscheck.json\") as stream:\n", " crosscheck_data = json.load(stream)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Lineshape comparison\n", "\n", "We compute a few lineshapes for the following point in phase space and compare it with the values from {cite}`LHCb-PAPER-2022-002`:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Load phase space point" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "σ1, σ2, σ3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", "lineshape_vars = crosscheck_data[\"mainvars\"]\n", "lineshape_subs = {\n", " σ1: lineshape_vars[\"m2kpi\"],\n", " σ2: lineshape_vars[\"m2pk\"],\n", " **DEFAULT_MODEL.parameter_defaults,\n", "}\n", "lineshape_vars" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The lineshapes are computed for the following decay chains:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "mystnb": { "code_prompt_show": "Load selected decay chains" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "K892_chain = DEFAULT_MODEL.decay.find_chain(\"K(892)\")\n", "L1405_chain = DEFAULT_MODEL.decay.find_chain(\"L(1405)\")\n", "L1690_chain = DEFAULT_MODEL.decay.find_chain(\"L(1690)\")\n", "Math(as_latex([K892_chain, L1405_chain, L1690_chain]))" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "mystnb": { "code_prompt_show": "Values for LHCb-PAPER-2022-002" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "crosscheck_data[\"lineshapes\"]" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "mystnb": { "code_prompt_show": "Values as computed by this framework" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "def build_dynamics(c):\n", " return builder.dynamics_choices.get_builder(c)(c)[0].doit()\n", "\n", "\n", "builder = load_model_builder(MODEL_FILE, PARTICLES, model_id=0)\n", "K892_bw_val = build_dynamics(K892_chain).xreplace(lineshape_subs).n()\n", "L1405_bw_val = build_dynamics(L1405_chain).xreplace(lineshape_subs).n()\n", "L1690_bw_val = build_dynamics(L1690_chain).xreplace(lineshape_subs).n()\n", "display_latex([K892_bw_val, L1405_bw_val, L1690_bw_val])" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Assert that these values are equal" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "lineshape_decimals = 13\n", "np.testing.assert_array_almost_equal(\n", " np.array(list(map(complex, crosscheck_data[\"lineshapes\"].values()))),\n", " np.array(list(map(complex, [K892_bw_val, L1405_bw_val, L1690_bw_val]))),\n", " decimal=lineshape_decimals,\n", ")\n", "src = f\"\"\"\n", ":::{{tip}}\n", "These values are **equal up to {lineshape_decimals} decimals**.\n", ":::\n", "\"\"\"\n", "Markdown(src)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Amplitude comparison\n", "\n", "The amplitude for each decay chain and each outer state helicity combination are evaluated on the following point in phase space:" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Load phase space point as in DPD coordinates" }, "tags": [ "hide-input" ] }, "outputs": [], "source": [ "amplitude_vars = dict(crosscheck_data[\"chainvars\"])\n", "transformer = create_data_transformer(DEFAULT_MODEL)\n", "input_data = {\n", " str(σ1): amplitude_vars[\"m2kpi\"],\n", " str(σ2): amplitude_vars[\"m2pk\"],\n", " str(σ3): amplitude_vars[\"m2ppi\"],\n", "}\n", "input_data = {k: float(v) for k, v in transformer(input_data).items()}\n", "display_latex({sp.Symbol(k): v for k, v in input_data.items()})" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Code for creating functions for each sub-amplitude" }, "tags": [ "scroll-input", "hide-input" ] }, "outputs": [], "source": [ "@lru_cache(maxsize=None)\n", "def create_amplitude_functions(\n", " model_id: int | str,\n", ") -> dict[tuple[sp.Rational, sp.Rational], sp.Expr]:\n", " model = load_model(MODEL_FILE, PARTICLES, model_id)\n", " production_couplings = get_production_couplings(model_id)\n", " fixed_parameters = {\n", " s: v\n", " for s, v in model.parameter_defaults.items()\n", " if s not in production_couplings\n", " }\n", " exprs = formulate_amplitude_expressions(model_id)\n", " return {\n", " k: perform_cached_lambdify(\n", " expr.xreplace(fixed_parameters),\n", " parameters=production_couplings,\n", " backend=\"numpy\",\n", " )\n", " for k, expr in tqdm(exprs.items(), desc=\"Performing doit\", disable=NO_TQDM)\n", " }\n", "\n", "\n", "@lru_cache(maxsize=None)\n", "def formulate_amplitude_expressions(\n", " model_id: int | str,\n", ") -> dict[tuple[sp.Rational, sp.Rational], sp.Expr]:\n", " builder = load_model_builder(MODEL_FILE, PARTICLES, model_id)\n", " half = sp.Rational(1, 2)\n", " exprs = {\n", " (λ_Λc, λ_p): builder.formulate_aligned_amplitude(λ_Λc, λ_p, 0, 0)[0]\n", " for λ_Λc in [-half, +half]\n", " for λ_p in [-half, +half]\n", " }\n", " model = load_model(MODEL_FILE, PARTICLES, model_id)\n", " return {\n", " k: perform_cached_doit(expr.doit().xreplace(model.amplitudes))\n", " for k, expr in tqdm(exprs.items(), desc=\"Lambdifying\", disable=NO_TQDM)\n", " }\n", "\n", "\n", "@lru_cache(maxsize=None)\n", "def get_production_couplings(model_id: int | str) -> dict[sp.Indexed, complex]:\n", " model = load_model(MODEL_FILE, PARTICLES, model_id)\n", " return {\n", " symbol: value\n", " for symbol, value in model.parameter_defaults.items()\n", " if isinstance(symbol, sp.Indexed)\n", " if \"production\" in str(symbol)\n", " }" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "jupyter": { "source_hidden": true }, "mystnb": { "code_prompt_show": "Code for creating a comparison table" }, "tags": [ "hide-input", "scroll-input" ] }, "outputs": [], "source": [ "def plusminus_to_helicity(plusminus: str) -> sp.Rational:\n", " half = sp.Rational(1, 2)\n", " if plusminus == \"+\":\n", " return +half\n", " if plusminus == \"-\":\n", " return -half\n", " raise NotImplementedError(plusminus)\n", "\n", "\n", "def create_comparison_table(\n", " model_id: int | str, decimals: int | None = None\n", ") -> Markdown:\n", " min_ls = not is_ls_model(model_id)\n", " amplitude_funcs = create_amplitude_functions(model_id)\n", " real_amp_crosscheck = {\n", " k: v\n", " for k, v in get_amplitude_crosscheck_data(model_id).items()\n", " if k.startswith(\"Ar\")\n", " }\n", " production_couplings = get_production_couplings(model_id)\n", " couplings_to_zero = {str(symbol): 0 for symbol in production_couplings}\n", "\n", " src = \"\"\n", " if decimals is not None:\n", " src += dedent(f\"\"\"\n", " :::{{tip}}\n", " Computed amplitudes are equal to LHCb amplitudes up to **{decimals} decimals**.\n", " :::\n", " \"\"\")\n", " src += dedent(\"\"\"\n", " | | Computed | Expected | Difference |\n", " | ---:| --------:| --------:| ----------:|\n", " \"\"\")\n", " for amp_identifier, entry in real_amp_crosscheck.items():\n", " coupling = parameter_key_to_symbol(\n", " amp_identifier.replace(\"Ar\", \"A\"),\n", " min_ls,\n", " particle_definitions=PARTICLES,\n", " )\n", " src += f\"| **`{amp_identifier}`** | ${sp.latex(coupling)}$ |\\n\"\n", " for matrix_key, expected in entry.items():\n", " matrix_suffix = matrix_key[1:] # ++, +-, -+, --\n", " λ_Λc, λ_p = map(plusminus_to_helicity, matrix_suffix)\n", " func = amplitude_funcs[(λ_Λc, -λ_p)]\n", " func.update_parameters(couplings_to_zero)\n", " func.update_parameters({str(coupling): 1})\n", " computed = complex(func(input_data))\n", " computed *= determine_conversion_factor(coupling, λ_p, min_ls)\n", " expected = complex(expected)\n", " if abs(expected) != 0.0:\n", " diff = abs(computed - expected) / abs(expected)\n", " if diff < 1e-6:\n", " diff = f\"{diff:.2e}\"\n", " else:\n", " diff = f'{diff:.2e}'\n", " else:\n", " diff = \"\"\n", " src += (\n", " f\"| `{matrix_key}` | {computed:>.6f} | {expected:>.6f} | {diff} |\\n\"\n", " )\n", " if decimals is not None:\n", " np.testing.assert_array_almost_equal(\n", " computed,\n", " expected,\n", " decimal=decimals,\n", " err_msg=f\" {amp_identifier} {matrix_key}\",\n", " )\n", " return Markdown(src)\n", "\n", "\n", "def determine_conversion_factor(\n", " coupling: sp.Indexed, λ_p: sp.Rational, min_ls: bool\n", ") -> int:\n", " resonance_name = coupling.indices[0]\n", " resonance = PARTICLES[str(resonance_name)]\n", " if min_ls:\n", " factor = get_conversion_factor(resonance)\n", " else:\n", " _, L, S = coupling.indices\n", " factor = get_conversion_factor_ls(resonance, L, S)\n", " half = sp.Rational(1, 2)\n", " factor *= int((-1) ** (half + λ_p)) # # additional sign flip for amplitude\n", " return factor\n", "\n", "\n", "def is_ls_model(model_id: int | str) -> bool:\n", " if isinstance(model_id, int):\n", " return model_id == 17\n", " return \"LS couplings\" in model_id\n", "\n", "\n", "def get_amplitude_crosscheck_data(model_id: int | str) -> dict[str, complex]:\n", " if is_ls_model(model_id):\n", " return crosscheck_data[\"chains_LS\"]\n", " return crosscheck_data[\"chains\"]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Default model" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "hide-input", "hide-output" ] }, "outputs": [], "source": [ "create_comparison_table(model_id=0, decimals=13)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "### LS-model" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "tags": [ "hide-input", "hide-output" ] }, "outputs": [], "source": [ "create_comparison_table(\n", " \"Alternative amplitude model obtained using LS couplings\",\n", " decimals=13,\n", ")" ] } ], "metadata": { "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.8.15" } }, "nbformat": 4, "nbformat_minor": 4 }