diff --git a/.constraints/py3.10.txt b/.constraints/py3.10.txt index cb918e7f..543f9d93 100644 --- a/.constraints/py3.10.txt +++ b/.constraints/py3.10.txt @@ -41,7 +41,7 @@ gitpython==3.1.43 graphviz==0.20.3 greenlet==3.1.1 h11==0.14.0 -httpcore==1.0.5 +httpcore==1.0.6 httpx==0.27.2 identify==2.6.1 idna==3.10 @@ -50,7 +50,7 @@ importlib-metadata==8.5.0 iniconfig==2.0.0 ipykernel==6.29.5 ipympl==0.9.4 -ipython==8.27.0 +ipython==8.28.0 ipython-genutils==0.2.0 ipywidgets==8.1.5 isoduration==20.11.0 @@ -155,6 +155,7 @@ sphinx-design==0.6.1 sphinx-hep-pdgref==0.2.0 sphinx-pybtex-etal-style==0.0.2 sphinx-remove-toctrees==1.0.0.post1 +sphinx-reredirects==0.1.5 sphinx-thebe==0.3.1 sphinx-togglebutton==0.3.2 sphinxcontrib-applehelp==2.0.0 @@ -170,7 +171,7 @@ starlette==0.39.2 tabulate==0.9.0 terminado==0.18.1 tinycss2==1.3.0 -tomli==2.0.1 +tomli==2.0.2 tornado==6.4.1 tox==4.21.0 traitlets==5.14.3 diff --git a/.constraints/py3.11.txt b/.constraints/py3.11.txt index d0d2a3d2..365a2237 100644 --- a/.constraints/py3.11.txt +++ b/.constraints/py3.11.txt @@ -40,7 +40,7 @@ gitpython==3.1.43 graphviz==0.20.3 greenlet==3.1.1 h11==0.14.0 -httpcore==1.0.5 +httpcore==1.0.6 httpx==0.27.2 identify==2.6.1 idna==3.10 @@ -49,7 +49,7 @@ importlib-metadata==8.5.0 iniconfig==2.0.0 ipykernel==6.29.5 ipympl==0.9.4 -ipython==8.27.0 +ipython==8.28.0 ipython-genutils==0.2.0 ipywidgets==8.1.5 isoduration==20.11.0 @@ -154,6 +154,7 @@ sphinx-design==0.6.1 sphinx-hep-pdgref==0.2.0 sphinx-pybtex-etal-style==0.0.2 sphinx-remove-toctrees==1.0.0.post1 +sphinx-reredirects==0.1.5 sphinx-thebe==0.3.1 sphinx-togglebutton==0.3.2 sphinxcontrib-applehelp==2.0.0 diff --git a/.constraints/py3.12.txt b/.constraints/py3.12.txt index f9680938..43466887 100644 --- a/.constraints/py3.12.txt +++ b/.constraints/py3.12.txt @@ -40,7 +40,7 @@ gitpython==3.1.43 graphviz==0.20.3 greenlet==3.1.1 h11==0.14.0 -httpcore==1.0.5 +httpcore==1.0.6 httpx==0.27.2 identify==2.6.1 idna==3.10 @@ -49,7 +49,7 @@ importlib-metadata==8.5.0 iniconfig==2.0.0 ipykernel==6.29.5 ipympl==0.9.4 -ipython==8.27.0 +ipython==8.28.0 ipython-genutils==0.2.0 ipywidgets==8.1.5 isoduration==20.11.0 @@ -154,6 +154,7 @@ sphinx-design==0.6.1 sphinx-hep-pdgref==0.2.0 sphinx-pybtex-etal-style==0.0.2 sphinx-remove-toctrees==1.0.0.post1 +sphinx-reredirects==0.1.5 sphinx-thebe==0.3.1 sphinx-togglebutton==0.3.2 sphinxcontrib-applehelp==2.0.0 diff --git a/.constraints/py3.9.txt b/.constraints/py3.9.txt index fa55a40b..55a18671 100644 --- a/.constraints/py3.9.txt +++ b/.constraints/py3.9.txt @@ -41,7 +41,7 @@ gitpython==3.1.43 graphviz==0.20.3 greenlet==3.1.1 h11==0.14.0 -httpcore==1.0.5 +httpcore==1.0.6 httpx==0.27.2 identify==2.6.1 idna==3.10 @@ -156,6 +156,7 @@ sphinx-design==0.6.1 sphinx-hep-pdgref==0.2.0 sphinx-pybtex-etal-style==0.0.2 sphinx-remove-toctrees==1.0.0.post1 +sphinx-reredirects==0.1.5 sphinx-thebe==0.3.1 sphinx-togglebutton==0.3.2 sphinxcontrib-applehelp==2.0.0 @@ -171,7 +172,7 @@ starlette==0.39.2 tabulate==0.9.0 terminado==0.18.1 tinycss2==1.3.0 -tomli==2.0.1 +tomli==2.0.2 tornado==6.4.1 tox==4.21.0 traitlets==5.14.3 diff --git a/.cspell.json b/.cspell.json index 291b3d52..f9587cf0 100644 --- a/.cspell.json +++ b/.cspell.json @@ -22,12 +22,8 @@ ], "ignorePaths": [ "**/*.bib", - "**/*.rst_t", "**/.cspell.json", - "**/Manifest.toml", - "**/Project.toml", "*.ico", - "*.rst_t", "*.svg", ".constraints/*.txt", ".editorconfig", @@ -39,388 +35,86 @@ ".readthedocs.yml", ".vscode/*", "CITATION.cff", - "docs/_templates/*", "docs/adr/*/*", "docs/conf.py", "pyproject.toml", "tox.ini" ], "language": "en-US", - "overrides": [ - { - "allowCompoundWords": true, - "filename": "**/*.jl", - "languageId": "julia" - } - ], "words": [ - "aitchison", - "ampform", - "analyticity", - "argand", - "Atlassian", - "autograd", - "blatt", - "breit", - "chromodynamics", - "Clebsch", - "compwa", - "conda", + "AmpForm", + "Blatt", + "Breit", + "Chromodynamics", + "Colab", + "ComPWA", "Curvenote", - "Dalitz", - "deadsnakes", - "defaultdict", - "direnv", - "dummifies", - "dummify", - "eigenstates", - "eval", - "flatté", - "functools", - "GlueX", - "Gordan", - "Hankel", - "helicities", + "docstrings", "helicity", - "Hippel", - "itertools", - "JHEP", - "JPAC", - "Ketzer", - "kwargs", - "Källén", - "lambdification", - "lambdified", - "lambdifies", "lambdify", - "lambdifying", - "LHCb", - "lineshape", - "lineshapes", "Mathematica", - "MathML", - "matplotlib", - "Mikhasenko", - "miniconda", - "mkdir", - "multivalued", - "mypy", + "Miniconda", "Numba", - "numpy", + "NumPy", "parametrizations", - "pathlib", - "permutate", - "Plotly", "PyPA", - "pyproject", + "PyPI", "pytest", - "PYTHONHASHSEED", - "qrules", - "Quigg", + "QRules", "Reana", - "roadmap", - "Schwarz", "Scikit", - "scipy", - "sympify", - "sympy", - "tensorflow", - "tensorwaves", - "textwrap", - "toolkits", + "SciPy", + "SymPy", + "TensorWaves", + "Theano", "TPUs", - "traceback", "unbinned", - "unitarity", - "unitless", - "unnormalized", - "unphysical", - "vectorize", - "weisskopf", - "wigners", + "Weisskopf", "Zenodo" ], "ignoreWords": [ - "Basdevant", - "Colab", - "Danilkin", - "Deineka", "MAINT", - "Tiator", - "absl", - "adrs", - "allclose", - "appmode", - "arange", - "arccos", - "arctan", - "asarray", - "asdot", - "aslatex", - "astype", - "autolaunch", - "autonumbering", - "autoscale", - "autoupdate", - "axhline", - "axvline", - "azim", - "bbox", - "bdist", - "bgcolor", - "boldsymbol", - "byckling", - "cbar", - "cbff", - "celltoolbar", - "clim", - "cmap", - "cmath", - "cmax", - "cmin", - "codegen", + "PYTHONHASHSEED", "codemirror", - "colorbar", - "colorscale", - "colspan", - "combi", "commitlint", - "compat", - "componentwise", - "concat", - "coolwarm", - "csqrt", - "cstride", "cxxcode", - "dalitzplot", - "darkred", - "dasharray", - "dataclass", - "dataclasses", - "dataframe", - "deepcopy", - "displaystyle", - "docnb", - "docstrings", + "deadsnakes", + "direnv", "dotprint", - "einsum", - "elif", - "endswith", "envrc", - "epsabs", - "epsrel", - "eqnarray", - "errordef", - "evaluatable", "expertsystem", - "facecolor", - "facecolors", "fcode", - "figsize", - "filterwarnings", "fontcolor", - "fontsize", - "forall", - "framealpha", - "funcs", - "fvector", - "getitem", - "getsource", "graphviz", - "griddata", - "gridspec", - "halign", - "hasattr", - "heatmap", - "heli", - "hepstats", - "histtype", - "hotpink", - "hoverinfo", - "hspace", - "hypotests", - "iframe", - "imag", - "iminuit", - "infty", - "ioff", - "iplt", "ipykernel", - "ipympl", - "ipynb", - "ipyplot", "ipython", - "ipywidgets", - "isdigit", - "isinstance", - "isnan", "isort", - "isospin", - "isrealobj", - "jaxlib", - "joinpath", - "jpsi", "juliaup", "jupyterlab", "kernelspec", - "kmatrix", - "kutschke", - "lambdifier", - "lambdifygenerated", - "lightgray", - "linecap", - "linejoin", - "linestyle", - "linewidth", - "linkcheck", - "linspace", - "livereveal", - "lstrip", - "makedirs", - "marangotto", - "matexpr", "mathbb", - "mathbf", - "mathcal", - "mathop", - "mathrm", - "mathtt", + "mathml", "maxdepth", - "maxsize", - "meshgrid", - "migrad", - "mmikhasenko", - "mname", - "mplot3d", - "msigma", - "multiline", - "mystnb", - "nanmax", - "nanmean", - "nanstd", - "nansum", "nbconvert", "nbformat", "nbmake", - "nbody", "nbsp", - "ncols", - "ndarray", - "nonlocal", - "nonumber", "nopython", - "noqa", - "noreply", - "nrows", - "nsimplify", - "nstar", - "numpycode", - "operatorname", - "pandoc", - "pbar", - "pcolormesh", - "pdg's", - "phasespace", - "phsp", - "pkpi", - "pmatrix", - "ppnp", - "preorder", "prereleased", - "println", - "pvalues", - "py's", "pycode", - "pyfunc", "pygments", - "pylance", - "pypi", - "pyplot", - "pyright", - "pythoncode", - "q", - "quadpack", - "quadpy", - "rankdir", - "rcdefaults", + "pyproject", "redeboer", - "relim", - "repr", - "richman", - "rightarrow", "royalblue", - "rpartition", - "rstride", - "rstrip", "rtfd", - "rtol", - "rules's", - "savefig", - "scalex", - "scaley", - "scimath", - "sdist", "seealso", - "seterr", "setuptools", - "sharex", - "sharey", - "showcode", - "showlegend", - "showmarkdowntxt", - "showscale", - "showtags", - "simplefilter", "spflueger", "startswith", - "staticmethod", - "subslide", - "substack", - "suptitle", - "surfacecolor", - "symplot", - "tbody", - "thead", - "theano", - "threebody", - "ticklabels", - "ticktext", - "tickvals", "timeit", "toctree", "toprettyxml", - "tqdm", - "treewise", - "twinx", - "unevaluatable", "unsrt", - "venv", - "viridis", - "vmax", - "vmin", - "wspace", - "xanchor", - "xaxis", - "xdata", - "xlabel", - "xlim", - "xlink", - "xreplace", - "xtick", - "xticklabels", - "xticks", - "yanchor", - "yaxis", - "ydata", - "ylabel", - "ylim", - "yticklabels", - "yticks", - "zaxis", - "zfit", - "zlabel", - "zlim", - "zorder", - "zticks" + "venv" ] } diff --git a/.editorconfig b/.editorconfig index 88befe69..5c33e7e6 100644 --- a/.editorconfig +++ b/.editorconfig @@ -15,6 +15,3 @@ indent_size = 4 [LICENSE] indent_size = unset - -[setup.cfg] -indent_size = 4 diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index 54ae3bc8..678aa9bb 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -10,22 +10,6 @@ jobs: milestone: if: startsWith(github.ref, 'refs/tags') uses: ComPWA/actions/.github/workflows/close-milestone.yml@v1 - package-name: - uses: ComPWA/actions/.github/workflows/get-pypi-name.yml@v1 - pypi: - environment: - name: PyPI - url: https://pypi.org/p/${{ needs.package-name.outputs.name }} - if: startsWith(github.ref, 'refs/tags') - name: Publish to PyPI - needs: - - package-name - permissions: - id-token: write - runs-on: ubuntu-22.04 - steps: - - uses: ComPWA/actions/build-pypi-distribution@v1 - - uses: pypa/gh-action-pypi-publish@release/v1 push: if: startsWith(github.ref, 'refs/tags') && !github.event.release.prerelease secrets: inherit diff --git a/.github/workflows/notebooks.yml b/.github/workflows/notebooks.yml deleted file mode 100644 index 133537fb..00000000 --- a/.github/workflows/notebooks.yml +++ /dev/null @@ -1,44 +0,0 @@ -name: Run all notebooks - -on: - workflow_dispatch: - inputs: - notebook-selector: - description: Relative path to notebooks - required: false - type: string - -jobs: - pytest: - name: Test all notebooks - runs-on: ubuntu-20.04 - steps: - - uses: actions/checkout@v4 - - uses: actions/cache@v4 - with: - key: | - ${{ github.workflow }}-${{ github.job }}-${{ runner.os }}-${{ hashFiles('.constraints/py3.*.txt', 'setup.cfg') }} - path: | - .pytest_cache - ~/.cache/pip/ - - uses: actions/setup-python@v5 - with: - python-version: "3.9" - - name: Install dependencies - run: | - sudo apt-get -y install graphviz - pip install -c .constraints/py3.9.txt .[test] - - run: | - function run-nbmake() { - set +e - pytest --nbmake ${{ github.event.inputs.notebook-selector }} - error_code=$? - set -e - echo "Pytest returned error code $error_code" - case $error_code in - 0|5) return 0;; - *) return $error_code;; - esac - } - run-nbmake - # cspell:ignore esac diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 24f3d6b1..4fa2c6c6 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -3,7 +3,6 @@ ci: autoupdate_schedule: quarterly skip: - check-jsonschema - - pin-nb-requirements - pyright repos: @@ -20,8 +19,8 @@ repos: - --doc-apt-packages=graphviz - --dev-python-version=3.12 - --github-pages - - --imports-on-top - --no-prettierrc + - --no-pypi - --pin-requirements=bimonthly - --repo-name=compwa.github.io - --repo-title=ComPWA Organization @@ -106,7 +105,6 @@ repos: hooks: - id: toml-sort args: [--in-place] - exclude: (?x)^(.*/Manifest\.toml|.*/Project\.toml)$ - repo: https://github.com/python-jsonschema/check-jsonschema rev: 0.29.2 @@ -147,39 +145,6 @@ repos: docs/_static/favicon.ico )$ - - repo: https://github.com/nbQA-dev/nbQA - rev: 1.9.0 - hooks: - - id: nbqa-isort - args: [--float-to-top] - exclude: >- - (?x)^( - docs/adr/.*\.ipynb| - docs/report/001\.ipynb| - docs/report/002\.ipynb| - docs/report/003\.ipynb| - docs/report/008\.ipynb| - docs/report/013\.ipynb| - docs/report/018\.ipynb| - docs/report/020\.ipynb| - docs/report/026\.ipynb| - docs/report/027\.ipynb| - docs/report/028\.ipynb| - docs/symbolics\.ipynb - )$ - - - repo: local - hooks: - - id: pin-nb-requirements - name: Check whether notebook contains a pip install line - description: - Specify which packages to install specifically in order to run - this notebook. - entry: python3 .pre-commit/pin_nb_requirements.py - language: system - types: - - jupyter - - repo: https://github.com/ComPWA/mirrors-pyright rev: v1.1.379 hooks: diff --git a/.pre-commit/pin_nb_requirements.py b/.pre-commit/pin_nb_requirements.py deleted file mode 100644 index bd933756..00000000 --- a/.pre-commit/pin_nb_requirements.py +++ /dev/null @@ -1,244 +0,0 @@ -"""Enforce adding a pip install statement in the notebook. - -In the `compwa.github.io repo `_, notebooks -should specify which package versions should be used to run the notebook. This hook -checks whether a notebook has such install statements and whether they comply with the -expected formatting. -""" -# cspell:ignore oneline precommit - -from __future__ import annotations - -import argparse -import re -import sys -from functools import lru_cache -from textwrap import dedent -from typing import TYPE_CHECKING, Callable, TypeVar - -import attr -import nbformat -from nbformat import NotebookNode - -if TYPE_CHECKING: - from collections.abc import Sequence - -if sys.version_info >= (3, 10): - from typing import ParamSpec -else: - from typing_extensions import ParamSpec - -__EXPECTED_PIP_INSTALL_LINE = "%pip install -q" - - -def main(argv: Sequence[str] | None = None) -> int: - parser = argparse.ArgumentParser(__doc__) - parser.add_argument("filenames", nargs="*", help="Filenames to check.") - args = parser.parse_args(argv) - - errors: list[PrecommitError] = [] - for filename in args.filenames: - try: - check_pinned_requirements(filename) - except PrecommitError as exception: # noqa: PERF203 - errors.append(exception) - if errors: - for error in errors: - error_msg = "\n ".join(error.args) - print(error_msg) # noqa: T201 - return 1 - return 0 - - -def check_pinned_requirements(filename: str) -> None: - notebook = load_notebook(filename) - if not __has_python_kernel(notebook): - return - for cell_id, cell in enumerate(notebook["cells"]): - if cell["cell_type"] != "code": - continue - source = __to_oneline(cell["source"]) - pip_requirements = extract_pip_requirements(source) - if pip_requirements is None: - continue - executor = Executor() - executor(_check_pip_requirements, filename, pip_requirements) - executor(_format_pip_requirements, filename, source, notebook, cell_id) - executor(_update_metadata, filename, cell["metadata"], notebook) - executor.finalize() - return - msg = ( - f'Notebook "{filename}" does not contain a pip install cell of the form' - f" {__EXPECTED_PIP_INSTALL_LINE} some-package==0.1.0 package2==3.2" - ) - raise PrecommitError(msg) - - -def __has_python_kernel(notebook: dict) -> bool: - # cspell:ignore kernelspec - metadata = notebook.get("metadata", {}) - kernel_specification = metadata.get("kernelspec", {}) - kernel_language = kernel_specification.get("language", "") - return "python" in kernel_language - - -@lru_cache(maxsize=1) -def __to_oneline(source: str) -> str: - src_lines = source.split("\n") - return "".join(s.rstrip().rstrip("\\") for s in src_lines) - - -@lru_cache(maxsize=1) -def extract_pip_requirements(source: str) -> list[str] | None: - r"""Check if the source in a cell is a pip install statement. - - >>> extract_pip_requirements("Not a pip install statement") - >>> extract_pip_requirements("pip install") - [] - >>> extract_pip_requirements("pip3 install attrs") - ['attrs'] - >>> extract_pip_requirements("pip3 install -q attrs") - ['attrs'] - >>> extract_pip_requirements("pip3 install attrs &> /dev/null") - ['attrs'] - >>> extract_pip_requirements("%pip install attrs numpy==1.24.4 ") - ['attrs', 'numpy==1.24.4'] - >>> extract_pip_requirements("!python3 -mpip install sympy") - ['sympy'] - >>> extract_pip_requirements(''' - ... python3 -m pip install \ - ... attrs numpy \ - ... sympy \ - ... tensorflow - ... ''') - ['attrs', 'numpy', 'sympy', 'tensorflow'] - """ - # cspell:ignore mpip - matches = re.match( - r"[%\!]?\s*(python3?\s+-m\s*)?pip3?\s+install\s*(-q)?(.*?)(&?>\s*/dev/null)?$", - __to_oneline(source).strip(), - ) - if matches is None: - return None - packages = matches.group(3).split(" ") - packages = [p.strip() for p in packages] - return [p for p in packages if p] - - -def _check_pip_requirements(filename: str, requirements: list[str]) -> None: - if len(requirements) == 0: - msg = f'At least one dependency required in install cell of "{filename}"' - raise PrecommitError(msg) - for req in requirements: - req = req.strip() - if not req: - continue - if "git+" in req: - continue - unpinned_requirements = [] - for req in requirements: - if req.startswith("git+"): - continue - if any(equal_sign in req for equal_sign in ["==", "~="]): - continue - package = req.split("<")[0].split(">")[0].strip() - unpinned_requirements.append(package) - if unpinned_requirements: - msg = ( - f'Install cell in notebook "{filename}" contains requirements without' - "pinning (== or ~=):" - ) - for req in unpinned_requirements: - msg += f"\n - {req}" - msg_unformatted = f""" - Get the currently installed versions with: - - python3 -m pip freeze | grep -iE '{"|".join(sorted(unpinned_requirements))}' - """ - msg += dedent(msg_unformatted) - raise PrecommitError(msg) - - -def _format_pip_requirements( - filename: str, install_statement: str, notebook: NotebookNode, cell_id: int -) -> None: - requirements = extract_pip_requirements(install_statement) - if requirements is None: - return - git_requirements = {r for r in requirements if r.startswith("git+")} - pip_requirements = set(requirements) - git_requirements - pip_requirements = {r.lower().replace("_", "-") for r in pip_requirements} - sorted_requirements = sorted(pip_requirements) + sorted(git_requirements) - expected = f"{__EXPECTED_PIP_INSTALL_LINE} {' '.join(sorted_requirements)}" - if install_statement != expected: - notebook["cells"][cell_id]["source"] = expected - nbformat.write(notebook, filename) - msg = f'Ordered and formatted pip install cell in "{filename}"' - raise PrecommitError(msg) - - -def _update_metadata(filename: str, metadata: dict, notebook: NotebookNode) -> None: - updated_metadata = False - jupyter_metadata = metadata.get("jupyter") - if jupyter_metadata is not None and jupyter_metadata.get("source_hidden"): - if len(jupyter_metadata) == 1: - metadata.pop("jupyter") - else: - jupyter_metadata.pop("source_hidden") - updated_metadata = True - tags = set(metadata.get("tags", [])) - expected_tags = {"remove-cell"} - if expected_tags != tags: - metadata["tags"] = sorted(expected_tags) - updated_metadata = True - if updated_metadata: - nbformat.write(notebook, filename) - msg = f'Updated metadata of pip install cell in notebook "{filename}"' - raise PrecommitError(msg) - - -def load_notebook(path: str) -> NotebookNode: - return nbformat.read(path, as_version=nbformat.NO_CONVERT) - - -T = TypeVar("T") -P = ParamSpec("P") - - -@attr.s(on_setattr=attr.setters.frozen) -class Executor: - # https://github.com/ComPWA/policy/blob/359331f/src/compwa_policy/utilities/executor.py - error_messages: list[str] = attr.ib(factory=list, init=False) - - def __call__( - self, function: Callable[P, T], *args: P.args, **kwargs: P.kwargs - ) -> T | None: - """Execute a function and collect any `.PrecommitError` exceptions.""" - try: - result = function(*args, **kwargs) - except PrecommitError as exception: - error_message = str("\n".join(exception.args)) - self.error_messages.append(error_message) - return None - else: - return result - - def finalize(self, exception: bool = True) -> int: - error_msg = self.merge_messages() - if error_msg: - if exception: - raise PrecommitError(error_msg) - print(error_msg) # noqa: T201 - return 1 - return 0 - - def merge_messages(self) -> str: - stripped_messages = (s.strip() for s in self.error_messages) - return "\n--------------------\n".join(stripped_messages) - - -class PrecommitError(RuntimeError): ... - - -if __name__ == "__main__": - raise SystemExit(main()) diff --git a/.readthedocs.yml b/.readthedocs.yml index 1429c63f..b6844a63 100644 --- a/.readthedocs.yml +++ b/.readthedocs.yml @@ -15,8 +15,6 @@ build: tools: python: "3.12" jobs: - pre_install: - - ./docs/install-julia-on-rtd.sh post_install: - python -m pip install 'uv>=0.2.0' - python -m uv pip install -c .constraints/py3.12.txt -e .[doc] diff --git a/.taplo.toml b/.taplo.toml index 16ab049a..0db1df7c 100644 --- a/.taplo.toml +++ b/.taplo.toml @@ -1,8 +1,3 @@ -exclude = [ - "**/Manifest.toml", - "**/Project.toml", -] - [formatting] align_comments = false align_entries = false diff --git a/.vscode/extensions.json b/.vscode/extensions.json index c1ac5a71..bbd4e772 100644 --- a/.vscode/extensions.json +++ b/.vscode/extensions.json @@ -8,7 +8,6 @@ "executablebookproject.myst-highlight", "github.vscode-github-actions", "github.vscode-pull-request-github", - "julialang.language-julia", "mhutchie.git-graph", "ms-python.python", "ms-python.vscode-pylance", @@ -16,6 +15,7 @@ "ms-vscode.live-server", "ms-vsliveshare.vsliveshare", "redhat.vscode-yaml", + "simonsiefke.svg-preview", "soulcode.vscode-unwanted-extensions", "stkb.rewrap", "streetsidesoftware.code-spell-checker", @@ -26,6 +26,7 @@ "bungcip.better-toml", "davidanson.vscode-markdownlint", "garaioag.garaio-vscode-unwanted-recommendations", + "julialang.language-julia", "ms-python.black-formatter", "ms-python.flake8", "ms-python.isort", diff --git a/.vscode/launch.json b/.vscode/launch.json deleted file mode 100644 index 84ad26a0..00000000 --- a/.vscode/launch.json +++ /dev/null @@ -1,20 +0,0 @@ -{ - "version": "0.2.0", - "configurations": [ - { - "name": "Debug current file", - "type": "python", - "request": "launch", - "program": "${file}", - "console": "integratedTerminal" - }, - { - "name": "Deep-debug current file", - "type": "python", - "request": "launch", - "program": "${file}", - "console": "integratedTerminal", - "justMyCode": false - } - ] -} diff --git a/.vscode/settings.json b/.vscode/settings.json index 9b178884..e717005e 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -57,7 +57,6 @@ "notebook.formatOnSave.enabled": true, "notebook.gotoSymbols.showAllSymbols": true, "python.analysis.autoImportCompletions": false, - "python.analysis.typeCheckingMode": "strict", "python.terminal.activateEnvironment": false, "python.testing.pytestEnabled": false, "redhat.telemetry.enabled": false, diff --git a/docs/.gitignore b/docs/.gitignore deleted file mode 100644 index 2047cf6a..00000000 --- a/docs/.gitignore +++ /dev/null @@ -1 +0,0 @@ -report-inventory.md diff --git a/docs/InstallIJulia.jl b/docs/InstallIJulia.jl deleted file mode 100644 index 5b1b88e2..00000000 --- a/docs/InstallIJulia.jl +++ /dev/null @@ -1,5 +0,0 @@ -import Pkg -Pkg.add("IJulia") - -import IJulia -IJulia.installkernel("julia-compwa.github.io") diff --git a/docs/_list_technical_reports.py b/docs/_list_technical_reports.py deleted file mode 100644 index cf45d57f..00000000 --- a/docs/_list_technical_reports.py +++ /dev/null @@ -1,128 +0,0 @@ -from __future__ import annotations - -import os -import re -from os.path import dirname -from textwrap import dedent -from typing import TYPE_CHECKING - -import nbformat - -if TYPE_CHECKING: - from nbformat import NotebookNode - - -def main() -> int: - table = _create_tr_table() - this_dir = dirname(__file__) - output_path = f"{this_dir}/report-inventory.md" - with open(output_path, "w") as f: - f.write(table) - return 0 - - -def _create_tr_table() -> str: - notebook_paths = _get_technical_report_paths() - src = dedent( - """ - | | TR | Title | Details | Tags | Status | - |:--:|:--:|:------|:--------|:-----|:-------| - """ - ).strip() - for notebook in notebook_paths: - card_info = _get_card_info(notebook) - tr = card_info["tr"] - title = card_info["title"] - details = re.sub( - r"\[([^\]]+)\]\((\./)?(\d\d\d)\.ipynb\)", - r"[\1](report/\3.ipynb)", - card_info.get("details", ""), - ) - details = re.sub(r"(
)?", "", details) - tags = " ".join(_to_badge(tag) for tag in sorted(card_info["tags"])) - status = re.sub( - r"\[([^\]]+)\-([^\]]+)\]\(([^\)]+)\)", - r"[\1‑\2](\3)", - card_info.get("footer", "\n").splitlines()[0], - ) - src += ( - f"\n| | [TR‑{tr}](report/{tr}.ipynb) | {title} | {details} | {tags} |" - f" {status} |" - ) - return src - - -def _to_badge(tag: str) -> str: - return f"{{bdg-info-line}}`{tag}`" - - -def _get_technical_report_paths() -> list[str]: - this_dir = dirname(__file__) - report_dir = f"{this_dir}/report" - return [ - f"{report_dir}/{file}" - for file in sorted(os.listdir(report_dir)) - if file.endswith(".ipynb") - ] - - -def _get_card_info(path: str) -> dict[str, str]: - notebook = _open_notebook(path) - for cell in notebook["cells"]: - if cell["cell_type"] != "markdown": - continue - src: list[str] = cell["source"].splitlines() - src = [s for s in src if s.strip() if not s.strip().startswith("" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q black==21.5b2 jax==0.4.28 jaxlib==0.4.28 numpy==1.23 sympy==1.8" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import inspect\n", - "\n", - "import jax\n", - "import jax.numpy as jnp\n", - "import numpy as np\n", - "import sympy as sp\n", - "from black import FileMode, format_str\n", - "from IPython.display import display" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Negative input values" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When using {mod}`numpy` as back-end, {mod}`sympy` lambdifies a {func}`~sympy.functions.elementary.miscellaneous.sqrt` to a {obj}`numpy.sqrt`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sqrt{x}$" - ], - "text/plain": [ - "sqrt(x)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "x = sp.Symbol(\"x\")\n", - "sqrt_expr = sp.sqrt(x)\n", - "sqrt_expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return (sqrt(x))\n", - "\n" - ] - } - ], - "source": [ - "np_sqrt = sp.lambdify(x, sqrt_expr, \"numpy\")\n", - "source = inspect.getsource(np_sqrt)\n", - "print(source)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As expected, if input values for the {obj}`numpy.sqrt` are negative, {mod}`numpy` raises a {class}`RuntimeWarning` and returns `NaN`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - ":2: RuntimeWarning: invalid value encountered in sqrt\n", - " return (sqrt(x))\n" - ] - }, - { - "data": { - "text/plain": [ - "array([ nan, nan, 0. , 0.70710678, 1. ])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sample = np.linspace(-1, 1, 5)\n", - "np_sqrt(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "If we want {mod}`numpy` to return imaginary numbers for negative input values, one can use {class}`complex` input data instead (e.g. {doc}`numpy.complex64 `). Negative values are then treated as lying just above the real axis, so that their square root is a positive imaginary number:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "array([0. +1.j , 0. +0.70710677j,\n", - " 0. +0.j , 0.70710677+0.j ,\n", - " 1. +0.j ], dtype=complex64)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "complex_sample = sample.astype(np.complex64)\n", - "np_sqrt(complex_sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A {func}`sympy.sqrt ` lambdified to [JAX](https://jax.rtfd.io) exhibits the same behavior:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return (sqrt(x))\n", - "\n" - ] - } - ], - "source": [ - "jax_sqrt = jax.jit(sp.lambdify(x, sqrt_expr, jnp))\n", - "source = inspect.getsource(jax_sqrt)\n", - "print(source)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n" - ] - }, - { - "data": { - "text/plain": [ - "DeviceArray([ nan, nan, 0. , 0.70710677, 1. ], dtype=float32)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jax_sqrt(sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "DeviceArray([-4.3711388e-08+1.j , -3.0908620e-08+0.70710677j,\n", - " 0.0000000e+00+0.j , 7.0710677e-01+0.j ,\n", - " 1.0000000e+00+0.j ], dtype=complex64)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jax_sqrt(complex_sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**There is a problem with this approach though**: once input data is complex, _all_ square roots in a larger expression (some amplitude model) compute imaginary solutions for negative values, while this is not always the desired behavior.\n", - "\n", - "Take for instance the two square roots appearing in {class}`~ampform.dynamics.phasespace.PhaseSpaceFactor` --- does the $\\sqrt{s}$ also have to be evaluatable for negative $s$?" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Complex square root" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Numpy also offers a special function that evaluates negative values even if the input values are real: {func}`numpy.emath.sqrt`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "1j" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np.emath.sqrt(-1)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Unfortunately, the {mod}`jax.numpy` API does not interface to {mod}`numpy.emath`. It is possible to decorate {func}`numpy.emath.sqrt` be decorated with {func}`jax.jit`, but that **only works with static, hashable arguments**:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception", - "hide-output", - "scroll-output" - ] - }, - "outputs": [ - { - "ename": "TracerArrayConversionError", - "evalue": "The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mUnfilteredStackTrace\u001b[0m Traceback (most recent call last)", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel_launcher.py:17\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch_new_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/traitlets/config/application.py:976\u001b[0m, in \u001b[0;36mApplication.launch_instance\u001b[0;34m(cls, argv, **kwargs)\u001b[0m\n\u001b[1;32m 975\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m--> 976\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelapp.py:712\u001b[0m, in \u001b[0;36mIPKernelApp.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 711\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 712\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/tornado/platform/asyncio.py:199\u001b[0m, in \u001b[0;36mBaseAsyncIOLoop.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 198\u001b[0m asyncio\u001b[38;5;241m.\u001b[39mset_event_loop(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop)\n\u001b[0;32m--> 199\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43masyncio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_forever\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 200\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mBaseEventLoop.run_forever\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_once\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 571\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36mBaseEventLoop._run_once\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 1858\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m \u001b[43mhandle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/events.py:81\u001b[0m, in \u001b[0;36mHandle._run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_callback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:510\u001b[0m, in \u001b[0;36mKernel.dispatch_queue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 509\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 510\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m 511\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:499\u001b[0m, in \u001b[0;36mKernel.process_one\u001b[0;34m(self, wait)\u001b[0m\n\u001b[1;32m 498\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 499\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:406\u001b[0m, in \u001b[0;36mKernel.dispatch_shell\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 406\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m 407\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:730\u001b[0m, in \u001b[0;36mKernel.execute_request\u001b[0;34m(self, stream, ident, parent)\u001b[0m\n\u001b[1;32m 729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 730\u001b[0m reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m 732\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/ipkernel.py:383\u001b[0m, in \u001b[0;36mIPythonKernel.do_execute\u001b[0;34m(self, code, silent, store_history, user_expressions, allow_stdin, cell_id)\u001b[0m\n\u001b[1;32m 382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 383\u001b[0m res \u001b[38;5;241m=\u001b[39m \u001b[43mshell\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 384\u001b[0m \u001b[43m \u001b[49m\u001b[43mcode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 385\u001b[0m \u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 386\u001b[0m \u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 387\u001b[0m \u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcell_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 388\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 389\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/zmqshell.py:528\u001b[0m, in \u001b[0;36mZMQInteractiveShell.run_cell\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 527\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2881\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m 2880\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2881\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 2882\u001b[0m \u001b[43m \u001b[49m\u001b[43mraw_cell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshell_futures\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\n\u001b[1;32m 2883\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2884\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2936\u001b[0m, in \u001b[0;36mInteractiveShell._run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m 2935\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2936\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoro\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2937\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m(coro)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3135\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_async\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple, cell_id)\u001b[0m\n\u001b[1;32m 3133\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3135\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m 3136\u001b[0m interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m 3138\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3338\u001b[0m, in \u001b[0;36mInteractiveShell.run_ast_nodes\u001b[0;34m(self, nodelist, cell_name, interactivity, compiler, result)\u001b[0m\n\u001b[1;32m 3337\u001b[0m asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m 3339\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n", - " \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n", - "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 1\u001b[0m jax_csqrt_error \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np\u001b[38;5;241m.\u001b[39memath\u001b[38;5;241m.\u001b[39msqrt, backend\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_csqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/_src/traceback_util.py:143\u001b[0m, in \u001b[0;36mapi_boundary..reraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/_src/api.py:426\u001b[0m, in \u001b[0;36m_cpp_jit..cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 425\u001b[0m flat_fun, out_tree \u001b[38;5;241m=\u001b[39m flatten_fun(f, in_tree)\n\u001b[0;32m--> 426\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxla\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxla_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 427\u001b[0m \u001b[43m \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 428\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 429\u001b[0m \u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 430\u001b[0m \u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 431\u001b[0m \u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mflat_fun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 432\u001b[0m \u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdonated_invars\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 433\u001b[0m out_pytree_def \u001b[38;5;241m=\u001b[39m out_tree()\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1565\u001b[0m, in \u001b[0;36mCallPrimitive.bind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[0;32m-> 1565\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1556\u001b[0m, in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1555\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m maybe_new_sublevel(top_trace):\n\u001b[0;32m-> 1556\u001b[0m outs \u001b[38;5;241m=\u001b[39m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, apply_todos(env_trace_todo(), outs))\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1568\u001b[0m, in \u001b[0;36mCallPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m 1567\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m-> 1568\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:609\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m 608\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 609\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/xla.py:578\u001b[0m, in \u001b[0;36m_xla_call_impl\u001b[0;34m(fun, device, backend, name, donated_invars, *args)\u001b[0m\n\u001b[1;32m 577\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_xla_call_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, device, backend, name, donated_invars):\n\u001b[0;32m--> 578\u001b[0m compiled_fun \u001b[38;5;241m=\u001b[39m \u001b[43m_xla_callable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 579\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43munsafe_map\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg_spec\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 580\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/linear_util.py:262\u001b[0m, in \u001b[0;36mcache..memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 262\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 263\u001b[0m cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/xla.py:652\u001b[0m, in \u001b[0;36m_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m 651\u001b[0m abstract_args, _ \u001b[38;5;241m=\u001b[39m unzip2(arg_specs)\n\u001b[0;32m--> 652\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mabstract_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mjit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(c, core\u001b[38;5;241m.\u001b[39mTracer) \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m consts):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/partial_eval.py:1209\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, transform_name)\u001b[0m\n\u001b[1;32m 1208\u001b[0m main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m () \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m-> 1209\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1210\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/partial_eval.py:1188\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals)\u001b[0m\n\u001b[1;32m 1187\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[0;32m-> 1188\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1189\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/linear_util.py:166\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 167\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m 168\u001b[0m \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m 169\u001b[0m \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks only\u001b[39;00m\n\u001b[1;32m 171\u001b[0m \u001b[38;5;66;03m# after this exception is handled, which can corrupt the global state.\u001b[39;00m\n", - "File \u001b[0;32m<__array_function__ internals>:180\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(*args, **kwargs)\u001b[0m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/numpy/lib/scimath.py:247\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 200\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 201\u001b[0m \u001b[38;5;124;03mCompute the square root of x.\u001b[39;00m\n\u001b[1;32m 202\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 245\u001b[0m \u001b[38;5;124;03m-2j\u001b[39;00m\n\u001b[1;32m 246\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 247\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43m_fix_real_lt_zero\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 248\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m nx\u001b[38;5;241m.\u001b[39msqrt(x)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/numpy/lib/scimath.py:134\u001b[0m, in \u001b[0;36m_fix_real_lt_zero\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 113\u001b[0m \u001b[38;5;124;03m\"\"\"Convert `x` to complex if it has real, negative components.\u001b[39;00m\n\u001b[1;32m 114\u001b[0m \n\u001b[1;32m 115\u001b[0m \u001b[38;5;124;03mOtherwise, output is just the array version of the input (via asarray).\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 132\u001b[0m \n\u001b[1;32m 133\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 134\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43masarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 135\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(isreal(x) \u001b[38;5;241m&\u001b[39m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m)):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:472\u001b[0m, in \u001b[0;36mTracer.__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__array__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw):\n\u001b[0;32m--> 472\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m TracerArrayConversionError(\u001b[38;5;28mself\u001b[39m)\n", - "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------", - "\nThe above exception was the direct cause of the following exception:\n", - "\u001b[0;31mTracerArrayConversionError\u001b[0m Traceback (most recent call last)", - "Input \u001b[0;32mIn [13]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 1\u001b[0m jax_csqrt_error \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np\u001b[38;5;241m.\u001b[39memath\u001b[38;5;241m.\u001b[39msqrt, backend\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcpu\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_csqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m<__array_function__ internals>:180\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(*args, **kwargs)\u001b[0m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/numpy/lib/scimath.py:247\u001b[0m, in \u001b[0;36msqrt\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 198\u001b[0m \u001b[38;5;129m@array_function_dispatch\u001b[39m(_unary_dispatcher)\n\u001b[1;32m 199\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21msqrt\u001b[39m(x):\n\u001b[1;32m 200\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m 201\u001b[0m \u001b[38;5;124;03m Compute the square root of x.\u001b[39;00m\n\u001b[1;32m 202\u001b[0m \n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 245\u001b[0m \u001b[38;5;124;03m -2j\u001b[39;00m\n\u001b[1;32m 246\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 247\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43m_fix_real_lt_zero\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 248\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m nx\u001b[38;5;241m.\u001b[39msqrt(x)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/numpy/lib/scimath.py:134\u001b[0m, in \u001b[0;36m_fix_real_lt_zero\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 112\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_fix_real_lt_zero\u001b[39m(x):\n\u001b[1;32m 113\u001b[0m \u001b[38;5;124;03m\"\"\"Convert `x` to complex if it has real, negative components.\u001b[39;00m\n\u001b[1;32m 114\u001b[0m \n\u001b[1;32m 115\u001b[0m \u001b[38;5;124;03m Otherwise, output is just the array version of the input (via asarray).\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 132\u001b[0m \n\u001b[1;32m 133\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 134\u001b[0m x \u001b[38;5;241m=\u001b[39m \u001b[43masarray\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 135\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(isreal(x) \u001b[38;5;241m&\u001b[39m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m)):\n\u001b[1;32m 136\u001b[0m x \u001b[38;5;241m=\u001b[39m _tocomplex(x)\n", - "\u001b[0;31mTracerArrayConversionError\u001b[0m: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)" - ] - } - ], - "source": [ - "jax_csqrt_error = jax.jit(np.emath.sqrt, backend=\"cpu\")\n", - "jax_csqrt_error(-1)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "DeviceArray(0.+1.j, dtype=complex64)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jax_csqrt = jax.jit(np.emath.sqrt, backend=\"cpu\", static_argnums=0)\n", - "jax_csqrt(-1)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception" - ] - }, - "outputs": [ - { - "ename": "ValueError", - "evalue": "Non-hashable static arguments are not supported. An error occured while trying to hash an object of type , [-1. -0.5 0. 0.5 1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mValueError\u001b[0m Traceback (most recent call last)", - "Input \u001b[0;32mIn [15]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax_csqrt\u001b[49m\u001b[43m(\u001b[49m\u001b[43msample\u001b[49m\u001b[43m)\u001b[49m\n", - "\u001b[0;31mValueError\u001b[0m: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type , [-1. -0.5 0. 0.5 1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n" - ] - } - ], - "source": [ - "jax_csqrt(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Conditional square root" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To be able to control which square roots in the complete expression should be evaluatable for negative values, one could use {class}`~sympy.functions.elementary.piecewise.Piecewise`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{cases} i \\sqrt{- x} & \\text{for}\\: x < 0 \\\\\\sqrt{x} & \\text{otherwise} \\end{cases}$" - ], - "text/plain": [ - "Piecewise((I*sqrt(-x), x < 0), (sqrt(x), True))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def complex_sqrt(x: sp.Symbol) -> sp.Expr:\n", - " return sp.Piecewise(\n", - " (sp.sqrt(-x) * sp.I, x < 0),\n", - " (sp.sqrt(x), True),\n", - " )\n", - "\n", - "\n", - "complex_sqrt(x)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle 2 i$" - ], - "text/plain": [ - "2*I" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle 2$" - ], - "text/plain": [ - "2" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "display(\n", - " complex_sqrt(-4),\n", - " complex_sqrt(+4),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Be careful though when lambdifying this expression: do not use the `__dict__` of the {mod}`numpy` module as backend, but use the module itself instead. When using `__dict__`, {func}`~sympy.utilities.lambdify.lambdify` will return an `if-else` statement, which is inefficient and, worse, will result in problems with {doc}`JAX `:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "\n", - "Do not use the module `__dict__` for the `modules` argument of {func}`~sympy.utilities.lambdify.lambdify`.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return (((1j*sqrt(-x)) if (x < 0) else (sqrt(x))))\n", - "\n" - ] - } - ], - "source": [ - "np_complex_sqrt_no_select = sp.lambdify(x, complex_sqrt(x), np.__dict__)\n", - "source = inspect.getsource(np_complex_sqrt_no_select)\n", - "print(source)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "1j" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_complex_sqrt_no_select(-1)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception", - "remove-output" - ] - }, - "outputs": [ - { - "ename": "ConcretizationTypeError", - "evalue": "Abstract tracer value encountered where concrete value is expected: Tracedwith\nThe problem arose with the `bool` function. \nWhile tracing the function _lambdifygenerated at :1, transformed by jit., this concrete value was not available in Python because it depends on the value of the arguments to _lambdifygenerated at :1, transformed by jit. at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).\n (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mUnfilteredStackTrace\u001b[0m Traceback (most recent call last)", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel_launcher.py:17\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch_new_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/traitlets/config/application.py:976\u001b[0m, in \u001b[0;36mApplication.launch_instance\u001b[0;34m(cls, argv, **kwargs)\u001b[0m\n\u001b[1;32m 975\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m--> 976\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelapp.py:712\u001b[0m, in \u001b[0;36mIPKernelApp.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 711\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 712\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/tornado/platform/asyncio.py:199\u001b[0m, in \u001b[0;36mBaseAsyncIOLoop.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 198\u001b[0m asyncio\u001b[38;5;241m.\u001b[39mset_event_loop(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop)\n\u001b[0;32m--> 199\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43masyncio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_forever\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 200\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mBaseEventLoop.run_forever\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_once\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 571\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36mBaseEventLoop._run_once\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 1858\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m \u001b[43mhandle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/events.py:81\u001b[0m, in \u001b[0;36mHandle._run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_callback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:510\u001b[0m, in \u001b[0;36mKernel.dispatch_queue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 509\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 510\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m 511\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:499\u001b[0m, in \u001b[0;36mKernel.process_one\u001b[0;34m(self, wait)\u001b[0m\n\u001b[1;32m 498\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 499\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:406\u001b[0m, in \u001b[0;36mKernel.dispatch_shell\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 406\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m 407\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:730\u001b[0m, in \u001b[0;36mKernel.execute_request\u001b[0;34m(self, stream, ident, parent)\u001b[0m\n\u001b[1;32m 729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 730\u001b[0m reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m 732\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/ipkernel.py:383\u001b[0m, in \u001b[0;36mIPythonKernel.do_execute\u001b[0;34m(self, code, silent, store_history, user_expressions, allow_stdin, cell_id)\u001b[0m\n\u001b[1;32m 382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 383\u001b[0m res \u001b[38;5;241m=\u001b[39m \u001b[43mshell\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 384\u001b[0m \u001b[43m \u001b[49m\u001b[43mcode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 385\u001b[0m \u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 386\u001b[0m \u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 387\u001b[0m \u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcell_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 388\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 389\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/zmqshell.py:528\u001b[0m, in \u001b[0;36mZMQInteractiveShell.run_cell\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 527\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2881\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m 2880\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2881\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 2882\u001b[0m \u001b[43m \u001b[49m\u001b[43mraw_cell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshell_futures\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\n\u001b[1;32m 2883\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2884\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2936\u001b[0m, in \u001b[0;36mInteractiveShell._run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m 2935\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2936\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoro\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2937\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m(coro)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3135\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_async\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple, cell_id)\u001b[0m\n\u001b[1;32m 3133\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3135\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m 3136\u001b[0m interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m 3138\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3338\u001b[0m, in \u001b[0;36mInteractiveShell.run_ast_nodes\u001b[0;34m(self, nodelist, cell_name, interactivity, compiler, result)\u001b[0m\n\u001b[1;32m 3337\u001b[0m asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m 3339\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n", - " \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n", - "Input \u001b[0;32mIn [20]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 1\u001b[0m jax_complex_sqrt_no_select \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np_complex_sqrt_no_select)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_complex_sqrt_no_select\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/_src/traceback_util.py:143\u001b[0m, in \u001b[0;36mapi_boundary..reraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/_src/api.py:426\u001b[0m, in \u001b[0;36m_cpp_jit..cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 425\u001b[0m flat_fun, out_tree \u001b[38;5;241m=\u001b[39m flatten_fun(f, in_tree)\n\u001b[0;32m--> 426\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxla\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxla_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 427\u001b[0m \u001b[43m \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 428\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 429\u001b[0m \u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 430\u001b[0m \u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 431\u001b[0m \u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mflat_fun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 432\u001b[0m \u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdonated_invars\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 433\u001b[0m out_pytree_def \u001b[38;5;241m=\u001b[39m out_tree()\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1565\u001b[0m, in \u001b[0;36mCallPrimitive.bind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[0;32m-> 1565\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1556\u001b[0m, in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1555\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m maybe_new_sublevel(top_trace):\n\u001b[0;32m-> 1556\u001b[0m outs \u001b[38;5;241m=\u001b[39m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, apply_todos(env_trace_todo(), outs))\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1568\u001b[0m, in \u001b[0;36mCallPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m 1567\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m-> 1568\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:609\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m 608\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 609\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/xla.py:578\u001b[0m, in \u001b[0;36m_xla_call_impl\u001b[0;34m(fun, device, backend, name, donated_invars, *args)\u001b[0m\n\u001b[1;32m 577\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_xla_call_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, device, backend, name, donated_invars):\n\u001b[0;32m--> 578\u001b[0m compiled_fun \u001b[38;5;241m=\u001b[39m \u001b[43m_xla_callable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 579\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43munsafe_map\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg_spec\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 580\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/linear_util.py:262\u001b[0m, in \u001b[0;36mcache..memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 262\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 263\u001b[0m cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/xla.py:652\u001b[0m, in \u001b[0;36m_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m 651\u001b[0m abstract_args, _ \u001b[38;5;241m=\u001b[39m unzip2(arg_specs)\n\u001b[0;32m--> 652\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mabstract_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mjit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(c, core\u001b[38;5;241m.\u001b[39mTracer) \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m consts):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/partial_eval.py:1209\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, transform_name)\u001b[0m\n\u001b[1;32m 1208\u001b[0m main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m () \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m-> 1209\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1210\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/partial_eval.py:1188\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals)\u001b[0m\n\u001b[1;32m 1187\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[0;32m-> 1188\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1189\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/linear_util.py:166\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 167\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m 168\u001b[0m \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m 169\u001b[0m \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks only\u001b[39;00m\n\u001b[1;32m 171\u001b[0m \u001b[38;5;66;03m# after this exception is handled, which can corrupt the global state.\u001b[39;00m\n", - "File \u001b[0;32m:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m (((\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx)) \u001b[38;5;28;01mif\u001b[39;00m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m (sqrt(x))))\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:538\u001b[0m, in \u001b[0;36mTracer.__bool__\u001b[0;34m(self)\u001b[0m\n\u001b[0;32m--> 538\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__bool__\u001b[39m(\u001b[38;5;28mself\u001b[39m): \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43maval\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_bool\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:960\u001b[0m, in \u001b[0;36mconcretization_function_error..error\u001b[0;34m(self, arg)\u001b[0m\n\u001b[1;32m 959\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21merror\u001b[39m(\u001b[38;5;28mself\u001b[39m, arg):\n\u001b[0;32m--> 960\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m ConcretizationTypeError(arg, fname_context)\n", - "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Tracedwith\nThe problem arose with the `bool` function. \nWhile tracing the function _lambdifygenerated at :1, transformed by jit., this concrete value was not available in Python because it depends on the value of the arguments to _lambdifygenerated at :1, transformed by jit. at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).\n (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError)\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------", - "\nThe above exception was the direct cause of the following exception:\n", - "\u001b[0;31mConcretizationTypeError\u001b[0m Traceback (most recent call last)", - "Input \u001b[0;32mIn [20]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 1\u001b[0m jax_complex_sqrt_no_select \u001b[38;5;241m=\u001b[39m jax\u001b[38;5;241m.\u001b[39mjit(np_complex_sqrt_no_select)\n\u001b[0;32m----> 2\u001b[0m \u001b[43mjax_complex_sqrt_no_select\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m (((\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx)) \u001b[38;5;28;01mif\u001b[39;00m (x \u001b[38;5;241m<\u001b[39m \u001b[38;5;241m0\u001b[39m) \u001b[38;5;28;01melse\u001b[39;00m (sqrt(x))))\n", - "\u001b[0;31mConcretizationTypeError\u001b[0m: Abstract tracer value encountered where concrete value is expected: Tracedwith\nThe problem arose with the `bool` function. \nWhile tracing the function _lambdifygenerated at :1, transformed by jit., this concrete value was not available in Python because it depends on the value of the arguments to _lambdifygenerated at :1, transformed by jit. at flattened positions [0], and the computation of these values is being staged out (that is, delayed rather than executed eagerly).\n (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError)" - ] - } - ], - "source": [ - "jax_complex_sqrt_no_select = jax.jit(np_complex_sqrt_no_select)\n", - "jax_complex_sqrt_no_select(-1)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When instead using the {mod}`numpy` module (or `\"numpy\"`), {func}`~sympy.utilities.lambdify.lambdify` correctly lambdifies to {func}`numpy.select` to represent the cases." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "np_complex_sqrt = sp.lambdify(x, complex_sqrt(x), np)\n", - "source = inspect.getsource(np_complex_sqrt)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return select(\n", - " [less(x, 0), True],\n", - " [1j * sqrt(-x), sqrt(x)],\n", - " default=nan,\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print(format_str(source.replace(\"nan)\", \"nan,)\"), mode=FileMode()))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Still, JAX does not handle this correctly. First, lambdifying JAX again results in this `if-else` syntax:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return (((1j*sqrt(-x)) if (x < 0) else (sqrt(x))))\n", - "\n" - ] - } - ], - "source": [ - "jnp_complex_sqrt = sp.lambdify(x, complex_sqrt(x), jnp)\n", - "source = inspect.getsource(jnp_complex_sqrt)\n", - "print(source)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "But even if we lambdify to {mod}`numpy` and decorate the result with a {func}`jax.jit` decorator, the resulting function does not work properly:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "jax_complex_sqrt_error = jax.jit(np_complex_sqrt)\n", - "source = inspect.getsource(jax_complex_sqrt_error)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return select(\n", - " [less(x, 0), True],\n", - " [1j * sqrt(-x), sqrt(x)],\n", - " default=nan,\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print(format_str(source.replace(\"nan)\", \"nan,)\"), mode=FileMode()))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception", - "hide-output" - ] - }, - "outputs": [ - { - "ename": "TracerArrayConversionError", - "evalue": "The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mUnfilteredStackTrace\u001b[0m Traceback (most recent call last)", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel_launcher.py:17\u001b[0m, in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 15\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mipykernel\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m kernelapp \u001b[38;5;28;01mas\u001b[39;00m app\n\u001b[0;32m---> 17\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mlaunch_new_instance\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/traitlets/config/application.py:976\u001b[0m, in \u001b[0;36mApplication.launch_instance\u001b[0;34m(cls, argv, **kwargs)\u001b[0m\n\u001b[1;32m 975\u001b[0m app\u001b[38;5;241m.\u001b[39minitialize(argv)\n\u001b[0;32m--> 976\u001b[0m \u001b[43mapp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelapp.py:712\u001b[0m, in \u001b[0;36mIPKernelApp.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 711\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 712\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mstart\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/tornado/platform/asyncio.py:199\u001b[0m, in \u001b[0;36mBaseAsyncIOLoop.start\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 198\u001b[0m asyncio\u001b[38;5;241m.\u001b[39mset_event_loop(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39masyncio_loop)\n\u001b[0;32m--> 199\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43masyncio_loop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_forever\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 200\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/base_events.py:570\u001b[0m, in \u001b[0;36mBaseEventLoop.run_forever\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 569\u001b[0m \u001b[38;5;28;01mwhile\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m:\n\u001b[0;32m--> 570\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_once\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 571\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stopping:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/base_events.py:1859\u001b[0m, in \u001b[0;36mBaseEventLoop._run_once\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 1858\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m-> 1859\u001b[0m \u001b[43mhandle\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1860\u001b[0m handle \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/asyncio/events.py:81\u001b[0m, in \u001b[0;36mHandle._run\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 80\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 81\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_context\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_callback\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_args\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 82\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m (\u001b[38;5;167;01mSystemExit\u001b[39;00m, \u001b[38;5;167;01mKeyboardInterrupt\u001b[39;00m):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:510\u001b[0m, in \u001b[0;36mKernel.dispatch_queue\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 509\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 510\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mprocess_one()\n\u001b[1;32m 511\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:499\u001b[0m, in \u001b[0;36mKernel.process_one\u001b[0;34m(self, wait)\u001b[0m\n\u001b[1;32m 498\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 499\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m dispatch(\u001b[38;5;241m*\u001b[39margs)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:406\u001b[0m, in \u001b[0;36mKernel.dispatch_shell\u001b[0;34m(self, msg)\u001b[0m\n\u001b[1;32m 405\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(result):\n\u001b[0;32m--> 406\u001b[0m \u001b[38;5;28;01mawait\u001b[39;00m result\n\u001b[1;32m 407\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/kernelbase.py:730\u001b[0m, in \u001b[0;36mKernel.execute_request\u001b[0;34m(self, stream, ident, parent)\u001b[0m\n\u001b[1;32m 729\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m inspect\u001b[38;5;241m.\u001b[39misawaitable(reply_content):\n\u001b[0;32m--> 730\u001b[0m reply_content \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m reply_content\n\u001b[1;32m 732\u001b[0m \u001b[38;5;66;03m# Flush output before sending the reply.\u001b[39;00m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/ipkernel.py:383\u001b[0m, in \u001b[0;36mIPythonKernel.do_execute\u001b[0;34m(self, code, silent, store_history, user_expressions, allow_stdin, cell_id)\u001b[0m\n\u001b[1;32m 382\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m with_cell_id:\n\u001b[0;32m--> 383\u001b[0m res \u001b[38;5;241m=\u001b[39m \u001b[43mshell\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 384\u001b[0m \u001b[43m \u001b[49m\u001b[43mcode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 385\u001b[0m \u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 386\u001b[0m \u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 387\u001b[0m \u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcell_id\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 388\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 389\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/ipykernel/zmqshell.py:528\u001b[0m, in \u001b[0;36mZMQInteractiveShell.run_cell\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 527\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_last_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m--> 528\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43msuper\u001b[39;49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_cell\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2881\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m 2880\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2881\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_run_cell\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 2882\u001b[0m \u001b[43m \u001b[49m\u001b[43mraw_cell\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mstore_history\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msilent\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mshell_futures\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mcell_id\u001b[49m\n\u001b[1;32m 2883\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2884\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:2936\u001b[0m, in \u001b[0;36mInteractiveShell._run_cell\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, cell_id)\u001b[0m\n\u001b[1;32m 2935\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m-> 2936\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mrunner\u001b[49m\u001b[43m(\u001b[49m\u001b[43mcoro\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 2937\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/async_helpers.py:129\u001b[0m, in \u001b[0;36m_pseudo_sync_runner\u001b[0;34m(coro)\u001b[0m\n\u001b[1;32m 128\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 129\u001b[0m \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m 130\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mStopIteration\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m exc:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3135\u001b[0m, in \u001b[0;36mInteractiveShell.run_cell_async\u001b[0;34m(self, raw_cell, store_history, silent, shell_futures, transformed_cell, preprocessing_exc_tuple, cell_id)\u001b[0m\n\u001b[1;32m 3133\u001b[0m interactivity \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mnone\u001b[39m\u001b[38;5;124m\"\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m silent \u001b[38;5;28;01melse\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mast_node_interactivity\n\u001b[0;32m-> 3135\u001b[0m has_raised \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_ast_nodes(code_ast\u001b[38;5;241m.\u001b[39mbody, cell_name,\n\u001b[1;32m 3136\u001b[0m interactivity\u001b[38;5;241m=\u001b[39minteractivity, compiler\u001b[38;5;241m=\u001b[39mcompiler, result\u001b[38;5;241m=\u001b[39mresult)\n\u001b[1;32m 3138\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mlast_execution_succeeded \u001b[38;5;241m=\u001b[39m \u001b[38;5;129;01mnot\u001b[39;00m has_raised\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/IPython/core/interactiveshell.py:3338\u001b[0m, in \u001b[0;36mInteractiveShell.run_ast_nodes\u001b[0;34m(self, nodelist, cell_name, interactivity, compiler, result)\u001b[0m\n\u001b[1;32m 3337\u001b[0m asy \u001b[38;5;241m=\u001b[39m compare(code)\n\u001b[0;32m-> 3338\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28;01mawait\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mrun_code(code, result, async_\u001b[38;5;241m=\u001b[39masy):\n\u001b[1;32m 3339\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n", - " \u001b[0;31m[... skipping hidden 1 frame]\u001b[0m\n", - "Input \u001b[0;32mIn [26]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax_complex_sqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/_src/traceback_util.py:143\u001b[0m, in \u001b[0;36mapi_boundary..reraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mfun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/_src/api.py:426\u001b[0m, in \u001b[0;36m_cpp_jit..cache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 425\u001b[0m flat_fun, out_tree \u001b[38;5;241m=\u001b[39m flatten_fun(f, in_tree)\n\u001b[0;32m--> 426\u001b[0m out_flat \u001b[38;5;241m=\u001b[39m \u001b[43mxla\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mxla_call\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 427\u001b[0m \u001b[43m \u001b[49m\u001b[43mflat_fun\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 428\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs_flat\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 429\u001b[0m \u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 430\u001b[0m \u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 431\u001b[0m \u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mflat_fun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[38;5;18;43m__name__\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 432\u001b[0m \u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdonated_invars\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 433\u001b[0m out_pytree_def \u001b[38;5;241m=\u001b[39m out_tree()\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1565\u001b[0m, in \u001b[0;36mCallPrimitive.bind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1564\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbind\u001b[39m(\u001b[38;5;28mself\u001b[39m, fun, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mparams):\n\u001b[0;32m-> 1565\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_bind\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1556\u001b[0m, in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1555\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m maybe_new_sublevel(top_trace):\n\u001b[0;32m-> 1556\u001b[0m outs \u001b[38;5;241m=\u001b[39m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtop_trace\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1557\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mmap\u001b[39m(full_lower, apply_todos(env_trace_todo(), outs))\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:1568\u001b[0m, in \u001b[0;36mCallPrimitive.process\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m 1567\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess\u001b[39m(\u001b[38;5;28mself\u001b[39m, trace, fun, tracers, params):\n\u001b[0;32m-> 1568\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mtrace\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mprocess_call\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:609\u001b[0m, in \u001b[0;36mEvalTrace.process_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m 608\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mprocess_call\u001b[39m(\u001b[38;5;28mself\u001b[39m, primitive, f, tracers, params):\n\u001b[0;32m--> 609\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mprimitive\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mimpl\u001b[49m\u001b[43m(\u001b[49m\u001b[43mf\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mtracers\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/xla.py:578\u001b[0m, in \u001b[0;36m_xla_call_impl\u001b[0;34m(fun, device, backend, name, donated_invars, *args)\u001b[0m\n\u001b[1;32m 577\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_xla_call_impl\u001b[39m(fun: lu\u001b[38;5;241m.\u001b[39mWrappedFun, \u001b[38;5;241m*\u001b[39margs, device, backend, name, donated_invars):\n\u001b[0;32m--> 578\u001b[0m compiled_fun \u001b[38;5;241m=\u001b[39m \u001b[43m_xla_callable\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdevice\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mbackend\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdonated_invars\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 579\u001b[0m \u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43munsafe_map\u001b[49m\u001b[43m(\u001b[49m\u001b[43marg_spec\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 580\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/linear_util.py:262\u001b[0m, in \u001b[0;36mcache..memoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 262\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mcall\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 263\u001b[0m cache[key] \u001b[38;5;241m=\u001b[39m (ans, fun\u001b[38;5;241m.\u001b[39mstores)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/xla.py:652\u001b[0m, in \u001b[0;36m_xla_callable\u001b[0;34m(fun, device, backend, name, donated_invars, *arg_specs)\u001b[0m\n\u001b[1;32m 651\u001b[0m abstract_args, _ \u001b[38;5;241m=\u001b[39m unzip2(arg_specs)\n\u001b[0;32m--> 652\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mpe\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mtrace_to_jaxpr_final\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mabstract_args\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mtransform_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mjit\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 653\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28many\u001b[39m(\u001b[38;5;28misinstance\u001b[39m(c, core\u001b[38;5;241m.\u001b[39mTracer) \u001b[38;5;28;01mfor\u001b[39;00m c \u001b[38;5;129;01min\u001b[39;00m consts):\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/partial_eval.py:1209\u001b[0m, in \u001b[0;36mtrace_to_jaxpr_final\u001b[0;34m(fun, in_avals, transform_name)\u001b[0m\n\u001b[1;32m 1208\u001b[0m main\u001b[38;5;241m.\u001b[39mjaxpr_stack \u001b[38;5;241m=\u001b[39m () \u001b[38;5;66;03m# type: ignore\u001b[39;00m\n\u001b[0;32m-> 1209\u001b[0m jaxpr, out_avals, consts \u001b[38;5;241m=\u001b[39m \u001b[43mtrace_to_subjaxpr_dynamic\u001b[49m\u001b[43m(\u001b[49m\u001b[43mfun\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmain\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43min_avals\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1210\u001b[0m \u001b[38;5;28;01mdel\u001b[39;00m fun, main\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/interpreters/partial_eval.py:1188\u001b[0m, in \u001b[0;36mtrace_to_subjaxpr_dynamic\u001b[0;34m(fun, main, in_avals)\u001b[0m\n\u001b[1;32m 1187\u001b[0m in_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mnew_arg, in_avals)\n\u001b[0;32m-> 1188\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[43mfun\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcall_wrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43min_tracers\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1189\u001b[0m out_tracers \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mmap\u001b[39m(trace\u001b[38;5;241m.\u001b[39mfull_raise, ans)\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/linear_util.py:166\u001b[0m, in \u001b[0;36mWrappedFun.call_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 166\u001b[0m ans \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mf\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;28;43mdict\u001b[39;49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mparams\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 167\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m 168\u001b[0m \u001b[38;5;66;03m# Some transformations yield from inside context managers, so we have to\u001b[39;00m\n\u001b[1;32m 169\u001b[0m \u001b[38;5;66;03m# interrupt them before reraising the exception. Otherwise they will only\u001b[39;00m\n\u001b[1;32m 170\u001b[0m \u001b[38;5;66;03m# get garbage-collected at some later time, running their cleanup tasks only\u001b[39;00m\n\u001b[1;32m 171\u001b[0m \u001b[38;5;66;03m# after this exception is handled, which can corrupt the global state.\u001b[39;00m\n", - "File \u001b[0;32m:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m (select([\u001b[43mless\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m,\u001b[38;5;28;01mTrue\u001b[39;00m], [\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx),sqrt(x)], default\u001b[38;5;241m=\u001b[39mnan))\n", - "File \u001b[0;32m~/miniconda3/envs/compwa.github.io/lib/python3.9/site-packages/jax/core.py:472\u001b[0m, in \u001b[0;36mTracer.__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__array__\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkw):\n\u001b[0;32m--> 472\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m TracerArrayConversionError(\u001b[38;5;28mself\u001b[39m)\n", - "\u001b[0;31mUnfilteredStackTrace\u001b[0m: jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)\n\nThe stack trace below excludes JAX-internal frames.\nThe preceding is the original exception that occurred, unmodified.\n\n--------------------", - "\nThe above exception was the direct cause of the following exception:\n", - "\u001b[0;31mTracerArrayConversionError\u001b[0m Traceback (most recent call last)", - "Input \u001b[0;32mIn [26]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[0;32m----> 1\u001b[0m \u001b[43mjax_complex_sqrt_error\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m-\u001b[39;49m\u001b[38;5;241;43m1\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m:2\u001b[0m, in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_lambdifygenerated\u001b[39m(x):\n\u001b[0;32m----> 2\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m (select([\u001b[43mless\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[43m)\u001b[49m,\u001b[38;5;28;01mTrue\u001b[39;00m], [\u001b[38;5;241m1\u001b[39mj\u001b[38;5;241m*\u001b[39msqrt(\u001b[38;5;241m-\u001b[39mx),sqrt(x)], default\u001b[38;5;241m=\u001b[39mnan))\n", - "\u001b[0;31mTracerArrayConversionError\u001b[0m: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Tracedwith (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError)" - ] - } - ], - "source": [ - "jax_complex_sqrt_error(-1)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The very same function in created purely with {mod}`jax.numpy` does work without problems, so it seems this is a SymPy problem:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "@jax.jit\n", - "def jax_complex_sqrt(x):\n", - " return jnp.select(\n", - " [jnp.less(x, 0), True],\n", - " [1j * jnp.sqrt(-x), jnp.sqrt(x)],\n", - " default=jnp.nan,\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "DeviceArray([0. +1.j , 0. +0.70710677j,\n", - " 0. +0.j , 0.70710677+0.j ,\n", - " 1. +0.j ], dtype=complex64)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jax_complex_sqrt(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A solution to this is presented in {ref}`report/001:Handle for JAX`." - ] - } - ], - "metadata": { - "colab": { - "toc_visible": true - }, - "keep_output": 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.10.14" - }, - "orphan": true - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/001.ipynb b/docs/report/001.ipynb deleted file mode 100644 index 186d8c16..00000000 --- a/docs/report/001.ipynb +++ /dev/null @@ -1,710 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "lambdification", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Custom lambdification\n", - "TR-001\n", - "^^^\n", - "See also {doc}`SymPy's tutorial page on the printing modules `.\n", - "+++\n", - "✅ [ampform#72](https://github.com/ComPWA/ampform/issues/72), [tensorwaves#284](https://github.com/ComPWA/tensorwaves/issues/284)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Custom lambdification\n", - "" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q black==21.5b2 jax==0.4.28 jaxlib==0.4.28 matplotlib==3.4.2 numpy==1.26.4 sympy==1.8" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import inspect\n", - "\n", - "import jax\n", - "import jax.numpy as jnp\n", - "import numpy as np\n", - "import sympy as sp\n", - "from black import FileMode, format_str" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Overwrite printer methods" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As noted in [TR-000](000.ipynb), it's hard to lambdify a {func}`sympy.sqrt ` to {doc}`JAX `. One possible way out is to define a custom class that derives from {class}`sympy.Expr ` and {doc}`overwrite its printer methods `." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from sympy.printing.printer import Printer\n", - "\n", - "\n", - "class ComplexSqrt(sp.Expr):\n", - " def __new__(cls, x, *args, **kwargs):\n", - " x = sp.sympify(x)\n", - " expr = sp.Expr.__new__(cls, x, *args, **kwargs)\n", - " if hasattr(x, \"free_symbols\") and not x.free_symbols:\n", - " return expr.evaluate()\n", - " return expr\n", - "\n", - " def evaluate(self):\n", - " x = self.args[0]\n", - " if not x.is_real:\n", - " return sp.sqrt(x)\n", - " return sp.Piecewise(\n", - " (sp.I * sp.sqrt(-x), x < 0),\n", - " (sp.sqrt(x), True),\n", - " )\n", - "\n", - " def _latex(self, printer: Printer, *args) -> str:\n", - " x = printer._print(self.args[0])\n", - " return Rf\"\\sqrt[\\mathrm{{c}}]{{{x}}}\"\n", - "\n", - " def _numpycode(self, printer: Printer, *args) -> str:\n", - " printer.module_imports[\"numpy.lib\"].add(\"scimath\")\n", - " x = printer._print(self.args[0])\n", - " return f\"scimath.sqrt({x})\"\n", - "\n", - " def _pythoncode(self, printer: Printer, *args) -> str:\n", - " printer.module_imports[\"cmath\"].add(\"sqrt as csqrt\")\n", - " x = printer._print(self.args[0])\n", - " return f\"csqrt({x})\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As opposed to the {doc}`derivation of a sympy.Expr `, this class evaluates directly, because the `evaluate` key-word argument is not used processed by the `__new__` method:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle 2 i$" - ], - "text/plain": [ - "2*I" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "ComplexSqrt(-4)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The `_latex()` method ensures that `ComplexSqrt` renders nicely in notebooks:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sqrt[\\mathrm{c}]{x}$" - ], - "text/plain": [ - "ComplexSqrt(x)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "x = sp.Symbol(\"x\")\n", - "ComplexSqrt(x)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Plot custom class" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In addition, one may modify this `Lambdifier` class, so that {func}`sympy.plot() ` also works on this custom class:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from sympy.plotting.experimental_lambdify import Lambdifier\n", - "\n", - "Lambdifier.builtin_functions_different[\"ComplexSqrt\"] = \"sqrt\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "x = sp.Symbol(\"x\")\n", - "expr = ComplexSqrt(x)\n", - "p1 = sp.plot(sp.re(expr), (x, -1, 2), show=False, line_color=\"red\")\n", - "p2 = sp.plot(sp.im(expr), (x, -1, 2), show=False)\n", - "p1.append(p2[0])\n", - "p1.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164990419-9cd03001-d6f3-44b1-a8f9-beed2c6bf69b.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Lambdifying" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The important part, lambdifying to {mod}`numpy` or {mod}`math` works well as well now:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return (csqrt(x))\n", - "\n" - ] - } - ], - "source": [ - "lambdified_py = sp.lambdify(x, ComplexSqrt(x), \"math\")\n", - "source = inspect.getsource(lambdified_py)\n", - "print(source)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return (scimath.sqrt(x))\n", - "\n" - ] - } - ], - "source": [ - "numpy_lambdified = sp.lambdify(x, ComplexSqrt(x), \"numpy\")\n", - "source = inspect.getsource(numpy_lambdified)\n", - "print(source)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([0. +1.j , 0. +0.70710678j,\n", - " 0. +0.j , 0.70710678+0.j ,\n", - " 1. +0.j ])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sample = np.linspace(-1, +1, 5)\n", - "numpy_lambdified(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Just as noted in {ref}`report/000:Complex square root` though, {mod}`numpy.emath` is not provided by the NumPy API of {doc}`JAX `. As discussed there, we can at most decorate the {mod}`numpy.emath` version with {func}`jax.jit` and work with static arguments only:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "DeviceArray(0.+1.j, dtype=complex64)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jax_lambdified = jax.jit(numpy_lambdified, backend=\"cpu\", static_argnums=0)\n", - "jax_lambdified(-1)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In that case, unhashable (non-static) input samples are still not accepted:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception", - "keep_output" - ] - }, - "outputs": [ - { - "ename": "ValueError", - "evalue": "Non-hashable static arguments are not supported. An error occured while trying to hash an object of type , [-1. -0.5 0. 0.5 1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mValueError\u001b[0m Traceback (most recent call last)", - "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjax_lambdified\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0msample\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;31mValueError\u001b[0m: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type , [-1. -0.5 0. 0.5 1. ]. The error was:\nTypeError: unhashable type: 'numpy.ndarray'\n" - ] - } - ], - "source": [ - "jax_lambdified(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Handle for JAX" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As concluded in {ref}`report/000:Conditional square root`, the alternative to lambdify to {mod}`numpy.emath` is to lambdify to {func}`numpy.select`. This has some caveats, though, like that you should not use `__dict__`. Worse, JAX is not immediately supported as backend. Fortunately, we now know how to overwrite lambdify methods.\n", - "\n", - "An additional tool we need now is to {doc}`define a new printer class ` for JAX, so that we can also define a special rendering method for `ComplexSqrt` in the case of JAX. Most of its printing methods should be the same as that of SymPy's `NumPyPrinter`, the rest we can overwrite:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "\n", - "Alternative would be to add a method `_jaxcode` to the `ComplexSqrt` class above. See {doc}`sympy:modules/printing`.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from sympy.printing.numpy import NumPyPrinter\n", - "\n", - "\n", - "class JaxPrinter(NumPyPrinter):\n", - " _module = \"jax\"\n", - "\n", - " def _print_ComplexSqrt(self, expr: sp.Expr) -> str:\n", - " arg = expr.args[0]\n", - " x = self._print(arg)\n", - " return (\n", - " f\"select([less({x}, 0), True], [1j * sqrt(-{x}), sqrt({x})],\"\n", - " \" default=nan,)\"\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "numpy_expr = sp.lambdify(x, ComplexSqrt(x), modules=np, printer=JaxPrinter)\n", - "source = inspect.getsource(numpy_expr)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return select(\n", - " [less(x, 0), True],\n", - " [1j * sqrt(-x), sqrt(x)],\n", - " default=nan,\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print(format_str(source, mode=FileMode()))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "jax_expr = sp.lambdify(x, ComplexSqrt(x), modules=jnp, printer=JaxPrinter)\n", - "source = inspect.getsource(jax_expr)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(x):\n", - " return select(\n", - " [less(x, 0), True],\n", - " [1j * sqrt(-x), sqrt(x)],\n", - " default=nan,\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print(format_str(source, mode=FileMode()))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-stderr", - "keep_output" - ] - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n" - ] - }, - { - "data": { - "text/plain": [ - "DeviceArray([0. +1.j , 0. +0.70710677j,\n", - " 0. +0.j , 0.70710677+0.j ,\n", - " 1. +0.j ], dtype=complex64)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jax_expr(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The lambdified function can of course also be decorated with {func}`jax.jit`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "jit_expr = jax.jit(jax_expr)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Performance check" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "rng = np.random.default_rng()\n", - "sample = rng.normal(size=1_000_000)\n", - "jax_sample = jnp.array(sample)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip} section\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "1.91 ms ± 116 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" - ] - } - ], - "source": [ - "%timeit jit_expr(jax_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "6.31 ms ± 42.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" - ] - } - ], - "source": [ - "%timeit jax_expr(jax_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output", - "remove-stderr" - ] - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - ":2: RuntimeWarning: invalid value encountered in sqrt\n", - " return (select([less(x, 0), True], [1j * sqrt(-x), sqrt(x)], default=nan,))\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "16.9 ms ± 614 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" - ] - } - ], - "source": [ - "%timeit numpy_expr(sample)" - ] - } - ], - "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.10.13" - }, - "orphan": true - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/002.ipynb b/docs/report/002.ipynb deleted file mode 100644 index 12519944..00000000 --- a/docs/report/002.ipynb +++ /dev/null @@ -1,1967 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "lambdification", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Faster lambdification by splitting expressions\n", - "TR-002\n", - "^^^\n", - "This notebook investigates how to speed up {func}`sympy.lambdify ` by splitting up the expression tree of a complicated expression into components, lambdifying those, and then combining them back again.\n", - "+++\n", - "✅ [tensorwaves#281](https://github.com/ComPWA/tensorwaves/issues/281)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Speed up lambdifying" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.9.1 numpy==1.19.5 pandas==1.2.4 qrules==0.8.2 sympy==1.8 tensorwaves[jax]==0.2.9 git+https://github.com/zfit/phasespace@7131fbd" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "from __future__ import annotations\n", - "\n", - "import inspect\n", - "import logging\n", - "import timeit\n", - "import warnings\n", - "from collections.abc import Generator, Sequence\n", - "from typing import Callable\n", - "\n", - "import ampform\n", - "import graphviz\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import pandas as pd\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.dynamics.builder import create_relativistic_breit_wigner_with_ff\n", - "from tensorwaves.data import generate_phsp\n", - "from tensorwaves.data.phasespace import TFUniformRealNumberGenerator\n", - "from tensorwaves.data.transform import HelicityTransformer\n", - "from tensorwaves.model import LambdifiedFunction, SympyModel\n", - "\n", - "LOGGER = logging.getLogger()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Create dummy expression" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "First, let's create an amplitude model with {mod}`ampform`. We'll use this model as complicated {class}`sympy.Expr ` in the rest of this notebooks." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "LOGGER.setLevel(logging.ERROR)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "result = qrules.generate_transitions(\n", - " initial_state=(\"J/psi(1S)\", [-1, +1]),\n", - " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", - " allowed_intermediate_particles=[\"f(0)(980)\"],\n", - " allowed_interaction_types=[\"strong\", \"EM\"],\n", - " formalism_type=\"canonical-helicity\",\n", - ")\n", - "dot = qrules.io.asdot(result, collapse_graphs=True, render_final_state_id=False)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "graphviz.Source(dot).render(\"002-f0(980)-graph\", format=\"svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164983331-6eb948fe-d360-40bd-a4f7-fa1aad3e296a.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "model_builder = ampform.get_builder(result)\n", - "for name in result.get_intermediate_particles().names:\n", - " model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", - "model = model_builder.generate()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "[C[J/\\psi(1S) \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980) \\to \\pi^{0}_{0} \\pi^{0}_{0}],\n", - " Gamma_f(0)(980),\n", - " d_f(0)(980),\n", - " m_1,\n", - " m_12,\n", - " m_2,\n", - " m_f(0)(980),\n", - " phi_1+2,\n", - " phi_1,1+2,\n", - " theta_1+2,\n", - " theta_1,1+2]" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "free_symbols = sorted(model.expression.free_symbols, key=lambda s: s.name)\n", - "free_symbols" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Helicity model components" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A {class}`~ampform.helicity.HelicityModel` has the benefit that it comes with {attr}`~ampform.helicity.HelicityModel.components` (intensities and amplitudes) that together form its {attr}`~ampform.helicity.HelicityModel.expression`. Let's separate these components into _amplitude_ and _intensity_." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output", - "full-width" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "['A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{+1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{+1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=0,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]',\n", - " 'A[J/\\\\psi(1S)_{-1} \\\\to f_{0}(980)_{0} \\\\gamma_{-1,L=2,S=1}; f_{0}(980)_{0} \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0,L=0,S=0}]']" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "amplitudes = {\n", - " name: expr for name, expr in model.components.items() if name.startswith(\"A\")\n", - "}\n", - "sorted(amplitudes)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensities = {\n", - " name: expr for name, expr in model.components.items() if name.startswith(\"I\")\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "assert len(amplitudes) + len(intensities) == len(model.components)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Component structure" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that each intensity consists of a subset of these amplitudes. This means that _intensities have a larger expression tree than amplitudes_." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "amplitude_to_symbol = {\n", - " expr: sp.Symbol(f\"A{i}\") for i, expr in enumerate(amplitudes.values(), 1)\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensity_to_symbol = {\n", - " expr: sp.Symbol(f\"I{i}\") for i, expr in enumerate(intensities.values(), 1)\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle I_{1} + I_{2} + I_{3} + I_{4}$" - ], - "text/plain": [ - "I1 + I2 + I3 + I4" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "intensity_expr = model.expression.subs(intensity_to_symbol, simultaneous=True)\n", - "intensity_expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot = sp.dotprint(intensity_expr)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "graphviz.Source(dot).render(\"002-collapsed-expression-tree\", format=\"svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164983184-fde89791-2e75-4bd1-9c03-9a45edf24216.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left|{A_{1} + A_{2}}\\right|^{2} + \\left|{A_{3} + A_{4}}\\right|^{2} + \\left|{A_{5} + A_{6}}\\right|^{2} + \\left|{A_{7} + A_{8}}\\right|^{2}$" - ], - "text/plain": [ - "Abs(A1 + A2)**2 + Abs(A3 + A4)**2 + Abs(A5 + A6)**2 + Abs(A7 + A8)**2" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "amplitude_expr = model.expression.subs(amplitude_to_symbol, simultaneous=True)\n", - "amplitude_expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot = sp.dotprint(amplitude_expr)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "graphviz.Source(dot).render(\"002-simple-expression-tree\", format=\"svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164983978-73d1b6a4-0f09-4a10-88d8-de9d6a055bf3.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Performance check" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Lambdifying the whole {attr}`HelicityModel.expression ` is slowest. The {func}`~sympy.utilities.lambdify.lambdify` function first prints the expression as a {obj}`str` (!) with (in this case) {mod}`numpy` syntax and then uses {func}`eval` to convert that back to actual {mod}`numpy` objects:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{margin}\n", - "\n", - "We store the time with {mod}`timeit` for section {ref}`report/002:Arbitrary expressions`.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "runtime = {}\n", - "start = timeit.default_timer()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 1.46 s, sys: 703 µs, total: 1.46 s\n", - "Wall time: 1.46 s\n" - ] - } - ], - "source": [ - "%%time\n", - "\n", - "np_complete_model = sp.lambdify(free_symbols, model.expression.doit(), \"numpy\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "stop = timeit.default_timer()\n", - "runtime[\"complete model\"] = stop - start" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Printing to {obj}`str` and converting back with {func}`eval` becomes exponentially slow the larger the expression tree. This means that it's more efficient to lambdify sub-trees of the expression tree separately. Lambdifying the four _intensities_ of this model separately, the effect is not noticeable:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "LOGGER.setLevel(logging.INFO)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:root:Lambdifying I1\n", - "INFO:root:Lambdifying I2\n", - "INFO:root:Lambdifying I3\n", - "INFO:root:Lambdifying I4\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 1.56 s, sys: 4.94 ms, total: 1.56 s\n", - "Wall time: 1.56 s\n" - ] - } - ], - "source": [ - "%%time\n", - "\n", - "for expr, symbol in intensity_to_symbol.items():\n", - " logging.info(f\"Lambdifying {symbol.name}\")\n", - " start = timeit.default_timer()\n", - " sp.lambdify(free_symbols, expr.doit(), \"numpy\")\n", - " stop = timeit.default_timer()\n", - " runtime[symbol.name] = stop - start" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "...but each of the eight _amplitudes_ separately does result in a significant speed-up:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "INFO:root:Lambdifying A1\n", - "INFO:root:Lambdifying A2\n", - "INFO:root:Lambdifying A3\n", - "INFO:root:Lambdifying A4\n", - "INFO:root:Lambdifying A5\n", - "INFO:root:Lambdifying A6\n", - "INFO:root:Lambdifying A7\n", - "INFO:root:Lambdifying A8\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 547 ms, sys: 3.85 ms, total: 550 ms\n", - "Wall time: 547 ms\n" - ] - } - ], - "source": [ - "%%time\n", - "\n", - "np_amplitudes = {}\n", - "for expr, symbol in amplitude_to_symbol.items():\n", - " logging.info(f\"Lambdifying {symbol.name}\")\n", - " start = timeit.default_timer()\n", - " np_expr = sp.lambdify(free_symbols, expr.doit(), \"numpy\")\n", - " stop = timeit.default_timer()\n", - " runtime[symbol.name] = stop - start\n", - " np_amplitudes[symbol] = np_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Recombining components" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{ref}`Recall ` what amplitude module expressed in its amplitude components looks like:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left|{A_{1} + A_{2}}\\right|^{2} + \\left|{A_{3} + A_{4}}\\right|^{2} + \\left|{A_{5} + A_{6}}\\right|^{2} + \\left|{A_{7} + A_{8}}\\right|^{2}$" - ], - "text/plain": [ - "Abs(A1 + A2)**2 + Abs(A3 + A4)**2 + Abs(A5 + A6)**2 + Abs(A7 + A8)**2" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "amplitude_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We have to lambdify that top expression as well:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sorted_amplitude_symbols = sorted(np_amplitudes, key=lambda s: s.name)\n", - "np_amplitude_expr = sp.lambdify(sorted_amplitude_symbols, amplitude_expr, \"numpy\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(A1, A2, A3, A4, A5, A6, A7, A8):\n", - " return (abs(A1 + A2)**2 + abs(A3 + A4)**2 + abs(A5 + A6)**2 + abs(A7 + A8)**2)\n", - "\n" - ] - } - ], - "source": [ - "source = inspect.getsource(np_amplitude_expr)\n", - "print(source)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now have a lambdified expression for the complete amplitude model, as well as lambdified expressions that are to be plugged in to its arguments." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def componentwise_lambdified(*args):\n", - " \"\"\"Lambdified amplitude model, recombined from its amplitude components.\n", - "\n", - " .. warning:: Order of the ``args`` has to be the same as that\n", - " of the ``args`` of the lambdified amplitude components.\n", - " \"\"\"\n", - " amplitude_values = []\n", - " for amp_symbol in sorted_amplitude_symbols:\n", - " np_amplitude = np_amplitudes[amp_symbol]\n", - " values = np_amplitude(*args)\n", - " amplitude_values.append(values)\n", - " return np_amplitude_expr(*amplitude_values)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Test with data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Okay, so does all this work? Let's first generate a phase space sample with good-old {mod}`tensorwaves`. We can then use this sample as input to the component-wise lambdified function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "LOGGER.setLevel(logging.ERROR)\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sympy_model = SympyModel(\n", - " expression=model.expression,\n", - " parameters=model.parameter_defaults,\n", - ")\n", - "intensity = LambdifiedFunction(sympy_model, backend=\"jax\")\n", - "data_converter = HelicityTransformer(model.adapter)\n", - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_sample = generate_phsp(10_000, model.adapter.reaction_info, random_generator=rng)\n", - "phsp_set = data_converter.transform(phsp_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(8, 4))\n", - "ax.hist(phsp_set[\"m_12\"], bins=50, alpha=0.5, density=True)\n", - "ax.hist(\n", - " phsp_set[\"m_12\"],\n", - " bins=50,\n", - " alpha=0.5,\n", - " density=True,\n", - " weights=np.array(intensity(phsp_set)),\n", - ")\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "fig.savefig(\"002-histogram-m12.svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164983924-9ecf9149-af1d-437b-b5f2-4a73a4d1d81b.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The arguments of the component-wise lambdified amplitude model should be covered by the entries in the phase space set and the provided parameter defaults:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "kinematic_variable_names = set(phsp_set)\n", - "parameter_names = {symbol.name for symbol in model.parameter_defaults}\n", - "free_symbol_names = {symbol.name for symbol in free_symbols}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "assert free_symbol_names <= kinematic_variable_names ^ parameter_names" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "That allows us to sort the input arrays and parameter defaults so that they can be used as positional argument input to the component-wise lambdified amplitude model:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "merged_par_var_values = {\n", - " symbol.name: value for symbol, value in model.parameter_defaults.items()\n", - "}\n", - "merged_par_var_values.update(phsp_set)\n", - "args_values = [merged_par_var_values[symbol.name] for symbol in free_symbols]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, here's the result of plugging that back into the component-wise lambdified expression:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([0.00048765, 0.00033425, 0.00524706, ..., 0.00140122, 0.00714365,\n", - " 0.00030117])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "componentwise_result = componentwise_lambdified(*args_values)\n", - "componentwise_result" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And it's indeed the same as that the intensity computed by {mod}`tensorwaves` (direct lambdify):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "-7.307471250984975e-11" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "tensorwaves_result = np.array(intensity(phsp_set))\n", - "mean_difference = (componentwise_result - tensorwaves_result).mean()\n", - "mean_difference" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert mean_difference < 1e-9" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Arbitrary expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The problem with {ref}`report/002:Test with data` is that it requires a {obj}`~ampform.helicity.HelicityModel`. In {mod}`tensorwaves`, we want to work with general {class}`sympy.Expr `s though (see [`SympyModel`](https://tensorwaves.readthedocs.io/en/0.2.7/api/tensorwaves.model.html#tensorwaves.model.SympyModel)), where we don't have sub-{attr}`ampform.helicity.HelicityModel.components` available.\n", - "\n", - "Instead, we have to split up the lambdifying in a more general way that can handle arbitrary {class}`sympy.core.expr.Expr`s. For that we need:\n", - "\n", - "1. A general method of traversing through a SymPy expression tree. This can be done with {doc}`sympy:tutorials/intro-tutorial/manipulation`.\n", - "2. A **fast** method to estimate the complexity of a model, so that we can decide whether a node in the expression tree is small enough to be lambdified without much runtime. The best measure for complexity is {func}`~sympy.core.function.count_ops` (\"count operations\"), see notes under {doc}`sympy:modules/simplify/simplify`." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Expression complexity" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Let's tackle 2. first and use the {attr}`HelicityModel.expression ` and its {attr}`~ampform.helicity.HelicityModel.components` that we lambdified earlier on. Here's an overview of the number of operations versus the time it took to lambdify each component:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "keep_output", - "hide-input" - ] - }, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
operationsruntime (s)
complete model8230.980456
I12090.279897
I22030.235227
I32070.215937
I42010.233635
A11030.045300
A21030.040710
A31000.039767
A41000.035684
A51020.036551
A61020.036198
A7990.042208
A8990.040205
\n", - "
" - ], - "text/plain": [ - " operations runtime (s)\n", - "complete model 823 0.980456\n", - "I1 209 0.279897\n", - "I2 203 0.235227\n", - "I3 207 0.215937\n", - "I4 201 0.233635\n", - "A1 103 0.045300\n", - "A2 103 0.040710\n", - "A3 100 0.039767\n", - "A4 100 0.035684\n", - "A5 102 0.036551\n", - "A6 102 0.036198\n", - "A7 99 0.042208\n", - "A8 99 0.040205" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df = pd.DataFrame(runtime.values(), index=runtime, columns=[\"runtime (s)\"])\n", - "operations = [sp.count_ops(model.expression)]\n", - "operations.extend(sp.count_ops(expr) for expr in intensity_to_symbol)\n", - "operations.extend(sp.count_ops(expr) for expr in amplitude_to_symbol)\n", - "df.insert(0, \"operations\", operations)\n", - "df" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From this we can already roughly see that the lambdify runtime scales roughly with the number of SymPy operations.\n", - "\n", - "To better visualize this, we can lambdify the expressions in {class}`~ampform.dynamics.BlattWeisskopfSquared` for each angular momentums and compute their runtime a number of times with {mod}`timeit`. Note that the {class}`~ampform.dynamics.BlattWeisskopfSquared` becomes increasingly complex the higher the angular momentum." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{cases} 1 & \\text{for}\\: L = 0 \\\\\\frac{2 z}{z + 1} & \\text{for}\\: L = 1 \\\\\\frac{13 z^{2}}{9 z + \\left(z - 3\\right)^{2}} & \\text{for}\\: L = 2 \\\\\\frac{277 z^{3}}{z \\left(z - 15\\right)^{2} + \\left(2 z - 5\\right) \\left(18 z - 45\\right)} & \\text{for}\\: L = 3 \\\\\\frac{12746 z^{4}}{25 z \\left(2 z - 21\\right)^{2} + \\left(z^{2} - 45 z + 105\\right)^{2}} & \\text{for}\\: L = 4 \\\\\\frac{998881 z^{5}}{z^{5} + 15 z^{4} + 315 z^{3} + 6300 z^{2} + 99225 z + 893025} & \\text{for}\\: L = 5 \\\\\\frac{118394977 z^{6}}{z^{6} + 21 z^{5} + 630 z^{4} + 18900 z^{3} + 496125 z^{2} + 9823275 z + 108056025} & \\text{for}\\: L = 6 \\\\\\frac{19727003738 z^{7}}{z^{7} + 28 z^{6} + 1134 z^{5} + 47250 z^{4} + 1819125 z^{3} + 58939650 z^{2} + 1404728325 z + 18261468225} & \\text{for}\\: L = 7 \\\\\\frac{4392846440677 z^{8}}{z^{8} + 36 z^{7} + 1890 z^{6} + 103950 z^{5} + 5457375 z^{4} + 255405150 z^{3} + 9833098275 z^{2} + 273922023375 z + 4108830350625} & \\text{for}\\: L = 8 \\end{cases}$" - ], - "text/plain": [ - "Piecewise((1, Eq(L, 0)), (2*z/(z + 1), Eq(L, 1)), (13*z**2/(9*z + (z - 3)**2), Eq(L, 2)), (277*z**3/(z*(z - 15)**2 + (2*z - 5)*(18*z - 45)), Eq(L, 3)), (12746*z**4/(25*z*(2*z - 21)**2 + (z**2 - 45*z + 105)**2), Eq(L, 4)), (998881*z**5/(z**5 + 15*z**4 + 315*z**3 + 6300*z**2 + 99225*z + 893025), Eq(L, 5)), (118394977*z**6/(z**6 + 21*z**5 + 630*z**4 + 18900*z**3 + 496125*z**2 + 9823275*z + 108056025), Eq(L, 6)), (19727003738*z**7/(z**7 + 28*z**6 + 1134*z**5 + 47250*z**4 + 1819125*z**3 + 58939650*z**2 + 1404728325*z + 18261468225), Eq(L, 7)), (4392846440677*z**8/(z**8 + 36*z**7 + 1890*z**6 + 103950*z**5 + 5457375*z**4 + 255405150*z**3 + 9833098275*z**2 + 273922023375*z + 4108830350625), Eq(L, 8)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from ampform.dynamics import BlattWeisskopfSquared\n", - "\n", - "angular_momentum, z = sp.symbols(\"L z\")\n", - "BlattWeisskopfSquared(angular_momentum, z).doit()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "operations = []\n", - "runtime = []\n", - "for angular_momentum in range(9):\n", - " ff2 = BlattWeisskopfSquared(angular_momentum, z)\n", - " operations.append(sp.count_ops(ff2.doit()))\n", - " n_iterations = 10\n", - " t = timeit.timeit(\n", - " setup=f\"\"\"\n", - "import sympy as sp\n", - "from ampform.dynamics import BlattWeisskopfSquared\n", - "z = sp.Symbol(\"z\")\n", - "ff2 = BlattWeisskopfSquared({angular_momentum}, z)\n", - " \"\"\",\n", - " stmt='sp.lambdify(z, ff2.doit(), \"numpy\")',\n", - " number=n_iterations,\n", - " )\n", - " runtime.append(t / n_iterations * 1_000)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
operationsruntime (ms)
000.81877
131.24712
271.64094
3122.52622
4142.29422
5161.88900
6192.24741
7222.72068
8253.01171
\n", - "
" - ], - "text/plain": [ - " operations runtime (ms)\n", - "0 0 0.81877\n", - "1 3 1.24712\n", - "2 7 1.64094\n", - "3 12 2.52622\n", - "4 14 2.29422\n", - "5 16 1.88900\n", - "6 19 2.24741\n", - "7 22 2.72068\n", - "8 25 3.01171" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "df = pd.DataFrame(\n", - " {\n", - " \"operations\": operations,\n", - " \"runtime (ms)\": runtime,\n", - " },\n", - ")\n", - "df" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(8, 4))\n", - "plt.scatter(x=df[\"operations\"], y=df[\"runtime (ms)\"])\n", - "ax.set_ylim(bottom=0)\n", - "ax.set_xlabel(\"operations\")\n", - "ax.set_ylabel(\"runtime (ms)\")\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "fig.savefig(\"002-runtime-vs-operations.svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164983940-3da1d1df-d740-42e4-8a6e-e899c5148034.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Identifying nodes" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now imagine that we don't know anything about the {attr}`~ampform.helicity.HelicityModel.expression` that we created before other than that it is a {class}`sympy.Expr `." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Approach 1: Generator" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A first attempt is to use a generator to recursively identify components in the expression that lie within a certain 'complexity' (as computed by {func}`~sympy.core.function.count_ops`)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def recurse_tree(\n", - " expression: sp.Expr, *, min_complexity: int = 0, max_complexity: int\n", - ") -> Generator[sp.Expr, None, None]:\n", - " for arg in expression.args:\n", - " complexity = sp.count_ops(arg)\n", - " if complexity < max_complexity and complexity > min_complexity:\n", - " yield arg\n", - " else:\n", - " yield from recurse_tree(\n", - " arg,\n", - " min_complexity=min_complexity,\n", - " max_complexity=max_complexity,\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can then use this generator function to create a mapping of these sub-expressions within the expression tree to {class}`~sympy.core.symbol.Symbol`s. That mapping can then be used in {meth}`~sympy.core.basic.Basic.xreplace` to replace the sub-expressions with those symbols." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{margin}\n", - "\n", - "The {meth}`~sympy.core.basic.Basic.xreplace` method is **much faster** than {meth}`~sympy.core.basic.Basic.subs`, because it doesn't do any evaluation.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 314 ms, sys: 135 µs, total: 314 ms\n", - "Wall time: 313 ms\n" - ] - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\left|{f_{0} + f_{1}}\\right|^{2} + \\left|{f_{2} + f_{3}}\\right|^{2} + \\left|{f_{4} + f_{5}}\\right|^{2} + \\left|{f_{6} + f_{7}}\\right|^{2}$" - ], - "text/plain": [ - "Abs(f0 + f1)**2 + Abs(f2 + f3)**2 + Abs(f4 + f5)**2 + Abs(f6 + f7)**2" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "%%time\n", - "\n", - "expression = model.expression.doit()\n", - "sub_expressions = {}\n", - "for i, expr in enumerate(recurse_tree(expression, max_complexity=100)):\n", - " symbol = sp.Symbol(f\"f{i}\")\n", - " complexity = sp.count_ops(expr)\n", - " sub_expressions[expr] = symbol\n", - "expression.xreplace(sub_expressions)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Approach 2: Direct substitution" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There is one problem though: {meth}`~sympy.core.basic.Basic.xreplace` is not accurate for larger expressions. It would therefore be better to directly substitute the sub-expression with a symbol while we loop over the nodes in the expression tree. The following function can do that:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def split_expression(\n", - " expression: sp.Expr,\n", - " max_complexity: int,\n", - " min_complexity: int = 0,\n", - ") -> tuple[sp.Expr, dict[sp.Symbol, sp.Expr]]:\n", - " i = 0\n", - " symbol_mapping = {}\n", - "\n", - " def recursive_split(sub_expression: sp.Expr) -> sp.Expr:\n", - " nonlocal i\n", - " for arg in sub_expression.args:\n", - " complexity = sp.count_ops(arg)\n", - " if complexity < max_complexity and complexity > min_complexity:\n", - " symbol = sp.Symbol(f\"f{i}\")\n", - " i += 1\n", - " symbol_mapping[symbol] = arg\n", - " sub_expression = sub_expression.xreplace({arg: symbol})\n", - " else:\n", - " new_arg = recursive_split(arg)\n", - " sub_expression = sub_expression.xreplace({arg: new_arg})\n", - " return sub_expression\n", - "\n", - " top_expression = recursive_split(expression)\n", - " return top_expression, symbol_mapping" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And indeed, this is **much faster** than {ref}`report/002:Approach 1: Generator` (it's even possible to parallelize this for loop):" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 7 µs, sys: 1 µs, total: 8 µs\n", - "Wall time: 15.5 µs\n" - ] - } - ], - "source": [ - "%time\n", - "\n", - "top_expression, sub_expressions = split_expression(expression, max_complexity=100)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left|{f_{0} + f_{1}}\\right|^{2} + \\left|{f_{2} + f_{3}}\\right|^{2} + \\left|{f_{4} + f_{5}}\\right|^{2} + \\left|{f_{6} + f_{7}}\\right|^{2}$" - ], - "text/plain": [ - "Abs(f0 + f1)**2 + Abs(f2 + f3)**2 + Abs(f4 + f5)**2 + Abs(f6 + f7)**2" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "top_expression" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\frac{C[J/\\psi(1S) \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980) \\to \\pi^{0}_{0} \\pi^{0}_{0}] \\Gamma_{f(0)(980)} m_{f(0)(980)} \\left(\\frac{\\cos{\\left(\\theta_{1+2} \\right)}}{2} + \\frac{1}{2}\\right) e^{i \\phi_{1+2}}}{- \\frac{i \\Gamma_{f(0)(980)} m_{f(0)(980)} \\sqrt{\\frac{\\left(m_{12}^{2} - \\left(m_{1} - m_{2}\\right)^{2}\\right) \\left(m_{12}^{2} - \\left(m_{1} + m_{2}\\right)^{2}\\right)}{m_{12}^{2}}} \\sqrt{m_{f(0)(980)}^{2}}}{\\sqrt{\\frac{\\left(m_{f(0)(980)}^{2} - \\left(m_{1} - m_{2}\\right)^{2}\\right) \\left(m_{f(0)(980)}^{2} - \\left(m_{1} + m_{2}\\right)^{2}\\right)}{m_{f(0)(980)}^{2}}} \\left|{m_{12}}\\right|} - m_{12}^{2} + m_{f(0)(980)}^{2}}$" - ], - "text/plain": [ - "C[J/\\psi(1S) \\to f_{0}(980)_{0} \\gamma_{+1}; f_{0}(980) \\to \\pi^{0}_{0} \\pi^{0}_{0}]*Gamma_f(0)(980)*m_f(0)(980)*(cos(theta_1+2)/2 + 1/2)*exp(I*phi_1+2)/(-I*Gamma_f(0)(980)*m_f(0)(980)*sqrt((m_12**2 - (m_1 - m_2)**2)*(m_12**2 - (m_1 + m_2)**2)/m_12**2)*sqrt(m_f(0)(980)**2)/(sqrt((m_f(0)(980)**2 - (m_1 - m_2)**2)*(m_f(0)(980)**2 - (m_1 + m_2)**2)/m_f(0)(980)**2)*Abs(m_12)) - m_12**2 + m_f(0)(980)**2)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sub_expressions[sp.Symbol(\"f0\")]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Lambdify and combine" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now that we have the machinery to split up arbitrary expressions by complexity, we need to lambdify the top expression as well as each of the sub-expressions and recombine them. The following function can do that and return a recombined {class}`~typing.Callable`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def optimized_lambdify(\n", - " args: Sequence[sp.Symbol],\n", - " expr: sp.Expr,\n", - " modules: str | None = None,\n", - " min_complexity: int = 0,\n", - " max_complexity: int = 100,\n", - ") -> Callable:\n", - " top_expression, definitions = split_expression(\n", - " expression,\n", - " min_complexity=min_complexity,\n", - " max_complexity=max_complexity,\n", - " )\n", - " top_symbols = sorted(definitions, key=lambda s: s.name)\n", - " top_lambdified = sp.lambdify(top_symbols, top_expression, modules)\n", - " sub_lambdified = [\n", - " sp.lambdify(args, definitions[symbol], modules) for symbol in top_symbols\n", - " ]\n", - "\n", - " def recombined_function(*args):\n", - " new_args = [sub_expr(*args) for sub_expr in sub_lambdified]\n", - " return top_lambdified(*new_args)\n", - "\n", - " return recombined_function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We can use the same input values as in {ref}`report/002:Test with data` to check that the resulting lambdified expression results in the same output." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 8 µs, sys: 1 µs, total: 9 µs\n", - "Wall time: 17.4 µs\n" - ] - } - ], - "source": [ - "%time\n", - "\n", - "treewise_lambdified = optimized_lambdify(free_symbols, expression, \"numpy\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([0.00048765, 0.00033425, 0.00524706, ..., 0.00140122, 0.00714365,\n", - " 0.00030117])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "treewise_result = treewise_lambdified(*args_values)\n", - "treewise_result" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And it's indeed the same as that the intensity computed by {mod}`tensorwaves` (direct lambdify):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "-7.307471274905997e-11" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "mean_difference = (treewise_result - tensorwaves_result).mean()\n", - "mean_difference" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert mean_difference < 1e-9" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Comparison" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now have a look at a slightly more complicated model:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "result = qrules.generate_transitions(\n", - " initial_state=(\"J/psi(1S)\", [+1]),\n", - " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", - " allowed_intermediate_particles=[\"f(0)\"],\n", - " allowed_interaction_types=[\"strong\", \"EM\"],\n", - " formalism_type=\"canonical-helicity\",\n", - ")\n", - "model_builder = ampform.get_builder(result)\n", - "for name in result.get_intermediate_particles().names:\n", - " model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", - "complex_model = model_builder.generate()\n", - "dot = qrules.io.asdot(result, collapse_graphs=True, render_final_state_id=False)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "graphviz.Source(dot).render(\"002-f0-graph\", format=\"svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164983377-f7c3c5a3-edfd-49aa-b449-08ee77cda67f.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This makes it clear that the functions defined in {ref}`report/002:Arbitrary expressions` results in a huge speed-up!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "new_expression = complex_model.expression.doit()\n", - "new_free_symbols = sorted(new_expression.free_symbols, key=lambda s: s.name)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip} section\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 4.57 s, sys: 3.16 ms, total: 4.57 s\n", - "Wall time: 4.57 s\n" - ] - } - ], - "source": [ - "%%time\n", - "\n", - "np_expr = sp.lambdify(new_free_symbols, new_expression)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 261 ms, sys: 87 µs, total: 262 ms\n", - "Wall time: 260 ms\n" - ] - } - ], - "source": [ - "%%time\n", - "\n", - "np_expr = optimized_lambdify(new_free_symbols, new_expression)" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/003.ipynb b/docs/report/003.ipynb deleted file mode 100644 index a80948b0..00000000 --- a/docs/report/003.ipynb +++ /dev/null @@ -1,1029 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Chew-Mandelstam dispersion integrals\n", - "TR-003\n", - "^^^\n", - "This report formulates a symbolic dispersion integral to approach the left-hand cut in the form factor for arbitrary angular momentum. The integral is evaluated with SciPy's [`quad_vec`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad_vec.html).\n", - "+++\n", - "🚧 [ampform#265](https://github.com/ComPWA/ampform/issues/265)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Chew-Mandelstam" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.10 black==23.12.1 mpl-interactions==0.24.1 scipy==1.12.0" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This report is an attempt formulate the Chew-Mandelstam function described in [PDG2023, §50.3.3 (Resonances), pp.14–15](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15) with [SymPy](https://docs.sympy.org), so that it can be implemented in [AmpForm](https://ampform.rtfd.io)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "import inspect\n", - "import os\n", - "import warnings\n", - "from functools import partial\n", - "from typing import Any\n", - "\n", - "import black\n", - "import matplotlib.pyplot as plt\n", - "import mpl_interactions.ipyplot as iplt\n", - "import numpy as np\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.dynamics import (\n", - " BlattWeisskopfSquared,\n", - " BreakupMomentumSquared,\n", - " PhaseSpaceFactor,\n", - " PhaseSpaceFactorComplex,\n", - ")\n", - "from ampform.io import aslatex\n", - "from ampform.sympy import unevaluated\n", - "from ampform.sympy.math import ComplexSqrt\n", - "from IPython.display import SVG, Markdown, Math, display\n", - "from ipywidgets import FloatSlider\n", - "from scipy.integrate import quad_vec\n", - "from sympy.printing.pycode import _unpack_integral_limits\n", - "from tqdm.auto import tqdm\n", - "\n", - "warnings.filterwarnings(\"ignore\")\n", - "PDG = qrules.load_pdg()\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## S-wave" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As can be seen in Eq. (50.44) on [PDG2023, §Resonances, p.15](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15), the Chew-Mandelstam function $\\Sigma_a$ for a particle $a$ decaying to particles $1, 2$ has a simple form for angular momentum $L=0$ ($S$-wave):\n", - "\n", - ":::{math}\n", - ":class: full-width\n", - ":label: chew-mandelstam\n", - "\\Sigma_a(s) = \\frac{1}{16\\pi^2}\n", - "\\left[\n", - " \\frac{2q_a}{\\sqrt{s}}\n", - " \\log\\frac{m_1^2+m_2^2-s+2\\sqrt{s}q_a}{2m_1m_2}\n", - " - \\left(m_1^2-m_2^2\\right)\n", - " \\left(\\frac{1}{s}-\\frac{1}{(m_1+m_2)^2}\\right)\n", - " \\log\\frac{m_1}{m_2}\n", - "\\right]\n", - ":::\n", - "\n", - "The only question is how to deal with negative values for the squared break-up momentum $q_a^2$. Here, we will use AmpForm's {class}`~ampform.sympy.math.ComplexSqrt`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "z = sp.Symbol(\"z\")\n", - "sqrt_expr = ComplexSqrt(z)\n", - "Math(aslatex({sqrt_expr: sqrt_expr.get_definition()}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class BreakupMomentum(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"q_a\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q_squared = BreakupMomentumSquared(s, m1, m2)\n", - " return ComplexSqrt(q_squared)\n", - "\n", - "\n", - "@unevaluated\n", - "class ChewMandelstamSWave(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " # evaluate=False in order to keep same style as PDG\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " left_term = sp.Mul(\n", - " 2 * q / sp.sqrt(s),\n", - " sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)),\n", - " evaluate=False,\n", - " )\n", - " right_term = (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " return sp.Mul(\n", - " 1 / (16 * sp.pi**2),\n", - " left_term - right_term,\n", - " evaluate=False,\n", - " )\n", - "\n", - "\n", - "s, m1, m2 = sp.symbols(\"s m1 m2 \")\n", - "cm_expr = ChewMandelstamSWave(s, m1, m2)\n", - "q_expr = BreakupMomentum(s, m1, m2)\n", - "Math(aslatex({e: e.doit(deep=False) for e in [cm_expr, q_expr]}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It should be noted that this equation is not well-defined along the real axis, that is, for $\\mathrm{Im}(s) = 0$. For this reason, we split $s$ into a real part $s'$ with a small imaginary offset (the PDG indicates this with $s+0i$). We parametrized this imaginary offset with $\\epsilon$, and for the interactive plot, we do so with a power of $10$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s_prime = sp.Symbol(R\"s^{\\prime}\", real=True)\n", - "epsilon = sp.Symbol(\"epsilon\", positive=True)\n", - "s_plus = s_prime + sp.I * sp.Pow(10, -epsilon)\n", - "Math(Rf\"{sp.latex(s)} \\to {sp.latex(s_plus)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "hide-input" - ] - }, - "source": [ - "We are now ready to use express the symbolic expressions above as a numerical function. For comparison, we will plot the Chew-Mandelstam function for $S$-waves next to AmpForm's {class}`~ampform.dynamics.phasespace.PhaseSpaceFactorComplex`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "rho_expr = PhaseSpaceFactorComplex(s, m1, m2)\n", - "Math(aslatex({rho_expr: rho_expr.doit(deep=False)}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "symbols = (s_prime, m1, m2, epsilon)\n", - "cm_func = sp.lambdify(symbols, cm_expr.doit().subs(s, s_plus))\n", - "rho_func = sp.lambdify(symbols, rho_expr.doit().subs(s, s_plus))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As starting values for the interactive plot, we assume $\\pi\\eta$ scattering (just like in the PDG section) and use their masses as values for $m_1$ and $m_1$, respectively." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define sliders and plot range" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "s_min, s_max = -0.15, 1.4\n", - "plot_domain = np.linspace(s_min, s_max, num=500)\n", - "\n", - "m1_val = PDG[\"pi0\"].mass\n", - "m2_val = PDG[\"eta\"].mass\n", - "sliders = dict(\n", - " m1=FloatSlider(\n", - " description=str(m1),\n", - " min=0,\n", - " max=1.2,\n", - " value=m1_val,\n", - " ),\n", - " m2=FloatSlider(\n", - " description=str(m2),\n", - " min=0,\n", - " max=1.2,\n", - " value=m2_val,\n", - " ),\n", - " epsilon=FloatSlider(\n", - " description=str(epsilon),\n", - " min=0,\n", - " max=8,\n", - " value=4,\n", - " ),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "scroll-input", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "fig, axes = plt.subplots(ncols=2, figsize=(11, 4.5), tight_layout=True)\n", - "ax1, ax2 = axes\n", - "for ax in axes:\n", - " ax.axhline(0, linewidth=0.5, c=\"black\")\n", - "\n", - "real_style = {\"label\": \"Real part\", \"c\": \"black\", \"linestyle\": \"dashed\"}\n", - "imag_style = {\"label\": \"Imag part\", \"c\": \"red\"}\n", - "threshold_style = {\"label\": R\"$s_\\mathrm{thr}$\", \"c\": \"grey\", \"linewidth\": 0.5}\n", - "\n", - "xlim = (plot_domain.min(), plot_domain.max())\n", - "ylim = (-1, +1)\n", - "y_factor = 16 * np.pi\n", - "controls = iplt.axvline(\n", - " lambda *args, **kwargs: (kwargs[\"m1\"] + kwargs[\"m2\"]) ** 2,\n", - " **sliders,\n", - " ax=ax1,\n", - " **threshold_style,\n", - ")\n", - "iplt.axvline(\n", - " lambda *args, **kwargs: (kwargs[\"m1\"] + kwargs[\"m2\"]) ** 2,\n", - " controls=controls,\n", - " ax=ax2,\n", - " **threshold_style,\n", - ")\n", - "iplt.plot(\n", - " plot_domain,\n", - " lambda *args, **kwargs: (y_factor * 1j * rho_func(*args, **kwargs)).real,\n", - " controls=controls,\n", - " ylim=ylim,\n", - " xlim=xlim,\n", - " alpha=0.7,\n", - " ax=ax1,\n", - " **real_style,\n", - ")\n", - "iplt.plot(\n", - " plot_domain,\n", - " lambda *args, **kwargs: (y_factor * 1j * rho_func(*args, **kwargs)).imag,\n", - " controls=controls,\n", - " ylim=ylim,\n", - " xlim=xlim,\n", - " alpha=0.7,\n", - " ax=ax1,\n", - " **imag_style,\n", - ")\n", - "\n", - "iplt.plot(\n", - " plot_domain,\n", - " lambda *args, **kwargs: y_factor * cm_func(*args, **kwargs).real,\n", - " controls=controls,\n", - " ylim=ylim,\n", - " xlim=xlim,\n", - " alpha=0.7,\n", - " ax=ax2,\n", - " **real_style,\n", - ")\n", - "iplt.plot(\n", - " plot_domain,\n", - " lambda *args, **kwargs: y_factor * cm_func(*args, **kwargs).imag,\n", - " controls=controls,\n", - " ylim=ylim,\n", - " xlim=xlim,\n", - " alpha=0.7,\n", - " ax=ax2,\n", - " **imag_style,\n", - ")\n", - "\n", - "for ax in axes:\n", - " ax.legend(loc=\"lower right\")\n", - " ax.set_xticks(np.arange(0, 1.21, 0.3))\n", - " ax.set_yticks(np.arange(-1, 1.1, 0.5))\n", - " ax.set_xlim()\n", - " ax.set_xlabel(\"$s$ (GeV$^2$)\")\n", - "\n", - "ax1.set_ylabel(R\"$16\\pi \\; i\\rho(s)$\")\n", - "ax2.set_ylabel(R\"$16\\pi \\; \\Sigma(s)$\")\n", - "ax1.set_title(R\"Complex phase space factor $\\rho$\")\n", - "ax2.set_title(\"Chew-Mandelstam $S$-wave ($L=0$)\")\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " filename = \"003/chew-mandelstam-s-wave.svg\"\n", - " os.makedirs(os.path.dirname(filename), exist_ok=True)\n", - " fig.savefig(filename)\n", - " display(*sliders.values(), SVG(filename))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "Compare the plots above with Figure 50.6 on [PDG2023, §Resonances, p.16](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15).\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## General dispersion integral" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For higher angular momenta, the PDG notes that one has to compute the dispersion integral given by Eq. (50.44) on [PDG2023, §Resonances, p.15](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=15):\n", - "\n", - "$$\n", - "\\Sigma_a(s+0i) =\n", - " \\frac{s-s_{\\mathrm{thr}_a}}{\\pi}\n", - " \\int^\\infty_{s_{\\mathrm{thr}_a}} \\frac{\n", - " \\rho_a(s')n_a^2(s')\n", - " }{\n", - " (s' - s_{\\mathrm{thr}_a})(s'-s-i0)\n", - " }\n", - " \\mathop{}\\!\\mathrm{d}s'\n", - "$$ (dispersion-integral)\n", - "\n", - "Equation {eq}`chew-mandelstam` is the analytic solution for $L=0$." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From Equations (50.33–34) on [PDG2023, §Resonances, p.12](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=12), it can be deduced that the function $n_a^2$ is the same as AmpForm's {class}`~ampform.dynamics.BlattWeisskopfSquared` (note that this function is normalized, whereas the PDG's $F_j$ function has $1$ in the nominator). For this reason, we simply use {class}`~ampform.dynamics.BlattWeisskopfSquared` for the definition of $n_a^2$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class FormFactor(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " L: Any\n", - " q0: Any = 1\n", - " _latex_repr_ = R\"n_a^2\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2, L, q0 = self.args\n", - " q_squared = BreakupMomentumSquared(s, m1, m2)\n", - " return BlattWeisskopfSquared(\n", - " z=q_squared / (q0**2),\n", - " angular_momentum=L,\n", - " )\n", - "\n", - "\n", - "L = sp.Symbol(\"L\", integer=True, positive=True)\n", - "q0 = sp.Symbol(\"q0\", real=True)\n", - "na_expr = FormFactor(s, m1, m2, L, q0)\n", - "bl_expr = BlattWeisskopfSquared(z, L)\n", - "Math(aslatex({e: e.doit(deep=False) for e in [na_expr, bl_expr]}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For $\\rho_a$, we use AmpForm's {class}`~ampform.dynamics.phasespace.PhaseSpaceFactor`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "Math(aslatex({rho_expr: rho_expr.doit(deep=False)}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The symbolic integrand is then formulated as:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "s_thr = (m1 + m2) ** 2\n", - "integrand = (PhaseSpaceFactor(s_prime, m1, m2) * FormFactor(s_prime, m1, m2, L, q0)) / (\n", - " (s_prime - s_thr) * (s_prime - s - epsilon * sp.I)\n", - ")\n", - "integrand" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, we {func}`~sympy.utilities.lambdify.lambdify` this integrand to a {mod}`numpy` expression so that we can integrate it efficiently:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "integrand_func = sp.lambdify(\n", - " args=(s_prime, s, L, epsilon, m1, m2, q0),\n", - " expr=integrand.doit(),\n", - " modules=\"numpy\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "Integrals can be expressed symbolically with SymPy, with some caveats. See {ref}`report/016:SymPy integral`.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As discussed in [TR-016](016.ipynb), {func}`scipy.integrate.quad` cannot integrate over complex-valued functions, but {func}`scipy.integrate.quad_vec` can. For comparison, we now compute this integral for a few values of $L>0$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "s_domain = np.linspace(s_min, s_max, num=50)\n", - "max_L = 3\n", - "l_values = list(range(1, max_L + 1))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It is handy to store the numerical results of each dispersion integral in a {obj}`dict` with $L$ as keys:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "s_thr_val = float(s_thr.subs({m1: m1_val, m2: m2_val}))\n", - "integral_values = {\n", - " l_val: quad_vec(\n", - " lambda x: integrand_func(\n", - " x,\n", - " s=s_domain,\n", - " L=l_val,\n", - " epsilon=1e-3,\n", - " m1=m1_val,\n", - " m2=m2_val,\n", - " q0=1.0,\n", - " ),\n", - " a=s_thr_val,\n", - " b=np.inf,\n", - " )[0]\n", - " for l_val in tqdm(l_values, desc=\"Evaluating integrals\")\n", - "}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, as can be seen from Eq. {eq}`dispersion-integral`, the resulting values from the integral have to be shifted with a factor $\\frac{s-s_{\\mathrm{thr}_a}}{\\pi}$ to get $\\Sigma_a$. We also scale the values with $16\\pi$ so that it can be compared with the plot generated in {ref}`report/003:S-wave`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sigma = {\n", - " l_val: (s_domain - s_thr_val) / np.pi * integral_values[l_val] for l_val in l_values\n", - "}\n", - "sigma_scaled = {l_val: 16 * np.pi * sigma[l_val] for l_val in l_values}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "fig, axes = plt.subplots(\n", - " nrows=len(l_values),\n", - " sharex=True,\n", - " figsize=(5, 2.5 * len(l_values)),\n", - " tight_layout=True,\n", - ")\n", - "fig.suptitle(f\"Dispersion integrals for $m_1={m1_val:.2f}, m_2={m2_val:.2f}$\")\n", - "for ax, l_val in zip(axes, l_values):\n", - " ax.axhline(0, linewidth=0.5, c=\"black\")\n", - " ax.axvline(s_thr_val, **threshold_style)\n", - " ax.plot(s_domain, sigma_scaled[l_val].real, **real_style)\n", - " ax.plot(s_domain, sigma_scaled[l_val].imag, **imag_style)\n", - " ax.set_title(f\"$L = {l_val}$\")\n", - " ax.set_ylabel(R\"$16\\pi \\; \\Sigma(s)$\")\n", - "axes[-1].set_xlabel(\"$s$ (GeV$^2$)\")\n", - "axes[0].legend()\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " filename = \"003/chew-mandelstam-l-non-zero.svg\"\n", - " os.makedirs(os.path.dirname(filename), exist_ok=True)\n", - " fig.savefig(filename)\n", - " display(SVG(filename))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## SymPy expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we attempt to implement Equation {eq}`dispersion-integral` using {ref}`report/016:SymPy integral`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Symbolic integral definition" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "class UnevaluatableIntegral(sp.Integral):\n", - " abs_tolerance = 1e-5\n", - " rel_tolerance = 1e-5\n", - " limit = 50\n", - " dummify = True\n", - "\n", - " def doit(self, **hints):\n", - " args = [arg.doit(**hints) for arg in self.args]\n", - " return self.func(*args)\n", - "\n", - " def _numpycode(self, printer, *args):\n", - " integration_vars, limits = _unpack_integral_limits(self)\n", - " if len(limits) != 1 or len(integration_vars) != 1:\n", - " msg = f\"Cannot handle {len(limits)}-dimensional integrals\"\n", - " raise ValueError(msg)\n", - " x = integration_vars[0]\n", - " a, b = limits[0]\n", - " expr = self.args[0]\n", - " if self.dummify:\n", - " dummy = sp.Dummy()\n", - " expr = expr.xreplace({x: dummy})\n", - " x = dummy\n", - " integrate_func = \"quad_vec\"\n", - " printer.module_imports[\"scipy.integrate\"].add(integrate_func)\n", - " return (\n", - " f\"{integrate_func}(lambda {printer._print(x)}: {printer._print(expr)},\"\n", - " f\" {printer._print(a)}, {printer._print(b)},\"\n", - " f\" epsabs={self.abs_tolerance}, epsrel={self.abs_tolerance},\"\n", - " f\" limit={self.limit})[0]\"\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Definition of the symbolic dispersion integral" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def dispersion_integral(\n", - " s,\n", - " m1,\n", - " m2,\n", - " angular_momentum,\n", - " meson_radius=1,\n", - " s_prime=sp.Symbol(\"x\", real=True),\n", - " epsilon=sp.Symbol(\"epsilon\", positive=True),\n", - "):\n", - " s_thr = (m1 + m2) ** 2\n", - " q_squared = BreakupMomentumSquared(s_prime, m1, m2)\n", - " ff_squared = BlattWeisskopfSquared(\n", - " angular_momentum=L, z=q_squared * meson_radius**2\n", - " )\n", - " phsp_factor = PhaseSpaceFactor(s_prime, m1, m2)\n", - " return sp.Mul(\n", - " (s - s_thr) / sp.pi,\n", - " UnevaluatableIntegral(\n", - " (phsp_factor * ff_squared)\n", - " / (s_prime - s_thr)\n", - " / (s_prime - s - sp.I * epsilon),\n", - " (s_prime, s_thr, sp.oo),\n", - " ),\n", - " evaluate=False,\n", - " )\n", - "\n", - "\n", - "integral_expr = dispersion_integral(s, m1, m2, angular_momentum=L, s_prime=s_prime)\n", - "integral_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "\n", - "We have to keep track of the integration variable ($s'$ in Equation {eq}`dispersion-integral`), so that we don't run into trouble if we use {func}`~sympy.utilities.lambdify.lambdify` with common sub-expressions. The problem is that the integration variable _should not_ be extracted as a common sub-expression, otherwise the lambdified {func}`scipy.integrate.quad_vec` expression cannot handle vectorized input.\n", - "\n", - ":::\n", - "\n", - "To keep the function under the integral simple, we substitute angular momentum $L$ with a definite value before we lambdify:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "UnevaluatableIntegral.abs_tolerance = 1e-4\n", - "UnevaluatableIntegral.rel_tolerance = 1e-4\n", - "integral_s_wave_func = sp.lambdify(\n", - " [s, m1, m2, epsilon],\n", - " integral_expr.subs(L, 0).doit(),\n", - " # integration symbol should not be extracted as common sub-expression!\n", - " cse=partial(sp.cse, ignore=[s_prime], list=False),\n", - ")\n", - "integral_s_wave_func = np.vectorize(integral_s_wave_func)\n", - "\n", - "integral_p_wave_func = sp.lambdify(\n", - " [s, m1, m2, epsilon],\n", - " integral_expr.subs(L, 1).doit(),\n", - " cse=partial(sp.cse, ignore=[s_prime], list=False),\n", - ")\n", - "integral_p_wave_func = np.vectorize(integral_p_wave_func)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "src = inspect.getsource(integral_s_wave_func.pyfunc)\n", - "src = f\"\"\"```python\n", - "{black.format_str(src, mode=black.FileMode()).strip()}\n", - "```\"\"\"\n", - "Markdown(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "s_values = np.linspace(-0.15, 1.4, num=200)\n", - "%time s_wave_values = integral_s_wave_func(s_values, m1_val, m2_val, epsilon=1e-5)\n", - "%time p_wave_values = integral_p_wave_func(s_values, m1_val, m2_val, epsilon=1e-5)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that the dispersion integral for $L=0$ indeed reproduces the same shape as in {ref}`report/003:S-wave`!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "s_wave_values *= 16 * np.pi\n", - "p_wave_values *= 16 * np.pi\n", - "\n", - "s_values = np.linspace(-0.15, 1.4, num=200)\n", - "fig, axes = plt.subplots(nrows=2, figsize=(6, 7), sharex=True)\n", - "ax1, ax2 = axes\n", - "fig.suptitle(f\"Symbolic dispersion integrals for $m_1={m1_val:.2f}, m_2={m2_val:.2f}$\")\n", - "for ax in axes:\n", - " ax.axhline(0, linewidth=0.5, c=\"black\")\n", - " ax.axvline(s_thr_val, **threshold_style)\n", - " ax.set_title(f\"$L = {l_val}$\")\n", - " ax.set_ylabel(R\"$16\\pi \\; \\Sigma(s)$\")\n", - "axes[-1].set_xlabel(\"$s$ (GeV$^2$)\")\n", - "\n", - "ax1.set_title(\"$S$-wave ($L=0$)\")\n", - "ax1.plot(s_values, s_wave_values.real, **real_style)\n", - "ax1.plot(s_values, s_wave_values.imag, **imag_style)\n", - "\n", - "ax2.set_title(\"$P$-wave ($L=1$)\")\n", - "ax2.plot(s_values, p_wave_values.real, **real_style)\n", - "ax2.plot(s_values, p_wave_values.imag, **imag_style)\n", - "\n", - "ax1.legend()\n", - "fig.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " filename = \"003/symbolic-chew-mandelstam.svg\"\n", - " os.makedirs(os.path.dirname(filename), exist_ok=True)\n", - " fig.savefig(filename)\n", - " display(SVG(filename))" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/004.ipynb b/docs/report/004.ipynb deleted file mode 100644 index 335077c8..00000000 --- a/docs/report/004.ipynb +++ /dev/null @@ -1,582 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "physics" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Investigation of analyticity\n", - "TR-004\n", - "^^^\n", - "+++\n", - "WIP\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Analyticity" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "source": [ - "" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.10 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import os\n", - "import warnings\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import sympy as sp\n", - "from ampform.dynamics import (\n", - " BreakupMomentumSquared,\n", - " PhaseSpaceFactor,\n", - " relativistic_breit_wigner_with_ff,\n", - ")\n", - "from IPython.display import Math, display\n", - "from ipywidgets import widgets\n", - "from matplotlib import cm\n", - "from mpl_interactions import heatmap_slicer\n", - "\n", - "warnings.filterwarnings(\"ignore\")\n", - "\n", - "plt.rcParams.update({\"font.size\": 14})\n", - "\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Branch points of $\\rho(s)$" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Investigation of Section 2.1.2 in {cite}`aitchisonUnitarityAnalyticityCrossing2015`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s = sp.Symbol(\"s\")\n", - "m1, m2 = sp.symbols(\"m1 m2\", real=True)\n", - "rho = 16 * sp.pi * PhaseSpaceFactor(s, m1, m2).doit()\n", - "rho" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Or, assuming both decay products to be of unit mass:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "rho.subs({\n", - " m1: 1,\n", - " m2: 1,\n", - "})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "np_rho = sp.lambdify((s, m1, m2), rho, \"numpy\")\n", - "\n", - "m1_val = 1.8\n", - "m2_val = 0.5\n", - "s_thr = (m1_val + m2_val) ** 2\n", - "s_diff = abs(m1_val - m2_val) ** 2\n", - "\n", - "x = np.linspace(-1, +7, num=100)\n", - "y = np.linspace(-2, +2, num=100)\n", - "X, Y = np.meshgrid(x, y)\n", - "s_values = X + Y * 1j\n", - "rho_values = np_rho(s_values, m1=m1_val, m2=m2_val)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "rho_min, rho_max = -5, +5\n", - "Z = rho_values.real\n", - "fig, ax = plt.subplots(\n", - " figsize=(10, 10),\n", - " subplot_kw={\"projection\": \"3d\"},\n", - " tight_layout=True,\n", - ")\n", - "fig.suptitle(R\"Solution sheet for $\\rho$ in the $s$-plane\")\n", - "colors = np.select(\n", - " [rho_min > Z, Z < 0, Z >= 0, rho_max < Z],\n", - " [rho_min, Z / np.abs(rho_min), Z / np.abs(rho_max), rho_max],\n", - ")\n", - "colors = (colors + 1) / 2\n", - "ax.plot_surface(\n", - " X,\n", - " Y,\n", - " Z,\n", - " facecolors=cm.bwr(colors),\n", - " linewidth=0,\n", - ")\n", - "ax.set_xlabel(\"Re($s$)\")\n", - "ax.set_ylabel(\"Im($s$)\")\n", - "ax.set_zlabel(R\"Re($\\rho$)\")\n", - "ax.set_zlim(rho_min, rho_max)\n", - "ax.view_init(elev=25, azim=-110)\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " from IPython.display import Image\n", - "\n", - " output_file = \"004-3D-plot.png\"\n", - " plt.savefig(output_file)\n", - " display(Image(output_file))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "fig, axes = heatmap_slicer(\n", - " x,\n", - " y,\n", - " (rho_values.real, rho_values.imag),\n", - " heatmap_names=(R\"Re($\\rho$)\", R\"Im($\\rho$)\"),\n", - " labels=(\"Re($s$)\", \"Im($s$)\"),\n", - " interaction_type=\"move\",\n", - " slices=\"both\",\n", - " vmin=-5,\n", - " vmax=5,\n", - " figsize=(12, 3),\n", - ")\n", - "for ax in axes[2:]:\n", - " ax.set_ylim(rho_min, rho_max)\n", - " tick_width = 5\n", - " tick_min = np.around(rho_min / tick_width, decimals=0) * tick_width\n", - " ax.set_yticks(np.arange(tick_min, rho_max + 0.1, 5))\n", - "axes[2].set_title(\"Re($s$)\")\n", - "axes[3].set_title(\"Im($s$)\")\n", - "for ax in axes[:3]:\n", - " ax.axvline(s_diff, c=\"black\", linewidth=0.3, linestyle=\"dotted\")\n", - " ax.axvline(s_thr, c=\"black\", linewidth=0.3, linestyle=\"dotted\")\n", - "for ax in axes:\n", - " ax.axvline(0, c=\"black\", linewidth=0.5)\n", - " ax.axhline(0, c=\"black\", linewidth=0.5)\n", - "axes[3].axvline(0, c=\"black\", linewidth=0.5)\n", - "plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " from IPython.display import Image\n", - "\n", - " output_file = \"004-interactive-locator.png\"\n", - " plt.savefig(output_file)\n", - " display(Image(output_file))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Physical vs. unphysical sheet" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Interactive reproduction of Figure 49.1 on {pdg-review}`2020; Resonances; p.2`. The formulas below come from a {func}`~ampform.dynamics.relativistic_breit_wigner_with_ff` with $L=0$. As phase space factor, we used the square root of {class}`~ampform.dynamics.phasespace.BreakupMomentumSquared` instead of the default {class}`~ampform.dynamics.phasespace.PhaseSpaceFactor`, because this introduces only one branch point in the $s$-plane (namely the one over the nominator)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def breakup_momentum(s: sp.Symbol, m_a: sp.Symbol, m_b: sp.Symbol) -> sp.Expr:\n", - " return sp.sqrt(BreakupMomentumSquared(s, m_a, m_b).doit())\n", - "\n", - "\n", - "s = sp.Symbol(\"s\")\n", - "m0, gamma0, m1, m2 = sp.symbols(\"m0 Gamma0 m1 m2\", real=True, positive=True)\n", - "\n", - "unphysical_amp = relativistic_breit_wigner_with_ff(\n", - " s,\n", - " m0,\n", - " gamma0,\n", - " m_a=m1,\n", - " m_b=m2,\n", - " angular_momentum=0,\n", - " meson_radius=1,\n", - " phsp_factor=breakup_momentum,\n", - ").doit()\n", - "\n", - "sqrt_term = unphysical_amp.args[2].args[0].args[2]\n", - "physical_amp = unphysical_amp.subs(sqrt_term, sp.sqrt(sqrt_term**2))\n", - "\n", - "display(\n", - " Math(R\"\\mathrm{Physical:} \\quad \" + sp.latex(physical_amp)),\n", - " Math(R\"\\mathrm{Unphysical:} \\quad \" + sp.latex(unphysical_amp)),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "args = (s, m0, gamma0, m1, m2)\n", - "np_amp_physical = sp.lambdify(args, physical_amp, \"numpy\")\n", - "np_amp_unphysical = sp.lambdify(args, unphysical_amp, \"numpy\")\n", - "\n", - "x_min, x_max = -0.2, 1.3\n", - "y_min, y_max = -1.8, +1.8\n", - "z_min, z_max = -2.5, +2.5\n", - "\n", - "x = np.linspace(x_min, x_max, num=50)\n", - "y_neg = np.linspace(y_min, -1e-4, num=30)\n", - "y_pos = np.linspace(1e-4, y_max, num=30)\n", - "\n", - "X, Y_neg = np.meshgrid(x, y_neg)\n", - "X, Y_pos = np.meshgrid(x, y_pos)\n", - "s_values_neg = X + Y_neg * 1j\n", - "s_values_pos = X + Y_pos * 1j\n", - "\n", - "z_cut_min = 0.75 * z_min\n", - "z_cut_max = 0.75 * z_max\n", - "cut_off_min = np.vectorize(lambda z: max(z_cut_min, z))\n", - "cut_off_max = np.vectorize(lambda z: min(z_cut_max, z))\n", - "\n", - "plot_style = {\n", - " \"linewidth\": 0,\n", - " \"alpha\": 0.7,\n", - " \"antialiased\": True,\n", - " \"rstride\": 1,\n", - " \"cstride\": 1,\n", - "}\n", - "axis_style = {\n", - " \"c\": \"black\",\n", - " \"linewidth\": 0.7,\n", - " \"linestyle\": \"dashed\",\n", - "}\n", - "\n", - "fig, axes = plt.subplots(\n", - " ncols=2,\n", - " figsize=(10, 6),\n", - " subplot_kw={\"projection\": \"3d\"},\n", - " tight_layout=True,\n", - ")\n", - "ax1, ax2 = axes\n", - "fig.suptitle(\"$S$-wave Breit-Wigner ($L=0$) plotted over the complex $s$-plane\")\n", - "\n", - "m0_min = np.sign(x_min) * np.sqrt(np.abs(x_min))\n", - "m0_max = np.sign(x_max) * np.sqrt(np.abs(x_max))\n", - "\n", - "sliders = {\n", - " \"m0\": widgets.FloatSlider(\n", - " min=m0_min,\n", - " max=m0_max,\n", - " value=0.8,\n", - " step=0.01,\n", - " description=\"$m_0$\",\n", - " ),\n", - " \"gamma0\": widgets.FloatSlider(\n", - " min=0.0,\n", - " max=y_max,\n", - " value=0.3,\n", - " step=0.01,\n", - " description=R\"$\\Gamma_0$\",\n", - " ),\n", - " \"m1\": widgets.FloatSlider(\n", - " min=1e-4,\n", - " max=m0_max / 2,\n", - " step=0.01,\n", - " description=\"$m_1$\",\n", - " ),\n", - " \"m2\": widgets.FloatSlider(\n", - " min=1e-4,\n", - " max=m0_max / 2,\n", - " step=0.01,\n", - " description=\"$m_2$\",\n", - " ),\n", - "}\n", - "\n", - "\n", - "@widgets.interact(**sliders)\n", - "def plot(m0, gamma0, m1, m2):\n", - " def plot_expression(ax, amp, neg_color=\"green\"):\n", - " ax.clear()\n", - " z_values_neg = amp(s_values_neg, m0, gamma0, m1, m2).imag\n", - " z_values_pos = amp(s_values_pos, m0, gamma0, m1, m2).imag\n", - " Z_neg = cut_off_min(cut_off_max(z_values_neg))\n", - " Z_pos = cut_off_min(cut_off_max(z_values_pos))\n", - "\n", - " s_thr = (m1 + m2) ** 2\n", - " x0 = x[x >= s_thr] + 1e-4j\n", - " y0 = np.zeros(len(x0))\n", - " z0 = amp(x0, m0, gamma0, m1, m2).imag\n", - "\n", - " ax.plot_surface(X, Y_neg, Z_neg, **plot_style, color=neg_color)\n", - " ax.plot_surface(X, Y_pos, Z_pos, **plot_style, color=\"green\")\n", - " ax.plot(x0, y0, z0, linewidth=2.5, c=\"darkred\", zorder=8)\n", - " ax.scatter([x0[0]], [0], [z0[0]], c=\"darkred\", s=20, zorder=9)\n", - "\n", - " ax.set_xlabel(\"Re($s$)\", labelpad=-15)\n", - " ax.set_ylabel(\"Im($s$)\", labelpad=-15)\n", - " ax.set_zlabel(\"Im($A$)\", labelpad=-15)\n", - " ax.set_xticks([])\n", - " ax.set_yticks([])\n", - " ax.set_zticks([])\n", - " ax.set_zlim(z_min, z_max)\n", - "\n", - " plot_expression(ax1, np_amp_physical)\n", - " plot_expression(ax2, np_amp_unphysical, neg_color=\"gold\")\n", - "\n", - " ax1.text(x_min, y_max, z_max / 2, \"physical sheet\", c=\"green\")\n", - " ax2.text(x_min, y_min, -z_max, \"unphysical sheet\", c=\"gold\")\n", - "\n", - " fig.canvas.draw_idle()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " from IPython.display import SVG\n", - "\n", - " output_file = \"004-physical-vs-unphysical.svg\"\n", - " plt.savefig(output_file)\n", - " display(SVG(output_file))" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/005.ipynb b/docs/report/005.ipynb deleted file mode 100644 index 3b177eea..00000000 --- a/docs/report/005.ipynb +++ /dev/null @@ -1,1051 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix", - "physics" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Symbolic K-matrix expressions\n", - "TR-005\n", - "^^^\n", - "Implementation of this report is tracked through [ampform#67](https://github.com/ComPWA/ampform/issues/67).\n", - "+++\n", - "To be implemented\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# K-matrix" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "source": [ - "" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This report investigates how to implement $K$-matrix dynamics with {doc}`SymPy `. We here describe only the version that is **not Lorentz-invariant**, because it is simplest and allows us to check whether the case $n_R=1, n=1$ (single resonance, single channel) reduces to a Breit-Wigner function. We followed the physics as described by {pdg-review}`Resonances` and {cite}`Chung:1995dx,petersPartialWaveAnalysis2004,meyerMatrixTutorial2008`. For the Lorentz-invariant version, see [TR-009](009.ipynb).\n", - "\n", - "A brief overview of the origin of the $\\boldsymbol{K}$-matrix is given first. This overview follows {cite}`Chung:1995dx`, but skips over quite a few details, as this is only an attempt to provide some context of what is going on." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.11.* matplotlib==3.5.1" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import warnings\n", - "\n", - "import graphviz\n", - "import matplotlib.pyplot as plt\n", - "import mpl_interactions.ipyplot as iplt\n", - "import numpy as np\n", - "import symplot\n", - "import sympy as sp\n", - "from IPython.display import Math, display\n", - "from ipywidgets import widgets as ipywidgets\n", - "from matplotlib import cm\n", - "from mpl_interactions.controller import Controls\n", - "\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Physics" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "source": [ - "The $\\boldsymbol{K}$-matrix formalism is used to describe **coupled, two-body scattering processes** of the type $c_id_i \\to R \\to a_ib_i$, with $i$ representing each separate channel and $R$ a number of resonances that these channels have in common." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot = \"\"\"\n", - "digraph {\n", - " rankdir=LR;\n", - " node [shape=point, width=0];\n", - " edge [arrowhead=none];\n", - " \"Na\" [shape=none, label=\"aᵢ\"];\n", - " \"Nb\" [shape=none, label=\"bᵢ\"];\n", - " \"Nc\" [shape=none, label=\"cᵢ\"];\n", - " \"Nd\" [shape=none, label=\"dᵢ\"];\n", - " { rank=same \"Nc\", \"Nd\" };\n", - " { rank=same \"Na\", \"Nb\" };\n", - " \"Nc\" -> \"N0\";\n", - " \"Nd\" -> \"N0\";\n", - " \"N1\" -> \"Na\";\n", - " \"N1\" -> \"Nb\";\n", - " \"N0\" -> \"N1\" [label=\"R\"];\n", - " \"N0\" [shape=none, label=\"\"];\n", - " \"N1\" [shape=none, label=\"\"];\n", - "}\n", - "\"\"\"\n", - "graph = graphviz.Source(dot)\n", - "graph" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{image} https://user-images.githubusercontent.com/29308176/164994485-fc4843c3-856b-4853-857a-679e258cf7c8.svg\n", - ":align: center\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "graph.render(\"005-two-body-scattering\", format=\"svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Partial wave expansion\n", - "\n", - "In amplitude analysis, the main aim is to express the differential cross section $\\frac{d\\sigma}{d\\Omega}$ (that is, the intensity distribution in each spherical direction $\\Omega=(\\phi,\\theta)$ as we can observe in experiments). This differential cross section can be expressed in terms of the **scattering amplitude** $A$ as:\n", - "\n", - "$$\n", - "\\frac{d\\sigma}{d\\Omega} = \\left|A(\\Omega)\\right|^2\n", - "$$ (differential cross section)\n", - "\n", - "We can now further express $A$ in terms of partial wave amplitudes by splitting it up in terms of its angular momentum components $J$:\n", - "\n", - "$$\n", - "A(\\Omega) = \\frac{1}{2q_i}\\sum_J\\left(2J+1\\right) T^J(s) {D^J_{\\lambda\\mu}}^*\\left(\\phi,\\theta,0\\right)\n", - "$$ (partial waves)\n", - "\n", - "with $\\lambda=\\lambda_a-\\lambda_b$ and $\\mu=\\lambda_c-\\lambda_d$ the helicity differences of the final and initial states $ab,cd$.\n", - "\n", - "The above sketch is just with one channel in mind, but the same holds true though for a number of channels $n$, with the only difference that the $T$ operator becomes a $\\boldsymbol{T}$-matrix of rank $n$.\n", - "\n", - "### Transition operator\n", - "\n", - "The important point is that we have now expressed $A$ in terms of an angular part (depending on $\\Omega$) and a dynamical part $\\boldsymbol{T}$ that depends on the {ref}`Mandelstam variable ` $s$.\n", - "\n", - "\n", - "The dynamical part $\\boldsymbol{T}$ is usually called the **transition operator**. The reason is that it describes the interacting part of the **scattering operator** $\\boldsymbol{S}$, which describes the (complex) amplitude $\\langle f|\\boldsymbol{S}|i\\rangle$ of an initial state $|i\\rangle$ transitioning to a final state $|f\\rangle$. The scattering operator describes both the non-interacting amplitude and the transition amplitude, so it relates to the transition operator as:[^1]\n", - "\n", - "[^1]: Some authors like {cite}`Chung:1995dx` multiply the transition operator by a factor 2.\n", - "\n", - "$$\n", - "\\boldsymbol{S} = \\boldsymbol{I} + i\\boldsymbol{T}\n", - "$$ (scattering operator)\n", - "\n", - "with $\\boldsymbol{I}$ the identity operator. With this in mind, there is an important restriction that the $T$-operator needs to comply with: **unitarity**. This means that $\\boldsymbol{S}$ should conserve probability, namely $\\boldsymbol{S}^\\dagger\\boldsymbol{S} = \\boldsymbol{I}$.\n", - "\n", - "### K-matrix formalism\n", - "\n", - "Now there is a trick to ensure unitarity of $\\boldsymbol{S}$. We can express $\\boldsymbol{S}$ in terms of an operator $\\boldsymbol{K}$ by applying a [Cayley transformation](https://en.wikipedia.org/wiki/Cayley_transform):\n", - "\n", - "$$\n", - "\\boldsymbol{S} = (\\boldsymbol{I} + i\\boldsymbol{K})(I - i\\boldsymbol{K})^{-1}\n", - "$$ (Cayley transformation)\n", - "\n", - "Unitarity is conserved if $K$ is real. Finally, the $\\boldsymbol{T}$-matrix can be expressed in terms of $\\boldsymbol{K}$ as follows:\n", - "\n", - "$$\n", - "\\boldsymbol{T} = \\boldsymbol{K} \\left(\\boldsymbol{I} - i\\boldsymbol{K}\\right)^{-1}\n", - "$$ (T-in-terms-of-K)\n", - "\n", - "### Resonances\n", - "\n", - "The challenge is now to choose a correct parametrization for the elements of $\\boldsymbol{K}$ so that it correctly describes the resonances we observe. There are several choices, but a common one is the following summation over the resonances $R$:\n", - "\n", - "$$\n", - "K_{ij} = \\sum_R\\frac{g_{R,i}^*g_{R,j}}{m_R^2-m^2}\n", - "$$ (K-matrix-parametrization)\n", - "\n", - "with $g_{R,i}$ the residue functions that can be further expressed as\n", - "\n", - "$$\n", - "g_{R,i}=\\gamma_{R,i}\\sqrt{m_R\\Gamma_R}\n", - "$$ (residue-function)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Implementation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The challenge is to generate a correct parametrization for an arbitrary **number of coupled channels $n$** and an arbitrary **number of resonances $n_R$**. Our approach is to construct an $n \\times n$ {obj}`sympy.Matrix ` with {class}`~sympy.core.symbol.Symbol`s as its elements. We then use substitute these {class}`~sympy.core.symbol.Symbol`s with certain parametrizations using {meth}`~sympy.core.basic.Basic.subs`. In order to generate symbols for $n_R$ resonances and $n$ channels, we use {doc}`indexed symbols `.\n", - "\n", - "This approach is less elegant and (theoretically) slower than using {class}`~sympy.matrices.expressions.MatrixSymbol`s. That approach is explored in [TR-007](007.ipynb)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It would be nice to use a {class}`~sympy.core.symbol.Symbol` to represent the number of channels $n$ and specify its value later." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "n_channels = sp.Symbol(\"n\", integer=True, positive=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Unfortunately, this does not work well in the {obj}`~sympy.matrices.dense.Matrix` class. We therefore set variables $n$ to a specific {obj}`int` value and define some other {class}`~sympy.core.symbol.Symbol`s for the rest of the implementation.[^2] The value we choose in this example is `n_channels=1`, because we want to see if this reproduces a Breit-Wigner function.[^3]\n", - "\n", - "[^2]: We use {class}`~sympy.core.symbol.Symbol`s as indices, because that renders more nicely.\n", - "\n", - "[^3]: Of course, there is no need to work with matrices in this $1 \\times 1$ case. To keeps things general, however, we keep using matrices.\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "n_channels = 1\n", - "i, j, R, n_resonances = sp.symbols(\"i j R n_R\", integer=True, negative=False)\n", - "m = sp.Symbol(\"m\", real=True)\n", - "M = sp.IndexedBase(\"m\", shape=(n_resonances,))\n", - "Gamma = sp.IndexedBase(\"Gamma\", shape=(n_resonances,))\n", - "gamma = sp.IndexedBase(\"gamma\", shape=(n_resonances, n_channels))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The parametrization of $K_{ij}$ from Eq. {eq}`K-matrix-parametrization` can be expressed as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def Kij(\n", - " m: sp.Symbol,\n", - " M: sp.IndexedBase,\n", - " Gamma: sp.IndexedBase,\n", - " gamma: sp.IndexedBase,\n", - " i: int,\n", - " j: int,\n", - " n_resonances: int | sp.Symbol,\n", - ") -> sp.Expr:\n", - " g_i = gamma[R, i] * sp.sqrt(M[R] * Gamma[R])\n", - " g_j = gamma[R, j] * sp.sqrt(M[R] * Gamma[R])\n", - " parametrization = (g_i * g_j) / (M[R] ** 2 - m**2)\n", - " return sp.Sum(parametrization, (R, 0, n_resonances - 1))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle K_{ij} = \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,i} {\\gamma}_{R,j} {m}_{R}}{- m^{2} + {m}_{R}^{2}}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "n_R = sp.Symbol(\"n_R\")\n", - "kij = Kij(m, M, Gamma, gamma, i, j, n_R)\n", - "Math(\"K_{ij} = \" + f\"{sp.latex(kij)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now define the $\\boldsymbol{K}$-matrix in terms of a {obj}`~sympy.matrices.dense.Matrix` with {class}`~sympy.tensor.indexed.IndexedBase` instances as elements that can serve as {class}`~sympy.core.symbol.Symbol`s. These {class}`~sympy.core.symbol.Symbol`s will be substituted with the parametrization later. We could of course have inserted the parametrization directly, but this slows down matrix multiplication in the following steps." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle {K}_{i,j}$" - ], - "text/plain": [ - "K[i, j]" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\left[\\begin{matrix}{K}_{0,0}\\end{matrix}\\right]$" - ], - "text/plain": [ - "Matrix([[K[0, 0]]])" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "K_symbol = sp.IndexedBase(\"K\", shape=(n_channels, n_channels))\n", - "K = sp.Matrix([[K_symbol[i, j] for j in range(n_channels)] for i in range(n_channels)])\n", - "display(K_symbol[i, j], K)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The $\\boldsymbol{T}$-matrix can now be computed from Eq. {eq}`T-in-terms-of-K`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left[\\begin{matrix}\\frac{{K}_{0,0}}{- i {K}_{0,0} + 1}\\end{matrix}\\right]$" - ], - "text/plain": [ - "Matrix([[K[0, 0]/(-I*K[0, 0] + 1)]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "T = K * (sp.eye(n_channels) - sp.I * K).inv()\n", - "T" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, we need to substitute the elements $K_{i,j}$ with the parametrization we defined above:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left[\\begin{matrix}\\frac{\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{- i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1}\\end{matrix}\\right]$" - ], - "text/plain": [ - "Matrix([[Sum(Gamma[R]*gamma[R, 0]**2*m[R]/(-m**2 + m[R]**2), (R, 0, n_R - 1))/(-I*Sum(Gamma[R]*gamma[R, 0]**2*m[R]/(-m**2 + m[R]**2), (R, 0, n_R - 1)) + 1)]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "T_subs = T.subs({\n", - " K[i, j]: Kij(m, M, Gamma, gamma, i, j, n_resonances)\n", - " for i in range(n_channels)\n", - " for j in range(n_channels)\n", - "})\n", - "T_subs" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "\n", - "It is important to perform {meth}`~sympy.core.basic.Basic.doit` _after_ {meth}`~sympy.core.basic.Basic.subs`, otherwise the {class}`~sympy.concrete.summations.Sum` cannot be evaluated and there will be no warning of a failed substitution.\n", - "\n", - ":::\n", - "\n", - "Now indeed, when taking $n_R=1$, the resulting element from the $\\boldsymbol{T}$-matrix looks like a Breit-Wigner function (compare {func}`~ampform.dynamics.relativistic_breit_wigner`)!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle - \\frac{{\\Gamma}_{0} {\\gamma}_{0,0}^{2} {m}_{0}}{m^{2} + i {\\Gamma}_{0} {\\gamma}_{0,0}^{2} {m}_{0} - {m}_{0}^{2}}$" - ], - "text/plain": [ - "-Gamma[0]*gamma[0, 0]**2*m[0]/(m**2 + I*Gamma[0]*gamma[0, 0]**2*m[0] - m[0]**2)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "n_resonances_val = 1\n", - "rel_bw = T_subs[0, 0].subs(n_resonances, n_resonances_val).doit()\n", - "if n_resonances_val == 1 or n == 2:\n", - " rel_bw = rel_bw.simplify()\n", - "rel_bw" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generalization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The above procedure has been condensed into a function that can handle an arbitrary number of resonances and an arbitrary number of channels." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def create_symbol_matrix(name: str, n: int) -> sp.Matrix:\n", - " symbol = sp.IndexedBase(\"K\", shape=(n, n))\n", - " return sp.Matrix([[symbol[i, j] for j in range(n)] for i in range(n)])\n", - "\n", - "\n", - "def k_matrix(n_resonances: int, n_channels: int) -> sp.Matrix:\n", - " # Define symbols\n", - " m = sp.Symbol(\"m\", real=True)\n", - " M = sp.IndexedBase(\"m\", shape=(n_resonances,))\n", - " Gamma = sp.IndexedBase(\"Gamma\", shape=(n_resonances,))\n", - " gamma = sp.IndexedBase(\"gamma\", shape=(n_resonances, n_channels))\n", - " # Define K-matrix and T-matrix\n", - " K = create_symbol_matrix(\"K\", n_channels)\n", - " T = K * (sp.eye(n_channels) - sp.I * K).inv()\n", - " # Substitute elements\n", - " return T.subs({\n", - " K[i, j]: Kij(m, M, Gamma, gamma, i, j, n_resonances)\n", - " for i in range(n_channels)\n", - " for j in range(n_channels)\n", - " })" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Single channel, single resonance:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle - \\frac{{\\Gamma}_{0} {\\gamma}_{0,0}^{2} {m}_{0}}{m^{2} + i {\\Gamma}_{0} {\\gamma}_{0,0}^{2} {m}_{0} - {m}_{0}^{2}}$" - ], - "text/plain": [ - "-Gamma[0]*gamma[0, 0]**2*m[0]/(m**2 + I*Gamma[0]*gamma[0, 0]**2*m[0] - m[0]**2)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "k_matrix(n_resonances=1, n_channels=1)[0, 0].doit().simplify()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Single channel, $n_R$ resonances" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\frac{\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{- i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1}$" - ], - "text/plain": [ - "Sum(Gamma[R]*gamma[R, 0]**2*m[R]/(-m**2 + m[R]**2), (R, 0, n_R - 1))/(-I*Sum(Gamma[R]*gamma[R, 0]**2*m[R]/(-m**2 + m[R]**2), (R, 0, n_R - 1)) + 1)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "k_matrix(n_resonances=sp.Symbol(\"n_R\"), n_channels=1)[0, 0]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two channels, one resonance (Flatté function):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle - \\frac{{\\Gamma}_{0} {\\gamma}_{0,0}^{2} {m}_{0}}{m^{2} + i {\\Gamma}_{0} {\\gamma}_{0,0}^{2} {m}_{0} + i {\\Gamma}_{0} {\\gamma}_{0,1}^{2} {m}_{0} - {m}_{0}^{2}}$" - ], - "text/plain": [ - "-Gamma[0]*gamma[0, 0]**2*m[0]/(m**2 + I*Gamma[0]*gamma[0, 0]**2*m[0] + I*Gamma[0]*gamma[0, 1]**2*m[0] - m[0]**2)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "k_matrix(n_resonances=1, n_channels=2)[0, 0].doit().simplify()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "keep_output" - ] - }, - "source": [ - "Two channels, $n_R$ resonances:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "\\mathtt{\\text{}} = & \\frac{\\left(i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - 1\\right) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{\\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0} {\\gamma}_{R,1} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right)^{2} - 1} \\\\\n", - "& + \\frac{i \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0} {\\gamma}_{R,1} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right)^{2}}{- \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - i \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma}_{R} {\\gamma}_{R,0} {\\gamma}_{R,1} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right)^{2} + 1} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "expr = k_matrix(n_resonances=sp.Symbol(\"n_R\"), n_channels=2)[0, 0]\n", - "Math(sp.multiline_latex(\"\", expr))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, let's use [`matplotlib`](https://matplotlib.org), [`mpl_interactions`](https://mpl-interactions.rtfd.io), and {mod}`symplot` to visualize the $\\boldsymbol{K}$-matrix for arbitrary $n$ and $n_R$.\n", - "\n", - ":::{margin}\n", - "\n", - "[TR-008](008.ipynb) explains the need for {func}`symplot.substitute_indexed_symbols`.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def plot_k_matrix(\n", - " n_channels: int,\n", - " n_resonances: int,\n", - " title: str = \"\",\n", - ") -> None:\n", - " # Convert to Symbol: symplot cannot handle IndexedBase\n", - " i = sp.Symbol(\"i\", integer=True, negative=False)\n", - " expr = k_matrix(n_resonances, n_channels)[i, i].doit()\n", - " expr = symplot.substitute_indexed_symbols(expr)\n", - " np_expr, sliders = symplot.prepare_sliders(expr, plot_symbol=m)\n", - " symbol_to_arg = {symbol: arg for arg, symbol in sliders._arg_to_symbol.items()}\n", - "\n", - " # Set plot domain\n", - " x_min, x_max = 1e-3, 3\n", - " y_min, y_max = -0.5, +0.5\n", - "\n", - " plot_domain = np.linspace(x_min, x_max, num=500)\n", - " x_values = np.linspace(x_min, x_max, num=160)\n", - " y_values = np.linspace(y_min, y_max, num=80)\n", - " X, Y = np.meshgrid(x_values, y_values)\n", - " plot_domain_complex = X + Y * 1j\n", - "\n", - " # Set slider values and ranges\n", - " m0_values = np.linspace(x_min, x_max, num=n_resonances + 2)\n", - " m0_values = m0_values[1:-1]\n", - " for R in range(n_resonances):\n", - " for i in range(n_channels):\n", - " sliders.set_ranges({\n", - " \"i\": (0, n_channels - 1),\n", - " f\"m{R}\": (0, 3, 100),\n", - " f\"Gamma{R}\": (-1, 1, 100),\n", - " Rf\"\\gamma_{{{R},{i}}}\": (0, 2, 100),\n", - " })\n", - " sliders.set_values({\n", - " f\"m{R}\": m0_values[R],\n", - " f\"Gamma{R}\": (R + 1) * 0.1,\n", - " Rf\"\\gamma_{{{R},{i}}}\": 1 - 0.1 * R + 0.1 * i,\n", - " })\n", - "\n", - " # Create interactive plots\n", - " controls = Controls(**sliders)\n", - " fig, (ax_2d, ax_3d) = plt.subplots(\n", - " nrows=2,\n", - " figsize=(8, 6),\n", - " sharex=True,\n", - " tight_layout=True,\n", - " )\n", - "\n", - " fig.canvas.toolbar_visible = False\n", - " fig.canvas.header_visible = False\n", - " fig.canvas.footer_visible = False\n", - " if not title:\n", - " title = (\n", - " Rf\"${n_channels} \\times {n_channels}$ $K$-matrix\"\n", - " f\" with {n_resonances} resonances\"\n", - " )\n", - " fig.suptitle(title)\n", - "\n", - " ax_2d.set_ylabel(\"$|T|^{2}$\")\n", - " ax_2d.set_yticks([])\n", - " ax_3d.set_xlabel(\"Re $m$\")\n", - " ax_3d.set_ylabel(\"Im $m$\")\n", - " ax_3d.set_xticks([])\n", - " ax_3d.set_yticks([])\n", - " ax_3d.set_facecolor(\"white\")\n", - "\n", - " ax_3d.axhline(0, linewidth=0.5, c=\"black\", linestyle=\"dotted\")\n", - "\n", - " # 2D plot\n", - " def plot(channel: int):\n", - " def wrapped(*args, **kwargs) -> sp.Expr:\n", - " kwargs[\"i\"] = channel\n", - " return np.abs(np_expr(*args, **kwargs)) ** 2\n", - "\n", - " return wrapped\n", - "\n", - " for i in range(n_channels):\n", - " iplt.plot(\n", - " plot_domain,\n", - " plot(i),\n", - " ax=ax_2d,\n", - " controls=controls,\n", - " ylim=\"auto\",\n", - " label=f\"channel {i}\",\n", - " )\n", - " if n_channels > 1:\n", - " ax_2d.legend(loc=\"upper right\")\n", - " mass_line_style = {\n", - " \"c\": \"red\",\n", - " \"alpha\": 0.3,\n", - " }\n", - " for name in controls.params:\n", - " if not name.startswith(\"m\"):\n", - " continue\n", - " iplt.axvline(controls[name], ax=ax_2d, **mass_line_style)\n", - "\n", - " # 3D plot\n", - " color_mesh = None\n", - " resonances_indicators = []\n", - "\n", - " def plot3(*, z_cutoff, complex_rendering, **kwargs):\n", - " nonlocal color_mesh\n", - " Z = np_expr(plot_domain_complex, **kwargs)\n", - " if complex_rendering == \"imag\":\n", - " Z_values = Z.imag\n", - " ax_title = \"Re $T$\"\n", - " elif complex_rendering == \"real\":\n", - " Z_values = Z.real\n", - " ax_title = \"Im $T$\"\n", - " elif complex_rendering == \"abs\":\n", - " Z_values = np.abs(Z)\n", - " ax_title = \"$|T|$\"\n", - " else:\n", - " raise NotImplementedError\n", - "\n", - " if n_channels == 1:\n", - " ax_3d.set_title(ax_title)\n", - " else:\n", - " i = kwargs[\"i\"]\n", - " ax_3d.set_title(f\"{ax_title}, channel {i}\")\n", - "\n", - " if color_mesh is None:\n", - " color_mesh = ax_3d.pcolormesh(X, Y, Z_values, cmap=cm.coolwarm)\n", - " else:\n", - " color_mesh.set_array(Z_values)\n", - " color_mesh.set_clim(vmin=-z_cutoff, vmax=+z_cutoff)\n", - "\n", - " if resonances_indicators:\n", - " for R, (line, text) in enumerate(resonances_indicators):\n", - " mass = kwargs[f\"m{R}\"]\n", - " line.set_xdata(mass)\n", - " text.set_x(mass + (x_max - x_min) * 0.008)\n", - " else:\n", - " for R in range(n_resonances):\n", - " mass = kwargs[f\"m{R}\"]\n", - " resonances_indicators.append(\n", - " (\n", - " ax_3d.axvline(mass, **mass_line_style),\n", - " ax_3d.text(\n", - " x=mass + (x_max - x_min) * 0.008,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{R}$\",\n", - " c=\"red\",\n", - " ),\n", - " ),\n", - " )\n", - "\n", - " # Create switch for imag/real/abs\n", - " name = \"complex_rendering\"\n", - " sliders._sliders[name] = ipywidgets.RadioButtons(\n", - " options=[\"imag\", \"real\", \"abs\"],\n", - " description=R\"\\(s\\)-plane plot\",\n", - " )\n", - " sliders._arg_to_symbol[name] = name\n", - "\n", - " # Create cut-off slider for z-direction\n", - " name = \"z_cutoff\"\n", - " sliders._sliders[name] = ipywidgets.FloatSlider(\n", - " value=1.5,\n", - " min=0.01,\n", - " max=10,\n", - " step=0.1,\n", - " description=R\"\\(z\\)-cutoff\",\n", - " )\n", - " sliders._arg_to_symbol[name] = name\n", - "\n", - " # Create GUI\n", - " sliders_copy = dict(sliders)\n", - " h_boxes = []\n", - " for R in range(n_resonances):\n", - " buttons = [\n", - " sliders_copy.pop(f\"m{R}\"),\n", - " sliders_copy.pop(f\"Gamma{R}\"),\n", - " ]\n", - " if n_channels == 1:\n", - " dummy_name = symbol_to_arg[Rf\"\\gamma_{{{R},0}}\"]\n", - " buttons.append(sliders_copy.pop(dummy_name))\n", - " h_box = ipywidgets.HBox(buttons)\n", - " h_boxes.append(h_box)\n", - " remaining_sliders = sorted(sliders_copy.values(), key=lambda s: s.description)\n", - " if n_channels == 1:\n", - " remaining_sliders.remove(sliders[\"i\"])\n", - " ui = ipywidgets.VBox(h_boxes + remaining_sliders)\n", - " output = ipywidgets.interactive_output(plot3, controls=sliders)\n", - " display(ui, output)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "plot_k_matrix(n_resonances=3, n_channels=1)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "![record](https://user-images.githubusercontent.com/29308176/164994739-c1d128cd-2689-4849-8fa5-5c1ca7909f21.gif)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "plot_k_matrix(n_resonances=2, n_channels=2)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "![](https://user-images.githubusercontent.com/29308176/164994885-9bc96678-bfb2-4750-8368-7651610a7b4a.gif)" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/006.ipynb b/docs/report/006.ipynb deleted file mode 100644 index 6be62a10..00000000 --- a/docs/report/006.ipynb +++ /dev/null @@ -1,567 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "tips" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Interactive 3D plots\n", - "TR-006\n", - "^^^\n", - "This report illustrates how to interact with [`matplotlib`](https://matplotlib.org) 3D plots through [Matplotlib sliders](https://matplotlib.org/stable/api/widgets_api.html) and [ipywidgets](https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20List.html).\n", - "+++\n", - "✅ [ampform#38](https://github.com/ComPWA/ampform/pull/38)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Interactive 3D plots" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ipywidgets==8.1.1 plotly==5.18.0 sympy==1.12" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import os\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import plotly.graph_objects as go\n", - "import sympy as sp\n", - "from IPython.display import display\n", - "from ipywidgets import widgets as ipywidgets\n", - "from matplotlib import cm\n", - "from matplotlib import widgets as mpl_widgets\n", - "\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This report illustrates how to interact with [`matplotlib`](https://matplotlib.org) 3D plots through [Matplotlib sliders](https://matplotlib.org/stable/api/widgets_api.html) and [ipywidgets](https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20List.html). This might be implemented later on in {mod}`symplot` and/or [`mpl_interactions`](https://mpl-interactions.readthedocs.io) (see [ianhi/mpl-interactions#89](https://github.com/ianhi/mpl-interactions/issues/89)).\n", - "\n", - "In this example, we create a surface plot (see {meth}`~mpl_toolkits.mplot3d.axes3d.Axes3D.plot_surface`) for the following function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sqrt{x^{a} + \\sin^{2}{\\left(\\frac{y}{b} \\right)}}$" - ], - "text/plain": [ - "sqrt(x**a + sin(y/b)**2)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "x, y, a, b = sp.symbols(\"x y a b\")\n", - "expression = sp.sqrt(x**a + sp.sin(y / b) ** 2)\n", - "expression" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The function is formulated with {mod}`sympy`, but we use {func}`~sympy.utilities.lambdify.lambdify` to express it as a {mod}`numpy` function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "numpy_function = sp.lambdify(\n", - " args=(x, y, a, b),\n", - " expr=expression,\n", - " modules=\"numpy\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A surface plot has to be generated over a {func}`numpy.meshgrid`. This defines the $xy$-plane over which we want to plot our function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "x_min, x_max = 0.1, 2\n", - "y_min, y_max = -50, +50\n", - "x_values = np.linspace(x_min, x_max, num=20)\n", - "y_values = np.linspace(y_min, y_max, num=40)\n", - "X, Y = np.meshgrid(x_values, y_values)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The $z$-values for {meth}`~mpl_toolkits.mplot3d.axes3d.Axes3D.plot_surface` can now be simply computed as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "a_init = -0.5\n", - "b_init = 20\n", - "Z = numpy_function(X, Y, a=a_init, b=b_init)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now want to create sliders for $a$ and $b$, so that we can live-update the surface plot through those sliders." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Matplotlib widgets" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Matplotlib provides its own way to define {mod}`matplotlib.widgets`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "fig1, ax1 = plt.subplots(ncols=1, subplot_kw={\"projection\": \"3d\"})\n", - "\n", - "# Create sliders and insert them within the figure\n", - "plt.subplots_adjust(bottom=0.25)\n", - "a_slider = mpl_widgets.Slider(\n", - " ax=plt.axes([0.2, 0.1, 0.65, 0.03]),\n", - " label=f\"${sp.latex(a)}$\",\n", - " valmin=-2,\n", - " valmax=2,\n", - " valinit=a_init,\n", - ")\n", - "b_slider = mpl_widgets.Slider(\n", - " ax=plt.axes([0.2, 0.05, 0.65, 0.03]),\n", - " label=f\"${sp.latex(b)}$\",\n", - " valmin=10,\n", - " valmax=50,\n", - " valinit=b_init,\n", - " valstep=1,\n", - ")\n", - "\n", - "\n", - "# Define what to do when a slider changes\n", - "def update_plot(val=None):\n", - " a = a_slider.val\n", - " b = b_slider.val\n", - " ax1.clear()\n", - " Z = numpy_function(X, Y, a, b)\n", - " ax1.plot_surface(\n", - " X,\n", - " Y,\n", - " Z,\n", - " rstride=3,\n", - " cstride=1,\n", - " cmap=cm.coolwarm,\n", - " antialiased=False,\n", - " )\n", - " ax1.set_xlabel(f\"${sp.latex(x)}$\")\n", - " ax1.set_ylabel(f\"${sp.latex(y)}$\")\n", - " ax1.set_zlabel(f\"${sp.latex(expression)}$\")\n", - " ax1.set_xticks([])\n", - " ax1.set_yticks([])\n", - " ax1.set_zticks([])\n", - " ax1.set_facecolor(\"white\")\n", - " fig1.canvas.draw_idle()\n", - "\n", - "\n", - "a_slider.on_changed(update_plot)\n", - "b_slider.on_changed(update_plot)\n", - "\n", - "# Plot the surface as initialization\n", - "update_plot()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "![Interactive inline matplotlib output](https://user-images.githubusercontent.com/29308176/164993434-da965bbb-459d-43b5-8294-eb64475f5192.gif)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## `ipywidgets`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As an alternative, you can use [ipywidgets](https://ipywidgets.readthedocs.io/en/latest/examples/Widget%20List.html). This package has lot more sliders to offer than Matplotlib, and they look nicer, but it only work within a Jupyter notebook.\n", - "\n", - "For more info, see [Using Interact](https://ipywidgets.readthedocs.io/en/stable/examples/Using%20Interact.html)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Using `interact`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Simplest option is to use the [`ipywidgets.interact()`](https://ipywidgets.readthedocs.io/en/stable/examples/Using%20Interact.html) function:\n", - "\n", - "{{ run_interactive }}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "fig2, ax2 = plt.subplots(ncols=1, subplot_kw={\"projection\": \"3d\"})\n", - "\n", - "\n", - "@ipywidgets.interact(a=(-2.0, 2.0), b=(10, 50))\n", - "def plot2(a=a_init, b=b_init):\n", - " ax2.clear()\n", - " Z = numpy_function(X, Y, a, b)\n", - " ax2.plot_surface(\n", - " X,\n", - " Y,\n", - " Z,\n", - " rstride=3,\n", - " cstride=1,\n", - " cmap=cm.coolwarm,\n", - " antialiased=False,\n", - " )\n", - " ax2.set_xlabel(f\"${sp.latex(x)}$\")\n", - " ax2.set_ylabel(f\"${sp.latex(y)}$\")\n", - " ax2.set_zlabel(f\"${sp.latex(expression)}$\")\n", - " ax2.set_xticks([])\n", - " ax2.set_yticks([])\n", - " ax2.set_zticks([])\n", - " ax2.set_facecolor(\"white\")\n", - " fig2.canvas.draw_idle()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Using `interactive_output`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "You can have more control with [`ipywidgets.interactive_output()`](https://ipywidgets.readthedocs.io/en/stable/examples/Using%20Interact.html). That allows defining the sliders independently, so that you can arrange them as a user interface:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "fig3, ax3 = plt.subplots(ncols=1, subplot_kw={\"projection\": \"3d\"})\n", - "a_ipyslider = ipywidgets.FloatSlider(\n", - " description=f\"${sp.latex(a)}$\",\n", - " value=a_init,\n", - " min=-2,\n", - " max=2,\n", - " step=0.1,\n", - " readout_format=\".1f\",\n", - ")\n", - "b_ipyslider = ipywidgets.IntSlider(\n", - " description=f\"${sp.latex(b)}$\",\n", - " value=b_init,\n", - " min=10,\n", - " max=50,\n", - ")\n", - "\n", - "\n", - "def plot3(a=a_init, b=b_init):\n", - " ax3.clear()\n", - " Z = numpy_function(X, Y, a, b)\n", - " ax3.plot_surface(\n", - " X,\n", - " Y,\n", - " Z,\n", - " rstride=3,\n", - " cstride=1,\n", - " cmap=cm.coolwarm,\n", - " antialiased=False,\n", - " )\n", - " ax3.set_xlabel(f\"${sp.latex(x)}$\")\n", - " ax3.set_ylabel(f\"${sp.latex(y)}$\")\n", - " ax3.set_zlabel(f\"${sp.latex(expression)}$\")\n", - " ax3.set_xticks([])\n", - " ax3.set_yticks([])\n", - " ax3.set_zticks([])\n", - " ax3.set_facecolor(\"white\")\n", - " fig3.canvas.draw_idle()\n", - "\n", - "\n", - "ui = ipywidgets.HBox([a_ipyslider, b_ipyslider])\n", - "output = ipywidgets.interactive_output(\n", - " plot3, controls={\"a\": a_ipyslider, \"b\": b_ipyslider}\n", - ")\n", - "display(ui, output)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "![ipywidgets interactive output with interactive_output()](https://user-images.githubusercontent.com/29308176/164993430-6f6b906a-dfb5-4c7c-bae5-d9951c02112b.gif)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Plotly with ipywidgets" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "3D plots with [Plotly](https://plotly.com/python) look a lot nicer and make it possible for the user to pan and zoom the 3D object. As an added bonus, Plotly figures [render as interactive 3D objects](https://myst-nb.readthedocs.io/en/v0.17.2/render/interactive.html#plotly) in the static HTML Sphinx build.\n", - "\n", - "Making 3D Plotly plots interactive with {mod}`ipywidgets` is quite similar to the previous examples with {mod}`matplotlib`. Two recommendations are:\n", - "\n", - "1. Set `continuous_update=False`, because {mod}`plotly` is slower than {mod}`matplotlib` in updating the figure.\n", - "2. Save the camera orientation and update it after calling `Figure.show()`.\n", - "3. When embedding the notebook a static webpage with [MyST-NB](https://myst-nb.readthedocs.io), avoid calling `Figure.show()` through [`ipywidgets.interactive_output()`](https://ipywidgets.readthedocs.io/en/stable/examples/Using%20Interact.html), because it causes the notebook to hang in some cycle (see CI for [ComPWA/compwa.github.io@d9240f1](https://github.com/ComPWA/compwa.github.io/pull/208/commits/d9240f1)). In the example below, the `update_plotly()` function is aborted if the notebook is run through Sphinx. " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "plotly_a = ipywidgets.FloatSlider(\n", - " description=f\"${sp.latex(a)}$\",\n", - " value=a_init,\n", - " min=-2,\n", - " max=2,\n", - " step=0.1,\n", - " continuous_update=False,\n", - " readout_format=\".1f\",\n", - ")\n", - "plotly_b = ipywidgets.IntSlider(\n", - " description=f\"${sp.latex(b)}$\",\n", - " value=b_init,\n", - " min=10,\n", - " max=50,\n", - " continuous_update=False,\n", - ")\n", - "plotly_controls = {\"a\": plotly_a, \"b\": plotly_b}\n", - "\n", - "plotly_surface = go.Surface(\n", - " x=X,\n", - " y=Y,\n", - " z=Z,\n", - " surfacecolor=Z,\n", - " colorscale=\"RdBu_r\",\n", - " name=\"Surface\",\n", - ")\n", - "plotly_fig = go.Figure(data=[plotly_surface])\n", - "plotly_fig.update_layout(height=500)\n", - "if STATIC_WEB_PAGE:\n", - " plotly_fig.show()\n", - "\n", - "\n", - "def update_plotly(a, b):\n", - " if STATIC_WEB_PAGE:\n", - " return\n", - " Z = numpy_function(X, Y, a, b)\n", - " camera_orientation = plotly_fig.layout.scene.camera\n", - " plotly_fig.update_traces(\n", - " x=X,\n", - " y=Y,\n", - " z=Z,\n", - " surfacecolor=Z,\n", - " selector=dict(name=\"Surface\"),\n", - " )\n", - " plotly_fig.show()\n", - " plotly_fig.update_layout(scene=dict(camera=camera_orientation))\n", - "\n", - "\n", - "plotly_ui = ipywidgets.HBox([plotly_a, plotly_b])\n", - "plotly_output = ipywidgets.interactive_output(update_plotly, plotly_controls)\n", - "display(plotly_ui, plotly_output)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso} {doc}`/report/023`\n", - ":::" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/007.ipynb b/docs/report/007.ipynb deleted file mode 100644 index 6748a966..00000000 --- a/docs/report/007.ipynb +++ /dev/null @@ -1,194 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} `MatrixSymbols`\n", - "TR-007\n", - "^^^\n", - "This report is a sequel to [TR-005](005.ipynb). In that report, the $\\boldsymbol{K}$ was constructed with a {obj}`sympy.Matrix `, but it might be more elegant to work with {class}`~sympy.matrices.expressions.MatrixSymbol`s.\n", - "+++\n", - "WIP\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# `MatrixSymbol`s" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q sympy==1.9" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import sympy as sp\n", - "from IPython.display import display" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here are some examples of computations with {class}`~sympy.matrices.expressions.MatrixSymbol`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "N = sp.Symbol(\"n\", integer=True, positive=True)\n", - "i, j = sp.symbols(\"i j\", integer=True, negative=False)\n", - "K = sp.MatrixSymbol(\"K\", N, N)\n", - "display(K, K[i, j], K[0, 0])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "A = sp.MatrixSymbol(\"A\", N, N)\n", - "(A * K)[0, 0]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The important thing is that elements of a {class}`~sympy.matrices.expressions.MatrixSymbol` can be substituted:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "K[0, 0].subs(K[0, 0], i)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "A * K" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "(A * K)[0, 0]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now make the matrices $2 \\times 2$ by specifying $n$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "A_n2 = A.subs(N, 2)\n", - "K_n2 = K.subs(N, 2)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "(A_n2 * K_n2)[0, 0]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "v, w, x, y = sp.symbols(\"v, w, x, y\", real=True)\n", - "substitutions = {\n", - " A_n2[0, 0]: v,\n", - " A_n2[0, 1]: w,\n", - " K_n2[0, 0]: x,\n", - " K_n2[1, 0]: y,\n", - "}\n", - "(A_n2 * K_n2)[0, 0].subs(substitutions)" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/008.ipynb b/docs/report/008.ipynb deleted file mode 100644 index 64f51b90..00000000 --- a/docs/report/008.ipynb +++ /dev/null @@ -1,271 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Indexed free symbols\n", - "TR-008\n", - "^^^\n", - "This report has been implemented in [ampform#111](https://github.com/ComPWA/ampform/issues/111). Additionally, [tensorwaves#427](https://github.com/ComPWA/tensorwaves/issues/427) makes it possible to lambdify {class}`sympy.Expr ` with {class}`~sympy.tensor.indexed.Indexed` symbols directly.\n", - "+++\n", - "✅ [ampform#111](https://github.com/ComPWA/ampform/issues/111)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# `Indexed` free symbols" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q sympy==1.8" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In [TR-005](005.ipynb), we made use of {doc}`indexed symbols ` to create a $\\boldsymbol{K}$-matrix. The problem with that approach is that {class}`~sympy.tensor.indexed.IndexedBase` and their resulting {class}`~sympy.tensor.indexed.Indexed` instances when taking indices behave strangely in an expression tree.\n", - "\n", - "The following {class}`~sympy.core.expr.Expr` uses a {class}`~sympy.core.symbol.Symbol` and a elements in {class}`~sympy.tensor.indexed.IndexedBase`s (an {class}`~sympy.tensor.indexed.Indexed` instance):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle x {\\alpha}_{2} + {c}_{0,1}$" - ], - "text/plain": [ - "x*alpha[2] + c[0, 1]" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import sympy as sp\n", - "\n", - "x = sp.Symbol(\"x\")\n", - "c = sp.IndexedBase(\"c\")\n", - "alpha = sp.IndexedBase(\"alpha\")\n", - "expression = c[0, 1] + alpha[2] * x\n", - "expression" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Although seemingly there are just **three** {attr}`~sympy.core.basic.Basic.free_symbols`, there are actually **five**:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{alpha, alpha[2], c, c[0, 1], x}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "expression.free_symbols" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This becomes problematic when using {func}`~sympy.utilities.lambdify.lambdify`, particularly through {func}`symplot.prepare_sliders`.\n", - "\n", - "In addition, while `c[0, 1]` and `alpha[2]` are {class}`~sympy.tensor.indexed.Indexed` as expected, `alpha` and `c` are {class}`~sympy.core.symbol.Symbol`s, not {class}`~sympy.tensor.indexed.IndexedBase`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "{c: sympy.core.symbol.Symbol,\n", - " alpha[2]: sympy.tensor.indexed.Indexed,\n", - " x: sympy.core.symbol.Symbol,\n", - " c[0, 1]: sympy.tensor.indexed.Indexed,\n", - " alpha: sympy.core.symbol.Symbol}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "{s: type(s) for s in expression.free_symbols}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The {doc}`expression tree ` partially explains this behavior:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "keep_output": false - }, - "outputs": [], - "source": [ - "import graphviz\n", - "\n", - "dot = sp.dotprint(expression)\n", - "graphviz.Source(dot);" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164993648-13c6b74a-b85f-4492-aaf2-c64cdc30e345.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We would like to collapse the nodes under `c[0, 1]` and `alpha[2]` to two single {class}`~sympy.core.symbol.Symbol` nodes that are **still nicely rendered as $c_{0,1}$ and $\\alpha_2$**. The following function does that and converts the `[]` into subscripts. It does that in such a way that the name of the {class}`~sympy.core.symbol.Symbol` remains as short as possible, that is, short enough that it still renders nicely as LaTeX:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from sympy.printing.latex import translate\n", - "\n", - "\n", - "def to_symbol(idx: sp.Indexed) -> sp.Symbol:\n", - " base_name, _, _ = str(idx).rpartition(\"[\")\n", - " subscript = \",\".join(map(str, idx.indices))\n", - " if len(idx.indices) > 1:\n", - " base_name = translate(base_name)\n", - " subscript = \"_{\" + subscript + \"}\"\n", - " return sp.Symbol(f\"{base_name}{subscript}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Next, we use {meth}`~sympy.core.basic.Basic.subs` to substitute the nodes `c[0, 1]` and `alpha[2]` with these {class}`~sympy.core.symbol.Symbol`s:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def replace_indexed_symbols(expression: sp.Expr) -> sp.Expr:\n", - " return expression.subs({\n", - " s: to_symbol(s) for s in expression.free_symbols if isinstance(s, sp.Indexed)\n", - " })" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And indeed, the expression tree has been simplified correctly!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "new_expression = replace_indexed_symbols(expression)\n", - "dot = sp.dotprint(new_expression)\n", - "graphviz.Source(dot);" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164993649-47231cf6-0ee2-4eed-a122-633e2cf5db1a.svg)" - ] - } - ], - "metadata": { - "colab": { - "toc_visible": true - }, - "keep_output": 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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/009.ipynb b/docs/report/009.ipynb deleted file mode 100644 index 53b6cabc..00000000 --- a/docs/report/009.ipynb +++ /dev/null @@ -1,1030 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "K-matrix", - "dynamics", - "physics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Symbolic expressions for Lorentz-invariant K-matrix\n", - "TR-009\n", - "^^^\n", - "This report is a sequel to [TR-005](005.ipynb).\n", - "+++\n", - "✅ [ampform#120](https://github.com/ComPWA/ampform/issues/120)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Lorentz-invariant K-matrix" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.10.5 sympy==1.8" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Physics" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "source": [ - "The **Lorentz-invariant description $\\boldsymbol{\\hat{T}}$ of the $\\boldsymbol{T}$-matrix** is:\n", - "\n", - "$$\n", - "\\boldsymbol{T} = \\sqrt{\\boldsymbol{\\rho^\\dagger}} \\, \\boldsymbol{\\hat{T}} \\sqrt{\\boldsymbol{\\rho}}\n", - "$$ (covariant-T-matrix)\n", - "\n", - "with the phase space factor matrix $\\boldsymbol{\\rho}$ defined as:\n", - "\n", - "$$\n", - "\\sqrt{\\boldsymbol{\\rho}} = \\begin{pmatrix}\n", - "\\rho_0 & \\cdots & 0 \\\\\n", - "\\vdots & \\ddots & \\vdots \\\\\n", - "0 & \\cdots & \\rho_{n-1}\n", - "\\end{pmatrix}\n", - "$$ (rho matrix)\n", - "\n", - "and\n", - "\n", - "$$\n", - "\\rho_i = \\frac{2q_i}{m} = \\sqrt{\n", - " \\left[1-\\left(\\frac{m_{i,a}+m_{i,b}}{m}\\right)^2\\right]\n", - " \\left[1-\\left(\\frac{m_{i,a}-m_{i,b}}{m}\\right)^2\\right]\n", - "}\n", - "$$ (phase space factor)\n", - "\n", - "This results in a similar transformation for the $\\boldsymbol{K}$-matrix\n", - "\n", - "$$\n", - "\\boldsymbol{K} = \\sqrt{\\boldsymbol{\\rho^\\dagger}} \\; \\boldsymbol{\\hat{K}} \\sqrt{\\boldsymbol{\\rho}}\n", - "$$ (covariant-K-matrix)\n", - "\n", - "with (compare Eq. {eq}`T-in-terms-of-K` in [TR-005](005.ipynb)):\n", - "\n", - "$$\n", - "\\boldsymbol{\\hat{T}} = \\boldsymbol{\\hat{K}}(\\boldsymbol{I} - i\\boldsymbol{\\rho}\\boldsymbol{\\hat{K}})^{-1}\n", - "$$ (covariant-T-matrix as K)\n", - "\n", - "It's common to integrate these phase space factors into the parametrization of $K_{ij}$ as well:\n", - "\n", - "$$\n", - "K_{ij} = \\sum_R \\frac{g_{R,i}(m)g_{R,j}(m)}{\\left(m_R^2-m^2\\right)\\sqrt{\\rho_i\\rho_j}}\n", - "$$ (covariant parametrization)\n", - "\n", - "Compare this with Eq. {eq}`K-matrix-parametrization` in [TR-005](005.ipynb).\n", - "\n", - "In addition, one often uses an \"energy dependent\" [`coupled_width()`](https://ampform.readthedocs.io/en/0.10.5/api/ampform.dynamics.html#ampform.dynamics.coupled_width) $\\Gamma_R(m)$ instead of a fixed width $\\Gamma_R$ as done in [TR-005](005.ipynb)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import os\n", - "import re\n", - "import warnings\n", - "from typing import TYPE_CHECKING\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import mpl_interactions.ipyplot as iplt\n", - "import numpy as np\n", - "import symplot\n", - "import sympy as sp\n", - "from ampform.dynamics import coupled_width, phase_space_factor_complex\n", - "from ampform.dynamics.decorator import (\n", - " UnevaluatedExpression,\n", - " create_expression,\n", - " implement_doit_method,\n", - ")\n", - "from IPython.display import Math, display\n", - "from ipywidgets import widgets as ipywidgets\n", - "from matplotlib import cm\n", - "from mpl_interactions.controller import Controls\n", - "\n", - "if TYPE_CHECKING:\n", - " from sympy.printing.latex import LatexPrinter\n", - "\n", - "warnings.filterwarnings(\"ignore\")\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Implementation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Wrapping expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To keep a nice rendering, we wrap the expressions for [`phase_space_factor()`](https://ampform.readthedocs.io/en/0.10.5/api/ampform.dynamics.html#ampform.dynamics.phase_space_factor) and [`coupled_width()`](https://ampform.readthedocs.io/en/0.10.5/api/ampform.dynamics.html#ampform.dynamics.coupled_width) into a class that derives from {class}`~sympy.core.expr.Expr` (see e.g. the implementation of {class}`~ampform.dynamics.BlattWeisskopfSquared`). Note that we need to use {func}`~symplot.partial_doit` to keep these expression symbols after evaluating the {class}`~sympy.concrete.summations.Sum`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@implement_doit_method()\n", - "class PhaseSpaceFactor(UnevaluatedExpression):\n", - " is_commutative = True\n", - "\n", - " def __new__(\n", - " cls,\n", - " s: sp.Symbol,\n", - " m_a: sp.Symbol,\n", - " m_b: sp.Symbol,\n", - " i: int,\n", - " **hints,\n", - " ) -> PhaseSpaceFactor:\n", - " return create_expression(cls, s, m_a, m_b, i, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m_a, m_b, *_ = self.args\n", - " return phase_space_factor_complex(s, m_a, m_b)\n", - "\n", - " def _latex(self, printer: LatexPrinter, *args) -> str:\n", - " s = printer._print(self.args[0])\n", - " i = self.args[-1]\n", - " return Rf\"\\rho_{{{i}}}({s})\"\n", - "\n", - "\n", - "@implement_doit_method()\n", - "class CoupledWidth(UnevaluatedExpression):\n", - " is_commutative = True\n", - "\n", - " def __new__(\n", - " cls,\n", - " s: sp.Symbol,\n", - " mass0: sp.IndexedBase,\n", - " gamma0: sp.IndexedBase,\n", - " m_a: sp.IndexedBase,\n", - " m_b: sp.IndexedBase,\n", - " angular_momentum: int,\n", - " R: int | sp.Symbol,\n", - " i: int,\n", - " **hints,\n", - " ) -> CoupledWidth:\n", - " return create_expression(\n", - " cls, s, mass0, gamma0, m_a, m_b, angular_momentum, R, i, **hints\n", - " )\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, mass0, gamma0, m_a, m_b, angular_momentum, R, i = self.args\n", - "\n", - " def phsp_factor(s, m_a, m_b):\n", - " return PhaseSpaceFactor(s, m_a, m_b, i)\n", - "\n", - " return coupled_width(\n", - " s,\n", - " mass0[R],\n", - " gamma0[R, i],\n", - " m_a[i],\n", - " m_b[i],\n", - " angular_momentum=angular_momentum,\n", - " meson_radius=1,\n", - " phsp_factor=phsp_factor,\n", - " )\n", - "\n", - " def _latex(self, printer: LatexPrinter, *args) -> str:\n", - " s = printer._print(self.args[0])\n", - " R = self.args[-2]\n", - " i = self.args[-1]\n", - " return Rf\"{{\\Gamma_{{{R},{i}}}}}({s})\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And here is what the equations look like:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "n_channels = 2\n", - "n_resonances, i, R, L = sp.symbols(\"n_R, i, R, L\", integer=True, negative=False)\n", - "m = sp.Symbol(\"m\", real=True)\n", - "M = sp.IndexedBase(\"m\", shape=(n_resonances,))\n", - "Gamma = sp.IndexedBase(\"Gamma\", shape=(n_resonances, n_channels))\n", - "gamma = sp.IndexedBase(\"gamma\", shape=(n_resonances, n_channels))\n", - "m_a = sp.IndexedBase(\"m_a\", shape=(n_channels,))\n", - "m_b = sp.IndexedBase(\"m_b\", shape=(n_channels,))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "width_expr = CoupledWidth(m**2, M, Gamma, m_a, m_b, 0, R, i)\n", - "phsp_expr = PhaseSpaceFactor(m**2, m_a[i], m_b[i], i)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "{\\Gamma_{R,i}}(m^{2}) = & \\frac{B_{0}^2\\left(\\frac{\\left(m^{2} - \\left({m_{a}}_{i} - {m_{b}}_{i}\\right)^{2}\\right) \\left(m^{2} - \\left({m_{a}}_{i} + {m_{b}}_{i}\\right)^{2}\\right)}{4 m^{2}}\\right) {\\Gamma}_{R,i} \\rho_{i}(m^{2})}{B_{0}^2\\left(\\frac{\\left(- \\left({m_{a}}_{i} - {m_{b}}_{i}\\right)^{2} + {m}_{R}^{2}\\right) \\left(- \\left({m_{a}}_{i} + {m_{b}}_{i}\\right)^{2} + {m}_{R}^{2}\\right)}{4 {m}_{R}^{2}}\\right) \\rho_{i}({m}_{R}^{2})} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "Math(\n", - " sp.multiline_latex(\n", - " lhs=width_expr,\n", - " rhs=width_expr.evaluate(),\n", - " )\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "\\rho_{i}(m^{2}) = & \\frac{\\sqrt[\\mathrm{c}]{\\frac{\\left(m^{2} - \\left({m_{a}}_{i} - {m_{b}}_{i}\\right)^{2}\\right) \\left(m^{2} - \\left({m_{a}}_{i} + {m_{b}}_{i}\\right)^{2}\\right)}{4 m^{2}}}}{8 \\pi m} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "Math(\n", - " sp.multiline_latex(\n", - " lhs=phsp_expr,\n", - " rhs=phsp_expr.doit().simplify().subs(sp.Abs(m), m),\n", - " )\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "\n", - "In `PhaseSpaceFactor`, we used {class}`~ampform.dynamics.phasespace.PhaseSpaceFactorComplex` instead of {class}`~ampform.dynamics.phasespace.PhaseSpaceFactor`, meaning that we choose the _positive_ square root when values under the square root are negative. The only reason for doing this is, so that there is output in the figure under {ref}`report/009:Visualization`. The choice for which square root to choose has to do with analyticity (see [TR-004](004.ipynb)) and choosing which Riemann sheet to connect to. This issue is ignored in this report.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generalization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The implementation is quite similar to {ref}`that of TR-005 `, with the only difference being additional $\\boldsymbol{\\rho}$-matrix and the insertion of coupled width. Don't forget to convert back to $\\boldsymbol{T}$ from $\\boldsymbol{\\hat{T}}$ with Eq. {eq}`covariant-T-matrix`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def Kij_relativistic(\n", - " m: sp.Symbol,\n", - " M: sp.IndexedBase,\n", - " Gamma: sp.IndexedBase,\n", - " gamma: sp.IndexedBase,\n", - " i: int,\n", - " j: int,\n", - " n_resonances: int | sp.Symbol,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - ") -> sp.Expr:\n", - " def residue_function(i):\n", - " return gamma[R, i] * sp.sqrt(\n", - " M[R] * CoupledWidth(m**2, M, Gamma, m_a, m_b, angular_momentum, R, i)\n", - " )\n", - "\n", - " g_i = residue_function(i)\n", - " g_j = residue_function(j)\n", - " parametrization = (g_i * g_j) / (M[R] ** 2 - m**2)\n", - " return sp.Sum(parametrization, (R, 0, n_resonances - 1))\n", - "\n", - "\n", - "def relativistic_k_matrix(\n", - " n_resonances: int,\n", - " n_channels: int,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - ") -> sp.Matrix:\n", - " # Define symbols\n", - " m = sp.Symbol(\"m\", real=True)\n", - " M = sp.IndexedBase(\"m\", shape=(n_resonances,))\n", - " Gamma = sp.IndexedBase(\"Gamma\", shape=(n_resonances, n_channels))\n", - " gamma = sp.IndexedBase(\"gamma\", shape=(n_resonances, n_channels))\n", - " m_a = sp.IndexedBase(\"m_a\", shape=(n_channels,))\n", - " m_b = sp.IndexedBase(\"m_b\", shape=(n_channels,))\n", - " # Define phase space matrix\n", - " sqrt_rho = sp.zeros(n_channels, n_channels)\n", - " sqrt_rho_dagger = sp.zeros(n_channels, n_channels)\n", - " for i in range(n_channels):\n", - " rho = PhaseSpaceFactor(m**2, m_a[i], m_b[i], i)\n", - " sqrt_rho[i, i] = sp.sqrt(rho)\n", - " sqrt_rho_dagger[i, i] = 1 / sp.conjugate(sqrt_rho[i, i])\n", - " # Define K-matrix and T-matrix\n", - " K = create_symbol_matrix(\"K\", n_channels)\n", - " T_hat = K * (sp.eye(n_channels) - sp.I * rho * K).inv()\n", - " T = sqrt_rho_dagger * T_hat * sqrt_rho\n", - " # Substitute elements\n", - " return T.subs({\n", - " K[i, j]: Kij_relativistic(\n", - " m=m,\n", - " M=M,\n", - " Gamma=Gamma,\n", - " gamma=gamma,\n", - " i=i,\n", - " j=j,\n", - " n_resonances=n_resonances,\n", - " angular_momentum=angular_momentum,\n", - " )\n", - " for i in range(n_channels)\n", - " for j in range(n_channels)\n", - " })\n", - "\n", - "\n", - "def create_symbol_matrix(name: str, n: int) -> sp.Matrix:\n", - " symbol = sp.IndexedBase(name, shape=(n, n))\n", - " return sp.Matrix([[symbol[i, j] for j in range(n)] for i in range(n)])" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "Single channel, one resonance (compare {func}`~ampform.dynamics.relativistic_breit_wigner_with_ff`):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "\\frac{\\sqrt{\\rho_{0}(m^{2})} \\sum_{R=0}^{0} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{\\left(- i \\rho_{0}(m^{2}) \\sum_{R=0}^{0} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1\\right) \\overline{\\sqrt{\\rho_{0}(m^{2})}}} = &- \\frac{{\\Gamma_{0,0}}(m^{2}) {\\gamma}_{0,0}^{2} {m}_{0} \\sqrt{\\rho_{0}(m^{2})}}{\\left(m^{2} + i {\\Gamma_{0,0}}(m^{2}) {\\gamma}_{0,0}^{2} {m}_{0} \\rho_{0}(m^{2}) - {m}_{0}^{2}\\right) \\overline{\\sqrt{\\rho_{0}(m^{2})}}} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "expr = relativistic_k_matrix(n_resonances=1, n_channels=1)[0, 0]\n", - "Math(\n", - " sp.multiline_latex(\n", - " lhs=expr,\n", - " rhs=symplot.partial_doit(expr, sp.Sum).simplify(doit=False),\n", - " )\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two channels, one resonance ('Flatté'):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\frac{{\\Gamma_{0,0}}(m^{2}) {\\gamma}_{0,0}^{2} {m}_{0} \\sqrt{\\rho_{0}(m^{2})}}{\\left(- m^{2} - i {\\Gamma_{0,0}}(m^{2}) {\\gamma}_{0,0}^{2} {m}_{0} \\rho_{1}(m^{2}) - i {\\Gamma_{0,1}}(m^{2}) {\\gamma}_{0,1}^{2} {m}_{0} \\rho_{1}(m^{2}) + {m}_{0}^{2}\\right) \\overline{\\sqrt{\\rho_{0}(m^{2})}}}$" - ], - "text/plain": [ - "CoupledWidth(m**2, m, Gamma, m_a, m_b, 0, 0, 0)*gamma[0, 0]**2*m[0]*sqrt(PhaseSpaceFactor(m**2, m_a[0], m_b[0], 0))/((-m**2 - I*CoupledWidth(m**2, m, Gamma, m_a, m_b, 0, 0, 0)*gamma[0, 0]**2*m[0]*PhaseSpaceFactor(m**2, m_a[1], m_b[1], 1) - I*CoupledWidth(m**2, m, Gamma, m_a, m_b, 0, 0, 1)*gamma[0, 1]**2*m[0]*PhaseSpaceFactor(m**2, m_a[1], m_b[1], 1) + m[0]**2)*conjugate(sqrt(PhaseSpaceFactor(m**2, m_a[0], m_b[0], 0))))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "expr = relativistic_k_matrix(n_resonances=1, n_channels=2)[0, 0]\n", - "symplot.partial_doit(expr, sp.Sum).simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Single channel, $n_R$ resonances:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\frac{\\sqrt{\\rho_{0}(m^{2})} \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{\\left(- i \\rho_{0}(m^{2}) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1\\right) \\overline{\\sqrt{\\rho_{0}(m^{2})}}}$" - ], - "text/plain": [ - "sqrt(PhaseSpaceFactor(m**2, m_a[0], m_b[0], 0))*Sum(CoupledWidth(m**2, m, Gamma, m_a, m_b, 0, R, 0)*gamma[R, 0]**2*m[R]/(-m**2 + m[R]**2), (R, 0, n_R - 1))/((-I*PhaseSpaceFactor(m**2, m_a[0], m_b[0], 0)*Sum(CoupledWidth(m**2, m, Gamma, m_a, m_b, 0, R, 0)*gamma[R, 0]**2*m[R]/(-m**2 + m[R]**2), (R, 0, n_R - 1)) + 1)*conjugate(sqrt(PhaseSpaceFactor(m**2, m_a[0], m_b[0], 0))))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "relativistic_k_matrix(n_resonances, n_channels=1)[0, 0]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Two channels, $n_R$ resonances:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "\\mathtt{\\text{}} = & \\frac{\\left(\\frac{\\left(- i \\rho_{1}(m^{2}) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,1}}(m^{2}) {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1\\right) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}}{- \\rho_{1}(m^{2})^{2} \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,1}}(m^{2}) {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + \\rho_{1}(m^{2})^{2} \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{\\sqrt{{\\Gamma_{R,0}}(m^{2}) {m}_{R}} \\sqrt{{\\Gamma_{R,1}}(m^{2}) {m}_{R}} {\\gamma}_{R,0} {\\gamma}_{R,1}}{- m^{2} + {m}_{R}^{2}}\\right)^{2} - i \\rho_{1}(m^{2}) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - i \\rho_{1}(m^{2}) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,1}}(m^{2}) {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1} + \\frac{i \\rho_{1}(m^{2}) \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{\\sqrt{{\\Gamma_{R,0}}(m^{2}) {m}_{R}} \\sqrt{{\\Gamma_{R,1}}(m^{2}) {m}_{R}} {\\gamma}_{R,0} {\\gamma}_{R,1}}{- m^{2} + {m}_{R}^{2}}\\right)^{2}}{- \\rho_{1}(m^{2})^{2} \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}}\\right) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,1}}(m^{2}) {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + \\rho_{1}(m^{2})^{2} \\left(\\sum_{R=0}^{n_{R} - 1} \\frac{\\sqrt{{\\Gamma_{R,0}}(m^{2}) {m}_{R}} \\sqrt{{\\Gamma_{R,1}}(m^{2}) {m}_{R}} {\\gamma}_{R,0} {\\gamma}_{R,1}}{- m^{2} + {m}_{R}^{2}}\\right)^{2} - i \\rho_{1}(m^{2}) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,0}}(m^{2}) {\\gamma}_{R,0}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} - i \\rho_{1}(m^{2}) \\sum_{R=0}^{n_{R} - 1} \\frac{{\\Gamma_{R,1}}(m^{2}) {\\gamma}_{R,1}^{2} {m}_{R}}{- m^{2} + {m}_{R}^{2}} + 1}\\right) \\sqrt{\\rho_{0}(m^{2})}}{\\overline{\\sqrt{\\rho_{0}(m^{2})}}} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "expr = relativistic_k_matrix(n_resonances, n_channels=2)[0, 0]\n", - "Math(sp.multiline_latex(\"\", expr))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{margin}\n", - "\n", - "[TR-008](008.ipynb) explains the need for {func}`symplot.substitute_indexed_symbols`.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def plot_relativistic_k_matrix(\n", - " n_channels: int,\n", - " n_resonances: int,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - " title: str = \"\",\n", - ") -> None:\n", - " # Convert to Symbol: symplot cannot handle IndexedBase\n", - " epsilon = sp.Symbol(\"epsilon\")\n", - " i, j = sp.symbols(\"i, j\", integer=True, negative=False)\n", - " j = i\n", - " expr = relativistic_k_matrix(\n", - " n_resonances, n_channels, angular_momentum=angular_momentum\n", - " ).doit()[i, j]\n", - " expr = symplot.substitute_indexed_symbols(expr)\n", - " expr = expr.subs(m, m + epsilon * sp.I)\n", - " np_expr, sliders = symplot.prepare_sliders(expr, m)\n", - " symbol_to_arg = {symbol: arg for arg, symbol in sliders._arg_to_symbol.items()}\n", - "\n", - " # Set plot domain\n", - " x_min, x_max = 1e-3, 3\n", - " y_min, y_max = -0.5, +0.5\n", - "\n", - " plot_domain = np.linspace(x_min, x_max, num=500)\n", - " x_values = np.linspace(x_min, x_max, num=160)\n", - " y_values = np.linspace(y_min, y_max, num=80)\n", - " X, Y = np.meshgrid(x_values, y_values)\n", - " plot_domain_complex = X + Y * 1j\n", - "\n", - " # Set slider values and ranges\n", - " m0_values = np.linspace(x_min, x_max, num=n_resonances + 2)\n", - " m0_values = m0_values[1:-1]\n", - " if \"L\" in sliders:\n", - " sliders.set_ranges(L=(0, 8))\n", - " for R in range(n_resonances):\n", - " for i in range(n_channels):\n", - " sliders.set_ranges({\n", - " \"i\": (0, n_channels - 1),\n", - " \"epsilon\": (y_min * 0.2, y_max * 0.2, 0.01),\n", - " f\"m{R}\": (0, 3, 100),\n", - " Rf\"\\Gamma_{{{R},{i}}}\": (-2, +2, 100),\n", - " Rf\"\\gamma_{{{R},{i}}}\": (0, 10, 100),\n", - " f\"m_a{i}\": (0, 1, 0.01),\n", - " f\"m_b{i}\": (0, 1, 0.01),\n", - " })\n", - " sliders.set_values({\n", - " f\"m{R}\": m0_values[R],\n", - " Rf\"\\Gamma_{{{R},{i}}}\": 2.0 * (0.4 + R * 0.2 - i * 0.3),\n", - " Rf\"\\gamma_{{{R},{i}}}\": 0.25 * (10 - R + i),\n", - " f\"m_a{i}\": (i + 1) * 0.25,\n", - " f\"m_b{i}\": (i + 1) * 0.25,\n", - " })\n", - "\n", - " # Create interactive plots\n", - " controls = Controls(**sliders)\n", - " fig, axes = plt.subplots(\n", - " nrows=2,\n", - " figsize=(8, 6),\n", - " sharex=True,\n", - " tight_layout=True,\n", - " )\n", - " fig.canvas.toolbar_visible = False\n", - " fig.canvas.header_visible = False\n", - " fig.canvas.footer_visible = False\n", - " if not title:\n", - " title = (\n", - " Rf\"${n_channels} \\times {n_channels}$ $K$-matrix\"\n", - " f\" with {n_resonances} resonances\"\n", - " )\n", - " fig.suptitle(title)\n", - "\n", - " for ax in axes:\n", - " ax.set_xlim(x_min, x_max)\n", - " ax_2d, ax_3d = axes\n", - " ax_2d.set_ylabel(\"$|T|^{2}$\")\n", - " ax_2d.set_yticks([])\n", - " ax_3d.set_xlabel(\"Re $m$\")\n", - " ax_3d.set_ylabel(\"Im $m$\")\n", - " ax_3d.set_xticks([])\n", - " ax_3d.set_yticks([])\n", - " ax_3d.set_facecolor(\"white\")\n", - "\n", - " ax_3d.axhline(0, linewidth=0.5, c=\"black\", linestyle=\"dotted\")\n", - "\n", - " # 2D plot\n", - " def plot(channel: int):\n", - " def wrapped(*args, **kwargs) -> sp.Expr:\n", - " kwargs[\"i\"] = channel\n", - " return np.abs(np_expr(*args, **kwargs)) ** 2\n", - "\n", - " return wrapped\n", - "\n", - " for i in range(n_channels):\n", - " iplt.plot(\n", - " plot_domain,\n", - " plot(i),\n", - " ax=axes[0],\n", - " controls=controls,\n", - " ylim=\"auto\",\n", - " label=f\"channel {i}\",\n", - " )\n", - " if n_channels > 1:\n", - " axes[0].legend(loc=\"upper right\")\n", - " mass_line_style = {\n", - " \"c\": \"red\",\n", - " \"alpha\": 0.3,\n", - " }\n", - " for name in controls.params:\n", - " if not re.match(r\"^m[0-9]+$\", name):\n", - " continue\n", - " iplt.axvline(controls[name], ax=axes[0], **mass_line_style)\n", - "\n", - " # 3D plot\n", - " color_mesh = None\n", - " epsilon_indicator = None\n", - " resonances_indicators = []\n", - " threshold_indicators = []\n", - "\n", - " def plot3(*, z_cutoff, complex_rendering, **kwargs):\n", - " nonlocal color_mesh, epsilon_indicator\n", - " epsilon = kwargs[\"epsilon\"]\n", - " kwargs[\"epsilon\"] = 0\n", - " Z = np_expr(plot_domain_complex, **kwargs)\n", - " if complex_rendering == \"imag\":\n", - " Z_values = Z.imag\n", - " ax_title = \"Re $T$\"\n", - " elif complex_rendering == \"real\":\n", - " Z_values = Z.real\n", - " ax_title = \"Im $T$\"\n", - " elif complex_rendering == \"abs\":\n", - " Z_values = np.abs(Z)\n", - " ax_title = \"$|T|$\"\n", - " else:\n", - " raise NotImplementedError\n", - "\n", - " if n_channels == 1:\n", - " axes[-1].set_title(ax_title)\n", - " else:\n", - " i = kwargs[\"i\"]\n", - " axes[-1].set_title(f\"{ax_title}, channel {i}\")\n", - "\n", - " if color_mesh is None:\n", - " color_mesh = ax_3d.pcolormesh(X, Y, Z_values, cmap=cm.coolwarm)\n", - " else:\n", - " color_mesh.set_array(Z_values)\n", - " color_mesh.set_clim(vmin=-z_cutoff, vmax=+z_cutoff)\n", - "\n", - " if resonances_indicators:\n", - " for R, (line, text) in enumerate(resonances_indicators):\n", - " mass = kwargs[f\"m{R}\"]\n", - " line.set_xdata(mass)\n", - " text.set_x(mass + (x_max - x_min) * 0.008)\n", - " else:\n", - " for R in range(n_resonances):\n", - " mass = kwargs[f\"m{R}\"]\n", - " line = ax_3d.axvline(mass, **mass_line_style)\n", - " text = ax_3d.text(\n", - " x=mass + (x_max - x_min) * 0.008,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{R}$\",\n", - " c=\"red\",\n", - " )\n", - " resonances_indicators.append((line, text))\n", - "\n", - " if epsilon_indicator is None:\n", - " line = ax.axhline(\n", - " epsilon,\n", - " linewidth=0.5,\n", - " c=\"blue\",\n", - " linestyle=\"dotted\",\n", - " label=R\"$\\epsilon$\",\n", - " )\n", - " text = axes[-1].text(\n", - " x=x_min + 0.008,\n", - " y=epsilon + 0.01,\n", - " s=R\"$\\epsilon$\",\n", - " c=\"blue\",\n", - " )\n", - " epsilon_indicator = line, text\n", - " else:\n", - " line, text = epsilon_indicator\n", - " line.set_xdata(epsilon)\n", - " text.set_y(epsilon + 0.01)\n", - "\n", - " x_offset = (x_max - x_min) * 0.015\n", - " if threshold_indicators:\n", - " for i, (line_thr, line_diff, text_thr, text_diff) in enumerate(\n", - " threshold_indicators\n", - " ):\n", - " m_a = kwargs[f\"m_a{i}\"]\n", - " m_b = kwargs[f\"m_b{i}\"]\n", - " s_thr = m_a + m_b\n", - " m_diff = m_a - m_b\n", - " line_thr.set_xdata(s_thr)\n", - " line_diff.set_xdata(m_diff)\n", - " text_thr.set_x(s_thr)\n", - " text_diff.set_x(m_diff - x_offset)\n", - " else:\n", - " colors = cm.plasma(np.linspace(0, 1, n_channels))\n", - " for i, color in enumerate(colors):\n", - " m_a = kwargs[f\"m_a{i}\"]\n", - " m_b = kwargs[f\"m_b{i}\"]\n", - " s_thr = m_a + m_b\n", - " m_diff = m_a - m_b\n", - " line_thr = ax.axvline(s_thr, c=color, linestyle=\"dotted\")\n", - " line_diff = ax.axvline(m_diff, c=color, linestyle=\"dashed\")\n", - " text_thr = ax.text(\n", - " x=s_thr,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{{a{i}}}+m_{{b{i}}}$\",\n", - " c=color,\n", - " rotation=-90,\n", - " )\n", - " text_diff = ax.text(\n", - " x=m_diff - x_offset,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{{a{i}}}-m_{{b{i}}}$\",\n", - " c=color,\n", - " rotation=+90,\n", - " )\n", - " threshold_indicators.append((line_thr, line_diff, text_thr, text_diff))\n", - " for i, (_, line_diff, _, text_diff) in enumerate(threshold_indicators):\n", - " m_a = kwargs[f\"m_a{i}\"]\n", - " m_b = kwargs[f\"m_b{i}\"]\n", - " s_thr = m_a + m_b\n", - " m_diff = m_a - m_b\n", - " if m_diff > x_offset + 0.01 and s_thr - abs(m_diff) > x_offset:\n", - " line_diff.set_alpha(0.5)\n", - " text_diff.set_alpha(0.5)\n", - " else:\n", - " line_diff.set_alpha(0)\n", - " text_diff.set_alpha(0)\n", - "\n", - " # Create switch for imag/real/abs\n", - " name = \"complex_rendering\"\n", - " sliders._sliders[name] = ipywidgets.RadioButtons(\n", - " options=[\"imag\", \"real\", \"abs\"],\n", - " description=R\"\\(s\\)-plane plot\",\n", - " )\n", - " sliders._arg_to_symbol[name] = name\n", - "\n", - " # Create cut-off slider for z-direction\n", - " name = \"z_cutoff\"\n", - " sliders._sliders[name] = ipywidgets.IntSlider(\n", - " value=30,\n", - " min=+1,\n", - " max=+100,\n", - " description=R\"\\(z\\)-cutoff\",\n", - " )\n", - " sliders._arg_to_symbol[name] = name\n", - "\n", - " # Create GUI\n", - " sliders_copy = dict(sliders)\n", - " h_boxes = []\n", - " for R in range(n_resonances):\n", - " buttons = [sliders_copy.pop(f\"m{R}\")]\n", - " if n_channels == 1:\n", - " buttons.append(sliders_copy.pop(symbol_to_arg[Rf\"\\Gamma_{{{R},0}}\"]))\n", - " buttons.append(sliders_copy.pop(symbol_to_arg[Rf\"\\gamma_{{{R},0}}\"]))\n", - " h_box = ipywidgets.HBox(buttons)\n", - " h_boxes.append(h_box)\n", - " remaining_sliders = sorted(\n", - " sliders_copy.values(), key=lambda s: (str(type(s)), s.description)\n", - " )\n", - " if n_channels == 1:\n", - " remaining_sliders.remove(sliders[\"i\"])\n", - " ui = ipywidgets.VBox(h_boxes + remaining_sliders)\n", - " output = ipywidgets.interactive_output(plot3, controls=sliders)\n", - " display(ui, output)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " L = 0" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "keep_output": false - }, - "outputs": [], - "source": [ - "plot_relativistic_k_matrix(\n", - " n_resonances=2,\n", - " n_channels=1,\n", - " angular_momentum=L,\n", - " title=\"Relativistic $K$-matrix, single channel\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "![](https://user-images.githubusercontent.com/29308176/164993776-43db5a5e-82b9-42f1-93c0-5d992d50477c.gif)" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/010.ipynb b/docs/report/010.ipynb deleted file mode 100644 index 24bbb9a3..00000000 --- a/docs/report/010.ipynb +++ /dev/null @@ -1,770 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix", - "physics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} P-vector\n", - "TR-010\n", - "^^^\n", - "This report is a sequel to [TR-005](005.ipynb) and [TR-009](009.ipynb).\n", - "+++\n", - "✅ [ampform#131](https://github.com/ComPWA/ampform/issues/131)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# P-vector" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.10.5 sympy==1.8" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import os\n", - "import re\n", - "import warnings\n", - "from typing import TYPE_CHECKING\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import mpl_interactions.ipyplot as iplt\n", - "import numpy as np\n", - "import symplot\n", - "import sympy as sp\n", - "from ampform.dynamics import (\n", - " BlattWeisskopfSquared,\n", - " breakup_momentum_squared,\n", - " coupled_width,\n", - " phase_space_factor_complex,\n", - ")\n", - "from ampform.dynamics.decorator import (\n", - " UnevaluatedExpression,\n", - " create_expression,\n", - " implement_doit_method,\n", - ")\n", - "from IPython.display import display\n", - "from ipywidgets import widgets as ipywidgets\n", - "from matplotlib import cm\n", - "from mpl_interactions.controller import Controls\n", - "\n", - "if TYPE_CHECKING:\n", - " from sympy.printing.latex import LatexPrinter\n", - "\n", - "warnings.filterwarnings(\"ignore\")\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Physics" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "source": [ - "As described in [TR-005](005.ipynb), the $\\boldsymbol{K}$-matrix describes **scattering processes** of the type $cd \\to ab$. The $P$-vector approach is one of two generalizations for **production processes** of the type $c \\to ab$. For more details on this approach, {cite}`Chung:1995dx` refers to {cite}`Aitchison:1972ay`.\n", - "\n", - "If we take the production vector $P$ to be:\n", - "\n", - "$$\n", - "P_i = \\sum_R \\frac{\\beta^0_R\\,g_{R,i}(m)}{m_R^2-m^2}\n", - "$$ (P-vector)\n", - "\n", - "and, in its invariant form,\n", - "\n", - "$$\n", - "\\hat{P}_i = \\sum_R \\frac{\\beta^0_R\\,g_{R,i}(m)}{\\left(m_R^2-m^2\\right)\\sqrt{\\rho_i}}\n", - "$$ (invariant P-vector)\n", - "\n", - "with $g_{R,i}(m)$ given by Eq. {eq}`residue-function` (possibly with [`coupled_width()`](https://ampform.readthedocs.io/en/0.10.5/api/ampform.dynamics.html#ampform.dynamics.coupled_width)), then the vector $F$ describes the resulting amplitudes by\n", - "\n", - "$$\n", - "\\begin{eqnarray}\n", - "F & = & \\left(\\boldsymbol{I}-i\\boldsymbol{K}\\right)^{-1}P \\\\\n", - "\\hat{F} & = & \\left(\\boldsymbol{I}-i\\boldsymbol{\\hat{K}\\boldsymbol{\\rho}}\\right)^{-1}\\hat{P}\n", - "\\end{eqnarray}\n", - "$$ (F-vector)\n", - "\n", - "with, from Eq. {eq}`covariant-K-matrix`:\n", - "\n", - "$$\n", - "\\hat{\\boldsymbol{K}} = \\sqrt{\\left(\\boldsymbol{\\rho}^\\dagger\\right)^{-1}} \\boldsymbol{K} \\sqrt{\\boldsymbol{\\rho}^{-1}}\n", - "$$ (K-hat in terms of K)\n", - "\n", - "\n", - "Just like with the residue functions in [TR-005](005.ipynb) and [TR-009](009.ipynb), $\\beta$ is often expressed in terms of resonance mass and 'width':\n", - "\n", - "$$\n", - "\\beta^0_R = \\beta_R\\sqrt{m_R\\Gamma^0_R}\n", - "$$ (beta functions)\n", - "\n", - "When in addition, we use a [`coupled_width()`](https://ampform.readthedocs.io/en/0.10.5/api/ampform.dynamics.html#ampform.dynamics.coupled_width), the $\\hat{P}$-vector becomes:\n", - "\n", - "$$\n", - "\\hat{P}_i = \\sum_R \\frac{\\beta_R\\gamma_{R,i}m_R\\Gamma^0_R B_{R,i}(m)}{m_R^2-m^2}\n", - "$$\n", - "\n", - "with $B_{R,i}(m)$ the ratio of Blatt-Weisskopf barrier factors ({class}`~ampform.dynamics.BlattWeisskopfSquared`) for channel $i$." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Implementation" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "@implement_doit_method()\n", - "class PhaseSpaceFactor(UnevaluatedExpression):\n", - " is_commutative = True\n", - "\n", - " def __new__(\n", - " cls,\n", - " s: sp.Symbol,\n", - " m_a: sp.Symbol,\n", - " m_b: sp.Symbol,\n", - " i: int,\n", - " **hints,\n", - " ) -> PhaseSpaceFactor:\n", - " return create_expression(cls, s, m_a, m_b, i, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m_a, m_b, *_ = self.args\n", - " return phase_space_factor_complex(s, m_a, m_b)\n", - "\n", - " def _latex(self, printer: LatexPrinter, *args) -> str:\n", - " s = printer._print(self.args[0])\n", - " i = self.args[-1]\n", - " return Rf\"\\rho_{{{i}}}({s})\"\n", - "\n", - "\n", - "@implement_doit_method()\n", - "class CoupledWidth(UnevaluatedExpression):\n", - " is_commutative = True\n", - "\n", - " def __new__(\n", - " cls,\n", - " s: sp.Symbol,\n", - " mass0: sp.IndexedBase,\n", - " gamma0: sp.IndexedBase,\n", - " m_a: sp.IndexedBase,\n", - " m_b: sp.IndexedBase,\n", - " angular_momentum: int,\n", - " R: int | sp.Symbol,\n", - " i: int,\n", - " **hints,\n", - " ) -> CoupledWidth:\n", - " return create_expression(\n", - " cls, s, mass0, gamma0, m_a, m_b, angular_momentum, R, i, **hints\n", - " )\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, mass0, gamma0, m_a, m_b, angular_momentum, R, i = self.args\n", - "\n", - " def phsp_factor(s, m_a, m_b):\n", - " return PhaseSpaceFactor(s, m_a, m_b, i)\n", - "\n", - " return coupled_width(\n", - " s,\n", - " mass0[R],\n", - " gamma0[R, i],\n", - " m_a[i],\n", - " m_b[i],\n", - " angular_momentum=angular_momentum,\n", - " meson_radius=1,\n", - " phsp_factor=phsp_factor,\n", - " )\n", - "\n", - " def _latex(self, printer: LatexPrinter, *args) -> str:\n", - " s = printer._print(self.args[0])\n", - " R = self.args[-2]\n", - " i = self.args[-1]\n", - " return Rf\"{{\\Gamma_{{{R},{i}}}}}({s})\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def Pi_relativistic(\n", - " m: sp.Symbol,\n", - " M: sp.IndexedBase,\n", - " Gamma: sp.IndexedBase,\n", - " gamma: sp.IndexedBase,\n", - " beta: sp.IndexedBase,\n", - " m_a: sp.IndexedBase,\n", - " m_b: sp.IndexedBase,\n", - " R: int | sp.Symbol,\n", - " i: int,\n", - " n_resonances: int | sp.Symbol,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - ") -> sp.Expr:\n", - " q_squared = breakup_momentum_squared(m**2, m_a[i], m_b[i])\n", - " ff2 = BlattWeisskopfSquared(z=q_squared, angular_momentum=angular_momentum)\n", - " parametrization = (beta[R] * gamma[R, i] * M[R] * Gamma[R, i] * sp.sqrt(ff2)) / (\n", - " M[R] ** 2 - m**2\n", - " )\n", - " return sp.Sum(parametrization, (R, 0, n_resonances - 1))\n", - "\n", - "\n", - "def Kij_relativistic(\n", - " m: sp.Symbol,\n", - " M: sp.IndexedBase,\n", - " Gamma: sp.IndexedBase,\n", - " gamma: sp.IndexedBase,\n", - " m_a: sp.IndexedBase,\n", - " m_b: sp.IndexedBase,\n", - " R: sp.IndexedBase,\n", - " i: int,\n", - " j: int,\n", - " n_resonances: int | sp.Symbol,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - ") -> sp.Expr:\n", - " def residue_function(i):\n", - " return gamma[R, i] * sp.sqrt(\n", - " M[R] * CoupledWidth(m**2, M, Gamma, m_a, m_b, angular_momentum, R, i)\n", - " )\n", - "\n", - " g_i = residue_function(i)\n", - " g_j = residue_function(j)\n", - " parametrization = (g_i * g_j) / (M[R] ** 2 - m**2)\n", - " return sp.Sum(parametrization, (R, 0, n_resonances - 1))\n", - "\n", - "\n", - "def f_vector(\n", - " n_resonances: int,\n", - " n_channels: int,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - ") -> sp.Matrix:\n", - " # Define symbols\n", - " R = sp.Symbol(\"R\")\n", - " m = sp.Symbol(\"m\", real=True)\n", - " M = sp.IndexedBase(\"m\", shape=(n_resonances,))\n", - " Gamma = sp.IndexedBase(\"Gamma\", shape=(n_resonances, n_channels))\n", - " gamma = sp.IndexedBase(\"gamma\", shape=(n_resonances, n_channels))\n", - " beta = sp.IndexedBase(\"beta\", shape=(n_resonances,))\n", - " m_a = sp.IndexedBase(\"m_a\", shape=(n_channels,))\n", - " m_b = sp.IndexedBase(\"m_b\", shape=(n_channels,))\n", - " # Define phase space matrix\n", - " rho = sp.zeros(n_channels, n_channels)\n", - " for i in range(n_channels):\n", - " rho[i, i] = PhaseSpaceFactor(m**2, m_a[i], m_b[i], i)\n", - " # Define P-vector, K-matrix and T-matrix\n", - " P = create_symbol_matrix(\"P\", n_channels, 1)\n", - " K = create_symbol_matrix(\"K\", n_channels, n_channels)\n", - " F = (sp.eye(n_channels) - sp.I * K * rho).inv() * P\n", - " # Substitute elements\n", - " return F.subs({\n", - " K[i, j]: Kij_relativistic(\n", - " m=m,\n", - " M=M,\n", - " Gamma=Gamma,\n", - " gamma=gamma,\n", - " m_a=m_a,\n", - " m_b=m_b,\n", - " i=i,\n", - " j=j,\n", - " R=R,\n", - " n_resonances=n_resonances,\n", - " angular_momentum=angular_momentum,\n", - " )\n", - " for i in range(n_channels)\n", - " for j in range(n_channels)\n", - " }).subs({\n", - " P[i]: Pi_relativistic(\n", - " m=m,\n", - " M=M,\n", - " Gamma=Gamma,\n", - " gamma=gamma,\n", - " beta=beta,\n", - " i=i,\n", - " m_a=m_a,\n", - " m_b=m_b,\n", - " R=R,\n", - " n_resonances=n_resonances,\n", - " angular_momentum=angular_momentum,\n", - " )\n", - " for i in range(n_channels)\n", - " })\n", - "\n", - "\n", - "def create_symbol_matrix(name: str, m: int, n: int) -> sp.Matrix:\n", - " symbol = sp.IndexedBase(name, shape=(m, n))\n", - " return sp.Matrix([[symbol[i, j] for j in range(n)] for i in range(m)])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\frac{\\sqrt{B_{L}^2\\left(\\frac{\\left(m^{2} - \\left({m_{a}}_{0} - {m_{b}}_{0}\\right)^{2}\\right) \\left(m^{2} - \\left({m_{a}}_{0} + {m_{b}}_{0}\\right)^{2}\\right)}{4 m^{2}}\\right)} {\\Gamma}_{0,0} {\\beta}_{0} {\\gamma}_{0,0} {m}_{0}}{\\left(1 - \\frac{i {\\Gamma_{0,0}}(m^{2}) {\\gamma}_{0,0}^{2} {m}_{0} \\rho_{0}(m^{2})}{- m^{2} + {m}_{0}^{2}}\\right) \\left(- m^{2} + {m}_{0}^{2}\\right)}$" - ], - "text/plain": [ - "sqrt(BlattWeisskopfSquared(L, (m**2 - (m_a[0] - m_b[0])**2)*(m**2 - (m_a[0] + m_b[0])**2)/(4*m**2)))*Gamma[0, 0]*beta[0]*gamma[0, 0]*m[0]/((1 - I*CoupledWidth(m**2, m, Gamma, m_a, m_b, L, 0, 0)*gamma[0, 0]**2*m[0]*PhaseSpaceFactor(m**2, m_a[0], m_b[0], 0)/(-m**2 + m[0]**2))*(-m**2 + m[0]**2))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "L = sp.Symbol(\"L\", integer=True)\n", - "symplot.partial_doit(\n", - " f_vector(n_resonances=1, n_channels=1, angular_momentum=L)[0, 0], sp.Sum\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "\n", - "Compare {func}`~ampform.dynamics.relativistic_breit_wigner_with_ff`\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def plot_f_vector(\n", - " n_channels: int,\n", - " n_resonances: int,\n", - " angular_momentum: int | sp.Symbol = 0,\n", - " title: str = \"\",\n", - ") -> None:\n", - " # Convert to Symbol: symplot cannot handle\n", - " m = sp.Symbol(\"m\", real=True)\n", - " epsilon = sp.Symbol(\"epsilon\", real=True)\n", - " i = sp.Symbol(\"i\", integer=True, negative=False)\n", - " expr = f_vector(n_resonances, n_channels, angular_momentum=angular_momentum).doit()[\n", - " i, 0\n", - " ]\n", - " expr = symplot.substitute_indexed_symbols(expr)\n", - " expr = expr.subs(m, m + epsilon * sp.I)\n", - " np_expr, sliders = symplot.prepare_sliders(expr, m)\n", - " symbol_to_arg = {symbol: arg for arg, symbol in sliders._arg_to_symbol.items()}\n", - "\n", - " # Set plot domain\n", - " x_min, x_max = 1e-3, 3\n", - " y_min, y_max = -0.5, +0.5\n", - "\n", - " plot_domain = np.linspace(x_min, x_max, num=500)\n", - " x_values = np.linspace(x_min, x_max, num=160)\n", - " y_values = np.linspace(y_min, y_max, num=80)\n", - " X, Y = np.meshgrid(x_values, y_values)\n", - " plot_domain_complex = X + Y * 1j\n", - "\n", - " # Set slider values and ranges\n", - " m0_values = np.linspace(x_min, x_max, num=n_resonances + 2)\n", - " m0_values = m0_values[1:-1]\n", - "\n", - " def set_default_values():\n", - " if \"L\" in sliders:\n", - " sliders.set_ranges(L=(0, 8))\n", - " sliders.set_ranges(\n", - " i=(0, n_channels - 1),\n", - " epsilon=(y_min * 0.2, y_max * 0.2, 0.01),\n", - " )\n", - " for R in range(n_resonances):\n", - " for i in range(n_channels):\n", - " sliders.set_ranges({\n", - " f\"m{R}\": (0, 3, 100),\n", - " f\"beta{R}\": (0, 5, 0.1),\n", - " Rf\"\\Gamma_{{{R},{i}}}\": (-5, +5, 100),\n", - " Rf\"\\gamma_{{{R},{i}}}\": (0, 20, 100),\n", - " f\"m_a{i}\": (0, 1, 0.01),\n", - " f\"m_b{i}\": (0, 1, 0.01),\n", - " })\n", - " sliders.set_values({\n", - " f\"m{R}\": m0_values[R],\n", - " f\"beta{R}\": 1,\n", - " Rf\"\\Gamma_{{{R},{i}}}\": 3 * (0.4 + R * 0.2 - i * 0.3),\n", - " Rf\"\\gamma_{{{R},{i}}}\": 0.2 * (10 - R + i),\n", - " f\"m_a{i}\": (i + 1) * 0.25,\n", - " f\"m_b{i}\": (i + 1) * 0.25,\n", - " })\n", - "\n", - " set_default_values()\n", - "\n", - " # Create interactive plots\n", - " controls = Controls(**sliders)\n", - " nrows = 2 # set to 3 for imag+real\n", - " fig, axes = plt.subplots(\n", - " nrows=nrows,\n", - " figsize=(8, nrows * 3.0),\n", - " sharex=True,\n", - " tight_layout=True,\n", - " )\n", - " fig.canvas.toolbar_visible = False\n", - " fig.canvas.header_visible = False\n", - " fig.canvas.footer_visible = False\n", - " for ax in axes:\n", - " ax.set_xlim(x_min, x_max)\n", - " if not title:\n", - " title = f\"{n_channels}-channel $F$-vector with {n_resonances} resonances\"\n", - " fig.suptitle(title)\n", - "\n", - " # 2D plot\n", - " axes[0].set_ylabel(\"$|T|^{2}$\")\n", - " axes[0].set_yticks([])\n", - "\n", - " def plot(channel: int):\n", - " def wrapped(*args, **kwargs) -> sp.Expr:\n", - " kwargs[\"i\"] = channel\n", - " return np.abs(np_expr(*args, **kwargs)) ** 2\n", - "\n", - " return wrapped\n", - "\n", - " for i in range(n_channels):\n", - " iplt.plot(\n", - " plot_domain,\n", - " plot(i),\n", - " ax=axes[0],\n", - " controls=controls,\n", - " ylim=\"auto\",\n", - " label=f\"channel {i}\",\n", - " )\n", - " if n_channels > 1:\n", - " axes[0].legend(loc=\"upper right\")\n", - " mass_line_style = {\n", - " \"c\": \"red\",\n", - " \"alpha\": 0.3,\n", - " }\n", - " for name in controls.params:\n", - " if not re.match(r\"^m[0-9]+$\", name):\n", - " continue\n", - " iplt.axvline(controls[name], ax=axes[0], **mass_line_style)\n", - "\n", - " # 3D plot\n", - " def plot3(**kwargs):\n", - " z_cutoff = kwargs.pop(\"z_cutoff\")\n", - " epsilon = kwargs[\"epsilon\"]\n", - " kwargs[\"epsilon\"] = 0\n", - " imag_real = kwargs.pop(\"imag_real\")\n", - " Z = np_expr(plot_domain_complex, **kwargs)\n", - " if imag_real == \"imag\":\n", - " Z_values = Z.imag\n", - " ax_title = \"Re $T$\"\n", - " elif imag_real == \"real\":\n", - " Z_values = Z.real\n", - " ax_title = \"Im $T$\"\n", - " elif imag_real == \"abs\":\n", - " Z_values = np.abs(Z)\n", - " ax_title = \"$|T|$\"\n", - " else:\n", - " raise NotImplementedError\n", - " for ax in axes[1:]:\n", - " ax.clear()\n", - " axes[-1].pcolormesh(\n", - " X, Y, Z_values, cmap=cm.coolwarm, vmin=-z_cutoff, vmax=+z_cutoff\n", - " )\n", - " i = kwargs[\"i\"]\n", - " if n_channels == 1:\n", - " axes[-1].set_title(ax_title)\n", - " else:\n", - " axes[-1].set_title(f\"{ax_title}, channel {i}\")\n", - " for ax in axes[1:]:\n", - " ax.axhline(0, linewidth=0.5, c=\"black\", linestyle=\"dotted\")\n", - " if epsilon != 0.0:\n", - " ax.axhline(\n", - " epsilon,\n", - " linewidth=0.5,\n", - " c=\"blue\",\n", - " linestyle=\"dotted\",\n", - " label=R\"$\\epsilon$\",\n", - " )\n", - " axes[-1].text(\n", - " x=x_min + 0.008,\n", - " y=epsilon + 0.01,\n", - " s=R\"$\\epsilon$\",\n", - " c=\"blue\",\n", - " )\n", - " for R in range(n_resonances):\n", - " mass = kwargs[f\"m{R}\"]\n", - " ax.axvline(mass, **mass_line_style)\n", - " if \"m_a0\" in kwargs:\n", - " colors = cm.plasma(np.linspace(0, 1, n_channels))\n", - " for i, color in enumerate(colors):\n", - " m_a = kwargs[f\"m_a{i}\"]\n", - " m_b = kwargs[f\"m_b{i}\"]\n", - " s_thr = m_a + m_b\n", - " ax.axvline(s_thr, c=color, linestyle=\"dotted\")\n", - " ax.text(\n", - " x=s_thr,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{{a{i}}}+m_{{b{i}}}$\",\n", - " c=color,\n", - " rotation=-90,\n", - " )\n", - " m_diff = m_a - m_b\n", - " x_offset = (x_max - x_min) * 0.015\n", - " if m_diff > x_offset + 0.01 and s_thr - abs(m_diff) > x_offset:\n", - " ax.axvline(\n", - " m_diff,\n", - " c=color,\n", - " linestyle=\"dashed\",\n", - " alpha=0.5,\n", - " )\n", - " ax.text(\n", - " x=m_diff - x_offset,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{{a{i}}}-m_{{b{i}}}$\",\n", - " c=color,\n", - " rotation=+90,\n", - " )\n", - " ax.set_ylabel(\"Im $m$\")\n", - " ax.set_xticks([])\n", - " ax.set_yticks([])\n", - " ax.set_facecolor(\"white\")\n", - " for R in range(n_resonances):\n", - " mass = kwargs[f\"m{R}\"]\n", - " axes[-1].text(\n", - " x=mass + (x_max - x_min) * 0.008,\n", - " y=0.95 * y_min,\n", - " s=f\"$m_{R}$\",\n", - " c=\"red\",\n", - " )\n", - " axes[-1].set_xlabel(\"Re $m$\")\n", - " fig.canvas.draw_idle()\n", - "\n", - " # Create switch for imag/real/abs\n", - " name = \"imag_real\"\n", - " sliders._sliders[name] = ipywidgets.RadioButtons(\n", - " options=[\"imag\", \"real\", \"abs\"],\n", - " description=R\"\\(s\\)-plane plot\",\n", - " )\n", - " sliders._arg_to_symbol[name] = name\n", - "\n", - " # Create cut-off slider for z-direction\n", - " name = \"z_cutoff\"\n", - " sliders._sliders[name] = ipywidgets.IntSlider(\n", - " value=10,\n", - " min=+1,\n", - " max=+50,\n", - " description=R\"\\(z\\)-cutoff\",\n", - " )\n", - " sliders._arg_to_symbol[name] = name\n", - "\n", - " # Create GUI\n", - " sliders_copy = dict(sliders)\n", - " h_boxes = []\n", - " for R in range(n_resonances):\n", - " buttons = [sliders_copy.pop(f\"m{R}\")]\n", - " if n_channels == 1:\n", - " buttons.append(sliders_copy.pop(symbol_to_arg[Rf\"\\Gamma_{{{R},0}}\"]))\n", - " buttons.append(sliders_copy.pop(symbol_to_arg[Rf\"\\gamma_{{{R},0}}\"]))\n", - " h_box = ipywidgets.HBox(buttons)\n", - " h_boxes.append(h_box)\n", - " remaining_sliders = sorted(\n", - " sliders_copy.values(), key=lambda s: (str(type(s)), s.description)\n", - " )\n", - " if n_channels == 1:\n", - " remaining_sliders.remove(sliders[\"i\"])\n", - " ui = ipywidgets.VBox(h_boxes + remaining_sliders)\n", - " output = ipywidgets.interactive_output(plot3, controls=sliders)\n", - " display(ui, output)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " L = 0" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "keep_output": false - }, - "outputs": [], - "source": [ - "plot_f_vector(\n", - " n_resonances=2,\n", - " n_channels=1,\n", - " angular_momentum=L,\n", - " title=\"Relativistic $F$-vector, single channel\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "![](https://user-images.githubusercontent.com/29308176/164993778-1f5987c2-4ff6-45e3-9ef4-cabbcb27b70a.gif)" - ] - } - ], - "metadata": { - "colab": { - "toc_visible": true - }, - "keep_output": 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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/011.ipynb b/docs/report/011.ipynb deleted file mode 100644 index b53146d2..00000000 --- a/docs/report/011.ipynb +++ /dev/null @@ -1,4822 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Helicity angles as symbolic expressions\n", - "TR-011\n", - "^^^\n", - "This report has been implemented in and [ampform#177](https://github.com/ComPWA/ampform/issues/177) and [tensorwaves#345](https://github.com/ComPWA/tensorwaves/issues/345). The report contains some bugs which were also addressed in these PRs.\n", - "+++\n", - "✅ [ampform#177](https://github.com/ComPWA/ampform/issues/177), [tensorwaves#345](https://github.com/ComPWA/tensorwaves/issues/345)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Symbolic kinematics" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.11.4 black==21.7b0 graphviz==0.17 numpy==1.19.5 qrules==0.9.2 git+https://github.com/ComPWA/sympy@20570-add-ArraySlice" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import inspect\n", - "from typing import TYPE_CHECKING\n", - "\n", - "import black\n", - "import graphviz\n", - "import numpy as np\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.data import EventCollection\n", - "from ampform.kinematics import (\n", - " _compute_helicity_angles,\n", - " determine_attached_final_state,\n", - " get_helicity_angle_label,\n", - ")\n", - "from ampform.sympy import (\n", - " UnevaluatedExpression,\n", - " create_expression,\n", - " implement_doit_method,\n", - ")\n", - "from IPython.display import Math, display\n", - "from qrules.topology import Topology, create_isobar_topologies\n", - "from sympy.printing.numpy import NumPyPrinter\n", - "from sympy.tensor.array.expressions.array_expressions import ArraySlice, ArraySymbol\n", - "\n", - "if TYPE_CHECKING:\n", - " from sympy.printing.printer import Printer" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This report investigates issue [compwa.github.io#56](https://github.com/ComPWA/compwa.github.io/issues/56). The ideal solution would be to use only SymPy in the existing [`ampform.kinematics`](https://ampform.readthedocs.io/en/0.11.4/api/ampform.kinematics.html) module. This has two benefits:\n", - "\n", - "1. It allows computing kinematic variables from four-momenta with different computational back-ends.\n", - "2. Expressions for kinematic variable can be inspected through their LaTeX representation.\n", - "\n", - "To simplify things, we investigate 1. by only lambdifying to NumPy. It should be relatively straightforward to lambdify to other back-ends like TensorFlow (as long as they support Einstein summation)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Test sample" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Data sample taken from [this test in AmpForm](https://github.com/ComPWA/ampform/tree/0.11.4/tests/conftest.py#L53-L115) and topology and expected angles taken from [here](https://github.com/ComPWA/ampform/tree/0.11.4/tests/test_kinematics.py#L15-L112)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "topologies = create_isobar_topologies(4)\n", - "topology = topologies[1]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot = qrules.io.asdot(topology)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "graphviz.Source(dot).render(\"011-topology\", format=\"svg\");" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164990755-ba9dfa9b-ed27-4bcf-b9cf-8fd62a71e3f4.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "events = EventCollection({\n", - " 0: np.array([ # pi0\n", - " (1.35527, 0.514208, -0.184219, 1.23296),\n", - " (0.841933, 0.0727385, -0.0528868, 0.826163),\n", - " (0.550927, -0.162529, 0.29976, -0.411133),\n", - " ]),\n", - " 1: np.array([ # gamma\n", - " (0.755744, -0.305812, 0.284, -0.630057),\n", - " (1.02861, 0.784483, 0.614347, -0.255334),\n", - " (0.356875, -0.20767, 0.272796, 0.0990739),\n", - " ]),\n", - " 2: np.array([ # pi0\n", - " (0.208274, -0.061663, -0.0211864, 0.144596),\n", - " (0.461193, -0.243319, -0.283044, -0.234866),\n", - " (1.03294, 0.82872, -0.0465425, -0.599834),\n", - " ]),\n", - " 3: np.array([ # pi0\n", - " (0.777613, -0.146733, -0.0785946, -0.747499),\n", - " (0.765168, -0.613903, -0.278416, -0.335962),\n", - " (1.15616, -0.458522, -0.526014, 0.911894),\n", - " ]),\n", - "})" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{'phi_1+2+3': ScalarSequence([ 2.79758029 2.51292308 -1.07396684],\n", - " 'theta_1+2+3': ScalarSequence([2.72456853 3.03316287 0.69240082],\n", - " 'phi_2+3,1+2+3': ScalarSequence([1.0436215 1.8734936 0.16073833],\n", - " 'theta_2+3,1+2+3': ScalarSequence([2.45361589 1.40639741 0.98079245],\n", - " 'phi_2,2+3,1+2+3': ScalarSequence([ 0.36955786 -1.68820498 0.63063002],\n", - " 'theta_2,2+3,1+2+3': ScalarSequence([1.0924374 1.99375767 1.31959621]}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "angles = _compute_helicity_angles(events, topology)\n", - "angles._DataSet__data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Einstein summation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "First challenge is to express the Einstein summation in [the existing implementation](https://github.com/ComPWA/ampform/tree/0.11.4/src/ampform/data.py#L207-L214) in terms of SymPy. The aim is to render the expression resulting nicely as LaTeX while at the same time being able to lambdify the expression to efficient NumPy code. We do this by deriving from [`UnevaluatedExpression`](https://ampform.readthedocs.io/en/0.11.4/api/ampform.sympy.html#ampform.sympy.UnevaluatedExpression) and using the decorator functions provided by the [`ampform.sympy`](https://ampform.readthedocs.io/en/0.11.4/api/ampform.sympy.html) module." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Define boost and rotation classes" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "First, wrap rotations and boosts in a class so with a nice LaTeX printer. {ref}`Later on `, a NumPy printer method will be defined externally for each of them." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "class BoostZ(UnevaluatedExpression):\n", - " def __new__(cls, beta: sp.Symbol, **hints) -> BoostZ:\n", - " return create_expression(cls, beta, **hints)\n", - "\n", - " def as_explicit(self) -> sp.Expr:\n", - " gamma = 1 / sp.sqrt(1 - beta**2)\n", - " return sp.Matrix([\n", - " [gamma, 0, 0, -gamma * beta],\n", - " [0, 1, 0, 0],\n", - " [0, 0, 1, 0],\n", - " [-gamma * beta, 0, 0, gamma],\n", - " ])\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " beta, *_ = self.args\n", - " beta = printer._print(beta)\n", - " return Rf\"\\boldsymbol{{B_z}}\\left({beta}\\right)\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "class RotationY(UnevaluatedExpression):\n", - " def __new__(cls, angle: sp.Symbol, **hints) -> RotationY:\n", - " return create_expression(cls, angle, **hints)\n", - "\n", - " def as_explicit(self) -> sp.Expr:\n", - " angle = self.args[0]\n", - " return sp.Matrix([\n", - " [1, 0, 0, 0],\n", - " [0, sp.cos(angle), 0, sp.sin(angle)],\n", - " [0, 0, 1, 0],\n", - " [0, -sp.sin(angle), 0, sp.cos(angle)],\n", - " ])\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " angle, *_ = self.args\n", - " angle = printer._print(angle)\n", - " return Rf\"\\boldsymbol{{R_y}}\\left({angle}\\right)\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "class RotationZ(UnevaluatedExpression):\n", - " def __new__(cls, angle: sp.Symbol, **hints) -> RotationZ:\n", - " return create_expression(cls, angle, **hints)\n", - "\n", - " def as_explicit(self) -> sp.Expr:\n", - " angle = self.args[0]\n", - " return sp.Matrix([\n", - " [1, 0, 0, 0],\n", - " [0, sp.cos(angle), -sp.sin(angle), 0],\n", - " [0, sp.sin(angle), sp.cos(angle), 0],\n", - " [0, 0, 0, 1],\n", - " ])\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " angle, *_ = self.args\n", - " angle = printer._print(angle)\n", - " return Rf\"\\boldsymbol{{R_z}}\\left({angle}\\right)\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Define Einstein summation class" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Similarly, we define a `ArrayMultiplication` class that will eventually be lambdified to {func}`numpy.einsum`.\n", - "\n", - "(011-array-multiplication)=" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "class ArrayMultiplication(sp.Expr):\n", - " def __new__(cls, *tensors: sp.Symbol, **hints):\n", - " return create_expression(cls, *tensors, **hints)\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " tensors_latex = map(printer._print, self.args)\n", - " return \" \".join(tensors_latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Indeed an expression involving these classes looks nice on the top-level:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\boldsymbol{B_z}\\left(\\beta\\right) \\boldsymbol{R_y}\\left(- \\theta\\right) \\boldsymbol{R_z}\\left(- \\phi\\right) p$" - ], - "text/plain": [ - "ArrayMultiplication(BoostZ(beta), RotationY(-theta), RotationZ(-phi), p)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "n_events = 3\n", - "momentum = sp.MatrixSymbol(\"p\", m=n_events, n=4)\n", - "beta = sp.Symbol(\"beta\")\n", - "phi = sp.Symbol(\"phi\")\n", - "theta = sp.Symbol(\"theta\")\n", - "\n", - "boosted_momentum = ArrayMultiplication(\n", - " BoostZ(beta),\n", - " RotationY(-theta),\n", - " RotationZ(-phi),\n", - " momentum,\n", - ")\n", - "boosted_momentum" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "\n", - "It could be that the above can be achieved with SymPy's `ArrayTensorProduct` and `ArrayContraction`. See for instance [sympy/sympy#22279](https://github.com/sympy/sympy/issues/22279).\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Define lambdification" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "# small helper function\n", - "\n", - "\n", - "def print_lambdify(symbols, expr):\n", - " np_expr = sp.lambdify(symbols, expr)\n", - " src = inspect.getsource(np_expr)\n", - " src = black.format_str(src, mode=black.Mode(line_length=79))\n", - " print(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we have a problem: lambdification does not work..." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(beta, theta, p):\n", - " return ( # Not supported in Python with SciPy and NumPy:\n", - " # ArrayMultiplication\n", - " ArrayMultiplication(\n", - " BoostZ(beta), RotationY(-theta), RotationZ(-phi), p\n", - " )\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print_lambdify([beta, theta, momentum], boosted_momentum)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "But lambdification can be defined externally to both the SymPy library and the expression classes. Here's an implementation for NumPy where we define the lambdification **through the expression class** (with a `_numpycode` method):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def print_as_numpy(self, printer: Printer, *args) -> str:\n", - " def multiply(matrix, vector):\n", - " return (\n", - " 'einsum(\"ij...,j...\",'\n", - " f\" transpose({matrix}, axes=(1, 2, 0)),\"\n", - " f\" transpose({vector}))\"\n", - " )\n", - "\n", - " def recursive_multiply(tensors):\n", - " if len(tensors) < 2:\n", - " msg = \"Need at least two tensors\"\n", - " raise ValueError(msg)\n", - " if len(tensors) == 2:\n", - " return multiply(tensors[0], tensors[1])\n", - " return multiply(tensors[0], recursive_multiply(tensors[1:]))\n", - "\n", - " printer.module_imports[\"numpy\"].update({\"einsum\", \"transpose\"})\n", - " tensors = list(map(printer._print, self.args))\n", - " if len(tensors) == 0:\n", - " return \"\"\n", - " if len(tensors) == 1:\n", - " return tensors[0]\n", - " return recursive_multiply(tensors)\n", - "\n", - "\n", - "ArrayMultiplication._numpycode = print_as_numpy" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(beta, theta, p):\n", - " return einsum(\n", - " \"ij...,j...\",\n", - " transpose(beta, axes=(1, 2, 0)),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\", transpose(theta, axes=(1, 2, 0)), transpose(p)\n", - " )\n", - " ),\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print_lambdify(\n", - " symbols=[beta, theta, momentum],\n", - " expr=ArrayMultiplication(beta, theta, momentum),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This also needs to be done for the rotation and boost classes:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(beta, theta, p):\n", - " return ( # Not supported in Python with SciPy and NumPy:\n", - " # BoostZ\n", - " # RotationY\n", - " # RotationZ\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(BoostZ(beta), axes=(1, 2, 0)),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(RotationY(-theta), axes=(1, 2, 0)),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(RotationZ(-phi), axes=(1, 2, 0)),\n", - " transpose(p),\n", - " )\n", - " ),\n", - " )\n", - " ),\n", - " )\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print_lambdify([beta, theta, momentum], boosted_momentum)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This time, we define the lambdification **through the printer class**:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def _print_BoostZ(self: NumPyPrinter, expr: BoostZ) -> str:\n", - " self.module_imports[\"numpy\"].update({\"array\", \"ones\", \"zeros\", \"sqrt\"})\n", - " arg = expr.args[0]\n", - " beta = self._print(arg)\n", - " gamma = f\"1 / sqrt(1 - ({beta}) ** 2)\"\n", - " n_events = f\"len({beta})\"\n", - " zeros = f\"zeros({n_events})\"\n", - " ones = f\"ones({n_events})\"\n", - " return f\"\"\"array(\n", - " [\n", - " [{gamma}, {zeros}, {zeros}, -{gamma} * {beta}],\n", - " [{zeros}, {ones}, {zeros}, {zeros}],\n", - " [{zeros}, {zeros}, {ones}, {zeros}],\n", - " [-{gamma} * {beta}, {zeros}, {zeros}, {gamma}],\n", - " ]\n", - " ).transpose(2, 0, 1)\"\"\"\n", - "\n", - "\n", - "NumPyPrinter._print_BoostZ = _print_BoostZ" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def _print_RotationY(self: NumPyPrinter, expr: RotationY) -> str:\n", - " self.module_imports[\"numpy\"].update({\"array\", \"cos\", \"ones\", \"zeros\", \"sin\"})\n", - " arg = expr.args[0]\n", - " angle = self._print(arg)\n", - " n_events = f\"len({angle})\"\n", - " zeros = f\"zeros({n_events})\"\n", - " ones = f\"ones({n_events})\"\n", - " return f\"\"\"array(\n", - " [\n", - " [{ones}, {zeros}, {zeros}, {zeros}],\n", - " [{zeros}, cos({angle}), {zeros}, sin({angle})],\n", - " [{zeros}, {zeros}, {ones}, {zeros}],\n", - " [{zeros}, -sin({angle}), {zeros}, cos({angle})],\n", - " ]\n", - " ).transpose(2, 0, 1)\"\"\"\n", - "\n", - "\n", - "NumPyPrinter._print_RotationY = _print_RotationY" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "def _print_RotationZ(self: NumPyPrinter, expr: RotationZ) -> str:\n", - " self.module_imports[\"numpy\"].update({\"array\", \"cos\", \"ones\", \"zeros\", \"sin\"})\n", - " arg = expr.args[0]\n", - " angle = self._print(arg)\n", - " n_events = f\"len({angle})\"\n", - " zeros = f\"zeros({n_events})\"\n", - " ones = f\"ones({n_events})\"\n", - " return f\"\"\"array(\n", - " [\n", - " [{ones}, {zeros}, {zeros}, {zeros}],\n", - " [{zeros}, cos({angle}), -sin({angle}), {zeros}],\n", - " [{zeros}, sin({angle}), cos({angle}), {zeros}],\n", - " [{zeros}, {zeros}, {zeros}, {ones}],\n", - " ]\n", - " ).transpose(2, 0, 1)\"\"\"\n", - "\n", - "\n", - "NumPyPrinter._print_RotationZ = _print_RotationZ" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-output", - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(beta, theta, p):\n", - " return einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " 1 / sqrt(1 - (beta) ** 2),\n", - " zeros(len(beta)),\n", - " zeros(len(beta)),\n", - " -1 / sqrt(1 - (beta) ** 2) * beta,\n", - " ],\n", - " [\n", - " zeros(len(beta)),\n", - " ones(len(beta)),\n", - " zeros(len(beta)),\n", - " zeros(len(beta)),\n", - " ],\n", - " [\n", - " zeros(len(beta)),\n", - " zeros(len(beta)),\n", - " ones(len(beta)),\n", - " zeros(len(beta)),\n", - " ],\n", - " [\n", - " -1 / sqrt(1 - (beta) ** 2) * beta,\n", - " zeros(len(beta)),\n", - " zeros(len(beta)),\n", - " 1 / sqrt(1 - (beta) ** 2),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(len(-theta)),\n", - " zeros(len(-theta)),\n", - " zeros(len(-theta)),\n", - " zeros(len(-theta)),\n", - " ],\n", - " [\n", - " zeros(len(-theta)),\n", - " cos(-theta),\n", - " zeros(len(-theta)),\n", - " sin(-theta),\n", - " ],\n", - " [\n", - " zeros(len(-theta)),\n", - " zeros(len(-theta)),\n", - " ones(len(-theta)),\n", - " zeros(len(-theta)),\n", - " ],\n", - " [\n", - " zeros(len(-theta)),\n", - " -sin(-theta),\n", - " zeros(len(-theta)),\n", - " cos(-theta),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(len(-phi)),\n", - " zeros(len(-phi)),\n", - " zeros(len(-phi)),\n", - " zeros(len(-phi)),\n", - " ],\n", - " [\n", - " zeros(len(-phi)),\n", - " cos(-phi),\n", - " -sin(-phi),\n", - " zeros(len(-phi)),\n", - " ],\n", - " [\n", - " zeros(len(-phi)),\n", - " sin(-phi),\n", - " cos(-phi),\n", - " zeros(len(-phi)),\n", - " ],\n", - " [\n", - " zeros(len(-phi)),\n", - " zeros(len(-phi)),\n", - " zeros(len(-phi)),\n", - " ones(len(-phi)),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(p),\n", - " )\n", - " ),\n", - " )\n", - " ),\n", - " )\n", - "\n" - ] - } - ], - "source": [ - "print_lambdify([beta, theta, momentum], boosted_momentum)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "\n", - "The code above contains a lot of duplicate code, such as `len(-phi)`. This could possibly be improved with {class}`~sympy.codegen.ast.CodeBlock`. See [ampform#166](https://github.com/ComPWA/ampform/issues/166).\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Angle computation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Computing phi" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The simplest angle to compute is $\\phi$, because it's simply $\\phi=\\arctan(p_y, p_x)$, with $p$ a four-momentum (see [existing implementation](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L161-L162)). This means we would need a way to represent $p_x$ and $p_y$ and some container for $\\phi$ itself. For convenience, we define $p_z$ and $E_p$ as well." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@implement_doit_method()\n", - "class FourMomentumX(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumX:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> ArraySlice:\n", - " return ArraySlice(self.momentum, (slice(None), 1))\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " momentum = printer._print(self.momentum)\n", - " return f\"{{{momentum}}}_{{x}}\"\n", - "\n", - "\n", - "@implement_doit_method()\n", - "class FourMomentumY(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumY:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> ArraySlice:\n", - " return ArraySlice(self.momentum, (slice(None), 2))\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " momentum = printer._print(self.momentum)\n", - " return f\"{{{momentum}}}_{{y}}\"\n", - "\n", - "\n", - "@implement_doit_method()\n", - "class FourMomentumZ(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> FourMomentumZ:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> ArraySlice:\n", - " return ArraySlice(self.momentum, (slice(None), 3))\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " momentum = printer._print(self.momentum)\n", - " return f\"{{{momentum}}}_{{y}}\"\n", - "\n", - "\n", - "@implement_doit_method()\n", - "class Energy(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> Energy:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> ArraySlice:\n", - " return ArraySlice(self.momentum, (slice(None), 0))\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " momentum = printer._print(self.momentum)\n", - " return f\"{{E}}_{{{momentum}}}\"\n", - "\n", - "\n", - "@implement_doit_method()\n", - "class Phi(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> Phi:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " p = self.momentum\n", - " return sp.atan2(FourMomentumY(p), FourMomentumX(p))\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " momentum = printer._print(self.momentum)\n", - " return Rf\"\\phi\\left({momentum}\\right)\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The classes indeed render nicely as LaTeX:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle p_{1}$" - ], - "text/plain": [ - "p_1" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\phi\\left(p_{1}\\right) & = & \\operatorname{atan_{2}}{\\left({p_{1}}_{y},{p_{1}}_{x} \\right)} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "{p_{1}}_{x} & = & p_{1}\\left[:, 1\\right] \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "{p_{1}}_{y} & = & p_{1}\\left[:, 2\\right] \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "{E}_{p_{1}} & = & p_{1}\\left[:, 0\\right] \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "p = ArraySymbol(\"p_1\")\n", - "phi = Phi(p)\n", - "p_x = FourMomentumX(p)\n", - "p_y = FourMomentumY(p)\n", - "energy = Energy(p)\n", - "\n", - "math_style = {\"environment\": \"eqnarray\"}\n", - "display(\n", - " p,\n", - " Math(sp.multiline_latex(Phi(p), phi.evaluate(), **math_style)),\n", - " Math(sp.multiline_latex(p_x, p_x.doit(), **math_style)),\n", - " Math(sp.multiline_latex(p_y, p_y.doit(), **math_style)),\n", - " Math(sp.multiline_latex(energy, energy.doit(), **math_style)),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that the four classes like `FourMomentumX` expect an `ArraySymbol` as input. This requires [sympy/sympy#22265](https://github.com/sympy/sympy/issues/22265). This allows lambdifying the above expressions to valid NumPy code. Let's compare this with the existing implementation using the {ref}`report/011:Test sample`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[ 1.35527 , 0.514208 , -0.184219 , 1.23296 ],\n", - " [ 0.841933 , 0.0727385, -0.0528868, 0.826163 ],\n", - " [ 0.550927 , -0.162529 , 0.29976 , -0.411133 ]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "momentum_sample = events[0]\n", - "np.array(momentum_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([ 0.514208 , 0.0727385, -0.162529 ])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(p, p_x.doit())\n", - "np_expr(momentum_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, momentum_sample.p_x)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([1.35527 , 0.841933, 0.550927])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(p, energy.doit())\n", - "np_expr(momentum_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, momentum_sample.energy)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([-0.34401236, -0.62867104, 2.06762909])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([-0.34401236, -0.62867104, 2.06762909])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(p, phi.doit())\n", - "display(np.array(momentum_sample.phi()))\n", - "np_expr(momentum_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, momentum_sample.phi())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Computing theta" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Computing $\\theta$ is more complicated, because requires the norm of the three-momentum of $p$ (see [existing implementation](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L164-L165)). In other words:\n", - "\n", - "$$\n", - "\\theta = \\arccos\\left(\\frac{p_z}{\\left|\\vec{p}\\right|}\\right)\n", - "\\quad \\mathrm{with} \\quad\n", - "\\left|\\vec{p}\\right| = \\sqrt{p_x^2+p_y^2+p_z^2}\n", - "$$\n", - "\n", - "The complication here is that $\\left|\\vec{p}\\right|$ [needs to be computed with `np.sum`](https://github.com/ComPWA/ampform/blob/f381794/src/ampform/data.py#L157-L159) over a slice of the arrays. As of writing, it is not yet possible to write a SymPy expression that can lambdify to this or an equivalent with `np.einsum` ([sympy/sympy#22279](https://github.com/sympy/sympy/issues/22279)). So for now, an intermediate `ArrayAxisSum` class has to be written for this. See also [this remark](https://github.com/sympy/sympy/discussions/22269#discussioncomment-1473005)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "class ArrayAxisSum(sp.Expr):\n", - " array: ArraySymbol = property(lambda self: self.args[0])\n", - " axis: int | None = property(lambda self: self.args[1])\n", - "\n", - " def __new__(\n", - " cls, array: ArraySymbol, axis: int | None = None, **hints\n", - " ) -> ArrayAxisSum:\n", - " if axis is not None and not isinstance(axis, (int, sp.Integer)):\n", - " msg = \"Only single digits allowed for axis\"\n", - " raise TypeError(msg)\n", - " return create_expression(cls, array, axis, **hints)\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " A = printer._print(self.array)\n", - " if self.axis is None:\n", - " return Rf\"\\sum{{{A}}}\"\n", - " axis = printer._print(self.axis)\n", - " return Rf\"\\sum_{{\\mathrm{{axis{axis}}}}}{{{A}}}\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Looks nice as LaTeX:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\mathrm{axis1}}{A}$" - ], - "text/plain": [ - "ArrayAxisSum(A, 1)" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum{A}$" - ], - "text/plain": [ - "ArrayAxisSum(A, None)" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "A = ArraySymbol(\"A\")\n", - "display(\n", - " ArrayAxisSum(A, axis=1),\n", - " ArrayAxisSum(A),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now let's define a printer method for NumPy:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def _print_ArrayAxisSum(self: NumPyPrinter, expr: ArrayAxisSum) -> str:\n", - " self.module_imports[\"numpy\"].add(\"sum\")\n", - " array = self._print(expr.array)\n", - " axis = self._print(expr.axis)\n", - " return f\"sum({array}, axis={axis})\"\n", - "\n", - "\n", - "NumPyPrinter._print_ArrayAxisSum = _print_ArrayAxisSum" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(A):\n", - " return sum(A, axis=1)\n", - "\n" - ] - } - ], - "source": [ - "print_lambdify(A, ArrayAxisSum(A, axis=1))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "...and let's check whether it works as expected for a 3-dimensional array:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[[ 0, 1],\n", - " [ 2, 3],\n", - " [ 4, 5]],\n", - "\n", - " [[ 6, 7],\n", - " [ 8, 9],\n", - " [10, 11]]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "array = np.array(range(12)).reshape(2, 3, 2)\n", - "array" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "66" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(A, ArrayAxisSum(A))\n", - "np_expr(array)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert _ == np.sum(array)\n", - "assert _ == sum(range(12))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[ 6, 8],\n", - " [10, 12],\n", - " [14, 16]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=0))\n", - "np_expr(array)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, np.sum(array, axis=0))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[ 6, 9],\n", - " [24, 27]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=1))\n", - "np_expr(array)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, np.sum(array, axis=1))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[ 1, 5, 9],\n", - " [13, 17, 21]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(A, ArrayAxisSum(A, axis=2))\n", - "np_expr(array)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, np.sum(array, axis=2))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we're ready to define a class that can represent $\\left|\\vec{p}\\right|$ and that lambdifies to the expressions [given by the existing implementation](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L153-L159)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "@implement_doit_method()\n", - "class ThreeMomentumNorm(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> ThreeMomentumNorm:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> ArraySlice:\n", - " three_momentum = ArraySlice(self.momentum, (slice(None), slice(1, None)))\n", - " norm_squared = ArrayAxisSum(three_momentum**2, axis=1)\n", - " return sp.sqrt(norm_squared)\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " three_momentum = printer._print(self.momentum)\n", - " return Rf\"\\left|\\vec{{{three_momentum}}}\\right|\"\n", - "\n", - " def _numpycode(self, printer, *args) -> str:\n", - " return printer._print(self.evaluate())" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\left|\\vec{p_{1}}\\right| & = & \\sqrt{\\sum_{\\mathrm{axis1}}{p_{1}\\left[:, 1:\\right]^{2}}} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "p_norm = ThreeMomentumNorm(p)\n", - "Math(sp.multiline_latex(p_norm, p_norm.doit(), **math_style))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([1.34853137, 0.83104344, 0.53413676])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(p, p_norm.doit())\n", - "np_expr(momentum_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.array_equal(_, momentum_sample.p_norm())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "With that, we're ready to define the `Theta` class!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "@implement_doit_method()\n", - "class Theta(UnevaluatedExpression):\n", - " momentum: ArraySymbol = property(lambda self: self.args[0])\n", - "\n", - " def __new__(cls, momentum: ArraySymbol, **hints) -> Theta:\n", - " return create_expression(cls, momentum, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " p = self.momentum\n", - " return sp.acos(FourMomentumZ(p) / ThreeMomentumNorm(p))\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " momentum = printer._print(self.momentum)\n", - " return Rf\"\\theta\\left({momentum}\\right)\"" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The math doesn't look the best, but this due to the {ref}`ArrayMultiplication <011-array-multiplication>` class." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\theta\\left(p_{1}\\right) & = & \\operatorname{acos}{\\left(p_{1}\\left[:, 3\\right] \\frac{1}{\\sqrt{\\sum_{\\mathrm{axis1}}{p_{1}\\left[:, 1:\\right]^{2}}}} \\right)} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "Math(sp.multiline_latex(Theta(p), Theta(p).doit(), **math_style))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "At any rate, when we lambdify the whole thing, we have exactly the same result as the [original `theta()` method](https://github.com/ComPWA/ampform/blob/f38179483ac20664584014d9eb659866966681c2/src/ampform/data.py#L164-L165) gave!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([0.41702412, 0.10842903, 2.4491907 ])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([0.41702412, 0.10842903, 2.4491907 ])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "np_expr = sp.lambdify(p, Theta(p).doit())\n", - "display(np.array(momentum_sample.theta()))\n", - "np_expr(momentum_sample)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert np.allclose(_, momentum_sample.theta())" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Recursive angle computation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, we are ready to compute kinematic angles recursively from a {class}`~qrules.topology.Topology`, just as in the [existing implementation](https://github.com/ComPWA/ampform/blob/307bbb1a09e75a320fe812945d7f0c7e0cdcc9f8/src/ampform/kinematics.py#L204-L280), but now with SymPy only." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{0: p0, 1: p1, 2: p2, 3: p3}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def create_event_collection_from_topology(\n", - " topology: Topology,\n", - ") -> dict[int, ArraySymbol]:\n", - " n_final_states = len(topology.outgoing_edge_ids)\n", - " return {i: ArraySymbol(f\"p{i}\") for i in range(n_final_states)}\n", - "\n", - "\n", - "momentum_symbols = create_event_collection_from_topology(topology)\n", - "momentum_symbols" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now it's quite trivial to rewrite the [existing implementation](https://github.com/ComPWA/ampform/blob/307bbb1a09e75a320fe812945d7f0c7e0cdcc9f8/src/ampform/kinematics.py#L204-L280) with the classes defined above.\n", - "\n", - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def compute_helicity_angles(\n", - " events: dict[int, ArraySymbol], topology: Topology\n", - ") -> dict[str, sp.Expr]:\n", - " if topology.outgoing_edge_ids != set(events):\n", - " msg = (\n", - " f\"Momentum IDs {set(events)} do not match final state edge IDs\"\n", - " f\" {set(topology.outgoing_edge_ids)}\"\n", - " )\n", - " raise ValueError(msg)\n", - "\n", - " def __recursive_helicity_angles(\n", - " events: dict[int, ArraySymbol], node_id: int\n", - " ) -> dict[str, sp.Expr]:\n", - " helicity_angles: dict[str, sp.Expr] = {}\n", - " child_state_ids = sorted(topology.get_edge_ids_outgoing_from_node(node_id))\n", - " if all(topology.edges[i].ending_node_id is None for i in child_state_ids):\n", - " state_id = child_state_ids[0]\n", - " four_momentum = events[state_id]\n", - " phi_label, theta_label = get_helicity_angle_label(topology, state_id)\n", - " helicity_angles[phi_label] = Phi(four_momentum)\n", - " helicity_angles[theta_label] = Theta(four_momentum)\n", - " for state_id in child_state_ids:\n", - " edge = topology.edges[state_id]\n", - " if edge.ending_node_id is not None:\n", - " # recursively determine all momenta ids in the list\n", - " sub_momenta_ids = determine_attached_final_state(topology, state_id)\n", - " if len(sub_momenta_ids) > 1:\n", - " # add all of these momenta together -> defines new subsystem\n", - " four_momentum = sum(events[i] for i in sub_momenta_ids)\n", - "\n", - " # boost all of those momenta into this new subsystem\n", - " phi = Phi(four_momentum)\n", - " theta = Theta(four_momentum)\n", - " p3_norm = ThreeMomentumNorm(four_momentum)\n", - " beta = p3_norm / Energy(four_momentum)\n", - " new_momentum_pool = {\n", - " k: ArrayMultiplication(\n", - " BoostZ(beta),\n", - " RotationY(-theta),\n", - " RotationZ(-phi),\n", - " p,\n", - " )\n", - " for k, p in events.items()\n", - " if k in sub_momenta_ids\n", - " }\n", - "\n", - " # register current angle variables\n", - " phi_label, theta_label = get_helicity_angle_label(\n", - " topology, state_id\n", - " )\n", - " helicity_angles[phi_label] = Phi(four_momentum)\n", - " helicity_angles[theta_label] = Theta(four_momentum)\n", - "\n", - " # call next recursion\n", - " angles = __recursive_helicity_angles(\n", - " new_momentum_pool,\n", - " edge.ending_node_id,\n", - " )\n", - " helicity_angles.update(angles)\n", - "\n", - " return helicity_angles\n", - "\n", - " initial_state_id = next(iter(topology.incoming_edge_ids))\n", - " initial_state_edge = topology.edges[initial_state_id]\n", - " assert initial_state_edge.ending_node_id is not None\n", - " return __recursive_helicity_angles(events, initial_state_edge.ending_node_id)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The computation works indeed and can be rendered to both LaTeX and {mod}`numpy`!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "['phi_1+2+3',\n", - " 'theta_1+2+3',\n", - " 'phi_2+3,1+2+3',\n", - " 'theta_2+3,1+2+3',\n", - " 'phi_2,2+3,1+2+3',\n", - " 'theta_2,2+3,1+2+3']" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "symbolic_angles = compute_helicity_angles(momentum_symbols, topology)\n", - "list(symbolic_angles)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-output", - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\phi\\left(p_{1} + p_{2} + p_{3}\\right) & = & \\operatorname{atan_{2}}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 2\\right],\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1\\right] \\right)} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(p0, p1, p2, p3):\n", - " return arctan2((p1 + p2 + p3)[:, 2], (p1 + p2 + p3)[:, 1])\n", - "\n", - "\n" - ] - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\theta\\left(p_{1} + p_{2} + p_{3}\\right) & = & \\operatorname{acos}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 3\\right] \\frac{1}{\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}}} \\right)} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(p0, p1, p2, p3):\n", - " return arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1) ** (-1 / 2)\n", - " )\n", - "\n", - "\n" - ] - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\phi\\left(\\boldsymbol{B_z}\\left(\\left|\\vec{p_{1} + p_{2} + p_{3}}\\right| {E}_{p_{1} + p_{2} + p_{3}}^{-1}\\right) \\boldsymbol{R_y}\\left(- \\theta\\left(p_{1} + p_{2} + p_{3}\\right)\\right) \\boldsymbol{R_z}\\left(- \\phi\\left(p_{1} + p_{2} + p_{3}\\right)\\right) p_{2} + \\boldsymbol{B_z}\\left(\\left|\\vec{p_{1} + p_{2} + p_{3}}\\right| {E}_{p_{1} + p_{2} + p_{3}}^{-1}\\right) \\boldsymbol{R_y}\\left(- \\theta\\left(p_{1} + p_{2} + p_{3}\\right)\\right) \\boldsymbol{R_z}\\left(- \\phi\\left(p_{1} + p_{2} + p_{3}\\right)\\right) p_{3}\\right) & = & \\operatorname{atan_{2}}{\\left(\\left(\\boldsymbol{B_z}\\left(\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}} \\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 0\\right]^{-1}\\right) \\boldsymbol{R_y}\\left(- \\operatorname{acos}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 3\\right] \\frac{1}{\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}}} \\right)}\\right) \\boldsymbol{R_z}\\left(- \\operatorname{atan_{2}}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 2\\right],\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1\\right] \\right)}\\right) p_{2} + \\boldsymbol{B_z}\\left(\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}} \\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 0\\right]^{-1}\\right) \\boldsymbol{R_y}\\left(- \\operatorname{acos}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 3\\right] \\frac{1}{\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}}} \\right)}\\right) \\boldsymbol{R_z}\\left(- \\operatorname{atan_{2}}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 2\\right],\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1\\right] \\right)}\\right) p_{3}\\right)\\left[:, 2\\right],\\left(\\boldsymbol{B_z}\\left(\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}} \\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 0\\right]^{-1}\\right) \\boldsymbol{R_y}\\left(- \\operatorname{acos}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 3\\right] \\frac{1}{\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}}} \\right)}\\right) \\boldsymbol{R_z}\\left(- \\operatorname{atan_{2}}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 2\\right],\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1\\right] \\right)}\\right) p_{2} + \\boldsymbol{B_z}\\left(\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}} \\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 0\\right]^{-1}\\right) \\boldsymbol{R_y}\\left(- \\operatorname{acos}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 3\\right] \\frac{1}{\\sqrt{\\sum_{\\mathrm{axis1}}{\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1:\\right]^{2}}}} \\right)}\\right) \\boldsymbol{R_z}\\left(- \\operatorname{atan_{2}}{\\left(\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 2\\right],\\left(p_{1} + p_{2} + p_{3}\\right)\\left[:, 1\\right] \\right)}\\right) p_{3}\\right)\\left[:, 1\\right] \\right)} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "def _lambdifygenerated(p0, p1, p2, p3):\n", - " return arctan2(\n", - " (\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " -sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " -sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(p2),\n", - " )\n", - " ),\n", - " )\n", - " ),\n", - " )\n", - " + einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " -sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " -sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(p3),\n", - " )\n", - " ),\n", - " )\n", - " ),\n", - " )\n", - " )[:, 2],\n", - " (\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " -sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " -sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(p2),\n", - " )\n", - " ),\n", - " )\n", - " ),\n", - " )\n", - " + einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " -1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " )\n", - " * sqrt(sum((p1 + p2 + p3)[:, 1:] ** 2, axis=1))\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ),\n", - " 1\n", - " / sqrt(\n", - " 1\n", - " - (\n", - " sqrt(\n", - " sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " )\n", - " * (p1 + p2 + p3)[:, 0] ** (-1.0)\n", - " )\n", - " ** 2\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " -sin(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:]\n", - " ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arccos(\n", - " (p1 + p2 + p3)[:, 3]\n", - " * sum(\n", - " (p1 + p2 + p3)[:, 1:] ** 2,\n", - " axis=1,\n", - " )\n", - " ** (-1 / 2)\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(\n", - " einsum(\n", - " \"ij...,j...\",\n", - " transpose(\n", - " array(\n", - " [\n", - " [\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " -sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " sin(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " cos(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[:, 2],\n", - " (p1 + p2 + p3)[:, 1],\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " [\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " zeros(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ones(\n", - " len(\n", - " -arctan2(\n", - " (p1 + p2 + p3)[\n", - " :, 2\n", - " ],\n", - " (p1 + p2 + p3)[\n", - " :, 1\n", - " ],\n", - " )\n", - " )\n", - " ),\n", - " ],\n", - " ]\n", - " ).transpose(2, 0, 1),\n", - " axes=(1, 2, 0),\n", - " ),\n", - " transpose(p3),\n", - " )\n", - " ),\n", - " )\n", - " ),\n", - " )\n", - " )[:, 1],\n", - " )\n", - "\n", - "\n" - ] - } - ], - "source": [ - "def display_kinematic_variable(name: str) -> None:\n", - " expr = symbolic_angles[name]\n", - " display(Math(sp.multiline_latex(expr, expr.doit(), **math_style)))\n", - " print_lambdify(momentum_symbols.values(), expr.doit())\n", - "\n", - "\n", - "for angle_name in list(symbolic_angles)[:3]:\n", - " display_kinematic_variable(angle_name)\n", - " print()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Comparison computed value" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, what's left is to compare computed data with the {ref}`report/011:Test sample`. Here's a little comparison function that can be called for each angle:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def compare(angle_name: str) -> None:\n", - " np_expr = sp.lambdify(\n", - " args=momentum_symbols.values(),\n", - " expr=symbolic_angles[angle_name].doit(),\n", - " modules=\"numpy\",\n", - " )\n", - " computed = np_expr(*events.values())\n", - " expected = np.array(angles[angle_name])\n", - " display(computed, expected)\n", - " np.testing.assert_allclose(computed, expected)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "phi_1+2+3\n" - ] - }, - { - "data": { - "text/plain": [ - "array([ 2.79758029, 2.51292308, -1.07396684])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([ 2.79758029, 2.51292308, -1.07396684])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 2.48 ms, sys: 6.56 ms, total: 9.04 ms\n", - "Wall time: 7.77 ms\n", - "\n", - "theta_1+2+3\n" - ] - }, - { - "data": { - "text/plain": [ - "array([2.72456853, 3.03316287, 0.69240082])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([2.72456853, 3.03316287, 0.69240082])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 5.13 ms, sys: 297 µs, total: 5.43 ms\n", - "Wall time: 4.93 ms\n", - "\n", - "phi_2+3,1+2+3\n" - ] - }, - { - "data": { - "text/plain": [ - "array([1.0436215 , 1.8734936 , 0.16073833])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([1.0436215 , 1.8734936 , 0.16073833])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 34.1 ms, sys: 0 ns, total: 34.1 ms\n", - "Wall time: 33.1 ms\n", - "\n", - "theta_2+3,1+2+3\n" - ] - }, - { - "data": { - "text/plain": [ - "array([2.45361589, 1.40639741, 0.98079245])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([2.45361589, 1.40639741, 0.98079245])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 36.1 ms, sys: 451 µs, total: 36.6 ms\n", - "Wall time: 35.5 ms\n", - "\n", - "phi_2,2+3,1+2+3\n" - ] - }, - { - "data": { - "text/plain": [ - "array([ 0.36955786, -1.68820498, 0.63063002])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([ 0.36955786, -1.68820498, 0.63063002])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 1.25 s, sys: 162 ms, total: 1.42 s\n", - "Wall time: 1.41 s\n", - "\n", - "theta_2,2+3,1+2+3\n" - ] - }, - { - "data": { - "text/plain": [ - "array([1.0924374 , 1.99375767, 1.31959621])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "array([1.0924374 , 1.99375767, 1.31959621])" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 1.22 s, sys: 166 ms, total: 1.38 s\n", - "Wall time: 1.38 s\n", - "\n" - ] - } - ], - "source": [ - "for angle_name in symbolic_angles:\n", - " print(angle_name)\n", - " %time compare(angle_name)\n", - " print()" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/012.ipynb b/docs/report/012.ipynb deleted file mode 100644 index 316dc9cb..00000000 --- a/docs/report/012.ipynb +++ /dev/null @@ -1,371 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "lambdification", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Extended DataSample performance\n", - "TR-012\n", - "^^^\n", - "[ampform#198](https://github.com/ComPWA/ampform/issues/198) makes it easier to generate expressions for kinematic variables that are not contained in the [`HelicityModel.expression`](https://ampform.readthedocs.io/en/0.12.1/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression). In TensorWaves, this results in a [`DataSample`](https://tensorwaves.readthedocs.io/en/0.4.x/api/tensorwaves.interface.html#tensorwaves.interface.DataSample) with more keys.\n", - "\n", - "A question was raised whether this affects the duration of fits. This report shows that this is not the case (see {ref}`report/012:Conclusion`).\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Extended `DataSample` performance" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Note: you may need to restart the kernel to use updated packages.\n" - ] - } - ], - "source": [ - "%pip install -q ampform==0.12.6 tensorwaves[jax,pwa]==0.4.8" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import logging\n", - "\n", - "import ampform\n", - "import numpy as np\n", - "import qrules\n", - "from ampform.dynamics.builder import (\n", - " create_non_dynamic_with_ff,\n", - " create_relativistic_breit_wigner_with_ff,\n", - ")\n", - "from tensorwaves.data import (\n", - " IntensityDistributionGenerator,\n", - " SympyDataTransformer,\n", - " TFPhaseSpaceGenerator,\n", - " TFUniformRealNumberGenerator,\n", - ")\n", - "from tensorwaves.function.sympy import create_parametrized_function\n", - "\n", - "LOGGER = logging.getLogger(\"absl\")\n", - "LOGGER.setLevel(logging.ERROR)\n", - "LOGGER = logging.getLogger()\n", - "LOGGER.setLevel(logging.ERROR)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate amplitude model" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Formulate a [`HelicityModel`](https://ampform.readthedocs.io/en/0.12.1/api/ampform.helicity.html#ampform.helicity.HelicityModel) just like in the [usual workflow](https://ampform.readthedocs.io/en/0.12.1/usage/amplitude.html):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "reaction = qrules.generate_transitions(\n", - " initial_state=(\"J/psi(1S)\", [-1, +1]),\n", - " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", - " allowed_intermediate_particles=[\"f(0)\"],\n", - " allowed_interaction_types=[\"strong\", \"EM\"],\n", - " formalism=\"helicity\",\n", - ")\n", - "\n", - "builder = ampform.get_builder(reaction)\n", - "builder.set_dynamics(\"J/psi(1S)\", create_non_dynamic_with_ff)\n", - "for name in reaction.get_intermediate_particles().names:\n", - " builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", - "model = builder.formulate()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now register more topologies with [`HelicityAdapter.permutate_registered_topologies()`](https://ampform.readthedocs.io/en/0.12.1/api/ampform.kinematics.html#ampform.kinematics.HelicityAdapter.permutate_registered_topologies) and formulate a new 'extended' model:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "builder.adapter.permutate_registered_topologies()\n", - "extended_model = builder.formulate()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Create computational functions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, create [`ParametrizedFunction`](https://tensorwaves.readthedocs.io/en/0.4.x/api/tensorwaves.interface.html#tensorwaves.interface.ParametrizedFunction)s for the normal model and the extended model:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensity = create_parametrized_function(\n", - " expression=model.expression.doit(),\n", - " parameters=model.parameter_defaults,\n", - " backend=\"jax\",\n", - ")\n", - "helicity_transformer = SympyDataTransformer.from_sympy(\n", - " model.kinematic_variables, backend=\"jax\"\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "extended_intensity = create_parametrized_function(\n", - " expression=extended_model.expression.doit(),\n", - " parameters=extended_model.parameter_defaults,\n", - " backend=\"jax\",\n", - ")\n", - "extended_helicity_transformer = SympyDataTransformer.from_sympy(\n", - " extended_model.kinematic_variables, backend=\"jax\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Generate phase space domain and hit-and-miss data sample with the normal intensity function and helicity transformer..." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "data_generator = IntensityDistributionGenerator(\n", - " function=intensity,\n", - " domain_generator=phsp_generator,\n", - " domain_transformer=helicity_transformer,\n", - ")\n", - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_momenta = phsp_generator.generate(100_000, rng)\n", - "data_momenta = data_generator.generate(10_000, rng)\n", - "phsp = helicity_transformer(phsp_momenta)\n", - "data = helicity_transformer(data_momenta)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "...and with the extended function and transformer:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "extended_phsp_generator = TFPhaseSpaceGenerator(\n", - " # actually same as phsp_generator\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "extended_data_generator = IntensityDistributionGenerator(\n", - " function=extended_intensity,\n", - " domain_generator=phsp_generator,\n", - " domain_transformer=helicity_transformer,\n", - ")\n", - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_momenta = extended_phsp_generator.generate(100_000, rng)\n", - "data_momenta = extended_data_generator.generate(10_000, rng)\n", - "extended_phsp = extended_helicity_transformer(phsp_momenta)\n", - "extended_data = extended_helicity_transformer(data_momenta)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Conclusion" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "data": { - "text/plain": [ - "(100000,)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "intensities = intensity(phsp)\n", - "extended_intensities = extended_intensity(extended_phsp)\n", - "extended_intensities.shape" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Computation time per iteration is the same:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-skip}\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "14.7 ms ± 761 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n", - "14.7 ms ± 669 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" - ] - } - ], - "source": [ - "%timeit -n10 intensity(phsp)\n", - "\n", - "%timeit -n10 extended_intensity(extended_phsp)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Output arrays are also the same:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "np.testing.assert_allclose(intensities, extended_intensities)\n", - "\n", - "assert set(data) < set(extended_data)\n", - "assert set(phsp) < set(extended_phsp)\n", - "for var in data:\n", - " np.testing.assert_allclose(phsp[var], extended_phsp[var])\n", - " np.testing.assert_allclose(data[var], extended_data[var])" - ] - } - ], - "metadata": { - "colab": { - "toc_visible": true - }, - "keep_output": 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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/013.ipynb b/docs/report/013.ipynb deleted file mode 100644 index 82e524af..00000000 --- a/docs/report/013.ipynb +++ /dev/null @@ -1,638 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Spin alignment with data\n", - "TR-013\n", - "^^^\n", - "In this report, we attempt to check the effect of activating spin alignment ([ampform#245](https://github.com/ComPWA/ampform/pull/245)) and compare it with [Figure 2](https://downloads.hindawi.com/journals/ahep/2020/6674595.pdf#page=9) in {cite}`Marangotto:2019ucc`.\n", - "\n", - "See also [TR-014](014.ipynb) and [TR-015](015.ipynb).\n", - "+++\n", - "WIP\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Spin alignment with data" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.13.3 qrules[viz]==0.9.7 tensorwaves[jax,pwa]==0.4.8" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "import logging\n", - "import warnings\n", - "\n", - "from IPython.display import display\n", - "\n", - "LOGGER = logging.getLogger()\n", - "LOGGER.setLevel(logging.ERROR)\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Phase space sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import qrules\n", - "\n", - "PDG = qrules.load_pdg()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator\n", - "\n", - "phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=PDG[\"Lambda(c)+\"].mass,\n", - " final_state_masses={\n", - " 0: PDG[\"p\"].mass,\n", - " 1: PDG[\"K-\"].mass,\n", - " 2: PDG[\"pi+\"].mass,\n", - " },\n", - ")\n", - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_momenta = phsp_generator.generate(1_000_000, rng)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate transitions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from qrules.particle import ParticleCollection, create_particle\n", - "\n", - "particle_db = ParticleCollection()\n", - "particle_db.add(PDG[\"Lambda(c)+\"])\n", - "particle_db.add(PDG[\"p\"])\n", - "particle_db.add(PDG[\"K-\"])\n", - "particle_db.add(PDG[\"pi+\"])\n", - "\n", - "particle_db.add(\n", - " create_particle(\n", - " PDG[\"K*(892)0\"],\n", - " name=\"K*\",\n", - " latex=\"K^*\",\n", - " )\n", - ")\n", - "particle_db.add(\n", - " create_particle(\n", - " PDG[\"Lambda(1405)\"],\n", - " name=\"Lambda*\",\n", - " latex=R\"\\Lambda^*\",\n", - " )\n", - ")\n", - "particle_db.add(\n", - " create_particle(\n", - " PDG[\"Delta(1232)++\"],\n", - " name=\"Delta*++\",\n", - " latex=R\"\\Delta^*\",\n", - " )\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "reaction = qrules.generate_transitions(\n", - " initial_state=(\"Lambda(c)+\", [-0.5, +0.5]),\n", - " final_state=[\"p\", \"K-\", \"pi+\"],\n", - " formalism=\"helicity\",\n", - " particle_db=particle_db,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "import graphviz\n", - "\n", - "n = len(reaction.transitions)\n", - "for i, t in enumerate(reaction.transitions[:: n // 3]):\n", - " dot = qrules.io.asdot([t], collapse_graphs=True, size=3.5)\n", - " graph = graphviz.Source(dot)\n", - " graph.render(f\"013-graph{i}\", format=\"svg\")\n", - " display(graph)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{container} full-width\n", - "![013_12_0](https://user-images.githubusercontent.com/29308176/164991353-d3228bee-4ce7-40f6-87c1-4ec9babba238.svg)\n", - "![013_12_1](https://user-images.githubusercontent.com/29308176/164991356-98885719-874c-486f-b70b-58e8cd2d9b09.svg)\n", - "![013_12_2](https://user-images.githubusercontent.com/29308176/164991358-c2b5e5f3-4b62-433e-af4b-fff44ad0822b.svg)\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Distribution without alignment" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Amplitude model formulated following [Appendix C](https://downloads.hindawi.com/journals/ahep/2020/6674595.pdf#page=13):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{m_{A}=-1/2}^{1/2} \\sum_{m_{0}=-1/2}^{1/2} \\sum_{m_{1}=0} \\sum_{m_{2}=0}{\\left|{{A^{01}}_{m_{A},m_{0},m_{1},m_{2}} + {A^{02}}_{m_{A},m_{0},m_{1},m_{2}} + {A^{12}}_{m_{A},m_{0},m_{1},m_{2}}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(A^01[m_A, m0, m1, m2] + A^02[m_A, m0, m1, m2] + A^12[m_A, m0, m1, m2])**2, (m_A, (1/2, -1/2)), (m0, (1/2, -1/2)), (m1, (0,)), (m2, (0,)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import ampform\n", - "from ampform.dynamics.builder import RelativisticBreitWignerBuilder\n", - "\n", - "builder = ampform.get_builder(reaction)\n", - "builder.align_spin = False\n", - "builder.stable_final_state_ids = list(reaction.final_state)\n", - "builder.scalar_initial_state_mass = True\n", - "bw_builder = RelativisticBreitWignerBuilder()\n", - "for name in reaction.get_intermediate_particles().names:\n", - " builder.set_dynamics(name, bw_builder)\n", - "standard_model = builder.formulate()\n", - "standard_model.intensity" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "import sympy as sp\n", - "from IPython.display import Math, display\n", - "\n", - "for i, (symbol, expr) in enumerate(standard_model.amplitudes.items()):\n", - " if i == 3:\n", - " display(Math(R\"\\dots\"))\n", - " break\n", - " latex = sp.multiline_latex(symbol, expr, environment=\"eqnarray\")\n", - " display(Math(latex))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Importing the parameter values given by [Table 1](https://downloads.hindawi.com/journals/ahep/2020/6674595.pdf#page=13):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from ampform.helicity import HelicityModel\n", - "\n", - "# fmt: off\n", - "parameter_table = {\n", - " # K*\n", - " R\"C_{\\Lambda_{c}^{+} \\to K^*_{0} p_{+1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}}\": 1,\n", - " R\"C_{\\Lambda_{c}^{+} \\to K^*_{+1} p_{+1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}}\": 0.5 + 0.5j,\n", - " R\"C_{\\Lambda_{c}^{+} \\to K^*_{-1} p_{-1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}}\": 1j,\n", - " R\"C_{\\Lambda_{c}^{+} \\to K^*_{0} p_{-1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}}\": -0.5 - 0.5j,\n", - " \"m_{K^*}\": 0.9, # GeV\n", - " R\"\\Gamma_{K^*}\": 0.2, # GeV\n", - " # Λ*\n", - " R\"C_{\\Lambda_{c}^{+} \\to \\Lambda^*_{-1/2} \\pi^{+}_{0}; \\Lambda^* \\to K^{-}_{0} p_{+1/2}}\": 1j,\n", - " R\"C_{\\Lambda_{c}^{+} \\to \\Lambda^*_{+1/2} \\pi^{+}_{0}; \\Lambda^* \\to K^{-}_{0} p_{+1/2}}\": 0.8 - 0.4j,\n", - " R\"m_{\\Lambda^*}\": 1.6, # GeV\n", - " R\"\\Gamma_{\\Lambda^*}\": 0.2, # GeV\n", - " # Δ*\n", - " R\"C_{\\Lambda_{c}^{+} \\to \\Delta^*_{+1/2} K^{-}_{0}; \\Delta^* \\to p_{+1/2} \\pi^{+}_{0}}\": 0.6 - 0.4j,\n", - " R\"C_{\\Lambda_{c}^{+} \\to \\Delta^*_{-1/2} K^{-}_{0}; \\Delta^* \\to p_{+1/2} \\pi^{+}_{0}}\": 0.1j,\n", - " R\"m_{\\Delta^*}\": 1.4, # GeV\n", - " R\"\\Gamma_{\\Delta^*}\": 0.2, # GeV\n", - "}\n", - "# fmt: on\n", - "\n", - "\n", - "def set_coefficients(model: HelicityModel) -> None:\n", - " for name, value in parameter_table.items():\n", - " model.parameter_defaults[name] = value" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{lc}\n", - " C_{\\Lambda_{c}^{+} \\to K^*_{0} p_{+1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}} & 1 \\\\\n", - " C_{\\Lambda_{c}^{+} \\to K^*_{+1} p_{+1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}} & 0.5+0.5i \\\\\n", - " C_{\\Lambda_{c}^{+} \\to K^*_{-1} p_{-1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}} & 1i \\\\\n", - " C_{\\Lambda_{c}^{+} \\to K^*_{0} p_{-1/2}; K^* \\to K^{-}_{0} \\pi^{+}_{0}} & -0.5-0.5i \\\\\n", - " m_{K^*} & 0.9 \\\\\n", - " \\Gamma_{K^*} & 0.2 \\\\\n", - " C_{\\Lambda_{c}^{+} \\to \\Lambda^*_{-1/2} \\pi^{+}_{0}; \\Lambda^* \\to K^{-}_{0} p_{+1/2}} & 1i \\\\\n", - " C_{\\Lambda_{c}^{+} \\to \\Lambda^*_{+1/2} \\pi^{+}_{0}; \\Lambda^* \\to K^{-}_{0} p_{+1/2}} & 0.8-0.4i \\\\\n", - " m_{\\Lambda^*} & 1.6 \\\\\n", - " \\Gamma_{\\Lambda^*} & 0.2 \\\\\n", - " C_{\\Lambda_{c}^{+} \\to \\Delta^*_{+1/2} K^{-}_{0}; \\Delta^* \\to p_{+1/2} \\pi^{+}_{0}} & 0.6-0.4i \\\\\n", - " C_{\\Lambda_{c}^{+} \\to \\Delta^*_{-1/2} K^{-}_{0}; \\Delta^* \\to p_{+1/2} \\pi^{+}_{0}} & 0.1i \\\\\n", - " m_{\\Delta^*} & 1.4 \\\\\n", - " \\Gamma_{\\Delta^*} & 0.2 \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "set_coefficients(standard_model)\n", - "\n", - "latex = R\"\\begin{array}{lc}\" + \"\\n\"\n", - "for par_name, value in parameter_table.items():\n", - " value = str(value).lstrip(\"(\").rstrip(\")\").replace(\"j\", \"i\")\n", - " symbol = sp.Symbol(par_name)\n", - " latex += Rf\" {sp.latex(symbol)} & {value} \\\\\" + \"\\n\"\n", - "latex += R\"\\end{array}\"\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Generate data" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from tensorwaves.data import SympyDataTransformer\n", - "from tensorwaves.function.sympy import create_function\n", - "\n", - "\n", - "def compute_sub_intensities(\n", - " model: HelicityModel, resonance_name: str, phsp, full_expression\n", - ") -> np.ndarray:\n", - " parameter_values = {}\n", - " for symbol, value in model.parameter_defaults.items():\n", - " if resonance_name not in symbol.name and symbol.name.startswith(\"C\"):\n", - " parameter_values[symbol] = 0\n", - " else:\n", - " parameter_values[symbol] = value\n", - " sub_expression = full_expression.subs(parameter_values)\n", - " sub_intensity = create_function(sub_expression, backend=\"jax\")\n", - " return np.array(sub_intensity(phsp).real)\n", - "\n", - "\n", - "def plot_distributions(model: HelicityModel) -> None:\n", - " helicity_transformer = SympyDataTransformer.from_sympy(\n", - " model.kinematic_variables, backend=\"jax\"\n", - " )\n", - " phsp = helicity_transformer(phsp_momenta)\n", - " phsp = {k: v.real for k, v in phsp.items()}\n", - "\n", - " full_expression = model.expression.doit()\n", - " substituted_expression = full_expression.xreplace(model.parameter_defaults)\n", - " intensity_func = create_function(substituted_expression, backend=\"jax\")\n", - " intensities_all = np.array(intensity_func(phsp).real)\n", - " intensities_k = compute_sub_intensities(model, \"K^*\", phsp, full_expression)\n", - " intensities_delta = compute_sub_intensities(model, \"Delta^*\", phsp, full_expression)\n", - " intensities_lambda = compute_sub_intensities(\n", - " model, \"Lambda^*\", phsp, full_expression\n", - " )\n", - "\n", - " fig, ax = plt.subplots(nrows=2, ncols=3, figsize=(8, 5))\n", - " hist_kwargs = {\n", - " \"bins\": 80,\n", - " \"histtype\": \"step\",\n", - " }\n", - "\n", - " for x in ax.flatten():\n", - " x.set_yticks([])\n", - "\n", - " ax[0, 0].set_xlabel(\"$m^2(pK^-)$ [GeV$^2/c^4$]\")\n", - " ax[0, 1].set_xlabel(R\"$m^2(K^-\\pi^+)$ [GeV$^2/c^4$]\")\n", - " ax[0, 2].set_xlabel(R\"$m^2(p\\pi^+)$ [GeV$^2/c^4$]\")\n", - " ax[1, 0].set_xlabel(R\"$\\cos\\theta(p)$\")\n", - " ax[1, 1].set_xlabel(R\"$\\phi(p)$\")\n", - " ax[1, 2].set_xlabel(R\"$\\chi$\")\n", - "\n", - " for x, xticks in {\n", - " ax[0, 0]: [2, 2.5, 3, 3.5, 4, 4.5],\n", - " ax[0, 1]: [0.4, 0.6, 0.8, 1, 1.2, 1.4, 1.6, 1.8, 2],\n", - " ax[0, 2]: [1, 1.5, 2, 2.5, 3],\n", - " ax[1, 0]: [-1, -0.5, 0, 0.5, 1],\n", - " ax[1, 1]: [-3, -2, -1, 0, 1, 2, 3],\n", - " }.items():\n", - " x.set_xticks(xticks)\n", - " x.set_xticklabels(xticks)\n", - "\n", - " for weights, color, label in [\n", - " (intensities_all, \"red\", \"Model\"),\n", - " (intensities_k, \"orange\", R\"$K^*\\to\\,K^{^-}\\pi^+$\"),\n", - " (intensities_delta, \"brown\", R\"$\\Delta^{*^{++}} \\to\\,p\\pi^+$\"),\n", - " (intensities_lambda, \"purple\", R\"$\\Lambda^* \\to\\,p K^{^-}$\"),\n", - " ]:\n", - " kwargs = dict(weights=weights, color=color, **hist_kwargs)\n", - " ax[0, 0].hist(np.array(phsp[\"m_01\"] ** 2), **kwargs)\n", - " ax[0, 1].hist(np.array(phsp[\"m_12\"] ** 2), **kwargs)\n", - " ax[0, 2].hist(np.array(phsp[\"m_02\"] ** 2), **kwargs)\n", - " ax[1, 0].hist(np.array(np.cos(phsp[\"theta_01\"])), **kwargs)\n", - " ax[1, 1].hist(np.array(phsp[\"phi_01\"]), **kwargs, label=label)\n", - "\n", - " ax[1, 2].remove()\n", - " handles, labels = ax[1, 1].get_legend_handles_labels()\n", - " fig.legend(handles, labels, loc=\"lower right\")\n", - "\n", - " ax[0, 2].set_xlim(1, 3.4)\n", - " ax[1, 0].set_xlim(-1, +1)\n", - " ax[1, 1].set_xlim(-np.pi, +np.pi)\n", - "\n", - " fig.tight_layout()\n", - "\n", - " plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "\n", - "It takes several minutes to lambdify the full expression and expressions for the Wigner rotation angles.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "plot_distributions(standard_model)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![013_23_0](https://user-images.githubusercontent.com/29308176/164991359-ee679063-0571-4beb-9ff4-b15e6dbf65a0.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Spin alignment sum" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, with the spin alignment sum from [ampform#245](https://github.com/ComPWA/ampform/pull/245) inserted:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{m_{A}=-1/2}^{1/2} \\sum_{m_{0}=-1/2}^{1/2} \\sum_{m_{1}=0} \\sum_{m_{2}=0}{\\left|{\\sum_{\\lambda^{01}_{0}=-1/2}^{1/2} \\sum_{\\mu^{01}_{0}=-1/2}^{1/2} \\sum_{\\nu^{01}_{0}=-1/2}^{1/2} \\sum_{\\lambda^{01}_{1}=0} \\sum_{\\mu^{01}_{1}=0} \\sum_{\\nu^{01}_{1}=0} \\sum_{\\lambda^{01}_{2}=0}{{A^{01}}_{m_{A},\\lambda^{01}_{0},- \\lambda^{01}_{1},- \\lambda^{01}_{2}} D^{0}_{m_{1},\\nu^{01}_{1}}\\left(\\alpha^{01}_{1},\\beta^{01}_{1},\\gamma^{01}_{1}\\right) D^{0}_{m_{2},\\lambda^{01}_{2}}\\left(\\phi_{01},\\theta_{01},0\\right) D^{0}_{\\mu^{01}_{1},\\lambda^{01}_{1}}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{0}_{\\nu^{01}_{1},\\mu^{01}_{1}}\\left(\\phi_{01},\\theta_{01},0\\right) D^{\\frac{1}{2}}_{m_{0},\\nu^{01}_{0}}\\left(\\alpha^{01}_{0},\\beta^{01}_{0},\\gamma^{01}_{0}\\right) D^{\\frac{1}{2}}_{\\mu^{01}_{0},\\lambda^{01}_{0}}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{\\frac{1}{2}}_{\\nu^{01}_{0},\\mu^{01}_{0}}\\left(\\phi_{01},\\theta_{01},0\\right)} + \\sum_{\\lambda^{02}_{0}=-1/2}^{1/2} \\sum_{\\mu^{02}_{0}=-1/2}^{1/2} \\sum_{\\nu^{02}_{0}=-1/2}^{1/2} \\sum_{\\lambda^{02}_{1}=0} \\sum_{\\lambda^{02}_{2}=0} \\sum_{\\mu^{02}_{2}=0} \\sum_{\\nu^{02}_{2}=0}{{A^{02}}_{m_{A},\\lambda^{02}_{0},- \\lambda^{02}_{1},- \\lambda^{02}_{2}} D^{0}_{m_{1},\\lambda^{02}_{1}}\\left(\\phi_{02},\\theta_{02},0\\right) D^{0}_{m_{2},\\nu^{02}_{2}}\\left(\\alpha^{02}_{2},\\beta^{02}_{2},\\gamma^{02}_{2}\\right) D^{0}_{\\mu^{02}_{2},\\lambda^{02}_{2}}\\left(\\phi^{02}_{0},\\theta^{02}_{0},0\\right) D^{0}_{\\nu^{02}_{2},\\mu^{02}_{2}}\\left(\\phi_{02},\\theta_{02},0\\right) D^{\\frac{1}{2}}_{m_{0},\\nu^{02}_{0}}\\left(\\alpha^{02}_{0},\\beta^{02}_{0},\\gamma^{02}_{0}\\right) D^{\\frac{1}{2}}_{\\mu^{02}_{0},\\lambda^{02}_{0}}\\left(\\phi^{02}_{0},\\theta^{02}_{0},0\\right) D^{\\frac{1}{2}}_{\\nu^{02}_{0},\\mu^{02}_{0}}\\left(\\phi_{02},\\theta_{02},0\\right)} + \\sum_{\\lambda^{12}_{0}=-1/2}^{1/2} \\sum_{\\lambda^{12}_{1}=0} \\sum_{\\mu^{12}_{1}=0} \\sum_{\\nu^{12}_{1}=0} \\sum_{\\lambda^{12}_{2}=0} \\sum_{\\mu^{12}_{2}=0} \\sum_{\\nu^{12}_{2}=0}{{A^{12}}_{m_{A},\\lambda^{12}_{0},\\lambda^{12}_{1},- \\lambda^{12}_{2}} D^{0}_{m_{1},\\nu^{12}_{1}}\\left(\\alpha^{12}_{1},\\beta^{12}_{1},\\gamma^{12}_{1}\\right) D^{0}_{m_{2},\\nu^{12}_{2}}\\left(\\alpha^{12}_{2},\\beta^{12}_{2},\\gamma^{12}_{2}\\right) D^{0}_{\\mu^{12}_{1},\\lambda^{12}_{1}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{0}_{\\mu^{12}_{2},\\lambda^{12}_{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{0}_{\\nu^{12}_{1},\\mu^{12}_{1}}\\left(\\phi_{0},\\theta_{0},0\\right) D^{0}_{\\nu^{12}_{2},\\mu^{12}_{2}}\\left(\\phi_{0},\\theta_{0},0\\right) D^{\\frac{1}{2}}_{m_{0},\\lambda^{12}_{0}}\\left(\\phi_{0},\\theta_{0},0\\right)}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(A^01[m_A, lambda_0^01, -lambda_1^01, -lambda_2^01]*WignerD(0, m1, nu_1^01, alpha_1^01, beta_1^01, gamma_1^01)*WignerD(0, m2, lambda_2^01, phi_01, theta_01, 0)*WignerD(0, mu_1^01, lambda_1^01, phi_0^01, theta_0^01, 0)*WignerD(0, nu_1^01, mu_1^01, phi_01, theta_01, 0)*WignerD(1/2, m0, nu_0^01, alpha_0^01, beta_0^01, gamma_0^01)*WignerD(1/2, mu_0^01, lambda_0^01, phi_0^01, theta_0^01, 0)*WignerD(1/2, nu_0^01, mu_0^01, phi_01, theta_01, 0), (lambda_0^01, (-1/2, 1/2)), (mu_0^01, (-1/2, 1/2)), (nu_0^01, (-1/2, 1/2)), (lambda_1^01, (0,)), (mu_1^01, (0,)), (nu_1^01, (0,)), (lambda_2^01, (0,))) + PoolSum(A^02[m_A, lambda_0^02, -lambda_1^02, -lambda_2^02]*WignerD(0, m1, lambda_1^02, phi_02, theta_02, 0)*WignerD(0, m2, nu_2^02, alpha_2^02, beta_2^02, gamma_2^02)*WignerD(0, mu_2^02, lambda_2^02, phi_0^02, theta_0^02, 0)*WignerD(0, nu_2^02, mu_2^02, phi_02, theta_02, 0)*WignerD(1/2, m0, nu_0^02, alpha_0^02, beta_0^02, gamma_0^02)*WignerD(1/2, mu_0^02, lambda_0^02, phi_0^02, theta_0^02, 0)*WignerD(1/2, nu_0^02, mu_0^02, phi_02, theta_02, 0), (lambda_0^02, (-1/2, 1/2)), (mu_0^02, (-1/2, 1/2)), (nu_0^02, (-1/2, 1/2)), (lambda_1^02, (0,)), (lambda_2^02, (0,)), (mu_2^02, (0,)), (nu_2^02, (0,))) + PoolSum(A^12[m_A, lambda_0^12, lambda_1^12, -lambda_2^12]*WignerD(0, m1, nu_1^12, alpha_1^12, beta_1^12, gamma_1^12)*WignerD(0, m2, nu_2^12, alpha_2^12, beta_2^12, gamma_2^12)*WignerD(0, mu_1^12, lambda_1^12, phi_1^12, theta_1^12, 0)*WignerD(0, mu_2^12, lambda_2^12, phi_1^12, theta_1^12, 0)*WignerD(0, nu_1^12, mu_1^12, phi_0, theta_0, 0)*WignerD(0, nu_2^12, mu_2^12, phi_0, theta_0, 0)*WignerD(1/2, m0, lambda_0^12, phi_0, theta_0, 0), (lambda_0^12, (-1/2, 1/2)), (lambda_1^12, (0,)), (mu_1^12, (0,)), (nu_1^12, (0,)), (lambda_2^12, (0,)), (mu_2^12, (0,)), (nu_2^12, (0,))))**2, (m_A, (1/2, -1/2)), (m0, (1/2, -1/2)), (m1, (0,)), (m2, (0,)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "builder.align_spin = True\n", - "aligned_model = builder.formulate()\n", - "set_coefficients(aligned_model)\n", - "aligned_model.intensity" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "plot_distributions(aligned_model)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![013_28_0](https://user-images.githubusercontent.com/29308176/164991360-35450331-9174-4abe-9715-0a07dbb164ac.svg)\n", - "\n", - "Compare with [Figure 2](https://downloads.hindawi.com/journals/ahep/2020/6674595.pdf#page=9). Note that the distributions differ close to threshold, because the distributions in the paper are produced [with form factors](https://ampform.readthedocs.io/en/0.12.x/api/ampform.dynamics.html#ampform.dynamics.relativistic_breit_wigner_with_ff) and an [energy-dependent width](https://ampform.readthedocs.io/en/0.12.x/api/ampform.dynamics.html#ampform.dynamics.EnergyDependentWidth)." - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/014.ipynb b/docs/report/014.ipynb deleted file mode 100644 index b97825cb..00000000 --- a/docs/report/014.ipynb +++ /dev/null @@ -1,1698 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Amplitude model with sum notation\n", - "TR-014\n", - "^^^\n", - "See also [TR-013](013.ipynb) and [TR-015](015.ipynb).\n", - "+++\n", - "✅ [ampform#245](https://github.com/ComPWA/ampform/issues/245)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Amplitude model with sum notation" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q qrules[viz]==0.9.7 sympy==1.9 git+https://github.com/ComPWA/ampform@3ed3ed5" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import inspect\n", - "import itertools\n", - "import logging\n", - "from collections import defaultdict\n", - "from collections.abc import Iterable, Sequence\n", - "from functools import cache\n", - "from typing import TYPE_CHECKING\n", - "\n", - "import ampform\n", - "import attrs\n", - "import graphviz\n", - "import qrules\n", - "import symplot\n", - "import sympy as sp\n", - "from ampform.dynamics.builder import (\n", - " ResonanceDynamicsBuilder,\n", - " create_non_dynamic,\n", - " create_relativistic_breit_wigner,\n", - ")\n", - "from ampform.helicity import (\n", - " _generate_kinematic_variable_set,\n", - " _generate_kinematic_variables,\n", - ")\n", - "from ampform.helicity.decay import TwoBodyDecay\n", - "from ampform.helicity.naming import HelicityAmplitudeNameGenerator\n", - "from ampform.sympy import (\n", - " UnevaluatedExpression,\n", - " create_expression,\n", - " implement_doit_method,\n", - ")\n", - "from IPython.display import Math, display\n", - "from qrules import ReactionInfo\n", - "from sympy.core.symbol import Str\n", - "from sympy.physics.quantum.spin import Rotation as Wigner\n", - "from sympy.printing.precedence import PRECEDENCE\n", - "\n", - "if TYPE_CHECKING:\n", - " from qrules.topology import Topology\n", - " from qrules.transition import StateTransition\n", - " from sympy.printing.latex import LatexPrinter\n", - "\n", - "LOGGER = logging.getLogger()\n", - "LOGGER.setLevel(logging.ERROR)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(tr-014-problem)=\n", - "## Problem description" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Challenge\n", - ":class: purple\n", - "\n", - "Formulate the [`HelicityModel`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel) in such a way that:\n", - "\n", - "1. The sum over the amplitudes is concise and expresses that the sum depends only on helicities.\n", - "2. It is the {class}`sympy.Expr ` of the amplitude model can easily and uniquely be constructed from the data in the [`HelicityModel`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel). (Currently, this is as simple as {meth}`HelicityModel.expression.doit `).\n", - "3. All parameters under [`parameter_defaults`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel.parameter_defaults) are of type {class}`sympy.Symbol `. This is important for a correct lambdification of the arguments with {func}`sympy.utilities.lambdify.lambdify`.\n", - "\n", - "This report presents two solutions:\n", - "- {ref}`tr-014-solution-1`\n", - "- {ref}`tr-014-solution-2`\n", - "\n", - ":::\n", - "\n", - "[ampform#245](https://github.com/ComPWA/ampform/pull/245) implements spin alignment, which results in large sum combinatorics for all helicity combinations. The result is an amplitude model expression that is too large to be rendered as LaTeX.\n", - "\n", - "To some extend, this is already the case with the [current implementation](https://ampform.readthedocs.io/en/0.12.3/usage/formalism.html) of the 'standard' helicity formalism {cite}`Jacob:1959at, Richman:1984gh, kutschkeAngularDistributionCookbook1996, chungSpinFormalismsUpdated2014`: many of the terms in the total intensity expression differ only by the helicities of the final and initial state." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# Simplify resonance notation\n", - "PDG = qrules.load_pdg()\n", - "delta_res = PDG[\"Delta(1600)++\"]\n", - "lambda_res = PDG[\"Lambda(1405)\"]\n", - "particles = set(PDG)\n", - "particles.remove(delta_res)\n", - "particles.remove(lambda_res)\n", - "particles.add(attrs.evolve(delta_res, latex=R\"\\Delta\"))\n", - "particles.add(attrs.evolve(lambda_res, latex=R\"\\Lambda\"))\n", - "MODIFIED_PDG = qrules.ParticleCollection(particles)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "reaction = qrules.generate_transitions(\n", - " initial_state=\"Lambda(c)+\",\n", - " final_state=[\"K-\", \"p\", \"pi+\"],\n", - " formalism=\"helicity\",\n", - " allowed_intermediate_particles=[\"Delta(1600)++\"],\n", - " particle_db=MODIFIED_PDG,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "display(*(graphviz.Source(qrules.io.asdot(t, size=3)) for t in reaction.transitions))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{container} full-width\n", - "![](https://user-images.githubusercontent.com/29308176/164992090-84c38c7d-9c1c-4e57-abef-1ad46d557eda.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992091-e24324cc-8e5c-46b6-a5ec-df1bd2b98cb3.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992092-f2b2d578-be50-4102-aa8f-780eea7cf8c7.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992095-75f21354-3e50-4919-b07c-9bb68ee71929.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992096-d18f2f49-d317-4874-b728-d14192cb876d.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992097-be7b64a7-0fa5-44a6-863a-120ccf648a5a.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992099-ccd9b237-27e7-46ac-88b5-69417f4102cc.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992101-99d6d3e4-5a30-49dc-888a-e9f6ee580eff.svg)\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "builder = ampform.get_builder(reaction)\n", - "model = builder.formulate()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "I = & \\left|{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \\\\\n", - "& + \\left|{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \\\\\n", - "& + \\left|{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \\\\\n", - "& + \\left|{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right) + D^{\\frac{1}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}\\right|^{2} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def remove_coefficients(expr: sp.Expr) -> sp.Expr:\n", - " coefficients = {s: 1 for s in expr.free_symbols if s.name.startswith(\"C_\")}\n", - " return expr.subs(coefficients)\n", - "\n", - "\n", - "model = builder.formulate()\n", - "full_expression = remove_coefficients(model.expression)\n", - "I = sp.Symbol(\"I\")\n", - "latex = sp.multiline_latex(I, full_expression)\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here, we did not insert any dynamics, but it is unusual that dynamics expressions depend on helicity or spin projection (see [`TwoBodyKinematicVariableSet`](https://ampform.readthedocs.io/en/stable/api/ampform.dynamics.builder.html#ampform.dynamics.builder.TwoBodyKinematicVariableSet))." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Simplified notation with `PoolSum`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Both {ref}`tr-014-solution-1` and {ref}`tr-014-solution-2` require the definition of a special \"`PoolSum`\" class to simplify the summation over the amplitudes. The class mimics {class}`sympy.Sum ` in that it substitutes certain {class}`~sympy.core.symbol.Symbol`s in an expression over which we symbol a range of values. The range of values in a `PoolSum` does not have to be a sequential range, but can be a collection of arbitrary items." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "@implement_doit_method\n", - "class PoolSum(UnevaluatedExpression):\n", - " precedence = PRECEDENCE[\"Add\"]\n", - "\n", - " def __new__(\n", - " cls,\n", - " expression: sp.Expr,\n", - " *indices: tuple[sp.Symbol, Iterable[sp.Float]],\n", - " **hints,\n", - " ) -> PoolSum:\n", - " indices = tuple((s, tuple(v)) for s, v in indices)\n", - " return create_expression(cls, expression, *indices, **hints)\n", - "\n", - " @property\n", - " def expression(self) -> sp.Expr:\n", - " return self.args[0]\n", - "\n", - " @property\n", - " def indices(self) -> list[tuple[sp.Symbol, tuple[sp.Float, ...]]]:\n", - " return self.args[1:]\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " indices = dict(self.indices)\n", - " return sum(\n", - " self.expression.subs(zip(indices, combi))\n", - " for combi in itertools.product(*indices.values())\n", - " )\n", - "\n", - " def _latex(self, printer: LatexPrinter, *args) -> str:\n", - " indices = dict(self.indices)\n", - " sum_symbols: list[str] = []\n", - " for idx, values in indices.items():\n", - " sum_symbols.append(_render_sum_symbol(printer, idx, values))\n", - " expression = printer._print(self.expression)\n", - " return R\" \".join(sum_symbols) + f\"{{{expression}}}\"\n", - "\n", - " def cleanup(self) -> sp.Expr | PoolSum:\n", - " substitutions = {}\n", - " new_indices = []\n", - " for idx, values in self.indices:\n", - " if idx not in self.expression.free_symbols:\n", - " continue\n", - " if len(values) == 0:\n", - " continue\n", - " if len(values) == 1:\n", - " substitutions[idx] = values[0]\n", - " else:\n", - " new_indices.append((idx, values))\n", - " new_expression = self.expression.xreplace(substitutions)\n", - " if len(new_indices) == 0:\n", - " return new_expression\n", - " return PoolSum(new_expression, *new_indices)\n", - "\n", - "\n", - "def _render_sum_symbol(\n", - " printer: LatexPrinter, idx: sp.Symbol, values: Sequence[float]\n", - ") -> str:\n", - " if len(values) == 0:\n", - " return \"\"\n", - " idx = printer._print(idx)\n", - " if len(values) == 1:\n", - " value = values[0]\n", - " return Rf\"\\sum_{{{idx}={value}}}\"\n", - " if _is_regular_series(values):\n", - " sorted_values = sorted(values)\n", - " first_value = sorted_values[0]\n", - " last_value = sorted_values[-1]\n", - " return Rf\"\\sum_{{{idx}={first_value}}}^{{{last_value}}}\"\n", - " idx_values = \",\".join(map(printer._print, values))\n", - " return Rf\"\\sum_{{{idx}\\in\\left\\{{{idx_values}\\right\\}}}}\"\n", - "\n", - "\n", - "def _is_regular_series(values: Sequence[float]) -> bool:\n", - " if len(values) <= 1:\n", - " return False\n", - " sorted_values = sorted(values)\n", - " for val, next_val in zip(sorted_values, sorted_values[1:]):\n", - " difference = float(next_val - val)\n", - " if difference != 1.0:\n", - " return False\n", - " return True" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here's a sketch of how to construct the amplitude model with a `PoolSum`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left. \\sum_{\\lambda_{\\Lambda_c}=-1/2}^{1/2} \\sum_{\\lambda_{p}=-1/2}^{1/2} \\sum_{\\lambda_{\\pi}=0} \\sum_{\\lambda_{K}=0}{\\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{s_{\\Delta}}_{\\lambda_{\\Delta},- \\lambda_{\\pi} + \\lambda_{p}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{s_{\\Lambda_c}}_{\\lambda_{\\Lambda_c},\\lambda_{K} - \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right)}}\\right|^{2}} \\right|_{\\substack{ s_{\\Lambda_c}=\\frac{1}{2}\\\\ s_{\\Delta}=\\frac{3}{2} }}$" - ], - "text/plain": [ - "Subs(PoolSum(Abs(PoolSum(WignerD(s_\\Delta, \\lambda_\\Delta, -\\lambda_\\pi + \\lambda_p, phi_1^12, theta_1^12, 0)*WignerD(s_{\\Lambda_c}, \\lambda_{\\Lambda_c}, \\lambda_K - \\lambda_\\Delta, phi_12, theta_12, 0), (\\lambda_\\Delta, (-1/2, 1/2))))**2, (\\lambda_{\\Lambda_c}, (-1/2, 1/2)), (\\lambda_p, (-1/2, 1/2)), (\\lambda_\\pi, (0,)), (\\lambda_K, (0,))), (s_{\\Lambda_c}, s_\\Delta), (1/2, 3/2))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda_{\\Lambda_c}=-1/2}^{1/2} \\sum_{\\lambda_{p}=-1/2}^{1/2}{\\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{s_{\\Delta}}_{\\lambda_{\\Delta},\\lambda_{p}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{s_{\\Lambda_c}}_{\\lambda_{\\Lambda_c},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right)}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(WignerD(s_\\Delta, \\lambda_\\Delta, \\lambda_p, phi_1^12, theta_1^12, 0)*WignerD(s_{\\Lambda_c}, \\lambda_{\\Lambda_c}, -\\lambda_\\Delta, phi_12, theta_12, 0), (\\lambda_\\Delta, (-1/2, 1/2))))**2, (\\lambda_{\\Lambda_c}, (-1/2, 1/2)), (\\lambda_p, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "I = & \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},- \\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n", - "& + \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},\\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n", - "& + \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},- \\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n", - "& + \\left|{\\sum_{\\lambda_{\\Delta}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\frac{1}{2},- \\lambda_{\\Delta}}\\left(\\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{\\Delta},\\frac{1}{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\frac{\\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\sin^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)}}{8} - \\frac{3 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\sin{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\sin{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{4} + \\frac{9 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\sin^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{\\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{3 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\cos{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{4} + \\frac{9 \\sin^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{\\sin^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)}}{8} - \\frac{3 \\sin{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\sin{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)}}{4} + \\frac{9 \\sin^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)}}{8} + \\frac{\\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)}}{8} + \\frac{3 \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos{\\left(\\frac{\\theta^{12}_{1}}{2} \\right)} \\cos{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{4} + \\frac{9 \\cos^{2}{\\left(\\frac{\\theta_{12}}{2} \\right)} \\cos^{2}{\\left(\\frac{3 \\theta^{12}_{1}}{2} \\right)}}{8}$" - ], - "text/plain": [ - "sin(theta_12/2)**2*sin(theta_1^12/2)**2/8 - 3*sin(theta_12/2)**2*sin(theta_1^12/2)*sin(3*theta_1^12/2)/4 + 9*sin(theta_12/2)**2*sin(3*theta_1^12/2)**2/8 + sin(theta_12/2)**2*cos(theta_1^12/2)**2/8 + 3*sin(theta_12/2)**2*cos(theta_1^12/2)*cos(3*theta_1^12/2)/4 + 9*sin(theta_12/2)**2*cos(3*theta_1^12/2)**2/8 + sin(theta_1^12/2)**2*cos(theta_12/2)**2/8 - 3*sin(theta_1^12/2)*sin(3*theta_1^12/2)*cos(theta_12/2)**2/4 + 9*sin(3*theta_1^12/2)**2*cos(theta_12/2)**2/8 + cos(theta_12/2)**2*cos(theta_1^12/2)**2/8 + 3*cos(theta_12/2)**2*cos(theta_1^12/2)*cos(3*theta_1^12/2)/4 + 9*cos(theta_12/2)**2*cos(3*theta_1^12/2)**2/8" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "half = sp.S.Half\n", - "\n", - "spin_parent = sp.Symbol(R\"s_{\\Lambda_c}\", real=True)\n", - "spin_resonance = sp.Symbol(R\"s_\\Delta\", real=True)\n", - "\n", - "phi_12, theta_12 = sp.symbols(\"phi_12 theta_12\", real=True)\n", - "phi_1_12, theta_1_12 = sp.symbols(R\"phi_1^12 theta_1^12\", real=True)\n", - "\n", - "lambda_parent = sp.Symbol(R\"\\lambda_{\\Lambda_c}\", real=True)\n", - "lambda_resonance = sp.Symbol(R\"\\lambda_\\Delta\", real=True)\n", - "lambda_p = sp.Symbol(R\"\\lambda_p\", real=True)\n", - "lambda_k = sp.Symbol(R\"\\lambda_K\", real=True)\n", - "lambda_pi = sp.Symbol(R\"\\lambda_\\pi\", real=True)\n", - "sum_expr = sp.Subs(\n", - " PoolSum(\n", - " sp.Abs(\n", - " PoolSum(\n", - " Wigner.D(\n", - " spin_parent,\n", - " lambda_parent,\n", - " lambda_k - lambda_resonance,\n", - " phi_12,\n", - " theta_12,\n", - " 0,\n", - " )\n", - " * Wigner.D(\n", - " spin_resonance,\n", - " lambda_resonance,\n", - " lambda_p - lambda_pi,\n", - " phi_1_12,\n", - " theta_1_12,\n", - " 0,\n", - " ),\n", - " (lambda_resonance, (-half, +half)),\n", - " )\n", - " )\n", - " ** 2,\n", - " (lambda_parent, (-half, +half)),\n", - " (lambda_p, (-half, +half)),\n", - " (lambda_pi, (0,)),\n", - " (lambda_k, (0,)),\n", - " ),\n", - " (spin_parent, spin_resonance),\n", - " (half, 3 * half),\n", - ")\n", - "display(\n", - " sum_expr,\n", - " sum_expr.expr.cleanup(),\n", - " Math(sp.multiline_latex(I, sum_expr.doit(deep=False))),\n", - " sum_expr.doit(deep=False).doit(deep=True),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "assert sum_expr.doit(deep=False).doit() == full_expression.doit()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(tr-014-solution-1)=\n", - "## Solution 1: `Indexed` coefficients" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{margin}\n", - "\n", - "When introducing spin alignment ([ampform#245](https://github.com/ComPWA/ampform/pull/245)), we have to distinguish the helicity symbols between different topologies.\n", - "\n", - ":::\n", - "\n", - "The current implementation of the [`HelicityAmplitudeBuilder`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityAmplitudeBuilder) has to be changed quite a bit to produce an amplitude model with `PoolSum`s. First of all, we have to introduce special {class}`~sympy.core.symbol.Symbol`s for the helicities, $\\lambda_i$, with $i$ the state ID (taking a sum of attached final state IDs in case of a resonance ID). Next, [`formulate_wigner_d()`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.formulate_wigner_d) has to be modified to insert these {class}`~sympy.core.symbol.Symbol`s into the {class}`~sympy.physics.quantum.spin.WignerD`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "def formulate_wigner_d(transition: StateTransition, node_id: int) -> sp.Expr:\n", - " from sympy.physics.quantum.spin import Rotation as Wigner\n", - "\n", - " decay = TwoBodyDecay.from_transition(transition, node_id)\n", - " topology = transition.topology\n", - " parent_helicity = create_helicity_symbol(topology, decay.parent.id)\n", - " child1_helicity = create_helicity_symbol(topology, decay.children[0].id)\n", - " child2_helicity = create_helicity_symbol(topology, decay.children[1].id)\n", - " _, phi, theta = _generate_kinematic_variables(transition, node_id)\n", - " return Wigner.D(\n", - " j=sp.Rational(decay.parent.particle.spin),\n", - " m=parent_helicity,\n", - " mp=child1_helicity - child2_helicity,\n", - " alpha=-phi,\n", - " beta=theta,\n", - " gamma=0,\n", - " )\n", - "\n", - "\n", - "def create_helicity_symbol(topology: Topology, state_id: int) -> sp.Symbol:\n", - " if state_id in topology.incoming_edge_ids:\n", - " suffix = \"\"\n", - " else:\n", - " suffix = f\"_{state_id}\"\n", - " return sp.Symbol(f\"lambda{suffix}\", rational=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)$" - ], - "text/plain": [ - "WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0)" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "wigner_functions = {\n", - " sp.Mul(*[\n", - " formulate_wigner_d(transition, node_id) for node_id in transition.topology.nodes\n", - " ])\n", - " for transition in reaction.transitions\n", - "}\n", - "display(*wigner_functions)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We also have to collect the allowed helicity values for each of these helicity symbols." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "if TYPE_CHECKING:\n", - " from typing import Literal\n", - "\n", - "\n", - "@cache\n", - "def get_helicities(\n", - " reaction: ReactionInfo, which: Literal[\"inner\", \"outer\"]\n", - ") -> dict[int, set[sp.Rational]]:\n", - " helicities = defaultdict(set)\n", - " initial_state_ids = set(reaction.initial_state)\n", - " final_state_ids = set(reaction.final_state)\n", - " intermediate_state_ids = (\n", - " set(reaction.transitions[0].states) - initial_state_ids - final_state_ids\n", - " )\n", - " if which == \"inner\":\n", - " state_ids = sorted(intermediate_state_ids)\n", - " elif which == \"outer\":\n", - " state_ids = sorted(initial_state_ids | final_state_ids)\n", - " for transition in reaction.transitions:\n", - " for state_id in state_ids:\n", - " state = transition.states[state_id]\n", - " helicity = sp.Rational(state.spin_projection)\n", - " symbol = create_helicity_symbol(transition.topology, state_id)\n", - " helicities[symbol].add(helicity)\n", - " return dict(helicities)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{lambda_3: {-1/2, 1/2}}" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "{lambda: {-1/2, 1/2}, lambda_0: {0}, lambda_1: {-1/2, 1/2}, lambda_2: {0}}" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "inner_helicities = get_helicities(reaction, which=\"inner\")\n", - "outer_helicities = get_helicities(reaction, which=\"outer\")\n", - "display(inner_helicities, outer_helicities)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "These collected helicity values can then be combined with the Wigner-$D$ expressions through a `PoolSum`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def formulate_intensity(reaction: ReactionInfo):\n", - " wigner_functions = {\n", - " sp.Mul(*[\n", - " formulate_wigner_d(transition, node_id)\n", - " for node_id in transition.topology.nodes\n", - " ])\n", - " for transition in reaction.transitions\n", - " }\n", - " inner_helicities = get_helicities(reaction, which=\"inner\")\n", - " outer_helicities = get_helicities(reaction, which=\"outer\")\n", - " return PoolSum(\n", - " sp.Abs(\n", - " PoolSum(\n", - " sum(wigner_functions),\n", - " *inner_helicities.items(),\n", - " )\n", - " )\n", - " ** 2,\n", - " *outer_helicities.items(),\n", - " )\n", - "\n", - "\n", - "formulate_intensity(reaction)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This is indeed identical to the model as formulated with the existing implementation:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "assert formulate_intensity(reaction).doit() == full_expression.doit()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note how this approach also works in case there are two decay topologies:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "reaction_two_resonances = qrules.generate_transitions(\n", - " initial_state=\"Lambda(c)+\",\n", - " final_state=[\"K-\", \"p\", \"pi+\"],\n", - " formalism=\"helicity\",\n", - " allowed_intermediate_particles=[\"Lambda(1405)\", \"Delta(1600)++\"],\n", - " particle_db=MODIFIED_PDG,\n", - ")\n", - "assert len(reaction_two_resonances.transition_groups) == 2\n", - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", - "display(*[\n", - " graphviz.Source(\n", - " qrules.io.asdot(g, collapse_graphs=True, size=4, render_resonance_id=True)\n", - " )\n", - " for g in reaction_two_resonances.transition_groups\n", - "])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{container} full-width\n", - "![](https://user-images.githubusercontent.com/29308176/164992102-fcac3af8-285a-49e8-b830-58fc947fef30.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992104-f37f1a89-3cf0-43bb-a013-9f20e3064aed.svg)\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{1}=-1/2}^{1/2}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\lambda,\\lambda_{3}}\\left(- \\phi_{01},\\theta_{01},0\\right) D^{\\frac{1}{2}}_{\\lambda_{3},- \\lambda_{1}}\\left(- \\phi^{01}_{0},\\theta^{01}_{0},0\\right) + D^{\\frac{1}{2}}_{\\lambda,\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(WignerD(1/2, lambda, lambda_3, -phi_01, theta_01, 0)*WignerD(1/2, lambda_3, -lambda_1, -phi_0^01, theta_0^01, 0) + WignerD(1/2, lambda, lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1, -phi_1^12, theta_1^12, 0), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_1, (-1/2, 1/2)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "formulate_intensity(reaction_two_resonances).cleanup()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "ampform_expr = ampform.get_builder(reaction_two_resonances).formulate().expression\n", - "ampform_expr = remove_coefficients(ampform_expr)\n", - "sum_expr = formulate_intensity(reaction_two_resonances)\n", - "assert ampform_expr.doit() == sum_expr.doit()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Inserting coefficients" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "There is a problem, though: the sums above could be written in sum form because the values over which we sum appear as arguments ({attr}`~sympy.core.basic.Basic.args` in the {class}`~sympy.physics.quantum.spin.WignerD` functions). This is only true because we previously set all coefficients to $1$. The coefficient _names_ are, however, also 'dependent' on the helicities in the final state over which we sum:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "[C_{\\Lambda_{c}^{+} \\to \\Delta_{-1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}},\n", - " C_{\\Lambda_{c}^{+} \\to \\Delta_{+1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}}]" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sorted(model.parameter_defaults, key=str)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We therefore have to somehow introduce a dependence in these {class}`~sympy.core.symbol.Symbol`s on the helicity values. An idea may be to use {class}`~sympy.tensor.indexed.IndexedBase`. Modifying the function introduced in {ref}`tr-014-solution-1`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2}} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(C[lambda_0, lambda_1, lambda_2]*WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "C = sp.IndexedBase(\"C\")\n", - "\n", - "\n", - "@cache\n", - "def formulate_intensity_with_coefficient(reaction: ReactionInfo):\n", - " amplitudes = {\n", - " sp.Mul(\n", - " C[\n", - " [\n", - " create_helicity_symbol(transition.topology, state_id)\n", - " for state_id in transition.final_states\n", - " ]\n", - " ],\n", - " *[\n", - " formulate_wigner_d(transition, node_id)\n", - " for node_id in transition.topology.nodes\n", - " ],\n", - " )\n", - " for transition in reaction.transitions\n", - " }\n", - " inner_helicities = get_helicities(reaction, which=\"inner\")\n", - " outer_helicities = get_helicities(reaction, which=\"outer\")\n", - " return PoolSum(\n", - " sp.Abs(\n", - " PoolSum(\n", - " sum(amplitudes),\n", - " *inner_helicities.items(),\n", - " )\n", - " )\n", - " ** 2,\n", - " *outer_helicities.items(),\n", - " )\n", - "\n", - "\n", - "indexed_coefficient_expr = formulate_intensity_with_coefficient(reaction)\n", - "indexed_coefficient_expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "I = & \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,- \\frac{1}{2},0} D^{\\frac{1}{2}}_{- \\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n", - "& + \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,- \\frac{1}{2},0} D^{\\frac{1}{2}}_{\\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n", - "& + \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,\\frac{1}{2},0} D^{\\frac{1}{2}}_{- \\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \\\\\n", - "& + \\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{C}_{0,\\frac{1}{2},0} D^{\\frac{1}{2}}_{\\frac{1}{2},\\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}}\\right|^{2} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "latex = sp.multiline_latex(I, indexed_coefficient_expr.doit(deep=False))\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Caveat\n", - "\n", - "Using {class}`~sympy.tensor.indexed.IndexedBase` makes the coefficient names concise, but harder to understand.\n", - "\n", - ":::\n", - "\n", - "This seems to work rather well, but there is a subtle problems introduced by writing the coefficients as a {class}`~sympy.tensor.indexed.IndexedBase`: the {class}`~sympy.tensor.indexed.IndexedBase` itself is considered listed under the {attr}`~sympy.core.basic.Basic.free_symbols` of the expression." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "[C, C[0, -1/2, 0], C[0, 1/2, 0], phi_12, phi_1^12, theta_12, theta_1^12]" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "free_symbols = sorted(indexed_coefficient_expr.doit().free_symbols, key=str)\n", - "free_symbols" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In addition, not all symbols in the expression are of type {class}`~sympy.core.symbol.Symbol` anymore:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{C: (sympy.core.symbol.Symbol, True),\n", - " C[0, -1/2, 0]: (sympy.tensor.indexed.Indexed, False),\n", - " C[0, 1/2, 0]: (sympy.tensor.indexed.Indexed, False),\n", - " phi_12: (sympy.core.symbol.Symbol, True),\n", - " phi_1^12: (sympy.core.symbol.Symbol, True),\n", - " theta_12: (sympy.core.symbol.Symbol, True),\n", - " theta_1^12: (sympy.core.symbol.Symbol, True)}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "{s: (type(s), isinstance(s, sp.Symbol)) for s in free_symbols}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This will become problematic when lambdifying, because it results in an additional argument in the signature of the generated function:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "func = sp.lambdify(free_symbols, indexed_coefficient_expr.doit())\n", - "inspect.signature(func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A solution may be to use [`symplot.substitute_indexed_symbols()`](https://ampform.readthedocs.io/en/0.12.3/usage/symplot.html#symplot.substitute_indexed_symbols):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{C_{0,-1/2,0}, C_{0,1/2,0}, phi_12, phi_1^12, theta_12, theta_1^12}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "indexed_coefficient_expr_symbols_only = symplot.substitute_indexed_symbols(\n", - " indexed_coefficient_expr.doit()\n", - ")\n", - "indexed_coefficient_expr_symbols_only.free_symbols" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "args = sorted(indexed_coefficient_expr_symbols_only.free_symbols, key=str)\n", - "func = sp.lambdify(args, indexed_coefficient_expr_symbols_only)\n", - "inspect.signature(func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Caveat\n", - "\n", - "This seems clumsy, because [`substitute_indexed_symbols()`](https://ampform.readthedocs.io/en/0.12.3/usage/symplot.html#symplot.substitute_indexed_symbols) would have to be actively called before creating a computational function with [TensorWaves](https://tensorwaves.rtfd.io). It also becomes a hassle to keep track of the correct {class}`~sympy.core.symbol.Symbol` names in [`HelicityModel.parameter_defaults`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityModel.parameter_defaults).\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "# One topology\n", - "expr = ampform.get_builder(reaction).formulate().expression\n", - "expr = remove_coefficients(expr.doit())\n", - "sum_expr = formulate_intensity_with_coefficient(reaction)\n", - "sum_expr = symplot.substitute_indexed_symbols(sum_expr.doit())\n", - "sum_expr = remove_coefficients(sum_expr)\n", - "assert sum_expr == expr\n", - "\n", - "# Two topologies\n", - "expr = ampform.get_builder(reaction_two_resonances).formulate().expression\n", - "expr = remove_coefficients(expr.doit())\n", - "sum_expr = formulate_intensity_with_coefficient(reaction_two_resonances)\n", - "sum_expr = symplot.substitute_indexed_symbols(sum_expr.doit())\n", - "sum_expr = remove_coefficients(sum_expr)\n", - "assert sum_expr == expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Inserting dynamics" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Dynamics pose a challenge that is similar to {ref}`report/014:Inserting coefficients` in that we have to introduce expressions that are dependent on spin. Still, as can be seen from the available attributes on a [`TwoBodyKinematicVariableSet`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.dynamics.builder.html#ampform.dynamics.builder.TwoBodyKinematicVariableSet) (which serves as input to [`ResonanceDynamicsBuilder`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.dynamics.builder.html#ampform.dynamics.builder.ResonanceDynamicsBuilder)s), dynamics (currently) cannot depend on helicities.\n", - "\n", - "What may become a problem are $LS$-combinations. So far we have only considered a [`ReactionInfo`](https://qrules.readthedocs.io/en/0.9.6/api/qrules.transition.html#qrules.transition.ReactionInfo) that was created with `formalism=\"helicity\"`, but we also have to sum over $LS$-combinations when using `formalism=\"canonical-helicity\"`. This is particularly important when using dynamics with form factors, which depend on [`angular_momentum`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.dynamics.builder.html#ampform.dynamics.builder.TwoBodyKinematicVariableSet.angular_momentum).\n", - "\n", - "\n", - ":::{note}\n", - "\n", - "The {class}`sympy.tensor.indexed.Indexed` now also contains the names of the resonances.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_intensity_with_dynamics(\n", - " reaction: ReactionInfo,\n", - " dynamics_choices: dict[str, ResonanceDynamicsBuilder],\n", - "):\n", - " amplitudes = set()\n", - " for transition in reaction.transitions:\n", - " final_state_helicities = [\n", - " create_helicity_symbol(transition.topology, state_id)\n", - " for state_id in transition.final_states\n", - " ]\n", - " resonances = [\n", - " Str(s.particle.latex) for s in transition.intermediate_states.values()\n", - " ]\n", - " indices = [*final_state_helicities, *resonances]\n", - " coefficient = C[indices]\n", - " expr: sp.Expr = coefficient\n", - " for node_id in sorted(transition.topology.nodes):\n", - " expr *= formulate_wigner_d(transition, node_id)\n", - " decay = TwoBodyDecay.from_transition(transition, node_id)\n", - " parent_particle = decay.parent.particle\n", - " dynamics_builder = dynamics_choices.get(\n", - " parent_particle.name, create_non_dynamic\n", - " )\n", - " variables = _generate_kinematic_variable_set(transition, node_id)\n", - " dynamics, _ = dynamics_builder(parent_particle, variables)\n", - " expr *= dynamics\n", - " amplitudes.add(expr)\n", - " inner_helicities = get_helicities(reaction, which=\"inner\")\n", - " outer_helicities = get_helicities(reaction, which=\"outer\")\n", - " return PoolSum(\n", - " sp.Abs(\n", - " PoolSum(sum(amplitudes), *inner_helicities.items()),\n", - " evaluate=False,\n", - " )\n", - " ** 2,\n", - " *outer_helicities.items(),\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output", - "full-width" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{\\frac{\\Gamma_{\\Delta} m_{\\Delta} {C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2},\\Delta} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}}}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(\\Gamma_{\\Delta}*m_{\\Delta}*C[lambda_0, lambda_1, lambda_2, \\Delta]*WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0)/(-I*\\Gamma_{\\Delta}*m_{\\Delta} - m_12**2 + m_{\\Delta}**2), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "formulate_intensity_with_dynamics(\n", - " reaction,\n", - " dynamics_choices={\n", - " resonance.name: create_relativistic_breit_wigner\n", - " for resonance in reaction.get_intermediate_particles()\n", - " },\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{\\frac{\\Gamma_{\\Delta} m_{\\Delta} {C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2},\\Delta} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{0} + \\lambda_{3}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\lambda_{3},\\lambda_{1} - \\lambda_{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} + \\frac{\\Gamma_{\\Lambda} m_{\\Lambda} {C}_{\\lambda_{0},\\lambda_{1},\\lambda_{2},\\Lambda} D^{\\frac{1}{2}}_{\\lambda,- \\lambda_{2} + \\lambda_{3}}\\left(- \\phi_{01},\\theta_{01},0\\right) D^{\\frac{1}{2}}_{\\lambda_{3},\\lambda_{0} - \\lambda_{1}}\\left(- \\phi^{01}_{0},\\theta^{01}_{0},0\\right)}{- i \\Gamma_{\\Lambda} m_{\\Lambda} - m_{01}^{2} + m_{\\Lambda}^{2}}}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(\\Gamma_{\\Delta}*m_{\\Delta}*C[lambda_0, lambda_1, lambda_2, \\Delta]*WignerD(1/2, lambda, -lambda_0 + lambda_3, -phi_12, theta_12, 0)*WignerD(3/2, lambda_3, lambda_1 - lambda_2, -phi_1^12, theta_1^12, 0)/(-I*\\Gamma_{\\Delta}*m_{\\Delta} - m_12**2 + m_{\\Delta}**2) + \\Gamma_{\\Lambda}*m_{\\Lambda}*C[lambda_0, lambda_1, lambda_2, \\Lambda]*WignerD(1/2, lambda, -lambda_2 + lambda_3, -phi_01, theta_01, 0)*WignerD(1/2, lambda_3, lambda_0 - lambda_1, -phi_0^01, theta_0^01, 0)/(-I*\\Gamma_{\\Lambda}*m_{\\Lambda} - m_01**2 + m_{\\Lambda}**2), (lambda_3, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (lambda_0, (0,)), (lambda_1, (-1/2, 1/2)), (lambda_2, (0,)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "formulate_intensity_with_dynamics(\n", - " reaction_two_resonances,\n", - " dynamics_choices={\n", - " resonance.name: create_relativistic_breit_wigner\n", - " for resonance in reaction_two_resonances.get_intermediate_particles()\n", - " },\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The resulting amplitude is again identical to the original [`HelicityModel.expression`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# One topology\n", - "b = ampform.get_builder(reaction)\n", - "for p in reaction.get_intermediate_particles():\n", - " b.set_dynamics(p.name, create_relativistic_breit_wigner)\n", - "expr = b.formulate().expression\n", - "expr = remove_coefficients(expr.doit())\n", - "sum_expr = formulate_intensity_with_dynamics(\n", - " reaction,\n", - " dynamics_choices={\n", - " resonance.name: create_relativistic_breit_wigner\n", - " for resonance in reaction.get_intermediate_particles()\n", - " },\n", - ")\n", - "sum_expr = symplot.substitute_indexed_symbols(sum_expr.doit())\n", - "sum_expr = remove_coefficients(sum_expr)\n", - "assert sum_expr == expr\n", - "\n", - "# Two topologies\n", - "b = ampform.get_builder(reaction_two_resonances)\n", - "for p in reaction_two_resonances.get_intermediate_particles():\n", - " b.set_dynamics(p.name, create_relativistic_breit_wigner)\n", - "expr = b.formulate().expression\n", - "expr = remove_coefficients(expr)\n", - "sum_expr = formulate_intensity_with_dynamics(\n", - " reaction_two_resonances,\n", - " dynamics_choices={\n", - " resonance.name: create_relativistic_breit_wigner\n", - " for resonance in reaction_two_resonances.get_intermediate_particles()\n", - " },\n", - ")\n", - "sum_expr = symplot.partial_doit(sum_expr, doit_classes=(PoolSum,))\n", - "sum_expr = symplot.partial_doit(sum_expr, doit_classes=(PoolSum,)) # recurse\n", - "sum_expr = symplot.substitute_indexed_symbols(sum_expr)\n", - "sum_expr = remove_coefficients(sum_expr)\n", - "assert sum_expr.free_symbols == expr.free_symbols\n", - "for intensity1, intensity2 in zip(sum_expr.args, expr.args):\n", - " # Annoyingly, Abs is rewritten with conjugates when using PoolSum...\n", - " amp1 = intensity1.args[0]\n", - " amp2 = intensity2.rewrite(sp.conjugate).args[0]\n", - " amp1 = sp.factor(amp1, deep=True, fraction=False)\n", - " amp2 = sp.factor(amp2, deep=True, fraction=False)\n", - " assert amp1 == amp2" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(tr-014-solution-2)=\n", - "## Solution 2: `Indexed` amplitude components" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The main problem with {ref}`tr-014-solution-2` is that it requires changing coefficient {class}`~sympy.core.symbol.Symbol`s to instances of {class}`~sympy.tensor.indexed.Indexed`, which have to be substituted using [`substitute_indexed_symbols()`](https://ampform.readthedocs.io/en/0.12.3/usage/symplot.html#symplot.substitute_indexed_symbols) (**after** calling {meth}`~sympy.core.basic.Basic.doit`).\n", - "\n", - "An alternative would be insert dynamics (and coefficients) into the `PoolSum`s over the helicities is to index the amplitude itself. The [`HelicityModel.expression`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression) would then contain {class}`~sympy.tensor.indexed.Indexed` symbols that represent specific amplitudes. A definition of these amplitudes can be provided through [`HelicityModel.components`](https://ampform.readthedocs.io/en/0.12.3/api/ampform.helicity.html#ampform.helicity.HelicityModel.components) or an equivalent attribute." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "A = sp.IndexedBase(R\"\\mathcal{A}\")\n", - "\n", - "\n", - "def formulate_intensity_indexed_amplitudes_only(\n", - " reaction: ReactionInfo,\n", - " dynamics_choices: dict[str, ResonanceDynamicsBuilder],\n", - ") -> tuple[sp.Expr, dict[sp.Indexed, sp.Expr]]:\n", - " name_generator = HelicityAmplitudeNameGenerator()\n", - " amplitudes = set()\n", - " amplitude_definitions = {}\n", - " for transition in reaction.transitions:\n", - " name_generator.register_amplitude_coefficient_name(transition)\n", - " for transition in reaction.transitions:\n", - " suffix = name_generator.generate_sequential_amplitude_suffix(transition)\n", - " expr: sp.Expr = sp.Symbol(f\"C_{{{suffix}}}\")\n", - " for node_id in sorted(transition.topology.nodes):\n", - " expr *= ampform.helicity.formulate_wigner_d(transition, node_id)\n", - " decay = TwoBodyDecay.from_transition(transition, node_id)\n", - " parent_particle = decay.parent.particle\n", - " dynamics_builder = dynamics_choices.get(\n", - " parent_particle.name, create_non_dynamic\n", - " )\n", - " variables = _generate_kinematic_variable_set(transition, node_id)\n", - " dynamics, _ = dynamics_builder(parent_particle, variables)\n", - " expr *= dynamics\n", - " resonances = [\n", - " Str(s.particle.latex) for s in transition.intermediate_states.values()\n", - " ]\n", - " helicity_symbols = [\n", - " create_helicity_symbol(transition.topology, state_id)\n", - " for state_id in sorted(transition.states)\n", - " ]\n", - " helicities = [\n", - " sp.Rational(transition.states[state_id].spin_projection)\n", - " for state_id in sorted(transition.states)\n", - " ]\n", - " amplitudes.add(A[[*helicity_symbols, *resonances]])\n", - " amplitude_definitions[A[[*helicities, *resonances]]] = expr\n", - " inner_helicities = get_helicities(reaction, which=\"inner\")\n", - " outer_helicities = get_helicities(reaction, which=\"outer\")\n", - " expression = PoolSum(\n", - " sp.Abs(\n", - " PoolSum(sum(amplitudes), *inner_helicities.items()),\n", - " evaluate=False,\n", - " )\n", - " ** 2,\n", - " *outer_helicities.items(),\n", - " )\n", - " return expression, amplitude_definitions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "expression, amplitudes = formulate_intensity_indexed_amplitudes_only(\n", - " reaction_two_resonances,\n", - " dynamics_choices={\n", - " resonance.name: create_relativistic_breit_wigner\n", - " for resonance in reaction_two_resonances.get_intermediate_particles()\n", - " },\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "I = & \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\lambda_{0}=0} \\sum_{\\lambda_{1}=-1/2}^{1/2} \\sum_{\\lambda_{2}=0}{\\left|{\\sum_{\\lambda_{3}=-1/2}^{1/2}{{\\mathcal{A}}_{\\lambda,\\lambda_{0},\\lambda_{1},\\lambda_{2},\\lambda_{3},\\Delta} + {\\mathcal{A}}_{\\lambda,\\lambda_{0},\\lambda_{1},\\lambda_{2},\\lambda_{3},\\Lambda}}}\\right|^{2}} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "I = & \\left|{{\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \\\\\n", - "& + \\left|{{\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \\\\\n", - "& + \\left|{{\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \\\\\n", - "& + \\left|{{\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Lambda} + {\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Delta} + {\\mathcal{A}}_{\\frac{1}{2},0,\\frac{1}{2},0,\\frac{1}{2},\\Lambda}}\\right|^{2} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "{\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,- \\frac{1}{2},\\Delta} = & \\frac{C_{\\Lambda_{c}^{+} \\to \\Delta_{-1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}} \\Gamma_{\\Delta} m_{\\Delta} D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "{\\mathcal{A}}_{- \\frac{1}{2},0,- \\frac{1}{2},0,\\frac{1}{2},\\Delta} = & \\frac{C_{\\Lambda_{c}^{+} \\to \\Delta_{+1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}} \\Gamma_{\\Delta} m_{\\Delta} D^{\\frac{1}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{\\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "{\\mathcal{A}}_{- \\frac{1}{2},0,\\frac{1}{2},0,- \\frac{1}{2},\\Delta} = & \\frac{C_{\\Lambda_{c}^{+} \\to \\Delta_{-1/2} K^{-}_{0}; \\Delta \\to p_{+1/2} \\pi^{+}_{0}} \\Gamma_{\\Delta} m_{\\Delta} D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi_{12},\\theta_{12},0\\right) D^{\\frac{3}{2}}_{- \\frac{1}{2},\\frac{1}{2}}\\left(- \\phi^{12}_{1},\\theta^{12}_{1},0\\right)}{- i \\Gamma_{\\Delta} m_{\\Delta} - m_{12}^{2} + m_{\\Delta}^{2}} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "display(Math(sp.multiline_latex(I, expression)))\n", - "display(Math(sp.multiline_latex(I, expression.doit())))\n", - "for i, (symbol, expr) in enumerate(amplitudes.items(), 1):\n", - " latex = sp.multiline_latex(symbol, expr)\n", - " display(Math(latex))\n", - " if i == 3:\n", - " break" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$\\dots$" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The resulting amplitude is indeed identical to the original [`HelicityModel.expression`](https://ampform.readthedocs.io/en/stable/api/ampform.helicity.html#ampform.helicity.HelicityModel.expression):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "b = ampform.get_builder(reaction_two_resonances)\n", - "for resonance in reaction_two_resonances.get_intermediate_particles():\n", - " b.set_dynamics(resonance.name, create_relativistic_breit_wigner)\n", - "model_two_res = b.formulate()\n", - "\n", - "assert model_two_res.expression == expression.doit().xreplace(amplitudes)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "\n", - "Currently, amplitudes with different resonances are put under a different amplitude symbol, identified by that resonance. Such resonances can be combined, e.g. $\\mathcal{A}_{\\lambda_i} = \\mathcal{A}_{\\lambda_i,\\Delta} + \\mathcal{A}_{\\lambda_i,\\Lambda}$. This would also make it easier to introduce correct interference terms through the [$K$-matrix](https://ampform.readthedocs.io/en/0.12.3/usage/dynamics/k-matrix.html).\n", - "\n", - ":::\n", - "\n", - ":::{admonition} Question\n", - "\n", - "The helicity of the intermediate state is also passed to the indexed amplitude. This is required for the coefficient name, which has a helicity subscript for the intermediate state, e.g. $C_{\\Lambda_{c}^{+} \\to \\Lambda_{\\color{red}+1/2} \\pi^{+}_{0}; \\Lambda \\to K^{-}_{0} p_{+1/2}}$. Does it really make sense to distinguish coefficients for different helicities of intermediate states?\n", - "\n", - ":::" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/015.ipynb b/docs/report/015.ipynb deleted file mode 100644 index 98a4bb7f..00000000 --- a/docs/report/015.ipynb +++ /dev/null @@ -1,1646 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Spin alignment implementation\n", - "TR-015\n", - "^^^\n", - "This report has been implemented through [ampform#245](https://github.com/ComPWA/ampform/pull/245). For details on how to use it, see [this notebook](https://ampform.readthedocs.io/en/0.13.3/usage/helicity/spin-alignment.html). See also [TR-013](013.ipynb) and [TR-014](014.ipynb).\n", - "+++\n", - "✅ [ampform#245](https://github.com/ComPWA/ampform/pull/245)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Spin alignment implementation" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.13.3 qrules[viz]==0.9.7 sympy==1.10" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import inspect\n", - "import logging\n", - "import warnings\n", - "\n", - "import ampform\n", - "import graphviz\n", - "import numpy as np\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.helicity import (\n", - " formulate_helicity_rotation_chain,\n", - " formulate_rotation_chain,\n", - " formulate_spin_alignment,\n", - " formulate_wigner_d,\n", - ")\n", - "from ampform.kinematics import (\n", - " compute_boost_chain,\n", - " compute_wigner_angles,\n", - " compute_wigner_rotation_matrix,\n", - " create_four_momentum_symbols,\n", - ")\n", - "from IPython.display import Math, display\n", - "from qrules.topology import create_isobar_topologies\n", - "\n", - "LOGGER = logging.getLogger()\n", - "LOGGER.setLevel(logging.ERROR)\n", - "warnings.filterwarnings(\"ignore\")\n", - "\n", - "\n", - "def show_transition(transition, **kwargs):\n", - " if \"size\" not in kwargs:\n", - " kwargs[\"size\"] = 5\n", - " dot = qrules.io.asdot(transition, **kwargs)\n", - " display(graphviz.Source(dot))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Helicity formalism" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Imagine we want to formulate the amplitude for the following **single** {external+qrules-0.9.x:class}`.StateTransition`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "full_reaction = qrules.generate_transitions(\n", - " initial_state=\"J/psi(1S)\",\n", - " final_state=[\"K0\", (\"Sigma+\", [+0.5]), (\"p~\", [+0.5])],\n", - " allowed_intermediate_particles=[\"Sigma(1660)~-\", \"N(1650)+\"],\n", - " allowed_interaction_types=\"strong\",\n", - " formalism=\"helicity\",\n", - ")\n", - "graphs = full_reaction.to_graphs()\n", - "single_transition_reaction = full_reaction.from_graphs(\n", - " [graphs[0]], formalism=full_reaction.formalism\n", - ")\n", - "transition = single_transition_reaction.transitions[0]\n", - "show_transition(transition)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992510-063aab30-aad8-4339-9152-46ce41da13c0.svg)\n", - "\n", - "The specific {attr}`~qrules.transition.State.spin_projection`s for each {attr}`~qrules.transition.State.particle` only make sense _given a specific reference frame_. AmpForm's {class}`~ampform.helicity.HelicityAmplitudeBuilder` interprets these projections as the **helicity** $\\lambda=\\vec{S}\\cdot\\vec{p}$ of each particle _in the rest frame of the parent particle_. For example, the helicity $\\lambda_2=+\\tfrac{1}{2}$ of $\\bar p$ is the helicity as measured in the rest frame of resonance $\\bar\\Sigma(1660)^-$. The reason is that these helicities are needed when formulating the two-particle state for the decay node $\\bar\\Sigma(1660)^- \\to K^0\\bar p$ (see {external+ampform-0.14.x:doc}`ampform:usage/helicity/formalism`).\n", - "\n", - "Ignoring dynamics and coefficients, the {class}`~ampform.helicity.HelicityModel` for this single transition is rather simple:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left|{D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{02}_{0},\\theta^{02}_{0},0\\right) D^{1}_{-1,-1}\\left(- \\phi_{02},\\theta_{02},0\\right)}\\right|^{2}$" - ], - "text/plain": [ - "Abs(WignerD(1/2, -1/2, -1/2, -phi_0^02, theta_0^02, 0)*WignerD(1, -1, -1, -phi_02, theta_02, 0))**2" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "builder = ampform.get_builder(single_transition_reaction)\n", - "model = builder.formulate()\n", - "model.expression.subs(model.parameter_defaults).subs(1.0, 1)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The two Wigner-$D$ functions come from the two **two-body decay nodes** that appear in the {external+qrules-0.9.x:class}`.StateTransition` above. They were formulated as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle D^{\\frac{1}{2}}_{- \\frac{1}{2},- \\frac{1}{2}}\\left(- \\phi^{02}_{0},\\theta^{02}_{0},0\\right) D^{1}_{-1,-1}\\left(- \\phi_{02},\\theta_{02},0\\right)$" - ], - "text/plain": [ - "WignerD(1/2, -1/2, -1/2, -phi_0^02, theta_0^02, 0)*WignerD(1, -1, -1, -phi_02, theta_02, 0)" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "sp.Mul(\n", - " formulate_wigner_d(transition, node_id=0),\n", - " formulate_wigner_d(transition, node_id=1),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, as {external+ampform-0.14.x:func}`~.formulate_wigner_d` explains, the numbers that appear in the Wigner-$D$ functions here are computed from the helicities of the decay products. But there's a subtle problem: these helicities are _assumed to be in the rest frame of each parent particle_. For the first node, this is fine, because the parent particle rest frame matches that of the initial state in the {external+qrules-0.9.x:class}`.StateTransition` above. In the second node, however, we are in a different rest frame. This can result in phase differences for the different amplitudes.\n", - "\n", - "If there is a single decay {class}`~qrules.topology.Topology` in the {class}`~qrules.transition.ReactionInfo` object for which we are formulating an amplitude model, the problem we identified here can be ignored. The reason is that the phase difference for each {external+qrules-0.9.x:class}`.StateTransition` (with each an identical decay {class}`~qrules.topology.Topology`) is the same and does not introduce interference effects within the coherent sum. It again becomes a problem, however, when we are formulating an amplitude model _with different topologies_. An example would be the following reaction:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "show_transition(full_reaction, collapse_graphs=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992511-98d8fa79-06dc-40ac-b91c-388ee2fb06f6.svg)\n", - "\n", - "\n", - "When formulating the amplitude model for this reaction, the {class}`~ampform.helicity.HelicityAmplitudeBuilder` implements the 'standard' helicity formalism as described in {cite}`Richman:1984gh, kutschkeAngularDistributionCookbook1996, chungSpinFormalismsUpdated2014` and simply sums over the different amplitudes to get the full amplitude:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "builder = ampform.get_builder(full_reaction)\n", - "model = builder.formulate()\n", - "latex = sp.multiline_latex(\n", - " sp.Symbol(\"I\"),\n", - " model.expression.subs(model.parameter_defaults).subs(1.0, 1),\n", - " environment=\"eqnarray\",\n", - ")\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As pointed out in {cite}`Marangotto:2019ucc, JPAC:2019ufm, Wang:2020giv`, this is wrong because of the mismatch in reference frames for the helicities." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Aligning reference frames" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the rest of this document, we follow {cite}`Marangotto:2019ucc` to align all amplitudes in the different topologies back to the initial state reference frame $A$, so that they can be correctly summed up. Specifically, we want to formulate a new, correctly aligned amplitude $\\mathcal{A}^{A\\to 0,1,\\dots}_{m_A,m_0,m_1,\\dots}$ from the original amplitudes $\\mathcal{A}^{A\\to R,S,i,...\\to 0,1,\\dots}_{\\lambda_A,\\lambda_0,\\lambda_1,\\dots}$ by applying Eq.(45) and Eq.(47) for generic, multi-body decays. Here, the $\\lambda$ values are the helicities in the parent rest frame of each two-body decay and the $m$ are the canonical[^canonical] spin projections in the rest frame of the mother particle that is the same no matter the {class}`~qrules.topology.Topology`.\n", - "\n", - "[^canonical]: The canonical rest frame differs from the 'helicity' rest frame in that it is reached by a direct Lorentz boost without rotating the coordinate system in such a way that the $z$-axis aligns with the momentum of one of the decay products.\n", - "\n", - "Just as in {cite}`Marangotto:2019ucc`, we test the implementation with 1-to-3 body decays. We use the notation from {func}`~ampform.helicity.naming.get_boost_chain_suffix` to indicate resonances $R,S,U$. This results in the following figure for the two alignments sums of Equations (45) and (46) in {cite}`Marangotto:2019ucc`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot1 = \"\"\"\n", - "digraph {\n", - " bgcolor=none\n", - " rankdir=LR\n", - " edge [arrowhead=none]\n", - " node [shape=none, width=0]\n", - " A\n", - " 0 [fontcolor=red]\n", - " 1 [fontcolor=green, label=<1>]\n", - " 2 [fontcolor=blue, label=<2>]\n", - " { rank=same A }\n", - " { rank=same 0, 1, 2 }\n", - " N0 [label=\"\"]\n", - " N1 [label=\"\"]\n", - " A -> N0 [style=dotted]\n", - " N0 -> N1 [label=\"R = 01\", fontcolor=orange]\n", - " N1 -> 0\n", - " N0 -> 2 [style=dashed]\n", - " N1 -> 1 [style=dashed]\n", - "}\n", - "\"\"\"\n", - "dot2 = \"\"\"\n", - "digraph {\n", - " bgcolor=none\n", - " rankdir=LR\n", - " edge [arrowhead=none]\n", - " node [shape=none, width=0]\n", - " A\n", - " 0 [label=0, fontcolor=red]\n", - " 1 [label=1, fontcolor=green, label=<1>]\n", - " 2 [label=2, fontcolor=blue, label=<2>]\n", - " { rank=same A }\n", - " { rank=same 0, 1, 2 }\n", - " N0 [label=\"\"]\n", - " N1 [label=\"\"]\n", - " A -> N0 [style=dotted]\n", - " N0 -> N1 [label=\"S = 02\", fontcolor=violet]\n", - " N1 -> 0\n", - " N0 -> 1 [style=dashed]\n", - " N1 -> 2 [style=dashed]\n", - "}\n", - "\"\"\"\n", - "display(*map(graphviz.Source, [dot1, dot2]))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992512-f99faee9-2ca2-415c-975d-07ce16326914.svg)\n", - "![](https://user-images.githubusercontent.com/29308176/164992514-d4c01c7c-7d2b-4842-8005-a8559523d001.svg)\n", - "\n", - "The dashed edges and bars above the state IDs indicate \"opposite helicity\" states. The helicity of an **opposite helicity state** gets a minus sign in the Wigner-$D$ function for a two-body state as formulated by {external+ampform-0.14.x:func}`.formulate_wigner_d` (see {ref}`report/015:Helicity formalism`) and therefore needs to be defined consistently. AmpForm does this with {external+ampform-0.14.x:func}`.is_opposite_helicity_state`.\n", - "\n", - "Opposite helicity states are also of importance in the spin alignment procedure sketched by {cite}`Marangotto:2019ucc`. The Wigner-$D$ functions that appear in Equations (45) and (46) from {cite}`Marangotto:2019ucc`, operate on the spin of the final state, but the angles in the Wigner-$D$ function are taken from the sibling state:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "\\begin{eqnarray}\n", - "\\mathcal{A}^{A \\to {\\color{orange}R},2 \\to 0,1,2}_{m_A,m_0,m_1,m_2}\n", - "&=&\n", - " \\sum_{\\lambda_0^{01},\\mu_0^{01},\\nu_0^{01}}\n", - " {\\color{red}{D^{s_0}_{m_0,\\nu_0^{01}}}}\\!\\left({\\color{red}{\\alpha_0^{01}, \\beta_0^{01}, \\gamma_0^{01}}}\\right)\n", - " {\\color{red}{D^{s_0}_{\\nu_0^{01},\\mu_0^{01}}}}\\!\\left({\\color{orange}{\\phi_{_{01}}, \\theta_{_{01}}}}, 0\\right)\n", - " {\\color{red}{D^{s_0}_{\\mu_0^{01},\\lambda_0^{01}}}}\\!\\left({\\color{red}{\\phi_0^{01}, \\theta_0^{01}}}\\right) \\\\\n", - "&\\times&\n", - " \\sum_{\\lambda_1^{01},\\mu_1^{01},\\nu_1^{01}}\n", - " {\\color{green}{D^{s_1}_{m_1,\\nu_1^{01}}}}\\!\\left({\\color{green}{\\alpha_1^{01}, \\beta_1^{01}, \\gamma_1^{01}}}\\right)\n", - " {\\color{green}{D^{s_1}_{\\nu_1^{01},\\mu_1^{01}}}}\\!\\left({\\color{orange}{\\phi_{_{01}}, \\theta_{_{01}}}}, 0\\right)\n", - " {\\color{green}{D^{s_1}_{\\mu_1^{01},\\lambda_1^{01}}}}\\!\\left({\\color{red}{\\phi_0^{01}, \\theta_0^{01}}}\\right) \\\\\n", - "&\\times&\n", - " \\sum_{\\lambda_2^{01}}\n", - " {\\color{blue}{D^{s_2}_{m_2,\\lambda_2^{01}}}}\\!\\left({\\color{orange}{\\phi_{_{01}}, \\theta_{_{01}}}}, 0\\right) \\\\\n", - "&\\times&\n", - " \\mathcal{A}^{A \\to {\\color{orange}R},2 \\to 0,1,2}_{m_A,\\lambda_0^{01},\\bar\\lambda_1^{01},\\bar\\lambda_2^{01}}\n", - "\\end{eqnarray}\n", - "$$ (alignment-R)\n", - "\n", - "$$\n", - "\\begin{eqnarray}\n", - "\\mathcal{A}^{A \\to {\\color{violet}S},1 \\to 0,1,2}_{m_A,m_0,m_1,m_2}\n", - "&=&\n", - " \\sum_{\\lambda_0^{02},\\mu_0^{02},\\nu_0^{02}}\n", - " {\\color{red}{D^{s_0}_{m_0,\\nu_0^{02}}}}\\!\\left({\\color{red}{\\alpha_0^{02}, \\beta_0^{02}, \\gamma_0^{02}}}\\right)\n", - " {\\color{red}{D^{s_0}_{\\nu_0^{02},\\mu_0^{02}}}}\\!\\left({\\color{violet}{\\phi_{_{02}}, \\theta_{_{02}}}}, 0\\right)\n", - " {\\color{red}{D^{s_0}_{\\mu_0^{02},\\lambda_0^{02}}}}\\!\\left({\\color{red}{\\phi_0^{02}, \\theta_0^{02}}}\\right) \\\\\n", - "&\\times&\n", - " \\sum_{\\lambda_1^{02}}\n", - " {\\color{green}{D^{s_1}_{m_1,\\lambda_1^{02}}}}\\!\\left({\\color{violet}{\\phi_{_{02}}, \\theta_{_{02}}}}, 0\\right) \\\\\n", - "&\\times&\n", - " \\sum_{\\lambda_2^{02},\\mu_2^{02},\\nu_2^{02}}\n", - " {\\color{blue}{D^{s_2}_{m_2,\\nu_2^{02}}}}\\!\\left({\\color{blue}{\\alpha_2^{02}, \\beta_2^{02}, \\gamma_2^{02}}}\\right)\n", - " {\\color{blue}{D^{s_2}_{\\nu_2^{02},\\mu_2^{02}}}}\\!\\left({\\color{violet}{\\phi_{_{02}}, \\theta_{_{02}}}}, 0\\right)\n", - " {\\color{blue}{D^{s_2}_{\\mu_2^{02},\\lambda_2^{02}}}}\\!\\left({\\color{red}{\\phi_0^{02}, \\theta_0^{02}}}\\right) \\\\\n", - "&\\times&\n", - " \\mathcal{A}^{A \\to {\\color{violet}S},2 \\to 0,1,2}_{m_A,\\lambda_0^{02},\\bar\\lambda_1^{02},\\bar\\lambda_2^{02}}\n", - "\\end{eqnarray}\n", - "$$ (alignment-S)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This procedure also allows us to formulate the alignment summation for $\\mathcal{A}^{A \\to {\\color{turquoise}U},0 \\to 0,1,2}_{m_A,m_0,m_1,m_2}$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot3 = \"\"\"\n", - "digraph {\n", - " bgcolor=none\n", - " rankdir=LR\n", - " edge [arrowhead=none]\n", - " node [shape=none, width=0]\n", - " 0 [shape=none, label=0, fontcolor=red]\n", - " 1 [shape=none, label=1, fontcolor=green]\n", - " 2 [shape=none, label=2, fontcolor=blue, label=<2>]\n", - " A [shape=none, label=A]\n", - " { rank=same A }\n", - " { rank=same 0, 1, 2 }\n", - " N0 [label=\"\"]\n", - " N1 [label=\"\"]\n", - " A -> N0 [style=dotted]\n", - " N0 -> N1 [label=12>, fontcolor=turquoise, style=dashed]\n", - " N0 -> 0\n", - " N1 -> 1\n", - " N1 -> 2 [style=dashed]\n", - "}\n", - "\"\"\"\n", - "graphviz.Source(dot3)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992516-0a53992d-8733-4d43-a6b9-510e0ec7e453.svg)\n", - "\n", - "$$\n", - "\\begin{eqnarray}\n", - "\\mathcal{A}^{A \\to {\\color{turquoise}U},0 \\to 0,1,2}_{m_A,m_0,m_1,m_2}\n", - "&=&\n", - " \\sum_{\\lambda_0^{12}}\n", - " {\\color{red}{D^{s_0}_{m_0,\\lambda_0^{12}}}}\\!\\left({\\color{red}{\\phi_0, \\theta_0}}, 0\\right) \\\\\n", - "&\\times&\n", - " \\sum_{\\lambda_1^{12},\\mu_1^{12},\\nu_1^{12}}\n", - " {\\color{green}{D^{s_0}_{m_0,\\nu_1^{12}}}}\\!\\left({\\color{green}{\\alpha_1^{12}, \\beta_1^{12}, \\gamma_1^{12}}}\\right)\n", - " {\\color{green}{D^{s_0}_{\\nu_1^{12},\\mu_1^{12}}}}\\!\\left({\\color{red}{\\phi_0, \\theta_0}}, 0\\right)\n", - " {\\color{green}{D^{s_0}_{\\mu_1^{12},\\lambda_1^{12}}}}\\!\\left({\\color{green}{\\phi_1^{12}, \\theta_1^{12}}}\\right) \\\\\n", - "&\\times&\n", - " \\sum_{\\lambda_2^{12},\\mu_2^{12},\\nu_2^{12}}\n", - " {\\color{blue}{D^{s_2}_{m_2,\\nu_2^{12}}}}\\!\\left({\\color{blue}{\\alpha_2^{12}, \\beta_2^{12}, \\gamma_2^{12}}}\\right)\n", - " {\\color{blue}{D^{s_2}_{\\nu_2^{12},\\mu_2^{12}}}}\\!\\left({\\color{red}{\\phi_0, \\theta_0}}, 0\\right)\n", - " {\\color{blue}{D^{s_2}_{\\mu_2^{12},\\lambda_2^{12}}}}\\!\\left({\\color{green}{\\phi_1^{12}, \\theta_1^{12}}}\\right) \\\\\n", - "&\\times&\n", - " \\mathcal{A}^{A \\to {\\color{turquoise}U},2 \\to 0,1,2}_{m_A,\\lambda_1^{12},\\bar\\lambda_1^{12},\\bar\\lambda_2^{12}}\n", - "\\end{eqnarray}\n", - "$$ (alignment-U)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, the total intensity can be computed from these amplitudes by incoherently summing over the initial and final state canonical spin projections (see [Equation (47)](https://downloads.hindawi.com/journals/ahep/2020/6674595.pdf#page=7) in {cite}`Marangotto:2019ucc`):\n", - "\n", - "$$\n", - "I = \\sum_{m_A,m_0,m_1,m_2}\\left|\n", - " \\mathcal{A}^{A \\to {\\color{orange}R},2 \\to 0,1,2}_{m_A,m_0,m_1,m_2} +\n", - " \\mathcal{A}^{A \\to {\\color{violet}S},1 \\to 0,1,2}_{m_A,m_0,m_1,m_2} +\n", - " \\mathcal{A}^{A \\to {\\color{turquoise}U},0 \\to 0,1,2}_{m_A,m_0,m_1,m_2}\n", - "\\right|^2\n", - "$$ (total-intensity)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### $J/\\psi \\to K^0 \\Sigma^+ \\bar{p}$" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "def show_all_spin_matrices(transition, functor, cleanup: bool) -> None:\n", - " for i in transition.final_states:\n", - " state = transition.states[i]\n", - " particle_name = state.particle.latex\n", - " s = sp.Rational(state.particle.spin)\n", - " m = sp.Rational(state.spin_projection)\n", - " display(Math(Rf\"|s_{i},m_{i}\\rangle=|{s},{m}\\rangle \\quad ({particle_name})\"))\n", - " if functor is formulate_rotation_chain:\n", - " args = (transition, i)\n", - " else:\n", - " args = (transition, i, state.spin_projection)\n", - " summation = functor(*args)\n", - " if cleanup:\n", - " summation = summation.cleanup()\n", - " display(summation)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In this section, we test some of the functions from the {mod}`~ampform.helicity` and {mod}`~ampform.kinematics` modules to see if they reproduce Equations {eq}`alignment-R`, {eq}`alignment-S`, and {eq}`alignment-U`. We perform this test on the channel $J/\\psi \\to K^0 \\Sigma^+ \\bar{p}$ with resonances generated for each of the three allowed three-body topologies. The transition that corresponds to Equation {eq}`alignment-R` is shown below.\n", - "\n", - "The first step is to use {external+ampform-0.14.x:func}`.formulate_helicity_rotation_chain` to generate the Wigner-$D$ functions for all **helicity rotations** for each final state. These helicity rotations \"undo\" all rotations that came from each Lorentz boosts when boosting from initial state $J/\\psi$ to each final state:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle |s_0,m_0\\rangle=|0,0\\rangle \\quad (K^{0})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle D^{0}_{0,0}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{0}_{\\nu^{01}_{0},0}\\left(\\phi_{01},\\theta_{01},0\\right)$" - ], - "text/plain": [ - "WignerD(0, 0, 0, phi_0^01, theta_0^01, 0)*WignerD(0, nu_0^01, 0, phi_01, theta_01, 0)" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_1,m_1\\rangle=|1/2,1/2\\rangle \\quad (\\Sigma^{+})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{01}_{1}=-1/2}^{1/2} \\sum_{\\mu^{01}_{1}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{\\mu^{01}_{1},\\lambda^{01}_{1}}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{\\frac{1}{2}}_{\\nu^{01}_{1},\\mu^{01}_{1}}\\left(\\phi_{01},\\theta_{01},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, mu_1^01, lambda_1^01, phi_0^01, theta_0^01, 0)*WignerD(1/2, nu_1^01, mu_1^01, phi_01, theta_01, 0), (lambda_1^01, (-1/2, 1/2)), (mu_1^01, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_2,m_2\\rangle=|1/2,1/2\\rangle \\quad (\\overline{p})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{01}_{2}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{0.5,\\lambda^{01}_{2}}\\left(\\phi_{01},\\theta_{01},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, 0.5, lambda_2^01, phi_01, theta_01, 0), (lambda_2^01, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "transition_r = full_reaction.transitions[-1]\n", - "show_all_spin_matrices(transition_r, formulate_helicity_rotation_chain, cleanup=True)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "show_transition(transition_r)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992517-98f11540-7eb3-47cb-b9ce-ecdf427451ab.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The function {external+ampform-0.14.x:func}`.formulate_rotation_chain` goes one step further. It adds a **Wigner rotation** to the generated list of helicity rotation Wigner-$D$ functions in case there are resonances in between the initial state and rotated final state. If there are no resonances in between (here, state `2`, the $\\bar p$), there is only one helicity rotation and there is no need for a Wigner rotation." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle |s_0,m_0\\rangle=|0,0\\rangle \\quad (K^{0})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{01}_{0}=0} \\sum_{\\mu^{01}_{0}=0} \\sum_{\\nu^{01}_{0}=0}{D^{0}_{m_{0},\\nu^{01}_{0}}\\left(\\alpha^{01}_{0},\\beta^{01}_{0},\\gamma^{01}_{0}\\right) D^{0}_{\\mu^{01}_{0},\\lambda^{01}_{0}}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{0}_{\\nu^{01}_{0},\\mu^{01}_{0}}\\left(\\phi_{01},\\theta_{01},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(0, m0, nu_0^01, alpha_0^01, beta_0^01, gamma_0^01)*WignerD(0, mu_0^01, lambda_0^01, phi_0^01, theta_0^01, 0)*WignerD(0, nu_0^01, mu_0^01, phi_01, theta_01, 0), (lambda_0^01, (0,)), (mu_0^01, (0,)), (nu_0^01, (0,)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_1,m_1\\rangle=|1/2,1/2\\rangle \\quad (\\Sigma^{+})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{01}_{1}=-1/2}^{1/2} \\sum_{\\mu^{01}_{1}=-1/2}^{1/2} \\sum_{\\nu^{01}_{1}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{m_{1},\\nu^{01}_{1}}\\left(\\alpha^{01}_{1},\\beta^{01}_{1},\\gamma^{01}_{1}\\right) D^{\\frac{1}{2}}_{\\mu^{01}_{1},\\lambda^{01}_{1}}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{\\frac{1}{2}}_{\\nu^{01}_{1},\\mu^{01}_{1}}\\left(\\phi_{01},\\theta_{01},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, m1, nu_1^01, alpha_1^01, beta_1^01, gamma_1^01)*WignerD(1/2, mu_1^01, lambda_1^01, phi_0^01, theta_0^01, 0)*WignerD(1/2, nu_1^01, mu_1^01, phi_01, theta_01, 0), (lambda_1^01, (-1/2, 1/2)), (mu_1^01, (-1/2, 1/2)), (nu_1^01, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_2,m_2\\rangle=|1/2,1/2\\rangle \\quad (\\overline{p})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{01}_{2}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{m_{2},\\lambda^{01}_{2}}\\left(\\phi_{01},\\theta_{01},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, m2, lambda_2^01, phi_01, theta_01, 0), (lambda_2^01, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "show_all_spin_matrices(transition_r, formulate_rotation_chain, cleanup=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "**These are indeed all the terms that we see in Equation {eq}`alignment-R`!**\n", - "\n", - "To create all sum combinations for all final states, we can use {external+ampform-0.14.x:func}`.formulate_spin_alignment`. This should give the sum of Eq.(45):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{01}_{1}=-1/2}^{1/2} \\sum_{\\mu^{01}_{1}=-1/2}^{1/2} \\sum_{\\nu^{01}_{1}=-1/2}^{1/2} \\sum_{\\lambda^{01}_{2}=-1/2}^{1/2}{D^{0}_{0,0}\\left(\\phi_{01},\\theta_{01},0\\right) D^{0}_{0,0}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{0}_{m_{0},0}\\left(\\alpha^{01}_{0},\\beta^{01}_{0},\\gamma^{01}_{0}\\right) D^{\\frac{1}{2}}_{m_{1},\\nu^{01}_{1}}\\left(\\alpha^{01}_{1},\\beta^{01}_{1},\\gamma^{01}_{1}\\right) D^{\\frac{1}{2}}_{m_{2},\\lambda^{01}_{2}}\\left(\\phi_{01},\\theta_{01},0\\right) D^{\\frac{1}{2}}_{\\mu^{01}_{1},\\lambda^{01}_{1}}\\left(\\phi^{01}_{0},\\theta^{01}_{0},0\\right) D^{\\frac{1}{2}}_{\\nu^{01}_{1},\\mu^{01}_{1}}\\left(\\phi_{01},\\theta_{01},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(0, 0, 0, phi_01, theta_01, 0)*WignerD(0, 0, 0, phi_0^01, theta_0^01, 0)*WignerD(0, m0, 0, alpha_0^01, beta_0^01, gamma_0^01)*WignerD(1/2, m1, nu_1^01, alpha_1^01, beta_1^01, gamma_1^01)*WignerD(1/2, m2, lambda_2^01, phi_01, theta_01, 0)*WignerD(1/2, mu_1^01, lambda_1^01, phi_0^01, theta_0^01, 0)*WignerD(1/2, nu_1^01, mu_1^01, phi_01, theta_01, 0), (lambda_1^01, (-1/2, 1/2)), (mu_1^01, (-1/2, 1/2)), (nu_1^01, (-1/2, 1/2)), (lambda_2^01, (-1/2, 1/2)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "alignment_summation = formulate_spin_alignment(transition_r)\n", - "alignment_summation.cleanup()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, here are the generated spin alignment terms for the other two decay chains. Notice that the first is indeed the same as {eq}`alignment-S`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle |s_0,m_0\\rangle=|0,0\\rangle \\quad (K^{0})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{02}_{0}=0} \\sum_{\\mu^{02}_{0}=0} \\sum_{\\nu^{02}_{0}=0}{D^{0}_{m_{0},\\nu^{02}_{0}}\\left(\\alpha^{02}_{0},\\beta^{02}_{0},\\gamma^{02}_{0}\\right) D^{0}_{\\mu^{02}_{0},\\lambda^{02}_{0}}\\left(\\phi^{02}_{0},\\theta^{02}_{0},0\\right) D^{0}_{\\nu^{02}_{0},\\mu^{02}_{0}}\\left(\\phi_{02},\\theta_{02},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(0, m0, nu_0^02, alpha_0^02, beta_0^02, gamma_0^02)*WignerD(0, mu_0^02, lambda_0^02, phi_0^02, theta_0^02, 0)*WignerD(0, nu_0^02, mu_0^02, phi_02, theta_02, 0), (lambda_0^02, (0,)), (mu_0^02, (0,)), (nu_0^02, (0,)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_1,m_1\\rangle=|1/2,1/2\\rangle \\quad (\\Sigma^{+})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{02}_{1}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{m_{1},\\lambda^{02}_{1}}\\left(\\phi_{02},\\theta_{02},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, m1, lambda_1^02, phi_02, theta_02, 0), (lambda_1^02, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_2,m_2\\rangle=|1/2,1/2\\rangle \\quad (\\overline{p})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{02}_{2}=-1/2}^{1/2} \\sum_{\\mu^{02}_{2}=-1/2}^{1/2} \\sum_{\\nu^{02}_{2}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{m_{2},\\nu^{02}_{2}}\\left(\\alpha^{02}_{2},\\beta^{02}_{2},\\gamma^{02}_{2}\\right) D^{\\frac{1}{2}}_{\\mu^{02}_{2},\\lambda^{02}_{2}}\\left(\\phi^{02}_{0},\\theta^{02}_{0},0\\right) D^{\\frac{1}{2}}_{\\nu^{02}_{2},\\mu^{02}_{2}}\\left(\\phi_{02},\\theta_{02},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, m2, nu_2^02, alpha_2^02, beta_2^02, gamma_2^02)*WignerD(1/2, mu_2^02, lambda_2^02, phi_0^02, theta_0^02, 0)*WignerD(1/2, nu_2^02, mu_2^02, phi_02, theta_02, 0), (lambda_2^02, (-1/2, 1/2)), (mu_2^02, (-1/2, 1/2)), (nu_2^02, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "reaction_s = qrules.generate_transitions(\n", - " initial_state=\"J/psi(1S)\",\n", - " final_state=[\"K0\", (\"Sigma+\", [+0.5]), (\"p~\", [+0.5])],\n", - " allowed_intermediate_particles=[\"N(1650)+\"],\n", - " allowed_interaction_types=\"strong\",\n", - " formalism=\"helicity\",\n", - ")\n", - "transition_s = reaction_s.transitions[0]\n", - "show_all_spin_matrices(transition_s, formulate_rotation_chain, cleanup=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "show_transition(transition_s)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992518-eeb93c12-642e-4094-a2c3-b72f2eb18de0.svg)\n", - "\n", - "...and that the second matches Equation {eq}`alignment-U`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle |s_0,m_0\\rangle=|0,0\\rangle \\quad (K^{0})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{12}_{0}=0}{D^{0}_{m_{0},\\lambda^{12}_{0}}\\left(\\phi_{0},\\theta_{0},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(0, m0, lambda_0^12, phi_0, theta_0, 0), (lambda_0^12, (0,)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_1,m_1\\rangle=|1/2,1/2\\rangle \\quad (\\Sigma^{+})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{12}_{1}=-1/2}^{1/2} \\sum_{\\mu^{12}_{1}=-1/2}^{1/2} \\sum_{\\nu^{12}_{1}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{m_{1},\\nu^{12}_{1}}\\left(\\alpha^{12}_{1},\\beta^{12}_{1},\\gamma^{12}_{1}\\right) D^{\\frac{1}{2}}_{\\mu^{12}_{1},\\lambda^{12}_{1}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{\\frac{1}{2}}_{\\nu^{12}_{1},\\mu^{12}_{1}}\\left(\\phi_{0},\\theta_{0},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, m1, nu_1^12, alpha_1^12, beta_1^12, gamma_1^12)*WignerD(1/2, mu_1^12, lambda_1^12, phi_1^12, theta_1^12, 0)*WignerD(1/2, nu_1^12, mu_1^12, phi_0, theta_0, 0), (lambda_1^12, (-1/2, 1/2)), (mu_1^12, (-1/2, 1/2)), (nu_1^12, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle |s_2,m_2\\rangle=|1/2,1/2\\rangle \\quad (\\overline{p})$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{12}_{2}=-1/2}^{1/2} \\sum_{\\mu^{12}_{2}=-1/2}^{1/2} \\sum_{\\nu^{12}_{2}=-1/2}^{1/2}{D^{\\frac{1}{2}}_{m_{2},\\nu^{12}_{2}}\\left(\\alpha^{12}_{2},\\beta^{12}_{2},\\gamma^{12}_{2}\\right) D^{\\frac{1}{2}}_{\\mu^{12}_{2},\\lambda^{12}_{2}}\\left(\\phi^{12}_{1},\\theta^{12}_{1},0\\right) D^{\\frac{1}{2}}_{\\nu^{12}_{2},\\mu^{12}_{2}}\\left(\\phi_{0},\\theta_{0},0\\right)}$" - ], - "text/plain": [ - "PoolSum(WignerD(1/2, m2, nu_2^12, alpha_2^12, beta_2^12, gamma_2^12)*WignerD(1/2, mu_2^12, lambda_2^12, phi_1^12, theta_1^12, 0)*WignerD(1/2, nu_2^12, mu_2^12, phi_0, theta_0, 0), (lambda_2^12, (-1/2, 1/2)), (mu_2^12, (-1/2, 1/2)), (nu_2^12, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "reaction_u = qrules.generate_transitions(\n", - " initial_state=\"J/psi(1S)\",\n", - " final_state=[\"K0\", (\"Sigma+\", [+0.5]), (\"p~\", [+0.5])],\n", - " allowed_intermediate_particles=[\"K*(1680)~0\"],\n", - " allowed_interaction_types=\"strong\",\n", - " formalism=\"helicity\",\n", - ")\n", - "transition_u = reaction_u.transitions[0]\n", - "show_all_spin_matrices(transition_u, formulate_rotation_chain, cleanup=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "show_transition(transition_u)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992519-e84aab1c-2895-481c-93fb-984f29f8b9e9.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Compute Wigner rotation angles" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now it's still a matter of computing the values for the angles $\\alpha,\\beta,\\gamma$ in the Wigner rotation matrices. These angles represents the difference between the canonical spin frame as attained by a direct boost from the initial state versus a chain of boosts through each resonance. See Equation (36) in {cite}`Marangotto:2019ucc`.\n", - "\n", - "The {mod}`~ampform.kinematics` module can generate an expression for the chain of Lorentz boosts from the initial state to the final state with {func}`~ampform.kinematics.lorentz.compute_boost_chain`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dot = qrules.io.asdot(transition_u)\n", - "topology = transition_u.topology\n", - "display(graphviz.Source(dot))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992519-e84aab1c-2895-481c-93fb-984f29f8b9e9.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\left[\\begin{matrix}\\boldsymbol{B}\\left(p_{0}\\right)\\end{matrix}\\right]$" - ], - "text/plain": [ - "[BoostMatrix(p0)]" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\left[\\begin{matrix}\\boldsymbol{B}\\left({p}_{12}\\right) & \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)\\end{matrix}\\right]$" - ], - "text/plain": [ - "[BoostMatrix(p1 + p2), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2), p1))]" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\left[\\begin{matrix}\\boldsymbol{B}\\left({p}_{12}\\right) & \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)\\end{matrix}\\right]$" - ], - "text/plain": [ - "[BoostMatrix(p1 + p2), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2), p2))]" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "momenta = create_four_momentum_symbols(topology)\n", - "for state_id in topology.outgoing_edge_ids:\n", - " boosts = compute_boost_chain(topology, momenta, state_id)\n", - " display(sp.Array(boosts))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This chain of Lorentz boosts needs to be 'undo' with a direct Lorentz boost back to the initial state. A contraction of inverse Lorentz boost with the chain of Lorentz boosts can be generated with {func}`~ampform.kinematics.angles.compute_wigner_rotation_matrix`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\boldsymbol{B}\\left(-\\left(p_{0}\\right)\\right) \\boldsymbol{B}\\left(p_{0}\\right)$" - ], - "text/plain": [ - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p0)), BoostMatrix(p0))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\boldsymbol{B}\\left(-\\left(p_{1}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)$" - ], - "text/plain": [ - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p1)), BoostMatrix(p1 + p2), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2), p1)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)$" - ], - "text/plain": [ - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p2)), BoostMatrix(p1 + p2), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2), p2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "for state_id in topology.outgoing_edge_ids:\n", - " expr = compute_wigner_rotation_matrix(topology, momenta, state_id)\n", - " display(expr)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The result of this matrix product is the rotation matrix for the Wigner rotation. The function {func}`~ampform.kinematics.angles.compute_wigner_angles` computes the required Euler angles from this rotation matrix by implementing Equations (B.2-3) from {cite}`Marangotto:2019ucc`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "angles = {}\n", - "for state_id in topology.outgoing_edge_ids:\n", - " angle_definitions = compute_wigner_angles(topology, momenta, state_id)\n", - " for name, expr in angle_definitions.items():\n", - " angle_symbol = sp.Symbol(name, real=True)\n", - " angles[angle_symbol] = expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\alpha^{12}_{0}&=&\\operatorname{atan_{2}}{\\left(\\boldsymbol{B}\\left(-\\left(p_{0}\\right)\\right) \\boldsymbol{B}\\left(p_{0}\\right)\\left[:, 3, 2\\right],\\boldsymbol{B}\\left(-\\left(p_{0}\\right)\\right) \\boldsymbol{B}\\left(p_{0}\\right)\\left[:, 3, 1\\right] \\right)}\\\\\n", - "\\beta^{12}_{0}&=&\\operatorname{acos}{\\left(\\boldsymbol{B}\\left(-\\left(p_{0}\\right)\\right) \\boldsymbol{B}\\left(p_{0}\\right)\\left[:, 3, 3\\right] \\right)}\\\\\n", - "\\gamma^{12}_{0}&=&\\operatorname{atan_{2}}{\\left(\\boldsymbol{B}\\left(-\\left(p_{0}\\right)\\right) \\boldsymbol{B}\\left(p_{0}\\right)\\left[:, 2, 3\\right],- \\boldsymbol{B}\\left(-\\left(p_{0}\\right)\\right) \\boldsymbol{B}\\left(p_{0}\\right)\\left[:, 1, 3\\right] \\right)}\\\\\n", - "\\alpha^{12}_{1}&=&\\operatorname{atan_{2}}{\\left(\\boldsymbol{B}\\left(-\\left(p_{1}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)\\left[:, 3, 2\\right],\\boldsymbol{B}\\left(-\\left(p_{1}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)\\left[:, 3, 1\\right] \\right)}\\\\\n", - "\\beta^{12}_{1}&=&\\operatorname{acos}{\\left(\\boldsymbol{B}\\left(-\\left(p_{1}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)\\left[:, 3, 3\\right] \\right)}\\\\\n", - "\\gamma^{12}_{1}&=&\\operatorname{atan_{2}}{\\left(\\boldsymbol{B}\\left(-\\left(p_{1}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)\\left[:, 2, 3\\right],- \\boldsymbol{B}\\left(-\\left(p_{1}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{1}\\right)\\left[:, 1, 3\\right] \\right)}\\\\\n", - "\\alpha^{12}_{2}&=&\\operatorname{atan_{2}}{\\left(\\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)\\left[:, 3, 2\\right],\\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)\\left[:, 3, 1\\right] \\right)}\\\\\n", - "\\beta^{12}_{2}&=&\\operatorname{acos}{\\left(\\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)\\left[:, 3, 3\\right] \\right)}\\\\\n", - "\\gamma^{12}_{2}&=&\\operatorname{atan_{2}}{\\left(\\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)\\left[:, 2, 3\\right],- \\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)\\left[:, 1, 3\\right] \\right)}\\\\\n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "latex_lines = [R\"\\begin{eqnarray}\"]\n", - "for symbol, expr in angles.items():\n", - " latex_lines.append(Rf\"{sp.latex(symbol)}&=&{sp.latex(expr)}\\\\\")\n", - "latex_lines.append(R\"\\end{eqnarray}\")\n", - "Math(\"\\n\".join(latex_lines))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "\n", - "In the topology underlying {eq}`alignment-U`, the Wigner rotation matrix with angles $\\alpha_0^{12}, \\beta_0^{12}, \\gamma_0^{12}$ is simply the identity matrix. This is the reason why it can be omitted in {external+ampform-0.14.x:func}`.formulate_rotation_chain` and we only have one helicity rotation.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Computational code" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Classes like {class}`~ampform.kinematics.lorentz.BoostMatrix` have been split up into smaller {func}`~ampform.sympy.unevaluated` expression classes so that lambdification to NumPy code results in relatively small and fast code, when using `cse=True` in {func}`~sympy.utilities.lambdify.lambdify` (see {class}`~ampform.sympy.NumPyPrintable`)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\beta^{12}_{1}:\\quad\\text{2,147 characters in generated code}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "beta = sp.Symbol(\"beta_1^12\", real=True)\n", - "beta_expr = angles[beta]\n", - "\n", - "func = sp.lambdify(momenta.values(), beta_expr.doit(), cse=True)\n", - "src = inspect.getsource(func)\n", - "n_characters = len(src)\n", - "latex = sp.latex(beta)\n", - "latex += Rf\":\\quad\\text{{{n_characters:,} characters in generated code}}\"\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Test on data sample" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "A test with a larger data distribution is being developed in [TR-013](./013.ipynb).\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The following phase space mini-sample of four-momenta has been generated for the decay $J/\\psi \\to K^0 \\Sigma^+ \\bar{p}$ with the [`tensorwaves.data`](https://tensorwaves.readthedocs.io/en/stable/api/tensorwaves.data.html) module." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "phsp = {\n", - " \"p0\": np.array([\n", - " [0.63140486, 0.13166435, -0.35734744, 0.07760603],\n", - " [0.65169531, 0.37242432, 0.12027178, 0.15467675],\n", - " [0.60647425, -0.22286205, -0.175258, 0.19952806],\n", - " [0.72744323, 0.05529811, 0.30502402, -0.43064999],\n", - " [0.76778868, -0.43557036, 0.35491651, -0.16185017],\n", - " ]),\n", - " \"p1\": np.array([\n", - " [1.37017117, 0.173769668, 0.355893315, -0.553093198],\n", - " [1.34556663, -5.272033e-04, -0.3074542, -0.54901747],\n", - " [1.41660182, 0.634007973, -0.0457976658, -0.433700564],\n", - " [1.38592340, 0.138369075, -0.258624859, 0.648189682],\n", - " [1.37858847, 0.551405385, -0.338705615, 0.259105737],\n", - " ]),\n", - " \"p2\": np.array([\n", - " [1.09532397, -0.30543402, 0.00145413, 0.47548716],\n", - " [1.09963805, -0.37189712, 0.18718247, 0.39434072],\n", - " [1.07382393, -0.41114592, 0.22105567, 0.2341725],\n", - " [0.98353336, -0.19366719, -0.04639917, -0.21753969],\n", - " [0.95052285, -0.11583502, -0.01621089, -0.09725557],\n", - " ]),\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\boldsymbol{B}\\left(-\\left(p_{2}\\right)\\right) \\boldsymbol{B}\\left({p}_{12}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{12}\\right) p_{2}\\right)$" - ], - "text/plain": [ - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p2)), BoostMatrix(p1 + p2), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2), p2)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "matrix_expr = compute_wigner_rotation_matrix(topology, momenta, state_id=2)\n", - "matrix_expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "array([[[ 1. , -0. , -0. , -0. ],\n", - " [-0. , 1. , 0.02, -0.02],\n", - " [-0. , -0.02, 1. , 0.03],\n", - " [ 0. , 0.02, -0.03, 1. ]],\n", - "\n", - " [[ 1. , 0. , -0. , 0. ],\n", - " [ 0. , 1. , -0.02, -0.04],\n", - " [ 0. , 0.02, 1. , -0. ],\n", - " [ 0. , 0.04, 0. , 1. ]],\n", - "\n", - " [[ 1. , 0. , 0. , 0. ],\n", - " [-0. , 1. , 0.02, -0.01],\n", - " [ 0. , -0.02, 1. , 0.02],\n", - " [ 0. , 0.01, -0.02, 1. ]],\n", - "\n", - " [[ 1. , 0. , -0. , 0. ],\n", - " [ 0. , 1. , -0.01, 0.02],\n", - " [ 0. , 0.01, 1. , 0.02],\n", - " [ 0. , -0.02, -0.02, 1. ]],\n", - "\n", - " [[ 1. , -0. , 0. , -0. ],\n", - " [-0. , 1. , -0.01, -0.01],\n", - " [ 0. , 0.01, 1. , 0.01],\n", - " [-0. , 0.01, -0.01, 1. ]]])" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "matrix_func = sp.lambdify(momenta.values(), matrix_expr.doit(), cse=True)\n", - "matrix_array = matrix_func(*phsp.values())\n", - "np.round(matrix_array, decimals=2).real" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\alpha^{12}_{0}&=&\\left[\\begin{matrix}-1.25 & 0.0 & 0.79 & 1.33 & 2.36\\end{matrix}\\right]\\\\\n", - "\\beta^{12}_{0}&=&\\left[\\begin{matrix}0.0 & \\text{NaN} & 0.0 & 0.0 & 0.0\\end{matrix}\\right]\\\\\n", - "\\gamma^{12}_{0}&=&\\left[\\begin{matrix}-1.89 & 3.14 & 2.36 & 1.82 & 0.79\\end{matrix}\\right]\\\\\n", - "\\alpha^{12}_{1}&=&\\left[\\begin{matrix}2.03 & -3.04 & 1.9 & 0.74 & 2.14\\end{matrix}\\right]\\\\\n", - "\\beta^{12}_{1}&=&\\left[\\begin{matrix}0.03 & 0.03 & 0.01 & 0.02 & 0.01\\end{matrix}\\right]\\\\\n", - "\\gamma^{12}_{1}&=&\\left[\\begin{matrix}-2.05 & 3.06 & -1.92 & -0.73 & -2.13\\end{matrix}\\right]\\\\\n", - "\\alpha^{12}_{2}&=&\\left[\\begin{matrix}-1.09 & 0.08 & -1.22 & -2.41 & -1.01\\end{matrix}\\right]\\\\\n", - "\\beta^{12}_{2}&=&\\left[\\begin{matrix}0.04 & 0.04 & 0.02 & 0.03 & 0.01\\end{matrix}\\right]\\\\\n", - "\\gamma^{12}_{2}&=&\\left[\\begin{matrix}1.11 & -0.1 & 1.25 & 2.4 & 1.0\\end{matrix}\\right]\\\\\n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "latex_lines = [R\"\\begin{eqnarray}\"]\n", - "for angle_symbol, angle_expr in angles.items():\n", - " angle_func = sp.lambdify(momenta.values(), angle_expr.doit(), cse=True)\n", - " angle_array = angle_func(*phsp.values())\n", - " rounded_values = np.round(angle_array, decimals=2).real\n", - " latex_lines.append(\n", - " Rf\"{sp.latex(angle_symbol)}&=&{sp.latex(sp.Array(rounded_values))}\\\\\"\n", - " )\n", - "latex_lines.append(R\"\\end{eqnarray}\")\n", - "Math(\"\\n\".join(latex_lines))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "dot = qrules.io.asdot(transition_u, collapse_graphs=True)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992522-be0a2ce4-3554-474a-b99c-2d901c07e247.svg)\n", - "\n", - ":::{note}\n", - "\n", - "The {obj}`~numpy.NAN` values above come from the fact that the inverse boost on a boost results in negative values under the square root of $\\gamma=\\sqrt{1-\\beta^2}$. This can be ignored, because the Wigner rotation is simply omitted when formulating the chain of rotation matrices, as noted in {ref}`report/015:Compute Wigner rotation angles`.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Four-body decay" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The algorithm for computing Euler angles for the Wigner rotation works an arbitrary number of final states. Here, we illustrate this by formulating an expression for the Wigner rotation matrix in a four-body decay." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\boldsymbol{B}\\left(-\\left(p_{3}\\right)\\right) \\boldsymbol{B}\\left({p}_{123}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{123}\\right) {p}_{23}\\right) \\boldsymbol{B}\\left(\\boldsymbol{B}\\left(\\boldsymbol{B}\\left({p}_{123}\\right) {p}_{23}\\right) \\boldsymbol{B}\\left({p}_{123}\\right) p_{3}\\right)$" - ], - "text/plain": [ - "MatrixMultiplication(BoostMatrix(NegativeMomentum(p3)), BoostMatrix(p1 + p2 + p3), BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 + p3), p2 + p3)), BoostMatrix(ArrayMultiplication(BoostMatrix(ArrayMultiplication(BoostMatrix(p1 + p2 + p3), p2 + p3)), ArrayMultiplication(BoostMatrix(p1 + p2 + p3), p3))))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "topology_4body = create_isobar_topologies(4)[1]\n", - "momenta_4body = create_four_momentum_symbols(topology_4body)\n", - "compute_wigner_rotation_matrix(topology_4body, momenta_4body, state_id=3)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dot = qrules.io.asdot(topology_4body)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/164992523-3eb56544-d1ea-4bdd-8f46-00e80354ff25.svg)" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/016.ipynb b/docs/report/016.ipynb deleted file mode 100644 index 26e5b05f..00000000 --- a/docs/report/016.ipynb +++ /dev/null @@ -1,509 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Symbolic integral\n", - "TR-016\n", - "^^^\n", - "This report investigates how to formulate a symbolic integral that correctly evaluates to\n", - "+++\n", - "To be implemented\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Lambdifying a symbolic integral" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.10 black==23.12.1 scipy==1.12.0 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "import inspect\n", - "\n", - "import black\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import sympy as sp\n", - "from ampform.io import aslatex\n", - "from ampform.sympy import unevaluated\n", - "from IPython.display import Markdown, Math\n", - "from scipy.integrate import quad, quad_vec\n", - "from sympy.printing.pycode import _unpack_integral_limits" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Numerical integration" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(quad)=\n", - "### SciPy's `quad()` function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SciPy's {func}`scipy.integrate.quad` cannot integrate complex-valued functions:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "def integrand(x):\n", - " return x * (x + 1j)\n", - "\n", - "\n", - "quad(integrand, 0.0, 2.0)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A [proposed solution](https://stackoverflow.com/a/5966088) is to wrap the {func}`~scipy.integrate.quad` function in a special integrate function that integrates the real and imaginary part of a function separately:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def complex_integrate(func, a, b, **quad_kwargs):\n", - " def real_func(x):\n", - " return func(x).real\n", - "\n", - " def imag_func(x):\n", - " return func(x).imag\n", - "\n", - " real_integral, real_integral_err = quad(real_func, a, b, **quad_kwargs)\n", - " imag_integral, imag_integral_err = quad(imag_func, a, b, **quad_kwargs)\n", - " return (\n", - " real_integral + 1j * imag_integral,\n", - " real_integral_err**2 + 1j * imag_integral_err,\n", - " )\n", - "\n", - "\n", - "complex_integrate(integrand, 0.0, 2.0)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "\n", - "The handling of uncertainties is incorrect.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(quad_vec)=\n", - "### SciPy's `quad_vec()` function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The easiest solution, however, seems to be {func}`scipy.integrate.quad_vec`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "quad_vec(integrand, 0.0, 2.0)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This has the added benefit that it can handle functions that return arrays:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def gaussian(x, mu, sigma):\n", - " return np.exp(-((x - mu) ** 2) / (2 * sigma**2)) / (sigma * np.sqrt(2 * np.pi))\n", - "\n", - "\n", - "mu_values = np.linspace(-2, +3, num=10)\n", - "result, _ = quad_vec(lambda x: gaussian(x, mu_values, sigma=0.5), 0, 2.0)\n", - "result" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Integrate with `quadpy`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "`quadpy` now requires a license. The examples below are only shown for documentation purposes.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[Alternatively](https://stackoverflow.com/a/42866568), one could use [`quadpy`](https://github.com/sigma-py/quadpy), which essentially does the same as in [`quad()`](#quad), but can also (to a large degree) handle vectorized input and properly handles uncertainties. For example:\n", - "\n", - "```python\n", - "from functools import partial\n", - "\n", - "\n", - "def parametrized_func(s_prime, s):\n", - " return s_prime * (s_prime + s + 1j)\n", - "\n", - "\n", - "s_array = np.linspace(-1, 1, num=10)\n", - "quadpy.quad(\n", - " partial(parametrized_func, s=s_array),\n", - " a=0.0,\n", - " b=2.0,\n", - ")\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "\n", - "One may need to play around with the tolerance if the function is not smooth, see [sigma-py/quadpy#255](https://github.com/sigma-py/quadpy/issues/255).\n", - "\n", - ":::\n", - "\n", - ":::{tip}\n", - "\n", - "[`quadpy`](https://github.com/sigma-py/quadpy) raises exceptions with {obj}`ModuleNotFoundError`s that are a bit unreadable. They are caused by the fact that [`orthopy`](https://pypi.org/project/orthopy) and [`ndim`](https://pypi.org/project/ndim) need to be installed separately.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## SymPy integral" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dispersion integral from Eq. {eq}`dispersion-integral` in **[TR-003](003.ipynb)** features a variable $s$ that is an argument to the function $\\Sigma_a$. This becomes a challenge when $s$ gets vectorized (in this case: gets an event-wise {obj}`numpy.array` of invariant masses). It seems that [`quad_vec()`](#quad_vec) can handle this well though." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "def parametrized_func(s_prime, s):\n", - " return s_prime * (s_prime + s + 1j)\n", - "\n", - "\n", - "s_array = np.linspace(-1, +1, num=10)\n", - "quad_vec(\n", - " lambda x: parametrized_func(x, s=s_array),\n", - " a=0.0,\n", - " b=2.0,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now attempt to design [SymPy](https://docs.sympy.org) expression classes that correctly {func}`~sympy.utilities.lambdify.lambdify` using this **vectorized** numerical integral for handles complex values. Note that this integral expression class derives from {class}`sympy.Integral ` and that:\n", - "\n", - "1. overwrites its {meth}`~sympy.core.basic.Basic.doit` method, so that the integral cannot be evaluated by SymPy,\n", - "2. provides a custom NumPy printer method (see **[TR-001](001.ipynb)**) that lambdifies this expression node to [`quadpy.quad()`](https://github.com/sigma-py/quadpy),\n", - "4. adds class variables that allow configuring [`quad_vec()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.integrate.quad_vec.html),\n", - "5. dummifies the integration variable in case it is not a valid Python variable name." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "class UnevaluatableIntegral(sp.Integral):\n", - " abs_tolerance = 1e-5\n", - " rel_tolerance = 1e-5\n", - " limit = 50\n", - " dummify = True\n", - "\n", - " def doit(self, **hints):\n", - " args = [arg.doit(**hints) for arg in self.args]\n", - " return self.func(*args)\n", - "\n", - " def _numpycode(self, printer, *args):\n", - " integration_vars, limits = _unpack_integral_limits(self)\n", - " if len(limits) != 1 or len(integration_vars) != 1:\n", - " msg = f\"Cannot handle {len(limits)}-dimensional integrals\"\n", - " raise ValueError(msg)\n", - " x = integration_vars[0]\n", - " a, b = limits[0]\n", - " expr = self.args[0]\n", - " if self.dummify:\n", - " dummy = sp.Dummy()\n", - " expr = expr.xreplace({x: dummy})\n", - " x = dummy\n", - " integrate_func = \"quad_vec\"\n", - " printer.module_imports[\"scipy.integrate\"].add(integrate_func)\n", - " return (\n", - " f\"{integrate_func}(lambda {printer._print(x)}: {printer._print(expr)},\"\n", - " f\" {printer._print(a)}, {printer._print(b)},\"\n", - " f\" epsabs={self.abs_tolerance}, epsrel={self.abs_tolerance},\"\n", - " f\" limit={self.limit})[0]\"\n", - " )" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To test whether this works, test this integral expression on another {func}`~ampform.sympy.unevaluated` expression:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class MyFunction(sp.Expr):\n", - " x: sp.Symbol\n", - " omega1: sp.Symbol\n", - " omega2: sp.Symbol\n", - " phi1: sp.Symbol\n", - " phi2: sp.Symbol\n", - " _latex_repr_ = R\"f\\left({x}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " x, omega1, omega2, phi1, phi2 = self.args\n", - " return sp.sin(omega1 * x + phi1) + sp.sin(omega2 * x + phi2)\n", - "\n", - "\n", - "x, omega1, omega2, phi1, phi2 = sp.symbols(\"x omega1 omega2 phi1 phi2\")\n", - "expr = MyFunction(x, omega1, omega2, phi1, phi2)\n", - "Math(aslatex({expr: expr.doit(deep=False)}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "w, a, b = sp.symbols(\"w a b\")\n", - "fourier_expr = UnevaluatableIntegral(expr * sp.exp(-sp.I * w * x), (x, a, b))\n", - "fourier_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Indeed the expression correctly lambdifies correctly, despite the {meth}`~sympy.core.basic.Basic.doit` call:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "func = sp.lambdify([x, omega1, omega2, phi1, phi2], expr.doit())\n", - "fourier_func = sp.lambdify([w, omega1, omega2, phi1, phi2, a, b], fourier_expr.doit())" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "src = inspect.getsource(fourier_func)\n", - "src = f\"\"\"```python\n", - "{black.format_str(src, mode=black.FileMode()).strip()}\n", - "```\"\"\"\n", - "Markdown(src)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "domain = np.linspace(-7, +7, num=500)\n", - "parameters = dict(\n", - " omega1=1.2,\n", - " omega2=2.3,\n", - " phi1=-1.2,\n", - " phi2=+0.4,\n", - ")\n", - "func_output = func(domain, **parameters)\n", - "fourier_output = fourier_func(domain, **parameters, a=-10, b=+10)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "fig, ax = plt.subplots()\n", - "ax.set_xlabel(\"$x,w$\")\n", - "ax.plot(domain, func_output, label=\"$f(x)$\")\n", - "ax.plot(domain, fourier_output.real, label=R\"$\\mathrm{Re}\\,F(w)$\")\n", - "ax.plot(domain, fourier_output.imag, label=R\"$\\mathrm{Im}\\,F(w)$\")\n", - "ax.legend()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "See how this integral expression class is applied to the phase space factor in **[TR-003](003.ipynb)**.\n", - ":::" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/017.ipynb b/docs/report/017.ipynb deleted file mode 100644 index 12777c0e..00000000 --- a/docs/report/017.ipynb +++ /dev/null @@ -1,769 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Symbolic phase space boundary for a three-body decay\n", - "TR-017\n", - "^^^\n", - "This reports shows how define the physical phase space region on a Dalitz plot using a Kibble function.\n", - "+++\n", - "✅ [compwa.github.io#139](https://github.com/ComPWA/compwa.github.io/issues/139)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Phase space for a three-body decay" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.0 sympy==1.10.1 tensorwaves==0.4.5" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import os\n", - "import warnings\n", - "from typing import TYPE_CHECKING\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import sympy as sp\n", - "from ampform.sympy import (\n", - " UnevaluatedExpression,\n", - " create_expression,\n", - " implement_doit_method,\n", - " make_commutative,\n", - ")\n", - "from IPython.display import Math\n", - "from ipywidgets import FloatSlider, VBox, interactive_output\n", - "from matplotlib.patches import Patch\n", - "from tensorwaves.function.sympy import create_parametrized_function\n", - "\n", - "if TYPE_CHECKING:\n", - " from matplotlib.axis import Axis\n", - " from matplotlib.contour import QuadContourSet\n", - " from matplotlib.lines import Line2D\n", - "\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Kinematics for a three-body decay $0 \\to 123$ can be fully described by two **Mandelstam variables** $\\sigma_1, \\sigma_2$, because the third variable $\\sigma_3$ can be expressed in terms $\\sigma_1, \\sigma_2$, the mass $m_0$ of the initial state, and the masses $m_1, m_2, m_3$ of the final state. As can be seen, the roles of $\\sigma_1, \\sigma_2, \\sigma_3$ are interchangeable.\n", - "\n", - "```{margin}\n", - "See Eq. (1.2) in {cite}`Byckling:1971vca`\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sigma_{3} = m_{0}^{2} + m_{1}^{2} + m_{2}^{2} + m_{3}^{2} - \\sigma_{1} - \\sigma_{2}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def compute_third_mandelstam(sigma1, sigma2, m0, m1, m2, m3) -> sp.Expr:\n", - " return m0**2 + m1**2 + m2**2 + m3**2 - sigma1 - sigma2\n", - "\n", - "\n", - "m0, m1, m2, m3 = sp.symbols(\"m:4\")\n", - "s1, s2, s3 = sp.symbols(\"sigma1:4\")\n", - "computed_s3 = compute_third_mandelstam(s1, s2, m0, m1, m2, m3)\n", - "Math(Rf\"{sp.latex(s3)} = {sp.latex(computed_s3)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The phase space is defined by the closed area that satisfies the condition $\\phi(\\sigma_1,\\sigma_2) \\leq 0$, where $\\phi$ is a **Kibble function**:\n", - "\n", - "\n", - "```{margin}\n", - "See §V.2 in {cite}`Byckling:1971vca`\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input", - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\phi\\left(\\sigma_{1}, \\sigma_{2}\\right) = \\lambda\\left(\\lambda\\left(\\sigma_{2}, m_{2}^{2}, m_{0}^{2}\\right), \\lambda\\left(\\sigma_{3}, m_{3}^{2}, m_{0}^{2}\\right), \\lambda\\left(\\sigma_{1}, m_{1}^{2}, m_{0}^{2}\\right)\\right)$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "@make_commutative\n", - "@implement_doit_method\n", - "class Kibble(UnevaluatedExpression):\n", - " def __new__(cls, sigma1, sigma2, sigma3, m0, m1, m2, m3, **hints) -> Kibble:\n", - " return create_expression(cls, sigma1, sigma2, sigma3, m0, m1, m2, m3, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " sigma1, sigma2, sigma3, m0, m1, m2, m3 = self.args\n", - " return Källén(\n", - " Källén(sigma2, m2**2, m0**2),\n", - " Källén(sigma3, m3**2, m0**2),\n", - " Källén(sigma1, m1**2, m0**2),\n", - " )\n", - "\n", - " def _latex(self, printer, *args):\n", - " sigma1, sigma2, *_ = map(printer._print, self.args)\n", - " return Rf\"\\phi\\left({sigma1}, {sigma2}\\right)\"\n", - "\n", - "\n", - "@make_commutative\n", - "@implement_doit_method\n", - "class Källén(UnevaluatedExpression):\n", - " def __new__(cls, x, y, z, **hints) -> Källén:\n", - " return create_expression(cls, x, y, z, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " x, y, z = self.args\n", - " return x**2 + y**2 + z**2 - 2 * x * y - 2 * y * z - 2 * z * x\n", - "\n", - " def _latex(self, printer, *args):\n", - " x, y, z = map(printer._print, self.args)\n", - " return Rf\"\\lambda\\left({x}, {y}, {z}\\right)\"\n", - "\n", - "\n", - "kibble = Kibble(s1, s2, s3, m0, m1, m2, m3)\n", - "Math(Rf\"{sp.latex(kibble)} = {sp.latex(kibble.doit(deep=False))}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "and $\\lambda$ is the **Källén function**:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\lambda\\left(x, y, z\\right) = x^{2} - 2 x y - 2 x z + y^{2} - 2 y z + z^{2}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "x, y, z = sp.symbols(\"x:z\")\n", - "expr = Källén(x, y, z)\n", - "Math(f\"{sp.latex(expr)} = {sp.latex(expr.doit())}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Any distribution over the phase space can now be defined using a two-dimensional grid over a Mandelstam pair $\\sigma_1,\\sigma_2$ of choice, with the condition $\\phi(\\sigma_1,\\sigma_2)<0$ selecting the values that are physically allowed." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{cases} 1 & \\text{for}\\: \\phi\\left(\\sigma_{1}, \\sigma_{2}\\right) \\leq 0 \\\\\\text{NaN} & \\text{otherwise} \\end{cases}$" - ], - "text/plain": [ - "Piecewise((1, Kibble(sigma1, sigma2, m0**2 + m1**2 + m2**2 + m3**2 - sigma1 - sigma2, m0, m1, m2, m3) <= 0), (nan, True))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def is_within_phasespace(\n", - " sigma1, sigma2, m0, m1, m2, m3, outside_value=sp.nan\n", - ") -> sp.Piecewise:\n", - " sigma3 = compute_third_mandelstam(sigma1, sigma2, m0, m1, m2, m3)\n", - " kibble = Kibble(sigma1, sigma2, sigma3, m0, m1, m2, m3)\n", - " return sp.Piecewise(\n", - " (1, sp.LessThan(kibble, 0)),\n", - " (outside_value, True),\n", - " )\n", - "\n", - "\n", - "is_within_phasespace(s1, s2, m0, m1, m2, m3)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "phsp_expr = is_within_phasespace(s1, s2, m0, m1, m2, m3, outside_value=0)\n", - "phsp_func = create_parametrized_function(\n", - " phsp_expr.doit(),\n", - " parameters={m0: 2.2, m1: 0.2, m2: 0.4, m3: 0.4},\n", - " backend=\"numpy\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input", - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "sliders = {\n", - " \"m0\": FloatSlider(description=\"m0\", max=3, value=2.1, step=0.01),\n", - " \"m1\": FloatSlider(description=\"m1\", max=2, value=0.2, step=0.01),\n", - " \"m2\": FloatSlider(description=\"m2\", max=2, value=0.4, step=0.01),\n", - " \"m3\": FloatSlider(description=\"m3\", max=2, value=0.4, step=0.01),\n", - "}\n", - "\n", - "resolution = 300\n", - "X, Y = np.meshgrid(\n", - " np.linspace(0, 4, num=resolution),\n", - " np.linspace(0, 4, num=resolution),\n", - ")\n", - "data = {\"sigma1\": X, \"sigma2\": Y}\n", - "\n", - "sidebar_ratio = 0.15\n", - "fig, ((ax1, _), (ax, ax2)) = plt.subplots(\n", - " figsize=(7, 7),\n", - " ncols=2,\n", - " nrows=2,\n", - " gridspec_kw={\n", - " \"height_ratios\": [sidebar_ratio, 1],\n", - " \"width_ratios\": [1, sidebar_ratio],\n", - " },\n", - ")\n", - "_.remove()\n", - "ax.set_xlim(0, 4)\n", - "ax.set_ylim(0, 4)\n", - "ax.set_xlabel(R\"$\\sigma_1$\")\n", - "ax.set_ylabel(R\"$\\sigma_2$\")\n", - "ax.set_xticks(range(5))\n", - "ax.set_yticks(range(5))\n", - "ax1.set_xlim(0, 4)\n", - "ax2.set_ylim(0, 4)\n", - "for a in [ax1, ax2]:\n", - " a.set_xticks([])\n", - " a.set_yticks([])\n", - " a.axis(\"off\")\n", - "fig.tight_layout()\n", - "fig.subplots_adjust(wspace=0, hspace=0)\n", - "\n", - "fig.canvas.toolbar_visible = False\n", - "fig.canvas.header_visible = False\n", - "fig.canvas.footer_visible = False\n", - "\n", - "MESH: QuadContourSet | None = None\n", - "PROJECTIONS: tuple[Line2D, Line2D] = None\n", - "BOUNDARIES: list[Line2D] | None = None\n", - "\n", - "\n", - "def plot(**parameters):\n", - " draw_boundaries(\n", - " parameters[\"m0\"],\n", - " parameters[\"m1\"],\n", - " parameters[\"m2\"],\n", - " parameters[\"m3\"],\n", - " )\n", - " global MESH, PROJECTIONS\n", - " if MESH is not None:\n", - " for coll in MESH.collections:\n", - " ax.collections.remove(coll)\n", - " phsp_func.update_parameters(parameters)\n", - " Z = phsp_func(data)\n", - " MESH = ax.contour(X, Y, Z, colors=\"black\")\n", - " contour = MESH.collections[0]\n", - " contour.set_facecolor(\"lightgray\")\n", - " x = X[0]\n", - " y = Y[:, 0]\n", - " Zx = np.nansum(Z, axis=0)\n", - " Zy = np.nansum(Z, axis=1)\n", - " if PROJECTIONS is None:\n", - " PROJECTIONS = (\n", - " ax1.plot(x, Zx, c=\"black\", lw=2)[0],\n", - " ax2.plot(Zy, y, c=\"black\", lw=2)[0],\n", - " )\n", - " else:\n", - " PROJECTIONS[0].set_data(x, Zx)\n", - " PROJECTIONS[1].set_data(Zy, y)\n", - " ax1.relim()\n", - " ax2.relim()\n", - " ax1.autoscale_view(scalex=False)\n", - " ax2.autoscale_view(scaley=False)\n", - " create_legend(ax)\n", - " fig.canvas.draw()\n", - "\n", - "\n", - "def draw_boundaries(m0, m1, m2, m3) -> None:\n", - " global BOUNDARIES\n", - " s1_min = (m2 + m3) ** 2\n", - " s1_max = (m0 - m1) ** 2\n", - " s2_min = (m1 + m3) ** 2\n", - " s2_max = (m0 - m2) ** 2\n", - " if BOUNDARIES is None:\n", - " BOUNDARIES = [\n", - " ax.axvline(s1_min, c=\"red\", ls=\"dotted\", label=\"$(m_2+m_3)^2$\"),\n", - " ax.axhline(s2_min, c=\"blue\", ls=\"dotted\", label=\"$(m_1+m_3)^2$\"),\n", - " ax.axvline(s1_max, c=\"red\", ls=\"dashed\", label=\"$(m_0-m_1)^2$\"),\n", - " ax.axhline(s2_max, c=\"blue\", ls=\"dashed\", label=\"$(m_0-m_2)^2$\"),\n", - " ]\n", - " else:\n", - " BOUNDARIES[0].set_data(get_line_data(s1_min))\n", - " BOUNDARIES[1].set_data(get_line_data(s2_min, horizontal=True))\n", - " BOUNDARIES[2].set_data(get_line_data(s1_max))\n", - " BOUNDARIES[3].set_data(get_line_data(s2_max, horizontal=True))\n", - "\n", - "\n", - "def create_legend(ax: Axis):\n", - " if ax.get_legend() is not None:\n", - " return\n", - " label = Rf\"${sp.latex(kibble)}\\leq0$\"\n", - " ax.legend(\n", - " handles=[\n", - " Patch(label=label, ec=\"black\", fc=\"lightgray\"),\n", - " *BOUNDARIES,\n", - " ],\n", - " loc=\"upper right\",\n", - " facecolor=\"white\",\n", - " framealpha=1,\n", - " )\n", - "\n", - "\n", - "def get_line_data(value, horizontal: bool = False) -> np.ndarray:\n", - " pair = (value, value)\n", - " if horizontal:\n", - " return np.array([(0, 1), pair])\n", - " return np.array([pair, (0, 1)])\n", - "\n", - "\n", - "output = interactive_output(plot, controls=sliders)\n", - "VBox([output, *sliders.values()])" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}\n", - "\n", - "```{image} https://user-images.githubusercontent.com/29308176/165257660-af66c9c1-2da9-475f-b4c4-4a52950acd00.gif\n", - ":alt: Cell output - interactive Dalitz plot\n", - ":align: center\n", - "```\n", - "
" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The phase space boundary can be described analytically in terms of $\\sigma_1$ or $\\sigma_2$, in which case there are two solutions:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sol1, sol2 = sp.solve(kibble.doit().subs(s3, computed_s3), s2)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "keep_output", - "full-width" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{c}\n", - " \\frac{- m_{0}^{2} m_{2}^{2} + m_{0}^{2} m_{3}^{2} + m_{0}^{2} \\sigma_{1} + m_{1}^{2} m_{2}^{2} - m_{1}^{2} m_{3}^{2} + m_{1}^{2} \\sigma_{1} + m_{2}^{2} \\sigma_{1} + m_{3}^{2} \\sigma_{1} - \\sigma_{1}^{2} - \\sqrt{\\left(m_{0}^{2} - 2 m_{0} m_{1} + m_{1}^{2} - \\sigma_{1}\\right) \\left(m_{0}^{2} + 2 m_{0} m_{1} + m_{1}^{2} - \\sigma_{1}\\right) \\left(m_{2}^{2} - 2 m_{2} m_{3} + m_{3}^{2} - \\sigma_{1}\\right) \\left(m_{2}^{2} + 2 m_{2} m_{3} + m_{3}^{2} - \\sigma_{1}\\right)}}{2 \\sigma_{1}} \\\\\n", - " \\frac{- m_{0}^{2} m_{2}^{2} + m_{0}^{2} m_{3}^{2} + m_{0}^{2} \\sigma_{1} + m_{1}^{2} m_{2}^{2} - m_{1}^{2} m_{3}^{2} + m_{1}^{2} \\sigma_{1} + m_{2}^{2} \\sigma_{1} + m_{3}^{2} \\sigma_{1} - \\sigma_{1}^{2} + \\sqrt{\\left(m_{0}^{2} - 2 m_{0} m_{1} + m_{1}^{2} - \\sigma_{1}\\right) \\left(m_{0}^{2} + 2 m_{0} m_{1} + m_{1}^{2} - \\sigma_{1}\\right) \\left(m_{2}^{2} - 2 m_{2} m_{3} + m_{3}^{2} - \\sigma_{1}\\right) \\left(m_{2}^{2} + 2 m_{2} m_{3} + m_{3}^{2} - \\sigma_{1}\\right)}}{2 \\sigma_{1}} \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "latex = R\"\\begin{array}{c}\" + \"\\n\"\n", - "for sol in [sol1, sol2]:\n", - " latex += Rf\" {sp.latex(sol)} \\\\\" + \"\\n\"\n", - "latex += R\"\\end{array}\"\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The boundary cannot be parametrized analytically in polar coordinates, but there is a numeric solution. The idea is to solve the condition $\\phi(\\sigma_1,\\sigma_2)=0$ after the following substitutions:\n", - "\n", - "```{margin}\n", - "See {cite}`Byckling:1971vca`, pp. 109–112\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "T0, T1, T2, T3 = sp.symbols(\"T0:4\")\n", - "r, theta = sp.symbols(\"r theta\", nonnegative=True)\n", - "substitutions = {\n", - " s1: (m2 + m3) ** 2 + T1,\n", - " s2: (m1 + m3) ** 2 + T2,\n", - " s3: (m1 + m2) ** 2 + T3,\n", - " T1: T0 / 3 - r * sp.cos(theta),\n", - " T2: T0 / 3 - r * sp.cos(theta + 2 * sp.pi / 3),\n", - " T3: T0 / 3 - r * sp.cos(theta + 4 * sp.pi / 3),\n", - " T0: (\n", - " m0**2 + m1**2 + m2**2 + m3**2 - (m2 + m3) ** 2 - (m1 + m3) ** 2 - (m1 + m2) ** 2\n", - " ),\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " \\sigma_{1} &=& T_{1} + \\left(m_{2} + m_{3}\\right)^{2} \\\\\n", - " \\sigma_{2} &=& T_{2} + \\left(m_{1} + m_{3}\\right)^{2} \\\\\n", - " \\sigma_{3} &=& T_{3} + \\left(m_{1} + m_{2}\\right)^{2} \\\\\n", - " T_{1} &=& \\frac{T_{0}}{3} - r \\cos{\\left(\\theta \\right)} \\\\\n", - " T_{2} &=& \\frac{T_{0}}{3} + r \\sin{\\left(\\theta + \\frac{\\pi}{6} \\right)} \\\\\n", - " T_{3} &=& \\frac{T_{0}}{3} + r \\cos{\\left(\\theta + \\frac{\\pi}{3} \\right)} \\\\\n", - " T_{0} &=& m_{0}^{2} + m_{1}^{2} + m_{2}^{2} + m_{3}^{2} - \\left(m_{1} + m_{2}\\right)^{2} - \\left(m_{1} + m_{3}\\right)^{2} - \\left(m_{2} + m_{3}\\right)^{2} \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "latex = R\"\\begin{array}{rcl}\" + \"\\n\"\n", - "for item in substitutions.items():\n", - " lhs, rhs = map(sp.latex, item)\n", - " latex += Rf\" {lhs} &=& {rhs} \\\\\" + \"\\n\"\n", - "latex += R\"\\end{array}\"\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For every value of $\\theta \\in [0, 2\\pi)$, the value of $r$ can now be found by solving the condition $\\phi(r, \\theta)=0$. Note that $\\phi(r, \\theta)$ is a **cubic polynomial of $r$**. For instance, if we take $m_0=5, m_1=2, m_{2,3}=1$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "phi_r = (\n", - " kibble.doit()\n", - " .subs(substitutions) # substitute sigmas\n", - " .subs(substitutions) # substitute T123\n", - " .subs(substitutions) # substitute T0\n", - " .subs({m0: 5, m1: 2, m2: 1, m3: 1})\n", - " .simplify()\n", - " .collect(r)\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\phi(r, \\theta) & = & r^{3} \\cdot \\left(56 \\sqrt{3} \\sin{\\left(\\theta \\right)} + 400 \\cos^{3}{\\left(\\theta \\right)} - 356 \\cos{\\left(\\theta \\right)} + 112 \\cos{\\left(\\theta + \\frac{\\pi}{3} \\right)}\\right) \\nonumber\\\\\n", - "& & + r^{2} \\cdot \\left(2000 \\cos^{2}{\\left(\\theta \\right)} + 2100\\right) \\nonumber\\\\\n", - "& & - 4800 r \\cos{\\left(\\theta \\right)} \\nonumber\\\\\n", - "& & + -25200 \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "lhs = sp.Symbol(R\"\\phi(r, \\theta)\")\n", - "latex = sp.multiline_latex(lhs, phi_r, environment=\"eqnarray\")\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The lowest value of $r$ that satisfies $\\phi(r,\\theta)=0$ defines the phase space boundary." - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/018.ipynb b/docs/report/018.ipynb deleted file mode 100644 index d56e1065..00000000 --- a/docs/report/018.ipynb +++ /dev/null @@ -1,609 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "tensorwaves" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Intensity distribution generator with importance sampling\n", - "TR-018\n", - "^^^\n", - "This reports sets out how data generation with TensorWaves works and what would be the best approach to tackle [tensorwaves#402](https://github.com/ComPWA/tensorwaves/issues/402).\n", - "+++\n", - "WIP\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Importance sampling" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.0 qrules[viz]==0.9.7 scipy==1.8.0 sympy==1.10.1 tensorwaves[jax,pwa]==0.4.8" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Model definition" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import logging\n", - "import os\n", - "import warnings\n", - "\n", - "import jax.numpy as jnp\n", - "import numpy as np\n", - "\n", - "logging.getLogger(\"absl\").setLevel(logging.ERROR) # no JAX warnings\n", - "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\" # no TF warnings\n", - "warnings.filterwarnings(\"ignore\") # sqrt negative argument" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We generate data for the reaction $J/\\psi \\to \\gamma \\pi^0\\pi^0$. We limit ourselves to two resonances, so that the amplitude model contains one narrow structure. This makes it hard to numerically compute the integral over the intensity distribution." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import qrules\n", - "\n", - "reaction = qrules.generate_transitions(\n", - " initial_state=(\"J/psi(1S)\", [-1, +1]),\n", - " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", - " allowed_intermediate_particles=[\"f(0)(980)\", \"omega(782)\"],\n", - " allowed_interaction_types=[\"strong\", \"EM\"],\n", - " formalism=\"canonical-helicity\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "import graphviz\n", - "\n", - "src = qrules.io.asdot(reaction, collapse_graphs=True)\n", - "_ = graphviz.Source(src).render(\"018-graph\", format=\"svg\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/183455676-25af2dcf-ecac-43af-aa11-b3f71fe095f8.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import ampform\n", - "from ampform.dynamics.builder import (\n", - " create_non_dynamic_with_ff,\n", - " create_relativistic_breit_wigner_with_ff,\n", - ")\n", - "\n", - "builder = ampform.get_builder(reaction)\n", - "builder.align_spin = False\n", - "builder.adapter.permutate_registered_topologies()\n", - "builder.scalar_initial_state_mass = True\n", - "builder.stable_final_state_ids = [0, 1, 2]\n", - "builder.set_dynamics(\"J/psi(1S)\", create_non_dynamic_with_ff)\n", - "for name in reaction.get_intermediate_particles().names:\n", - " builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", - "model = builder.formulate()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Phase space distribution" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "An evenly distributed phase space sample can be generated with a {class}`~tensorwaves.data.phasespace.TFPhaseSpaceGenerator`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import (\n", - " SympyDataTransformer,\n", - " TFPhaseSpaceGenerator,\n", - " TFUniformRealNumberGenerator,\n", - ")\n", - "\n", - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "transformer = SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", - "phsp = phsp_generator.generate(1_000_000, rng)\n", - "phsp = transformer(phsp)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "\n", - "\n", - "def convert_zero_to_nan(array):\n", - " array = np.array(array).astype(\"float\")\n", - " array[array == 0] = np.nan\n", - " return jnp.array(array)\n", - "\n", - "\n", - "Z, x_edges, y_edges = jnp.histogram2d(\n", - " phsp[\"m_01\"].real ** 2,\n", - " phsp[\"m_12\"].real ** 2,\n", - " bins=100,\n", - ")\n", - "X, Y = jnp.meshgrid(x_edges, y_edges)\n", - "Z = convert_zero_to_nan(Z)\n", - "\n", - "bin_width_x = X[0, 1] - X[0, 0]\n", - "bin_width_y = Y[1, 0] - Y[0, 0]\n", - "bar_title = (\n", - " Rf\"events per ${1e3 * bin_width_x:.0f} \\times {1e3 * bin_width_y:.0f}$ MeV$^2/c^4$\"\n", - ")\n", - "xlabel = R\"$M^2\\left(\\gamma\\pi^0\\right)$\"\n", - "ylabel = R\"$M^2\\left(\\pi^0\\pi^0\\right)$\"\n", - "\n", - "plt.ioff()\n", - "fig, ax = plt.subplots(dpi=200, figsize=(4.5, 4))\n", - "ax.set_title(\"TFPhaseSpaceGenerator sample\")\n", - "ax.set_xlabel(xlabel)\n", - "ax.set_ylabel(ylabel)\n", - "mesh = ax.pcolormesh(X, Y, Z)\n", - "c_bar = plt.colorbar(mesh, ax=ax)\n", - "c_bar.ax.set_ylabel(bar_title)\n", - "fig.savefig(\"018-TFPhaseSpaceGenerator.png\")\n", - "plt.ion()\n", - "plt.close(fig)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/183455681-31745b0c-d1ec-48c1-a6d6-388f4725df42.png)\n", - "\n", - "This {class}`~tensorwaves.data.phasespace.TFPhaseSpaceGenerator` actually uses a **hit-and-miss** strategy on a distribution and its weights generated by a {class}`~tensorwaves.data.phasespace.TFWeightedPhaseSpaceGenerator`. That generator interfaces to the [`phasespace`](https://phasespace.readthedocs.io) package. We have a short look at the distribution and its weights generated by a {class}`~tensorwaves.data.phasespace.TFWeightedPhaseSpaceGenerator`. The 'unweighted' distribution is uneven, because four-momenta events are generated using a certain decay algorithm. The weights cause these events to be normalized, so that we again have the same, evenly distributed distribution from above when we combine them.\n", - "\n", - "::::{margin}\n", - ":::{seealso}\n", - "[tensorwaves#16](https://github.com/ComPWA/tensorwaves/issues/16) on a Python interface for [`EvtGen`](https://gitlab.cern.ch/evtgen/evtgen).\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import TFWeightedPhaseSpaceGenerator\n", - "\n", - "weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "unweighted_phsp = weighted_phsp_generator.generate(1_000_000, rng)\n", - "phsp_weights = unweighted_phsp[\"weights\"]\n", - "unweighted_phsp = transformer(unweighted_phsp)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "from typing import TYPE_CHECKING\n", - "\n", - "from scipy.interpolate import griddata\n", - "\n", - "if TYPE_CHECKING:\n", - " from tensorwaves.interface import DataSample\n", - "\n", - "\n", - "def plot_distribution_and_weights(phsp: DataSample, weights: np.ndarray) -> None:\n", - " n_bins = 100\n", - " x = phsp[\"m_01\"].real ** 2\n", - " y = phsp[\"m_12\"].real ** 2\n", - " X, Y = jnp.meshgrid(\n", - " jnp.linspace(x.min(), x.max(), num=n_bins),\n", - " jnp.linspace(y.min(), y.max(), num=n_bins),\n", - " )\n", - "\n", - " Z_weights = griddata(np.transpose([x, y]), weights, (X, Y))\n", - " Z_unweighted, x_edges, y_edges = jnp.histogram2d(x, y, bins=n_bins)\n", - " Z_weighted, x_edges, y_edges = jnp.histogram2d(x, y, bins=n_bins, weights=weights)\n", - " # https://numpy.org/doc/stable/reference/generated/numpy.histogram2d.html\n", - " Z_unweighted = Z_unweighted.T\n", - " Z_weighted = Z_weighted.T\n", - "\n", - " X_edges, Y_edges = jnp.meshgrid(x_edges, y_edges)\n", - " Z_unweighted = convert_zero_to_nan(Z_unweighted)\n", - " Z_weighted = convert_zero_to_nan(Z_weighted)\n", - "\n", - " _, axes = plt.subplots(\n", - " dpi=200,\n", - " figsize=(16, 5),\n", - " ncols=3,\n", - " tight_layout=True,\n", - " )\n", - " for ax in axes:\n", - " ax.set_xlabel(xlabel)\n", - " ax.set_ylabel(ylabel)\n", - " axes[0].set_title(\"Unweighted distribution\")\n", - " axes[1].set_title(\"Weights\")\n", - " axes[2].set_title(\"Weighted phase space distribution\")\n", - "\n", - " mesh = axes[0].pcolormesh(X_edges, Y_edges, Z_unweighted)\n", - " c_bar = plt.colorbar(mesh, ax=axes[0])\n", - " c_bar.ax.set_ylabel(bar_title)\n", - "\n", - " mesh = axes[1].pcolormesh(X, Y, Z_weights)\n", - " c_bar = plt.colorbar(mesh, ax=axes[1])\n", - " c_bar.ax.set_ylabel(\"phase space weight\")\n", - "\n", - " mesh = axes[2].pcolormesh(X_edges, Y_edges, Z_weighted)\n", - " c_bar = plt.colorbar(mesh, ax=axes[2])\n", - " c_bar.ax.set_ylabel(bar_title)\n", - "\n", - "\n", - "plot_distribution_and_weights(unweighted_phsp, phsp_weights)\n", - "plt.gcf().suptitle(\"TFWeightedPhaseSpaceGenerator sample\")\n", - "plt.savefig(\"018-TFWeightedPhaseSpaceGenerator.png\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{container} full-width\n", - "![](https://user-images.githubusercontent.com/29308176/183455686-5fb52b80-bff3-4508-b2dd-69d9b2888678.png)\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Intensity distribution" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now use a {class}`~tensorwaves.data.IntensityDistributionGenerator` to generate a hit-and-miss data sample based on the amplitude model that we formulated for this $J/\\psi \\to \\gamma\\pi^0\\pi^0$ reaction." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.function.sympy import create_parametrized_function\n", - "\n", - "intensity_expr = model.expression.doit()\n", - "intensity_func = create_parametrized_function(\n", - " expression=intensity_expr,\n", - " parameters=model.parameter_defaults,\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import IntensityDistributionGenerator\n", - "\n", - "data_generator = IntensityDistributionGenerator(\n", - " domain_generator=weighted_phsp_generator,\n", - " function=intensity_func,\n", - " domain_transformer=transformer,\n", - ")\n", - "data = data_generator.generate(100_000, rng)\n", - "data = transformer(data)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that it takes a long time to generate a distribution for amplitude model. This is because most phase space points are outside the region where the intensity is highest and therefore result in a 'miss'." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "Z, x_edges, y_edges = jnp.histogram2d(\n", - " data[\"m_01\"].real ** 2,\n", - " data[\"m_12\"].real ** 2,\n", - " bins=100,\n", - ")\n", - "X, Y = jnp.meshgrid(x_edges, y_edges)\n", - "Z = Z.T # https://numpy.org/doc/stable/reference/generated/numpy.histogram2d.html\n", - "Z = convert_zero_to_nan(Z)\n", - "\n", - "plt.ioff()\n", - "fig, ax = plt.subplots(dpi=200, figsize=(4.5, 4))\n", - "ax.set_xlabel(xlabel)\n", - "ax.set_ylabel(ylabel)\n", - "mesh = ax.pcolormesh(X, Y, Z)\n", - "c_bar = plt.colorbar(mesh, ax=ax)\n", - "c_bar.ax.set_ylabel(\"intensity\")\n", - "fig.savefig(\"018-intensity-distribution.png\")\n", - "plt.ion()\n", - "plt.close(fig)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/183458258-cae90095-e0d7-4251-a2e6-d7c71750094c.png)\n", - "\n", - "The $\\omega$ resonance appears as a narrow structure on the Dalitz plot. This is problematic when computing the integral over this distribution, which is important when performing an {class}`~tensorwaves.estimator.UnbinnedNLL` fit. The integral that appears in the log-likelihood has to be computed in each fit iteration and this can be done most efficiently when there are more points on which to evaluate the amplitude model in the phase space regions where the intensity is high.\n", - "\n", - "The solution is to evaluate the intensity over an **importance-sampled phase space sample**. This is a phase space sample with more events in the regions where the intensity is high. Each point $\\tau$ carries a weight that is set to $1/I(\\tau)$. In fact, all this is, is the intensity-based sample from the previous step, with the weights computed posteriorly by simply evaluating the a amplitude model over the sample (and taking the inverse)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from copy import deepcopy\n", - "\n", - "importance_phsp = deepcopy(data)\n", - "importance_weights = 1 / intensity_func(importance_phsp)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Of course, we could define a special class for this.\n", - "\n", - "As expected, the inverse-intensity weights flatten the distribution again to a flat phase space sample:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "plot_distribution_and_weights(importance_phsp, importance_weights)\n", - "plt.gcf().suptitle(\"Importance-sampled phase space distribution\")\n", - "plt.savefig(\"018-importance-sampling.png\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{container} full-width\n", - "![](https://user-images.githubusercontent.com/29308176/183457906-ea95716f-1a82-4c0c-8992-62aa19131c33.png)\n", - "```\n", - "\n", - "Now, aren't we duplicating things here? Not really. First, in an actual analysis, there would be no intensity-based data sample. Second, the importance-sampled phase space sample is generated _with a specific parameter values_. During a fit, the parameters change and the integral over the (importance-sampled) phase space changes. So after updating parameters during a fit iteration, we have to multiply the new intensities with the importance weights (the inverse of the _original_ intensity distribution) in order to get the new distribution. This needs to be done in particular when computing the negative log likelihood ({class}`~tensorwaves.estimator.UnbinnedNLL`).[^1]\n", - "\n", - "[^1]: As of [TensorWaves v0.4.*](https://tensorwaves.rtfd.io/en/0.4.x), weights are not part of a {obj}`~tensorwaves.interface.DataSample` and are therefore not passed to the {class}`~tensorwaves.estimator.UnbinnedNLL`, where would be needed to correctly compute the integral. This requires an interface change.\n", - "\n", - "In the following, extreme example, we move the mass of the $f_0(980)$ resonance far from its original position. As can be seen in the distribution below, the narrow structure has indeed moved, but the structure is still visible as a blur in the original position, because there are many more phase space points in that region." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensity_func.update_parameters({\"m_{f_{0}(980)}\": 2.0})\n", - "new_intensities = intensity_func(importance_phsp)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "Z, x_edges, y_edges = jnp.histogram2d(\n", - " importance_phsp[\"m_01\"].real ** 2,\n", - " importance_phsp[\"m_12\"].real ** 2,\n", - " bins=100,\n", - " weights=new_intensities * importance_weights,\n", - ")\n", - "X, Y = jnp.meshgrid(x_edges, y_edges)\n", - "Z = Z.T # https://numpy.org/doc/stable/reference/generated/numpy.histogram2d.html\n", - "Z = convert_zero_to_nan(Z)\n", - "\n", - "plt.ioff()\n", - "fig, ax = plt.subplots(dpi=200, figsize=(4.5, 4))\n", - "ax.set_xlabel(xlabel)\n", - "ax.set_ylabel(ylabel)\n", - "mesh = ax.pcolormesh(X, Y, Z)\n", - "c_bar = plt.colorbar(mesh, ax=ax)\n", - "c_bar.ax.set_ylabel(R\"new intensity $\\times$ importance weight\")\n", - "fig.savefig(\"018-importance-sampling-after-modification.png\")\n", - "plt.ion()\n", - "plt.close(fig)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/183459123-ab1f3bb5-d51d-4122-97f1-0b51065b94b8.png)" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/019.ipynb b/docs/report/019.ipynb deleted file mode 100644 index 11b8bffd..00000000 --- a/docs/report/019.ipynb +++ /dev/null @@ -1,248 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "DX", - "tips" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Integrating Jupyter notebook with Julia notebooks in MyST-NB\n", - "TR-019\n", - "^^^\n", - "This report shows how to define a Julia kernel for Jupyter notebooks, so that it can be executed and converted to static pages with [MyST-NB](https://myst-nb.readthedocs.io).\n", - "+++\n", - "✅ [compwa.github.io#174](https://github.com/ComPWA/compwa.github.io/issues/174)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Jupyter notebook with Julia kernel" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This notebook shows that the instructions provided in the {ref}`Julia installation` and {ref}`IJulia instructions` work correctly. The cell outputs below are generated automatically with [MyST-NB](https://myst-nb.readthedocs.io) from the Julia code input.\n", - "\n", - "Simple example:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Hello world!\n" - ] - } - ], - "source": [ - "println(\"Hello world!\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here's an [example](https://rosettacode.org/wiki/Mandelbrot_set#Julia) that prints a Mandelbrot set!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " \n", - " \n", - " \n", - " ** \n", - " ****** \n", - " ******** \n", - " ****** \n", - " ******** ** * \n", - " *** ***************** \n", - " ************************ *** \n", - " **************************** \n", - " ****************************** \n", - " ****************************** \n", - " ************************************ \n", - " * ********************************** \n", - " ** ***** * ********************************** \n", - " *********** ************************************ \n", - " ************** ************************************ \n", - " *************************************************** \n", - " ***************************************************** \n", - " *********************************************************************** \n", - " ***************************************************** \n", - " *************************************************** \n", - " ************** ************************************ \n", - " *********** ************************************ \n", - " ** ***** * ********************************** \n", - " * ********************************** \n", - " ************************************ \n", - " ****************************** \n", - " ****************************** \n", - " **************************** \n", - " ************************ *** \n", - " *** ***************** \n", - " ******** ** * \n", - " ****** \n", - " ******** \n", - " ****** \n", - " ** \n", - " \n", - " \n", - " \n" - ] - } - ], - "source": [ - "function mandelbrot(a)\n", - " z = 0\n", - " for i=1:50\n", - " z = z^2 + a\n", - " end\n", - " return z\n", - "end\n", - "\n", - "for y=1.0:-0.05:-1.0\n", - " for x=-2.0:0.0315:0.5\n", - " abs(mandelbrot(complex(x, y))) < 2 ? print(\"*\") : print(\" \")\n", - " end\n", - " println()\n", - "end" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It's also possible to work with a local [environment](https://pkgdocs.julialang.org/v1/environments) from the notebook. In this case, we activate the environment defined by the file [`019/Project.toml`](./019/Project.toml) and instantiate it so that the exact versions of the dependencies as defined in [`019/Manifest.toml`](./019/Manifest.toml) are installed." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "using Pkg\n", - "Pkg.activate(joinpath(@__DIR__, \"019\"))\n", - "Pkg.instantiate()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "using Images\n", - " \n", - "@inline function hsv2rgb(h, s, v)\n", - " c = v * s\n", - " x = c * (1 - abs(((h/60) % 2) - 1))\n", - " m = v - c\n", - " r,g,b = if h < 60 (c, x, 0)\n", - " elseif h < 120 (x, c, 0)\n", - " elseif h < 180 (0, c, x)\n", - " elseif h < 240 (0, x, c)\n", - " elseif h < 300 (x, 0, c)\n", - " else (c, 0, x) end\n", - " (r + m), (b + m), (g + m)\n", - "end\n", - " \n", - "function mandelbrot()\n", - " w = 1600\n", - " h = 1200\n", - " zoom = 0.5\n", - " moveX = -0.5\n", - " moveY = 0\n", - " maxIter = 30\n", - " img = Array{RGB{Float64},2}(undef,h,w)\n", - " for x in 1:w\n", - " for y in 1:h\n", - " i = maxIter\n", - " z = c = Complex( (2*x - w) / (w * zoom) + moveX,\n", - " (2*y - h) / (h * zoom) + moveY )\n", - " while abs(z) < 2 && (i -= 1) > 0\n", - " z = z^2 + c\n", - " end\n", - " r,g,b = hsv2rgb(i / maxIter * 360, 1, i / maxIter)\n", - " img[y,x] = RGB{Float64}(r, g, b)\n", - " end\n", - " end\n", - " return img\n", - "end\n", - " \n", - "mandelbrot()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://github.com/ComPWA/compwa.github.io/assets/29308176/f8e604b9-b37c-4b5a-8114-91627da93d37)" - ] - } - ], - "metadata": { - "colab": { - "toc_visible": true - }, - "kernelspec": { - "display_name": "julia-compwa.github.io 1.10.3", - "language": "julia", - "name": "julia-compwa.github.io-1.10" - }, - "language_info": { - "file_extension": ".jl", - "mimetype": "application/julia", - "name": "julia", - "version": "1.10.3" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/019/Manifest.toml b/docs/report/019/Manifest.toml deleted file mode 100644 index 77929f0d..00000000 --- a/docs/report/019/Manifest.toml +++ /dev/null @@ -1,1101 +0,0 @@ -# This file is machine-generated - editing it directly is not advised - -julia_version = "1.10.3" -manifest_format = "2.0" -project_hash = "4ba49797108fc71b365784ebeeb1834809629bf2" - -[[deps.AbstractFFTs]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "d92ad398961a3ed262d8bf04a1a2b8340f915fef" -uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c" -version = "1.5.0" - - [deps.AbstractFFTs.extensions] - AbstractFFTsChainRulesCoreExt = "ChainRulesCore" - AbstractFFTsTestExt = "Test" - - [deps.AbstractFFTs.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[[deps.Adapt]] -deps = ["LinearAlgebra", "Requires"] -git-tree-sha1 = "6a55b747d1812e699320963ffde36f1ebdda4099" -uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -version = "4.0.4" -weakdeps = ["StaticArrays"] - - [deps.Adapt.extensions] - AdaptStaticArraysExt = "StaticArrays" - -[[deps.ArgTools]] -uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f" -version = "1.1.1" - -[[deps.ArnoldiMethod]] -deps = ["LinearAlgebra", "Random", "StaticArrays"] -git-tree-sha1 = "d57bd3762d308bded22c3b82d033bff85f6195c6" -uuid = "ec485272-7323-5ecc-a04f-4719b315124d" -version = "0.4.0" - -[[deps.ArrayInterface]] -deps = ["Adapt", "LinearAlgebra", "SparseArrays", "SuiteSparse"] -git-tree-sha1 = "133a240faec6e074e07c31ee75619c90544179cf" -uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" -version = "7.10.0" - - [deps.ArrayInterface.extensions] - ArrayInterfaceBandedMatricesExt = "BandedMatrices" - ArrayInterfaceBlockBandedMatricesExt = "BlockBandedMatrices" - ArrayInterfaceCUDAExt = "CUDA" - ArrayInterfaceCUDSSExt = "CUDSS" - ArrayInterfaceChainRulesExt = "ChainRules" - ArrayInterfaceGPUArraysCoreExt = "GPUArraysCore" - ArrayInterfaceReverseDiffExt = "ReverseDiff" - ArrayInterfaceStaticArraysCoreExt = "StaticArraysCore" - ArrayInterfaceTrackerExt = "Tracker" - - [deps.ArrayInterface.weakdeps] - BandedMatrices = "aae01518-5342-5314-be14-df237901396f" - BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0" - CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" - CUDSS = "45b445bb-4962-46a0-9369-b4df9d0f772e" - ChainRules = "082447d4-558c-5d27-93f4-14fc19e9eca2" - GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" - ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" - StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" - Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" - -[[deps.Artifacts]] -uuid = "56f22d72-fd6d-98f1-02f0-08ddc0907c33" - -[[deps.AxisAlgorithms]] -deps = ["LinearAlgebra", "Random", "SparseArrays", "WoodburyMatrices"] -git-tree-sha1 = "01b8ccb13d68535d73d2b0c23e39bd23155fb712" -uuid = "13072b0f-2c55-5437-9ae7-d433b7a33950" -version = "1.1.0" - -[[deps.AxisArrays]] -deps = ["Dates", "IntervalSets", "IterTools", "RangeArrays"] -git-tree-sha1 = "16351be62963a67ac4083f748fdb3cca58bfd52f" -uuid = "39de3d68-74b9-583c-8d2d-e117c070f3a9" -version = "0.4.7" - -[[deps.Base64]] -uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" - -[[deps.BitTwiddlingConvenienceFunctions]] -deps = ["Static"] -git-tree-sha1 = "0c5f81f47bbbcf4aea7b2959135713459170798b" -uuid = "62783981-4cbd-42fc-bca8-16325de8dc4b" -version = "0.1.5" - -[[deps.CEnum]] -git-tree-sha1 = "389ad5c84de1ae7cf0e28e381131c98ea87d54fc" -uuid = "fa961155-64e5-5f13-b03f-caf6b980ea82" -version = "0.5.0" - -[[deps.CPUSummary]] -deps = ["CpuId", "IfElse", "PrecompileTools", "Static"] -git-tree-sha1 = "601f7e7b3d36f18790e2caf83a882d88e9b71ff1" -uuid = "2a0fbf3d-bb9c-48f3-b0a9-814d99fd7ab9" -version = "0.2.4" - -[[deps.CatIndices]] -deps = ["CustomUnitRanges", "OffsetArrays"] -git-tree-sha1 = "a0f80a09780eed9b1d106a1bf62041c2efc995bc" -uuid = "aafaddc9-749c-510e-ac4f-586e18779b91" -version = "0.2.2" - -[[deps.ChainRulesCore]] -deps = ["Compat", "LinearAlgebra"] -git-tree-sha1 = "575cd02e080939a33b6df6c5853d14924c08e35b" -uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -version = "1.23.0" -weakdeps = ["SparseArrays"] - - [deps.ChainRulesCore.extensions] - ChainRulesCoreSparseArraysExt = "SparseArrays" - -[[deps.CloseOpenIntervals]] -deps = ["Static", "StaticArrayInterface"] -git-tree-sha1 = "70232f82ffaab9dc52585e0dd043b5e0c6b714f1" -uuid = "fb6a15b2-703c-40df-9091-08a04967cfa9" -version = "0.1.12" - -[[deps.Clustering]] -deps = ["Distances", "LinearAlgebra", "NearestNeighbors", "Printf", "Random", "SparseArrays", "Statistics", "StatsBase"] -git-tree-sha1 = "9ebb045901e9bbf58767a9f34ff89831ed711aae" -uuid = "aaaa29a8-35af-508c-8bc3-b662a17a0fe5" -version = "0.15.7" - -[[deps.ColorSchemes]] -deps = ["ColorTypes", "ColorVectorSpace", "Colors", "FixedPointNumbers", "PrecompileTools", "Random"] -git-tree-sha1 = "4b270d6465eb21ae89b732182c20dc165f8bf9f2" -uuid = "35d6a980-a343-548e-a6ea-1d62b119f2f4" -version = "3.25.0" - -[[deps.ColorTypes]] -deps = ["FixedPointNumbers", "Random"] -git-tree-sha1 = "b10d0b65641d57b8b4d5e234446582de5047050d" -uuid = "3da002f7-5984-5a60-b8a6-cbb66c0b333f" -version = "0.11.5" - -[[deps.ColorVectorSpace]] -deps = ["ColorTypes", "FixedPointNumbers", "LinearAlgebra", "Requires", "Statistics", "TensorCore"] -git-tree-sha1 = "a1f44953f2382ebb937d60dafbe2deea4bd23249" -uuid = "c3611d14-8923-5661-9e6a-0046d554d3a4" -version = "0.10.0" - - [deps.ColorVectorSpace.extensions] - SpecialFunctionsExt = "SpecialFunctions" - - [deps.ColorVectorSpace.weakdeps] - SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" - -[[deps.Colors]] -deps = ["ColorTypes", "FixedPointNumbers", "Reexport"] -git-tree-sha1 = "fc08e5930ee9a4e03f84bfb5211cb54e7769758a" -uuid = "5ae59095-9a9b-59fe-a467-6f913c188581" -version = "0.12.10" - -[[deps.Compat]] -deps = ["TOML", "UUIDs"] -git-tree-sha1 = "b1c55339b7c6c350ee89f2c1604299660525b248" -uuid = "34da2185-b29b-5c13-b0c7-acf172513d20" -version = "4.15.0" -weakdeps = ["Dates", "LinearAlgebra"] - - [deps.Compat.extensions] - CompatLinearAlgebraExt = "LinearAlgebra" - -[[deps.CompilerSupportLibraries_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "e66e0078-7015-5450-92f7-15fbd957f2ae" -version = "1.1.1+0" - -[[deps.ComputationalResources]] -git-tree-sha1 = "52cb3ec90e8a8bea0e62e275ba577ad0f74821f7" -uuid = "ed09eef8-17a6-5b46-8889-db040fac31e3" -version = "0.3.2" - -[[deps.CoordinateTransformations]] -deps = ["LinearAlgebra", "StaticArrays"] -git-tree-sha1 = "f9d7112bfff8a19a3a4ea4e03a8e6a91fe8456bf" -uuid = "150eb455-5306-5404-9cee-2592286d6298" -version = "0.6.3" - -[[deps.CpuId]] -deps = ["Markdown"] -git-tree-sha1 = "fcbb72b032692610bfbdb15018ac16a36cf2e406" -uuid = "adafc99b-e345-5852-983c-f28acb93d879" -version = "0.3.1" - -[[deps.CustomUnitRanges]] -git-tree-sha1 = "1a3f97f907e6dd8983b744d2642651bb162a3f7a" -uuid = "dc8bdbbb-1ca9-579f-8c36-e416f6a65cce" -version = "1.0.2" - -[[deps.DataAPI]] -git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe" -uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a" -version = "1.16.0" - -[[deps.DataStructures]] -deps = ["Compat", "InteractiveUtils", "OrderedCollections"] -git-tree-sha1 = "1d0a14036acb104d9e89698bd408f63ab58cdc82" -uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8" -version = "0.18.20" - -[[deps.Dates]] -deps = ["Printf"] -uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" - -[[deps.Distances]] -deps = ["LinearAlgebra", "Statistics", "StatsAPI"] -git-tree-sha1 = "66c4c81f259586e8f002eacebc177e1fb06363b0" -uuid = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7" -version = "0.10.11" -weakdeps = ["ChainRulesCore", "SparseArrays"] - - [deps.Distances.extensions] - DistancesChainRulesCoreExt = "ChainRulesCore" - DistancesSparseArraysExt = "SparseArrays" - -[[deps.Distributed]] -deps = ["Random", "Serialization", "Sockets"] -uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b" - -[[deps.DocStringExtensions]] -deps = ["LibGit2"] -git-tree-sha1 = "2fb1e02f2b635d0845df5d7c167fec4dd739b00d" -uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" -version = "0.9.3" - -[[deps.Downloads]] -deps = ["ArgTools", "FileWatching", "LibCURL", "NetworkOptions"] -uuid = "f43a241f-c20a-4ad4-852c-f6b1247861c6" -version = "1.6.0" - -[[deps.FFTViews]] -deps = ["CustomUnitRanges", "FFTW"] -git-tree-sha1 = "cbdf14d1e8c7c8aacbe8b19862e0179fd08321c2" -uuid = "4f61f5a4-77b1-5117-aa51-3ab5ef4ef0cd" -version = "0.3.2" - -[[deps.FFTW]] -deps = ["AbstractFFTs", "FFTW_jll", "LinearAlgebra", "MKL_jll", "Preferences", "Reexport"] -git-tree-sha1 = "4820348781ae578893311153d69049a93d05f39d" -uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" -version = "1.8.0" - -[[deps.FFTW_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "c6033cc3892d0ef5bb9cd29b7f2f0331ea5184ea" -uuid = "f5851436-0d7a-5f13-b9de-f02708fd171a" -version = "3.3.10+0" - -[[deps.FileIO]] -deps = ["Pkg", "Requires", "UUIDs"] -git-tree-sha1 = "82d8afa92ecf4b52d78d869f038ebfb881267322" -uuid = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" -version = "1.16.3" - -[[deps.FileWatching]] -uuid = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee" - -[[deps.FixedPointNumbers]] -deps = ["Statistics"] -git-tree-sha1 = "335bfdceacc84c5cdf16aadc768aa5ddfc5383cc" -uuid = "53c48c17-4a7d-5ca2-90c5-79b7896eea93" -version = "0.8.4" - -[[deps.Ghostscript_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "43ba3d3c82c18d88471cfd2924931658838c9d8f" -uuid = "61579ee1-b43e-5ca0-a5da-69d92c66a64b" -version = "9.55.0+4" - -[[deps.Graphics]] -deps = ["Colors", "LinearAlgebra", "NaNMath"] -git-tree-sha1 = "d61890399bc535850c4bf08e4e0d3a7ad0f21cbd" -uuid = "a2bd30eb-e257-5431-a919-1863eab51364" -version = "1.1.2" - -[[deps.Graphs]] -deps = ["ArnoldiMethod", "Compat", "DataStructures", "Distributed", "Inflate", "LinearAlgebra", "Random", "SharedArrays", "SimpleTraits", "SparseArrays", "Statistics"] -git-tree-sha1 = "4f2b57488ac7ee16124396de4f2bbdd51b2602ad" -uuid = "86223c79-3864-5bf0-83f7-82e725a168b6" -version = "1.11.0" - -[[deps.HistogramThresholding]] -deps = ["ImageBase", "LinearAlgebra", "MappedArrays"] -git-tree-sha1 = "7194dfbb2f8d945abdaf68fa9480a965d6661e69" -uuid = "2c695a8d-9458-5d45-9878-1b8a99cf7853" -version = "0.3.1" - -[[deps.HostCPUFeatures]] -deps = ["BitTwiddlingConvenienceFunctions", "IfElse", "Libdl", "Static"] -git-tree-sha1 = "eb8fed28f4994600e29beef49744639d985a04b2" -uuid = "3e5b6fbb-0976-4d2c-9146-d79de83f2fb0" -version = "0.1.16" - -[[deps.IfElse]] -git-tree-sha1 = "debdd00ffef04665ccbb3e150747a77560e8fad1" -uuid = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173" -version = "0.1.1" - -[[deps.ImageAxes]] -deps = ["AxisArrays", "ImageBase", "ImageCore", "Reexport", "SimpleTraits"] -git-tree-sha1 = "2e4520d67b0cef90865b3ef727594d2a58e0e1f8" -uuid = "2803e5a7-5153-5ecf-9a86-9b4c37f5f5ac" -version = "0.6.11" - -[[deps.ImageBase]] -deps = ["ImageCore", "Reexport"] -git-tree-sha1 = "eb49b82c172811fd2c86759fa0553a2221feb909" -uuid = "c817782e-172a-44cc-b673-b171935fbb9e" -version = "0.1.7" - -[[deps.ImageBinarization]] -deps = ["HistogramThresholding", "ImageCore", "LinearAlgebra", "Polynomials", "Reexport", "Statistics"] -git-tree-sha1 = "f5356e7203c4a9954962e3757c08033f2efe578a" -uuid = "cbc4b850-ae4b-5111-9e64-df94c024a13d" -version = "0.3.0" - -[[deps.ImageContrastAdjustment]] -deps = ["ImageBase", "ImageCore", "ImageTransformations", "Parameters"] -git-tree-sha1 = "eb3d4365a10e3f3ecb3b115e9d12db131d28a386" -uuid = "f332f351-ec65-5f6a-b3d1-319c6670881a" -version = "0.3.12" - -[[deps.ImageCore]] -deps = ["ColorVectorSpace", "Colors", "FixedPointNumbers", "MappedArrays", "MosaicViews", "OffsetArrays", "PaddedViews", "PrecompileTools", "Reexport"] -git-tree-sha1 = "b2a7eaa169c13f5bcae8131a83bc30eff8f71be0" -uuid = "a09fc81d-aa75-5fe9-8630-4744c3626534" -version = "0.10.2" - -[[deps.ImageCorners]] -deps = ["ImageCore", "ImageFiltering", "PrecompileTools", "StaticArrays", "StatsBase"] -git-tree-sha1 = "24c52de051293745a9bad7d73497708954562b79" -uuid = "89d5987c-236e-4e32-acd0-25bd6bd87b70" -version = "0.1.3" - -[[deps.ImageDistances]] -deps = ["Distances", "ImageCore", "ImageMorphology", "LinearAlgebra", "Statistics"] -git-tree-sha1 = "08b0e6354b21ef5dd5e49026028e41831401aca8" -uuid = "51556ac3-7006-55f5-8cb3-34580c88182d" -version = "0.2.17" - -[[deps.ImageFiltering]] -deps = ["CatIndices", "ComputationalResources", "DataStructures", "FFTViews", "FFTW", "ImageBase", "ImageCore", "LinearAlgebra", "OffsetArrays", "PrecompileTools", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "TiledIteration"] -git-tree-sha1 = "432ae2b430a18c58eb7eca9ef8d0f2db90bc749c" -uuid = "6a3955dd-da59-5b1f-98d4-e7296123deb5" -version = "0.7.8" - -[[deps.ImageIO]] -deps = ["FileIO", "IndirectArrays", "JpegTurbo", "LazyModules", "Netpbm", "OpenEXR", "PNGFiles", "QOI", "Sixel", "TiffImages", "UUIDs"] -git-tree-sha1 = "bca20b2f5d00c4fbc192c3212da8fa79f4688009" -uuid = "82e4d734-157c-48bb-816b-45c225c6df19" -version = "0.6.7" - -[[deps.ImageMagick]] -deps = ["FileIO", "ImageCore", "ImageMagick_jll", "InteractiveUtils"] -git-tree-sha1 = "8e2eae13d144d545ef829324f1f0a5a4fe4340f3" -uuid = "6218d12a-5da1-5696-b52f-db25d2ecc6d1" -version = "1.3.1" - -[[deps.ImageMagick_jll]] -deps = ["Artifacts", "Ghostscript_jll", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "OpenJpeg_jll", "Pkg", "Zlib_jll", "libpng_jll"] -git-tree-sha1 = "8d2e786fd090199a91ecbf4a66d03aedd0fb24d4" -uuid = "c73af94c-d91f-53ed-93a7-00f77d67a9d7" -version = "6.9.11+4" - -[[deps.ImageMetadata]] -deps = ["AxisArrays", "ImageAxes", "ImageBase", "ImageCore"] -git-tree-sha1 = "355e2b974f2e3212a75dfb60519de21361ad3cb7" -uuid = "bc367c6b-8a6b-528e-b4bd-a4b897500b49" -version = "0.9.9" - -[[deps.ImageMorphology]] -deps = ["DataStructures", "ImageCore", "LinearAlgebra", "LoopVectorization", "OffsetArrays", "Requires", "TiledIteration"] -git-tree-sha1 = "6f0a801136cb9c229aebea0df296cdcd471dbcd1" -uuid = "787d08f9-d448-5407-9aad-5290dd7ab264" -version = "0.4.5" - -[[deps.ImageQualityIndexes]] -deps = ["ImageContrastAdjustment", "ImageCore", "ImageDistances", "ImageFiltering", "LazyModules", "OffsetArrays", "PrecompileTools", "Statistics"] -git-tree-sha1 = "783b70725ed326340adf225be4889906c96b8fd1" -uuid = "2996bd0c-7a13-11e9-2da2-2f5ce47296a9" -version = "0.3.7" - -[[deps.ImageSegmentation]] -deps = ["Clustering", "DataStructures", "Distances", "Graphs", "ImageCore", "ImageFiltering", "ImageMorphology", "LinearAlgebra", "MetaGraphs", "RegionTrees", "SimpleWeightedGraphs", "StaticArrays", "Statistics"] -git-tree-sha1 = "3ff0ca203501c3eedde3c6fa7fd76b703c336b5f" -uuid = "80713f31-8817-5129-9cf8-209ff8fb23e1" -version = "1.8.2" - -[[deps.ImageShow]] -deps = ["Base64", "ColorSchemes", "FileIO", "ImageBase", "ImageCore", "OffsetArrays", "StackViews"] -git-tree-sha1 = "3b5344bcdbdc11ad58f3b1956709b5b9345355de" -uuid = "4e3cecfd-b093-5904-9786-8bbb286a6a31" -version = "0.3.8" - -[[deps.ImageTransformations]] -deps = ["AxisAlgorithms", "CoordinateTransformations", "ImageBase", "ImageCore", "Interpolations", "OffsetArrays", "Rotations", "StaticArrays"] -git-tree-sha1 = "e0884bdf01bbbb111aea77c348368a86fb4b5ab6" -uuid = "02fcd773-0e25-5acc-982a-7f6622650795" -version = "0.10.1" - -[[deps.Images]] -deps = ["Base64", "FileIO", "Graphics", "ImageAxes", "ImageBase", "ImageBinarization", "ImageContrastAdjustment", "ImageCore", "ImageCorners", "ImageDistances", "ImageFiltering", "ImageIO", "ImageMagick", "ImageMetadata", "ImageMorphology", "ImageQualityIndexes", "ImageSegmentation", "ImageShow", "ImageTransformations", "IndirectArrays", "IntegralArrays", "Random", "Reexport", "SparseArrays", "StaticArrays", "Statistics", "StatsBase", "TiledIteration"] -git-tree-sha1 = "12fdd617c7fe25dc4a6cc804d657cc4b2230302b" -uuid = "916415d5-f1e6-5110-898d-aaa5f9f070e0" -version = "0.26.1" - -[[deps.Imath_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "0936ba688c6d201805a83da835b55c61a180db52" -uuid = "905a6f67-0a94-5f89-b386-d35d92009cd1" -version = "3.1.11+0" - -[[deps.IndirectArrays]] -git-tree-sha1 = "012e604e1c7458645cb8b436f8fba789a51b257f" -uuid = "9b13fd28-a010-5f03-acff-a1bbcff69959" -version = "1.0.0" - -[[deps.Inflate]] -git-tree-sha1 = "ea8031dea4aff6bd41f1df8f2fdfb25b33626381" -uuid = "d25df0c9-e2be-5dd7-82c8-3ad0b3e990b9" -version = "0.1.4" - -[[deps.IntegralArrays]] -deps = ["ColorTypes", "FixedPointNumbers", "IntervalSets"] -git-tree-sha1 = "be8e690c3973443bec584db3346ddc904d4884eb" -uuid = "1d092043-8f09-5a30-832f-7509e371ab51" -version = "0.1.5" - -[[deps.IntelOpenMP_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "be50fe8df3acbffa0274a744f1a99d29c45a57f4" -uuid = "1d5cc7b8-4909-519e-a0f8-d0f5ad9712d0" -version = "2024.1.0+0" - -[[deps.InteractiveUtils]] -deps = ["Markdown"] -uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" - -[[deps.Interpolations]] -deps = ["Adapt", "AxisAlgorithms", "ChainRulesCore", "LinearAlgebra", "OffsetArrays", "Random", "Ratios", "Requires", "SharedArrays", "SparseArrays", "StaticArrays", "WoodburyMatrices"] -git-tree-sha1 = "88a101217d7cb38a7b481ccd50d21876e1d1b0e0" -uuid = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59" -version = "0.15.1" - - [deps.Interpolations.extensions] - InterpolationsUnitfulExt = "Unitful" - - [deps.Interpolations.weakdeps] - Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" - -[[deps.IntervalSets]] -git-tree-sha1 = "dba9ddf07f77f60450fe5d2e2beb9854d9a49bd0" -uuid = "8197267c-284f-5f27-9208-e0e47529a953" -version = "0.7.10" -weakdeps = ["Random", "RecipesBase", "Statistics"] - - [deps.IntervalSets.extensions] - IntervalSetsRandomExt = "Random" - IntervalSetsRecipesBaseExt = "RecipesBase" - IntervalSetsStatisticsExt = "Statistics" - -[[deps.IrrationalConstants]] -git-tree-sha1 = "630b497eafcc20001bba38a4651b327dcfc491d2" -uuid = "92d709cd-6900-40b7-9082-c6be49f344b6" -version = "0.2.2" - -[[deps.IterTools]] -git-tree-sha1 = "42d5f897009e7ff2cf88db414a389e5ed1bdd023" -uuid = "c8e1da08-722c-5040-9ed9-7db0dc04731e" -version = "1.10.0" - -[[deps.JLD2]] -deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Printf", "Reexport", "Requires", "TranscodingStreams", "UUIDs"] -git-tree-sha1 = "5ea6acdd53a51d897672edb694e3cc2912f3f8a7" -uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819" -version = "0.4.46" - -[[deps.JLLWrappers]] -deps = ["Artifacts", "Preferences"] -git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca" -uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210" -version = "1.5.0" - -[[deps.JpegTurbo]] -deps = ["CEnum", "FileIO", "ImageCore", "JpegTurbo_jll", "TOML"] -git-tree-sha1 = "fa6d0bcff8583bac20f1ffa708c3913ca605c611" -uuid = "b835a17e-a41a-41e7-81f0-2f016b05efe0" -version = "0.1.5" - -[[deps.JpegTurbo_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "3336abae9a713d2210bb57ab484b1e065edd7d23" -uuid = "aacddb02-875f-59d6-b918-886e6ef4fbf8" -version = "3.0.2+0" - -[[deps.LERC_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Pkg"] -git-tree-sha1 = "bf36f528eec6634efc60d7ec062008f171071434" -uuid = "88015f11-f218-50d7-93a8-a6af411a945d" -version = "3.0.0+1" - -[[deps.LayoutPointers]] -deps = ["ArrayInterface", "LinearAlgebra", "ManualMemory", "SIMDTypes", "Static", "StaticArrayInterface"] -git-tree-sha1 = "62edfee3211981241b57ff1cedf4d74d79519277" -uuid = "10f19ff3-798f-405d-979b-55457f8fc047" -version = "0.1.15" - -[[deps.LazyArtifacts]] -deps = ["Artifacts", "Pkg"] -uuid = "4af54fe1-eca0-43a8-85a7-787d91b784e3" - -[[deps.LazyModules]] -git-tree-sha1 = "a560dd966b386ac9ae60bdd3a3d3a326062d3c3e" -uuid = "8cdb02fc-e678-4876-92c5-9defec4f444e" -version = "0.3.1" - -[[deps.LibCURL]] -deps = ["LibCURL_jll", "MozillaCACerts_jll"] -uuid = "b27032c2-a3e7-50c8-80cd-2d36dbcbfd21" -version = "0.6.4" - -[[deps.LibCURL_jll]] -deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll", "Zlib_jll", "nghttp2_jll"] -uuid = "deac9b47-8bc7-5906-a0fe-35ac56dc84c0" -version = "8.4.0+0" - -[[deps.LibGit2]] -deps = ["Base64", "LibGit2_jll", "NetworkOptions", "Printf", "SHA"] -uuid = "76f85450-5226-5b5a-8eaa-529ad045b433" - -[[deps.LibGit2_jll]] -deps = ["Artifacts", "LibSSH2_jll", "Libdl", "MbedTLS_jll"] -uuid = "e37daf67-58a4-590a-8e99-b0245dd2ffc5" -version = "1.6.4+0" - -[[deps.LibSSH2_jll]] -deps = ["Artifacts", "Libdl", "MbedTLS_jll"] -uuid = "29816b5a-b9ab-546f-933c-edad1886dfa8" -version = "1.11.0+1" - -[[deps.Libdl]] -uuid = "8f399da3-3557-5675-b5ff-fb832c97cbdb" - -[[deps.Libtiff_jll]] -deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "LERC_jll", "Libdl", "Pkg", "Zlib_jll", "Zstd_jll"] -git-tree-sha1 = "3eb79b0ca5764d4799c06699573fd8f533259713" -uuid = "89763e89-9b03-5906-acba-b20f662cd828" -version = "4.4.0+0" - -[[deps.LinearAlgebra]] -deps = ["Libdl", "OpenBLAS_jll", "libblastrampoline_jll"] -uuid = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" - -[[deps.LittleCMS_jll]] -deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Libtiff_jll", "Pkg"] -git-tree-sha1 = "110897e7db2d6836be22c18bffd9422218ee6284" -uuid = "d3a379c0-f9a3-5b72-a4c0-6bf4d2e8af0f" -version = "2.12.0+0" - -[[deps.LogExpFunctions]] -deps = ["DocStringExtensions", "IrrationalConstants", "LinearAlgebra"] -git-tree-sha1 = "18144f3e9cbe9b15b070288eef858f71b291ce37" -uuid = "2ab3a3ac-af41-5b50-aa03-7779005ae688" -version = "0.3.27" - - [deps.LogExpFunctions.extensions] - LogExpFunctionsChainRulesCoreExt = "ChainRulesCore" - LogExpFunctionsChangesOfVariablesExt = "ChangesOfVariables" - LogExpFunctionsInverseFunctionsExt = "InverseFunctions" - - [deps.LogExpFunctions.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - ChangesOfVariables = "9e997f8a-9a97-42d5-a9f1-ce6bfc15e2c0" - InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112" - -[[deps.Logging]] -uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" - -[[deps.LoopVectorization]] -deps = ["ArrayInterface", "CPUSummary", "CloseOpenIntervals", "DocStringExtensions", "HostCPUFeatures", "IfElse", "LayoutPointers", "LinearAlgebra", "OffsetArrays", "PolyesterWeave", "PrecompileTools", "SIMDTypes", "SLEEFPirates", "Static", "StaticArrayInterface", "ThreadingUtilities", "UnPack", "VectorizationBase"] -git-tree-sha1 = "8f6786d8b2b3248d79db3ad359ce95382d5a6df8" -uuid = "bdcacae8-1622-11e9-2a5c-532679323890" -version = "0.12.170" - - [deps.LoopVectorization.extensions] - ForwardDiffExt = ["ChainRulesCore", "ForwardDiff"] - SpecialFunctionsExt = "SpecialFunctions" - - [deps.LoopVectorization.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" - SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b" - -[[deps.MKL_jll]] -deps = ["Artifacts", "IntelOpenMP_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "oneTBB_jll"] -git-tree-sha1 = "80b2833b56d466b3858d565adcd16a4a05f2089b" -uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7" -version = "2024.1.0+0" - -[[deps.MacroTools]] -deps = ["Markdown", "Random"] -git-tree-sha1 = "2fa9ee3e63fd3a4f7a9a4f4744a52f4856de82df" -uuid = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09" -version = "0.5.13" - -[[deps.ManualMemory]] -git-tree-sha1 = "bcaef4fc7a0cfe2cba636d84cda54b5e4e4ca3cd" -uuid = "d125e4d3-2237-4719-b19c-fa641b8a4667" -version = "0.1.8" - -[[deps.MappedArrays]] -git-tree-sha1 = "2dab0221fe2b0f2cb6754eaa743cc266339f527e" -uuid = "dbb5928d-eab1-5f90-85c2-b9b0edb7c900" -version = "0.4.2" - -[[deps.Markdown]] -deps = ["Base64"] -uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" - -[[deps.MbedTLS_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "c8ffd9c3-330d-5841-b78e-0817d7145fa1" -version = "2.28.2+1" - -[[deps.MetaGraphs]] -deps = ["Graphs", "JLD2", "Random"] -git-tree-sha1 = "1130dbe1d5276cb656f6e1094ce97466ed700e5a" -uuid = "626554b9-1ddb-594c-aa3c-2596fe9399a5" -version = "0.7.2" - -[[deps.Missings]] -deps = ["DataAPI"] -git-tree-sha1 = "ec4f7fbeab05d7747bdf98eb74d130a2a2ed298d" -uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28" -version = "1.2.0" - -[[deps.Mmap]] -uuid = "a63ad114-7e13-5084-954f-fe012c677804" - -[[deps.MosaicViews]] -deps = ["MappedArrays", "OffsetArrays", "PaddedViews", "StackViews"] -git-tree-sha1 = "7b86a5d4d70a9f5cdf2dacb3cbe6d251d1a61dbe" -uuid = "e94cdb99-869f-56ef-bcf0-1ae2bcbe0389" -version = "0.3.4" - -[[deps.MozillaCACerts_jll]] -uuid = "14a3606d-f60d-562e-9121-12d972cd8159" -version = "2023.1.10" - -[[deps.NaNMath]] -deps = ["OpenLibm_jll"] -git-tree-sha1 = "0877504529a3e5c3343c6f8b4c0381e57e4387e4" -uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3" -version = "1.0.2" - -[[deps.NearestNeighbors]] -deps = ["Distances", "StaticArrays"] -git-tree-sha1 = "ded64ff6d4fdd1cb68dfcbb818c69e144a5b2e4c" -uuid = "b8a86587-4115-5ab1-83bc-aa920d37bbce" -version = "0.4.16" - -[[deps.Netpbm]] -deps = ["FileIO", "ImageCore", "ImageMetadata"] -git-tree-sha1 = "d92b107dbb887293622df7697a2223f9f8176fcd" -uuid = "f09324ee-3d7c-5217-9330-fc30815ba969" -version = "1.1.1" - -[[deps.NetworkOptions]] -uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" -version = "1.2.0" - -[[deps.OffsetArrays]] -git-tree-sha1 = "e64b4f5ea6b7389f6f046d13d4896a8f9c1ba71e" -uuid = "6fe1bfb0-de20-5000-8ca7-80f57d26f881" -version = "1.14.0" -weakdeps = ["Adapt"] - - [deps.OffsetArrays.extensions] - OffsetArraysAdaptExt = "Adapt" - -[[deps.OpenBLAS_jll]] -deps = ["Artifacts", "CompilerSupportLibraries_jll", "Libdl"] -uuid = "4536629a-c528-5b80-bd46-f80d51c5b363" -version = "0.3.23+4" - -[[deps.OpenEXR]] -deps = ["Colors", "FileIO", "OpenEXR_jll"] -git-tree-sha1 = "327f53360fdb54df7ecd01e96ef1983536d1e633" -uuid = "52e1d378-f018-4a11-a4be-720524705ac7" -version = "0.3.2" - -[[deps.OpenEXR_jll]] -deps = ["Artifacts", "Imath_jll", "JLLWrappers", "Libdl", "Zlib_jll"] -git-tree-sha1 = "8292dd5c8a38257111ada2174000a33745b06d4e" -uuid = "18a262bb-aa17-5467-a713-aee519bc75cb" -version = "3.2.4+0" - -[[deps.OpenJpeg_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Libtiff_jll", "LittleCMS_jll", "Pkg", "libpng_jll"] -git-tree-sha1 = "76374b6e7f632c130e78100b166e5a48464256f8" -uuid = "643b3616-a352-519d-856d-80112ee9badc" -version = "2.4.0+0" - -[[deps.OpenLibm_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "05823500-19ac-5b8b-9628-191a04bc5112" -version = "0.8.1+2" - -[[deps.OrderedCollections]] -git-tree-sha1 = "dfdf5519f235516220579f949664f1bf44e741c5" -uuid = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" -version = "1.6.3" - -[[deps.PNGFiles]] -deps = ["Base64", "CEnum", "ImageCore", "IndirectArrays", "OffsetArrays", "libpng_jll"] -git-tree-sha1 = "67186a2bc9a90f9f85ff3cc8277868961fb57cbd" -uuid = "f57f5aa1-a3ce-4bc8-8ab9-96f992907883" -version = "0.4.3" - -[[deps.PaddedViews]] -deps = ["OffsetArrays"] -git-tree-sha1 = "0fac6313486baae819364c52b4f483450a9d793f" -uuid = "5432bcbf-9aad-5242-b902-cca2824c8663" -version = "0.5.12" - -[[deps.Parameters]] -deps = ["OrderedCollections", "UnPack"] -git-tree-sha1 = "34c0e9ad262e5f7fc75b10a9952ca7692cfc5fbe" -uuid = "d96e819e-fc66-5662-9728-84c9c7592b0a" -version = "0.12.3" - -[[deps.Pkg]] -deps = ["Artifacts", "Dates", "Downloads", "FileWatching", "LibGit2", "Libdl", "Logging", "Markdown", "Printf", "REPL", "Random", "SHA", "Serialization", "TOML", "Tar", "UUIDs", "p7zip_jll"] -uuid = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" -version = "1.10.0" - -[[deps.PkgVersion]] -deps = ["Pkg"] -git-tree-sha1 = "f9501cc0430a26bc3d156ae1b5b0c1b47af4d6da" -uuid = "eebad327-c553-4316-9ea0-9fa01ccd7688" -version = "0.3.3" - -[[deps.PolyesterWeave]] -deps = ["BitTwiddlingConvenienceFunctions", "CPUSummary", "IfElse", "Static", "ThreadingUtilities"] -git-tree-sha1 = "240d7170f5ffdb285f9427b92333c3463bf65bf6" -uuid = "1d0040c9-8b98-4ee7-8388-3f51789ca0ad" -version = "0.2.1" - -[[deps.Polynomials]] -deps = ["LinearAlgebra", "RecipesBase"] -git-tree-sha1 = "3aa2bb4982e575acd7583f01531f241af077b163" -uuid = "f27b6e38-b328-58d1-80ce-0feddd5e7a45" -version = "3.2.13" - - [deps.Polynomials.extensions] - PolynomialsChainRulesCoreExt = "ChainRulesCore" - PolynomialsMakieCoreExt = "MakieCore" - PolynomialsMutableArithmeticsExt = "MutableArithmetics" - - [deps.Polynomials.weakdeps] - ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" - MakieCore = "20f20a25-4f0e-4fdf-b5d1-57303727442b" - MutableArithmetics = "d8a4904e-b15c-11e9-3269-09a3773c0cb0" - -[[deps.PrecompileTools]] -deps = ["Preferences"] -git-tree-sha1 = "5aa36f7049a63a1528fe8f7c3f2113413ffd4e1f" -uuid = "aea7be01-6a6a-4083-8856-8a6e6704d82a" -version = "1.2.1" - -[[deps.Preferences]] -deps = ["TOML"] -git-tree-sha1 = "9306f6085165d270f7e3db02af26a400d580f5c6" -uuid = "21216c6a-2e73-6563-6e65-726566657250" -version = "1.4.3" - -[[deps.Printf]] -deps = ["Unicode"] -uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" - -[[deps.ProgressMeter]] -deps = ["Distributed", "Printf"] -git-tree-sha1 = "763a8ceb07833dd51bb9e3bbca372de32c0605ad" -uuid = "92933f4c-e287-5a05-a399-4b506db050ca" -version = "1.10.0" - -[[deps.QOI]] -deps = ["ColorTypes", "FileIO", "FixedPointNumbers"] -git-tree-sha1 = "18e8f4d1426e965c7b532ddd260599e1510d26ce" -uuid = "4b34888f-f399-49d4-9bb3-47ed5cae4e65" -version = "1.0.0" - -[[deps.Quaternions]] -deps = ["LinearAlgebra", "Random", "RealDot"] -git-tree-sha1 = "994cc27cdacca10e68feb291673ec3a76aa2fae9" -uuid = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0" -version = "0.7.6" - -[[deps.REPL]] -deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] -uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" - -[[deps.Random]] -deps = ["SHA"] -uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" - -[[deps.RangeArrays]] -git-tree-sha1 = "b9039e93773ddcfc828f12aadf7115b4b4d225f5" -uuid = "b3c3ace0-ae52-54e7-9d0b-2c1406fd6b9d" -version = "0.3.2" - -[[deps.Ratios]] -deps = ["Requires"] -git-tree-sha1 = "1342a47bf3260ee108163042310d26f2be5ec90b" -uuid = "c84ed2f1-dad5-54f0-aa8e-dbefe2724439" -version = "0.4.5" -weakdeps = ["FixedPointNumbers"] - - [deps.Ratios.extensions] - RatiosFixedPointNumbersExt = "FixedPointNumbers" - -[[deps.RealDot]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "9f0a1b71baaf7650f4fa8a1d168c7fb6ee41f0c9" -uuid = "c1ae055f-0cd5-4b69-90a6-9a35b1a98df9" -version = "0.1.0" - -[[deps.RecipesBase]] -deps = ["PrecompileTools"] -git-tree-sha1 = "5c3d09cc4f31f5fc6af001c250bf1278733100ff" -uuid = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" -version = "1.3.4" - -[[deps.Reexport]] -git-tree-sha1 = "45e428421666073eab6f2da5c9d310d99bb12f9b" -uuid = "189a3867-3050-52da-a836-e630ba90ab69" -version = "1.2.2" - -[[deps.RegionTrees]] -deps = ["IterTools", "LinearAlgebra", "StaticArrays"] -git-tree-sha1 = "4618ed0da7a251c7f92e869ae1a19c74a7d2a7f9" -uuid = "dee08c22-ab7f-5625-9660-a9af2021b33f" -version = "0.3.2" - -[[deps.Requires]] -deps = ["UUIDs"] -git-tree-sha1 = "838a3a4188e2ded87a4f9f184b4b0d78a1e91cb7" -uuid = "ae029012-a4dd-5104-9daa-d747884805df" -version = "1.3.0" - -[[deps.Rotations]] -deps = ["LinearAlgebra", "Quaternions", "Random", "StaticArrays"] -git-tree-sha1 = "2a0a5d8569f481ff8840e3b7c84bbf188db6a3fe" -uuid = "6038ab10-8711-5258-84ad-4b1120ba62dc" -version = "1.7.0" -weakdeps = ["RecipesBase"] - - [deps.Rotations.extensions] - RotationsRecipesBaseExt = "RecipesBase" - -[[deps.SHA]] -uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" -version = "0.7.0" - -[[deps.SIMDTypes]] -git-tree-sha1 = "330289636fb8107c5f32088d2741e9fd7a061a5c" -uuid = "94e857df-77ce-4151-89e5-788b33177be4" -version = "0.1.0" - -[[deps.SLEEFPirates]] -deps = ["IfElse", "Static", "VectorizationBase"] -git-tree-sha1 = "3aac6d68c5e57449f5b9b865c9ba50ac2970c4cf" -uuid = "476501e8-09a2-5ece-8869-fb82de89a1fa" -version = "0.6.42" - -[[deps.Serialization]] -uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" - -[[deps.SharedArrays]] -deps = ["Distributed", "Mmap", "Random", "Serialization"] -uuid = "1a1011a3-84de-559e-8e89-a11a2f7dc383" - -[[deps.SimpleTraits]] -deps = ["InteractiveUtils", "MacroTools"] -git-tree-sha1 = "5d7e3f4e11935503d3ecaf7186eac40602e7d231" -uuid = "699a6c99-e7fa-54fc-8d76-47d257e15c1d" -version = "0.9.4" - -[[deps.SimpleWeightedGraphs]] -deps = ["Graphs", "LinearAlgebra", "Markdown", "SparseArrays"] -git-tree-sha1 = "4b33e0e081a825dbfaf314decf58fa47e53d6acb" -uuid = "47aef6b3-ad0c-573a-a1e2-d07658019622" -version = "1.4.0" - -[[deps.Sixel]] -deps = ["Dates", "FileIO", "ImageCore", "IndirectArrays", "OffsetArrays", "REPL", "libsixel_jll"] -git-tree-sha1 = "2da10356e31327c7096832eb9cd86307a50b1eb6" -uuid = "45858cf5-a6b0-47a3-bbea-62219f50df47" -version = "0.1.3" - -[[deps.Sockets]] -uuid = "6462fe0b-24de-5631-8697-dd941f90decc" - -[[deps.SortingAlgorithms]] -deps = ["DataStructures"] -git-tree-sha1 = "66e0a8e672a0bdfca2c3f5937efb8538b9ddc085" -uuid = "a2af1166-a08f-5f64-846c-94a0d3cef48c" -version = "1.2.1" - -[[deps.SparseArrays]] -deps = ["Libdl", "LinearAlgebra", "Random", "Serialization", "SuiteSparse_jll"] -uuid = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" -version = "1.10.0" - -[[deps.StackViews]] -deps = ["OffsetArrays"] -git-tree-sha1 = "46e589465204cd0c08b4bd97385e4fa79a0c770c" -uuid = "cae243ae-269e-4f55-b966-ac2d0dc13c15" -version = "0.1.1" - -[[deps.Static]] -deps = ["IfElse"] -git-tree-sha1 = "d2fdac9ff3906e27f7a618d47b676941baa6c80c" -uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3" -version = "0.8.10" - -[[deps.StaticArrayInterface]] -deps = ["ArrayInterface", "Compat", "IfElse", "LinearAlgebra", "PrecompileTools", "Requires", "SparseArrays", "Static", "SuiteSparse"] -git-tree-sha1 = "5d66818a39bb04bf328e92bc933ec5b4ee88e436" -uuid = "0d7ed370-da01-4f52-bd93-41d350b8b718" -version = "1.5.0" -weakdeps = ["OffsetArrays", "StaticArrays"] - - [deps.StaticArrayInterface.extensions] - StaticArrayInterfaceOffsetArraysExt = "OffsetArrays" - StaticArrayInterfaceStaticArraysExt = "StaticArrays" - -[[deps.StaticArrays]] -deps = ["LinearAlgebra", "PrecompileTools", "Random", "StaticArraysCore"] -git-tree-sha1 = "bf074c045d3d5ffd956fa0a461da38a44685d6b2" -uuid = "90137ffa-7385-5640-81b9-e52037218182" -version = "1.9.3" -weakdeps = ["ChainRulesCore", "Statistics"] - - [deps.StaticArrays.extensions] - StaticArraysChainRulesCoreExt = "ChainRulesCore" - StaticArraysStatisticsExt = "Statistics" - -[[deps.StaticArraysCore]] -git-tree-sha1 = "36b3d696ce6366023a0ea192b4cd442268995a0d" -uuid = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" -version = "1.4.2" - -[[deps.Statistics]] -deps = ["LinearAlgebra", "SparseArrays"] -uuid = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" -version = "1.10.0" - -[[deps.StatsAPI]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "1ff449ad350c9c4cbc756624d6f8a8c3ef56d3ed" -uuid = "82ae8749-77ed-4fe6-ae5f-f523153014b0" -version = "1.7.0" - -[[deps.StatsBase]] -deps = ["DataAPI", "DataStructures", "LinearAlgebra", "LogExpFunctions", "Missings", "Printf", "Random", "SortingAlgorithms", "SparseArrays", "Statistics", "StatsAPI"] -git-tree-sha1 = "5cf7606d6cef84b543b483848d4ae08ad9832b21" -uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" -version = "0.34.3" - -[[deps.SuiteSparse]] -deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"] -uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9" - -[[deps.SuiteSparse_jll]] -deps = ["Artifacts", "Libdl", "libblastrampoline_jll"] -uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c" -version = "7.2.1+1" - -[[deps.TOML]] -deps = ["Dates"] -uuid = "fa267f1f-6049-4f14-aa54-33bafae1ed76" -version = "1.0.3" - -[[deps.Tar]] -deps = ["ArgTools", "SHA"] -uuid = "a4e569a6-e804-4fa4-b0f3-eef7a1d5b13e" -version = "1.10.0" - -[[deps.TensorCore]] -deps = ["LinearAlgebra"] -git-tree-sha1 = "1feb45f88d133a655e001435632f019a9a1bcdb6" -uuid = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50" -version = "0.1.1" - -[[deps.ThreadingUtilities]] -deps = ["ManualMemory"] -git-tree-sha1 = "eda08f7e9818eb53661b3deb74e3159460dfbc27" -uuid = "8290d209-cae3-49c0-8002-c8c24d57dab5" -version = "0.5.2" - -[[deps.TiffImages]] -deps = ["ColorTypes", "DataStructures", "DocStringExtensions", "FileIO", "FixedPointNumbers", "IndirectArrays", "Inflate", "Mmap", "OffsetArrays", "PkgVersion", "ProgressMeter", "UUIDs"] -git-tree-sha1 = "34cc045dd0aaa59b8bbe86c644679bc57f1d5bd0" -uuid = "731e570b-9d59-4bfa-96dc-6df516fadf69" -version = "0.6.8" - -[[deps.TiledIteration]] -deps = ["OffsetArrays", "StaticArrayInterface"] -git-tree-sha1 = "1176cc31e867217b06928e2f140c90bd1bc88283" -uuid = "06e1c1a7-607b-532d-9fad-de7d9aa2abac" -version = "0.5.0" - -[[deps.TranscodingStreams]] -git-tree-sha1 = "5d54d076465da49d6746c647022f3b3674e64156" -uuid = "3bb67fe8-82b1-5028-8e26-92a6c54297fa" -version = "0.10.8" - - [deps.TranscodingStreams.extensions] - TestExt = ["Test", "Random"] - - [deps.TranscodingStreams.weakdeps] - Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" - Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[[deps.UUIDs]] -deps = ["Random", "SHA"] -uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4" - -[[deps.UnPack]] -git-tree-sha1 = "387c1f73762231e86e0c9c5443ce3b4a0a9a0c2b" -uuid = "3a884ed6-31ef-47d7-9d2a-63182c4928ed" -version = "1.0.2" - -[[deps.Unicode]] -uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" - -[[deps.VectorizationBase]] -deps = ["ArrayInterface", "CPUSummary", "HostCPUFeatures", "IfElse", "LayoutPointers", "Libdl", "LinearAlgebra", "SIMDTypes", "Static", "StaticArrayInterface"] -git-tree-sha1 = "6129a4faf6242e7c3581116fbe3270f3ab17c90d" -uuid = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f" -version = "0.21.67" - -[[deps.WoodburyMatrices]] -deps = ["LinearAlgebra", "SparseArrays"] -git-tree-sha1 = "c1a7aa6219628fcd757dede0ca95e245c5cd9511" -uuid = "efce3f68-66dc-5838-9240-27a6d6f5f9b6" -version = "1.0.0" - -[[deps.Zlib_jll]] -deps = ["Libdl"] -uuid = "83775a58-1f1d-513f-b197-d71354ab007a" -version = "1.2.13+1" - -[[deps.Zstd_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "e678132f07ddb5bfa46857f0d7620fb9be675d3b" -uuid = "3161d3a3-bdf6-5164-811a-617609db77b4" -version = "1.5.6+0" - -[[deps.libblastrampoline_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "8e850b90-86db-534c-a0d3-1478176c7d93" -version = "5.8.0+1" - -[[deps.libpng_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl", "Zlib_jll"] -git-tree-sha1 = "d7015d2e18a5fd9a4f47de711837e980519781a4" -uuid = "b53b4c65-9356-5827-b1ea-8c7a1a84506f" -version = "1.6.43+1" - -[[deps.libsixel_jll]] -deps = ["Artifacts", "JLLWrappers", "JpegTurbo_jll", "Libdl", "Pkg", "libpng_jll"] -git-tree-sha1 = "d4f63314c8aa1e48cd22aa0c17ed76cd1ae48c3c" -uuid = "075b6546-f08a-558a-be8f-8157d0f608a5" -version = "1.10.3+0" - -[[deps.nghttp2_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "8e850ede-7688-5339-a07c-302acd2aaf8d" -version = "1.52.0+1" - -[[deps.oneTBB_jll]] -deps = ["Artifacts", "JLLWrappers", "Libdl"] -git-tree-sha1 = "7d0ea0f4895ef2f5cb83645fa689e52cb55cf493" -uuid = "1317d2d5-d96f-522e-a858-c73665f53c3e" -version = "2021.12.0+0" - -[[deps.p7zip_jll]] -deps = ["Artifacts", "Libdl"] -uuid = "3f19e933-33d8-53b3-aaab-bd5110c3b7a0" -version = "17.4.0+2" diff --git a/docs/report/019/Project.toml b/docs/report/019/Project.toml deleted file mode 100644 index e369a863..00000000 --- a/docs/report/019/Project.toml +++ /dev/null @@ -1,2 +0,0 @@ -[deps] -Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0" diff --git a/docs/report/020.ipynb b/docs/report/020.ipynb deleted file mode 100644 index 4a218fb9..00000000 --- a/docs/report/020.ipynb +++ /dev/null @@ -1,1377 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "sympy", - "tensorwaves" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Amplitude analysis with zfit\n", - "TR-020\n", - "^^^\n", - "This reports builds a [simple symbolic amplitude model](https://tensorwaves.readthedocs.io/en/0.4.5/amplitude-analysis.html) with {mod}`qrules` and {mod}`ampform` and feeds it to [zfit](https://zfit.rtfd.io) instead of {mod}`tensorwaves`.\n", - "+++\n", - "✅ [compwa.github.io#151](https://github.com/ComPWA/compwa.github.io/issues/151)
\n", - "WIP [ComPWA/.github#14](https://github.com/ComPWA/.github/issues/14)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Amplitude analysis with zfit" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform[viz]~=0.14.1 hepstats~=0.6.0 pandas~=1.4.2 sympy~=1.10.1 tensorwaves[jax,pwa]~=0.4.8 zfit~=0.10.0" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "import logging\n", - "import os\n", - "import warnings\n", - "\n", - "JAX_LOGGER = logging.getLogger(\"absl\")\n", - "JAX_LOGGER.setLevel(logging.ERROR)\n", - "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", - "warnings.simplefilter(\"ignore\", UserWarning)\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Formulating the model" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import qrules\n", - "\n", - "reaction = qrules.generate_transitions(\n", - " initial_state=(\"J/psi(1S)\", [-1, +1]),\n", - " final_state=[\"gamma\", \"pi0\", \"pi0\"],\n", - " allowed_intermediate_particles=[\"f(0)\"],\n", - " allowed_interaction_types=[\"strong\", \"EM\"],\n", - " formalism=\"helicity\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "import graphviz\n", - "\n", - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![image](https://user-images.githubusercontent.com/29308176/194346498-ff9b9379-90f1-4348-81ba-d5a57ddccc83.png)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import ampform\n", - "from ampform.dynamics.builder import (\n", - " create_non_dynamic_with_ff,\n", - " create_relativistic_breit_wigner_with_ff,\n", - ")\n", - "\n", - "model_builder = ampform.get_builder(reaction)\n", - "model_builder.scalar_initial_state_mass = True\n", - "model_builder.stable_final_state_ids = [0, 1, 2]\n", - "model_builder.set_dynamics(\"J/psi(1S)\", create_non_dynamic_with_ff)\n", - "for name in reaction.get_intermediate_particles().names:\n", - " model_builder.set_dynamics(name, create_relativistic_breit_wigner_with_ff)\n", - "model = model_builder.formulate()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phase space sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator\n", - "\n", - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "phsp_momenta = phsp_generator.generate(100_000, rng)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Intensity-based sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.function.sympy import create_function\n", - "\n", - "unfolded_expression = model.expression.doit()\n", - "fixed_intensity_func = create_function(\n", - " unfolded_expression.xreplace(model.parameter_defaults),\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import SympyDataTransformer\n", - "\n", - "transform_momenta = SympyDataTransformer.from_sympy(\n", - " model.kinematic_variables, backend=\"jax\"\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.data import (\n", - " IntensityDistributionGenerator,\n", - " TFWeightedPhaseSpaceGenerator,\n", - ")\n", - "\n", - "weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "data_generator = IntensityDistributionGenerator(\n", - " domain_generator=weighted_phsp_generator,\n", - " function=fixed_intensity_func,\n", - " domain_transformer=transform_momenta,\n", - ")\n", - "data_momenta = data_generator.generate(10_000, rng)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
m_12phi_0phi_1^12theta_0theta_1^12
01.499845+0.000000j2.941350-0.9844192.3446171.064114
10.580070+0.000000j1.4221270.1837251.0866671.535691
21.495937+0.000000j2.6955853.0636220.7779781.730394
31.172263+0.000000j0.5278501.5156851.3435300.602596
41.581282+0.000000j-0.678981-2.9515562.9874701.959462
..................
99951.486016+0.000000j-1.271331-1.3874952.7925712.565453
99960.584599+0.000000j-2.452912-1.9570861.0708892.313677
99971.956302+0.000000j0.3783142.7114960.5889871.551541
99981.585024+0.000000j-0.816920-1.1663152.0760681.807813
99991.712966+0.000000j0.6046570.5533471.2641402.079405
\n", - "

10000 rows × 5 columns

\n", - "
" - ], - "text/plain": [ - " m_12 phi_0 phi_1^12 theta_0 theta_1^12\n", - "0 1.499845+0.000000j 2.941350 -0.984419 2.344617 1.064114\n", - "1 0.580070+0.000000j 1.422127 0.183725 1.086667 1.535691\n", - "2 1.495937+0.000000j 2.695585 3.063622 0.777978 1.730394\n", - "3 1.172263+0.000000j 0.527850 1.515685 1.343530 0.602596\n", - "4 1.581282+0.000000j -0.678981 -2.951556 2.987470 1.959462\n", - "... ... ... ... ... ...\n", - "9995 1.486016+0.000000j -1.271331 -1.387495 2.792571 2.565453\n", - "9996 0.584599+0.000000j -2.452912 -1.957086 1.070889 2.313677\n", - "9997 1.956302+0.000000j 0.378314 2.711496 0.588987 1.551541\n", - "9998 1.585024+0.000000j -0.816920 -1.166315 2.076068 1.807813\n", - "9999 1.712966+0.000000j 0.604657 0.553347 1.264140 2.079405\n", - "\n", - "[10000 rows x 5 columns]" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import pandas as pd\n", - "\n", - "phsp = transform_momenta(phsp_momenta)\n", - "data = transform_momenta(data_momenta)\n", - "pd.DataFrame(data)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from matplotlib import cm\n", - "\n", - "resonances = sorted(\n", - " reaction.get_intermediate_particles(),\n", - " key=lambda p: p.mass,\n", - ")\n", - "evenly_spaced_interval = np.linspace(0, 1, len(resonances))\n", - "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n", - "fig, ax = plt.subplots(figsize=(9, 4))\n", - "ax.hist(\n", - " np.real(data[\"m_12\"]),\n", - " bins=100,\n", - " alpha=0.5,\n", - " density=True,\n", - ")\n", - "ax.set_xlabel(\"$m$ [GeV]\")\n", - "for p, color in zip(resonances, colors):\n", - " ax.axvline(x=p.mass, linestyle=\"dotted\", label=p.name, color=color)\n", - "ax.legend()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/194346791-c6e9f66f-ccf8-48b0-9b4b-5ea964db88d4.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Fit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Determine free parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "initial_parameters = {\n", - " R\"C_{J/\\psi(1S) \\to {f_{0}(1500)}_{0} \\gamma_{+1}; f_{0}(1500) \\to \\pi^{0}_{0} \\pi^{0}_{0}}\": (\n", - " 1.0 + 0.0j\n", - " ),\n", - " \"m_{f_{0}(500)}\": 0.4,\n", - " \"m_{f_{0}(980)}\": 0.88,\n", - " \"m_{f_{0}(1370)}\": 1.22,\n", - " \"m_{f_{0}(1500)}\": 1.45,\n", - " \"m_{f_{0}(1710)}\": 1.83,\n", - " R\"\\Gamma_{f_{0}(500)}\": 0.3,\n", - " R\"\\Gamma_{f_{0}(980)}\": 0.1,\n", - " R\"\\Gamma_{f_{0}(1710)}\": 0.3,\n", - "}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Parametrized function and caching" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.function.sympy import create_parametrized_function\n", - "\n", - "intensity_func = create_parametrized_function(\n", - " expression=unfolded_expression,\n", - " parameters=model.parameter_defaults,\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.estimator import create_cached_function\n", - "\n", - "free_parameter_symbols = [\n", - " symbol\n", - " for symbol in model.parameter_defaults\n", - " if symbol.name in set(initial_parameters)\n", - "]\n", - "cached_intensity_func, transform_to_cache = create_cached_function(\n", - " unfolded_expression,\n", - " parameters=model.parameter_defaults,\n", - " free_parameters=free_parameter_symbols,\n", - " backend=\"jax\",\n", - ")\n", - "cached_data = transform_to_cache(data)\n", - "cached_phsp = transform_to_cache(phsp)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Estimator" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from tensorwaves.estimator import UnbinnedNLL\n", - "\n", - "estimator = UnbinnedNLL(\n", - " intensity_func,\n", - " data=data,\n", - " phsp=phsp,\n", - " backend=\"jax\",\n", - ")\n", - "estimator_with_caching = UnbinnedNLL(\n", - " cached_intensity_func,\n", - " data=cached_data,\n", - " phsp=cached_phsp,\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Optimize fit parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "from matplotlib import cm\n", - "\n", - "reaction_info = model.reaction_info\n", - "resonances = sorted(\n", - " reaction_info.get_intermediate_particles(),\n", - " key=lambda p: p.mass,\n", - ")\n", - "\n", - "evenly_spaced_interval = np.linspace(0, 1, len(resonances))\n", - "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n", - "\n", - "\n", - "def indicate_masses(ax):\n", - " ax.set_xlabel(\"$m$ [GeV]\")\n", - " for color, resonance in zip(colors, resonances):\n", - " ax.axvline(\n", - " x=resonance.mass,\n", - " linestyle=\"dotted\",\n", - " label=resonance.name,\n", - " color=color,\n", - " )\n", - "\n", - "\n", - "def compare_model(\n", - " variable_name,\n", - " data,\n", - " phsp,\n", - " function,\n", - " bins=100,\n", - "):\n", - " intensities = function(phsp)\n", - " _, ax = plt.subplots(figsize=(9, 4))\n", - " data_projection = np.real(data[variable_name])\n", - " ax = plt.gca()\n", - " ax.hist(\n", - " data_projection,\n", - " bins=bins,\n", - " alpha=0.5,\n", - " label=\"data\",\n", - " density=True,\n", - " )\n", - " phsp_projection = np.real(phsp[variable_name])\n", - " ax.hist(\n", - " phsp_projection,\n", - " weights=np.array(intensities),\n", - " bins=bins,\n", - " histtype=\"step\",\n", - " color=\"red\",\n", - " label=\"fit model\",\n", - " density=True,\n", - " )\n", - " indicate_masses(ax)\n", - " ax.legend()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "original_parameters = intensity_func.parameters\n", - "intensity_func.update_parameters(initial_parameters)\n", - "compare_model(\"m_12\", data, phsp, intensity_func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/194347028-0fe1354e-0f2d-4684-a6d7-172ab53db8d5.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "20e7fb7c0b5a482f801be4eb16b4d6b7", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "0it [00:00, ?it/s]" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "FitResult(\n", - " minimum_valid=True,\n", - " execution_time=7.060763359069824,\n", - " function_calls=539,\n", - " estimator_value=-4891.01730754809,\n", - " parameter_values={\n", - " 'm_{f_{0}(500)}': 0.6102707294724865,\n", - " 'm_{f_{0}(980)}': 0.9902119846615327,\n", - " 'm_{f_{0}(1370)}': 1.3456300421915652,\n", - " 'm_{f_{0}(1500)}': 1.50502995100389,\n", - " 'm_{f_{0}(1710)}': 1.7096496843682751,\n", - " '\\\\Gamma_{f_{0}(500)}': 0.4226040807774344,\n", - " '\\\\Gamma_{f_{0}(980)}': 0.06479339507889993,\n", - " '\\\\Gamma_{f_{0}(1710)}': 0.13301019075808046,\n", - " 'C_{J/\\\\psi(1S) \\\\to {f_{0}(1500)}_{0} \\\\gamma_{+1}; f_{0}(1500) \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0}}': (1.0699249014701417-0.018664035501929042j),\n", - " },\n", - " parameter_errors={\n", - " 'm_{f_{0}(500)}': 0.006168655466103817,\n", - " 'm_{f_{0}(980)}': 0.0016283609785222876,\n", - " 'm_{f_{0}(1370)}': 0.005122588422790316,\n", - " 'm_{f_{0}(1500)}': 0.0033157863330869892,\n", - " 'm_{f_{0}(1710)}': 0.0025660827305775034,\n", - " '\\\\Gamma_{f_{0}(500)}': 0.023838186430050128,\n", - " '\\\\Gamma_{f_{0}(980)}': 0.003556673018336295,\n", - " '\\\\Gamma_{f_{0}(1710)}': 0.007573518980113613,\n", - " 'C_{J/\\\\psi(1S) \\\\to {f_{0}(1500)}_{0} \\\\gamma_{+1}; f_{0}(1500) \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0}}': (0.04106392764099969+0.07043808181098646j),\n", - " },\n", - ")" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from tensorwaves.optimizer import Minuit2\n", - "from tensorwaves.optimizer.callbacks import CSVSummary\n", - "\n", - "minuit2 = Minuit2(\n", - " callback=CSVSummary(\"fit_traceback.csv\"),\n", - " use_analytic_gradient=False,\n", - ")\n", - "fit_result = minuit2.optimize(estimator, initial_parameters)\n", - "fit_result" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "d2728281b6e842908798f10cfc6e1a30", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "0it [00:00, ?it/s]" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/plain": [ - "FitResult(\n", - " minimum_valid=True,\n", - " execution_time=3.6658225059509277,\n", - " function_calls=539,\n", - " estimator_value=-4891.01730754809,\n", - " parameter_values={\n", - " 'm_{f_{0}(500)}': 0.6102707294731716,\n", - " 'm_{f_{0}(980)}': 0.9902119846618569,\n", - " 'm_{f_{0}(1370)}': 1.3456300421927978,\n", - " 'm_{f_{0}(1500)}': 1.5050299510041418,\n", - " 'm_{f_{0}(1710)}': 1.7096496843680975,\n", - " '\\\\Gamma_{f_{0}(500)}': 0.42260408077678696,\n", - " '\\\\Gamma_{f_{0}(980)}': 0.06479339507977673,\n", - " '\\\\Gamma_{f_{0}(1710)}': 0.13301019075895135,\n", - " 'C_{J/\\\\psi(1S) \\\\to {f_{0}(1500)}_{0} \\\\gamma_{+1}; f_{0}(1500) \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0}}': (1.069924901473717-0.018664035486070114j),\n", - " },\n", - " parameter_errors={\n", - " 'm_{f_{0}(500)}': 0.006168655451483166,\n", - " 'm_{f_{0}(980)}': 0.0016283609759060128,\n", - " 'm_{f_{0}(1370)}': 0.005122588414282541,\n", - " 'm_{f_{0}(1500)}': 0.0033157863009583644,\n", - " 'm_{f_{0}(1710)}': 0.0025660827200538303,\n", - " '\\\\Gamma_{f_{0}(500)}': 0.023838186345858253,\n", - " '\\\\Gamma_{f_{0}(980)}': 0.00355667300785808,\n", - " '\\\\Gamma_{f_{0}(1710)}': 0.007573518972833387,\n", - " 'C_{J/\\\\psi(1S) \\\\to {f_{0}(1500)}_{0} \\\\gamma_{+1}; f_{0}(1500) \\\\to \\\\pi^{0}_{0} \\\\pi^{0}_{0}}': (0.04106392765352627+0.07043808113241967j),\n", - " },\n", - ")" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "minuit2 = Minuit2()\n", - "fit_result_with_caching = minuit2.optimize(estimator_with_caching, initial_parameters)\n", - "fit_result_with_caching" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Fit result analysis" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "intensity_func.update_parameters(fit_result.parameter_values)\n", - "compare_model(\"m_12\", data, phsp, intensity_func)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/194347843-0ce5e251-a78a-4123-a9b3-09fdd632a524.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "fit_traceback = pd.read_csv(\"fit_traceback.csv\")\n", - "fig, (ax1, ax2) = plt.subplots(\n", - " 2, figsize=(7, 9), sharex=True, gridspec_kw={\"height_ratios\": [1, 2]}\n", - ")\n", - "fit_traceback.plot(\"function_call\", \"estimator_value\", ax=ax1)\n", - "fit_traceback.plot(\"function_call\", sorted(initial_parameters), ax=ax2)\n", - "fig.tight_layout()\n", - "ax2.set_xlabel(\"function call\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/194347685-7bd12eee-ea65-44e1-b6b5-3a615234040a.svg)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Zfit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### PDF definition" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import jax.numpy as jnp\n", - "import zfit # suppress tf warnings\n", - "import zfit.z.numpy as znp\n", - "from zfit import supports, z\n", - "\n", - "zfit.run.set_graph_mode(False) # We cannot (yet) compile through the function\n", - "zfit.run.set_autograd_mode(False)\n", - "\n", - "\n", - "class TensorWavesPDF(zfit.pdf.BasePDF):\n", - " def __init__(self, intensity, norm, obs, params=None, name=\"tensorwaves\"):\n", - " \"\"\"tensorwaves intensity normalized over the *norm* dataset.\"\"\"\n", - " super().__init__(obs, params, name)\n", - " self.intensity = intensity\n", - " norm = {ob: jnp.asarray(ar) for ob, ar in zip(self.obs, z.unstack_x(norm))}\n", - " self.norm_sample = norm\n", - "\n", - " @supports(norm=True)\n", - " def _pdf(self, x, norm):\n", - " # we can also use better mechanics, where it automatically normalizes or not\n", - " # this here is rather to take full control, it is always possible\n", - "\n", - " # updating the parameters of the model. This seems not very TF compatible?\n", - " self.intensity.update_parameters({\n", - " p.name: float(p) for p in self.params.values()\n", - " })\n", - "\n", - " # converting the data to a dict for tensorwaves\n", - " data = {ob: jnp.asarray(ar) for ob, ar in zip(self.obs, z.unstack_x(x))}\n", - "\n", - " non_normalized_pdf = self.intensity(data)\n", - " # this is not really needed, but can be useful for e.g. sampling with `pdf(..., norm_range=False)`\n", - " if norm is False:\n", - " out = non_normalized_pdf\n", - " else:\n", - " out = non_normalized_pdf / jnp.mean(self.intensity(self.norm_sample))\n", - " return znp.asarray(out)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "params = [\n", - " zfit.param.convert_to_parameter(val, name, prefer_constant=False)\n", - " for name, val in model.parameter_defaults.items()\n", - "]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def reset_parameters():\n", - " for p in params_fit:\n", - " if p.name in initial_parameters:\n", - " p.set_value(initial_parameters[p.name])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "obs = [\n", - " zfit.Space(ob, limits=(np.min(data[ob]) - 1, np.max(data[ob]) + 1))\n", - " for ob in pd.DataFrame(phsp)\n", - "]\n", - "obs_all = zfit.dimension.combine_spaces(*obs)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Data conversion" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "phsp_zfit = zfit.Data.from_pandas(pd.DataFrame(phsp), obs=obs_all)\n", - "data_zfit = zfit.Data.from_pandas(pd.DataFrame(data), obs=obs_all)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Perform fit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{obj}`complex` parameters need to be removed first:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "params_fit = [p for p in params if p.name in initial_parameters if p.independent]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "jax_intensity_func = create_parametrized_function(\n", - " expression=unfolded_expression,\n", - " parameters=model.parameter_defaults,\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "reset_parameters()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pdf = TensorWavesPDF(\n", - " obs=obs_all,\n", - " intensity=jax_intensity_func,\n", - " norm=phsp_zfit,\n", - " params={f\"{p.name}\": p for i, p in enumerate(params_fit)},\n", - ")\n", - "loss = zfit.loss.UnbinnedNLL(pdf, data_zfit)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "minimizer = zfit.minimize.Minuit(gradient=True, mode=0)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "You can also try different minimizers, like {class}`~zfit.minimizers.minimizers_scipy.ScipyTrustConstrV1`, but {class}`~zfit.minimizers.minimizer_minuit.Minuit` seems to perform best.\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 22 s, sys: 188 ms, total: 22.2 s\n", - "Wall time: 8.56 s\n" - ] - }, - { - "data": { - "text/plain": [ - "\u001b[1mFitResult\u001b[22m of\n", - " params=[\\Gamma_{f_{0}(1710)}, \\Gamma_{f_{0}(500)}, \\Gamma_{f_{0}(980)}, m_{f_{0}(1370)}, m_{f_{0}(1500)}, m_{f_{0}(1710)}, m_{f_{0}(500)}, m_{f_{0}(980)}]] data=[] constraints=[]> \n", - "with\n", - "\n", - "\n", - "╒═════════╤═════════════╤══════════════════╤═════════╤═════════════╕\n", - "│ valid │ converged │ param at limit │ edm │ min value │\n", - "╞═════════╪═════════════╪══════════════════╪═════════╪═════════════╡\n", - "│ True\u001b[0m │ True\u001b[0m │ False\u001b[0m │ 0.00041 │ -1871.035 │\n", - "╘═════════╧═════════════╧══════════════════╧═════════╧═════════════╛\n", - "\n", - "\u001b[1mParameters\n", - "\u001b[22mname value (rounded) at limit\n", - "-------------------- ------------------ ----------\n", - "m_{f_{0}(500)} 0.608864 False\u001b[0m\n", - "\\Gamma_{f_{0}(500)} 0.419716 False\u001b[0m\n", - "m_{f_{0}(980)} 0.990038 False\u001b[0m\n", - "\\Gamma_{f_{0}(980)} 0.0643328 False\u001b[0m\n", - "m_{f_{0}(1370)} 1.35137 False\u001b[0m\n", - "m_{f_{0}(1500)} 1.50627 False\u001b[0m\n", - "m_{f_{0}(1710)} 1.70956 False\u001b[0m\n", - "\\Gamma_{f_{0}(1710)} 0.132484 False\u001b[0m" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "%%time\n", - "\n", - "result = minimizer.minimize(loss)\n", - "result" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 2.5 s, sys: 12.6 ms, total: 2.51 s\n", - "Wall time: 953 ms\n" - ] - }, - { - "data": { - "text/plain": [ - "\u001b[1mFitResult\u001b[22m of\n", - " params=[\\Gamma_{f_{0}(1710)}, \\Gamma_{f_{0}(500)}, \\Gamma_{f_{0}(980)}, m_{f_{0}(1370)}, m_{f_{0}(1500)}, m_{f_{0}(1710)}, m_{f_{0}(500)}, m_{f_{0}(980)}]] data=[] constraints=[]> \n", - "with\n", - "\n", - "\n", - "╒═════════╤═════════════╤══════════════════╤═════════╤═════════════╕\n", - "│ valid │ converged │ param at limit │ edm │ min value │\n", - "╞═════════╪═════════════╪══════════════════╪═════════╪═════════════╡\n", - "│ True\u001b[0m │ True\u001b[0m │ False\u001b[0m │ 0.00041 │ -1871.035 │\n", - "╘═════════╧═════════════╧══════════════════╧═════════╧═════════════╛\n", - "\n", - "\u001b[1mParameters\n", - "\u001b[22mname value (rounded) hesse at limit\n", - "-------------------- ------------------ ----------- ----------\n", - "m_{f_{0}(500)} 0.608864 +/- 0.0061 False\u001b[0m\n", - "\\Gamma_{f_{0}(500)} 0.419716 +/- 0.024 False\u001b[0m\n", - "m_{f_{0}(980)} 0.990038 +/- 0.0016 False\u001b[0m\n", - "\\Gamma_{f_{0}(980)} 0.0643328 +/- 0.0035 False\u001b[0m\n", - "m_{f_{0}(1370)} 1.35137 +/- 0.0039 False\u001b[0m\n", - "m_{f_{0}(1500)} 1.50627 +/- 0.002 False\u001b[0m\n", - "m_{f_{0}(1710)} 1.70956 +/- 0.0023 False\u001b[0m\n", - "\\Gamma_{f_{0}(1710)} 0.132484 +/- 0.007 False\u001b[0m" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "%%time\n", - "\n", - "result.hesse(name=\"hesse\")\n", - "result" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "CPU times: user 45.3 s, sys: 393 ms, total: 45.7 s\n", - "Wall time: 17.2 s\n" - ] - }, - { - "data": { - "text/plain": [ - "\u001b[1mFitResult\u001b[22m of\n", - " params=[\\Gamma_{f_{0}(1710)}, \\Gamma_{f_{0}(500)}, \\Gamma_{f_{0}(980)}, m_{f_{0}(1370)}, m_{f_{0}(1500)}, m_{f_{0}(1710)}, m_{f_{0}(500)}, m_{f_{0}(980)}]] data=[] constraints=[]> \n", - "with\n", - "\n", - "\n", - "╒═════════╤═════════════╤══════════════════╤═════════╤═════════════╕\n", - "│ valid │ converged │ param at limit │ edm │ min value │\n", - "╞═════════╪═════════════╪══════════════════╪═════════╪═════════════╡\n", - "│ True\u001b[0m │ True\u001b[0m │ False\u001b[0m │ 0.00041 │ -1871.035 │\n", - "╘═════════╧═════════════╧══════════════════╧═════════╧═════════════╛\n", - "\n", - "\u001b[1mParameters\n", - "\u001b[22mname value (rounded) hesse errors at limit\n", - "-------------------- ------------------ ----------- ------------------- ----------\n", - "m_{f_{0}(500)} 0.608864 +/- 0.0061 - 0.006 + 0.0063 False\u001b[0m\n", - "\\Gamma_{f_{0}(500)} 0.419716 +/- 0.024 - 0.024 + 0.023 False\u001b[0m\n", - "m_{f_{0}(980)} 0.990038 +/- 0.0016 - 0.0016 + 0.0016 False\u001b[0m\n", - "\\Gamma_{f_{0}(980)} 0.0643328 +/- 0.0035 - 0.0034 + 0.0036 False\u001b[0m\n", - "m_{f_{0}(1370)} 1.35137 +/- 0.0039 - 0.0039 + 0.0039 False\u001b[0m\n", - "m_{f_{0}(1500)} 1.50627 +/- 0.002 - 0.002 + 0.002 False\u001b[0m\n", - "m_{f_{0}(1710)} 1.70956 +/- 0.0023 - 0.0024 + 0.0024 False\u001b[0m\n", - "\\Gamma_{f_{0}(1710)} 0.132484 +/- 0.007 - 0.0068 + 0.0073 False\u001b[0m" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "%%time\n", - "\n", - "result.errors(name=\"errors\")\n", - "result" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Statistical inference using the hepstats library\n", - "\n", - "{mod}`hepstats` is built on top of [`zfit-interface`](https://zfit-interface.readthedocs.io):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from hepstats.hypotests import ConfidenceInterval\n", - "from hepstats.hypotests.calculators import AsymptoticCalculator\n", - "from hepstats.hypotests.parameters import POIarray" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "calculator = AsymptoticCalculator(result, minimizer)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We take one of the parameters as POI:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "poi = pdf.params[r\"\\Gamma_{f_{0}(500)}\"]\n", - "poi" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "\n", - "Confidence interval on \\Gamma_{f_{0}(500)}:\n", - "\t0.3964206394323228 < \\Gamma_{f_{0}(500)} < 0.44257337109434974 at 67.2% C.L.\n" - ] - } - ], - "source": [ - "poi_null = POIarray(poi, np.linspace(poi - 0.1, poi + 0.1, 50))\n", - "ci = ConfidenceInterval(calculator, poi_null)\n", - "alpha = 0.328\n", - "ci.interval(alpha=alpha);" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A helper function to plot the result:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def one_minus_cl_plot(x, pvalues, alpha=None, ax=None):\n", - " if alpha is None:\n", - " alpha = [0.32]\n", - " if isinstance(alpha, (float, int)):\n", - " alpha = [alpha]\n", - " if ax is None:\n", - " ax = plt.gca()\n", - "\n", - " ax.plot(x, pvalues, \".--\")\n", - " for a in alpha:\n", - " ax.axhline(a, color=\"red\", label=\"$\\\\alpha = \" + str(a) + \"$\")\n", - " ax.set_ylabel(\"1-CL\")\n", - "\n", - " return ax" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-output" - ] - }, - "outputs": [], - "source": [ - "plt.figure(figsize=(9, 8))\n", - "one_minus_cl_plot(poi_null.values, ci.pvalues(), alpha=alpha)\n", - "plt.xlabel(f\"${poi.name}$\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://user-images.githubusercontent.com/29308176/194348046-c0b9026e-d4e9-434a-830d-351e6ba6e635.svg)" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/021.ipynb b/docs/report/021.ipynb deleted file mode 100644 index e1c94407..00000000 --- a/docs/report/021.ipynb +++ /dev/null @@ -1,2487 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "polarimetry", - "polarization" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Polarimeter vector field\n", - "TR-021\n", - "^^^\n", - "_Mikhail Mikhasenko [@mmikhasenko](https://github.com/mmikhasenko),\n", - "Remco de Boer [@redeboer](https://github.com/redeboer)_\n", - "\n", - "This report formulates the polarimeter vector field for in $\\Lambda_c \\to p\\pi K$ with [SymPy](https://docs.sympy.org) and visualizes it as an interactive widget with [TensorWaves](https://tensorwaves.rtfd.io) and [ipywidgets](https://ipywidgets.readthedocs.io).\n", - "+++\n", - "✅ [compwa.github.io#129](https://github.com/ComPWA/compwa.github.io/issues/129)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Polarimeter vector field" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.4 qrules==0.9.7 sympy==1.11.1 tensorwaves[jax]==0.4.8" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import itertools\n", - "import logging\n", - "import os\n", - "from typing import TYPE_CHECKING\n", - "\n", - "import jax.numpy as jnp\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.sympy import (\n", - " PoolSum,\n", - " UnevaluatedExpression,\n", - " create_expression,\n", - " implement_doit_method,\n", - " make_commutative,\n", - " perform_cached_doit,\n", - ")\n", - "from attrs import frozen\n", - "from IPython.display import HTML, Image, Math, display\n", - "from ipywidgets import Button, Combobox, HBox, HTMLMath, Tab, VBox, interactive_output\n", - "from matplotlib.colors import LogNorm\n", - "from symplot import create_slider\n", - "from sympy.core.symbol import Str\n", - "from sympy.physics.matrices import msigma\n", - "from sympy.physics.quantum.spin import Rotation as Wigner\n", - "from tensorwaves.data.transform import SympyDataTransformer\n", - "from tensorwaves.function.sympy import create_function, create_parametrized_function\n", - "\n", - "if TYPE_CHECKING:\n", - " from qrules.particle import Particle\n", - " from tensorwaves.interface import DataSample, ParametrizedFunction\n", - "\n", - "LOGGER = logging.getLogger()\n", - "LOGGER.setLevel(logging.ERROR)\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)\n", - "\n", - "PDG = qrules.load_pdg()\n", - "\n", - "\n", - "def display_definitions(definitions: dict[sp.Symbol, sp.Expr]) -> None:\n", - " latex = R\"\\begin{array}{rcl}\" + \"\\n\"\n", - " for symbol, expr in definitions.items():\n", - " symbol = sp.sympify(symbol)\n", - " expr = sp.sympify(expr)\n", - " lhs = sp.latex(symbol)\n", - " rhs = sp.latex(expr)\n", - " latex += Rf\" {lhs} & = & {rhs} \\\\\" + \"\\n\"\n", - " latex += R\"\\end{array}\"\n", - " display(Math(latex))\n", - "\n", - "\n", - "def display_doit(\n", - " expr: UnevaluatedExpression, deep=False, terms_per_line: int = 10\n", - ") -> None:\n", - " latex = sp.multiline_latex(\n", - " lhs=expr,\n", - " rhs=expr.doit(deep=deep),\n", - " terms_per_line=terms_per_line,\n", - " environment=\"eqnarray\",\n", - " )\n", - " display(Math(latex))\n", - "\n", - "\n", - "# hack for moving Indexed indices below superscript of the base\n", - "def _print_Indexed_latex(self, printer, *args):\n", - " base = printer._print(self.base)\n", - " indices = \", \".join(map(printer._print, self.indices))\n", - " return f\"{base}_{{{indices}}}\"\n", - "\n", - "\n", - "sp.Indexed._latex = _print_Indexed_latex" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Amplitude model\n", - "\n", - "The helicity amplitude for the $\\Lambda_c \\to p K \\pi$ decay reads as a sum of three partial wave series, incorporating $\\Lambda^{**}$ resonances, $\\Delta^{**}$ resonances, and $K^{**}$ resonances. The particles are ordered as $\\Lambda_c(\\mathbf{0}) \\to p(\\mathbf{1}) \\pi(\\mathbf{2}) K(\\mathbf{3})$.\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\mathcal{A}_{\\nu,\\lambda}(m_{K\\pi},m_{pK}) &=\n", - " \\sum_{\\nu',\\lambda'} \\big[\\\\\n", - " &\\qquad d_{\\nu,\\nu'}^{1/2}(\\zeta_{1(1)}^0) \\mathcal{A}_{\\nu',\\lambda'}^{K} d_{\\lambda',\\lambda}^{1/2}(\\zeta_{1(1)}^1)\n", - " &\\mathbf{subsystem\\,1\\;}(\\to 23)\\\\\n", - " &\\qquad + d_{\\nu,\\nu'}^{1/2}(\\zeta_{2(1)}^0) \\mathcal{A}_{\\nu',\\lambda'}^{\\Lambda} d_{\\lambda',\\lambda}^{1/2}(\\zeta_{2(1)}^1)\n", - " &\\mathbf{subsystem\\,2\\;}(\\to 31)\\\\\n", - " &\\qquad + d_{\\nu,\\nu'}^{1/2}(\\zeta_{3(1)}^0) \\mathcal{A}_{\\nu',\\lambda'}^{\\Delta} d_{\\lambda',\\lambda}^{1/2}(\\zeta_{3(1)}^1)\\big]\\,.\n", - " &\\mathbf{subsystem\\,3\\;}(\\to 12)\\\\\n", - "\\end{align}\n", - "$$ (aligned-amplitude)\n", - "\n", - "where $\\zeta^{i}_{j(k)}$ is the Wigner rotation for particle $i,k$ and chain $j$. The number in brackets indicates the overall definition of the helicity states $|1/2,\\nu\\rangle$ and $|1/2,\\lambda\\rangle$ for $\\Lambda_c$ and proton, respectively.\n", - "\n", - "We use the particle-2 convention for the helicity couplings, which leads to the phase factor in the transitions $\\Lambda_c\\to K^{**}p$, and $\\Lambda\\to K p$:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\mathcal{A}^{K}_{\\nu,\\lambda} &= \\sum_{j,\\tau} \\delta_{\\nu,\\tau - \\lambda}\\mathcal{H}^{\\Lambda_c \\to K^{**} p}_{\\tau,\\lambda} (-)^{1/2 - \\lambda} \\,d^{j}_{\\lambda,0} (\\theta_{23}) \\, \\mathcal{H}^{K^{**} \\to \\pi K}_{0,0}\\\\\n", - " %\n", - " \\mathcal{A}^{\\Lambda}_{\\nu,\\lambda} &= \\sum_{j,\\tau}\n", - " \\delta_{\\nu,\\tau}\n", - " \\mathcal{H}^{\\Lambda_c \\to \\Lambda^{**} \\pi}_{\\tau,0} d^{j}_{\\tau,-\\lambda} (\\theta_{31}) \\mathcal{H}^{\\Lambda^{**} \\to K p}_{0,\\lambda} (-)^{1/2-\\lambda} \\\\\n", - " %\n", - " \\mathcal{A}^{\\Delta}_{\\nu,\\lambda} &= \\sum_{j,\\tau}\n", - " \\delta_{\\nu,\\tau}\n", - " \\mathcal{H}^{\\Lambda_c \\to \\Delta^{**} K}_{\\tau,0} d^{j}_{\\tau,\\lambda}(\\theta_{12}) \\mathcal{H}^{\\Delta^{**} \\to p\\pi}_{\\lambda,0} \\,.\n", - "\\end{align}\n", - "$$\n", - "\n", - "The helicity couplings in the particle-2 convention obey simple properties with respect to the parity transformation:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\mathcal{H}^{A\\to BC}_{-\\lambda,-\\lambda'} = P_A P_B P_C (-)^{j_A-j_B-j_C} \\mathcal{H}^{A\\to BC}_{\\lambda,\\lambda'}\n", - "\\end{align}\n", - "$$\n", - "\n", - "It reduced amount of the couplings in the strong decay of isobars. Moreover the magnitude of the couplings cannot be determined separately, therefore, it is set to 1:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\mathcal{H}^{\\Lambda^{**} \\to K p}_{0,1/2} &= 1\\,, &\n", - " \\mathcal{H}^{\\Delta^{**} \\to p\\pi}_{1/2,0} &= 1\\,,&\n", - " \\mathcal{H}^{K^{**} \\to \\pi K}_{0,0} &= 1, \\\\\n", - " \\mathcal{H}^{\\Lambda^{**} \\to K p}_{0,-1/2} &= -P_\\Lambda (-)^{j-1/2}\\,, &\n", - " \\mathcal{H}^{\\Delta^{**} \\to p\\pi}_{-1/2,0} &= -P_\\Delta (-)^{j-1/2}\\,, &&\n", - "\\end{align}\n", - "$$\n", - "\n", - "The helicity couplings for the $\\Lambda_c^+$ decay are fit parameters. There are four of them for the $K^{**}$ chain, and two for both the $\\Delta^{**}$ and $\\Lambda^{**}$ chains." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Resonances and LS-scheme" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define decay product and sub-system IDs" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "Λc = PDG[\"Lambda(c)+\"]\n", - "p = PDG[\"p\"]\n", - "K = PDG[\"K-\"]\n", - "π = PDG[\"pi+\"]\n", - "decay_products = {\n", - " 1: (π, K),\n", - " 2: (p, K),\n", - " 3: (p, π),\n", - "}\n", - "siblings = {\n", - " 1: p,\n", - " 2: π,\n", - " 3: K,\n", - "}\n", - "chain_ids = {\n", - " 1: \"K\",\n", - " 2: \"L\",\n", - " 3: \"D\",\n", - "}\n", - "chain_labels = {\n", - " 1: \"K^{**}\",\n", - " 2: R\"\\Lambda^{**}\",\n", - " 3: R\"\\Delta^{**}\",\n", - "}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Resonance choices and their $LS$-couplings are defined as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Select resonances" - } - }, - "outputs": [], - "source": [ - "resonance_names = {\n", - " 1: [\"K*(892)0\"],\n", - " 2: [\"Lambda(1520)\", \"Lambda(1670)\"],\n", - " 3: [\"Delta(1232)++\"],\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Compute LS-couplings" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/html": [ - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "
resonance\\(j^P\\)\\(m\\) (MeV)\\(\\Gamma_0\\) (MeV)\\(l_R\\)\\(l_{\\Lambda_c}^\\mathrm{min}\\)
\\(K^{*}(892)^{0} \\to \\pi^{+} K^{-}\\)\\(1^-\\)8954710
\\(\\Lambda(1520) \\to p K^{-}\\)\\(3/2^-\\)15191621
\\(\\Lambda(1670) \\to p K^{-}\\)\\(1/2^-\\)16743000
\\(\\Delta(1232)^{++} \\to p \\pi^{+}\\)\\(3/2^+\\)123211711
\n" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "@frozen\n", - "class Resonance:\n", - " particle: Particle\n", - " l_R: int\n", - " l_Λc: int\n", - "\n", - " @staticmethod\n", - " def generate_ls(particle: Particle, chain_id: int) -> Resonance:\n", - " LS_prod = generate_ls(Λc, particle, siblings[chain_id], strong=False)\n", - " LS_prod = [L for L, S in LS_prod]\n", - " LS_dec = generate_ls(particle, *decay_products[chain_id])\n", - " LS_dec = [L for L, S in LS_dec]\n", - " return Resonance(particle, l_R=min(LS_dec), l_Λc=min(LS_prod))\n", - "\n", - "\n", - "def generate_ls(\n", - " parent: Particle,\n", - " child1: Particle,\n", - " child2: Particle,\n", - " strong: bool = True,\n", - " max_L: int = 3,\n", - "):\n", - " s1 = child1.spin\n", - " s2 = child2.spin\n", - " s_values = arange(abs(s1 - s2), s1 + s2)\n", - " LS_values = set()\n", - " for S in s_values:\n", - " for L in arange(0, max_L):\n", - " if not abs(L - S) <= parent.spin <= L + S:\n", - " continue\n", - " η0, η1, η2 = [\n", - " int(parent.parity),\n", - " int(child1.parity),\n", - " int(child2.parity),\n", - " ]\n", - " if strong and η0 != η1 * η2 * (-1) ** L:\n", - " continue\n", - " LS_values.add((L, S))\n", - " return sorted(LS_values)\n", - "\n", - "\n", - "def arange(x1, x2):\n", - " spin_range = np.arange(float(x1), +float(x2) + 0.5)\n", - " return list(map(sp.Rational, spin_range))\n", - "\n", - "\n", - "resonance_particles = {\n", - " chain_id: [PDG[name] for name in names]\n", - " for chain_id, names in resonance_names.items()\n", - "}\n", - "ls_resonances = {\n", - " chain_id: [Resonance.generate_ls(particle, chain_id) for particle in particles]\n", - " for chain_id, particles in resonance_particles.items()\n", - "}\n", - "\n", - "\n", - "def jp(particle: Particle):\n", - " p = \"+\" if particle.parity > 0 else \"-\"\n", - " j = sp.Rational(particle.spin)\n", - " return Rf\"\\({j}^{p}\\)\"\n", - "\n", - "\n", - "def create_html_table_row(*items, typ=\"td\"):\n", - " items = (f\"<{typ}>{i}\" for i in items)\n", - " return \"\" + \"\".join(items) + \"\\n\"\n", - "\n", - "\n", - "column_names = [\n", - " \"resonance\",\n", - " R\"\\(j^P\\)\",\n", - " R\"\\(m\\) (MeV)\",\n", - " R\"\\(\\Gamma_0\\) (MeV)\",\n", - " R\"\\(l_R\\)\",\n", - " R\"\\(l_{\\Lambda_c}^\\mathrm{min}\\)\",\n", - "]\n", - "src = \"\\n\"\n", - "src += create_html_table_row(*column_names, typ=\"th\")\n", - "for chain_id, resonance_list in ls_resonances.items():\n", - " child1, child2 = decay_products[chain_id]\n", - " for resonance in resonance_list:\n", - " src += create_html_table_row(\n", - " Rf\"\\({resonance.particle.latex} \\to {child1.latex} {child2.latex}\\)\",\n", - " jp(resonance.particle),\n", - " int(1e3 * resonance.particle.mass),\n", - " int(1e3 * resonance.particle.width),\n", - " resonance.l_R,\n", - " resonance.l_Λc,\n", - " )\n", - "src += \"
\\n\"\n", - "HTML(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Aligned amplitude" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "hide-input", - "scroll-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda^{\\prime}=-1/2}^{1/2} \\sum_{\\nu^{\\prime}=-1/2}^{1/2}{A^{K}_{\\nu^{\\prime}, \\lambda^{\\prime}} d^{\\frac{1}{2}}_{\\lambda^{\\prime},\\lambda}\\left(\\zeta^1_{1(1)}\\right) d^{\\frac{1}{2}}_{\\nu,\\nu^{\\prime}}\\left(\\zeta^0_{1(1)}\\right) + A^{\\Delta}_{\\nu^{\\prime}, \\lambda^{\\prime}} d^{\\frac{1}{2}}_{\\lambda^{\\prime},\\lambda}\\left(\\zeta^1_{3(1)}\\right) d^{\\frac{1}{2}}_{\\nu,\\nu^{\\prime}}\\left(\\zeta^0_{3(1)}\\right) + A^{\\Lambda}_{\\nu^{\\prime}, \\lambda^{\\prime}} d^{\\frac{1}{2}}_{\\lambda^{\\prime},\\lambda}\\left(\\zeta^1_{2(1)}\\right) d^{\\frac{1}{2}}_{\\nu,\\nu^{\\prime}}\\left(\\zeta^0_{2(1)}\\right)}$" - ], - "text/plain": [ - "PoolSum(A^K[\\nu^{\\prime}, \\lambda^{\\prime}]*WignerD(1/2, \\lambda^{\\prime}, lambda, 0, \\zeta^1_{1(1)}, 0)*WignerD(1/2, nu, \\nu^{\\prime}, 0, \\zeta^0_{1(1)}, 0) + A^{\\Delta}[\\nu^{\\prime}, \\lambda^{\\prime}]*WignerD(1/2, \\lambda^{\\prime}, lambda, 0, \\zeta^1_{3(1)}, 0)*WignerD(1/2, nu, \\nu^{\\prime}, 0, \\zeta^0_{3(1)}, 0) + A^{\\Lambda}[\\nu^{\\prime}, \\lambda^{\\prime}]*WignerD(1/2, \\lambda^{\\prime}, lambda, 0, \\zeta^1_{2(1)}, 0)*WignerD(1/2, nu, \\nu^{\\prime}, 0, \\zeta^0_{2(1)}, 0), (\\lambda^{\\prime}, (-1/2, 1/2)), (\\nu^{\\prime}, (-1/2, 1/2)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "A_K = sp.IndexedBase(R\"A^K\")\n", - "A_Λ = sp.IndexedBase(R\"A^{\\Lambda}\")\n", - "A_Δ = sp.IndexedBase(R\"A^{\\Delta}\")\n", - "\n", - "half = sp.S.Half\n", - "\n", - "ζ_0_11 = sp.Symbol(R\"\\zeta^0_{1(1)}\", real=True)\n", - "ζ_0_21 = sp.Symbol(R\"\\zeta^0_{2(1)}\", real=True)\n", - "ζ_0_31 = sp.Symbol(R\"\\zeta^0_{3(1)}\", real=True)\n", - "ζ_1_11 = sp.Symbol(R\"\\zeta^1_{1(1)}\", real=True)\n", - "ζ_1_21 = sp.Symbol(R\"\\zeta^1_{2(1)}\", real=True)\n", - "ζ_1_31 = sp.Symbol(R\"\\zeta^1_{3(1)}\", real=True)\n", - "\n", - "\n", - "def formulate_aligned_amplitude(λ_Λc, λ_p):\n", - " _ν = sp.Symbol(R\"\\nu^{\\prime}\", rational=True)\n", - " _λ = sp.Symbol(R\"\\lambda^{\\prime}\", rational=True)\n", - " return PoolSum(\n", - " A_K[_ν, _λ] * Wigner.d(half, λ_Λc, _ν, ζ_0_11) * Wigner.d(half, _λ, λ_p, ζ_1_11)\n", - " + A_Λ[_ν, _λ]\n", - " * Wigner.d(half, λ_Λc, _ν, ζ_0_21)\n", - " * Wigner.d(half, _λ, λ_p, ζ_1_21)\n", - " + A_Δ[_ν, _λ]\n", - " * Wigner.d(half, λ_Λc, _ν, ζ_0_31)\n", - " * Wigner.d(half, _λ, λ_p, ζ_1_31),\n", - " (_λ, [-half, +half]),\n", - " (_ν, [-half, +half]),\n", - " )\n", - "\n", - "\n", - "ν = sp.Symbol(\"nu\")\n", - "λ = sp.Symbol(\"lambda\")\n", - "formulate_aligned_amplitude(λ_Λc=ν, λ_p=λ)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Dynamics\n", - "\n", - "The lineshape function is factored out of the $\\Lambda_c$ helicity coupling:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\mathcal{H}^{\\Lambda_c \\to R x}_{\\lambda,\\lambda'} = \\hat{\\mathcal{H}}^{\\Lambda_c \\to R x}_{\\lambda,\\lambda'}\\,\\mathcal{R}(s)\\,.\n", - "\\end{align}\n", - "$$\n", - "\n", - "The relativistic Breit-Wigner parametrization reads:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\mathcal{R}(s) =\n", - " \\left(\\frac{q}{q_0}\\right)^{l_{\\Lambda_c}^\\text{min}} \\frac{F_{l_{\\Lambda_c}^\\text{min}}(q R)}{F_{l_{\\Lambda_c}^\\text{min}}(q_0 R)}\\,\n", - " \\frac{1}{m^2-s-im\\Gamma(s)}\n", - " \\left(\\frac{p}{p_0}\\right)^{l_R} \\frac{F_{l_R}(pR)}{F_{l_R}(p_0R)},\n", - "\\end{align}\n", - "$$\n", - "\n", - "with energy-dependent width given by\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\Gamma(s) = \\Gamma_0 \\left(\\frac{p}{p_0}\\right)^{2l_R+1} \\frac{m}{\\sqrt{s}} \\, \\frac{F_{l_R}^2(pR)}{F_{l_R}^2(p_0R)}\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "The form-factor $F$ is the Blatt-Weisskopf factor with the length factor $R=5\\,$GeV$^{-1}$:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " F_0(pR) &= 1\\,,&\n", - " F_1(pR) &= \\sqrt{\\frac{1}{1+(pR)^2}}\\,,&\n", - " F_2(pR) &= \\sqrt{\\frac{1}{9+3(pR)^2+(pR)^4}}\\,.\n", - "\\end{align}\n", - "$$\n", - "\n", - "The break-up momenta is calculated for every decay chain separately. Using the notations $0->R(\\to ij) k$, one writes:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " p &= \\lambda^{1/2}(s,m_i^2,m_j^2)/(2\\sqrt{s})\\,, & \n", - " q &= \\lambda^{1/2}(s,m_0^2,m_k^2)/(2m_0)\\,.\n", - "\\end{align}\n", - "$$\n", - "\n", - "The momenta with subindex zero are computed for nominal mass of the resonance, $s=m^2$. The three-argument Källén function reads:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\lambda(x,y,z) = x^2+y^2+z^2 - 2xy-2yz-2zx\\,.\n", - "\\end{align}\n", - "$$ (Kallen)\n", - "\n", - "```{rubric} Formulation with SymPy\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define Blatt-Weisskopf form factor" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{align*}\n", - "F_{L}\\left(z\\right) = & \\begin{cases} 1 & \\text{for}\\: L = 0 \\\\\\frac{1}{\\sqrt{z^{2} + 1}} & \\text{for}\\: L = 1 \\\\\\frac{1}{\\sqrt{z^{4} + 3 z^{2} + 9}} & \\text{for}\\: L = 2 \\end{cases} \n", - "\\end{align*}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "@make_commutative\n", - "@implement_doit_method\n", - "class BlattWeisskopf(UnevaluatedExpression):\n", - " def __new__(cls, z, L, **hints):\n", - " return create_expression(cls, z, L, **hints)\n", - "\n", - " def evaluate(self):\n", - " z, L = self.args\n", - " cases = {\n", - " 0: 1,\n", - " 1: 1 / (1 + z**2),\n", - " 2: 1 / (9 + 3 * z**2 + z**4),\n", - " }\n", - " return sp.Piecewise(*[\n", - " (sp.sqrt(expr), sp.Eq(L, l_val)) for l_val, expr in cases.items()\n", - " ])\n", - "\n", - " def _latex(self, printer, *args):\n", - " z, L = map(printer._print, self.args)\n", - " return Rf\"F_{{{L}}}\\left({z}\\right)\"\n", - "\n", - "\n", - "z = sp.Symbol(\"z\", positive=True)\n", - "L = sp.Symbol(\"L\", integer=True, nonnegative=True)\n", - "latex = sp.multiline_latex(BlattWeisskopf(z, L), BlattWeisskopf(z, L).doit())\n", - "Math(latex)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define Källén function" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\lambda\\left(x, y, z\\right) & = & x^{2} - 2 x y - 2 x z + y^{2} - 2 y z + z^{2} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "@make_commutative\n", - "@implement_doit_method\n", - "class Källén(UnevaluatedExpression):\n", - " def __new__(cls, x, y, z, **hints):\n", - " return create_expression(cls, x, y, z, **hints)\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " x, y, z = self.args\n", - " return x**2 + y**2 + z**2 - 2 * x * y - 2 * y * z - 2 * z * x\n", - "\n", - " def _latex(self, printer, *args):\n", - " x, y, z = map(printer._print, self.args)\n", - " return Rf\"\\lambda\\left({x}, {y}, {z}\\right)\"\n", - "\n", - "\n", - "x, y, z = sp.symbols(\"x:z\")\n", - "display_doit(Källén(x, y, z))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define break-up momenta p and q" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "p_{s} & = & \\frac{\\sqrt{\\lambda\\left(s, m_{i}^{2}, m_{j}^{2}\\right)}}{2 \\sqrt{s}} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "q_{s} & = & \\frac{\\sqrt{\\lambda\\left(s, m_{0}^{2}, m_{k}^{2}\\right)}}{2 m_{0}} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "@make_commutative\n", - "@implement_doit_method\n", - "class P(UnevaluatedExpression):\n", - " def __new__(cls, s, mi, mj, **hints):\n", - " return create_expression(cls, s, mi, mj, **hints)\n", - "\n", - " def evaluate(self):\n", - " s, mi, mj = self.args\n", - " return sp.sqrt(Källén(s, mi**2, mj**2)) / (2 * sp.sqrt(s))\n", - "\n", - " def _latex(self, printer, *args):\n", - " s = printer._print(self.args[0])\n", - " return Rf\"p_{{{s}}}\"\n", - "\n", - "\n", - "@make_commutative\n", - "@implement_doit_method\n", - "class Q(UnevaluatedExpression):\n", - " def __new__(cls, s, m0, mk, **hints):\n", - " return create_expression(cls, s, m0, mk, **hints)\n", - "\n", - " def evaluate(self):\n", - " s, m0, mk = self.args\n", - " return sp.sqrt(Källén(s, m0**2, mk**2)) / (2 * m0) # <-- not s!\n", - "\n", - " def _latex(self, printer, *args):\n", - " s = printer._print(self.args[0])\n", - " return Rf\"q_{{{s}}}\"\n", - "\n", - "\n", - "s, m0, mi, mj, mk = sp.symbols(\"s m0 m_i:k\", nonnegative=True)\n", - "display_doit(P(s, mi, mj))\n", - "display_doit(Q(s, m0, mk))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\Gamma\\left(s\\right) & = & Γ_{0} \\frac{m}{\\sqrt{s}} \\frac{F_{l_{R}}\\left(R p_{s}\\right)^{2}}{F_{l_{R}}\\left(R p_{m^{2}}\\right)^{2}} \\left(\\frac{p_{s}}{p_{m^{2}}}\\right)^{2 l_{R} + 1} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "R = sp.Symbol(\"R\")\n", - "parameter_defaults = {\n", - " R: 5, # GeV^{-1} (length factor)\n", - "}\n", - "\n", - "\n", - "@make_commutative\n", - "@implement_doit_method\n", - "class EnergyDependentWidth(UnevaluatedExpression):\n", - " def __new__(cls, s, m0, Γ0, m1, m2, L, R):\n", - " return create_expression(cls, s, m0, Γ0, m1, m2, L, R)\n", - "\n", - " def evaluate(self):\n", - " s, m0, Γ0, m1, m2, L, R = self.args\n", - " p = P(s, m1, m2)\n", - " p0 = P(m0**2, m1, m2)\n", - " ff = BlattWeisskopf(p * R, L) ** 2\n", - " ff0 = BlattWeisskopf(p0 * R, L) ** 2\n", - " return sp.Mul(\n", - " Γ0,\n", - " (p / p0) ** (2 * L + 1),\n", - " m0 / sp.sqrt(s),\n", - " ff / ff0,\n", - " evaluate=False,\n", - " )\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " s = printer._print(self.args[0])\n", - " return Rf\"\\Gamma\\left({s}\\right)\"\n", - "\n", - "\n", - "l_R = sp.Symbol(\"l_R\", integer=True, positive=True)\n", - "m, Γ0, m1, m2 = sp.symbols(\"m Γ0 m1 m2\", nonnegative=True)\n", - "display_doit(EnergyDependentWidth(s, m, Γ0, m1, m2, l_R, R))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define relativistic Breit-Wigner" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{eqnarray}\n", - "\\mathcal{R}\\left(s\\right) & = & \\frac{\\frac{F_{l_{R}}\\left(R p_{s}\\right)}{F_{l_{R}}\\left(R p_{m^{2}}\\right)} \\frac{F_{l_{\\Lambda_c}}\\left(R q_{s}\\right)}{F_{l_{\\Lambda_c}}\\left(R q_{m^{2}}\\right)} \\left(\\frac{p_{s}}{p_{m^{2}}}\\right)^{l_{R}} \\left(\\frac{q_{s}}{q_{m^{2}}}\\right)^{l_{\\Lambda_c}}}{m^{2} - i m \\Gamma\\left(s\\right) - s} \n", - "\\end{eqnarray}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "@make_commutative\n", - "@implement_doit_method\n", - "class RelativisticBreitWigner(UnevaluatedExpression):\n", - " def __new__(cls, s, m0, Γ0, m1, m2, l_R, l_Λc, R):\n", - " return create_expression(cls, s, m0, Γ0, m1, m2, l_R, l_Λc, R)\n", - "\n", - " def evaluate(self):\n", - " s, m0, Γ0, m1, m2, l_R, l_Λc, R = self.args\n", - " q = Q(s, m1, m2)\n", - " q0 = Q(m0**2, m1, m2)\n", - " p = P(s, m1, m2)\n", - " p0 = P(m0**2, m1, m2)\n", - " width = EnergyDependentWidth(s, m0, Γ0, m1, m2, l_R, R)\n", - " return sp.Mul(\n", - " (q / q0) ** l_Λc,\n", - " BlattWeisskopf(q * R, l_Λc) / BlattWeisskopf(q0 * R, l_Λc),\n", - " 1 / (m0**2 - s - sp.I * m0 * width),\n", - " (p / p0) ** l_R,\n", - " BlattWeisskopf(p * R, l_R) / BlattWeisskopf(p0 * R, l_R),\n", - " evaluate=False,\n", - " )\n", - "\n", - " def _latex(self, printer, *args) -> str:\n", - " s = printer._print(self.args[0])\n", - " return Rf\"\\mathcal{{R}}\\left({s}\\right)\"\n", - "\n", - "\n", - "l_Λc = sp.Symbol(R\"l_{\\Lambda_c}\", integer=True, positive=True)\n", - "display_doit(RelativisticBreitWigner(s, m, Γ0, m1, m2, l_R, l_Λc, R))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Decay chain amplitudes" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Formulate unaligned amplitudes for each chain" - }, - "tags": [ - "scroll-input", - "hide-input", - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\tau=-1}^{1}{\\left(-1\\right)^{\\frac{1}{2} - \\lambda} \\delta_{\\nu, - \\lambda + \\tau} \\mathcal{H}^\\mathrm{decay}_{K^{*}(892)^{0}, 0, 0} \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, \\tau, \\lambda} \\mathcal{R}\\left(\\sigma_{1}\\right) d^{1}_{\\tau,0}\\left(\\theta_{23}\\right)}$" - ], - "text/plain": [ - "PoolSum((-1)**(1/2 - lambda)*KroneckerDelta(nu, -lambda + tau)*\\mathcal{H}^\\mathrm{decay}[K^{*}(892)^{0}, 0, 0]*\\mathcal{H}^\\mathrm{production}[K^{*}(892)^{0}, tau, lambda]*RelativisticBreitWigner(sigma1, m_{K^{*}(892)^{0}}, \\Gamma_{K^{*}(892)^{0}}, m_pi, m_K, 1, 0, R)*WignerD(1, tau, 0, 0, theta23, 0), (tau, (-1, 0, 1)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\tau=-3/2}^{3/2}{\\left(-1\\right)^{\\frac{1}{2} - \\lambda} \\delta_{\\nu \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, \\lambda} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{3}{2}}_{\\tau,- \\lambda}\\left(\\theta_{31}\\right)} + \\sum_{\\tau=-1/2}^{1/2}{\\left(-1\\right)^{\\frac{1}{2} - \\lambda} \\delta_{\\nu \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, \\lambda} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{1}{2}}_{\\tau,- \\lambda}\\left(\\theta_{31}\\right)}$" - ], - "text/plain": [ - "PoolSum((-1)**(1/2 - lambda)*KroneckerDelta(nu, tau)*\\mathcal{H}^\\mathrm{decay}[\\Lambda(1520), 0, lambda]*\\mathcal{H}^\\mathrm{production}[\\Lambda(1520), tau, 0]*RelativisticBreitWigner(sigma2, m_{\\Lambda(1520)}, \\Gamma_{\\Lambda(1520)}, m_p, m_K, 2, 1, R)*WignerD(3/2, tau, -lambda, 0, theta31, 0), (tau, (-3/2, -1/2, 1/2, 3/2))) + PoolSum((-1)**(1/2 - lambda)*KroneckerDelta(nu, tau)*\\mathcal{H}^\\mathrm{decay}[\\Lambda(1670), 0, lambda]*\\mathcal{H}^\\mathrm{production}[\\Lambda(1670), tau, 0]*RelativisticBreitWigner(sigma2, m_{\\Lambda(1670)}, \\Gamma_{\\Lambda(1670)}, m_p, m_K, 0, 0, R)*WignerD(1/2, tau, -lambda, 0, theta31, 0), (tau, (-1/2, 1/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - }, - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\tau=-3/2}^{3/2}{\\delta_{\\nu \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, \\lambda, 0} \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, \\tau, 0} \\mathcal{R}\\left(\\sigma_{3}\\right) d^{\\frac{3}{2}}_{\\tau,\\lambda}\\left(\\theta_{12}\\right)}$" - ], - "text/plain": [ - "PoolSum(KroneckerDelta(nu, tau)*\\mathcal{H}^\\mathrm{decay}[\\Delta(1232)^{++}, lambda, 0]*\\mathcal{H}^\\mathrm{production}[\\Delta(1232)^{++}, tau, 0]*RelativisticBreitWigner(sigma3, m_{\\Delta(1232)^{++}}, \\Gamma_{\\Delta(1232)^{++}}, m_p, m_pi, 1, 1, R)*WignerD(3/2, tau, lambda, 0, theta12, 0), (tau, (-3/2, -1/2, 1/2, 3/2)))" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "def formulate_chain_amplitude(chain_id: int, λ_Λc, λ_p):\n", - " resonances = ls_resonances[chain_id]\n", - " if chain_id == 1:\n", - " return formulate_K_amplitude(λ_Λc, λ_p, resonances)\n", - " if chain_id == 2:\n", - " return formulate_Λ_amplitude(λ_Λc, λ_p, resonances)\n", - " if chain_id == 3:\n", - " return formulate_Δ_amplitude(λ_Λc, λ_p, resonances)\n", - " raise NotImplementedError\n", - "\n", - "\n", - "H_prod = sp.IndexedBase(R\"\\mathcal{H}^\\mathrm{production}\")\n", - "H_dec = sp.IndexedBase(R\"\\mathcal{H}^\\mathrm{decay}\")\n", - "\n", - "θ23 = sp.Symbol(\"theta23\", real=True)\n", - "θ31 = sp.Symbol(\"theta31\", real=True)\n", - "θ12 = sp.Symbol(\"theta12\", real=True)\n", - "\n", - "σ1, σ2, σ3 = sp.symbols(\"sigma1:4\", nonnegative=True)\n", - "m1, m2, m3 = sp.symbols(R\"m_p m_pi m_K\", nonnegative=True)\n", - "\n", - "\n", - "def formulate_K_amplitude(λ_Λc, λ_p, resonances: list[Resonance]):\n", - " τ = sp.Symbol(\"tau\", rational=True)\n", - " return sp.Add(*[\n", - " PoolSum(\n", - " sp.KroneckerDelta(λ_Λc, τ - λ_p)\n", - " * H_prod[stringify(res), τ, λ_p]\n", - " * formulate_dynamics(res, σ1, m2, m3)\n", - " * (-1) ** (half - λ_p)\n", - " * Wigner.d(sp.Rational(res.particle.spin), τ, 0, θ23)\n", - " * H_dec[stringify(res), 0, 0],\n", - " (τ, create_spin_range(res.particle.spin)),\n", - " )\n", - " for res in resonances\n", - " ])\n", - "\n", - "\n", - "def formulate_Λ_amplitude(λ_Λc, λ_p, resonances: list[Resonance]):\n", - " τ = sp.Symbol(\"tau\", rational=True)\n", - " return sp.Add(*[\n", - " PoolSum(\n", - " sp.KroneckerDelta(λ_Λc, τ)\n", - " * H_prod[stringify(res), τ, 0]\n", - " * formulate_dynamics(res, σ2, m1, m3)\n", - " * Wigner.d(sp.Rational(res.particle.spin), τ, -λ_p, θ31)\n", - " * H_dec[stringify(res), 0, λ_p]\n", - " * (-1) ** (half - λ_p),\n", - " (τ, create_spin_range(res.particle.spin)),\n", - " )\n", - " for res in resonances\n", - " ])\n", - "\n", - "\n", - "def formulate_Δ_amplitude(λ_Λc, λ_p, resonances: list[Resonance]):\n", - " τ = sp.Symbol(\"tau\", rational=True)\n", - " return sp.Add(*[\n", - " PoolSum(\n", - " sp.KroneckerDelta(λ_Λc, τ)\n", - " * H_prod[stringify(res), τ, 0]\n", - " * formulate_dynamics(res, σ3, m1, m2)\n", - " * Wigner.d(sp.Rational(res.particle.spin), τ, λ_p, θ12)\n", - " * H_dec[stringify(res), λ_p, 0],\n", - " (τ, create_spin_range(res.particle.spin)),\n", - " )\n", - " for res in resonances\n", - " ])\n", - "\n", - "\n", - "def formulate_dynamics(decay: Resonance, s, m1, m2):\n", - " l_R = sp.Rational(decay.l_R)\n", - " l_Λc = sp.Rational(decay.l_Λc)\n", - " mass = sp.Symbol(f\"m_{{{decay.particle.latex}}}\")\n", - " width = sp.Symbol(Rf\"\\Gamma_{{{decay.particle.latex}}}\")\n", - " parameter_defaults[mass] = decay.particle.mass\n", - " parameter_defaults[width] = decay.particle.width\n", - " return RelativisticBreitWigner(s, mass, width, m1, m2, l_R, l_Λc, R)\n", - "\n", - "\n", - "def stringify(particle: Particle | Resonance) -> Str:\n", - " if isinstance(particle, Resonance):\n", - " particle = particle.particle\n", - " return Str(particle.latex)\n", - "\n", - "\n", - "def create_spin_range(j):\n", - " return arange(-j, +j)\n", - "\n", - "\n", - "display(\n", - " formulate_chain_amplitude(1, ν, λ),\n", - " formulate_chain_amplitude(2, ν, λ),\n", - " formulate_chain_amplitude(3, ν, λ),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Angle definitions\n", - "\n", - "Angles with repeated lower indices are trivial. The other angles are computed from the invariants and these are the angles with sign\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\zeta_{1(1)}^{0} &= \\hat{\\theta}_{1(1)}^{0} = 0\\,, & \\zeta_{1(1)}^{1} &= 0\\,,\\\\\n", - " \\zeta_{2(1)}^0 &=\\hat{\\theta}_{2(1)} =-\\hat{\\theta}_{1(2)}\\,,\\\\\n", - " \\zeta_{3(1)}^0 &= \\hat{\\theta}_{3(1)}\\,, &\\zeta_{3(1)}^1 &= -\\zeta_{1(3)}^1\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "The expressions for the cosine of the positive (anticlockwise) angles, $\\theta_{12}, \\theta_{23}, \\theta_{13}$ and $\\hat\\theta_{1(2)}, \\hat\\theta_{3(1)}, \\zeta^1_{1(3)}$ can be expressed in terms of Mandelstam variables $\\sigma_1, \\sigma_2, \\sigma_3$ using {cite}`JPAC:2019ufm`, Appendix A:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Formulate expressions for DPD-angles" - }, - "tags": [ - "scroll-input", - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " \\theta_{12} & = & \\operatorname{acos}{\\left(\\frac{2 \\sigma_{3} \\left(- m_{K}^{2} - m_{p}^{2} + \\sigma_{2}\\right) - \\left(- m_{K}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{3}\\right) \\left(m_{p}^{2} - m_{\\pi}^{2} + \\sigma_{3}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{K}^{2}, \\sigma_{3}\\right)} \\sqrt{\\lambda\\left(\\sigma_{3}, m_{p}^{2}, m_{\\pi}^{2}\\right)}} \\right)} \\\\\n", - " \\theta_{23} & = & \\operatorname{acos}{\\left(\\frac{2 \\sigma_{1} \\left(- m_{p}^{2} - m_{\\pi}^{2} + \\sigma_{3}\\right) - \\left(- m_{K}^{2} + m_{\\pi}^{2} + \\sigma_{1}\\right) \\left(- m_{p}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{1}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{p}^{2}, \\sigma_{1}\\right)} \\sqrt{\\lambda\\left(\\sigma_{1}, m_{\\pi}^{2}, m_{K}^{2}\\right)}} \\right)} \\\\\n", - " \\theta_{31} & = & \\operatorname{acos}{\\left(\\frac{2 \\sigma_{2} \\left(- m_{K}^{2} - m_{\\pi}^{2} + \\sigma_{1}\\right) - \\left(m_{K}^{2} - m_{p}^{2} + \\sigma_{2}\\right) \\left(- m_{\\pi}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{2}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{\\pi}^{2}, \\sigma_{2}\\right)} \\sqrt{\\lambda\\left(\\sigma_{2}, m_{K}^{2}, m_{p}^{2}\\right)}} \\right)} \\\\\n", - " \\zeta^0_{1(1)} & = & 0 \\\\\n", - " \\zeta^0_{2(1)} & = & - \\operatorname{acos}{\\left(\\frac{- 2 m_{\\Lambda_c}^{2} \\left(- m_{p}^{2} - m_{\\pi}^{2} + \\sigma_{3}\\right) + \\left(m_{p}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{1}\\right) \\left(m_{\\pi}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{2}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{\\pi}^{2}, \\sigma_{2}\\right)} \\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, \\sigma_{1}, m_{p}^{2}\\right)}} \\right)} \\\\\n", - " \\zeta^0_{3(1)} & = & \\operatorname{acos}{\\left(\\frac{- 2 m_{\\Lambda_c}^{2} \\left(- m_{K}^{2} - m_{p}^{2} + \\sigma_{2}\\right) + \\left(m_{K}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{3}\\right) \\left(m_{p}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{1}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{p}^{2}, \\sigma_{1}\\right)} \\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, \\sigma_{3}, m_{K}^{2}\\right)}} \\right)} \\\\\n", - " \\zeta^1_{1(1)} & = & 0 \\\\\n", - " \\zeta^1_{2(1)} & = & \\operatorname{acos}{\\left(\\frac{2 m_{p}^{2} \\left(- m_{K}^{2} - m_{\\Lambda_c}^{2} + \\sigma_{3}\\right) + \\left(- m_{K}^{2} - m_{p}^{2} + \\sigma_{2}\\right) \\left(m_{p}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{1}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{p}^{2}, \\sigma_{1}\\right)} \\sqrt{\\lambda\\left(\\sigma_{2}, m_{p}^{2}, m_{K}^{2}\\right)}} \\right)} \\\\\n", - " \\zeta^1_{3(1)} & = & - \\operatorname{acos}{\\left(\\frac{2 m_{p}^{2} \\left(- m_{\\pi}^{2} - m_{\\Lambda_c}^{2} + \\sigma_{2}\\right) + \\left(- m_{p}^{2} - m_{\\pi}^{2} + \\sigma_{3}\\right) \\left(m_{p}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{1}\\right)}{\\sqrt{\\lambda\\left(m_{\\Lambda_c}^{2}, m_{p}^{2}, \\sigma_{1}\\right)} \\sqrt{\\lambda\\left(\\sigma_{3}, m_{p}^{2}, m_{\\pi}^{2}\\right)}} \\right)} \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "m0 = sp.Symbol(R\"m_{\\Lambda_c}\", nonnegative=True)\n", - "angles = {\n", - " θ12: sp.acos(\n", - " (2 * σ3 * (σ2 - m3**2 - m1**2) - (σ3 + m1**2 - m2**2) * (m0**2 - σ3 - m3**2))\n", - " / (sp.sqrt(Källén(m0**2, m3**2, σ3)) * sp.sqrt(Källén(σ3, m1**2, m2**2)))\n", - " ),\n", - " θ23: sp.acos(\n", - " (2 * σ1 * (σ3 - m1**2 - m2**2) - (σ1 + m2**2 - m3**2) * (m0**2 - σ1 - m1**2))\n", - " / (sp.sqrt(Källén(m0**2, m1**2, σ1)) * sp.sqrt(Källén(σ1, m2**2, m3**2)))\n", - " ),\n", - " θ31: sp.acos(\n", - " (2 * σ2 * (σ1 - m2**2 - m3**2) - (σ2 + m3**2 - m1**2) * (m0**2 - σ2 - m2**2))\n", - " / (sp.sqrt(Källén(m0**2, m2**2, σ2)) * sp.sqrt(Källén(σ2, m3**2, m1**2)))\n", - " ),\n", - " ζ_0_11: sp.S.Zero, # = \\hat\\theta^0_{1(1)}\n", - " ζ_0_21: -sp.acos( # = -\\hat\\theta^{1(2)}\n", - " ((m0**2 + m1**2 - σ1) * (m0**2 + m2**2 - σ2) - 2 * m0**2 * (σ3 - m1**2 - m2**2))\n", - " / (sp.sqrt(Källén(m0**2, m2**2, σ2)) * sp.sqrt(Källén(m0**2, σ1, m1**2)))\n", - " ),\n", - " ζ_0_31: sp.acos( # = \\hat\\theta^{3(1)}\n", - " ((m0**2 + m3**2 - σ3) * (m0**2 + m1**2 - σ1) - 2 * m0**2 * (σ2 - m3**2 - m1**2))\n", - " / (sp.sqrt(Källén(m0**2, m1**2, σ1)) * sp.sqrt(Källén(m0**2, σ3, m3**2)))\n", - " ),\n", - " ζ_1_11: sp.S.Zero,\n", - " ζ_1_21: sp.acos(\n", - " (2 * m1**2 * (σ3 - m0**2 - m3**2) + (m0**2 + m1**2 - σ1) * (σ2 - m1**2 - m3**2))\n", - " / (sp.sqrt(Källén(m0**2, m1**2, σ1)) * sp.sqrt(Källén(σ2, m1**2, m3**2)))\n", - " ),\n", - " ζ_1_31: -sp.acos( # = -\\zeta^1_{1(3)}\n", - " (2 * m1**2 * (σ2 - m0**2 - m2**2) + (m0**2 + m1**2 - σ1) * (σ3 - m1**2 - m2**2))\n", - " / (sp.sqrt(Källén(m0**2, m1**2, σ1)) * sp.sqrt(Källén(σ3, m1**2, m2**2)))\n", - " ),\n", - "}\n", - "\n", - "display_definitions(angles)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "where $m_0$ is the mass of the initial state $\\Lambda_c$ and $m_1, m_2, m_3$ are the masses of $p, \\pi, K$, respectively:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define initial and final state masses" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " m_{\\Lambda_c} & = & 2.28646 \\\\\n", - " m_{p} & = & 0.938272081 \\\\\n", - " m_{\\pi} & = & 0.13957039 \\\\\n", - " m_{K} & = & 0.493677 \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "masses = {\n", - " m0: Λc.mass,\n", - " m1: p.mass,\n", - " m2: π.mass,\n", - " m3: K.mass,\n", - "}\n", - "parameter_defaults.update(masses)\n", - "display_definitions(masses)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Helicity coupling values" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Compute helicity couplings of the decay node" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " \\mathcal{H}^\\mathrm{decay}_{K^{*}(892)^{0}, 0, 0} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, \\frac{1}{2}} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, - \\frac{1}{2}} & = & -1 \\\\\n", - " \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, \\frac{1}{2}} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, - \\frac{1}{2}} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, \\frac{1}{2}, 0} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, - \\frac{1}{2}, 0} & = & 1 \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "dec_couplings = {}\n", - "for res in ls_resonances[1]:\n", - " i = stringify(res)\n", - " dec_couplings[H_dec[i, 0, 0]] = 1\n", - "for res in ls_resonances[2]:\n", - " i = stringify(res.particle)\n", - " dec_couplings[H_dec[i, 0, half]] = 1\n", - " dec_couplings[H_dec[i, 0, -half]] = int(\n", - " int(res.particle.parity)\n", - " * int(K.parity)\n", - " * int(p.parity)\n", - " * (-1) ** (res.particle.spin - K.spin - p.spin)\n", - " )\n", - "for res in ls_resonances[3]:\n", - " i = stringify(res.particle)\n", - " dec_couplings[H_dec[i, half, 0]] = 1\n", - " dec_couplings[H_dec[i, -half, 0]] = int(\n", - " int(res.particle.parity)\n", - " * int(p.parity)\n", - " * int(π.parity)\n", - " * (-1) ** (res.particle.spin - p.spin - π.spin)\n", - " )\n", - "parameter_defaults.update(dec_couplings)\n", - "display_definitions(dec_couplings)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define helicity couplings of the production node" - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, 0, - \\frac{1}{2}} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, -1, - \\frac{1}{2}} & = & 1.0 - 1.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, 1, \\frac{1}{2}} & = & -3.0 - 3.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, 0, \\frac{1}{2}} & = & -1.0 - 4.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{0}^{*}(1430)^{0}, 0, - \\frac{1}{2}} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{0}^{*}(1430)^{0}, -1, - \\frac{1}{2}} & = & 1.0 - 1.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{0}^{*}(1430)^{0}, 1, \\frac{1}{2}} & = & -3.0 - 3.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{0}^{*}(1430)^{0}, 0, \\frac{1}{2}} & = & -1.0 - 4.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{2}^{*}(1430)^{0}, 0, - \\frac{1}{2}} & = & 1 \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{2}^{*}(1430)^{0}, -1, - \\frac{1}{2}} & = & 1.0 - 1.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{2}^{*}(1430)^{0}, 1, \\frac{1}{2}} & = & -3.0 - 3.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{K_{2}^{*}(1430)^{0}, 0, \\frac{1}{2}} & = & -1.0 - 4.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), \\frac{1}{2}, 0} & = & 1.5 \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), - \\frac{1}{2}, 0} & = & 0.3 \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), \\frac{1}{2}, 0} & = & -0.5 + 1.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), - \\frac{1}{2}, 0} & = & -0.3 - 0.1 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, \\frac{1}{2}, 0} & = & -13.0 + 5.0 i \\\\\n", - " \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, - \\frac{1}{2}, 0} & = & -7.0 + 3.0 i \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "prod_couplings = {\n", - " # chain 23:\n", - " H_prod[Str(\"K^{*}(892)^{0}\"), 0, -half]: 1,\n", - " H_prod[Str(\"K^{*}(892)^{0}\"), -1, -half]: 1 - 1j,\n", - " H_prod[Str(\"K^{*}(892)^{0}\"), +1, +half]: -3 - 3j,\n", - " H_prod[Str(\"K^{*}(892)^{0}\"), 0, +half]: -1 - 4j,\n", - " H_prod[Str(\"K_{0}^{*}(1430)^{0}\"), 0, -half]: 1,\n", - " H_prod[Str(\"K_{0}^{*}(1430)^{0}\"), -1, -half]: 1 - 1j,\n", - " H_prod[Str(\"K_{0}^{*}(1430)^{0}\"), +1, +half]: -3 - 3j,\n", - " H_prod[Str(\"K_{0}^{*}(1430)^{0}\"), 0, +half]: -1 - 4j,\n", - " H_prod[Str(\"K_{2}^{*}(1430)^{0}\"), 0, -half]: 1,\n", - " H_prod[Str(\"K_{2}^{*}(1430)^{0}\"), -1, -half]: 1 - 1j,\n", - " H_prod[Str(\"K_{2}^{*}(1430)^{0}\"), +1, +half]: -3 - 3j,\n", - " H_prod[Str(\"K_{2}^{*}(1430)^{0}\"), 0, +half]: -1 - 4j,\n", - " #\n", - " # chain 31:\n", - " H_prod[Str(R\"\\Lambda(1520)\"), +half, 0]: 1.5,\n", - " H_prod[Str(R\"\\Lambda(1520)\"), -half, 0]: 0.3,\n", - " H_prod[Str(R\"\\Lambda(1670)\"), +half, 0]: -0.5 + 1j,\n", - " H_prod[Str(R\"\\Lambda(1670)\"), -half, 0]: -0.3 - 0.1j,\n", - " # chain 12:\n", - " H_prod[Str(R\"\\Delta(1232)^{++}\"), +half, 0]: -13 + 5j,\n", - " H_prod[Str(R\"\\Delta(1232)^{++}\"), -half, 0]: -7 + 3j,\n", - "}\n", - "display_definitions(prod_couplings)\n", - "couplings = dict(dec_couplings)\n", - "couplings.update(prod_couplings)\n", - "parameter_defaults.update(prod_couplings)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Intensity expression\n", - "\n", - "Incoherent sum of the amplitudes defined by {ref}`report/021:Aligned amplitude`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Formulate expression for the DPD-aligned intensity" - }, - "tags": [ - "hide-input", - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\sum_{\\lambda=-1/2}^{1/2} \\sum_{\\nu=-1/2}^{1/2}{\\left|{\\sum_{\\lambda^{\\prime}=-1/2}^{1/2} \\sum_{\\nu^{\\prime}=-1/2}^{1/2}{A^{K}_{\\nu^{\\prime}, \\lambda^{\\prime}} d^{\\frac{1}{2}}_{\\lambda^{\\prime},\\lambda}\\left(\\zeta^1_{1(1)}\\right) d^{\\frac{1}{2}}_{\\nu,\\nu^{\\prime}}\\left(\\zeta^0_{1(1)}\\right) + A^{\\Delta}_{\\nu^{\\prime}, \\lambda^{\\prime}} d^{\\frac{1}{2}}_{\\lambda^{\\prime},\\lambda}\\left(\\zeta^1_{3(1)}\\right) d^{\\frac{1}{2}}_{\\nu,\\nu^{\\prime}}\\left(\\zeta^0_{3(1)}\\right) + A^{\\Lambda}_{\\nu^{\\prime}, \\lambda^{\\prime}} d^{\\frac{1}{2}}_{\\lambda^{\\prime},\\lambda}\\left(\\zeta^1_{2(1)}\\right) d^{\\frac{1}{2}}_{\\nu,\\nu^{\\prime}}\\left(\\zeta^0_{2(1)}\\right)}}\\right|^{2}}$" - ], - "text/plain": [ - "PoolSum(Abs(PoolSum(A^K[\\nu^{\\prime}, \\lambda^{\\prime}]*WignerD(1/2, \\lambda^{\\prime}, lambda, 0, \\zeta^1_{1(1)}, 0)*WignerD(1/2, nu, \\nu^{\\prime}, 0, \\zeta^0_{1(1)}, 0) + A^{\\Delta}[\\nu^{\\prime}, \\lambda^{\\prime}]*WignerD(1/2, \\lambda^{\\prime}, lambda, 0, \\zeta^1_{3(1)}, 0)*WignerD(1/2, nu, \\nu^{\\prime}, 0, \\zeta^0_{3(1)}, 0) + A^{\\Lambda}[\\nu^{\\prime}, \\lambda^{\\prime}]*WignerD(1/2, \\lambda^{\\prime}, lambda, 0, \\zeta^1_{2(1)}, 0)*WignerD(1/2, nu, \\nu^{\\prime}, 0, \\zeta^0_{2(1)}, 0), (\\lambda^{\\prime}, (-1/2, 1/2)), (\\nu^{\\prime}, (-1/2, 1/2))))**2, (lambda, (-1/2, 1/2)), (nu, (-1/2, 1/2)))" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "intensity_expr = PoolSum(\n", - " sp.Abs(formulate_aligned_amplitude(ν, λ)) ** 2,\n", - " (λ, [-half, +half]),\n", - " (ν, [-half, +half]),\n", - ")\n", - "intensity_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Remaining {attr}`~sympy.core.basic.Basic.free_symbols` are indeed the specific amplitudes as defined by {ref}`report/021:Decay chain amplitudes`:\n", - "\n", - "The specific amplitudes from {ref}`report/021:Decay chain amplitudes` need to be formulated for each value of $\\nu, \\lambda$, so that they can be substituted in the top expression:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Formulate unaligned amplitude expressions" - }, - "tags": [ - "hide-input", - "full-width", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " A^{K}_{- \\frac{1}{2}, - \\frac{1}{2}} & = & \\sum_{\\tau=-1}^{1}{- \\delta_{- \\frac{1}{2}, \\tau + \\frac{1}{2}} \\mathcal{H}^\\mathrm{decay}_{K^{*}(892)^{0}, 0, 0} \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, \\tau, - \\frac{1}{2}} \\mathcal{R}\\left(\\sigma_{1}\\right) d^{1}_{\\tau,0}\\left(\\theta_{23}\\right)} \\\\\n", - " A^{K}_{- \\frac{1}{2}, \\frac{1}{2}} & = & \\sum_{\\tau=-1}^{1}{\\delta_{- \\frac{1}{2}, \\tau - \\frac{1}{2}} \\mathcal{H}^\\mathrm{decay}_{K^{*}(892)^{0}, 0, 0} \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, \\tau, \\frac{1}{2}} \\mathcal{R}\\left(\\sigma_{1}\\right) d^{1}_{\\tau,0}\\left(\\theta_{23}\\right)} \\\\\n", - " A^{K}_{\\frac{1}{2}, - \\frac{1}{2}} & = & \\sum_{\\tau=-1}^{1}{- \\delta_{\\frac{1}{2}, \\tau + \\frac{1}{2}} \\mathcal{H}^\\mathrm{decay}_{K^{*}(892)^{0}, 0, 0} \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, \\tau, - \\frac{1}{2}} \\mathcal{R}\\left(\\sigma_{1}\\right) d^{1}_{\\tau,0}\\left(\\theta_{23}\\right)} \\\\\n", - " A^{K}_{\\frac{1}{2}, \\frac{1}{2}} & = & \\sum_{\\tau=-1}^{1}{\\delta_{\\frac{1}{2}, \\tau - \\frac{1}{2}} \\mathcal{H}^\\mathrm{decay}_{K^{*}(892)^{0}, 0, 0} \\mathcal{H}^\\mathrm{production}_{K^{*}(892)^{0}, \\tau, \\frac{1}{2}} \\mathcal{R}\\left(\\sigma_{1}\\right) d^{1}_{\\tau,0}\\left(\\theta_{23}\\right)} \\\\\n", - " A^{\\Lambda}_{- \\frac{1}{2}, - \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{- \\delta_{- \\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, - \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{3}{2}}_{\\tau,\\frac{1}{2}}\\left(\\theta_{31}\\right)} + \\sum_{\\tau=-1/2}^{1/2}{- \\delta_{- \\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, - \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{1}{2}}_{\\tau,\\frac{1}{2}}\\left(\\theta_{31}\\right)} \\\\\n", - " A^{\\Lambda}_{- \\frac{1}{2}, \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{\\delta_{- \\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{3}{2}}_{\\tau,- \\frac{1}{2}}\\left(\\theta_{31}\\right)} + \\sum_{\\tau=-1/2}^{1/2}{\\delta_{- \\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{1}{2}}_{\\tau,- \\frac{1}{2}}\\left(\\theta_{31}\\right)} \\\\\n", - " A^{\\Lambda}_{\\frac{1}{2}, - \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{- \\delta_{\\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, - \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{3}{2}}_{\\tau,\\frac{1}{2}}\\left(\\theta_{31}\\right)} + \\sum_{\\tau=-1/2}^{1/2}{- \\delta_{\\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, - \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{1}{2}}_{\\tau,\\frac{1}{2}}\\left(\\theta_{31}\\right)} \\\\\n", - " A^{\\Lambda}_{\\frac{1}{2}, \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{\\delta_{\\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1520), 0, \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1520), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{3}{2}}_{\\tau,- \\frac{1}{2}}\\left(\\theta_{31}\\right)} + \\sum_{\\tau=-1/2}^{1/2}{\\delta_{\\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Lambda(1670), 0, \\frac{1}{2}} \\mathcal{H}^\\mathrm{production}_{\\Lambda(1670), \\tau, 0} \\mathcal{R}\\left(\\sigma_{2}\\right) d^{\\frac{1}{2}}_{\\tau,- \\frac{1}{2}}\\left(\\theta_{31}\\right)} \\\\\n", - " A^{\\Delta}_{- \\frac{1}{2}, - \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{\\delta_{- \\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, - \\frac{1}{2}, 0} \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, \\tau, 0} \\mathcal{R}\\left(\\sigma_{3}\\right) d^{\\frac{3}{2}}_{\\tau,- \\frac{1}{2}}\\left(\\theta_{12}\\right)} \\\\\n", - " A^{\\Delta}_{- \\frac{1}{2}, \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{\\delta_{- \\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, \\frac{1}{2}, 0} \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, \\tau, 0} \\mathcal{R}\\left(\\sigma_{3}\\right) d^{\\frac{3}{2}}_{\\tau,\\frac{1}{2}}\\left(\\theta_{12}\\right)} \\\\\n", - " A^{\\Delta}_{\\frac{1}{2}, - \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{\\delta_{\\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, - \\frac{1}{2}, 0} \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, \\tau, 0} \\mathcal{R}\\left(\\sigma_{3}\\right) d^{\\frac{3}{2}}_{\\tau,- \\frac{1}{2}}\\left(\\theta_{12}\\right)} \\\\\n", - " A^{\\Delta}_{\\frac{1}{2}, \\frac{1}{2}} & = & \\sum_{\\tau=-3/2}^{3/2}{\\delta_{\\frac{1}{2} \\tau} \\mathcal{H}^\\mathrm{decay}_{\\Delta(1232)^{++}, \\frac{1}{2}, 0} \\mathcal{H}^\\mathrm{production}_{\\Delta(1232)^{++}, \\tau, 0} \\mathcal{R}\\left(\\sigma_{3}\\right) d^{\\frac{3}{2}}_{\\tau,\\frac{1}{2}}\\left(\\theta_{12}\\right)} \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "A = {1: A_K, 2: A_Λ, 3: A_Δ}\n", - "amp_definitions = {}\n", - "for chain_id in chain_ids:\n", - " for Λc_heli, p_heli in itertools.product([-half, +half], [-half, +half]):\n", - " symbol = A[chain_id][Λc_heli, p_heli]\n", - " expr = formulate_chain_amplitude(chain_id, ν, λ)\n", - " amp_definitions[symbol] = expr.subs({ν: Λc_heli, λ: p_heli})\n", - "display_definitions(amp_definitions)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Assert that all symbols are defined" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "unfolded_intensity_expr = perform_cached_doit(\n", - " perform_cached_doit(intensity_expr).xreplace(amp_definitions)\n", - ")\n", - "expr = unfolded_intensity_expr.xreplace(angles).doit()\n", - "expr = expr.xreplace(parameter_defaults)\n", - "assert expr.free_symbols == {σ1, σ2, σ3}\n", - "del expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Polarization sensitivity\n", - "\n", - "We introduce the **polarimeter vector field** (polarization sensitivity) of the $\\Lambda_c$ decay. It is defined by three quantities $(\\alpha_x,\\alpha_y,\\alpha_z)$ forming a three-dimensional vector $\\vec\\alpha$ dependent on just two decay variables, $\\sigma_1=m_{K\\pi}^2$, and $\\sigma_2=m_{pK}^2$.\n", - "\n", - "The polarimeter vector field is computed by averaging the Pauli matrices $\\vec\\sigma$ contracted with the $\\Lambda_c^+$ helicity indices given the transition amplitude.\n", - "\n", - "$$\n", - "\\begin{align}\n", - " \\vec\\alpha(m_{K\\pi},m_{pK}) =\n", - " \\sum_{\\lambda,\\nu,\\nu'}\n", - " A^{*}_{\\nu,\\lambda}\\vec\\sigma_{\\nu,\\nu'}\n", - " A_{\\nu',\\lambda} \\,\\big / \\sum_{\\lambda,\\nu}\n", - " \\left|A_{\\nu,\\lambda}\\right|^2\n", - "\\end{align}\n", - "$$ (polarimeter-field)\n", - "\n", - "The quantities $\\vec\\alpha(m_{K\\pi},m_{pK})$ give the model-independent representation of the $\\Lambda_c^+$ decay process. It can be used to study $\\Lambda_c^+$ production polarization using\n", - "\n", - "$$\n", - "\\begin{align}\n", - " I(\\alpha,\\beta,\\gamma,m_{K\\pi},m_{pK}) =\n", - " I_0(m_{K\\pi},m_{pK})\\,\n", - " \\left(1 + \\sum_{i,j} P_i R_{ij}(\\alpha,\\beta,\\gamma) \\alpha_j(m_{K\\pi},m_{pK}) \\right)\\,,\n", - "\\end{align}\n", - "$$ (master-intensity)\n", - "\n", - "where $R_{ij}(\\alpha,\\beta,\\gamma)$ is a three-dimensional rotation matrix:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " R(\\alpha,\\beta,\\gamma) = R_z(\\alpha)R_y(\\beta)R_z(\\gamma)\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "and $I_0$ is the averaged decay rate\n", - "\n", - "$$\n", - "\\begin{align}\n", - " I_0(m_{K\\pi},m_{pK}) = \\sum_{\\lambda,\\nu}\\left|A_{\\nu,\\lambda}\\right|^2\\,.\n", - "\\end{align}\n", - "$$" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def to_index(helicity):\n", - " \"\"\"Symbolic conversion of half-value helicities to Pauli matrix indices.\"\"\"\n", - " # https://github.com/ComPWA/compwa.github.io/pull/129#issuecomment-1096599896\n", - " return sp.Piecewise(\n", - " (1, sp.LessThan(helicity, 0)),\n", - " (0, True),\n", - " )\n", - "\n", - "\n", - "ν_prime = sp.Symbol(R\"\\nu^{\\prime}\")\n", - "polarimetry_exprs = tuple(\n", - " PoolSum(\n", - " formulate_aligned_amplitude(ν, λ).conjugate()\n", - " * msigma(i)[to_index(ν), to_index(ν_prime)]\n", - " * formulate_aligned_amplitude(ν_prime, λ),\n", - " (λ, [-half, +half]),\n", - " (ν, [-half, +half]),\n", - " (ν_prime, [-half, +half]),\n", - " )\n", - " / intensity_expr\n", - " for i in (1, 2, 3)\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Unfold polarimeter expressions" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "unfolded_polarimetry_exprs = tuple(\n", - " perform_cached_doit(perform_cached_doit(x).xreplace(amp_definitions))\n", - " for x in polarimetry_exprs\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Properties of the vector $\\vec\\alpha$\n", - "\n", - "The vector $\\vec \\alpha$ introduced in Eq. {eq}`polarimeter-field` obeys the following properties:\n", - "\n", - "1. It is a three-dimensional vector defined in the rest frame of the decaying particle.\n", - " Particularly, it is transformed as a regular vector in case initial (alignment) configuration change.\n", - "2. The length of the vector is limited by 1: $|\\vec{\\alpha}| < 1$ \n", - "3. $\\alpha_y=0$ for the decays of a fermion to a fermions and (pseudo)scalar\n", - "\n", - "Here is the prove of the second statement:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " I_{\\nu',\\nu} = \\sum_{\\lambda} A_{\\nu',\\lambda}^* A_{\\nu,\\lambda} =\n", - " \\begin{pmatrix}\n", - " a & c^*\\\\\n", - " c & b\n", - " \\end{pmatrix} = \\frac{a+b}{2}\\left(\n", - " \\mathbb{I} + (\\vec{\\sigma} \\cdot \\vec{\\alpha})\n", - " \\right)\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "where\n", - "\n", - "$$\n", - "\\begin{align}\n", - "a &= \\left|A_{+,+}\\right|^2+\\left|A_{+,-}\\right|^2\\,,\\\\ \\nonumber\n", - "b &= \\left|A_{-,+}\\right|^2+\\left|A_{-,-}\\right|^2\\,,\\\\ \\nonumber\n", - "c &= A_{+,+}^*A_{-,+} + A_{+,-}^*A_{-,-}\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "and\n", - "\n", - "$$\n", - "\\begin{align}\n", - "\\alpha_x &= \\frac{\\text{Re}\\,c}{a+b}\\,,&\n", - "\\alpha_x &= \\frac{\\text{Im}\\,c}{a+b}\\,,&\n", - "\\alpha_z &= \\frac{a-b}{a+b}\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "To constraint the length of the $\\vec\\alpha$, one notices $ab - c \\ge 0$. Therefore,\n", - "\n", - "$$\n", - "\\begin{align}\n", - "|\\vec\\alpha|^2 &= \\frac{(a-b)^2+c^2}{(a+b)^2} = \\frac{(a+b)^2-4ab+c^2}{(a+b)^2} \\leq \\frac{(a+b)^2-3ab}{(a+b)^2} \\leq 1\\,\n", - "\\end{align}\n", - "$$\n", - "\n", - "since $a,b \\geq 0$." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Computations with TensorWaves" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Conversion to computational backend\n", - "\n", - "The full [expression tree](https://docs.sympy.org/latest/tutorial/manipulation.html) can be converted to a computational, _parametrized_ function as follows. Note that identify all coupling symbols are interpreted as parameters. The remaining symbols (the angles) become arguments to the function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "free_parameters = {\n", - " symbol: value\n", - " for symbol, value in parameter_defaults.items()\n", - " if (symbol.name.startswith(\"m_\") and symbol not in masses)\n", - " or symbol.name.startswith(R\"\\Gamma_\")\n", - " or symbol in prod_couplings\n", - "}\n", - "fixed_parameters = {\n", - " symbol: value\n", - " for symbol, value in parameter_defaults.items()\n", - " if symbol not in free_parameters\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensity_func = create_parametrized_function(\n", - " unfolded_intensity_expr.xreplace(fixed_parameters),\n", - " parameters=free_parameters,\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "polarimetry_funcs = tuple(\n", - " create_parametrized_function(\n", - " expr.xreplace(fixed_parameters),\n", - " parameters=free_parameters,\n", - " backend=\"jax\",\n", - " )\n", - " for expr in unfolded_polarimetry_exprs\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phase space\n", - "\n", - "The $\\Lambda_c^+ \\to p K \\pi$ kinematics is fully described by two dynamic variables, $m_{K\\pi}$ and $m_{pK}$ (see {doc}`/report/017`). The third Mandelstam variable can be computed from the other two and the masses of the initial and final state:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{rcl}\n", - " \\sigma_{3} & = & m_{K}^{2} + m_{p}^{2} + m_{\\pi}^{2} + m_{\\Lambda_c}^{2} - \\sigma_{1} - \\sigma_{2} \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "computed_σ3 = m0**2 + m1**2 + m2**2 + m3**2 - σ1 - σ2\n", - "compute_third_mandelstam = create_function(computed_σ3.subs(masses), backend=\"jax\")\n", - "display_definitions({σ3: computed_σ3})" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Values for the angles will be computed form the Mandelstam values with a data transformer for the symbolic angle definitions:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "kinematic_variables = {\n", - " symbol: expression.doit().xreplace(masses).xreplace(fixed_parameters)\n", - " for symbol, expression in angles.items()\n", - "}\n", - "kinematic_variables.update({s: s for s in [σ1, σ2, σ3]}) # include identity\n", - "transformer = SympyDataTransformer.from_sympy(kinematic_variables, backend=\"jax\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now define phase space over a grid that contains the space in the Dalitz plane that is kinematically 'available' to the decay:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "m0_val, m1_val, m2_val, m3_val = masses.values()\n", - "σ1_min = (m2_val + m3_val) ** 2\n", - "σ1_max = (m0_val - m1_val) ** 2\n", - "σ2_min = (m1_val + m3_val) ** 2\n", - "σ2_max = (m0_val - m2_val) ** 2\n", - "\n", - "\n", - "def generate_phsp_grid(resolution: int):\n", - " x = np.linspace(σ1_min, σ1_max, num=resolution)\n", - " y = np.linspace(σ2_min, σ2_max, num=resolution)\n", - " X, Y = np.meshgrid(x, y)\n", - " Z = compute_third_mandelstam.function(X, Y)\n", - " phsp = {\"sigma1\": X, \"sigma2\": Y, \"sigma3\": Z}\n", - " return X, Y, transformer(phsp)\n", - "\n", - "\n", - "X, Y, phsp = generate_phsp_grid(resolution=500)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Intensity distribution\n", - "\n", - "Finally, all intensities can be computed as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", - "s1_label = R\"$\\sigma_1=m^2\\left(K\\pi\\right)$\"\n", - "s2_label = R\"$\\sigma_2=m^2\\left(pK\\right)$\"\n", - "s3_label = R\"$\\sigma_3=m^2\\left(p\\pi\\right)$\"\n", - "\n", - "plt.rc(\"font\", size=15)\n", - "fig, ax = plt.subplots(\n", - " figsize=(9, 8),\n", - " tight_layout=True,\n", - ")\n", - "ax.set_box_aspect(1)\n", - "ax.set_title(\"Intensity distribution\")\n", - "ax.set_xlabel(s1_label)\n", - "ax.set_ylabel(s2_label)\n", - "\n", - "total_intensities = intensity_func(phsp)\n", - "mesh = ax.pcolormesh(X, Y, total_intensities, norm=LogNorm())\n", - "fig.colorbar(mesh, ax=ax, fraction=0.05, pad=0.02)\n", - "fig.savefig(\"021-intensity-distribution.png\", dpi=200)\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{only} html\n", - "![](https://user-images.githubusercontent.com/29308176/199110362-07543795-facd-4482-9973-596b6c0ee2c4.png)\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input", - "full-width", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "\n", - "\n", - "def compute_sub_function(\n", - " func: ParametrizedFunction, phsp: DataSample, non_zero_couplings: str\n", - ") -> jnp.ndarray:\n", - " zero_couplings = {\n", - " par: 0\n", - " for par in func.parameters\n", - " if par.startswith(R\"\\mathcal{H}\")\n", - " if \"production\" in par\n", - " if not any(s in par for s in non_zero_couplings)\n", - " }\n", - " original_parameters = dict(func.parameters)\n", - " func.update_parameters(zero_couplings)\n", - " computed_values = func(phsp)\n", - " func.update_parameters(original_parameters)\n", - " return computed_values\n", - "\n", - "\n", - "def set_ylim_to_zero(ax):\n", - " _, y_max = ax.get_ylim()\n", - " ax.set_ylim(0, y_max)\n", - "\n", - "\n", - "fig, (ax1, ax2) = plt.subplots(\n", - " ncols=2,\n", - " figsize=(12, 5),\n", - " sharey=True,\n", - " tight_layout=True,\n", - ")\n", - "ax1.set_xlabel(s1_label)\n", - "ax2.set_xlabel(s2_label)\n", - "ax1.set_yticks([])\n", - "\n", - "x = X[0]\n", - "y = Y[:, 0]\n", - "ax1.fill(x, np.nansum(total_intensities, axis=0), alpha=0.3)\n", - "ax2.fill(y, np.nansum(total_intensities, axis=1), alpha=0.3)\n", - "for chain_id, chain_name in chain_ids.items():\n", - " label = f\"${chain_labels[chain_id]}$\"\n", - " sub_intensities = compute_sub_function(\n", - " intensity_func, phsp, non_zero_couplings=[chain_name]\n", - " )\n", - " ax1.plot(x, np.nansum(sub_intensities, axis=0), label=label)\n", - " ax2.plot(y, np.nansum(sub_intensities, axis=1), label=label)\n", - "set_ylim_to_zero(ax1)\n", - "set_ylim_to_zero(ax2)\n", - "ax2.legend()\n", - "fig.savefig(\"021-intensity-projections.svg\")\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{container} full-width\n", - "![](https://user-images.githubusercontent.com/29308176/199110474-d47016bc-6200-4ecc-a017-e77df2646ee6.svg)\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Fit fractions\n", - "\n", - "The total decay rate for $\\Lambda_c^+ \\to pK\\pi$ can be broken into fractions that correspond to the different decay chains and interference terms. The total rate is computed as an integral of the intensity over decay kinematics:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " I_\\text{tot}(\\{\\mathcal{H}\\}) = \\int d m_{pK}^2 d m_{K\\pi}^2\\,\n", - " I_0(m_{pK}, m_{K\\pi} | \\{\\mathcal{H}\\})\n", - " \\approx \\frac{\\Phi_0}{N_\\text{MC}} \\sum_{e=1}^{N_\\text{MC}}\\,\\,I_0(m_{pK,e}, m_{K\\pi,e} | \\{\\mathcal{H}\\})\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "where $\\Phi_0$ is an (irrelevant) constant equal to the flat phase-space integral, $(m_{pK,e}, m_{K\\pi,e})$ is a vector of the kinematic variables for the $e$-th point in the MC sample.\n", - "\n", - "The conditional argument $\\{\\mathcal{H}\\}$ indicates dependence of the rate on the value of the couplings. The individual fractions are found by computing the total rate for a subset of couplings set to zero,\n", - "\n", - "$$\n", - "\\begin{align}\n", - " I_\\text{tot}^{K} &= I_\\text{tot}\\left(\\{\\mathcal{H}^{\\Lambda_c^+\\to\\Delta^{**} K}, \\mathcal{H}^{\\Lambda_c^+\\to\\Lambda^{**} \\pi} = 0\\}\\right)\\,,\\\\\n", - " I_\\text{tot}^{\\Delta} &= I_\\text{tot}\\left(\\{\\mathcal{H}^{\\Lambda_c^+\\to K^{**} p}, \\mathcal{H}^{\\Lambda_c^+\\to\\Lambda^{**} \\pi} = 0\\}\\right)\\,,\\\\\n", - " I_\\text{tot}^{\\Lambda} &= I_\\text{tot}\\left(\\{\\mathcal{H}^{\\Lambda_c^+\\to\\Delta^{**} K}, \\mathcal{H}^{\\Lambda_c^+\\to K^{**} p} = 0\\}\\right)\\,,\\\\\n", - " I_\\text{tot}^{K/\\Lambda} &= I_\\text{tot}\\left(\\{\\mathcal{H}^{\\Lambda_c^+\\to\\Delta^{**} K} = 0\\}\\right) - I_\\text{tot}^{K} - I_\\text{tot}^{\\Lambda}\\,,\\\\\n", - " & \\dots\\,,\n", - "\\end{align}\n", - "$$\n", - "\n", - "where the terms with a single chain index are the rate of the decay chain. The sum of all fractions should give the total rate:\n", - "\n", - "$$\n", - "\\begin{align}\n", - " I_\\text{tot}\\left(\\{\\mathcal{H}\\}\\right)\n", - " = \\sum_{R} I_\\text{tot}^{R} + \\sum_{R < R'} I_\\text{tot}^{R/R'}\n", - "\\end{align}\n", - "$$" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_hide": "Code for computing decay rates" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def integrate_intensity(\n", - " intensity_func: ParametrizedFunction,\n", - " phsp: DataSample,\n", - " non_zero_couplings: list[str] | None = None,\n", - ") -> float:\n", - " if non_zero_couplings is None:\n", - " intensities = intensity_func(phsp)\n", - " else:\n", - " intensities = compute_sub_function(intensity_func, phsp, non_zero_couplings)\n", - " return np.nansum(intensities) / len(intensities)\n", - "\n", - "\n", - "def compute_interference(\n", - " intensity_func: ParametrizedFunction,\n", - " phsp: DataSample,\n", - " chain1: list[str],\n", - " chain2: list[str],\n", - ") -> float:\n", - " I_interference = integrate_intensity(intensity_func, phsp, chain1 + chain2)\n", - " I_chain1 = integrate_intensity(intensity_func, phsp, chain1)\n", - " I_chain2 = integrate_intensity(intensity_func, phsp, chain2)\n", - " return I_interference - I_chain1 - I_chain2\n", - "\n", - "\n", - "I_tot = integrate_intensity(intensity_func, phsp)\n", - "np.testing.assert_allclose(\n", - " I_tot,\n", - " integrate_intensity(intensity_func, phsp, [\"K\", R\"\\Lambda\", R\"\\Delta\"]),\n", - ")\n", - "I_K = integrate_intensity(intensity_func, phsp, non_zero_couplings=[\"K\"])\n", - "I_Λ = integrate_intensity(intensity_func, phsp, non_zero_couplings=[\"Lambda\"])\n", - "I_Δ = integrate_intensity(intensity_func, phsp, non_zero_couplings=[\"Delta\"])\n", - "I_ΛΔ = compute_interference(intensity_func, phsp, [\"Lambda\"], [\"Delta\"])\n", - "I_KΔ = compute_interference(intensity_func, phsp, [\"K\"], [\"Delta\"])\n", - "I_KΛ = compute_interference(intensity_func, phsp, [\"K\"], [\"Lambda\"])\n", - "np.testing.assert_allclose(I_tot, I_K + I_Λ + I_Δ + I_ΛΔ + I_KΔ + I_KΛ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{crr}\n", - "& I_\\mathrm{sub}\\,/\\,I \\\\\n", - "\\hline K^{**} & 0.371 \\\\\n", - " \\Lambda^{**} & 0.051 \\\\\n", - " \\color{gray}{\\Lambda(1520)} & \\color{gray}{0.031} \\\\\n", - " \\color{gray}{\\Lambda(1670)} & \\color{gray}{0.020} \\\\\n", - " \\Delta^{**} & 0.582 \\\\\n", - " \\Delta/\\Lambda & 0.013 \\\\\n", - " K/\\Delta & -0.018 \\\\\n", - " K/\\Lambda & 0.001 \\\\\n", - "\\hline \\mathrm{total} & 1.000 \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def render_resonance_row(chain_id):\n", - " rows = [\n", - " (\n", - " Rf\"\\color{{gray}}{{{p.latex}}}\",\n", - " (\n", - " Rf\"\\color{{gray}}{{{integrate_intensity(intensity_func, phsp, [p.name]) / I_tot:.3f}}}\"\n", - " ),\n", - " )\n", - " for p in resonance_particles[chain_id]\n", - " ]\n", - " if len(rows) > 1:\n", - " return rows\n", - " return []\n", - "\n", - "\n", - "rows = [\n", - " R\"\\hline\",\n", - " (\"K^{**}\", f\"{I_K / I_tot:.3f}\"),\n", - " *render_resonance_row(chain_id=1),\n", - " (R\"\\Lambda^{**}\", f\"{I_Λ / I_tot:.3f}\"),\n", - " *render_resonance_row(chain_id=2),\n", - " (R\"\\Delta^{**}\", f\"{I_Δ / I_tot:.3f}\"),\n", - " *render_resonance_row(chain_id=3),\n", - " (R\"\\Delta/\\Lambda\", f\"{I_ΛΔ / I_tot:.3f}\"),\n", - " (R\"K/\\Delta\", f\"{I_KΔ / I_tot:.3f}\"),\n", - " (R\"K/\\Lambda\", f\"{I_KΛ / I_tot:.3f}\"),\n", - " R\"\\hline\",\n", - " (\n", - " R\"\\mathrm{total}\",\n", - " f\"{(I_K + I_Λ + I_Δ + I_ΛΔ + I_KΔ + I_KΛ) / I_tot:.3f}\",\n", - " ),\n", - "]\n", - "\n", - "latex = R\"\\begin{array}{crr}\" + \"\\n\"\n", - "latex += R\"& I_\\mathrm{sub}\\,/\\,I \\\\\" + \"\\n\"\n", - "for row in rows:\n", - " if row == R\"\\hline\":\n", - " latex += R\"\\hline\"\n", - " else:\n", - " latex += \" \" + \" & \".join(row) + R\" \\\\\" + \"\\n\"\n", - "latex += R\"\\end{array}\"\n", - "Math(latex)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Polarimetry distributions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "for func in polarimetry_funcs:\n", - " assert np.nanmax(func(phsp).imag) < 1e-10" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "{{ run_interactive }}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/latex": [ - "$\\displaystyle \\begin{array}{cccc}\n", - "& \\bar{|\\alpha|} & \\bar\\alpha_x & \\bar\\alpha_y & \\bar\\alpha_z \\\\\n", - " K^{**} & 0.873 \\pm 0.040 & 0.000 \\pm 0.554 & 0.000 \\pm 0.396 & +0.100 \\pm 0.538 \\\\\n", - " \\Lambda^{**} & 0.906 \\pm 0.101 & -0.529 \\pm 0.214 & 0.000 \\pm 0.190 & -0.338 \\pm 0.596 \\\\\n", - " \\Delta^{**} & 0.540 \\pm 0.000 & +0.320 \\pm 0.153 & -0.000 \\pm 0.000 & -0.324 \\pm 0.246 \\\\\n", - "\\end{array}$" - ], - "text/plain": [ - "" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def render_mean(array, plus=True):\n", - " array = array.real\n", - " mean = f\"{np.nanmean(array):.3f}\"\n", - " std = f\"{np.nanstd(array):.3f}\"\n", - " if plus and float(mean) > 0:\n", - " mean = f\"+{mean}\"\n", - " return Rf\"{mean} \\pm {std}\"\n", - "\n", - "\n", - "latex = R\"\\begin{array}{cccc}\" + \"\\n\"\n", - "latex += R\"& \\bar{|\\alpha|} & \\bar\\alpha_x & \\bar\\alpha_y & \\bar\\alpha_z \\\\\" + \"\\n\"\n", - "for chain_id, chain_name in chain_ids.items():\n", - " latex += f\" {chain_labels[chain_id]} & \"\n", - " x, y, z = tuple(\n", - " compute_sub_function(func, phsp, non_zero_couplings=[chain_name])\n", - " for func in polarimetry_funcs\n", - " )\n", - " latex += render_mean(np.sqrt(x**2 + y**2 + z**2), plus=False) + \" & \"\n", - " latex += \" & \".join(map(render_mean, [x, y, z]))\n", - " latex += R\" \\\\\" + \"\\n\"\n", - "latex += R\"\\end{array}\"\n", - "Math(latex)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define sliders for the widget" - }, - "tags": [ - "scroll-input", - "hide-cell", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "# Slider construction\n", - "sliders = {}\n", - "for symbol in free_parameters:\n", - " if symbol.name.startswith(R\"\\mathcal{H}\"):\n", - " real_slider = create_slider(symbol)\n", - " imag_slider = create_slider(symbol)\n", - " sliders[f\"{symbol.name}_real\"] = real_slider\n", - " sliders[f\"{symbol.name}_imag\"] = imag_slider\n", - " real_slider.description = R\"\\(\\mathrm{Re}\\)\"\n", - " imag_slider.description = R\"\\(\\mathrm{Im}\\)\"\n", - " else:\n", - " slider = create_slider(symbol)\n", - " sliders[symbol.name] = slider\n", - "\n", - "# Slider ranges\n", - "σ3_max = (m0_val - m3_val) ** 2\n", - "σ3_min = (m1_val + m2_val) ** 2\n", - "\n", - "for name, slider in sliders.items():\n", - " slider.continuous_update = True\n", - " slider.step = 0.01\n", - " if name.startswith(\"m_\"):\n", - " if \"K\" in name:\n", - " slider.min = np.sqrt(σ1_min)\n", - " slider.max = np.sqrt(σ1_max)\n", - " elif R\"\\Lambda\" in name:\n", - " slider.min = np.sqrt(σ2_min)\n", - " slider.max = np.sqrt(σ2_max)\n", - " elif R\"\\Delta\" in name:\n", - " slider.min = np.sqrt(σ3_min)\n", - " slider.max = np.sqrt(σ3_max)\n", - " elif name.startswith(R\"\\Gamma_\"):\n", - " slider.min = 0\n", - " slider.max = max(0.5, 2 * slider.value)\n", - " elif name.startswith(R\"\\mathcal{H}\"):\n", - " slider.min = -15\n", - " slider.max = +15\n", - "\n", - "\n", - "# Slider values\n", - "def reset_sliders(click_event):\n", - " for symbol, value in free_parameters.items():\n", - " if symbol.name.startswith(R\"\\mathcal{H}\"):\n", - " set_slider(sliders[symbol.name + \"_real\"], value)\n", - " set_slider(sliders[symbol.name + \"_imag\"], value)\n", - " else:\n", - " set_slider(sliders[symbol.name], value)\n", - "\n", - "\n", - "def set_coupling_to_zero(filter_pattern):\n", - " if isinstance(filter_pattern, Combobox):\n", - " filter_pattern = filter_pattern.value\n", - " for name, _slider in sliders.items():\n", - " if not name.startswith(R\"\\mathcal{H}\"):\n", - " continue\n", - " if filter_pattern not in name:\n", - " continue\n", - " set_slider(_slider, 0)\n", - "\n", - "\n", - "def set_slider(slider, value):\n", - " if slider.description == R\"\\(\\mathrm{Im}\\)\":\n", - " value = complex(value).imag\n", - " else:\n", - " value = complex(value).real\n", - " n_decimals = -round(np.log10(slider.step))\n", - " if slider.value != round(value, n_decimals): # widget performance\n", - " slider.value = value\n", - "\n", - "\n", - "reset_sliders(click_event=None)\n", - "reset_button = Button(description=\"Reset slider values\")\n", - "reset_button.on_click(reset_sliders)\n", - "\n", - "all_resonances = [r.latex for r_list in resonance_particles.values() for r in r_list]\n", - "filter_button = Combobox(\n", - " placeholder=\"Enter coupling filter pattern\",\n", - " options=all_resonances,\n", - " description=R\"$\\mathcal{H}=0$\",\n", - ")\n", - "filter_button.on_submit(set_coupling_to_zero)\n", - "\n", - "# UI design\n", - "latex = {symbol.name: sp.latex(symbol) for symbol in free_parameters}\n", - "mass_sliders = [sliders[n] for n in sliders if n.startswith(\"m_\")]\n", - "width_sliders = [sliders[n] for n in sliders if n.startswith(R\"\\Gamma_\")]\n", - "coupling_sliders = {}\n", - "for res_list in resonance_particles.values():\n", - " for res in res_list:\n", - " coupling_sliders[res.name] = (\n", - " [s for n, s in sliders.items() if n.endswith(\"_real\") and res.latex in n],\n", - " [s for n, s in sliders.items() if n.endswith(\"_imag\") and res.latex in n],\n", - " [\n", - " HTMLMath(f\"${latex[n[:-5]]}$\")\n", - " for n in sliders\n", - " if n.endswith(\"_real\") and res.latex in n\n", - " ],\n", - " )\n", - "slider_tabs = Tab(\n", - " children=[\n", - " Tab(\n", - " children=[\n", - " VBox([HBox(s) for s in zip(*pair)])\n", - " for pair in coupling_sliders.values()\n", - " ],\n", - " titles=tuple(coupling_sliders),\n", - " ),\n", - " VBox([HBox([r, i]) for r, i in zip(mass_sliders, width_sliders)]),\n", - " ],\n", - " titles=(\"Couplings\", \"Masses and widths\"),\n", - ")\n", - "ui = VBox([slider_tabs, HBox([reset_button, filter_button])])" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Create interactive plot" - }, - "tags": [ - "scroll-input", - "remove-output", - "full-width", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "\n", - "fig, axes = plt.subplots(\n", - " figsize=(12, 6.2),\n", - " ncols=2,\n", - " sharey=True,\n", - ")\n", - "ax1, ax2 = axes\n", - "ax1.set_title(\"Intensity distribution\")\n", - "ax2.set_title(\"Polarimeter vector field\")\n", - "ax1.set_xlabel(Rf\"${s1_label[1:-1]}, \\alpha_x$\")\n", - "ax2.set_xlabel(Rf\"${s1_label[1:-1]}, \\alpha_x$\")\n", - "ax1.set_ylabel(Rf\"${s2_label[1:-1]}, \\alpha_z$\")\n", - "for ax in axes:\n", - " ax.set_box_aspect(1)\n", - "fig.canvas.toolbar_visible = False\n", - "fig.canvas.header_visible = False\n", - "fig.canvas.footer_visible = False\n", - "\n", - "mesh = None\n", - "quiver = None\n", - "Xα, Yα, phsp_α = generate_phsp_grid(resolution=35)\n", - "XI, YI, phsp_I = generate_phsp_grid(resolution=200)\n", - "\n", - "\n", - "def plot3(**kwargs):\n", - " global quiver, mesh\n", - " kwargs = to_complex_kwargs(**kwargs)\n", - " for func in [*list(polarimetry_funcs), intensity_func]:\n", - " func.update_parameters(kwargs)\n", - " intensity = intensity_func(phsp_I)\n", - " αx, αy, αz = tuple(func(phsp_α).real for func in polarimetry_funcs)\n", - " abs_α = jnp.sqrt(αx**2 + αy**2 + αz**2)\n", - " if mesh is None:\n", - " mesh = ax1.pcolormesh(XI, YI, intensity, cmap=plt.cm.Reds)\n", - " c_bar = fig.colorbar(mesh, ax=ax1, pad=0.01, fraction=0.0473)\n", - " c_bar.ax.set_yticks([])\n", - " else:\n", - " mesh.set_array(intensity)\n", - " if quiver is None:\n", - " quiver = ax2.quiver(Xα, Yα, αx, αz, abs_α, cmap=plt.cm.viridis_r, clim=(0, 1))\n", - " c_bar = fig.colorbar(quiver, ax=ax2, pad=0.01, fraction=0.0473)\n", - " c_bar.ax.set_ylabel(R\"$\\left|\\vec\\alpha\\right|$\")\n", - " else:\n", - " quiver.set_UVC(αx, αz, abs_α)\n", - " fig.canvas.draw_idle()\n", - "\n", - "\n", - "def to_complex_kwargs(**kwargs):\n", - " complex_valued_kwargs = {}\n", - " for key, value in dict(kwargs).items():\n", - " if key.endswith(\"real\"):\n", - " symbol_name = key[:-5]\n", - " imag = kwargs[f\"{symbol_name}_imag\"]\n", - " complex_valued_kwargs[symbol_name] = complex(value, imag)\n", - " elif key.endswith(\"imag\"):\n", - " continue\n", - " else:\n", - " complex_valued_kwargs[key] = value\n", - " return complex_valued_kwargs\n", - "\n", - "\n", - "output = interactive_output(plot3, controls=sliders)\n", - "fig.tight_layout()\n", - "display(ui, output)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " filename = \"021-polarimeter-field.png\"\n", - " plt.savefig(filename, dpi=200)\n", - " display(ui, Image(filename))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{container} full-width\n", - "![](https://user-images.githubusercontent.com/29308176/199110481-c37d5bad-746e-4670-8367-697bd3d71520.png)\n", - ":::" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/022.ipynb b/docs/report/022.ipynb deleted file mode 100644 index f3db0266..00000000 --- a/docs/report/022.ipynb +++ /dev/null @@ -1,740 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "physics", - "polarimetry", - "polarization" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Polarimetry: Computing the B-matrix for Λc→pKπ\n", - "TR-022\n", - "^^^\n", - "The $B$-matrix forms an extension of the polarimeter vector field $\\vec\\alpha$ ([arXiv:2301.07010](https://arxiv.org/abs/2301.07010), see also [TR-021](021.ipynb)) that takes the polarization of the proton into account. See [arXiv:2302.07665](https://arxiv.org/abs/2302.07665), Eq. (B6).\n", - "+++\n", - "✅ [compwa.github.io#196](https://github.com/ComPWA/compwa.github.io/pull/196)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# B-matrix extension of polarimeter" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q git+https://github.com/ComPWA/polarimetry@0.0.9" - ] - }, - { - "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 logging\n", - "from pathlib import Path\n", - "from warnings import filterwarnings\n", - "\n", - "import jax.numpy as jnp\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import polarimetry\n", - "import sympy as sp\n", - "from ampform.sympy import PoolSum\n", - "from IPython.display import display\n", - "from matplotlib import cm\n", - "from polarimetry import _to_index\n", - "from polarimetry.data import create_data_transformer, generate_meshgrid_sample\n", - "from polarimetry.io import (\n", - " mute_jax_warnings,\n", - " perform_cached_doit,\n", - " perform_cached_lambdify,\n", - ")\n", - "from polarimetry.lhcb import load_model_builder, load_model_parameters\n", - "from polarimetry.lhcb.particle import load_particles\n", - "from sympy.physics.matrices import msigma\n", - "from tqdm.auto import tqdm\n", - "\n", - "filterwarnings(\"ignore\")\n", - "logging.getLogger(\"polarimetry.function\").setLevel(logging.INFO)\n", - "mute_jax_warnings()\n", - "POLARIMETRY_DIR = Path(polarimetry.__file__).parent" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Formulate expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Reference subsystem 1 is defined as:\n", - "\n", - "![](022/orientation-K.svg)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Formulate amplitude models" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "MODEL_CHOICE = 0\n", - "MODEL_FILE = POLARIMETRY_DIR / \"lhcb/model-definitions.yaml\"\n", - "PARTICLES = load_particles(POLARIMETRY_DIR / \"lhcb/particle-definitions.yaml\")\n", - "BUILDER = load_model_builder(MODEL_FILE, PARTICLES, model_id=MODEL_CHOICE)\n", - "IMPORTED_PARAMETER_VALUES = load_model_parameters(\n", - " MODEL_FILE, BUILDER.decay, MODEL_CHOICE, PARTICLES\n", - ")\n", - "REFERENCE_SUBSYSTEM = 1\n", - "MODEL = BUILDER.formulate(REFERENCE_SUBSYSTEM, cleanup_summations=True)\n", - "MODEL.parameter_defaults.update(IMPORTED_PARAMETER_VALUES)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "\\vec\\alpha = \\sum_{\\nu',\\nu,\\lambda} A^*_{\\nu',\\lambda}\\vec\\sigma_{\\nu',\\nu} A_{\\nu,\\lambda} / I_0 \\\\\n", - "\\vec\\beta = \\sum_{\\nu,\\lambda',\\lambda} A^*_{\\nu,\\lambda'} \\vec\\sigma_{\\lambda',\\lambda} A^*_{\\nu,\\lambda} / I_0 \\\\\n", - "B_{\\tau,\\rho} = \\sum_{\\nu,\\nu',\\lambda',\\lambda} A^*_{\\nu',\\lambda'} \\sigma_{\\nu',\\nu}^\\tau A_{\\nu,\\lambda} \\sigma_{\\lambda',\\lambda}^\\rho\n", - "$$" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Symbol definitions" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "half = sp.Rational(1, 2)\n", - "λ, λp = sp.symbols(R\"lambda \\lambda^{\\prime}\", rational=True)\n", - "v, vp = sp.symbols(R\"nu \\nu^{\\prime}\", rational=True)\n", - "σ = [sp.Matrix([[1, 0], [0, 1]])]\n", - "σ.extend(msigma(i) for i in (1, 2, 3))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ref = REFERENCE_SUBSYSTEM\n", - "B = tuple(\n", - " tuple(\n", - " PoolSum(\n", - " BUILDER.formulate_aligned_amplitude(vp, λp, 0, 0, ref)[0].conjugate()\n", - " * σ[τ][_to_index(vp), _to_index(v)]\n", - " * BUILDER.formulate_aligned_amplitude(v, λ, 0, 0, ref)[0]\n", - " * σ[ρ][_to_index(λp), _to_index(λ)],\n", - " (v, [-half, +half]),\n", - " (vp, [-half, +half]),\n", - " (λ, [-half, +half]),\n", - " (λp, [-half, +half]),\n", - " ).cleanup()\n", - " for ρ in range(4)\n", - " )\n", - " for τ in range(4)\n", - ")\n", - "del ref\n", - "B = sp.Matrix(B)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "B_test = tuple(\n", - " tuple(sp.Symbol(Rf\"\\tau={τ}, \\rho={p}\") for ρ in range(4)) for τ in range(4)\n", - ")\n", - "B_test = sp.Matrix(B_test)\n", - "B_test" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "B_test[0, 2]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Functions and data" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Unfold symbolic expressions" - }, - "tags": [ - "hide-cell", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "cb7caa0c47e04b638794fa84a08f5d4c", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Unfolding expressions: 0%| | 0/16 [00:00 - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - diff --git a/docs/report/023.ipynb b/docs/report/023.ipynb deleted file mode 100644 index b45cd652..00000000 --- a/docs/report/023.ipynb +++ /dev/null @@ -1,140 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "documentation", - "jupyter", - "sphinx", - "3d" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Support for Plotly plots in Technical Reports\n", - "TR-023\n", - "^^^\n", - "+++\n", - "✅ [compwa.github.io#206](https://github.com/ComPWA/compwa.github.io/issues/206)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# 3D plots with Plotly\n", - "" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "source": [ - "" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q plotly==5.18.0" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This TR tests whether the HTML build of the TR notebooks supports [Plotly](https://plotly.com/python) figures. It's a follow-up to [TR-006](006.ipynb), without the interactivity of `ipywidgets`, but with the better 3D rendering of Plotly. For more info on how Plotly figures can be embedded in Sphinx HTML builds, see [this page](https://myst-nb.readthedocs.io/en/v0.17.2/render/interactive.html#plotly) of MyST-NB (particularly the remark on [`html_js_files`](https://www.sphinx-doc.org/en/master/usage/configuration.html#confval-html_js_files).\n", - "\n", - "The following example is copied from [this tutorial](https://plotly.com/python/3d-isosurface-plots)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import numpy as np\n", - "import plotly.graph_objects as go\n", - "\n", - "X, Y, Z = np.mgrid[-5:5:40j, -5:5:40j, -5:5:40j]\n", - "ellipsoid = X * X * 0.5 + Y * Y + Z * Z * 2\n", - "fig = go.Figure(\n", - " data=go.Isosurface(\n", - " x=X.flatten(),\n", - " y=Y.flatten(),\n", - " z=Z.flatten(),\n", - " value=ellipsoid.flatten(),\n", - " isomin=5,\n", - " isomax=50,\n", - " surface_fill=0.4,\n", - " caps=dict(x_show=False, y_show=False),\n", - " slices_z=dict(\n", - " show=True,\n", - " locations=[\n", - " -1,\n", - " -3,\n", - " ],\n", - " ),\n", - " slices_y=dict(show=True, locations=[0]),\n", - " )\n", - ")\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso} {ref}`report/006:Plotly with ipywidgets`\n", - ":::" - ] - } - ], - "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.10.13" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/024.ipynb b/docs/report/024.ipynb deleted file mode 100644 index 3322ba3d..00000000 --- a/docs/report/024.ipynb +++ /dev/null @@ -1,821 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "documentation" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Symbolic expressions and model serialization\n", - "TR-024\n", - "^^^\n", - "\n", - "Investigation into dumping SymPy expressions to human-readable format for model preservation. The notebook was motivated by the [COMAP-V workshop on analysis preservation](https://indico.cern.ch/event/1348003/). See also SymPy [printing](https://docs.sympy.org/latest/modules/printing.html), [parsing](https://docs.sympy.org/latest/modules/parsing.html), and [expression manipulation](https://docs.sympy.org/latest/tutorials/intro-tutorial/manipulation.html).\n", - "+++\n", - "🚧 [polarimetry#319](https://github.com/ComPWA/polarimetry/pull/319)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Symbolic model serialization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.10 graphviz==0.20.1 polarimetry-lc2pkpi==0.0.10 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "scroll-input", - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "from pathlib import Path\n", - "from textwrap import shorten\n", - "\n", - "import graphviz\n", - "import polarimetry\n", - "import sympy as sp\n", - "from ampform.io import aslatex\n", - "from ampform.sympy import unevaluated\n", - "from IPython.display import Markdown, Math\n", - "from polarimetry.amplitude import simplify_latex_rendering\n", - "from polarimetry.io import perform_cached_doit\n", - "from polarimetry.lhcb import load_model\n", - "from polarimetry.lhcb.particle import load_particles\n", - "from sympy.printing.mathml import MathMLPresentationPrinter\n", - "\n", - "simplify_latex_rendering()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Expression trees" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SymPy expressions are built up from symbols and mathematical operations as follows:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "x, y, z = sp.symbols(\"x y z\")\n", - "expression = sp.sin(x * y) / 2 - x**2 + 1 / z\n", - "expression" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the back, SymPy represents these expressions as **trees**. There are a few ways to visualize this for this particular example:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sp.printing.tree.print_tree(expression, assumptions=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "src = sp.dotprint(\n", - " expression,\n", - " styles=[\n", - " (sp.Number, {\"color\": \"grey\", \"fontcolor\": \"grey\"}),\n", - " (sp.Symbol, {\"color\": \"royalblue\", \"fontcolor\": \"royalblue\"}),\n", - " ],\n", - ")\n", - "graphviz.Source(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Expression trees are powerful, because we can use them as templates for any human-readable presentation we are interested in. In fact, the LaTeX representation that we saw when constructing the expression was generated by SymPy's LaTeX printer." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "src = sp.latex(expression)\n", - "Markdown(f\"```latex\\n{src}\\n```\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{hint} SymPy expressions can serve as a template for generating code!\n", - ":::\n", - "\n", - "Here's a number of other representations:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-output", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def to_mathml(expr: sp.Expr) -> str:\n", - " printer = MathMLPresentationPrinter()\n", - " xml = printer._print(expr)\n", - " return xml.toprettyxml().replace(\"\\t\", \" \")\n", - "\n", - "\n", - "Markdown(\n", - " f\"\"\"\n", - "```python\n", - "# Python\n", - "{sp.pycode(expression)}\n", - "```\n", - "```cpp\n", - "// C++\n", - "{sp.cxxcode(expression, standard=\"c++17\")}\n", - "```\n", - "```fortran\n", - "! Fortran\n", - "{sp.fcode(expression).strip()}\n", - "```\n", - "```matlab\n", - "% Matlab / Octave\n", - "{sp.octave_code(expression)}\n", - "```\n", - "```julia\n", - "# Julia\n", - "{sp.julia_code(expression)}\n", - "```\n", - "```rust\n", - "// Rust\n", - "{sp.rust_code(expression)}\n", - "```\n", - "```xml\n", - "\n", - "{to_mathml(expression)}\n", - "```\n", - "\"\"\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Foldable expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The previous example is quite simple, but SymPy works just as well with huge expressions, as we will see in [Large expressions](#large-expressions). Before, though, let's have a look how to define these larger expressions in such a way that we can still read them. A nice solution is to define {class}`sp.Expr ` classes with the `@unevaluated` decorator (see [ComPWA/ampform#364](https://github.com/ComPWA/ampform/issues/364)). Here, we define a Chew-Mandelstam function $\\rho^\\text{CM}$ for $S$-waves. This function requires the definition of a break-up momentum $q$." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "@unevaluated(real=False)\n", - "class PhspFactorSWave(sp.Expr):\n", - " s: sp.Symbol\n", - " m1: sp.Symbol\n", - " m2: sp.Symbol\n", - " _latex_repr_ = R\"\\rho^\\text{{CM}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " cm = (\n", - " (2 * q / sp.sqrt(s))\n", - " * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n", - " - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " ) / (16 * sp.pi**2)\n", - " return 16 * sp.pi * sp.I * cm\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class BreakupMomentum(sp.Expr):\n", - " s: sp.Symbol\n", - " m1: sp.Symbol\n", - " m2: sp.Symbol\n", - " _latex_repr_ = R\"q\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt((s - (m1 + m2) ** 2) * (s - (m1 - m2) ** 2) / (s * 4))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now have a very clean mathematical representation of how the $\\rho^\\text{CM}$ function is defined in terms of $q$:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "s, m1, m2 = sp.symbols(\"s m1 m2\")\n", - "q_expr = BreakupMomentum(s, m1, m2)\n", - "ρ_expr = PhspFactorSWave(s, m1, m2)\n", - "Math(aslatex({e: e.evaluate() for e in [ρ_expr, q_expr]}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now, let's build up a more complicated expression that contains this phase space factor. Here, we use SymPy to derive a Breit-Wigner using a single-channel [$K$ matrix](https://doi.org/10.1002/andp.19955070504) {cite}`Chung:1995dx`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "I = sp.Identity(n=1)\n", - "K = sp.MatrixSymbol(\"K\", m=1, n=1)\n", - "ρ = sp.MatrixSymbol(\"rho\", m=1, n=1)\n", - "T = (I - sp.I * K * ρ).inv() * K\n", - "T" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T.as_explicit()[0, 0]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here we need to provide definitions for the matrix elements of $K$ and $\\rho$. A suitable choice is our phase space factor for $S$ waves we defined above:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "m0, Γ0, γ0 = sp.symbols(\"m0 Gamma0 gamma0\")\n", - "K_expr = (γ0**2 * m0 * Γ0) / (s - m0**2)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "substitutions = {\n", - " K[0, 0]: K_expr,\n", - " ρ[0, 0]: ρ_expr,\n", - "}\n", - "Math(aslatex(substitutions))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "And there we have it! After some [algebraic simplifications](https://docs.sympy.org/latest/tutorials/intro-tutorial/simplification.html), we get a Breit-Wigner with Chew-Mandelstam phase space factor for $S$ waves:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_expr = T.as_explicit().xreplace(substitutions)\n", - "BW_expr = T_expr[0, 0].simplify(doit=False)\n", - "BW_expr" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The expression tree now has a node that is 'folded':" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dot_style = [\n", - " (sp.Basic, {\"style\": \"filled\", \"fillcolor\": \"white\"}),\n", - " (sp.Atom, {\"color\": \"gray\", \"style\": \"filled\", \"fillcolor\": \"white\"}),\n", - " (sp.Symbol, {\"color\": \"dodgerblue1\"}),\n", - " (PhspFactorSWave, {\"color\": \"indianred2\"}),\n", - "]\n", - "dot = sp.dotprint(BW_expr, bgcolor=None, styles=dot_style)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "After unfolding, we get the full expression tree of fundamental mathematical operations:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dot = sp.dotprint(BW_expr.doit(), bgcolor=None, styles=dot_style)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Large expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Here, we import the large symbolic intensity expression that was used for [![10.1007/JHEP07(2023)228]()]() and see how well SymPy serialization performs on a much more complicated model." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DATA_DIR = Path(polarimetry.__file__).parent / \"lhcb\"\n", - "PARTICLES = load_particles(DATA_DIR / \"particle-definitions.yaml\")\n", - "MODEL = load_model(DATA_DIR / \"model-definitions.yaml\", PARTICLES, model_id=0)\n", - "unfolded_intensity_expr = perform_cached_doit(MODEL.full_expression)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "Markdown(\n", - " f\"\"\"\n", - "The model contains **{sp.count_ops(unfolded_intensity_expr):,d}** mathematical operations.\n", - "See [ComPWA/polarimetry#319](https://github.com/ComPWA/polarimetry/pull/319) for the origin\n", - "of this investigation.\n", - "\"\"\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Serialization with `srepr`" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SymPy expressions can directly be serialized to Python code as well, with the function [`srepr()`](https://docs.sympy.org/latest/modules/printing.html#sympy.printing.repr.srepr). For the full intensity expression, we can do so with:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "\n", - "eval_str = sp.srepr(unfolded_intensity_expr)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "n_nodes = sp.count_ops(unfolded_intensity_expr)\n", - "byt = len(eval_str.encode(\"utf-8\"))\n", - "mb = f\"{1e-6 * byt:.2f}\"\n", - "rendering = shorten(eval_str, placeholder=\" ...\", width=85)\n", - "src = f\"\"\"\n", - "This serializes the intensity expression of {n_nodes:,d} nodes\n", - "to a string of **{mb} MB**.\n", - "\n", - "```python\n", - "{rendering} {\")\" * (rendering.count(\"(\") - rendering.count(\")\"))}\n", - "```\n", - "\"\"\"\n", - "Markdown(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It is up to the user, however, to import the classes of each exported node before the string can be unparsed with [`eval()`](https://docs.python.org/3/library/functions.html#eval) (see [this comment](https://github.com/ComPWA/polarimetry/issues/20#issuecomment-1809840854))." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "raises-exception" - ] - }, - "outputs": [], - "source": [ - "imported_intensity_expr = eval(eval_str)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the case of this intensity expression, it is sufficient to import all definition from the main `sympy` module and the `Str` class. Optionally, the required `import` statements can be embedded into the string:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "exec_str = f\"\"\"\\\n", - "from sympy import *\n", - "from sympy.core.symbol import Str\n", - "\n", - "def get_intensity_function() -> Expr:\n", - " return {eval_str}\n", - "\"\"\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "exec_filename = Path(\"../_static/exported_intensity_model.py\")\n", - "with open(exec_filename, \"w\") as f:\n", - " f.write(exec_str)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "Markdown(f\"See [`{exec_filename.name}`]({exec_filename}) for the exported model.\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The parsing is then done with [`exec()`](https://docs.python.org/3/library/functions.html#exec) instead of the [`eval()`](https://docs.python.org/3/library/functions.html#eval) function:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "%%time\n", - "\n", - "exec(exec_str)\n", - "imported_intensity_expr = get_intensity_function()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Notice how the imported expression is **exactly the same** as the serialized one, including assumptions:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "assert imported_intensity_expr == unfolded_intensity_expr\n", - "assert hash(imported_intensity_expr) == hash(unfolded_intensity_expr)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Common sub-expressions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A problem is that the expression exported generated with [`srepr()`](https://docs.sympy.org/latest/modules/printing.html#sympy.printing.repr.srepr) is not human-readable in practice for large expressions. One way out may be to extract common components of the main expression with [Foldable expressions](#foldable-expressions). Another may be to use SymPy to [detect and collect common sub-expressions](https://docs.sympy.org/latest/modules/rewriting.html#common-subexpression-detection-and-collection)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "sub_exprs, common_expr = sp.cse(unfolded_intensity_expr, order=\"none\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "Math(sp.multiline_latex(sp.Symbol(\"I\"), common_expr[0], environment=\"eqnarray\"))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "Math(aslatex(dict(sub_exprs[:10])))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This already works quite well with {func}`sp.lambdify ` (without `cse=True`, this would takes minutes):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "\n", - "args = sorted(unfolded_intensity_expr.free_symbols, key=str)\n", - "_ = sp.lambdify(args, unfolded_intensity_expr, cse=True, dummify=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Still, as can be seen above, there are many sub-expressions that have exactly the same form. It would be better to find those expressions that have a similar structure, so that we can serialize them to functions or custom sub-definitions." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In SymPy, the equivalence between the expressions can be determined by the [`match()`](https://docs.sympy.org/latest/modules/core.html#sympy.core.basic.Basic.match) method using [`Wild`](https://docs.sympy.org/latest/modules/core.html#sympy.core.symbol.Wild) symbols. We therefore first have to make all symbols in the common sub-expressions 'wild'. In addition, in the case of this intensity expression, some of symbols are [indexed](https://docs.sympy.org/latest/modules/tensor/indexed.html) and need to be replaced first." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "pure_symbol_expr = unfolded_intensity_expr.replace(\n", - " query=lambda z: isinstance(z, sp.Indexed),\n", - " value=lambda z: sp.Symbol(sp.latex(z), **z.assumptions0),\n", - ")\n", - "sub_exprs, common_expr = sp.cse(pure_symbol_expr, order=\"none\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that for example the following two common sub-expressions are equivalent:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "Math(aslatex({k: v for i, (k, v) in enumerate(sub_exprs) if i in {5, 8}}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "[`Wild`](https://docs.sympy.org/latest/modules/core.html#sympy.core.symbol.Wild) symbols now allow us to find how these expressions relate to each other." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "is_symbol = lambda z: isinstance(z, sp.Symbol)\n", - "make_wild = lambda z: sp.Wild(z.name)\n", - "X = [x.replace(is_symbol, make_wild) for _, x in sub_exprs]\n", - "Math(aslatex(X[5].match(X[8])))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{hint}\n", - "This can be used to define functions for larger, common expression blocks.\n", - ":::" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/025.ipynb b/docs/report/025.ipynb deleted file mode 100644 index c5ffc368..00000000 --- a/docs/report/025.ipynb +++ /dev/null @@ -1,471 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "K-matrix" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Rotated square root cut\n", - "TR-025\n", - "^^^\n", - "Investigation of the branch cut in the two Riemann sheets of a square root and what happens if the cut is rotated around $z=0$.\n", - "+++\n", - "✅ [compwa.github.io#236](https://github.com/ComPWA/compwa.github.io/pull/236)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Rotating square root cuts" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.14.10 ipywidgets==8.1.1 plotly==5.18.0 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "import os\n", - "from typing import Any\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import plotly.graph_objects as go\n", - "import sympy as sp\n", - "from ampform.io import aslatex\n", - "from ampform.sympy import unevaluated\n", - "from IPython.display import Image, Math, display\n", - "from ipywidgets import FloatSlider, VBox, interactive_output\n", - "from plotly.colors import DEFAULT_PLOTLY_COLORS\n", - "from plotly.subplots import make_subplots\n", - "\n", - "STATIC_WEB_PAGE = {\"EXECUTE_NB\", \"READTHEDOCS\"}.intersection(os.environ)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso} [Lecture 17](https://compwa.github.io/strong2020-salamanca/lecture17.html) on collision theory of the [STRONG2020 HaSP School](https://indico.ific.uv.es/event/6803) by Miguel Albaladejo.\n", - "\n", - ":::\n", - "\n", - "There are multiple solutions for $x$ to the equation $y^2 = x$. The fact that we usually take $y = \\sqrt{x}$ with $\\sqrt{-1} = i$ to be 'the' solution to this equation is just a matter of convention. It would be more complete to represent the solution as a set of points in the complex plane, that is, the set $S = \\left\\{\\left(z, w\\right)\\in\\mathbb{C}^2 | w^2=z\\right\\}$. This is set forms a **Riemann surface** in $\\mathbb{C}^2$ space." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the figure below we see the Riemann surface of a square root in $\\mathbb{C}^2$ space. The $xy$ plane forms the complex domain $\\mathbb{C}$, the $z$ axis indicates the imaginary part of the Riemann surface and the color indicates the real part." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "resolution = 30\n", - "R, Θ = np.meshgrid(\n", - " np.linspace(0, 1, num=resolution),\n", - " np.linspace(-np.pi, +np.pi, num=resolution),\n", - ")\n", - "X = R * np.cos(Θ)\n", - "Y = R * np.sin(Θ)\n", - "Z = X + Y * 1j\n", - "T = np.sqrt(Z)\n", - "style = lambda t: dict(\n", - " cmin=-1,\n", - " cmax=+1,\n", - " colorscale=\"RdBu_r\",\n", - " surfacecolor=t.real,\n", - ")\n", - "fig = go.Figure([\n", - " go.Surface(x=X, y=Y, z=+T.imag, **style(+T), name=\"+√z\"),\n", - " go.Surface(x=X, y=Y, z=-T.imag, **style(-T), name=\"-√z\", showscale=False),\n", - "])\n", - "fig.update_traces(selector=0, colorbar=dict(title=\"Re ±√z\"))\n", - "fig.update_layout(\n", - " height=550,\n", - " margin=dict(l=0, r=0, t=30, b=0, pad=0),\n", - " title_text=\"Riemann surface of a square root\",\n", - " title_x=0.5,\n", - ")\n", - "fig.update_scenes(\n", - " camera_center=dict(z=-0.1),\n", - " camera_eye=dict(x=1.4, y=1.4, z=1.4),\n", - " xaxis_title=\"Re z\",\n", - " yaxis_title=\"Im z\",\n", - " zaxis_title=\"Im ±√z\",\n", - ")\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "From this figure it becomes clear that it is impossible to define one **single-valued** function that gives the solution to $w^2 = u$ is $w \\neq 0$. The familiar single-valued square root operation $\\sqrt{}$ covers only one segment, or **sheet**, of the Riemann surface and it is defined in such a way that $\\sqrt{-1}=i$. The other half of the surface is covered by $-\\sqrt{}$.\n", - "\n", - "Notice, however, that the sheets for the imaginary component of $\\sqrt{}$ are not smoothly connected at each point. The sign flips around $z\\in\\mathbb{R^-}$, because we have $\\sqrt{-1+0i}=-1$ and $\\sqrt{-1+0i}=+1$. We call this discontinuity in the Riemann sheet a **branch cut**." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "x = np.linspace(-1, 0, num=resolution // 2)\n", - "y = np.zeros(resolution // 2)\n", - "t = np.sqrt(x + 1e-8j)\n", - "T = np.sqrt(Z)\n", - "\n", - "C0 = DEFAULT_PLOTLY_COLORS[0]\n", - "C1 = DEFAULT_PLOTLY_COLORS[1]\n", - "\n", - "style = lambda color, legend: dict(\n", - " colorscale=[[0, color], [1, color]],\n", - " showlegend=legend,\n", - " showscale=False,\n", - " surfacecolor=np.ones(T.shape),\n", - ")\n", - "linestyle = dict(\n", - " line_color=\"crimson\",\n", - " line_showscale=False,\n", - " line_width=15,\n", - " mode=\"lines\",\n", - " name=\"Branch cut\",\n", - ")\n", - "\n", - "fig = make_subplots(\n", - " rows=1,\n", - " cols=2,\n", - " horizontal_spacing=0.01,\n", - " subplot_titles=(\"Re ±√z\", \"Im ±√z\"),\n", - " specs=[[{\"type\": \"surface\"}, {\"type\": \"surface\"}]],\n", - ")\n", - "fig.add_traces(\n", - " [\n", - " go.Surface(x=X, y=Y, z=+T.real, **style(C0, True), name=\"+√z\"),\n", - " go.Surface(x=X, y=Y, z=-T.real, **style(C1, True), name=\"-√z\"),\n", - " ],\n", - " cols=1,\n", - " rows=1,\n", - ")\n", - "fig.add_traces(\n", - " [\n", - " go.Surface(x=X, y=Y, z=+T.imag, **style(C0, False), name=\"+√z\"),\n", - " go.Surface(x=X, y=Y, z=-T.imag, **style(C1, False), name=\"-√z\"),\n", - " go.Scatter3d(x=x, y=y, z=-t.imag, **linestyle, showlegend=True),\n", - " go.Scatter3d(x=x, y=y, z=+t.imag, **linestyle, showlegend=False),\n", - " ],\n", - " cols=2,\n", - " rows=1,\n", - ")\n", - "ticks = dict(\n", - " tickvals=[-1, 0, +1],\n", - " ticktext=[\"-1\", \"0\", \"+1\"],\n", - ")\n", - "fig.update_layout(\n", - " height=400,\n", - " margin=dict(l=2, r=2, t=20, b=0, pad=0),\n", - ")\n", - "fig.update_scenes(\n", - " camera_center=dict(z=-0.1),\n", - " camera_eye=dict(x=1.4, y=1.4, z=1.4),\n", - " xaxis=dict(title=\"Re z\", **ticks),\n", - " yaxis=dict(title=\"Im z\", **ticks),\n", - ")\n", - "fig.update_scenes(selector=0, zaxis=dict(title=\"Re ±√z\", **ticks))\n", - "fig.update_scenes(selector=1, zaxis=dict(title=\"Im ±√z\", **ticks))\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "By definition, the branch cut of $\\sqrt{}$ is located at $\\mathbb{R}^-$. There is no requirement about this definition though: we can segment the Riemann surface in any way into two sheets, as long as the sheets remain single-valued. One option is to **rotate** the cut. With the following definition, we have a single-value square-root function, where the cut is rotated over an angle $\\phi$ around $z=0$." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class RotatedSqrt(sp.Expr):\n", - " z: Any\n", - " phi: Any = 0\n", - " _latex_repr_ = R\"\\sqrt[{phi}]{{{z}}}\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " z, phi = self.args\n", - " return sp.exp(-phi * sp.I / 2) * sp.sqrt(z * sp.exp(phi * sp.I))\n", - "\n", - "\n", - "z, phi = sp.symbols(\"z phi\")\n", - "expr = RotatedSqrt(z, phi)\n", - "Math(aslatex({expr: expr.doit(deep=False)}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following widget, we see what the new **rotated square root** looks like in the complex plane. The left panes show the imaginary part and the right side shows the real part. The upper figures show the value of the rotated square root on the real axis, $\\mathrm{Re}\\,z$." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "symbols = (z, phi)\n", - "func = sp.lambdify(symbols, expr.doit())\n", - "\n", - "mpl_fig, axes = plt.subplots(\n", - " figsize=(12, 8.5),\n", - " gridspec_kw=dict(\n", - " height_ratios=[1, 2],\n", - " width_ratios=[1, 1, 0.03],\n", - " ),\n", - " ncols=3,\n", - " nrows=2,\n", - ")\n", - "mpl_fig.canvas.toolbar_visible = False\n", - "mpl_fig.canvas.header_visible = False\n", - "mpl_fig.canvas.footer_visible = False\n", - "axes[0, 2].remove()\n", - "ax1re, ax2re = axes[:, 0]\n", - "ax1im, ax2im = axes[:, 1]\n", - "ax_bar = axes[1, 2]\n", - "ax1re.set_ylabel(f\"${sp.latex(expr)}$\")\n", - "ax1im.set_title(Rf\"$\\mathrm{{Im}}\\,{sp.latex(expr)}$\")\n", - "ax1re.set_title(Rf\"$\\mathrm{{Re}}\\,{sp.latex(expr)}$\")\n", - "ax2re.set_ylabel(R\"$\\mathrm{Im}\\,z$\")\n", - "for ax in (ax1im, ax1re):\n", - " ax.set_yticks([-1, -0.5, 0, +0.5, +1])\n", - " ax.set_yticklabels([\"-1\", R\"$-\\frac{1}{2}$\", \"0\", R\"$+\\frac{1}{2}$\", \"+1\"])\n", - "for ax in axes[:, :2].flatten():\n", - " ax.set_xlabel(R\"$\\mathrm{Re}\\,z$\")\n", - " ax.set_xticks([-1, 0, +1])\n", - " ax.set_xticklabels([\"-1\", \"0\", \"+1\"])\n", - " ax.set_yticks([-1, 0, +1])\n", - " ax.set_yticklabels([\"-1\", \"0\", \"+1\"])\n", - "for i, ax in enumerate((ax2im, ax2re)):\n", - " ax.axhline(0, c=f\"C{i}\", ls=\"dotted\", zorder=99)\n", - " ax.set_ylim(-1, +1)\n", - "\n", - "data = None\n", - "x = np.linspace(-1, +1, num=400)\n", - "X_mpl, Y_mpl = np.meshgrid(x, x)\n", - "Z_mpl = X_mpl + Y_mpl * 1j\n", - "\n", - "\n", - "def plot(phi):\n", - " global data\n", - " mpl_fig.suptitle(Rf\"$\\phi={phi / np.pi:.4g}\\pi$\")\n", - " t_mpl = func(x, phi)\n", - " T_mpl = func(Z_mpl, phi)\n", - " if data is None:\n", - " data = {\n", - " \"im\": ax1im.plot(x, t_mpl.imag, label=\"imag\", c=\"C0\", ls=\"dotted\")[0],\n", - " \"re\": ax1re.plot(x, t_mpl.real, label=\"real\", c=\"C1\", ls=\"dotted\")[0],\n", - " \"im2D\": ax2im.pcolormesh(X_mpl, Y_mpl, T_mpl.imag, cmap=plt.cm.coolwarm),\n", - " \"re2D\": ax2re.pcolormesh(X_mpl, Y_mpl, T_mpl.real, cmap=plt.cm.coolwarm),\n", - " }\n", - " else:\n", - " data[\"re\"].set_ydata(t_mpl.real)\n", - " data[\"im\"].set_ydata(t_mpl.imag)\n", - " data[\"im2D\"].set_array(T_mpl.imag)\n", - " data[\"re2D\"].set_array(T_mpl.real)\n", - " data[\"im2D\"].set_clim(vmin=-1, vmax=+1)\n", - " data[\"re2D\"].set_clim(vmin=-1, vmax=+1)\n", - " ax1im.set_ylim(-1.2, +1.2)\n", - " ax1re.set_ylim(-1.2, +1.2)\n", - " mpl_fig.canvas.draw_idle()\n", - "\n", - "\n", - "sliders = dict(\n", - " phi=FloatSlider(\n", - " min=-3 * np.pi,\n", - " max=+3 * np.pi,\n", - " step=np.pi / 8,\n", - " description=\"phi\",\n", - " value=-np.pi / 4,\n", - " ),\n", - ")\n", - "ui = VBox(tuple(sliders.values()))\n", - "output = interactive_output(plot, controls=sliders)\n", - "cbar = plt.colorbar(data[\"re2D\"], cax=ax_bar)\n", - "cbar.ax.set_xlabel(f\"${sp.latex(expr)}$\")\n", - "cbar.ax.set_yticks([-1, 0, +1])\n", - "mpl_fig.tight_layout()\n", - "display(ui, output)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "if STATIC_WEB_PAGE:\n", - " filename = \"025/rotated-sqrt-complex-plane.png\"\n", - " os.makedirs(os.path.dirname(filename), exist_ok=True)\n", - " mpl_fig.savefig(filename, dpi=200)\n", - " display(ui, Image(filename))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "The real part does not have a cut if $\\phi = 2\\pi n, n \\in \\mathbb{Z}$. The cut in the imaginary part disappears if $\\phi = \\pi + 2\\pi n$.\n", - ":::" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/026.ipynb b/docs/report/026.ipynb deleted file mode 100644 index f347f5db..00000000 --- a/docs/report/026.ipynb +++ /dev/null @@ -1,638 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Visualization of the Riemann sheets for the single-channel $T$ matrix with one resonance pole\n", - "TR-026\n", - "^^^\n", - "This report investigates and reproduces the Riemann sheets shown in [Fig. 50.1](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=2) and [50.2](https://pdg.lbl.gov/2023/reviews/rpp2023-rev-resonances.pdf#page=4) of the PDG. The lineshape parametrization is directly derived with the $K$-matrix formalism. The transition from the first physical sheet to the second unphysical sheet is derived using analytic continuation.\n", - "+++\n", - "🚧 [ampform#67](https://github.com/ComPWA/ampform/issues/67)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Single-channel Riemann sheets" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The $T$ function can be extended into the complex plane. This results in $2^n$ Riemann sheets for $n$ channels, each starting at the threshold $s_{thr}=(m_1+m_2)^{2}$ of the two final state particles, the so-called branching point of the respective channel going along the so-called branch cut along the real axis where the function is not uniquely defined to $+\\infty$. This choice of the direction of the brach cut is most commonly used in particle physics. The physical Riemann sheet is defined for positive imaginary part (1st quadrant of the complex plane) and the unphysical Riemann sheets are only defined for negative imaginary part (4th quadrant of the complex plane). For the single-channel case there are two Riemann sheets, one physical and one unphysical." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.15.0 plotly==5.18.0 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "from __future__ import annotations\n", - "\n", - "import warnings\n", - "from typing import Any\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import plotly.graph_objects as go\n", - "import sympy as sp\n", - "from ampform.io import aslatex\n", - "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import unevaluated\n", - "from IPython.display import Math\n", - "\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Phase space factor definitions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "@unevaluated(real=False)\n", - "class PhaseSpaceFactor(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho_{{{m1}, {m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt((s - ((m1 + m2) ** 2)) * (s - (m1 - m2) ** 2) / s**2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class PhaseSpaceCM(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class ChewMandelstam(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " return (\n", - " 1\n", - " / (16 * sp.pi**2)\n", - " * (\n", - " (2 * q / sp.sqrt(s))\n", - " * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n", - " - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " )\n", - " )\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class BreakupMomentum(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"q\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))\n", - "\n", - "\n", - "s, m1, m2 = sp.symbols(\"s m1 m2\")\n", - "rho_expr = PhaseSpaceFactor(s, m1, m2)\n", - "rho_cm_expr = PhaseSpaceCM(s, m1, m2)\n", - "cm_expr = ChewMandelstam(s, m1, m2)\n", - "q_expr = BreakupMomentum(s, m1, m2)\n", - "kallen = Kallen(*sp.symbols(\"x:z\"))\n", - "src = aslatex({\n", - " e: e.doit(deep=False) for e in [rho_expr, rho_cm_expr, cm_expr, q_expr, kallen]\n", - "})\n", - "Math(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## T matrix definition with K matrix" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The dynamical part of the scattering amplitude is calculated via $K$ matrix formalism. In this report the single-channel case with one resonance pole is assumed." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form" - }, - "outputs": [], - "source": [ - "# @title\n", - "n = 1\n", - "I = sp.Identity(n)\n", - "K = sp.MatrixSymbol(\"K\", n, n)\n", - "CM = sp.MatrixSymbol(R\"{\\rho_{cm}}\", n, n)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T1 = (I + sp.I * K * CM).inv() * K\n", - "T1" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T1_explicit = T1.as_explicit()\n", - "T1_explicit[0, 0]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "g0, m0 = sp.symbols(R\"g^{0} m0\")\n", - "k_expr = (g0**2) / (s - m0**2)\n", - "definitions_I = {\n", - " K[0, 0]: k_expr,\n", - " CM[0, 0]: PhaseSpaceCM(s, m1, m2),\n", - "}\n", - "Math(aslatex(definitions_I))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T1_expr = T1_explicit[0, 0].xreplace(definitions_I)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T1_expr.simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Calculation of the second Riemann sheet\n", - "Since the $T$ function is real below the branch cut it can be shown that the discontinuity above and below the threshold reads as:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "CM(s+i\\epsilon)-CM(s-i\\epsilon)= i\\rho -(-i\\rho) =2i\\rho\n", - "$$" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "when $\\epsilon$ goes to zero.
\n", - "Which leads to:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "CM^{-1}_{\\mathrm{II}}(s-i\\epsilon)= Re(CM^{-1}_{\\mathrm{I}}(s-i\\epsilon))-i\\rho+2i\\rho\n", - "$$" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "For the the Amplitude for the second sheet is defined as:\n", - "\n", - ":::{card}\n", - "$$\n", - "A^{-1}_{\\mathrm{II}}(s)= A^{-1}_{\\mathrm{I}}(s)-2i\\rho\n", - "$$\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "rho = sp.MatrixSymbol(\"rho\", n, n)\n", - "T2 = (T1.inv() + 2 * sp.I * rho).inv()\n", - "T2" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "definitions_II = {\n", - " **definitions_I,\n", - " rho[0, 0]: PhaseSpaceFactor(s, m1, m2),\n", - "}\n", - "T2_explicit = T2.as_explicit()\n", - "T2_expr = T2_explicit[0, 0].xreplace(definitions_II)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T2_expr.simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualization of the 2 dimensional lineshape" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Define numerical functions" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "symbols = sp.Tuple(s, m1, m2, m0, g0)\n", - "T1_func = sp.lambdify(symbols, T1_expr.doit())\n", - "T2_func = sp.lambdify(symbols, T2_expr.doit())" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "mystnb": { - "code_prompt_show": "Define meshgrid and parameter values" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "epsilon = 1e-5\n", - "x = np.linspace(0, 6, num=200)\n", - "y = np.linspace(epsilon, 1, num=100)\n", - "X, Y = np.meshgrid(x, y)\n", - "Zn = X - Y * 1j\n", - "Zp = X + Y * 1j\n", - "\n", - "values = {\n", - " m1: 0.9,\n", - " m2: 0.8,\n", - " m0: 3.1,\n", - " g0: 1.5,\n", - "}\n", - "args = eval(str(symbols[1:].xreplace(values)))\n", - "\n", - "T1n = T1_func(Zn**2, *args)\n", - "T1p = T1_func(Zp**2, *args)\n", - "\n", - "T2n = T2_func(Zn**2, *args)\n", - "T2p = T2_func(Zp**2, *args)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "%config InlineBackend.figure_formats = [\"svg\"]\n", - "\n", - "plt.rcParams.update({\"font.size\": 16})\n", - "fig, axes = plt.subplots(figsize=(15, 6), ncols=2, sharey=True)\n", - "ax1, ax2 = axes\n", - "for ax in axes:\n", - " ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n", - "ax1.set_ylabel(R\"$\\mathrm{Im}\\,T$\")\n", - "\n", - "ax1.plot(x, T1n[0].imag, label=R\"$T_\\mathrm{I}(s-0i)$\")\n", - "ax1.plot(x, T1p[0].imag, label=R\"$T_\\mathrm{I}(s+0i)$\")\n", - "ax1.set_title(f\"${sp.latex(rho_cm_expr)}$\")\n", - "ax1.set_title(R\"$T_\\mathrm{I}$\")\n", - "\n", - "ax2.plot(x, T2n[0].imag, label=R\"$T_\\mathrm{II}(s-0i)$\")\n", - "ax2.plot(x, T2p[0].imag, label=R\"$T_\\mathrm{II}(s+0i)$\")\n", - "ax2.set_title(R\"$T_\\mathrm{II}$\")\n", - "\n", - "for ax in axes:\n", - " ax.legend()\n", - "\n", - "fig.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The Amplitude for the second sheet is only defined for $s$ positive real part and negative complex part. It inherits the analytic structure of the phasespace factor $\\rho$ (the branch cut starting form zero and from $s=s_{thr}$ on the real axis). So it is only defined up to the closest branch cut which is in this case the cut at $s=s_{thr}$." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualization of the Riemann sheets" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "def sty(sheet_name: str) -> dict:\n", - " sheet_color = sheet_colors[sheet_name]\n", - " n_lines = 12\n", - " return dict(\n", - " cmin=-vmax,\n", - " cmax=+vmax,\n", - " colorscale=[[0, \"rgb(0, 0, 0)\"], [1, sheet_color]],\n", - " contours=dict(\n", - " x=dict(\n", - " show=True,\n", - " start=x.min(),\n", - " end=x.max(),\n", - " size=(x.max() - x.min()) / n_lines,\n", - " color=\"black\",\n", - " width=1,\n", - " ),\n", - " y=dict(\n", - " show=True,\n", - " start=-y.max(),\n", - " end=+y.max(),\n", - " size=(y.max() - y.min()) / (n_lines // 2),\n", - " color=\"black\",\n", - " width=1,\n", - " ),\n", - " ),\n", - " name=sheet_name,\n", - " opacity=0.4,\n", - " showscale=False,\n", - " )\n", - "\n", - "\n", - "vmax = 1.6\n", - "project = np.imag\n", - "sheet_colors = {\n", - " \"Physical (T1)\": \"blue\",\n", - " \"Unphysical (T2)\": \"red\",\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jp-MarkdownHeadingCollapsed": true, - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "Sp = go.Surface(x=X, y=Y, z=-T1p.imag, **sty(\"Physical (T1)\"))\n", - "Sn = go.Surface(x=X, y=-Y, z=-T2n.imag, **sty(\"Unphysical (T2)\"))\n", - "Sp.name = \"Physical sheet I\"\n", - "\n", - "s_thr = values[m1] + values[m2]\n", - "threshold_filter = x >= s_thr\n", - "lineshape = go.Scatter3d(\n", - " x=x[threshold_filter],\n", - " y=np.zeros(threshold_filter.shape),\n", - " z=project(-T1p[0])[threshold_filter],\n", - " line=dict(color=\"yellow\", width=10),\n", - " mode=\"lines\",\n", - " name=\"Lineshape\",\n", - ")\n", - "point = go.Scatter3d(\n", - " x=[s_thr],\n", - " y=[0],\n", - " z=[0],\n", - " mode=\"markers\",\n", - " marker=dict(color=\"black\", size=6),\n", - " name=\"Branch point\",\n", - ")\n", - "\n", - "fig = go.Figure(data=[Sn, Sp, lineshape, point])\n", - "fig.update_layout(\n", - " height=550,\n", - " margin=dict(l=0, r=0, t=30, b=0),\n", - " showlegend=True,\n", - " legend=dict(\n", - " orientation=\"v\",\n", - " xanchor=\"left\",\n", - " yanchor=\"top\",\n", - " x=0.05,\n", - " y=0.95,\n", - " font=dict(size=24),\n", - " ),\n", - " title_text=\"Im(T) with Chew-Mandelstam phase space factor\",\n", - " title_font=dict(size=28),\n", - " title=dict(y=0.989),\n", - ")\n", - "fig.update_scenes(\n", - " camera_center=dict(z=-0.2),\n", - " xaxis_title_text=\"Re √s\",\n", - " yaxis_title_text=\"Im √s\",\n", - " zaxis_title_text=\"Im T(s)\",\n", - " zaxis_range=[-vmax, +vmax],\n", - ")\n", - "\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The lineshape, the part that is observed within the experiment, is given as the intersection of the Riemann sheets with real plane. Also note that the second Riemann sheets transitions smoothly into the first one. " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{attention}\n", - ":name: Discontinuity\n", - "Not that the second Riemann sheet also inherits the singularity at $s=0$, as it is derived from the common phasespace factor.\n", - ":::" - ] - } - ], - "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.10.14" - }, - "orphan": true - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/027.ipynb b/docs/report/027.ipynb deleted file mode 100644 index f3ec3b97..00000000 --- a/docs/report/027.ipynb +++ /dev/null @@ -1,1142 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Visualization of the Riemann sheets for the two-channel $T$-matrix with one pole\n", - "TR-027\n", - "^^^\n", - "Following **[TR-026](./026.ipynb)**, the Riemann sheets for the amplitude calculated within the $K$-matrix formalism for the two-channel case are visualized. The method of transitioning from the first physical sheet to the unphysical sheets is extended to the two dimensional case using [Eur. Phys. J. C (2023) 83:850](https://juser.fz-juelich.de/record/1017534/files/s10052-023-11953-6.pdf) in order to visualize the third and the fourth unphysical sheet.\n", - "+++\n", - "🚧 [ampform#67](https://github.com/ComPWA/ampform/issues/67)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Coupled channel Riemann sheets\n" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.15.0 plotly==5.18.0 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "from __future__ import annotations\n", - "\n", - "import os\n", - "import warnings\n", - "from typing import Any\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import plotly.graph_objects as go\n", - "import sympy as sp\n", - "from ampform.io import aslatex\n", - "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import unevaluated\n", - "from IPython.display import Math, display\n", - "from ipywidgets import widgets as w\n", - "from plotly.colors import DEFAULT_PLOTLY_COLORS\n", - "from plotly.subplots import make_subplots\n", - "\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "if \"COLAB_RELEASE_TAG\" in os.environ:\n", - " import subprocess\n", - "\n", - " from google.colab import output\n", - "\n", - " output.enable_custom_widget_manager()\n", - " subprocess.run(\"pip install -q ipympl\".split(), check=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Expression definitions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "@unevaluated(real=False)\n", - "class PhaseSpaceFactor(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho_{{{m1}, {m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt((s - ((m1 + m2) ** 2)) * (s - (m1 - m2) ** 2) / s**2)\n", - "\n", - "\n", - "s, m1, m2 = sp.symbols(\"s m1 m2\")\n", - "rho_expr = PhaseSpaceFactor(s, m1, m2)\n", - "Math(aslatex({rho_expr: rho_expr.doit(deep=False)}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "@unevaluated(real=False)\n", - "class PhaseSpaceFactorKallen(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho_{{{m1}, {m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return 2 * BreakupMomentum(s, m1, m2) / sp.sqrt(s)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class PhaseSpaceCM(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class ChewMandelstam(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " return (\n", - " 1\n", - " / (16 * sp.pi**2)\n", - " * (\n", - " (2 * q / sp.sqrt(s))\n", - " * sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2))\n", - " - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " )\n", - " )\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class BreakupMomentum(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"q\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))\n", - "\n", - "\n", - "s, m1, m2 = sp.symbols(\"s m1 m2\")\n", - "rho_expr_kallen = PhaseSpaceFactorKallen(s, m1, m2)\n", - "rho_cm_expr = PhaseSpaceCM(s, m1, m2)\n", - "cm_expr = ChewMandelstam(s, m1, m2)\n", - "q_expr = BreakupMomentum(s, m1, m2)\n", - "kallen = Kallen(*sp.symbols(\"x:z\"))\n", - "Math(\n", - " aslatex({\n", - " e: e.doit(deep=False)\n", - " for e in [rho_expr_kallen, rho_cm_expr, cm_expr, q_expr, kallen]\n", - " })\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Riemann sheet I" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Matrix definition" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "class DiagonalMatrix(sp.DiagonalMatrix):\n", - " def _latex(self, printer, *args):\n", - " return printer._print(self.args[0])\n", - "\n", - "\n", - "n = 2\n", - "I = sp.Identity(n)\n", - "K = sp.MatrixSymbol(\"K\", n, n)\n", - "CM = DiagonalMatrix(sp.MatrixSymbol(R\"\\rho^\\Sigma\", n, n))\n", - "Math(aslatex({CM: CM.as_explicit()}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_I = (I - sp.I * K * CM).inv() * K\n", - "T_I" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_I_explicit = T_I.as_explicit()\n", - "T_I_explicit[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Parametrization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Symbol definitions" - }, - "tags": [ - "hide-cell" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "s = sp.Symbol(\"s\")\n", - "ma1 = sp.Symbol(\"m_{a1}\")\n", - "mb1 = sp.Symbol(\"m_{b1}\")\n", - "ma2 = sp.Symbol(\"m_{a2}\")\n", - "mb2 = sp.Symbol(\"m_{b2}\")\n", - "m0 = sp.Symbol(\"m0\")\n", - "w0 = sp.Symbol(\"Gamma0\")\n", - "g1 = sp.Symbol(R\"g^{0}_1\")\n", - "g2 = sp.Symbol(R\"g^{0}_2\")\n", - "symbols = sp.Tuple(s, ma1, mb1, ma2, mb2, m0, g1, g2)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "k_expr_00 = (g1 * g1 * m0) / (s - m0**2)\n", - "k_expr_10 = (g1 * g2 * m0) / (s - m0**2)\n", - "k_expr_11 = (g2 * g2 * m0) / (s - m0**2)\n", - "cm_expressions = {\n", - " K[0, 0]: k_expr_00,\n", - " K[1, 1]: k_expr_11,\n", - " K[0, 1]: k_expr_10,\n", - " K[1, 0]: k_expr_10,\n", - " CM[0, 0]: -PhaseSpaceCM(s, ma1, mb1),\n", - " CM[1, 1]: -PhaseSpaceCM(s, ma2, mb2),\n", - "}\n", - "Math(aslatex(cm_expressions))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_I_cm_expr = T_I_explicit.xreplace(cm_expressions)\n", - "T_I_cm_expr[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Sheets II, III, and IV" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the case of two channels, there are four Riemann sheets. The first sheet ([Sheet I](#riemann-sheet-i)) is physical and three unphysical ones. The physical sheet is calculated using the analytic solution of the Chew-Mandelstam function.\n", - "\n", - "$$\n", - "\\begin{eqnarray}\n", - "\\operatorname{Disc}_{\\mathrm{I,II}} T_K^{-1}\n", - "&=& 2 i\\left[\\begin{array}{rr}\\rho_1 & 0 \\\\ 0 & 0 \\end{array}\\right], \\\\\n", - "\\operatorname{Disc}_{\\mathrm{I,III}} T_K^{-1}\n", - "&=& 2 i\\left[\\begin{array}{rr}\\rho_1 & 0 \\\\ 0 & \\rho_2 \\end{array}\\right], \\\\\n", - "\\operatorname{Disc}_{\\mathrm{I,IV}} T_K^{-1}\n", - "&=& 2 i\\left[\\begin{array}{rr}0 & 0 \\\\ 0& \\rho_2 \\end{array}\\right].\n", - "\\end{eqnarray}\n", - "$$\n", - "\n", - "Depending on the centre-of-mass energy, different Riemann sheets connect smoothly to the physical one. Therefore, two cases are studied: one where the resonance mass is above the threshold of the second and first channel, and another where the resonance mass is between the threshold of the first and second channel." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "rho = DiagonalMatrix(sp.MatrixSymbol(\"rho\", n, n))\n", - "Math(aslatex({rho: rho.as_explicit()}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_II = (T_I.inv() + 2 * sp.I * rho).inv()\n", - "T_III = (T_I.inv() + 2 * sp.I * rho).inv()\n", - "T_IV = (-T_I.inv() - 2 * sp.I * rho).inv()\n", - "Math(aslatex([T_II, T_III, T_IV]))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "T_II_explicit = T_II.as_explicit()\n", - "T_II_explicit[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "T_III_explicit = T_III.as_explicit()\n", - "T_III_explicit[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "T_IV_explicit = T_IV.as_explicit()\n", - "T_IV_explicit[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "rho_expressions_II = {\n", - " **cm_expressions,\n", - " rho[0, 0]: PhaseSpaceFactor(s, ma1, mb1),\n", - " rho[1, 1]: 0,\n", - "}\n", - "rho_expressions_III = {\n", - " **cm_expressions,\n", - " rho[0, 0]: PhaseSpaceFactor(s, ma1, mb1),\n", - " rho[1, 1]: PhaseSpaceFactor(s, ma2, mb2),\n", - "}\n", - "rho_expressions_IV = {\n", - " **cm_expressions,\n", - " rho[0, 0]: 0,\n", - " rho[1, 1]: PhaseSpaceFactor(s, ma2, mb2),\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form" - }, - "outputs": [], - "source": [ - "# @title\n", - "T_II_rho_expr = T_II_explicit.xreplace(rho_expressions_II)\n", - "T_III_rho_expr = T_III_explicit.xreplace(rho_expressions_III)\n", - "T_IV_rho_expr = T_IV_explicit.xreplace(rho_expressions_IV)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_II_rho_expr[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "T_III_rho_expr[0, 0].simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualizations" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Lineshapes (real axis)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "%config InlineBackend.figure_formats = [\"svg\"]\n", - "\n", - "T_I_func = sp.lambdify(symbols, T_I_cm_expr[0, 0].doit())\n", - "T_II_func = sp.lambdify(symbols, T_II_rho_expr[0, 0].doit())\n", - "T_III_func = sp.lambdify(symbols, T_III_rho_expr[0, 0].doit())\n", - "T_IV_func = sp.lambdify(symbols, T_IV_rho_expr[0, 0].doit())\n", - "parameter_defaults1 = {\n", - " ma1: 1.0,\n", - " mb1: 1.5,\n", - " ma2: 1.5,\n", - " mb2: 2.0,\n", - " m0: 4.0,\n", - " g1: 0.7,\n", - " g2: 0.7,\n", - "}\n", - "parameter_defaults2 = {\n", - " **parameter_defaults1,\n", - " m0: 3.0,\n", - "}\n", - "args1 = eval(str(symbols[1:].xreplace(parameter_defaults1)))\n", - "args2 = eval(str(symbols[1:].xreplace(parameter_defaults2)))\n", - "\n", - "epsilon = 1e-5\n", - "x = np.linspace(0, 8, num=300)\n", - "y = np.linspace(epsilon, 1, num=100)\n", - "X, Y = np.meshgrid(x, y)\n", - "Zn = X - Y * 1j\n", - "Zp = X + Y * 1j\n", - "\n", - "T1n_res1 = T_I_func(Zn**2, *args1)\n", - "T1p_res1 = T_I_func(Zp**2, *args1)\n", - "\n", - "T2n_res1 = T_II_func(Zn**2, *args1)\n", - "T2p_res1 = T_II_func(Zp**2, *args1)\n", - "\n", - "T3n_res1 = T_III_func(Zn**2, *args1)\n", - "T3p_res1 = T_III_func(Zp**2, *args1)\n", - "\n", - "T4n_res1 = T_IV_func(Zn**2, *args1)\n", - "T4p_res1 = T_IV_func(Zp**2, *args1)\n", - "\n", - "T1n_res2 = T_I_func(Zn**2, *args2)\n", - "T1p_res2 = T_I_func(Zp**2, *args2)\n", - "\n", - "T2n_res2 = T_II_func(Zn**2, *args2)\n", - "T2p_res2 = T_II_func(Zp**2, *args2)\n", - "\n", - "T3n_res2 = T_III_func(Zn**2, *args2)\n", - "T3p_res2 = T_III_func(Zp**2, *args2)\n", - "\n", - "T4n_res2 = T_IV_func(Zn**2, *args2)\n", - "T4p_res2 = T_IV_func(Zp**2, *args2)\n", - "\n", - "fig, axes = plt.subplots(figsize=(11, 6), ncols=4, sharey=True)\n", - "ax1, ax2, ax3, ax4 = axes.flatten()\n", - "\n", - "ax1.plot(x, T1n_res1[0].imag, label=R\"$T_\\mathrm{I}(s-0i)$\")\n", - "ax1.plot(x, T1p_res1[0].imag, label=R\"$T_\\mathrm{I}(s+0i)$\")\n", - "ax1.set_title(f\"${sp.latex(rho_cm_expr)}$\")\n", - "ax1.set_title(R\"$T_\\mathrm{I}$\")\n", - "\n", - "ax2.plot(x, T2n_res1[0].imag, label=R\"$T_\\mathrm{II}(s-0i)$\")\n", - "ax2.plot(x, T2p_res1[0].imag, label=R\"$T_\\mathrm{II}(s+0i)$\")\n", - "ax2.set_title(R\"$T_\\mathrm{II}$\")\n", - "\n", - "ax3.plot(x, T3n_res1[0].imag, label=R\"$T_\\mathrm{III}(s-0i)$\")\n", - "ax3.plot(x, T3p_res1[0].imag, label=R\"$T_\\mathrm{III}(s+0i)$\")\n", - "ax3.set_title(R\"$T_\\mathrm{III}$\")\n", - "\n", - "ax4.plot(x, T4n_res1[0].imag, label=R\"$T_\\mathrm{III}(s-0i)$\")\n", - "ax4.plot(x, T4p_res1[0].imag, label=R\"$T_\\mathrm{IV}(s+0i)$\")\n", - "ax4.set_title(R\"$T_\\mathrm{III}$\")\n", - "\n", - "for ax in axes:\n", - " ax.legend()\n", - " ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n", - " ax.set_ylim(-1, +1)\n", - "ax1.set_ylabel(R\"$\\mathrm{Im}\\,T(s)$ (a.u.)\")\n", - "\n", - "fig.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Complex plane (2D)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "It can be shown that if the resonance mass is above both thresholds the third sheet connects smoothly to the first sheet. If the resonance mass is above the first and below the second threshold the second sheet transitions smoothly into the first sheet." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "full-width", - "scroll-input", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "%config InlineBackend.figure_formats = [\"png\"]\n", - "\n", - "fig, axes = plt.subplots(figsize=(12, 8), ncols=2, nrows=2, sharey=True)\n", - "ax1, ax2, ax3, ax4 = axes.flatten()\n", - "\n", - "for ax in axes.flatten():\n", - " ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n", - "for ax in axes[:, 0]:\n", - " ax.set_ylabel(R\"$\\mathrm{Im}\\,\\sqrt{s}$\")\n", - "\n", - "ax1.set_title(\"I and II\")\n", - "ax2.set_title(\"I and III\")\n", - "ax3.set_title(\"I and II\")\n", - "ax4.set_title(\"I and III\")\n", - "\n", - "T_max = 2\n", - "\n", - "style = dict(vmin=-T_max, vmax=+T_max, cmap=plt.cm.coolwarm)\n", - "mesh = ax1.pcolormesh(X, Y, T1p_res1.imag, **style)\n", - "ax1.pcolormesh(X, -Y, T2n_res1.imag, **style)\n", - "ax2.pcolormesh(X, +Y, T1p_res1.imag, **style)\n", - "ax2.pcolormesh(X, -Y, T3n_res1.imag, **style)\n", - "ax3.pcolormesh(X, +Y, T1p_res2.imag, **style)\n", - "ax3.pcolormesh(X, -Y, T2n_res2.imag, **style)\n", - "ax4.pcolormesh(X, +Y, T1p_res2.imag, **style)\n", - "ax4.pcolormesh(X, -Y, T3n_res2.imag, **style)\n", - "\n", - "s_thr1 = parameter_defaults1[ma1] + parameter_defaults1[mb1]\n", - "s_thr2 = parameter_defaults1[ma2] + parameter_defaults1[mb2]\n", - "linestyle = dict(ls=\"dotted\", lw=1)\n", - "for ax in axes.flatten():\n", - " ax.axhline(0, c=\"black\", **linestyle)\n", - " ax.axvline(s_thr1, c=\"C0\", **linestyle, label=R\"$\\sqrt{s_\\mathrm{thr1}}$\")\n", - " ax.axvline(s_thr2, c=\"C1\", **linestyle, label=R\"$\\sqrt{s_\\mathrm{thr2}}$\")\n", - "linestyle = dict(c=\"r\", ls=\"dotted\", label=R\"$m_\\mathrm{res}$\")\n", - "for ax in axes[0]:\n", - " ax.axvline(parameter_defaults1[m0], **linestyle)\n", - "for ax in axes[1]:\n", - " ax.axvline(parameter_defaults2[m0], **linestyle)\n", - "ax2.legend()\n", - "\n", - "fig.text(0.5, 0.93, R\"$s_{thr1} dict:\n", - " sheet_color = sheet_colors[sheet_name]\n", - " n_lines = 16\n", - " return dict(\n", - " cmin=-vmax,\n", - " cmax=+vmax,\n", - " colorscale=[[0, \"rgb(0, 0, 0)\"], [1, sheet_color]],\n", - " contours=dict(\n", - " x=dict(\n", - " show=True,\n", - " start=x.min(),\n", - " end=x.max(),\n", - " size=(x.max() - x.min()) / n_lines,\n", - " color=\"black\",\n", - " ),\n", - " y=dict(\n", - " show=True,\n", - " start=-y.max(),\n", - " end=+y.max(),\n", - " size=(y.max() - y.min()) / (n_lines // 2),\n", - " color=\"black\",\n", - " ),\n", - " ),\n", - " name=sheet_name,\n", - " opacity=0.4,\n", - " showscale=False,\n", - " )\n", - "\n", - "\n", - "vmax = 2.0\n", - "project = np.imag\n", - "sheet_colors = {\n", - " \"T1 (physical)\": \"blue\",\n", - " \"T2 (unphysical)\": \"red\",\n", - " \"T3 (unphysical)\": \"green\",\n", - " \"T4 (unphysical)\": \"yellow\",\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "Sp_I_res1 = go.Surface(x=X, y=+Y, z=T1p_res1.imag, **sty(\"T1 (physical)\"))\n", - "Sn_II_res1 = go.Surface(x=X, y=-Y, z=T2n_res1.imag, **sty(\"T2 (unphysical)\"))\n", - "Sn_III_res1 = go.Surface(x=X, y=-Y, z=T3n_res1.imag, **sty(\"T3 (unphysical)\"))\n", - "\n", - "Sp_I_res2 = go.Surface(x=X, y=+Y, z=T1p_res2.imag, **sty(\"T1 (physical)\"))\n", - "Sn_II_res2 = go.Surface(x=X, y=-Y, z=T2n_res2.imag, **sty(\"T2 (unphysical)\"))\n", - "Sn_III_res2 = go.Surface(x=X, y=-Y, z=T3n_res2.imag, **sty(\"T3 (unphysical)\"))\n", - "\n", - "thr1_filter = x >= s_thr1\n", - "thr2_filter = x >= s_thr2\n", - "\n", - "line_kwargs = dict(\n", - " line=dict(color=\"yellow\", width=8),\n", - " mode=\"lines\",\n", - " name=\"Lineshape\",\n", - ")\n", - "lineshape_res1_z = project(T1p_res1[0])\n", - "lineshape_res2_z = project(T1p_res2[0])\n", - "lineshape_res1 = go.Scatter3d(\n", - " x=x[thr1_filter],\n", - " y=np.zeros(thr1_filter.shape),\n", - " z=lineshape_res1_z[thr1_filter],\n", - " **line_kwargs,\n", - ")\n", - "lineshape_res2 = go.Scatter3d(\n", - " x=x[thr1_filter],\n", - " y=np.zeros(thr1_filter.shape),\n", - " z=lineshape_res2_z[thr1_filter],\n", - " **line_kwargs,\n", - ")\n", - "\n", - "point_kwargs = dict(\n", - " hoverinfo=\"text\",\n", - " marker=dict(color=DEFAULT_PLOTLY_COLORS[:2], size=6),\n", - " mode=\"markers\",\n", - " text=[\"threshold 1\", \"threshold 2\"],\n", - ")\n", - "thr_points_res1 = go.Scatter3d(\n", - " x=[s_thr1, s_thr2],\n", - " y=[0, 0],\n", - " z=[lineshape_res1_z[thr1_filter][0], lineshape_res1_z[thr2_filter][0]],\n", - " **point_kwargs,\n", - ")\n", - "thr_points_res2 = go.Scatter3d(\n", - " x=[s_thr1, s_thr2],\n", - " y=[0, 0],\n", - " z=[lineshape_res2_z[thr1_filter][0], lineshape_res2_z[thr2_filter][0]],\n", - " **point_kwargs,\n", - ")\n", - "\n", - "plotly_fig = make_subplots(\n", - " rows=2,\n", - " cols=2,\n", - " horizontal_spacing=0.01,\n", - " vertical_spacing=0.05,\n", - " specs=[\n", - " [{\"type\": \"surface\"}, {\"type\": \"surface\"}],\n", - " [{\"type\": \"surface\"}, {\"type\": \"surface\"}],\n", - " ],\n", - " subplot_titles=[\n", - " \"thr₁ < thr₂ < mᵣ\",\n", - " \"thr₁ < mᵣ < thr₂\",\n", - " ],\n", - ")\n", - "\n", - "# thr₁ < thr₂ < mᵣ\n", - "selector = dict(col=1, row=1)\n", - "plotly_fig.add_trace(Sp_I_res1, **selector)\n", - "plotly_fig.add_trace(Sn_III_res1, **selector)\n", - "plotly_fig.add_trace(lineshape_res1, **selector)\n", - "plotly_fig.add_trace(thr_points_res1, **selector)\n", - "selector = dict(col=1, row=2)\n", - "plotly_fig.add_trace(Sp_I_res1, **selector)\n", - "plotly_fig.add_trace(Sn_II_res1, **selector)\n", - "plotly_fig.add_trace(lineshape_res1, **selector)\n", - "plotly_fig.add_trace(thr_points_res1, **selector)\n", - "\n", - "# thr₁ < mᵣ < thr₂\n", - "selector = dict(col=2, row=1)\n", - "plotly_fig.add_trace(Sp_I_res2, **selector)\n", - "plotly_fig.add_trace(Sn_II_res2, **selector)\n", - "plotly_fig.add_trace(lineshape_res2, **selector)\n", - "plotly_fig.add_trace(thr_points_res2, **selector)\n", - "selector = dict(col=2, row=2)\n", - "plotly_fig.add_trace(Sp_I_res2, **selector)\n", - "plotly_fig.add_trace(Sn_III_res2, **selector)\n", - "plotly_fig.add_trace(lineshape_res2, **selector)\n", - "plotly_fig.add_trace(thr_points_res2, **selector)\n", - "\n", - "plotly_fig.update_layout(\n", - " height=600,\n", - " margin=dict(l=0, r=0, t=20, b=0),\n", - " showlegend=False,\n", - ")\n", - "\n", - "plotly_fig.update_scenes(\n", - " camera_center=dict(z=-0.1),\n", - " camera_eye=dict(x=1.4, y=1.4, z=1.4),\n", - " xaxis_range=(2.0, 5.0),\n", - " xaxis_title_text=\"Re √s\",\n", - " yaxis_title_text=\"Im √s\",\n", - " zaxis_title_text=\"Im T(s)\",\n", - " zaxis_range=[-vmax, +vmax],\n", - ")\n", - "plotly_fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Complex plane widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%matplotlib widget" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "cellView": "form", - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "# @title\n", - "sliders = dict(\n", - " g01=w.FloatSlider(\n", - " description=R\"$g^{0}_1$\",\n", - " value=0.5,\n", - " min=-2,\n", - " max=+2,\n", - " ),\n", - " g02=w.FloatSlider(\n", - " description=\"$g^{0}_2$\",\n", - " value=0.5,\n", - " min=-2,\n", - " max=+2,\n", - " ),\n", - " m0=w.FloatSlider(\n", - " description=\"$m_0$\",\n", - " value=4,\n", - " min=0,\n", - " max=+6,\n", - " ),\n", - " T_max=w.FloatSlider(\n", - " description=R\"$T_\\mathrm{max}$\",\n", - " value=1.0,\n", - " min=0.1,\n", - " max=6.0,\n", - " step=0.1,\n", - " ),\n", - ")\n", - "\n", - "\n", - "fig, axes = plt.subplots(\n", - " figsize=(11, 6),\n", - " ncols=2,\n", - " nrows=2,\n", - " sharex=True,\n", - " gridspec_kw={\"height_ratios\": [1, 2]},\n", - ")\n", - "fig.canvas.toolbar_visible = False\n", - "fig.canvas.header_visible = False\n", - "fig.canvas.footer_visible = False\n", - "ax1d1, ax1d2, ax2d1, ax2d2 = axes.flatten()\n", - "\n", - "for ax in axes[1]:\n", - " ax.set_xlabel(R\"$\\mathrm{Re}\\,\\sqrt{s}$\")\n", - "ax1d1.set_ylabel(\"Intensity (a.u.)\")\n", - "ax2d1.set_ylabel(R\"$\\mathrm{Im}\\,\\sqrt{s}$\")\n", - "\n", - "ax1d1.set_title(\"I and II\")\n", - "ax1d2.set_title(\"I and III\")\n", - "\n", - "R_color = \"C4\"\n", - "T1_color = \"C0\"\n", - "T2_color = \"C3\"\n", - "T3_color = \"C2\"\n", - "\n", - "\n", - "LINES = None\n", - "MESH = None\n", - "style = dict(cmap=plt.cm.coolwarm)\n", - "\n", - "\n", - "def plot(m0, g01, g02, T_max):\n", - " global LINES, MESH\n", - " local_args = args1[:-3] + (m0, g01, g02)\n", - " T1p_res1 = T_I_func(Zp**2, *local_args)\n", - " T2n_res1 = T_II_func(Zn**2, *local_args)\n", - " T3n_res1 = T_III_func(Zn**2, *local_args)\n", - " T1y = np.abs(T1p_res1[0]) ** 2\n", - " T2y = np.abs(T2n_res1[0]) ** 2\n", - " T3y = np.abs(T3n_res1[0]) ** 2\n", - " if MESH is None and LINES is None:\n", - " LINES = [\n", - " ax1d1.axvline(m0, c=R_color, ls=\"dashed\"),\n", - " ax1d2.axvline(m0, c=R_color, ls=\"dashed\"),\n", - " ax2d1.axvline(m0, c=R_color, ls=\"dashed\", label=R\"$m_\\mathrm{res}$\"),\n", - " ax2d2.axvline(m0, c=R_color, ls=\"dashed\", label=R\"$m_\\mathrm{res}$\"),\n", - " ax1d1.plot(x, T1y, c=T1_color, label=R\"$\\left|T_\\mathrm{I}\\right|^2$\")[0],\n", - " ax1d1.plot(\n", - " x, T2y, c=T2_color, label=R\"$\\left|T_\\mathrm{II}\\right|^2$\", ls=\"dotted\"\n", - " )[0],\n", - " ax1d2.plot(x, T1y, c=T1_color, label=R\"$\\left|T_\\mathrm{I}\\right|^2$\")[0],\n", - " ax1d2.plot(\n", - " x,\n", - " T3y,\n", - " c=T3_color,\n", - " label=R\"$\\left|T_\\mathrm{III}\\right|^2$\",\n", - " ls=\"dotted\",\n", - " )[0],\n", - " ]\n", - " MESH = [\n", - " ax2d1.pcolormesh(X, Y, T1p_res1.imag, **style),\n", - " ax2d1.pcolormesh(X, -Y, T2n_res1.imag, **style),\n", - " ax2d2.pcolormesh(X, +Y, T1p_res1.imag, **style),\n", - " ax2d2.pcolormesh(X, -Y, T3n_res1.imag, **style),\n", - " ]\n", - " else:\n", - " MESH[0].set_array(T1p_res1.imag)\n", - " MESH[1].set_array(T2n_res1.imag)\n", - " MESH[2].set_array(T1p_res1.imag)\n", - " MESH[3].set_array(T3n_res1.imag)\n", - " LINES[0].set_xdata(m0)\n", - " LINES[1].set_xdata(m0)\n", - " LINES[2].set_xdata(m0)\n", - " LINES[3].set_xdata(m0)\n", - " LINES[4].set_ydata(T1y)\n", - " LINES[5].set_ydata(T2y)\n", - " LINES[6].set_ydata(T1y)\n", - " LINES[7].set_ydata(T3y)\n", - " for mesh in MESH:\n", - " mesh.set_clim(-T_max, +T_max)\n", - " for ax in axes[0]:\n", - " ax.set_ylim(0, max(T1y) * 1.05)\n", - " fig.canvas.draw()\n", - "\n", - "\n", - "for ax in axes[:, 1]:\n", - " ax.set_yticks([])\n", - "for ax in axes[1]:\n", - " ax.axhline(0, c=\"black\", ls=\"dotted\", lw=1)\n", - "for ax in axes[0]:\n", - " ax.axvline(s_thr1, c=\"C0\", ls=\"dotted\", lw=1)\n", - " ax.axvline(s_thr2, c=\"C1\", ls=\"dotted\", lw=1)\n", - "for ax in axes[1]:\n", - " ax.axvline(s_thr1, c=\"C0\", label=R\"$\\sqrt{s_\\mathrm{thr1}}$\", ls=\"dotted\", lw=1)\n", - " ax.axvline(s_thr2, c=\"C1\", label=R\"$\\sqrt{s_\\mathrm{thr2}}$\", ls=\"dotted\", lw=1)\n", - "ax2d1.text(0.5, +0.7, R\"$T_\\mathrm{I}$\", color=T1_color, size=20)\n", - "ax2d1.text(0.5, -0.75, R\"$T_\\mathrm{II}$\", color=T2_color, size=20)\n", - "ax2d2.text(0.5, +0.7, R\"$T_\\mathrm{I}$\", color=T1_color, size=20)\n", - "ax2d2.text(0.5, -0.75, R\"$T_\\mathrm{III}$\", color=T3_color, size=20)\n", - "\n", - "output = w.interactive_output(plot, controls=sliders)\n", - "UI = w.VBox(list(sliders.values()))\n", - "fig.tight_layout()\n", - "ax1d1.legend()\n", - "ax1d2.legend()\n", - "ax2d2.legend()\n", - "display(output, UI)" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/028.ipynb b/docs/report/028.ipynb deleted file mode 100644 index 89add1d7..00000000 --- a/docs/report/028.ipynb +++ /dev/null @@ -1,548 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "PDG" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Example of how to query the PDG Python API for decay\n", - "TR-028\n", - "^^^\n", - "This report shows how to search all known decays in the PDG using [its new Python API](https://pdgapi.lbl.gov/doc) and search three-body decays that have three equal particles in the final state.\n", - "+++\n", - "🚧 [compwa.github.io#271](https://github.com/ComPWA/compwa.github.io/issues/271)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# PDG Python API: decay query" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q pdg==0.1.2 tqdm==4.66.5" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "import pdg\n", - "\n", - "PDG = pdg.connect()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This creates a [`PdgApi`](https://pdgapi.lbl.gov/doc/pdg.api.html#pdg.api.PdgApi) instance containing the following type of objects:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "{pdg.data.PdgLifetime,\n", - " pdg.data.PdgMass,\n", - " pdg.data.PdgProperty,\n", - " pdg.data.PdgWidth,\n", - " pdg.decay.PdgBranchingFraction,\n", - " pdg.particle.PdgParticleList}" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "{type(obj) for obj in PDG.get_all()}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In this example, we ask the question **which particles can decay to a final state with three equal particles?** For this, we use [`PdgBranchingFraction`](https://pdgapi.lbl.gov/doc/pdg.decay.html#pdg.decay.PdgBranchingFraction)s, which contain information about particle decays in their [`description`](https://pdgapi.lbl.gov/doc/pdg.data.html#pdg.data.PdgData.description):" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "pdg.decay.PdgBranchingFraction" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jpsi_decay = PDG.get(\"M070.313/2023\")\n", - "type(jpsi_decay)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "'J/psi(1S) --> rho(1700) pi --> pi+ pi- pi0'" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "jpsi_decay.description" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "So, let's pull all the decay descriptions from the PDG and do some clean up with {meth}`str.strip` and {obj}`set`:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "7243" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "from pdg.decay import PdgBranchingFraction\n", - "\n", - "all_decays = {obj for obj in PDG.get_all() if isinstance(obj, PdgBranchingFraction)}\n", - "decay_descriptions = {dec.description.strip() for dec in all_decays}\n", - "len(decay_descriptions)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To get more insight into the decay products, we create a new {obj}`set` of decay descriptions, but now describe each item as an initial state with a {obj}`tuple` of decay products. We again have to do a bit of cleaning here. The final state description sometimes contains digits, like `\"3pi0\"`, which we want to be rendered as `(\"pi0\", \"pi0\", \"pi0\")`.\n", - "\n", - "Note that we decay all state descriptions in the decay chain into account. For example,\n", - "```python\n", - "\"J/psi(1S) --> rho(1700) pi --> pi+ pi- pi0\"\n", - "```\n", - "\n", - "has two 'final' states:\n", - "```python\n", - "(\"rho(1700)\", \"pi\")\n", - "(\"pi+\", \"pi-\", \"pi0\")\n", - "```" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "7360" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "def create_final_state(description: str) -> tuple[str, ...]:\n", - " items = []\n", - " for particle in description.split():\n", - " particle = particle.strip()\n", - " if particle in {\"\", \",\"}:\n", - " continue\n", - " multiplier = particle[0]\n", - " if multiplier.isdigit():\n", - " particles = int(multiplier) * particle[1:]\n", - " items.extend(particles)\n", - " else:\n", - " items.append(particle)\n", - " return tuple(sorted(items))\n", - "\n", - "\n", - "decays: set[tuple[str, tuple[str, ...]]] = set()\n", - "for description in decay_descriptions:\n", - " initial_state, *final_states = description.split(\" --> \")\n", - " initial_state = initial_state.strip()\n", - " decays.update(\n", - " (initial_state, create_final_state(final_state)) for final_state in final_states\n", - " )\n", - "len(decays)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now selecting the three-body decays is an easy matter using filters on comprehensions." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "2181" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "three_body_decays = {\n", - " (initial_state, final_state)\n", - " for initial_state, final_state in decays\n", - " if len(final_state) == 3\n", - "}\n", - "len(three_body_decays)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "[('B0', ('K0S', 'K0S', 'K0S')),\n", - " ('B0', ('a', 'a', 'a')),\n", - " ('B_s()0', ('a', 'a', 'a')),\n", - " ('B_s()0', ('phi', 'phi', 'phi')),\n", - " ('J/psi(1S)', ('g', 'g', 'g')),\n", - " ('J/psi(1S)', ('gamma', 'gamma', 'gamma')),\n", - " ('Upsilon(1S)', ('g', 'g', 'g')),\n", - " ('Upsilon(2S)', ('g', 'g', 'g')),\n", - " ('Upsilon(3S)', ('g', 'g', 'g')),\n", - " ('Z', ('g', 'g', 'g')),\n", - " ('Z', ('gamma', 'gamma', 'gamma')),\n", - " ('a_1(1260)', ('pi0', 'pi0', 'pi0')),\n", - " ('a_1(1640)', ('pi', 'pi', 'pi')),\n", - " ('pi_1(1600)', ('pi', 'pi', 'pi')),\n", - " ('pi_2(1670)', ('pi0', 'pi0', 'pi0')),\n", - " ('psi(2S)', ('g', 'g', 'g'))]" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "equal_state_3body_decays = {\n", - " (initial_state, final_state)\n", - " for initial_state, final_state in three_body_decays\n", - " if len(set(final_state)) == 1\n", - "}\n", - "sorted(equal_state_3body_decays)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Finally, and optionally, we can filter out final states that are not well defined, such as `g g g`, by checking whether they are defined in the PDG database." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " J/psi(1S) → g g g\n", - " J/psi(1S) → gamma gamma gamma\n", - " Upsilon(1S) → g g g\n", - " Upsilon(2S) → g g g\n", - " Upsilon(3S) → g g g\n", - " Z → g g g\n", - " Z → gamma gamma gamma\n", - " psi(2S) → g g g\n" - ] - } - ], - "source": [ - "from pdg.errors import PdgAmbiguousValueError, PdgNoDataError\n", - "\n", - "for initial_state, final_state in sorted(equal_state_3body_decays):\n", - " try:\n", - " for name in (initial_state, *final_state):\n", - " PDG.get_particle_by_name(name)\n", - " except (PdgAmbiguousValueError, PdgNoDataError):\n", - " pass\n", - " else:\n", - " print(f\"{initial_state:>20} → {' '.join(final_state)}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{warning}\n", - "Not all final state in the [`description`](https://pdgapi.lbl.gov/doc/pdg.data.html#pdg.data.PdgData.description)s can be programmatically deciphered as individual particles. One could try to use [regular expressions](https://docs.python.org/3/howto/regex.html), but it's hard to cover all cases. Consider for instance the following case which contains $S$ and $D$ waves.\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "['a_1(1260) --> K^*(892) K',\n", - " 'a_1(1260) --> ( rho(1450) pi )(S-wave) , rho --> pi pi',\n", - " 'a_1(1260) --> f_2(1270) pi , f_2() --> pi pi',\n", - " 'a_1(1260) --> f_0(500) pi , f_0() --> pi pi',\n", - " 'a_1(1260) --> f_0(980) pi , f_0() --> pi pi',\n", - " 'a_1(1260) --> pi0 pi0 pi0',\n", - " 'a_1(1260) --> pi+ pi- pi0',\n", - " 'a_1(1260) --> ( rho(1450) pi )(D-wave) , rho --> pi pi',\n", - " 'a_1(1260) --> pi gamma',\n", - " 'a_1(1260) --> ( rho pi )(D-wave) , rho --> pi pi',\n", - " 'a_1(1260) --> 3 pi',\n", - " 'a_1(1260) --> f_0(1370) pi , f_0() --> pi pi',\n", - " 'a_1(1260) --> K K pi',\n", - " 'a_1(1260) --> ( rho pi )(S-wave) , rho --> pi pi']" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "[dec for dec in decay_descriptions if dec.startswith(\"a_1(1260)\")]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Additionally, not all decays seem to be included. Here is an attempt to find $J/\\psi \\to \\pi^0 \\pi^0 \\pi^0$." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/plain": [ - "['J/psi(1S) --> 2(pi+ pi- pi0)',\n", - " 'J/psi(1S) --> 2(pi+ pi- pi0) eta',\n", - " 'J/psi(1S) --> 2(pi+ pi-) 3pi0',\n", - " 'J/psi(1S) --> 2(pi+ pi-) pi0',\n", - " 'J/psi(1S) --> 3(pi+ pi-) pi0',\n", - " 'J/psi(1S) --> 4(pi+ pi-) pi0',\n", - " 'J/psi(1S) --> K+ K- pi0 pi0 pi0',\n", - " 'J/psi(1S) --> K0S K+- pi-+ pi+ pi-',\n", - " 'J/psi(1S) --> K0S K+- pi-+ pi0 pi0',\n", - " 'J/psi(1S) --> K^*(892)+ K0S pi- + c.c. --> K0S K0S pi+ pi-',\n", - " 'J/psi(1S) --> K^*(892)0 K- pi+ + c.c. --> K+ K- pi+ pi-',\n", - " 'J/psi(1S) --> K_2^*(1430)0 K- pi+ + c.c. --> K+ K- pi+ pi-',\n", - " 'J/psi(1S) --> a_2(1320)+ pi- pi0 + c.c --> 2 (pi+ pi- ) pi0',\n", - " 'J/psi(1S) --> a_2(1320)0 pi+ pi- --> 2 (pi+ pi- ) pi0',\n", - " 'J/psi(1S) --> eta pi+ pi- 3 pi0',\n", - " 'J/psi(1S) --> eta pi+ pi- pi0',\n", - " 'J/psi(1S) --> gamma pi+ pi- 2pi0',\n", - " 'J/psi(1S) --> omega 3 pi0',\n", - " 'J/psi(1S) --> omega pi+ pi+ pi- pi-',\n", - " 'J/psi(1S) --> omega pi+ pi- 2pi0',\n", - " 'J/psi(1S) --> omega pi+ pi- pi0',\n", - " 'J/psi(1S) --> omega pi0 --> pi+ pi- pi0',\n", - " 'J/psi(1S) --> p pbar pi+ pi- pi0',\n", - " 'J/psi(1S) --> phi f_1(1285) --> phi pi0 f_0(980) --> phi 3pi0',\n", - " 'J/psi(1S) --> phi f_1(1285) --> phi pi0 f_0(980) --> phi pi0 pi+ pi-',\n", - " 'J/psi(1S) --> phi pi0 f_0(980) --> phi pi0 p0 pi0',\n", - " 'J/psi(1S) --> phi pi0 f_0(980) --> phi pi0 pi+ pi-',\n", - " 'J/psi(1S) --> pi+ pi- 3pi0',\n", - " 'J/psi(1S) --> pi+ pi- 4 pi0',\n", - " 'J/psi(1S) --> pi+ pi- pi0',\n", - " 'J/psi(1S) --> pi+ pi- pi0 K+ K-',\n", - " 'J/psi(1S) --> pi+ pi- pi0 pi0 eta',\n", - " 'J/psi(1S) --> rho(1450) pi --> pi+ pi- pi0',\n", - " 'J/psi(1S) --> rho(1700) pi --> pi+ pi- pi0',\n", - " 'J/psi(1S) --> rho(2150) pi --> pi+ pi- pi0',\n", - " 'J/psi(1S) --> rho+ K+ K- pi- + c.c --> K+ K- pi+ pi- pi0',\n", - " 'J/psi(1S) --> rho+ rho- pi+ pi- pi0',\n", - " 'J/psi(1S) --> rho+- pi-+ pi+ pi- 2pi0',\n", - " 'J/psi(1S) --> rho+- pi-+ pi0 pi0',\n", - " 'J/psi(1S) --> rho_3(1690) pi --> pi+ pi- pi0']" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "import re\n", - "\n", - "sorted(\n", - " decay\n", - " for decay in decay_descriptions\n", - " if decay.startswith(\"J/psi\") and re.match(r\".*(3 ?pi|pi.*pi.*pi).*\", decay)\n", - ")" - ] - } - ], - "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.12.6" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/029.ipynb b/docs/report/029.ipynb deleted file mode 100644 index 406bc065..00000000 --- a/docs/report/029.ipynb +++ /dev/null @@ -1,542 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "sympy" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Definition of the normalized Blatt–Weisskopf form factor from Hankel functions of the first kind.\n", - "TR-029\n", - "^^^\n", - "This report investigates how to implement [ComPWA/ampform#417](https://github.com/ComPWA/ampform/issues/417), where it was suggested to define the 'normalized' Blatt–Weisskopf function $B_L^2(z)$ from a Hankel function of the first kind, $h_l^{(1)}$.\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Blatt–Weisskopf from Hankel function" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q ampform==0.15.1 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "from functools import cache\n", - "\n", - "import sympy as sp\n", - "from ampform.dynamics.phasespace import BreakupMomentumSquared\n", - "from ampform.io import aslatex\n", - "from ampform.sympy import unevaluated\n", - "from IPython.display import Math, display" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Introduction" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As of AmpForm [v0.15](https://github.com/ComPWA/ampform/releases/tag/0.15.1), the implementation of [`BlattWeisskopfSquared`](https://ampform.readthedocs.io/0.15.x/api/ampform.dynamics/#ampform.dynamics.BlattWeisskopfSquared) contains hard-coded polynomials, see implementation [here](https://github.com/ComPWA/ampform/blob/0.15.1/src/ampform/dynamics/__init__.py#L66-L134).\n", - "The motivation for this can be found in the citations mentioned in [its API documentation](https://ampform.readthedocs.io/0.15.x/api/ampform.dynamics/#ampform.dynamics.BlattWeisskopfSquared).\n", - "However, as noted by [@mmikhasenko](https://github.com/mmikhasenko) in [ComPWA/ampform#417](https://github.com/ComPWA/ampform/issues/417), the polynomials can be derived from the spherical[^1] Hankel functions of the first kind.\n", - "Von Hippel and Quigg[^2] derived a generalization of the centrifugal barrier factor $F_L$, also called form factor, that was introduced by {cite}`Blatt:1952ije`, showing that\n", - "\n", - "[^1]: See [this page](https://mathworld.wolfram.com/SphericalHankelFunctionoftheFirstKind.html) on Wolfram MathWorld for an explanation about the difference between $h_\\ell^{(1)}$ and $H_\\ell^{(1)}$.\n", - "[^2]: See {cite}`VonHippel:1972fg`, pp. 626 and 637, and a review by COMPASS, {cite}`Ketzer:2019wmd`, p. 31." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{math}\n", - "F_\\ell^2(z^2) = \\frac{1}{z^2\\left|h^{(1)}_\\ell\\left(z\\right)\\right|^2}\\,,\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "where $h_\\ell^{(1)}$ is a Hankel function of the first kind. They also noted that, if $z\\in\\mathbb{R}$," - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "h_\\ell^{(1)}(z) =\n", - " \\left(- i\\right)^{\\ell+1}\n", - " \\frac{e^{iz}}{z}\n", - " \\sum_{k=0}^\\ell\n", - " \\frac{(\\ell+k)!}{(\\ell-k)! \\, k!}\n", - " \\left(\\frac{i}{2z}\\right)^k.\n", - "$$ (hankel-sum)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we call $F_\\ell(z)$ the _unnormalized_ Blatt–Weisskopf form factor.\n", - "Following Chung and other resources (see e.g. {cite}`Chung:1995dx`, p. 415), AmpForm implements a unitless, _normalized_ Blatt–Weisskopf factor $B_L$, meaning that $B_L(1)=1$.[^3]\n", - "It can be defined in terms of $F_L$ as\n", - "\n", - "[^3]: We switch to notating angular momentum with $L$ instead of $\\ell$ here to indicate that we are talking about a normalized function here." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{math}\n", - "B_L^2(z^2)\n", - " = \\frac{F_L^2(z^2)}{F_L^2(1)}\n", - " = \\frac{\\left|h^{(1)}_L(1)\\right|^2}{z^2\\left|h^{(1)}_L(z)\\right|^2}\\,.\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "As of writing, AmpForm uses $z$ as argument in [`BlattWeisskopfSquared`](https://ampform.readthedocs.io/0.15.x/api/ampform.dynamics/#ampform.dynamics.BlattWeisskopfSquared).\n", - "This means we have to work with a square root and assume that $z \\geq 0$, meaning\n", - "\n", - "$$\n", - "B_L^2(z) = \\frac{\\left|h^{(1)}_L(1)\\right|^2}{z\\left|h^{(1)}_L\\left(\\sqrt{z}\\right)\\right|^2}\\,.\n", - "$$ (blatt-weisskopf)\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Hankel function of the first kind" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Built-in SymPy function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "SymPy offers a Hankel function of the first kind, [`scipy.special.hankel1`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.hankel1.html)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "z = sp.Symbol(\"z\", nonnegative=True, real=True)\n", - "ell = sp.Symbol(R\"\\ell\", integer=True, nonnegative=True)\n", - "sp.hankel1(ell, z)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This function is the general[^1] Hankel function $H_\\ell$ and the class does not offer algebraic simplifications for specific values or assumptions of $\\ell$ and $z$." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "display(\n", - " sp.hankel1(ell, z).doit(),\n", - " sp.hankel1(ell, 0).doit(),\n", - " sp.hankel1(0, z).doit(),\n", - " sp.hankel1(0, 0).doit(),\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Custom class definition" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To implement Equation {eq}`hankel-sum` for the _spherical_ Hankel function, we have to define a custom [`@unevaluated`](https://ampform.readthedocs.io/0.15.x/api/ampform.sympy/#ampform.sympy.unevaluated) expression class.\n", - "The following class evaluates to the sum given in Equation {eq}`hankel-sum`.\n", - "We introduce a special [`sympy.Sum`](https://docs.sympy.org/latest/modules/concrete.html#sympy.concrete.summations.Sum) class that does not 'unfold' on symbolic input for $\\ell$ if [`doit()`](https://docs.sympy.org/latest/modules/core.html#sympy.core.basic.Basic.doit) is called (see [](#nested-doit))." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated\n", - "class SphericalHankel1(sp.Expr):\n", - " l: sp.Symbol | int\n", - " z: sp.Symbol | float\n", - " _latex_repr_ = R\"h_{{{l}}}^{{(1)}}\\left({z}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " l, z = self.args\n", - " k = sp.Dummy(\"k\", integer=True, nonnegative=True)\n", - " return (\n", - " (-sp.I) ** (1 + l)\n", - " * (sp.exp(z * sp.I) / z)\n", - " * SymbolicSum(\n", - " sp.factorial(l + k)\n", - " / (sp.factorial(l - k) * sp.factorial(k))\n", - " * (sp.I / (2 * z)) ** k,\n", - " (k, 0, l),\n", - " )\n", - " )\n", - "\n", - "\n", - "class SymbolicSum(sp.Sum):\n", - " def doit(self, deep: bool = True, **kwargs) -> sp.Expr:\n", - " if _get_indices(self):\n", - " expression = self.args[0]\n", - " indices = self.args[1:]\n", - " return SymbolicSum(expression.doit(deep=deep, **kwargs), *indices)\n", - " return super().doit(deep=deep, **kwargs)\n", - "\n", - "\n", - "@cache\n", - "def _get_indices(expr: sp.Sum) -> set[sp.Symbol]:\n", - " free_symbols = set()\n", - " for index in expr.args[1:]:\n", - " free_symbols.update(index.free_symbols)\n", - " return {s for s in free_symbols if not isinstance(s, sp.Dummy)}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "h1lz = SphericalHankel1(ell, z)\n", - "Math(aslatex({h1lz: h1lz.doit()}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Indeed, the absolute squared value $\\left|h_\\ell^{(1)}\\right|^2$ results in a clean fraction of polynomials (after some algebraic [simplifications](https://docs.sympy.org/latest/tutorials/intro-tutorial/simplification.html))." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "exprs = [sp.Abs(h1lz.xreplace({ell: i})) ** 2 for i in range(3)]\n", - "Math(aslatex({e: e.doit().simplify() for e in exprs}))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "exprs = [sp.Abs(h1lz.xreplace({ell: i, z: 1})) ** 2 for i in range(3)]\n", - "Math(aslatex({e: e.doit() for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Normalized Blatt–Weisskopf form factor" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now have the required expression classes for re-implementing [`BlattWeisskopfSquared`](https://ampform.readthedocs.io/0.15.x/api/ampform.dynamics/#ampform.dynamics.BlattWeisskopfSquared) using Equation {eq}`blatt-weisskopf` (with $z$ as input, instead of $z^2$)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "@unevaluated\n", - "class BlattWeisskopfSquared(sp.Expr):\n", - " L: sp.Symbol | int\n", - " z: sp.Symbol | float\n", - " _latex_repr_ = R\"B^2_{{{L}}}\\left({z}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " L = self.L\n", - " z = sp.Dummy(\"z\", nonnegative=True, real=True)\n", - " expr = (\n", - " sp.Abs(SphericalHankel1(L, 1)) ** 2\n", - " / sp.Abs(SphericalHankel1(L, sp.sqrt(z))) ** 2\n", - " / z\n", - " )\n", - " if not L.free_symbols:\n", - " expr = expr.doit().simplify()\n", - " return expr.xreplace({z: self.z})" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "An explicit [`simplify()`](https://docs.sympy.org/latest/tutorials/intro-tutorial/simplification.html#simplify) is required in order to reproduce the polynomial form upon evaluation.\n", - "To make the simplification as fast as possible, it is done internally within `evaluate()` with $z$ as a dummy variable.\n", - "This is to avoid performing nested simplifications if $z$ is in itself an expression (see [](#nested-doit)).\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "L = sp.Symbol(\"L\", integer=True, nonnegative=True)\n", - "BL2 = BlattWeisskopfSquared(L, z)\n", - "Math(aslatex({BL2: BL2.doit(deep=False)}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Indeed the polynomials are exactly the same as the [original `BlattWeisskopfSquared`](https://ampform.readthedocs.io/0.15.x/api/ampform.dynamics/#ampform.dynamics.BlattWeisskopfSquared)!" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "exprs = [BL2.xreplace({L: i}) for i in range(9)]\n", - "Math(aslatex({e: e.doit() for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Nested doit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Eventually, the barrier factors take $z=q/q_R$, with $q$ the break-up momentum and $q_R$ an impact factor. Here it becomes crucial that only $\\left|h_\\ell^{(1)}(z)\\right|^2$ is simplified to a polynomial fraction, not $q$ itself. The break-up momentum does need to unfold though." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s, m1, m2, qR = sp.symbols(\"s m1 m2 q_R\", nonnegative=True)\n", - "q2 = BreakupMomentumSquared(s, m1, m2)\n", - "Math(aslatex({q2: q2.doit()}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Symbolic angular momentum" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "BlattWeisskopfSquared(L, z=q2 / qR**2).doit(deep=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "BlattWeisskopfSquared(L, z=q2 / qR**2).doit(deep=True)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Numeric angular momentum" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "BlattWeisskopfSquared(L=2, z=q2 / qR**2).doit(deep=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "BlattWeisskopfSquared(L=2, z=q2 / qR**2).doit(deep=True)" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/030.ipynb b/docs/report/030.ipynb deleted file mode 100644 index ed906bb3..00000000 --- a/docs/report/030.ipynb +++ /dev/null @@ -1,1345 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Sub-intensities of P-vector amplitude model\n", - "TR-030\n", - "^^^\n", - "Sub-intensity plots for a model with $P$-vector dynamics. Also includes an investigation of phases in a $P$-vector lineshape.\n", - "+++\n", - "🚧 [compwa.github.io#278](https://github.com/ComPWA/compwa.github.io/pull/278)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Sub-intensities of P vector" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q 'qrules[viz]==0.10.2' 'tensorwaves[jax,phsp]==0.4.12' ampform==0.15.4 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import logging\n", - "import os\n", - "import re\n", - "from collections import defaultdict\n", - "from functools import lru_cache\n", - "from typing import Any\n", - "\n", - "import ampform\n", - "import attrs\n", - "import graphviz\n", - "import jax.numpy as jnp\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n", - "from ampform.helicity import ParameterValues\n", - "from ampform.io import aslatex, improve_latex_rendering\n", - "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import perform_cached_doit, unevaluated\n", - "from attrs import define, field\n", - "from IPython.display import Math\n", - "from qrules.particle import Particle, ParticleCollection\n", - "from sympy import Abs\n", - "from tensorwaves.data import (\n", - " SympyDataTransformer,\n", - " TFPhaseSpaceGenerator,\n", - " TFUniformRealNumberGenerator,\n", - ")\n", - "from tensorwaves.function.sympy import create_parametrized_function\n", - "from tensorwaves.interface import DataSample, Function, ParametrizedFunction\n", - "\n", - "improve_latex_rendering()\n", - "logging.getLogger(\"absl\").setLevel(logging.ERROR)\n", - "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", - "plt.rc(\"font\", size=12)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Studied decay" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define N* resonances" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@lru_cache(maxsize=1)\n", - "def create_particle_database() -> ParticleCollection:\n", - " particles = qrules.load_default_particles()\n", - " for nstar in particles.filter(lambda p: p.name.startswith(\"N\")):\n", - " particles.remove(nstar)\n", - " particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)\n", - " particles += create_nstar(mass=1.92, width=0.6, parity=+1, spin=1.5, idx=2)\n", - " return particles\n", - "\n", - "\n", - "def create_nstar(\n", - " mass: float, width: float, parity: int, spin: float, idx: int\n", - ") -> Particle:\n", - " spin = sp.Rational(spin)\n", - " parity_symbol = \"⁺\" if parity > 0 else \"⁻\"\n", - " unicode_subscripts = list(\"₀₁₂₃₄₅₆₇₈₉\")\n", - " return Particle(\n", - " name=f\"N{unicode_subscripts[idx]}({spin}{parity_symbol})\",\n", - " latex=Rf\"N_{idx}({spin.numerator}/{spin.denominator}^-)\",\n", - " pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,\n", - " mass=mass,\n", - " width=width,\n", - " baryon_number=1,\n", - " charge=+1,\n", - " isospin=(0.5, +0.5),\n", - " parity=parity,\n", - " spin=1.5,\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "reaction = qrules.generate_transitions(\n", - " initial_state=\"J/psi(1S)\",\n", - " final_state=[\"eta\", \"p\", \"p~\"],\n", - " allowed_intermediate_particles=[\"N\"],\n", - " allowed_interaction_types=[\"strong\"],\n", - " formalism=\"helicity\",\n", - " particle_db=create_particle_database(),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Amplitude builder" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Dynamics builder with X symbols of J^PC channels" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@define\n", - "class DynamicsSymbolBuilder:\n", - " collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (\n", - " field(factory=lambda: defaultdict(set))\n", - " )\n", - "\n", - " def __call__(\n", - " self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n", - " ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", - " jp = render_jp(resonance)\n", - " charge = resonance.charge\n", - " if variable_pool.angular_momentum is not None:\n", - " L = sp.Rational(variable_pool.angular_momentum)\n", - " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}\")\n", - " else:\n", - " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}\")\n", - " self.collected_symbols[X].add((resonance, variable_pool))\n", - " parameter_defaults = {}\n", - " return X, parameter_defaults\n", - "\n", - "\n", - "def render_jp(particle: Particle) -> str:\n", - " spin = sp.Rational(particle.spin)\n", - " j = (\n", - " str(spin)\n", - " if spin.denominator == 1\n", - " else Rf\"\\frac{{{spin.numerator}}}{{{spin.denominator}}}\"\n", - " )\n", - " if particle.parity is None:\n", - " return f\"J={j}\"\n", - " p = \"-\" if particle.parity < 0 else \"+\"\n", - " return f\"J^P={{{j}}}^{{{p}}}\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "model_builder = ampform.get_builder(reaction)\n", - "model_builder.adapter.permutate_registered_topologies()\n", - "model_builder.config.scalar_initial_state_mass = True\n", - "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", - "create_dynamics_symbol = DynamicsSymbolBuilder()\n", - "for resonance in reaction.get_intermediate_particles():\n", - " model_builder.set_dynamics(resonance.name, create_dynamics_symbol)\n", - "model = model_builder.formulate()\n", - "model.intensity.cleanup()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "selected_amplitudes = {\n", - " k: v for i, (k, v) in enumerate(model.amplitudes.items()) if i == 0\n", - "}\n", - "Math(aslatex(selected_amplitudes, terms_per_line=1))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "src = R\"\\begin{array}{cll}\" \"\\n\"\n", - "for symbol, resonances in create_dynamics_symbol.collected_symbols.items():\n", - " src += Rf\" {symbol} \\\\\" \"\\n\"\n", - " for p, _ in resonances:\n", - " src += Rf\" {p.latex} & m={p.mass:g}\\text{{ GeV}} & \\Gamma={p.width:g}\\text{{ GeV}} \\\\\"\n", - " src += \"\\n\"\n", - "src += R\"\\end{array}\"\n", - "Math(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Dynamics parametrization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phasespace factor" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso}\n", - "**[TR-026](./026.ipynb)** and **[TR-027](./027.ipynb)** on analyticity and Riemann sheets.\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Expression classes for phase space factors" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated(real=False)\n", - "class PhaseSpaceCM(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class ChewMandelstam(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " return (\n", - " (2 * q / sp.sqrt(s))\n", - " * sp.log(Abs((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)))\n", - " - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " ) / (16 * sp.pi**2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class BreakupMomentum(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"q\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s, m1, m2 = sp.symbols(\"s m1 m2\", nonnegative=True)\n", - "exprs = [\n", - " PhaseSpaceCM(s, m1, m2),\n", - " ChewMandelstam(s, m1, m2),\n", - " BreakupMomentum(s, m1, m2),\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Relativistic Breit-Wigner" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "PARAMETERS_BW = dict(model.parameter_defaults)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_breit_wigner(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " m1 = variables.outgoing_state_mass1\n", - " m2 = variables.outgoing_state_mass2\n", - " ρ = PhaseSpaceCM(s, m1, m2)\n", - " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", - " Γ0 = [sp.Symbol(Rf\"\\Gamma_{{{p.latex}}}\") for p, _ in resonances]\n", - " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", - " expr = sum(\n", - " (β_ * m_ * Γ0_) / (m_**2 - s - m_ * Γ0_ * ρ) for m_, Γ0_, β_ in zip(m, Γ0, β)\n", - " )\n", - " for i, (resonance, _) in enumerate(resonances):\n", - " PARAMETERS_BW[β[i]] = 1 + 0j\n", - " PARAMETERS_BW[m[i]] = resonance.mass\n", - " PARAMETERS_BW[Γ0[i]] = resonance.width\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "dynamics_expressions_bw = {\n", - " symbol: formulate_breit_wigner(resonances)\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - "}\n", - "model_bw = attrs.evolve(\n", - " model,\n", - " parameter_defaults=ParameterValues({\n", - " **model.parameter_defaults,\n", - " **PARAMETERS_BW,\n", - " }),\n", - ")\n", - "Math(aslatex(dynamics_expressions_bw))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### $P$ vector" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "PARAMETERS_F = dict(model.parameter_defaults)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_k_matrix(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", - " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", - " expr = sum((g_**2) / (m_**2 - s) for m_, g_ in zip(m, g))\n", - " for i, (resonance, _) in enumerate(resonances):\n", - " PARAMETERS_F[m[i]] = resonance.mass\n", - " PARAMETERS_F[g[i]] = 1\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_p_vector(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", - " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", - " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", - " expr = sum((g_ * β_) / (m_**2 - s) for m_, g_, β_ in zip(m, g, β))\n", - " for i, (resonance, _) in enumerate(resonances):\n", - " PARAMETERS_F[β[i]] = 1 + 0j\n", - " PARAMETERS_F[m[i]] = resonance.mass\n", - " PARAMETERS_F[g[i]] = 1\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_f_vector(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " m1 = variables.outgoing_state_mass1\n", - " m2 = variables.outgoing_state_mass2\n", - " rho = PhaseSpaceCM(s, m1, m2)\n", - " K = formulate_k_matrix(resonances)\n", - " P = formulate_p_vector(resonances)\n", - " return (1 / (1 - rho * K)) * P" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dynamics_expressions_fvector = {\n", - " symbol: formulate_f_vector(resonances)\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - "}\n", - "model_fvector = attrs.evolve(\n", - " model,\n", - " parameter_defaults=ParameterValues({\n", - " **model.parameter_defaults,\n", - " **PARAMETERS_F,\n", - " }),\n", - ")\n", - "Math(aslatex(dynamics_expressions_fvector))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Create numerical functions" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Amplitude model function" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "full_expression_bw = perform_cached_doit(model_bw.expression).xreplace(\n", - " dynamics_expressions_bw\n", - ")\n", - "intensity_func_bw = create_parametrized_function(\n", - " expression=perform_cached_doit(full_expression_bw),\n", - " backend=\"jax\",\n", - " parameters=PARAMETERS_BW,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "full_expression_fvector = perform_cached_doit(model_fvector.expression).xreplace(\n", - " dynamics_expressions_fvector\n", - ")\n", - "intensity_func_fvector = create_parametrized_function(\n", - " expression=perform_cached_doit(full_expression_fvector),\n", - " backend=\"jax\",\n", - " parameters=PARAMETERS_F,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Dynamics function" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Breit–Wigner parametrization" - }, - "tags": [ - "full-width", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dynamics_expr_bw, *_ = dynamics_expressions_bw.values()\n", - "dynamics_expr_bw" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "F-vector parametrization" - }, - "tags": [ - "full-width", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dynamics_expr_fvector, *_ = dynamics_expressions_fvector.values()\n", - "dynamics_expr_fvector.simplify(doit=False)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dynamics_func_bw = create_parametrized_function(\n", - " expression=perform_cached_doit(dynamics_expr_bw),\n", - " backend=\"jax\",\n", - " parameters=model_bw.parameter_defaults,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "dynamics_func_fvector = create_parametrized_function(\n", - " expression=perform_cached_doit(dynamics_expr_fvector),\n", - " backend=\"jax\",\n", - " parameters=model_fvector.parameter_defaults,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Generate phase space sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "phsp_momenta = phsp_generator.generate(500_000, rng)\n", - "\n", - "ε = 1e-8\n", - "transformer = SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", - "phsp = transformer(phsp_momenta)\n", - "phsp = {k: v + ε * 1j if re.match(r\"^m_\\d\\d$\", k) else v for k, v in phsp.items()}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Update function parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "m_res1 = 1.82\n", - "m_res2 = 1.92\n", - "g_res1 = 1\n", - "g_res2 = 1" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "toy_parameters_bw = {\n", - " R\"m_{N_1(3/2^-)}\": m_res1,\n", - " R\"m_{N_2(3/2^-)}\": m_res2,\n", - " R\"\\Gamma_{N_1(3/2^-)}\": g_res1 / m_res1,\n", - " R\"\\Gamma_{N_2(3/2^-)}\": g_res2 / m_res2,\n", - "}\n", - "dynamics_func_bw.update_parameters(toy_parameters_bw)\n", - "intensity_func_bw.update_parameters(toy_parameters_bw)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "toy_parameters_fvector = {\n", - " R\"\\beta_{N_1(3/2^-)}\": 1 + 0j,\n", - " R\"\\beta_{N_2(3/2^-)}\": 1 + 0j,\n", - " R\"m_{N_1(3/2^-)}\": m_res1,\n", - " R\"m_{N_2(3/2^-)}\": m_res2,\n", - " R\"g_{N_1(3/2^-)}\": g_res1,\n", - " R\"g_{N_2(3/2^-)}\": g_res2,\n", - "}\n", - "dynamics_func_fvector.update_parameters(toy_parameters_fvector)\n", - "intensity_func_fvector.update_parameters(toy_parameters_fvector)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Plots" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Sub-intensities" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function for computing sub-intensities" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def compute_sub_intensity(\n", - " func: ParametrizedFunction,\n", - " input_data: DataSample,\n", - " resonances: list[str],\n", - " coupling_pattern: str,\n", - "):\n", - " original_parameters = dict(func.parameters)\n", - " negative_lookahead = f\"(?!{'|'.join(map(re.escape, resonances))})\"\n", - " # https://regex101.com/r/WrgGyD/1\n", - " pattern = rf\"^{coupling_pattern}({negative_lookahead}.)*$\"\n", - " set_parameters_to_zero(func, pattern)\n", - " array = func(input_data)\n", - " func.update_parameters(original_parameters)\n", - " return array\n", - "\n", - "\n", - "def set_parameters_to_zero(func: ParametrizedFunction, name_pattern: str) -> None:\n", - " toy_parameters = dict(func.parameters)\n", - " for par_name in func.parameters:\n", - " if re.match(name_pattern, par_name) is not None:\n", - " toy_parameters[par_name] = 0\n", - " func.update_parameters(toy_parameters)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "total_intensities_bw = intensity_func_bw(phsp)\n", - "sub_intensities_bw = {\n", - " p: compute_sub_intensity(\n", - " intensity_func_bw,\n", - " phsp,\n", - " resonances=[p.latex],\n", - " coupling_pattern=r\"\\\\beta\",\n", - " )\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " for p, _ in resonances\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "total_intensities_fvector = intensity_func_fvector(phsp)\n", - "sub_intensities_fvector = {\n", - " p: compute_sub_intensity(\n", - " intensity_func_fvector,\n", - " phsp,\n", - " resonances=[p.latex],\n", - " coupling_pattern=r\"\\\\beta\",\n", - " )\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " for p, _ in resonances\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function for plotting histograms with JAX" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def fast_histogram(\n", - " data: jnp.ndarray,\n", - " weights: jnp.ndarray | None = None,\n", - " bins: int = 100,\n", - " density: bool | None = None,\n", - " fill: bool = True,\n", - " ax=plt,\n", - " **plot_kwargs,\n", - ") -> None:\n", - " bin_values, bin_edges = jnp.histogram(\n", - " data,\n", - " bins=bins,\n", - " density=density,\n", - " weights=weights,\n", - " )\n", - " if fill:\n", - " bin_rights = bin_edges[1:]\n", - " ax.fill_between(bin_rights, bin_values, step=\"pre\", **plot_kwargs)\n", - " else:\n", - " bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2\n", - " ax.step(bin_mids, bin_values, **plot_kwargs)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(8, 5))\n", - "ax.set_xlim(2, 5)\n", - "ax.set_xlabel(R\"$m_{p\\eta}^{2}$ [GeV$^2$]\")\n", - "ax.set_ylabel(R\"Intensity [a. u.]\")\n", - "ax.set_yticks([])\n", - "\n", - "bins = 150\n", - "phsp_projection = np.real(phsp[\"m_01\"]) ** 2\n", - "fast_histogram(\n", - " phsp_projection,\n", - " weights=total_intensities_fvector,\n", - " alpha=0.2,\n", - " bins=bins,\n", - " color=\"hotpink\",\n", - " label=\"Full intensity $F$ vector\",\n", - " ax=ax,\n", - ")\n", - "fast_histogram(\n", - " phsp_projection,\n", - " weights=total_intensities_bw,\n", - " alpha=0.2,\n", - " bins=bins,\n", - " color=\"grey\",\n", - " label=\"Full intensity Breit-Wigner\",\n", - " ax=ax,\n", - ")\n", - "for i, (p, v) in enumerate(sub_intensities_fvector.items()):\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=v,\n", - " alpha=0.6,\n", - " bins=bins,\n", - " color=f\"C{i}\",\n", - " fill=False,\n", - " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV}}$ $F$ vector\",\n", - " linewidth=2,\n", - " ax=ax,\n", - " )\n", - "for i, (p, v) in enumerate(sub_intensities_bw.items()):\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=v,\n", - " alpha=0.6,\n", - " bins=bins,\n", - " color=f\"C{i}\",\n", - " fill=False,\n", - " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV^2}}$ Breit-Wigner\",\n", - " linestyle=\"dashed\",\n", - " ax=ax,\n", - " )\n", - "\n", - "ax.set_ylim(0, None)\n", - "fig.legend(loc=\"upper right\")\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Argand plots" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "ε = 1e-8\n", - "x = np.linspace(2, 5, num=400)\n", - "plot_data = {\"m_01\": np.sqrt(x) + ε * 1j}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "total_dynamics_bw = dynamics_func_bw(plot_data)\n", - "sub_dynamics_bw = {\n", - " p: compute_sub_intensity(\n", - " dynamics_func_bw,\n", - " plot_data,\n", - " resonances=[p.latex],\n", - " coupling_pattern=r\"\\\\beta\",\n", - " )\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " for p, _ in resonances\n", - "}\n", - "total_dynamics_fvector = dynamics_func_fvector(plot_data)\n", - "sub_dynamics_fvector = {\n", - " p: compute_sub_intensity(\n", - " dynamics_func_fvector,\n", - " plot_data,\n", - " resonances=[p.latex],\n", - " coupling_pattern=r\"\\\\beta\",\n", - " )\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " for p, _ in resonances\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "x1 = np.linspace(2.0, (m_res1**2 + m_res2**2) / 2, num=500)\n", - "x2 = np.linspace((m_res1**2 + m_res2**2) / 2, 5.0, num=500)\n", - "plot_data1 = {\"m_01\": np.sqrt(x1) + ε * 1j}\n", - "plot_data2 = {\"m_01\": np.sqrt(x2) + ε * 1j}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function definition for plotting Argand plots" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def plot_argand(\n", - " total_func: Function, sub_funcs: dict[Particle, Function], title: str\n", - ") -> None:\n", - " fig, axes = plt.subplots(1, 2, figsize=(10, 5), sharey=True)\n", - " fig.subplots_adjust(wspace=0.05)\n", - " fig.suptitle(title, y=0.99)\n", - " ax1, ax2 = axes\n", - " ax1.set_title(\"Total amplitude\")\n", - " ax2.set_title(\"Amplitude for resonance only\")\n", - " ax1.set_ylabel(R\"$\\text{Im}\\,F$\")\n", - " for ax in axes:\n", - " ax.axhline(0, color=\"black\", linewidth=0.5)\n", - " ax.axvline(0, color=\"black\", linewidth=0.5)\n", - " ax.set_xlabel(R\"$\\text{Re}\\,F$\")\n", - "\n", - " y1 = total_func(plot_data1)\n", - " ax1.plot(\n", - " y1.real,\n", - " y1.imag,\n", - " label=f\"Domain of {m_res1}-GeV resonance \",\n", - " color=\"C0\",\n", - " )\n", - " y2 = total_func(plot_data2)\n", - " ax1.plot(\n", - " y2.real,\n", - " y2.imag,\n", - " label=f\"Domain of {m_res2}-GeV resonance \",\n", - " color=\"C1\",\n", - " )\n", - " for i, (k, v) in enumerate(sub_funcs.items()):\n", - " ax2.plot(\n", - " v.real,\n", - " v.imag,\n", - " color=f\"C{i}\",\n", - " label=f\"Resonance at {k.mass} GeV $F$-vector\",\n", - " )\n", - "\n", - " ax1.legend(loc=\"upper left\")\n", - " fig.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "plot_argand(\n", - " dynamics_func_fvector,\n", - " sub_dynamics_fvector,\n", - " title=\"F vector\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "plot_argand(\n", - " dynamics_func_bw,\n", - " sub_dynamics_bw,\n", - " title=\"Breit-Wigner\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phase" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "total_phase_bw = np.angle(total_dynamics_bw)\n", - "total_phase_fvector = np.angle(total_dynamics_fvector)\n", - "sub_phase_bw = {p: np.angle(v) for p, v in sub_dynamics_bw.items()}\n", - "sub_phase_fvector = {p: np.angle(v) for p, v in sub_dynamics_fvector.items()}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function definition for plotting phases" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def plot_phases(\n", - " total_intensity_array: np.ndarray,\n", - " sub_intensity_arrays: dict[Particle, np.ndarray],\n", - " total_phase_array: np.ndarray,\n", - " sub_phase_arrays: dict[Particle, np.ndarray],\n", - " title: str,\n", - ") -> None:\n", - " fig, ax1 = plt.subplots(figsize=(10, 6))\n", - " ax1.set_title(title)\n", - " ax2 = ax1.twinx()\n", - " ax1.set_xlim(2.0, 5.0)\n", - " ax1.set_xlabel(R\"$m_{p\\eta}^{2}$ [GeV$^{2}$]\")\n", - " ax1.set_ylabel(\"Intensity [a. u.]\")\n", - " ax2.set_ylabel(\"Angle\")\n", - " ax1.set_yticks([])\n", - " ax2.set_ylim([-np.pi, +np.pi])\n", - " ax2.set_yticks([\n", - " -np.pi,\n", - " -np.pi / 2,\n", - " 0,\n", - " +np.pi / 2,\n", - " +np.pi,\n", - " ])\n", - " ax2.set_yticklabels([\n", - " R\"$-\\pi$\",\n", - " R\"$-\\frac{\\pi}{2}$\",\n", - " \"0\",\n", - " R\"$+\\frac{\\pi}{2}$\",\n", - " R\"$+\\pi$\",\n", - " ])\n", - " ax2.axhline(0, c=\"black\", lw=0.5)\n", - "\n", - " # Plot background histograms\n", - " phsp_projection = np.real(phsp[\"m_01\"]) ** 2\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=total_intensity_array,\n", - " bins=bins,\n", - " alpha=0.2,\n", - " color=\"gray\",\n", - " label=\"Full intensity\",\n", - " ax=ax1,\n", - " )\n", - " for i, (k, v) in enumerate(sub_intensity_arrays.items()):\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=v,\n", - " bins=bins,\n", - " alpha=0.2,\n", - " color=f\"C{i}\",\n", - " label=Rf\"Resonance at ${k.mass}\\,\\mathrm{{GeV}}$\",\n", - " ax=ax1,\n", - " )\n", - " ax1.set_ylim(0, None)\n", - "\n", - " # Plot phases\n", - " ax2.scatter(\n", - " x,\n", - " total_phase_array,\n", - " color=\"gray\",\n", - " label=\"Total Phase\",\n", - " s=18,\n", - " )\n", - " for i, (k, v) in enumerate(sub_phase_arrays.items()):\n", - " ax2.scatter(\n", - " x,\n", - " v,\n", - " alpha=0.5,\n", - " color=f\"C{i}\",\n", - " label=f\"Resonance at {k.mass} GeV\",\n", - " s=8,\n", - " )\n", - " ax2.axvline(k.mass**2, linestyle=\"dotted\", color=f\"C{i}\")\n", - "\n", - " # Add legends\n", - " fig.legend(bbox_to_anchor=(0.1, 0.9), loc=\"upper left\")\n", - " fig.tight_layout()\n", - " fig.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "plot_phases(\n", - " total_intensity_array=total_intensities_fvector,\n", - " sub_intensity_arrays=sub_intensities_fvector,\n", - " total_phase_array=total_phase_fvector,\n", - " sub_phase_arrays=sub_phase_fvector,\n", - " title=\"F vector\",\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "plot_phases(\n", - " total_intensity_array=total_intensities_bw,\n", - " sub_intensity_arrays=sub_intensities_bw,\n", - " total_phase_array=total_phase_bw,\n", - " sub_phase_arrays=sub_phase_bw,\n", - " title=\"Breit-Wigner\",\n", - ")" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/031.ipynb b/docs/report/031.ipynb deleted file mode 100644 index 3f36fe4d..00000000 --- a/docs/report/031.ipynb +++ /dev/null @@ -1,1383 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Single-channel amplitude model fit with $P$-vector dynamics\n", - "TR-031\n", - "^^^\n", - "Comparison between fit performance for an amplitude model with Breit–Wigner and $P$-vector dynamics. In both cases, data is generated with $P$-vector dynamics.\n", - "+++\n", - "🚧 [compwa.github.io#278](https://github.com/ComPWA/compwa.github.io/pull/278)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# P-vector model fit, single channel" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q 'qrules[viz]==0.10.2' 'tensorwaves[jax,phsp]==0.4.12' ampform==0.15.4 pandas==2.2.2 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import logging\n", - "import os\n", - "import re\n", - "from collections import defaultdict\n", - "from functools import lru_cache\n", - "from typing import Any\n", - "\n", - "import ampform\n", - "import attrs\n", - "import graphviz\n", - "import jax.numpy as jnp\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import pandas as pd\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n", - "from ampform.helicity import ParameterValues\n", - "from ampform.io import aslatex, improve_latex_rendering\n", - "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import perform_cached_doit, unevaluated\n", - "from attrs import define, field\n", - "from IPython.display import Math\n", - "from matplotlib import cm\n", - "from qrules.particle import Particle, ParticleCollection\n", - "from sympy import Abs\n", - "from tensorwaves.data import (\n", - " IntensityDistributionGenerator,\n", - " SympyDataTransformer,\n", - " TFPhaseSpaceGenerator,\n", - " TFUniformRealNumberGenerator,\n", - " TFWeightedPhaseSpaceGenerator,\n", - ")\n", - "from tensorwaves.estimator import UnbinnedNLL\n", - "from tensorwaves.function.sympy import create_parametrized_function\n", - "from tensorwaves.interface import DataSample, FitResult, ParametrizedFunction\n", - "from tensorwaves.optimizer import Minuit2\n", - "\n", - "improve_latex_rendering()\n", - "logging.getLogger(\"absl\").setLevel(logging.ERROR)\n", - "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", - "plt.rc(\"font\", size=12)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Studied decay" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define N* resonances" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@lru_cache(maxsize=1)\n", - "def create_particle_database() -> ParticleCollection:\n", - " particles = qrules.load_default_particles()\n", - " for nstar in particles.filter(lambda p: p.name.startswith(\"N\")):\n", - " particles.remove(nstar)\n", - " particles += create_nstar(mass=1.65, width=0.6, parity=-1, spin=0.5, idx=1)\n", - " particles += create_nstar(mass=1.75, width=0.6, parity=-1, spin=0.5, idx=2)\n", - " particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)\n", - " particles += create_nstar(mass=1.92, width=0.6, parity=+1, spin=1.5, idx=2)\n", - " return particles\n", - "\n", - "\n", - "def create_nstar(\n", - " mass: float, width: float, parity: int, spin: float, idx: int\n", - ") -> Particle:\n", - " spin = sp.Rational(spin)\n", - " parity_symbol = \"⁺\" if parity > 0 else \"⁻\"\n", - " unicode_subscripts = list(\"₀₁₂₃₄₅₆₇₈₉\")\n", - " return Particle(\n", - " name=f\"N{unicode_subscripts[idx]}({spin}{parity_symbol})\",\n", - " latex=Rf\"N_{idx}({spin.numerator}/{spin.denominator}^-)\",\n", - " pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,\n", - " mass=mass,\n", - " width=width,\n", - " baryon_number=1,\n", - " charge=+1,\n", - " isospin=(0.5, +0.5),\n", - " parity=parity,\n", - " spin=1.5,\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "reaction = qrules.generate_transitions(\n", - " initial_state=\"J/psi(1S)\",\n", - " final_state=[\"eta\", \"p\", \"p~\"],\n", - " allowed_intermediate_particles=[\"N\"],\n", - " allowed_interaction_types=[\"strong\"],\n", - " formalism=\"helicity\",\n", - " particle_db=create_particle_database(),\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dot = qrules.io.asdot(reaction, collapse_graphs=True)\n", - "graphviz.Source(dot)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Amplitude builder" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Dynamics builder with X symbols of J^PC channels" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@define\n", - "class DynamicsSymbolBuilder:\n", - " collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (\n", - " field(factory=lambda: defaultdict(set))\n", - " )\n", - "\n", - " def __call__(\n", - " self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n", - " ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", - " jp = render_jp(resonance)\n", - " charge = resonance.charge\n", - " if variable_pool.angular_momentum is not None:\n", - " L = sp.Rational(variable_pool.angular_momentum)\n", - " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}\")\n", - " else:\n", - " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}\")\n", - " self.collected_symbols[X].add((resonance, variable_pool))\n", - " parameter_defaults = {}\n", - " return X, parameter_defaults\n", - "\n", - "\n", - "def render_jp(particle: Particle) -> str:\n", - " spin = sp.Rational(particle.spin)\n", - " j = (\n", - " str(spin)\n", - " if spin.denominator == 1\n", - " else Rf\"\\frac{{{spin.numerator}}}{{{spin.denominator}}}\"\n", - " )\n", - " if particle.parity is None:\n", - " return f\"J={j}\"\n", - " p = \"-\" if particle.parity < 0 else \"+\"\n", - " return f\"J^P={{{j}}}^{{{p}}}\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "model_builder = ampform.get_builder(reaction)\n", - "model_builder.adapter.permutate_registered_topologies()\n", - "model_builder.config.scalar_initial_state_mass = True\n", - "model_builder.config.stable_final_state_ids = [0, 1, 2]\n", - "create_dynamics_symbol = DynamicsSymbolBuilder()\n", - "for resonance in reaction.get_intermediate_particles():\n", - " model_builder.set_dynamics(resonance.name, create_dynamics_symbol)\n", - "model = model_builder.formulate()\n", - "model.intensity.cleanup()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "selected_amplitudes = {\n", - " k: v for i, (k, v) in enumerate(model.amplitudes.items()) if i == 0\n", - "}\n", - "Math(aslatex(selected_amplitudes, terms_per_line=1))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "src = R\"\\begin{array}{cll}\" \"\\n\"\n", - "for symbol, resonances in create_dynamics_symbol.collected_symbols.items():\n", - " src += Rf\" {symbol} \\\\\" \"\\n\"\n", - " for p, _ in resonances:\n", - " src += Rf\" {p.latex} & m={p.mass:g}\\text{{ GeV}} & \\Gamma={p.width:g}\\text{{ GeV}} \\\\\"\n", - " src += \"\\n\"\n", - "src += R\"\\end{array}\"\n", - "Math(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Dynamics parametrization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phasespace factor" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso}\n", - "**[TR-026](./026.ipynb)** and **[TR-027](./027.ipynb)** on analyticity and Riemann sheets.\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Expression classes for phase space factors" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated(real=False)\n", - "class PhaseSpaceCM(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class ChewMandelstam(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " return (\n", - " (2 * q / sp.sqrt(s))\n", - " * sp.log(Abs((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)))\n", - " - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " ) / (16 * sp.pi**2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class BreakupMomentum(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"q\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s, m1, m2 = sp.symbols(\"s m1 m2\", nonnegative=True)\n", - "exprs = [\n", - " PhaseSpaceCM(s, m1, m2),\n", - " ChewMandelstam(s, m1, m2),\n", - " BreakupMomentum(s, m1, m2),\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Relativistic Breit-Wigner" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "PARAMETERS_BW = dict(model.parameter_defaults)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_breit_wigner(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " m1 = variables.outgoing_state_mass1\n", - " m2 = variables.outgoing_state_mass2\n", - " ρ = PhaseSpaceCM(s, m1, m2)\n", - " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", - " Γ0 = [sp.Symbol(Rf\"\\Gamma_{{{p.latex}}}\") for p, _ in resonances]\n", - " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", - " expr = sum(\n", - " (β_ * m_ * Γ0_) / (m_**2 - s - m_ * Γ0_ * ρ) for m_, Γ0_, β_ in zip(m, Γ0, β)\n", - " )\n", - " for i, (resonance, _) in enumerate(resonances):\n", - " PARAMETERS_BW[β[i]] = 1 + 0j\n", - " PARAMETERS_BW[m[i]] = resonance.mass\n", - " PARAMETERS_BW[Γ0[i]] = resonance.width\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "dynamics_expressions_bw = {\n", - " symbol: formulate_breit_wigner(resonances)\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - "}\n", - "model_bw = attrs.evolve(\n", - " model,\n", - " parameter_defaults=ParameterValues({\n", - " **model.parameter_defaults,\n", - " **PARAMETERS_BW,\n", - " }),\n", - ")\n", - "Math(aslatex(dynamics_expressions_bw))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### $P$ vector" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "PARAMETERS_F = dict(model.parameter_defaults)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_k_matrix(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", - " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", - " expr = sum((g_**2) / (m_**2 - s) for m_, g_ in zip(m, g))\n", - " for i, (resonance, _) in enumerate(resonances):\n", - " PARAMETERS_F[m[i]] = resonance.mass\n", - " PARAMETERS_F[g[i]] = 1\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_p_vector(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " g = [sp.Symbol(Rf\"g_{{{p.latex}}}\") for p, _ in resonances]\n", - " m = [sp.Symbol(Rf\"m_{{{p.latex}}}\") for p, _ in resonances]\n", - " β = [sp.Symbol(Rf\"\\beta_{{{p.latex}}}\") for p, _ in resonances]\n", - " expr = sum((g_ * β_) / (m_**2 - s) for m_, g_, β_ in zip(m, g, β))\n", - " for i, (resonance, _) in enumerate(resonances):\n", - " PARAMETERS_F[β[i]] = 1 + 0j\n", - " PARAMETERS_F[m[i]] = resonance.mass\n", - " PARAMETERS_F[g[i]] = 1\n", - " return expr" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def formulate_f_vector(\n", - " resonances: list[tuple[Particle, TwoBodyKinematicVariableSet]],\n", - ") -> sp.Expr:\n", - " (_, variables), *_ = resonances\n", - " s = variables.incoming_state_mass**2\n", - " m1 = variables.outgoing_state_mass1\n", - " m2 = variables.outgoing_state_mass2\n", - " rho = PhaseSpaceCM(s, m1, m2)\n", - " K = formulate_k_matrix(resonances)\n", - " P = formulate_p_vector(resonances)\n", - " return (1 / (1 - rho * K)) * P" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "dynamics_expressions_fvector = {\n", - " symbol: formulate_f_vector(resonances)\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - "}\n", - "model_fvector = attrs.evolve(\n", - " model,\n", - " parameter_defaults=ParameterValues({\n", - " **model.parameter_defaults,\n", - " **PARAMETERS_F,\n", - " }),\n", - ")\n", - "Math(aslatex(dynamics_expressions_fvector))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Create numerical functions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "full_expression_bw = perform_cached_doit(model_bw.expression).xreplace(\n", - " dynamics_expressions_bw\n", - ")\n", - "intensity_func_bw = create_parametrized_function(\n", - " expression=perform_cached_doit(full_expression_bw),\n", - " backend=\"jax\",\n", - " parameters=PARAMETERS_BW,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "full_expression_fvector = perform_cached_doit(model_fvector.expression).xreplace(\n", - " dynamics_expressions_fvector\n", - ")\n", - "intensity_func_fvector = create_parametrized_function(\n", - " expression=perform_cached_doit(full_expression_fvector),\n", - " backend=\"jax\",\n", - " parameters=PARAMETERS_F,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Generate phase space sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "rng = TFUniformRealNumberGenerator(seed=0)\n", - "phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=reaction.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in reaction.final_state.items()},\n", - ")\n", - "phsp_momenta = phsp_generator.generate(100_000, rng)\n", - "\n", - "ε = 1e-8\n", - "transformer = SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", - "phsp = transformer(phsp_momenta)\n", - "phsp = {k: v + ε * 1j if re.match(r\"^m_\\d\\d$\", k) else v for k, v in phsp.items()}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Update function parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "toy_parameters_bw = {\n", - " R\"m_{N_1(1/2^-)}\": 1.65,\n", - " R\"m_{N_2(1/2^-)}\": 1.75,\n", - " R\"m_{N_1(3/2^-)}\": 1.85,\n", - " R\"m_{N_2(3/2^-)}\": 1.9,\n", - " R\"\\Gamma_{N_1(1/2^-)}\": 1 / 1.65,\n", - " R\"\\Gamma_{N_2(1/2^-)}\": 1 / 1.75,\n", - " R\"\\Gamma_{N_1(3/2^-)}\": 1 / 1.85,\n", - " R\"\\Gamma_{N_2(3/2^-)}\": 1 / 1.9,\n", - "}\n", - "intensity_func_bw.update_parameters(toy_parameters_bw)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "toy_parameters_fvector = {\n", - " R\"\\beta_{N_1(1/2^-)}\": 1 + 0j,\n", - " R\"\\beta_{N_2(1/2^-)}\": 1 + 0j,\n", - " R\"\\beta_{N_1(3/2^-)}\": 1 + 0j,\n", - " R\"\\beta_{N_2(3/2^-)}\": 1 + 0j,\n", - " R\"m_{N_1(1/2^-)}\": 1.65,\n", - " R\"m_{N_2(1/2^-)}\": 1.75,\n", - " R\"m_{N_1(3/2^-)}\": 1.95,\n", - " R\"m_{N_2(3/2^-)}\": 1.9,\n", - " R\"g_{N_1(1/2^-)}\": 1.65,\n", - " R\"g_{N_2(1/2^-)}\": 1,\n", - " R\"g_{N_1(3/2^-)}\": 1,\n", - " R\"g_{N_2(3/2^-)}\": 1,\n", - "}\n", - "intensity_func_fvector.update_parameters(toy_parameters_fvector)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Plot sub-intensities" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function for computing sub-intensities" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def compute_sub_intensity(\n", - " func: ParametrizedFunction,\n", - " input_data: DataSample,\n", - " resonances: list[str],\n", - " coupling_pattern: str,\n", - "):\n", - " original_parameters = dict(func.parameters)\n", - " negative_lookahead = f\"(?!{'|'.join(map(re.escape, resonances))})\"\n", - " # https://regex101.com/r/WrgGyD/1\n", - " pattern = rf\"^{coupling_pattern}({negative_lookahead}.)*$\"\n", - " set_parameters_to_zero(func, pattern)\n", - " array = func(input_data)\n", - " func.update_parameters(original_parameters)\n", - " return array\n", - "\n", - "\n", - "def set_parameters_to_zero(func: ParametrizedFunction, name_pattern: str) -> None:\n", - " toy_parameters = dict(func.parameters)\n", - " for par_name in func.parameters:\n", - " if re.match(name_pattern, par_name) is not None:\n", - " toy_parameters[par_name] = 0\n", - " func.update_parameters(toy_parameters)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "total_intensities_bw = intensity_func_bw(phsp)\n", - "sub_intensities_bw = {\n", - " p: compute_sub_intensity(\n", - " intensity_func_bw,\n", - " phsp,\n", - " resonances=[p.latex],\n", - " coupling_pattern=r\"\\\\beta\",\n", - " )\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " for p, _ in resonances\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "total_intensities_fvector = intensity_func_fvector(phsp)\n", - "sub_intensities_fvector = {\n", - " p: compute_sub_intensity(\n", - " intensity_func_fvector,\n", - " phsp,\n", - " resonances=[p.latex],\n", - " coupling_pattern=r\"\\\\beta\",\n", - " )\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " for p, _ in resonances\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function for plotting histograms with JAX" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def fast_histogram(\n", - " data: jnp.ndarray,\n", - " weights: jnp.ndarray | None = None,\n", - " bins: int = 100,\n", - " density: bool | None = None,\n", - " fill: bool = True,\n", - " ax=plt,\n", - " **plot_kwargs,\n", - ") -> None:\n", - " bin_values, bin_edges = jnp.histogram(\n", - " data,\n", - " bins=bins,\n", - " density=density,\n", - " weights=weights,\n", - " )\n", - " if fill:\n", - " bin_rights = bin_edges[1:]\n", - " ax.fill_between(bin_rights, bin_values, step=\"pre\", **plot_kwargs)\n", - " else:\n", - " bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2\n", - " ax.step(bin_mids, bin_values, **plot_kwargs)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(8, 5))\n", - "ax.set_xlim(2, 5)\n", - "ax.set_xlabel(R\"$m_{p\\eta}^{2}$ [GeV$^2$]\")\n", - "ax.set_ylabel(R\"Intensity [a. u.]\")\n", - "ax.set_yticks([])\n", - "\n", - "bins = 150\n", - "phsp_projection = np.real(phsp[\"m_01\"]) ** 2\n", - "fast_histogram(\n", - " phsp_projection,\n", - " weights=total_intensities_fvector,\n", - " alpha=0.2,\n", - " bins=bins,\n", - " color=\"hotpink\",\n", - " label=\"Full intensity $F$ vector\",\n", - " ax=ax,\n", - ")\n", - "fast_histogram(\n", - " phsp_projection,\n", - " weights=total_intensities_bw,\n", - " alpha=0.2,\n", - " bins=bins,\n", - " color=\"grey\",\n", - " label=\"Full intensity Breit-Wigner\",\n", - " ax=ax,\n", - ")\n", - "for i, (p, v) in enumerate(sub_intensities_fvector.items()):\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=v,\n", - " alpha=0.6,\n", - " bins=bins,\n", - " color=f\"C{i}\",\n", - " fill=False,\n", - " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV}}$ $F$ vector\",\n", - " linewidth=2,\n", - " ax=ax,\n", - " )\n", - "for i, (p, v) in enumerate(sub_intensities_bw.items()):\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=v,\n", - " alpha=0.6,\n", - " bins=bins,\n", - " color=f\"C{i}\",\n", - " fill=False,\n", - " label=Rf\"Resonance at ${p.mass}\\,\\mathrm{{GeV^2}}$ Breit-Wigner\",\n", - " linestyle=\"dashed\",\n", - " ax=ax,\n", - " )\n", - "\n", - "ax.set_ylim(0, None)\n", - "fig.legend(loc=\"upper right\")\n", - "plt.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Weighted data with $F$ vector " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(6, 5))\n", - "fast_histogram(\n", - " phsp[\"m_01\"].real,\n", - " bins=100,\n", - " weights=np.real(intensity_func_fvector(phsp)),\n", - " ax=ax,\n", - ")\n", - "ax.set_xlabel(R\"$M^2\\left(\\eta p\\right)\\, \\mathrm{[(GeV/c)^2]}$\")\n", - "ax.set_ylabel(R\"Intensity [a.u.]\")\n", - "ax.set_ylim(0, None)\n", - "fig.tight_layout()\n", - "fig.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n", - " initial_state_mass=model.reaction_info.initial_state[-1].mass,\n", - " final_state_masses={i: p.mass for i, p in model.reaction_info.final_state.items()},\n", - ")\n", - "data_generator = IntensityDistributionGenerator(\n", - " domain_generator=weighted_phsp_generator,\n", - " function=intensity_func_fvector,\n", - " domain_transformer=transformer,\n", - ")\n", - "data_momenta = data_generator.generate(50_000, rng)\n", - "data = transformer(data_momenta)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(9, 4))\n", - "fast_histogram(\n", - " np.real(data[\"m_01\"]),\n", - " bins=200,\n", - " alpha=0.5,\n", - " density=True,\n", - " ax=ax,\n", - ")\n", - "mass_parameters = {p: v for p, v in toy_parameters_bw.items() if p.startswith(\"m_{\")}\n", - "evenly_spaced_interval = np.linspace(0, 1, num=len(mass_parameters))\n", - "colors = [cm.rainbow(x) for x in evenly_spaced_interval]\n", - "ax.set_xlabel(\"$m$ [GeV]\")\n", - "for (k, v), color in zip(mass_parameters.items(), colors):\n", - " ax.axvline(v, c=color, label=f\"${k}$\", ls=\"dotted\")\n", - "ax.set_ylim(0, None)\n", - "ax.legend()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Perform fit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Estimator definition" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "estimator_bw = UnbinnedNLL(\n", - " intensity_func_bw,\n", - " data=data,\n", - " phsp=phsp,\n", - " backend=\"jax\",\n", - ")\n", - "estimator_fvector = UnbinnedNLL(\n", - " intensity_func_fvector,\n", - " data=data,\n", - " phsp=phsp,\n", - " backend=\"jax\",\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Initial parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Functions for comparing model to data" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def indicate_masses(ax, intensity_func, ls: str, lw: float, typ: str):\n", - " mass_pars = {\n", - " k: v for k, v in intensity_func.parameters.items() if k.startswith(\"m_{\")\n", - " }\n", - " for i, (k, v) in enumerate(mass_pars.items()):\n", - " ax.axvline(v, c=f\"C{i}\", label=f\"${k}$ ({typ})\", ls=ls, lw=lw)\n", - "\n", - "\n", - "def compare_model(\n", - " variable_name,\n", - " data,\n", - " phsp,\n", - " function1,\n", - " function2,\n", - " bins=100,\n", - ") -> None:\n", - " intensities1 = function1(phsp)\n", - " intensities2 = function2(phsp)\n", - " fig, ax = plt.subplots(figsize=(11, 4))\n", - " fig.subplots_adjust(right=0.85, top=0.95)\n", - " ax.set_xlabel(R\"$m_{p\\eta}$ [GeV]\")\n", - " ax.set_ylabel(\"Intensity [a. u.]\")\n", - " ax.set_yticks([])\n", - " data_projection = np.real(data[variable_name])\n", - " fast_histogram(\n", - " data_projection,\n", - " bins=bins,\n", - " alpha=0.5,\n", - " label=\"data\",\n", - " density=True,\n", - " ax=ax,\n", - " )\n", - " phsp_projection = np.real(phsp[variable_name])\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=np.array(intensities1),\n", - " bins=bins,\n", - " fill=False,\n", - " color=\"red\",\n", - " label=\"Fit model with F vector\",\n", - " density=True,\n", - " ax=ax,\n", - " )\n", - " fast_histogram(\n", - " phsp_projection,\n", - " weights=np.array(intensities2),\n", - " bins=bins,\n", - " fill=False,\n", - " color=\"blue\",\n", - " label=\"Fit model with Breit-Wigner\",\n", - " density=True,\n", - " ax=ax,\n", - " )\n", - " ax.set_ylim(0, None)\n", - " indicate_masses(ax, function1, ls=\"dashed\", lw=1, typ=\"F vector\")\n", - " indicate_masses(ax, function2, ls=\"dotted\", lw=1, typ=\"Breit-Wigner\")\n", - " fig.legend(loc=\"outside upper right\")\n", - " fig.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "initial_parameters_beta = {\n", - " R\"\\beta_{N_2(1/2^-)}\": 1 + 0j,\n", - " R\"\\beta_{N_2(3/2^-)}\": 1 + 0j,\n", - "}\n", - "initial_parameters_masses = {\n", - " R\"m_{N_1(1/2^-)}\": 1.6,\n", - " R\"m_{N_2(1/2^-)}\": 1.7,\n", - " R\"m_{N_1(3/2^-)}\": 1.8,\n", - " R\"m_{N_2(3/2^-)}\": 1.93,\n", - "}\n", - "initial_parameters_bw = {\n", - " **initial_parameters_beta,\n", - " **initial_parameters_masses,\n", - " R\"\\Gamma_{N_1(1/2^-)}\": 1 / 1.6,\n", - " R\"\\Gamma_{N_2(1/2^-)}\": 1 / 1.65,\n", - " R\"\\Gamma_{N_1(3/2^-)}\": 1 / 1.85,\n", - " R\"\\Gamma_{N_2(3/2^-)}\": 1 / 1.93,\n", - "}\n", - "initial_parameters_fvector = {\n", - " **initial_parameters_beta,\n", - " **initial_parameters_masses,\n", - " R\"g_{N_1(1/2^-)}\": 1.6,\n", - " R\"g_{N_2(1/2^-)}\": 1,\n", - " R\"g_{N_1(3/2^-)}\": 1.0,\n", - " R\"g_{N_2(3/2^-)}\": 1.0,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "original_parameters_bw = dict(intensity_func_bw.parameters)\n", - "intensity_func_bw.update_parameters(initial_parameters_bw)\n", - "original_parameters_fvector = dict(intensity_func_fvector.parameters)\n", - "intensity_func_fvector.update_parameters(initial_parameters_fvector)\n", - "compare_model(\"m_01\", data, phsp, intensity_func_fvector, intensity_func_bw)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Optimize parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "minuit2 = Minuit2()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-output" - ] - }, - "outputs": [], - "source": [ - "fit_result_bw = minuit2.optimize(estimator_bw, initial_parameters_bw)\n", - "assert fit_result_bw.minimum_valid\n", - "fit_result_bw" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "scroll-output" - ] - }, - "outputs": [], - "source": [ - "fit_result_fvector = minuit2.optimize(estimator_fvector, initial_parameters_fvector)\n", - "assert fit_result_fvector.minimum_valid\n", - "fit_result_fvector" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "intensity_func_fvector.update_parameters(fit_result_fvector.parameter_values)\n", - "intensity_func_bw.update_parameters(fit_result_bw.parameter_values)\n", - "compare_model(\"m_01\", data, phsp, intensity_func_fvector, intensity_func_bw)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Fit result comparison" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Functions for inspecting fit result" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def compute_aic_bic(fit_result: FitResult) -> tuple[float, float]:\n", - " n_real_par = fit_result.count_number_of_parameters(complex_twice=True)\n", - " n_events = len(next(iter(data.values())))\n", - " log_likelihood = -fit_result.estimator_value\n", - " aic = 2 * n_real_par - 2 * log_likelihood\n", - " bic = n_real_par * np.log(n_events) - 2 * log_likelihood\n", - " return aic, bic\n", - "\n", - "\n", - "def compare_parameters(initial: dict, optimized: dict, expected: dict) -> pd.DataFrame:\n", - " parameters = sorted(set(initial) | set(optimized))\n", - " df = pd.DataFrame(\n", - " {\n", - " f\"${p}$\": (\n", - " f\"{initial[p]:.3g}\",\n", - " f\"{optimized[p]:.3g}\",\n", - " f\"{expected[p]:.3g}\",\n", - " f\"{100 * abs((optimized[p] - expected[p]) / expected[p]):.1f}%\",\n", - " )\n", - " for p in parameters\n", - " },\n", - " ).T\n", - " df.columns = (\"initial\", \"fit result\", \"expected\", \"deviation\")\n", - " return df" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### P vector" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "compute_aic_bic(fit_result_fvector)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "compare_parameters(\n", - " initial=initial_parameters_fvector,\n", - " optimized=fit_result_fvector.parameter_values,\n", - " expected=original_parameters_fvector,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Breit–Wigner" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "compute_aic_bic(fit_result_bw)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "compare_parameters(\n", - " initial=initial_parameters_bw,\n", - " optimized=fit_result_bw.parameter_values,\n", - " expected=original_parameters_bw,\n", - ")" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/032.ipynb b/docs/report/032.ipynb deleted file mode 100644 index a93e90bd..00000000 --- a/docs/report/032.ipynb +++ /dev/null @@ -1,1345 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "dynamics", - "K-matrix" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} Coupled-channel fit with $P$-vector dynamics for one single pole\n", - "TR-032\n", - "^^^\n", - "Illustration of how to formulate an amplitude model for two channels with P-vector dynamics. A combined fit is performed over the sum of the likelihood over both distributions. The example uses a single pole, but can easily be extended to multiple poles.\n", - "+++\n", - "🚧 [compwa.github.io#278](https://github.com/ComPWA/compwa.github.io/pull/278)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# P-vector fit comparison" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q 'qrules[viz]==0.10.2' 'tensorwaves[jax,phsp]==0.4.12' ampform==0.15.4 pandas==2.2.2 sympy==1.12" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import logging\n", - "import os\n", - "import re\n", - "from collections import defaultdict\n", - "from collections.abc import Iterable, Mapping\n", - "from functools import lru_cache\n", - "from itertools import product\n", - "from typing import Any\n", - "\n", - "import ampform\n", - "import attrs\n", - "import graphviz\n", - "import jax.numpy as jnp\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import pandas as pd\n", - "import qrules\n", - "import sympy as sp\n", - "from ampform.dynamics.builder import TwoBodyKinematicVariableSet\n", - "from ampform.helicity import HelicityModel\n", - "from ampform.io import aslatex, improve_latex_rendering\n", - "from ampform.kinematics.phasespace import Kallen\n", - "from ampform.sympy import perform_cached_doit, unevaluated\n", - "from attrs import define, field, frozen\n", - "from IPython.display import Math, display\n", - "from qrules.particle import Particle, ParticleCollection\n", - "from qrules.transition import ReactionInfo\n", - "from sympy import Abs\n", - "from sympy.matrices.expressions.matexpr import MatrixElement\n", - "from tensorwaves.data import (\n", - " IntensityDistributionGenerator,\n", - " SympyDataTransformer,\n", - " TFPhaseSpaceGenerator,\n", - " TFUniformRealNumberGenerator,\n", - " TFWeightedPhaseSpaceGenerator,\n", - ")\n", - "from tensorwaves.estimator import UnbinnedNLL\n", - "from tensorwaves.function.sympy import create_parametrized_function\n", - "from tensorwaves.interface import (\n", - " DataSample,\n", - " Estimator,\n", - " FitResult,\n", - " Function,\n", - " ParameterValue,\n", - ")\n", - "from tensorwaves.optimizer import Minuit2\n", - "\n", - "improve_latex_rendering()\n", - "logging.getLogger(\"absl\").setLevel(logging.ERROR)\n", - "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", - "plt.rc(\"font\", size=12)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-input" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Studied decay" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Define N* resonances" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@lru_cache(maxsize=1)\n", - "def create_particle_database() -> ParticleCollection:\n", - " particles = qrules.load_default_particles()\n", - " for nstar in particles.filter(lambda p: p.name.startswith(\"N\")):\n", - " particles.remove(nstar)\n", - " particles += create_nstar(mass=1.82, width=0.6, parity=+1, spin=1.5, idx=1)\n", - " return particles\n", - "\n", - "\n", - "def create_nstar(\n", - " mass: float, width: float, parity: int, spin: float, idx: int\n", - ") -> Particle:\n", - " spin = sp.Rational(spin)\n", - " parity_symbol = \"⁺\" if parity > 0 else \"⁻\"\n", - " unicode_subscripts = list(\"₀₁₂₃₄₅₆₇₈₉\")\n", - " return Particle(\n", - " name=f\"N{unicode_subscripts[idx]}({spin}{parity_symbol})\",\n", - " latex=Rf\"N_{idx}({spin.numerator}/{spin.denominator}^-)\",\n", - " pid=2024_05_00_00 + 100 * bool(parity + 1) + idx,\n", - " mass=mass,\n", - " width=width,\n", - " baryon_number=1,\n", - " charge=+1,\n", - " isospin=(0.5, +0.5),\n", - " parity=parity,\n", - " spin=1.5,\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FINAL_STATES: list[tuple[str, ...]] = [\n", - " [\"K0\", \"Sigma+\", \"p~\"],\n", - " [\"eta\", \"p\", \"p~\"],\n", - "]\n", - "REACTIONS: list[ReactionInfo] = [\n", - " qrules.generate_transitions(\n", - " initial_state=\"J/psi(1S)\",\n", - " final_state=final_state,\n", - " allowed_intermediate_particles=[\"N\"],\n", - " allowed_interaction_types=[\"strong\"],\n", - " formalism=\"helicity\",\n", - " particle_db=create_particle_database(),\n", - " )\n", - " for final_state in FINAL_STATES\n", - "]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "for reaction in REACTIONS:\n", - " src = qrules.io.asdot(reaction, collapse_graphs=True)\n", - " graph = graphviz.Source(src)\n", - " display(graph)\n", - " del reaction, src, graph" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Amplitude builder" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Dynamics builder with X symbols of J^PC channels" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@define\n", - "class DynamicsSymbolBuilder:\n", - " collected_symbols: set[sp.Symbol, tuple[Particle, TwoBodyKinematicVariableSet]] = (\n", - " field(factory=lambda: defaultdict(set))\n", - " )\n", - "\n", - " def __call__(\n", - " self, resonance: Particle, variable_pool: TwoBodyKinematicVariableSet\n", - " ) -> tuple[sp.Expr, dict[sp.Symbol, float]]:\n", - " jp = render_jp(resonance)\n", - " charge = resonance.charge\n", - " if variable_pool.angular_momentum is not None:\n", - " L = sp.Rational(variable_pool.angular_momentum)\n", - " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}^{{l={L}}}\")\n", - " else:\n", - " X = sp.Symbol(Rf\"X_{{{jp}, Q={charge:+d}}}\")\n", - " self.collected_symbols[X].add((resonance, variable_pool))\n", - " parameter_defaults = {}\n", - " return X, parameter_defaults\n", - "\n", - "\n", - "def render_jp(particle: Particle) -> str:\n", - " spin = sp.Rational(particle.spin)\n", - " j = (\n", - " str(spin)\n", - " if spin.denominator == 1\n", - " else Rf\"\\frac{{{spin.numerator}}}{{{spin.denominator}}}\"\n", - " )\n", - " if particle.parity is None:\n", - " return f\"J={j}\"\n", - " p = \"-\" if particle.parity < 0 else \"+\"\n", - " return f\"J^P={{{j}}}^{{{p}}}\"" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "MODELS: list[HelicityModel] = []\n", - "for reaction in REACTIONS:\n", - " builder = ampform.get_builder(reaction)\n", - " builder.adapter.permutate_registered_topologies()\n", - " builder.config.scalar_initial_state_mass = True\n", - " builder.config.stable_final_state_ids = [0, 1, 2]\n", - " create_dynamics_symbol = DynamicsSymbolBuilder()\n", - " for resonance in reaction.get_intermediate_particles():\n", - " builder.set_dynamics(resonance.name, create_dynamics_symbol)\n", - " MODELS.append(builder.formulate())\n", - " del builder, reaction, resonance" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "full-width" - ] - }, - "outputs": [], - "source": [ - "selected_amplitudes = {\n", - " k: v for i, (k, v) in enumerate(MODELS[0].amplitudes.items()) if i == 0\n", - "}\n", - "Math(aslatex(selected_amplitudes, terms_per_line=1))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "src = R\"\\begin{array}{cll}\" \"\\n\"\n", - "for symbol, resonances in create_dynamics_symbol.collected_symbols.items():\n", - " src += Rf\" {symbol} \\\\\" \"\\n\"\n", - " for p, _ in resonances:\n", - " src += Rf\" {p.latex} & m={p.mass:g}\\text{{ GeV}} & \\Gamma={p.width:g}\\text{{ GeV}} \\\\\"\n", - " src += \"\\n\"\n", - "src += R\"\\end{array}\"\n", - "Math(src)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Dynamics parametrization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phasespace factor" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso}\n", - "**[TR-026](./026.ipynb)** and **[TR-027](./027.ipynb)** on analyticity and Riemann sheets.\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Expression classes for phase space factors" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "@unevaluated(real=False)\n", - "class PhaseSpaceCM(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\rho^\\mathrm{{CM}}_{{{m1},{m2}}}\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return -16 * sp.pi * sp.I * ChewMandelstam(s, m1, m2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class ChewMandelstam(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"\\Sigma\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " q = BreakupMomentum(s, m1, m2)\n", - " return (\n", - " (2 * q / sp.sqrt(s))\n", - " * sp.log(Abs((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)))\n", - " - (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\n", - " ) / (16 * sp.pi**2)\n", - "\n", - "\n", - "@unevaluated(real=False)\n", - "class BreakupMomentum(sp.Expr):\n", - " s: Any\n", - " m1: Any\n", - " m2: Any\n", - " _latex_repr_ = R\"q\\left({s}\\right)\"\n", - "\n", - " def evaluate(self) -> sp.Expr:\n", - " s, m1, m2 = self.args\n", - " return sp.sqrt(Kallen(s, m1**2, m2**2)) / (2 * sp.sqrt(s))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "s, m1, m2 = sp.symbols(\"s m1 m2\", nonnegative=True)\n", - "exprs = [\n", - " PhaseSpaceCM(s, m1, m2),\n", - " ChewMandelstam(s, m1, m2),\n", - " BreakupMomentum(s, m1, m2),\n", - "]\n", - "Math(aslatex({e: e.doit(deep=False) for e in exprs}))" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### $K$-matrix formalism" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "n_channels = len(REACTIONS)\n", - "I = sp.Identity(n_channels)\n", - "K = sp.MatrixSymbol(\"K\", n_channels, n_channels)\n", - "P = sp.MatrixSymbol(\"P\", n_channels, 1)\n", - "F = sp.MatrixSymbol(\"F\", n_channels, 1)\n", - "rho = sp.MatrixSymbol(\"rho\", n_channels, n_channels)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Find decay products per channel" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def get_decay_products(reaction: ReactionInfo) -> DecayProducts:\n", - " some_transition, *_ = reaction.transitions\n", - " decay_product_ids = some_transition.topology.get_edge_ids_outgoing_from_node(1)\n", - " for transition in reaction.transitions:\n", - " if decay_product_ids != transition.topology.get_edge_ids_outgoing_from_node(1):\n", - " msg = \"Reaction contains multiple sub-systems\"\n", - " raise ValueError(msg)\n", - " child1_id, child2_id = sorted(decay_product_ids)\n", - " return DecayProducts(\n", - " child1=reaction.final_state[child1_id],\n", - " child2=reaction.final_state[child2_id],\n", - " )\n", - "\n", - "\n", - "@frozen\n", - "class DecayProducts:\n", - " child1: Particle\n", - " child2: Particle\n", - "\n", - " @property\n", - " def children(self) -> tuple[Particle, Particle]:\n", - " return self.child1, self.child2\n", - "\n", - "\n", - "DECAYS = tuple(get_decay_products(m.reaction_info) for m in MODELS)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "PARAMETERS_DEFAULTS = {}\n", - "for model in MODELS:\n", - " PARAMETERS_DEFAULTS.update(model.parameter_defaults)\n", - " del model\n", - "PARAMETERS_DEFAULTS = {\n", - " par: value\n", - " for par, value in PARAMETERS_DEFAULTS.items()\n", - " if not re.match(r\"^m_\\d+$\", par.name)\n", - "}" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### $K$-matrix parametrization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def formulate_k_matrix(\n", - " resonances: list[tuple[Particle, int]], n_channels: int\n", - ") -> dict[MatrixElement, sp.Expr]:\n", - " expressions = {}\n", - " for i, j in product(range(n_channels), range(n_channels)):\n", - " resonance_contributions = []\n", - " for res, _ in resonances:\n", - " s = sp.Symbol(\"m_01\", real=True) ** 2\n", - " g_Ri = sp.Symbol(Rf\"g_{{{res.latex},{i}}}\")\n", - " g_Rj = sp.Symbol(Rf\"g_{{{res.latex},{j}}}\")\n", - " m_R = sp.Symbol(Rf\"m_{{{res.latex}}}\")\n", - " parameter_defaults = {\n", - " m_R: res.mass,\n", - " g_Ri: 1,\n", - " g_Rj: 0.1,\n", - " }\n", - " PARAMETERS_DEFAULTS.update(parameter_defaults)\n", - " expr = (g_Ri * g_Rj) / (m_R**2 - s)\n", - " resonance_contributions.append(expr)\n", - " expressions[K[i, j]] = sum(resonance_contributions)\n", - " return expressions\n", - "\n", - "\n", - "K_expressions = formulate_k_matrix(resonances, n_channels=len(REACTIONS))\n", - "K_matrix = K.as_explicit()\n", - "K.as_explicit().xreplace(K_expressions)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### $P$-vector parametrization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def formulate_p_vector(\n", - " resonances: list[tuple[Particle, int]], n_channels: int\n", - ") -> dict[MatrixElement, sp.Expr]:\n", - " expressions = {}\n", - " for i in range(n_channels):\n", - " resonance_contributions = []\n", - " for res, _ in resonances:\n", - " s = sp.Symbol(\"m_01\", real=True) ** 2\n", - " g_Ri = sp.Symbol(Rf\"g_{{{res.latex},{i}}}\")\n", - " beta_R = sp.Symbol(Rf\"\\beta_{{{res.latex}}}\")\n", - " m_R = sp.Symbol(Rf\"m_{{{res.latex}}}\")\n", - " parameter_defaults = {\n", - " m_R: res.mass,\n", - " beta_R: 1 + 0j,\n", - " g_Ri: 1,\n", - " }\n", - " PARAMETERS_DEFAULTS.update(parameter_defaults)\n", - " expr = (beta_R * g_Ri) / (m_R**2 - s)\n", - " resonance_contributions.append(expr)\n", - " expressions[P[i, 0]] = sum(resonance_contributions)\n", - " return expressions\n", - "\n", - "\n", - "P_expressions = formulate_p_vector(resonances, n_channels=len(REACTIONS))\n", - "P_vector = P.as_explicit()\n", - "P.as_explicit().xreplace(P_expressions)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Phase space factor parametrization" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def formulate_phsp_factor_matrix(n_channels: int) -> dict[sp.MatrixElement, sp.Expr]:\n", - " expressions = {}\n", - " for i in range(n_channels):\n", - " for j in range(n_channels):\n", - " if i == j:\n", - " m_a_i = sp.Symbol(Rf\"m_{{0,{i}}}\")\n", - " m_b_i = sp.Symbol(Rf\"m_{{1,{i}}}\")\n", - " s = sp.Symbol(\"m_01\", real=True) ** 2\n", - " rho_i = PhaseSpaceCM(s, m_a_i, m_b_i)\n", - " expressions[rho[i, j]] = rho_i\n", - " parameter_defaults = {\n", - " m_a_i: DECAYS[i].child1.mass,\n", - " m_b_i: DECAYS[i].child2.mass,\n", - " }\n", - " PARAMETERS_DEFAULTS.update(parameter_defaults)\n", - " else:\n", - " expressions[rho[i, j]] = 0\n", - " return expressions\n", - "\n", - "\n", - "rho_expressions = formulate_phsp_factor_matrix(n_channels=len(REACTIONS))\n", - "rho.as_explicit().xreplace(rho_expressions)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### $F$-vector construction" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "For some reason one has to leave out the multiplication of $\\rho$ by $i$ within the calculation of the $F$ vector\n", - ":::" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "F = (I - sp.I * K * rho).inv() * P\n", - "F" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "F_vector = F.as_explicit()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "full-width" - ] - }, - "outputs": [], - "source": [ - "parametrizations = {**K_expressions, **rho_expressions, **P_expressions}\n", - "F_exprs = F_vector.xreplace(parametrizations)\n", - "F_exprs[0].simplify(doit=False)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Create numerical functions" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "F_unfolded_exprs = np.array([perform_cached_doit(expr) for expr in F_exprs])\n", - "DYNAMICS_EXPRESSIONS_FVECTOR = [\n", - " {\n", - " symbol: F_unfolded_exprs[i]\n", - " for symbol, resonances in create_dynamics_symbol.collected_symbols.items()\n", - " }\n", - " for i in range(n_channels)\n", - "]\n", - "MODELS_FVECTOR = [\n", - " attrs.evolve(\n", - " model,\n", - " parameter_defaults=PARAMETERS_DEFAULTS,\n", - " )\n", - " for model in MODELS\n", - "]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "FULL_EXPRESSIONS_FVECTOR = [\n", - " perform_cached_doit(MODELS_FVECTOR[i].expression).xreplace(\n", - " DYNAMICS_EXPRESSIONS_FVECTOR[i]\n", - " )\n", - " for i in range(n_channels)\n", - "]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "INTENSITY_FUNCS_FVECTOR = [\n", - " create_parametrized_function(\n", - " expression=perform_cached_doit(FULL_EXPRESSIONS_FVECTOR[i]),\n", - " backend=\"jax\",\n", - " parameters=MODELS_FVECTOR[i].parameter_defaults,\n", - " )\n", - " for i in range(n_channels)\n", - "]" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Generate data" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Phase space sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "HELICITY_TRANSFORMERS = [\n", - " SympyDataTransformer.from_sympy(model.kinematic_variables, backend=\"jax\")\n", - " for model in MODELS_FVECTOR\n", - "]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "PHSP = []\n", - "ε = 1e-8\n", - "for i in range(n_channels):\n", - " rng = TFUniformRealNumberGenerator(seed=0)\n", - " phsp_generator = TFPhaseSpaceGenerator(\n", - " initial_state_mass=REACTIONS[i].initial_state[-1].mass,\n", - " final_state_masses={it: p.mass for it, p in REACTIONS[i].final_state.items()},\n", - " )\n", - " phsp_momenta = phsp_generator.generate(100_000, rng)\n", - " phsp = HELICITY_TRANSFORMERS[i](phsp_momenta)\n", - " phsp = {k: v.real for k, v in phsp.items()}\n", - " phsp = {k: v + ε * 1j if re.match(r\"^m_\\d\\d$\", k) else v for k, v in phsp.items()}\n", - " PHSP.append(phsp)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Set parameters for toy model" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function for plotting histograms with JAX" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def fast_histogram(\n", - " data: jnp.ndarray,\n", - " weights: jnp.ndarray | None = None,\n", - " bins: int = 100,\n", - " density: bool | None = None,\n", - " fill: bool = True,\n", - " ax=plt,\n", - " **plot_kwargs,\n", - ") -> None:\n", - " bin_values, bin_edges = jnp.histogram(\n", - " data,\n", - " bins=bins,\n", - " density=density,\n", - " weights=weights,\n", - " )\n", - " if fill:\n", - " bin_rights = bin_edges[1:]\n", - " ax.fill_between(bin_rights, bin_values, step=\"pre\", **plot_kwargs)\n", - " else:\n", - " bin_mids = (bin_edges[:-1] + bin_edges[1:]) / 2\n", - " ax.step(bin_mids, bin_values, **plot_kwargs)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Functions for indicated resonances and thresholds" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def indicate_masses(ax, intensity_func, set_labels: bool = True):\n", - " mass_pars = {\n", - " k: v for k, v in intensity_func.parameters.items() if k.startswith(\"m_{N\")\n", - " }\n", - " for i, (k, v) in enumerate(mass_pars.items()):\n", - " label = f\"${k}$\" if set_labels else None\n", - " ax.axvline(v, c=f\"C{i + n_channels}\", label=label, ls=\"dashed\")\n", - "\n", - "\n", - "def indicate_thresholds(ax, set_labels: bool = True) -> None:\n", - " for i, decay in enumerate(DECAYS):\n", - " m_thr = sum(p.mass for p in decay.children)\n", - " label = None\n", - " if set_labels:\n", - " label = f\"${'+'.join(f'm_{{{p.latex}}}' for p in decay.children)}$\"\n", - " ax.axvline(m_thr, c=f\"C{i}\", label=label, ls=\"dotted\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "toy_parameters = {\n", - " R\"\\beta_{N_1(3/2^-)}\": 1 + 0j,\n", - " R\"m_{N_1(3/2^-)}\": 1.71,\n", - " R\"g_{N_1(3/2^-),0}\": 3.2,\n", - " R\"g_{N_1(3/2^-),1}\": 1.5,\n", - "}\n", - "for func in INTENSITY_FUNCS_FVECTOR:\n", - " func.update_parameters(toy_parameters)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(9, 4))\n", - "ax.set_title(\"Model rendering from phase space\")\n", - "ax.set_xlabel(R\"$m_{p\\eta/K\\Sigma}$ [GeV]\")\n", - "for i in range(n_channels):\n", - " intensity = np.real(INTENSITY_FUNCS_FVECTOR[i](PHSP[i]))\n", - " fast_histogram(\n", - " np.real(PHSP[i][\"m_01\"]),\n", - " weights=intensity,\n", - " alpha=0.5,\n", - " bins=200,\n", - " density=True,\n", - " label=f\"${' '.join(p.latex for p in DECAYS[i].children)}$\",\n", - " ax=ax,\n", - " )\n", - "indicate_thresholds(ax)\n", - "indicate_masses(ax, INTENSITY_FUNCS_FVECTOR[i])\n", - "ax.legend()\n", - "ax.set_ylim(0, None)\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Toy data sample" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "DATA = []\n", - "for i in range(n_channels):\n", - " weighted_phsp_generator = TFWeightedPhaseSpaceGenerator(\n", - " initial_state_mass=MODELS[i].reaction_info.initial_state[-1].mass,\n", - " final_state_masses={\n", - " i: p.mass for i, p in MODELS[i].reaction_info.final_state.items()\n", - " },\n", - " )\n", - " data_generator = IntensityDistributionGenerator(\n", - " domain_generator=weighted_phsp_generator,\n", - " function=INTENSITY_FUNCS_FVECTOR[i],\n", - " domain_transformer=HELICITY_TRANSFORMERS[i],\n", - " )\n", - " data_momenta = data_generator.generate(50_000, rng)\n", - " data = HELICITY_TRANSFORMERS[i](data_momenta)\n", - " DATA.append(data)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "fig, ax = plt.subplots(figsize=(9, 4))\n", - "ax.set_title(\"Toy data sample\")\n", - "ax.set_xlabel(R\"$m_{p\\eta/K\\Sigma}$ [GeV]\")\n", - "for i in range(n_channels):\n", - " fast_histogram(\n", - " np.real(DATA[i][\"m_01\"]),\n", - " alpha=0.5,\n", - " bins=200,\n", - " density=True,\n", - " label=f\"${' '.join(p.latex for p in DECAYS[i].children)}$\",\n", - " ax=ax,\n", - " )\n", - "indicate_thresholds(ax)\n", - "indicate_masses(ax, INTENSITY_FUNCS_FVECTOR[i])\n", - "ax.legend()\n", - "ax.set_ylim(0, None)\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Perform fit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Estimator definition" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "class EstimatorSum(Estimator):\n", - " def __init__(self, estimators: Iterable[Estimator]) -> None:\n", - " self.__estimators = tuple(estimators)\n", - "\n", - " def __call__(self, parameters: Mapping[str, ParameterValue]) -> float:\n", - " return sum(estimator(parameters) for estimator in self.__estimators)\n", - "\n", - " def gradient(\n", - " self, parameters: Mapping[str, ParameterValue]\n", - " ) -> dict[str, ParameterValue]:\n", - " raise NotImplementedError" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "combined_estimators = EstimatorSum(\n", - " UnbinnedNLL(\n", - " INTENSITY_FUNCS_FVECTOR[i],\n", - " data=DATA[i],\n", - " phsp=PHSP[i],\n", - " backend=\"jax\",\n", - " )\n", - " for i in range(n_channels)\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Initial parameters " - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Functions for comparing model to data" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def compare_models(functions: list[Function], title: str, bins: int = 100):\n", - " fig, axes = plt.subplots(figsize=(8.5, 4.5), nrows=2, sharex=True)\n", - " axes[0].set_title(title)\n", - " for ax in axes:\n", - " ax.set_yticks([])\n", - " for i in range(n_channels):\n", - " _plot_comparison(\n", - " axes[i],\n", - " decay_id=i,\n", - " variable_name=\"m_01\",\n", - " function=functions[i],\n", - " bins=bins,\n", - " color=f\"C{i}\",\n", - " legend=(i == 1),\n", - " )\n", - " fig.legend()\n", - " fig.tight_layout()\n", - " fig.show()\n", - "\n", - "\n", - "def _plot_comparison(\n", - " ax,\n", - " decay_id: int,\n", - " variable_name: str,\n", - " function: Function[DataSample, np.ndarray],\n", - " bins: int,\n", - " color: str,\n", - " legend: bool,\n", - "):\n", - " phsp = PHSP[decay_id]\n", - " fast_histogram(\n", - " DATA[decay_id][variable_name].real,\n", - " alpha=0.5,\n", - " bins=bins,\n", - " color=color,\n", - " density=True,\n", - " label=f\"Data ${' '.join(p.latex for p in DECAYS[decay_id].children)}$\",\n", - " ax=ax,\n", - " )\n", - " fast_histogram(\n", - " phsp[variable_name].real,\n", - " weights=function(phsp),\n", - " bins=bins,\n", - " color=\"red\",\n", - " density=True,\n", - " fill=False,\n", - " label=\"Fit model\" if legend else None,\n", - " ax=ax,\n", - " )\n", - " indicate_thresholds(ax, set_labels=legend)\n", - " indicate_masses(ax, function, set_labels=legend)\n", - " ax.set_ylim(0, None)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "initial_parameters = {\n", - " R\"m_{N_1(3/2^-)}\": 1.9,\n", - " R\"g_{N_1(3/2^-),0}\": 2.8,\n", - " R\"g_{N_1(3/2^-),1}\": 1.6,\n", - "}" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "ORIGINAL_PARAMETERS_F = []\n", - "for i in range(n_channels):\n", - " ORIGINAL_PARAMETERS_F.append(dict(INTENSITY_FUNCS_FVECTOR[i].parameters))\n", - " INTENSITY_FUNCS_FVECTOR[i].update_parameters(initial_parameters)\n", - "compare_models(INTENSITY_FUNCS_FVECTOR, title=\"Model with starting parameters\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Optimize parameters" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "minuit2 = Minuit2()\n", - "fit_result = minuit2.optimize(combined_estimators, initial_parameters)\n", - "assert fit_result.minimum_valid\n", - "fit_result" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "for i in range(n_channels):\n", - " INTENSITY_FUNCS_FVECTOR[i].update_parameters(fit_result.parameter_values)\n", - "compare_models(INTENSITY_FUNCS_FVECTOR, title=\"Model with optimized parameters\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Fit quality check" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Functions for inspecting fit result" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "def compute_aic_bic(fit_result: FitResult) -> tuple[float, float]:\n", - " n_real_par = fit_result.count_number_of_parameters(complex_twice=True)\n", - " n_events = len(next(iter(data.values())))\n", - " log_likelihood = -fit_result.estimator_value\n", - " aic = 2 * n_real_par - 2 * log_likelihood\n", - " bic = n_real_par * np.log(n_events) - 2 * log_likelihood\n", - " return aic, bic\n", - "\n", - "\n", - "def compare_parameters(initial: dict, optimized: dict, expected: dict) -> pd.DataFrame:\n", - " parameters = sorted(set(initial) | set(optimized))\n", - " df = pd.DataFrame(\n", - " {\n", - " f\"${p}$\": (\n", - " f\"{initial[p]:.3g}\",\n", - " f\"{optimized[p]:.3g}\",\n", - " f\"{expected[p]:.3g}\",\n", - " f\"{100 * abs((optimized[p] - expected[p]) / expected[p]):.1f}%\",\n", - " )\n", - " for p in parameters\n", - " },\n", - " ).T\n", - " df.columns = (\"initial\", \"fit result\", \"expected\", \"deviation\")\n", - " return df" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "compute_aic_bic(fit_result)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "compare_parameters(\n", - " initial=initial_parameters,\n", - " optimized=fit_result.parameter_values,\n", - " expected={\n", - " **ORIGINAL_PARAMETERS_F[0],\n", - " **ORIGINAL_PARAMETERS_F[1],\n", - " },\n", - ")" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/report/033.ipynb b/docs/report/033.ipynb deleted file mode 100644 index ccde37f2..00000000 --- a/docs/report/033.ipynb +++ /dev/null @@ -1,3458 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "source": [ - "```{autolink-concat}\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "3d", - "kinematics", - "physics", - "tutorial" - ] - }, - "source": [ - "::::{margin}\n", - ":::{card} PWA101: Amplitude analysis with Python basics\n", - "TR-033\n", - "^^^\n", - "This tutorial introduces amplitude analysis, and specifically the technique called Partial Wave Analysis (PWA), by demonstrating its application to a specific reaction channel and amplitude model. Basic Python programming and libraries (e.g. [`numpy`](https://numpy.org/doc/stable), [`scipy`](https://docs.scipy.org/doc/scipy), etc.) are used to illustrate the more fundamental steps of PWA in hadron physics.\n", - "+++\n", - "✅ [ComPWA/RUB-EP1-AG#93](https://github.com/ComPWA/RUB-EP1-AG/issues/93), [compwa.github.io#217](https://github.com/ComPWA/compwa.github.io/pull/217)\n", - ":::\n", - "::::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Amplitude Analysis 101 (PWA 101)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install -q iminuit==2.26.0 matplotlib==3.9.1 numpy==1.26.4 pandas==2.2.2 particle==0.24.0 phasespace==1.10.3 scipy==1.14.0 tqdm==4.66.4 vector==1.4.1" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Abstract\n", - "This document introduces Amplitude Analysis / Partial Wave Analysis (PWA) by demonstrating its application to a specific reaction channel and amplitude model. It aims to equip readers with a basic understanding of the full workflow and methodologies of PWA in hadron physics through a practical, hands-on example. Only basic Python programming and libraries (e.g. [`numpy`](https://numpy.org/doc/stable), [`scipy`](https://docs.scipy.org/doc/scipy), etc.) are used to illustrate the more fundamental steps in a PWA. Calculations with 4-vectors in this report are performed with the [`vector`](https://vector.readthedocs.io/en/latest/usage/intro.html) package.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso}\n", - "A follow-up tutorial, [PWA101 v2.0](https://compwa.github.io/gluex-nstar), is being prepared in [ComPWA/gluex-nstar#13](https://github.com/ComPWA/gluex-nstar/pull/13). Whereas this report focuses on common, numerical Python libraries, v2.0 formulates the amplitude model with a {doc}`symbolic approach`.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "When performing an amplitude analysis, one usually starts with a long process of selecting data ('events') from a particle collider experiment for specific decay channels. It is a delicate process with the intend to select as much background noise from other decay channels, while maintaining sufficient signal events that an amplitude model can be 'fit' to it. This tutorial will not consider this event selection, but will focus on amplitude model formulation and fitting it to a data sample that is assumed to be come from experiment. As such, the tutorial is built up of the following main steps:\n", - "\n", - "1. [Formulate amplitude model](#amplitude-model).\n", - "2. [Visualize and inspect model](#visualization).\n", - "3. [Generate toy data](#generate-data).\n", - "4. [Fit model to data distribution](#fit-model).\n", - "\n", - "The following Python packages are all that we require to go through each of the steps." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Import Python libraries" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "from __future__ import annotations\n", - "\n", - "import logging\n", - "import os\n", - "import warnings\n", - "\n", - "import matplotlib as mpl\n", - "import matplotlib.pyplot as plt\n", - "import numpy as np\n", - "import pandas as pd\n", - "import phasespace\n", - "import scipy as sp\n", - "import tensorflow as tf\n", - "import vector\n", - "from iminuit import Minuit\n", - "from matplotlib import gridspec\n", - "from tqdm.auto import tqdm\n", - "from vector.backends.numpy import MomentumNumpy4D\n", - "\n", - "os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n", - "logging.disable(logging.WARNING)\n", - "warnings.filterwarnings(\"ignore\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(amplitude-model)=\n", - "## Amplitude model" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Theory" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "\n", - "This amplitude model is adapted from the [Lecture 11 in STRONG2020 HaSP School](https://indico.ific.uv.es/event/6803/contributions/21223/) by Vincent Mathieu.\n", - "\n", - "The reaction $ \\gamma p \\to \\eta \\pi^0 p$ is one of the reaction channels that are studied in photo-production experiments like [GlueX](http://www.gluex.org). For simplicity, we assume that the decay proceeds through three possible resonances—$a_2$, $\\Delta^+$, and $N^*$—with each in a different subsystem." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "```{image} https://github.com/ComPWA/compwa-org/assets/17490173/ec6bf191-bd5f-43b0-a6cb-da470b071630\n", - ":width: 100%\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we denote $1 \\equiv \\eta, 2 \\equiv \\pi^0, 3 \\equiv p$. Given these three subsystems in this particle transition, we can construct three amplitudes," - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "\\begin{eqnarray}\n", - "A^{12} &=& \\frac{\\sum a_m Y_2^m (\\Omega_1)}{s_{12}-m^2_{a_2}+im_{a_2} \\Gamma_{a_2}} \\times s_{12}^{0.5+0.9u_3} \\,, \\nonumber \\\\\n", - "A^{23} &=& \\frac{\\sum b_m Y_1^m (\\Omega_2)}{s_{23}-m^2_{\\Delta}+im_{\\Delta} \\Gamma_{\\Delta}} \\times s_{23}^{0.5+0.9t_1} \\,, \\nonumber \\\\\n", - "A^{31} &=& \\frac{c_0}{s_{31}-m^2_{N^*}+im_{N^*} \\Gamma_{N^*}} \\times s_{31}^{1.08+0.2t_2} \\,,\n", - "\\end{eqnarray}\n", - "$$ (full-model-with-exponential)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "where $s, t, u$ are the Mandelstam variables ($s_{ij}=(p_i+p_j)^2$, $t_i=(p_a-p_i)^2$, and $u_i=(p_b-p_i)^2$), $m$ is the resonance mass, $\\Gamma$ is the resonance width, $Y^m_l$ are spherical harmonics functions, $\\Omega_i$ are pairs of Euler angles (polar angle $\\theta$ and azimuthal angle $\\phi$) that represent the helicity decay angles, and $a_i$, $b_i$, and $c_i$ are complex-valued coefficients. Note that the Mandelstam variables and angles come from measured events, while the other symbols are parameters that need to be modified in order to have the amplitude model match the data." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{note}\n", - "All that an amplitude model requires as data input, are **four-momenta**. For this amplitude model, there are just three required four-momenta per event, one for each final state in $p\\gamma \\to \\eta\\pi^0 p$.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The original full amplitude model from the [Lecture 11 in STRONG2020 HaSP School](https://indico.ific.uv.es/event/6803/contributions/21223/) is shown in Equation {eq}`full-model-with-exponential`.\n", - "*In this report, only the Breit–Wigner and Spherical harmonics terms are kept, while the exponential factor is abandoned, i.e.*" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\n", - "\\begin{eqnarray}\n", - "A^{12} &=& \\frac{\\sum a_m Y_2^m (\\Omega_1)}{s_{12}-m^2_{a_2}+im_{a_2} \\Gamma_{a_2}} \\,, \\\\\n", - "A^{23} &=& \\frac{\\sum b_m Y_1^m (\\Omega_2)}{s_{23}-m^2_{\\Delta}+im_{\\Delta} \\Gamma_{\\Delta}} \\,, \\\\\n", - "A^{31} &=& \\frac{c_0}{s_{31}-m^2_{N^*}+im_{N^*} \\Gamma_{N^*}} \\,.\n", - "\\end{eqnarray}\n", - "$$ (model-dynamics-and-angular)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The intensity $I$ that describes our measured distributions is then expressed as a coherent sum of the amplitudes $A^{ij}$,\n", - "\n", - "$$\n", - "\\begin{eqnarray}\n", - "I &=& |A^{12} + A^{23} + A^{31}|^2 \\,.\n", - "\\end{eqnarray}\n", - "$$ (intensity-coherent-sum)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Ultimately, the amplitude model in Equations {eq}`model-dynamics-and-angular` and {eq}`intensity-coherent-sum` in this tutorial consists of three resonances, and each of them are formed by two components: a Breit-Wigner times some spherical harmonics ($l = 2, 1, 0$)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Assumption: spinless final state!\n", - ":class: dropdown\n", - "While the **spin** of the $\\eta$ meson and the $\\pi^0$ meson are both $0$, the spin of the proton is spin-$\\frac{1}{2}$.\n", - "However, in order to simplify the amplitude model, we treat the proton as a spin-0 particle.\n", - "Overall, the $\\eta$, $\\pi^0$ and $p$ are therefore all treated as spin-0 particles.\n", - "\n", - "The primary motivation for assuming proton to be spin-$0$ is to avoid the necessity of aligning the amplitudes (see e.g. [ComPWA/ampform#6](https://github.com/ComPWA/ampform/issues/6)).\n", - "This assumption enables the intensity $I$ to be written as a coherent sum of the amplitudes of the subsystems without the need for additional Wigner rotations.\n", - "In addition, the amplitude for each decay chain contains only one spherical harmonic, namely that of the resonance.\n", - "\n", - "The spherical harmonics in Equation {eq}`model-dynamics-and-angular` are therefore relevant only to the resonances. \n", - "Here, $l$ represents the spin of the resonances and $m$ represents its spin projection in the decay chain.\n", - "The total angular momentum and coupled spin (for the two-body state of the two decay products of the resonance) are not considered.\n", - "According to {cite}`Richman:1984gh` and other classical references on helicity, this is known as the **helicity basis**.\n", - "In contrast, the **canonical basis** does not sum over $L$ and $S$, resulting in a more complex coherent sum of amplitudes. \n", - "The transformation between these bases is also discussed [here](https://ampform.rtfd.io/0.15.x/usage/helicity/formalism.html) on the AmpForm website.\n", - "\n", - "In our case: \n", - "- $A^{12}$ represents a **d-wave** interaction, because we assume there is a $a_2$ resonance (spin 2) in this subsystem. The possible $m$ values are $−2,−1,0,1,2$. Each of these values corresponds to different orientations of the d-wave.\n", - "- $A^{23}$ represents a **p-wave** interaction, because we assume this subsystem has a (spin-1) $\\Delta^+$ resonance. The possible $m$ values are $−1,0,1$.\n", - "- $A^{31}$ represents an **s-wave** interaction, because we assume there is one spin-0 $N^*$ resonance in this subsystem. The only possible $m$ value is 0 and, since $Y_0^0=0$, the amplitude only consists of a Breit–Wigner.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Implementation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we define a few functions that implement parts of the amplitude model of Equation {eq}`model-dynamics-and-angular`. Later on in [](#visualization), we can use these functions to visualize each components of the model individually." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "(breit-wigner-model)=\n", - "#### Breit-Wigner (only) Model" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "The following functions define the Breit–Wigner function, as well as the following intensity function that only contains the Breit–Wigner (dynamics) component of each amplitude.\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "I &= &\n", - "\\left\\lvert \\frac{1}{s-m^2_{a_2}+im_{a_2} \\Gamma_{a_2}} +\n", - "\\frac{1}{s-m^2_{\\Delta}+im_{\\Delta} \\Gamma_{\\Delta}} +\n", - "\\frac{1}{s-m^2_{N^*}+im_{N^*} \\Gamma_{N^*}} \\right\\rvert ^2\n", - "\\end{array}\n", - "$$ (model-dynamics-only)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def BW_model(s12, s23, s31, *, M12, Gamma12, M23, Gamma23, M31, Gamma31, **kwargs):\n", - " A12 = BW(s12, M12, Gamma12)\n", - " A23 = BW(s23, M23, Gamma23)\n", - " A31 = BW(s31, M31, Gamma31)\n", - " return np.abs(A12 + A23 + A31) ** 2\n", - "\n", - "\n", - "def BW(s, m, Gamma):\n", - " return 1 / (s - m**2 + m * Gamma * 1j)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "(spherical-harmonics-model)=\n", - "#### Spherical Harmonics (only) Model" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The calculation of $Y_l^m(\\phi, \\theta)$ is done via [`scipy.special.sph_harm()`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.sph_harm.html). However, we use a different definition of $\\phi$ and $\\theta$, following a notation that is more common in hadron physics.\n", - "- $\\phi$ is the **azimuthal angle** and ranges from -$\\pi$ to $\\pi$. SciPy represents this as $\\theta$, ranging from $0$ to $2\\pi$.\n", - "- $\\theta$ is the **polar angle** and ranges from $0$ to $\\pi$. SciPy represents this as $\\phi$ with the same range.\n", - "\n", - "This leads to\n", - "\n", - "$$\n", - "Y_\\ell^m(\\phi, \\theta) = \\sqrt{\\frac{2n+1}{4\\pi}\\frac{(n-m)!}{(n+m)!}}e^{im\\phi}P_\\ell^m\\left(\\cos\\theta\\right) \\,.\n", - "$$ (spherical-harmonics-definition)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "Alternatively, we can formulate the spherical harmonics in terms of a Wigner-$d$ or Wigner-$D$ function, as\n", - "\n", - "$$\n", - "\\begin{eqnarray}\n", - "Y_\\ell^m(\\theta,\\phi) = \\sqrt{\\frac{2\\ell+1}{4\\pi}}D^\\ell_{m0}(-\\phi, \\theta, 0) = \\sqrt{\\frac{2\\ell+1}{4\\pi}}e^{im\\phi} d^\\ell_{m0}(\\theta) \\,.\n", - "\\end{eqnarray}\n", - "$$ (spherical-harmonics-wigner-d)\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the following, we define functions to compute the spherical harmonics of the three subsystem amplitudes in Equation {eq}`model-dynamics-and-angular`. Note how the function signature consists of two input data columns, `theta` and `phi`, and how the rest of the arguments are parameters. The final `kwargs` (key-word arguments) is there so that we can compose larger functions from these function definitions." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def Ylm12(\n", - " theta: np.ndarray,\n", - " phi: np.ndarray,\n", - " *,\n", - " a_minus2,\n", - " a_minus1,\n", - " a_0,\n", - " a_plus1,\n", - " a_plus2,\n", - " **kwargs,\n", - ") -> np.ndarray:\n", - " return (\n", - " a_plus2 * sp.special.sph_harm(2, 2, phi, theta)\n", - " + a_plus1 * sp.special.sph_harm(1, 2, phi, theta)\n", - " + a_0 * sp.special.sph_harm(0, 2, phi, theta)\n", - " + a_minus1 * sp.special.sph_harm(-1, 2, phi, theta)\n", - " + a_minus2 * sp.special.sph_harm(-2, 2, phi, theta)\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def Ylm23(\n", - " theta: np.ndarray, phi: np.ndarray, *, b_minus1, b_0, b_plus1, **kwargs\n", - ") -> np.ndarray:\n", - " return (\n", - " b_plus1 * sp.special.sph_harm(1, 1, phi, theta)\n", - " + b_0 * sp.special.sph_harm(0, 1, phi, theta)\n", - " + b_minus1 * sp.special.sph_harm(-1, 1, phi, theta)\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def Ylm31(theta: np.ndarray, phi: np.ndarray, c_0: complex, **kwargs) -> np.ndarray:\n", - " return c_0 * sp.special.sph_harm(0, 0, phi, theta)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "We now have the ingredients to define an intensity function that only contains spherical harmonics, that is, the angular part of the amplitude model:\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "I &=&\n", - "\\left\\lvert \\sum a_m Y_2^m (\\Omega_1) +\n", - "\\sum b_m Y_1^m (\\Omega_2) +\n", - "c_0 \\right\\rvert ^2 \\,.\n", - "\\end{array}\n", - "$$ (model-angular-only)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def SH_model(phi1, theta1, phi2, theta2, *, c_0, **pars):\n", - " return np.abs(Ylm12(phi1, theta1, **pars) + Ylm23(phi2, theta2, **pars) + c_0) ** 2" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Breit-Wigner $\\times$ Spherical Harmonics Model" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "Finally, we combine the [](#breit-wigner-model) and [](#spherical-harmonics-model) to get an implementation for Equation {eq}`model-dynamics-and-angular`,\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "I &=& \\left\\lvert A^{12}+A^{23}+A^{31} \\right\\rvert ^2 \\nonumber \\\\\n", - "&=&\n", - "\\left\\lvert \\frac{\\sum a_m Y_2^m (\\Omega_1)}{s-m^2_{a_2}+im_{a_2} \\Gamma_{a_2}} +\n", - "\\frac{\\sum b_m Y_1^m (\\Omega_2)}{s-m^2_{\\Delta}+im_{\\Delta} \\Gamma_{\\Delta}} +\n", - "\\frac{c_0}{s-m^2_{N^*}+im_{N^*} \\Gamma_{N^*}} \\right\\rvert ^2 \\,.\n", - "\\end{array}\n", - "$$" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def BW_SH_model(\n", - " s12,\n", - " s23,\n", - " s31,\n", - " phi1,\n", - " theta1,\n", - " phi2,\n", - " theta2,\n", - " *,\n", - " M12,\n", - " Gamma12,\n", - " M23,\n", - " Gamma23,\n", - " M31,\n", - " Gamma31,\n", - " c_0,\n", - " **parameters,\n", - "):\n", - " A12 = BW(s12, M12, Gamma12) * Ylm12(phi1, theta1, **parameters)\n", - " A23 = BW(s23, M23, Gamma23) * Ylm23(phi2, theta2, **parameters)\n", - " A31 = BW(s31, M31, Gamma31) * c_0\n", - " return np.abs(A12 + A23 + A31) ** 2" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "## Visualization" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As discussed in [](#theory), the amplitude model contains variables that are provided by experiment (\"events\" as data input) as well as parameters that are to be optimized. The data input is usually in the form of **four-momenta**, while the model is formulated in the form of Mandelstam variables and helicity angles. We therefore have to compute these variables from the four-momenta." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Phase space generation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "First, however, we need to generate a phase space sample of four-momenta in order to plot the amplitude model as a distribution over each of the variables. In this section, we use the [`phasespace`](https://github.com/zfit/phasespace) package for generating four-momenta for the reaction $p\\gamma \\to p\\eta\\pi^0$. The phase space sample will also be used later on to normalize the model when calculating the likelihood over the data sample (see [](#fit-model))." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Lab frame and CM frame" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In Center-of-Mass (CM) frame, the 4-momentum of the total system can be acquired by 4-momentum conservation:" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "$$\\begin{pmatrix}\n", - " E_0 \\\\\n", - " 0 \\\\\n", - " 0 \\\\\n", - " 0\n", - "\\end{pmatrix}\n", - "= \\begin{pmatrix}\n", - " E_{\\gamma} \\\\\n", - " 0 \\\\\n", - " 0 \\\\\n", - " p_z\n", - "\\end{pmatrix} +\n", - "\\begin{pmatrix}\n", - " E_p \\\\\n", - " 0 \\\\\n", - " 0 \\\\\n", - " -p_z\n", - "\\end{pmatrix}\n", - "= \\begin{pmatrix}\n", - " E_{\\eta} \\\\\n", - " p_{\\eta,x} \\\\\n", - " p_{\\eta,y} \\\\\n", - " p_{\\eta,z}\n", - "\\end{pmatrix} +\n", - "\\begin{pmatrix}\n", - " E_{\\pi} \\\\\n", - " p_{\\pi,x} \\\\\n", - " p_{\\pi,y} \\\\\n", - " p_{\\pi,z}\n", - "\\end{pmatrix} +\n", - "\\begin{pmatrix}\n", - " E_p \\\\\n", - " p_{p,x} \\\\\n", - " p_{p,y} \\\\\n", - " p_{p,z}\n", - "\\end{pmatrix}\n", - "$$ (four-momentum-conservation)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "source": [ - "From this we can see that the system depends mainly on beam momentum $p_z$ and CM total energy $E_0$." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{caution}\n", - "The calculation here involved using values from the CM frame. While this frame is commonly used for theoretical calculations, experimental data is often analyzed in the lab frame. However, it's worth noting that oin some collider experiments that do not have a fixed target, the CM frame can coincide with the lab frame.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The [GlueX](http://www.gluex.org/) experiment at Jefferson Lab uses a fixed proton target with a linearly polarized photon beam, and the beam energy range in the lab frame is typically from [**8 to 9 GeV**](https://doi.org/10.7566/JPSCP.26.022002)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We use the $\\gamma$ beam energy in the lab frame as input, e.g. $E_{\\gamma, \\text{lab}} = 8.5 \\; \\text{GeV}$, and want to know the collision energy in the CM frame.\n", - "\n", - "We can calculate the energy of the photon in the CM frame as follows. The four-momentum of photon (beam) in the lab frame is\n", - "\n", - "$$\n", - "p_{\\gamma,\\text{lab}} = \\left(E_{\\gamma, \\text{lab}}, \\vec{p}_{\\gamma,\\text{lab}}\\right) \\,.\n", - "$$ (beam-four-momentum)\n", - "\n", - "Since the photon is massless, $m_\\gamma=0$, and\n", - "\n", - "$$\n", - "m^2 = E^2 - \\left|\\vec{p}\\right|^2 \\,,\n", - "$$ (invariant-mass-definition)\n", - "\n", - "we get\n", - "\n", - "$$\n", - "E_{\\gamma} = |\\vec{p_{\\gamma}}| .\n", - "$$ (gamma-energy)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The proton (target) is at rest in the lab frame, so we have\n", - "\n", - "$$\n", - "p_{p,\\text{lab}} = (m_p, \\vec{0})\\,,\n", - "$$ (four-momentum-target)\n", - "\n", - "where $m_p$ is the proton mass. We have a total four-momentum in the lab frame of\n", - "\n", - "$$\n", - "p_{\\text{tot},\\text{lab}} = p_{\\gamma,\\text{lab}} + p_{p,\\text{lab}}\n", - "$$ (total-four-momentum-lab)\n", - "\n", - "and we know\n", - "\n", - "$$\n", - "E_p = \\sqrt{p_p^2+ m_p^2} \\,.\n", - "$$ (four-momentum-proton)\n", - "\n", - "The CM total energy $E_0$ expressed in terms of quantities from the **lab frame** is thus\n", - "\n", - "$$\n", - "m_0 \\equiv E_0 = \\sqrt{2 E_{\\gamma,\\text{lab}} m_p + m_p^2} \\,.\n", - "$$ (total-energy-cm)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Equivalently, from the **CM frame** perspective, since $\\vec{p}_{\\gamma} = -\\vec{p}_{p}$ and $|\\vec{p}_{\\gamma}| = |\\vec{p}_{p}|= p_{z}$, we find\n", - "\n", - "$$\n", - "E_0 = m_0 = |p_{z}|+\\sqrt{p_{z}^2 + m_p^2} \\,.\n", - "$$ (total_energy_CM_frame_still)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "remove-cell" - ] - }, - "source": [ - "Therefore, the energy of the photon ($E_{\\gamma}$) and its momentum ($p_{\\gamma}$) in the **CM frame** are:\n", - "\n", - "$$\n", - "E_{\\gamma} = |\\vec{p}_{\\gamma}| = p_{z} = 1.944 \\; \\text{GeV} \\,.\n", - "$$" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Our implementation based on Equation {eq}`total-energy-cm` thus becomes:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "E_lab_gamma = 8.5\n", - "m_proton = 0.938\n", - "m_eta = 0.548\n", - "m_pi = 0.135\n", - "m_0 = np.sqrt(2 * E_lab_gamma * m_proton + m_proton**2)\n", - "m_0" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "p_beam = 1.943718\n", - "np.testing.assert_almost_equal(\n", - " p_beam + np.sqrt(p_beam**2 + m_proton**2), m_0, decimal=6\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Thus, we then have the mass of the system $m_0$ (or the mass of a 'virtual' particle $p\\gamma$) in CM frame of\n", - "\n", - "$$\n", - "E_{0} = m_0\\approx 4.102 \\;\\; \\text{GeV} \\,.\n", - "$$" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Final state four-momenta" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The [`phasespace`](https://github.com/zfit/phasespace) library is a Python package designed to simulate particle decays according to the principles of relativistic kinematics and phase space distributions. We first use the [`phasespace`](https://github.com/zfit/phasespace) to generate decay particles." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "phsp_events = 500_000\n", - "weights, particles = phasespace.nbody_decay(m_0, [m_eta, m_pi, m_proton]).generate(\n", - " n_events=phsp_events\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that each event comes with a statistical weights for each generated event. These weights represent how likely each particular event configuration (set of momenta and energies) is, based on phase space considerations. In order to generate a flat distribution, we will have to use a hit-and-miss method over these weights." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "(Click to show) These are some functions that we defined for a structured process for generating phase space samples of particle decays and converting them into four-momentum vectors using TensorFlow and the phasespace Python package." - }, - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def generate_phsp_decay(\n", - " size: int, seed: int | None = None, bunch_size: int = 500_000\n", - ") -> tuple[MomentumNumpy4D, MomentumNumpy4D, MomentumNumpy4D]:\n", - " rng = np.random.default_rng(seed=seed)\n", - " generated_seed = rng.integers(1_000_000)\n", - " phsp_sample = generate_phsp_bunch(bunch_size, generated_seed)\n", - " while get_size(phsp_sample) < size:\n", - " bunch = generate_phsp_bunch(bunch_size, generated_seed)\n", - " phsp_sample = concatenate(phsp_sample, bunch)\n", - " phsp_sample = remove_overflow(phsp_sample, size)\n", - " return tuple(to_vector(tensor) for tensor in phsp_sample)\n", - "\n", - "\n", - "def generate_phsp_bunch(\n", - " size: int, seed: int | None = None\n", - ") -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:\n", - " rng = np.random.default_rng(seed=seed)\n", - " final_state = [m_eta, m_pi, m_proton]\n", - " weights, particles = phasespace.nbody_decay(m_0, final_state).generate(\n", - " n_events=size, seed=seed\n", - " )\n", - " random_weights = rng.uniform(0, weights.numpy().max(), size=weights.shape)\n", - " selector = weights > random_weights\n", - " return tuple(particles[f\"p_{i}\"][selector] for i in range(len(final_state)))\n", - "\n", - "\n", - "def get_size(phsp: tuple[tf.Tensor, tf.Tensor, tf.Tensor]):\n", - " return len(phsp[0])\n", - "\n", - "\n", - "def concatenate(\n", - " phsp1: tuple[tf.Tensor, tf.Tensor, tf.Tensor],\n", - " phsp2: tuple[tf.Tensor, tf.Tensor, tf.Tensor],\n", - ") -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]:\n", - " return tuple(tf.concat([phsp1[i], phsp2[i]], axis=0) for i in range(3))\n", - "\n", - "\n", - "def remove_overflow(phsp: tuple[tf.Tensor, tf.Tensor, tf.Tensor], size: int):\n", - " return tuple(tensor[:size] for tensor in phsp)\n", - "\n", - "\n", - "def to_vector(tensor: tf.Tensor) -> MomentumNumpy4D:\n", - " return vector.array({\n", - " key: tensor.numpy().T[j] for j, key in enumerate([\"px\", \"py\", \"pz\", \"E\"])\n", - " })" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Hit-and-miss sampling\n", - ":class: dropdown\n", - "- **Step 1**: Generate a small phase space sample bunch with `generate_phsp_bunch()`\n", - " - A NumPy random number generator (`rng`) is initialized with a fixed seed for reproducibility. This ensures that every time the simulation is run, it produces the same random numbers.\n", - " - The `phasespace` package is used to simulate a decay process where a parent particle decays into three daughter particles. This is set up by specifying the masses of the parent particle (`m_0`) and the daughter particles (`m_eta`, `m_pi`, `m_proton`). The function `nbody_decay` followed by generate is used to create the phase space and return both the weights of each decay event and the momenta of the decay products.\n", - " - As mentioned, the `phasespace` package generates four-momenta that are not evenly distributed .The distributions appear even only when multiplied by the generated weights. To get a sample bunch that is uniformly distributed, rejection sampling (\"hit-and-miss\") is applied over the weights. To achieve this, a random set of weights (`random_weights`) is generated, and only those events where the original weight exceeds the random weight are selected.\n", - "\n", - "- **Step 2**, Ensure sufficient sample size with `generate_phsp_decay()`\n", - " - The function starts by generating an initial phase space sample. If this initial sample does not meet the required size, more samples are generated and concatenated until the requested size is reached.\n", - " - Once the target sample size is reached or exceeded, `remove_overflow()` trims the sample to precisely match the requested size.\n", - "\n", - "- **Step 3**, Convert tensors to [four-momentum vectors](https://vector.readthedocs.io/en/latest/api/backends/vector.backends.numpy.html#vector.backends.numpy.MomentumNumpy4D) objects\n", - " The `phasespace` package generates four-momenta in the form of $4\\times N$ TensorFlow tensors, with $N$ the number of events. For each tensor in the phase space tuple (representing the momentum components like $E, p_x, p_y, p_z$), the `to_vector` function converts them into a structured `MomentumNumpy4D` array. This structured array includes keys for each momentum component and possibly the energy component and provides convenient methods for computing boosts etc.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Initial state four-momenta" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Because of the simple relation in Equation {eq}`four-momentum-conservation`, the four-momenta for the initial state, $\\gamma$ (a) and $p$ (b), do not have to be generated by a phase space generator." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To find $p_a$ and $p_b$ by 4-momentum conservation, we can use\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "m_a &=& m_{\\gamma} =0 \\\\\n", - "m_b &=& m_p \\\\\n", - "p_{a,x} &=& p_{a,y} = 0 = p_{b,x} = p_{b,y} \\\\\n", - "p_{a,z} &=& - p_{b,z} \\,.\n", - "\\end{array}\n", - "$$\n", - "\n", - "Due to energy conservation, we have\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "E_a + E_b &=& E_1 + E_2 + E_3 = E_0 \\\\\n", - "E_a + E_b &=& \\sqrt{m_a^2 + p_a^2} + \\sqrt{m_b^2 + p_b^2} = E_0 \\,.\n", - "\\end{array}\n", - "$$\n", - "\n", - "Since $m_a=0$ and $p_a = -p_b$, we have\n", - "\n", - "$$\n", - "p_a + \\sqrt{m_b^2 + (-p_a)^2} = E_0 \\,.\n", - "$$\n", - "\n", - "Reorganizing, we get,\n", - "\n", - "$$\n", - "p_{a,z} + \\sqrt{m_p^2 + p_a^2} - E_0 = 0\n", - "$$\n", - "\n", - "and\n", - "\n", - "$$\n", - "p_{a,z} = \\frac{E_0^2 - m_p^2}{2E_0}\n", - "$$\n", - "\n", - "Finally, this gives us\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "p_{b,z} &=& -p_{a,z} \\\\\n", - "E_a &=& p_a \\\\\n", - "E_{b} &=& \\sqrt{m_b^2+p_b^2} \\,.\n", - "\\end{array}\n", - "$$ (pb-pa-relation)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "With this, our numerical implementation to compute four-momentum of $p_a$ and $p_b$ from Equation {eq}`pb-pa-relation` becomes:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - } - }, - "outputs": [], - "source": [ - "def compute_pa_pb(\n", - " p1: MomentumNumpy4D, p2: MomentumNumpy4D, p3: MomentumNumpy4D\n", - ") -> tuple[MomentumNumpy4D, MomentumNumpy4D]:\n", - " shape = p1.shape\n", - " E0 = (p1 + p2 + p3).e\n", - " px = np.zeros(shape)\n", - " py = np.zeros(shape)\n", - " pz = np.ones(shape) * (E0**2 - p3.m.mean() ** 2) / (2 * E0)\n", - " E = np.ones(shape) * np.sqrt(p3.m.mean() ** 2 + pz.mean() ** 2)\n", - " pa = MomentumNumpy4D({\"E\": pz, \"px\": px, \"py\": py, \"pz\": pz})\n", - " pb = MomentumNumpy4D({\"E\": E, \"px\": px, \"py\": py, \"pz\": -pz})\n", - " return pa, pb" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Four-momenta of all particles" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now create a function `generate_phsp_all()` to create and compute all particles in phase space all-at-once, combining the two functions from the previous sections into one function." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def generate_phsp_all(\n", - " size: int, seed: int | None = None, bunch_size: int = 500_000\n", - ") -> tuple[\n", - " MomentumNumpy4D, MomentumNumpy4D, MomentumNumpy4D, MomentumNumpy4D, MomentumNumpy4D\n", - "]:\n", - " p1, p2, p3 = generate_phsp_decay(size, seed, bunch_size)\n", - " pa, pb = compute_pa_pb(p1, p2, p3)\n", - " return p1, p2, p3, pa, pb" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "%%time\n", - "p1_phsp, p2_phsp, p3_phsp, pa_phsp, pb_phsp = generate_phsp_all(\n", - " size=phsp_events, seed=42, bunch_size=1_000_000\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Kinematic variable calculation" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Spherical coordinate system" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Before introducing CM and helicity angles, we first introduce **polar angles** and **azimuthal angles** in spherical coordinates." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In the spherical coordinates, the **polar angle** $\\theta$ and **azimuthal angle** $\\phi$ are defined as\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "\\theta &=& \\arccos \\frac{p_z}{|p|} \\nonumber \\\\\n", - "\\phi &=& \\arctan2(p_y , p_x) \\,.\n", - "\\end{array}\n", - "$$ (polar-azimuthal-angle)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "```{image} https://github.com/ComPWA/compwa.github.io/assets/29308176/89e1ad6e-3a7b-4895-b747-8bd644652504\n", - ":align: center\n", - ":width: 40%\n", - "```" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In Equation {eq}`polar-azimuthal-angle`, $p_z$ is equivalent to $z$, and $|p|$ is $r$ in figure above, while $p_y$ equivalent to $y$, and $p_x$ is $x$." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The sample is plotted to check whether the distribution looks uniformly distributed in the Dalitz plane. The Mandelstam variable $s$ of each of the three subsystems can be easily computed with from the four-momentum objects as follows. Here, we use the methods and attributes provided by the [`vector`](https://vector.readthedocs.io) package." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "p12_phsp = p1_phsp + p2_phsp\n", - "p23_phsp = p2_phsp + p3_phsp\n", - "p31_phsp = p3_phsp + p1_phsp\n", - "s12_phsp = p12_phsp.m2\n", - "s23_phsp = p23_phsp.m2\n", - "s31_phsp = p31_phsp.m2" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "fig = plt.figure(figsize=(12, 5))\n", - "fig.suptitle(\"Phase Space Dalitz Plot\")\n", - "gs = gridspec.GridSpec(1, 3, width_ratios=[1, 1, 0.05]) # The last column for colorbar\n", - "ax1 = plt.subplot(gs[0])\n", - "ax2 = plt.subplot(gs[1])\n", - "cax = plt.subplot(gs[2]) # For colorbar\n", - "\n", - "hist2 = ax2.hist2d(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " bins=100,\n", - " cmin=1e-2,\n", - " density=True,\n", - " cmap=\"jet\",\n", - " range=[[0, 10.3], [0.5, 13]],\n", - ")\n", - "ax2.set_title(\"2D Histogram\")\n", - "ax2.set_xlabel(R\"$m^2(\\eta \\pi^0)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", - "\n", - "ax1.scatter(s12_phsp, s23_phsp, s=1e-4, c=\"black\", norm=mpl.colors.Normalize())\n", - "ax1.set_xlabel(R\"$m^2(\\eta \\pi^0)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", - "ax1.set_ylabel(R\"$m^2(\\pi^0 p)\\;\\left[\\mathrm{GeV}^2\\right]$\")\n", - "ax1.set_title(\"Scatter Plot\")\n", - "\n", - "fig.colorbar(hist2[3], cax=cax)\n", - "fig.tight_layout()\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://github.com/user-attachments/assets/88379dc9-235a-428c-a8e1-d0bcff6d8648)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "There are different ways to represent the Dalitz plot, each with its advantages.\n", - "\n", - "- *Scatter Plot*: This method plots individual events as points, offering a clear view of the density and distribution of events within the phase space. It is particularly useful for visualizing smaller datasets or when high resolution is needed to identify specific features or clusters.\n", - "\n", - "- *2D Histogram*: This approach divides the phase space into bins and counts the number of events within each bin, representing the density of events using a color scale. It is effective for large datasets, providing a smooth and continuous representation of the phase space that highlights overall trends and structures.\n", - "\n", - "In conclusion, to check if the phase space sample is evenly distributed, a scatter plot is typically more straightforward and visually clear.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### CM Angles " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Angles in the CM frame are the polar and azimuthal angles of the spatial components of the four-momenta in the CM frame (i.e. the frame that satisfies the relations in Equation {eq}`four-momentum-conservation`). They are different than the [helicity angles](#helicity-angles) in each subsystem (which is after rotation and boost into the subsystem)." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The values for phase space can be obtained directly in the CM frame, without boosting into a different frame after the generation.\n", - "We denote these angles as\n", - "\n", - "$$\n", - "\\theta_1^\\text{CM}, \\theta_2^\\text{CM}, \\theta_3^\\text{CM}, \\phi_1^\\text{CM}, \\phi_2^\\text{CM}, \\phi_3^\\text{CM} \\,.\n", - "$$" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "theta1_CM_phsp = p1_phsp.theta\n", - "theta2_CM_phsp = p2_phsp.theta\n", - "theta3_CM_phsp = p3_phsp.theta\n", - "phi1_CM_phsp = p1_phsp.phi\n", - "phi2_CM_phsp = p2_phsp.phi\n", - "phi3_CM_phsp = p3_phsp.phi" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Helicity angles" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{seealso}\n", - "- [Section of helicity formalism in **TR-015** (Spin alignment implementation)](https://compwa.github.io/report/015.html#helicity-formalism)\n", - "- [Helicity versus canonical in Ampform](https://ampform.readthedocs.io/stable/usage/helicity/formalism.html)\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The amplitude model shown in Equation {eq}`model-dynamics-and-angular` derives from the helicity formalism. This spin formalism breaks up the transition amplitudes of a decay chain into two-body decays (**isobar model**), which allows separating angular (spin) dependence from dynamics (e.g. Breit–Wigner).\n", - "\n", - "Crucially, the helicity formalism builds on a chain of boosts and rotations that align the reference frame of each two-body decay node such that the particle moves along the $z$ direction. In the amplitude model of Equation {eq}`model-dynamics-and-angular`, the rotations are still visible in the form of **Wigner-$D$ matrices**, which derive from rotations of spin states. (See Equation {eq}`spherical-harmonics-wigner-d` for the relation between $Y^m_l$ and $d$ and $D$ functions.)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Helicity states and Wigner matrices\n", - ":class: dropdown\n", - "##### Helicity\n", - "Helicity $\\lambda$ is the projection of a particle's spin $S$ onto its direction of motion $\\vec{p}$. Mathematically, it's given by\n", - "```{math}\n", - "\\lambda = \\frac{S \\cdot \\vec{p}}{|\\vec{p}|} \\,.\n", - "```\n", - "For massless particles (photons in particular), helicity is a Lorentz-invariant quantity, meaning it remains unchanged under boosts along the direction of motion.\n", - "\n", - "##### Helicity state\n", - "Helicity states are eigenstates of the helicity operator. For a particle moving in the $z$ direction, the helicity operator is $S \\cdot \\hat{p}$.\n", - "\n", - "##### Wigner-$D$ matrix\n", - "Wigner-$D$ matrices are used in the helicity formalism to describe the rotation of spin states. These matrices depend on the Euler angles of the rotation and are denoted by \n", - "```{math}\n", - "D^j_{m'm} (\\alpha,\\beta,\\gamma) = e^{-im'\\alpha} d^j_{m'm}(\\beta) e^{-im\\gamma} \\,,\n", - "```\n", - "where $j$ is the spin, $m'$ and $m$ are the magnetic quantum numbers, $\\alpha, \\beta, \\gamma$ are the Euler angles of the rotation, and $d^j_{m'm}(\\beta)$ is an element of the orthogonal Wigner's (small) $d$-matrix.\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As can be seen in the amplitude model of Equation {eq}`model-dynamics-and-angular`, that the rotations (represented by $Y^m_l$) contain solid angles $\\Omega=\\phi,\\theta$ (see [spherical coordinates](#spherical-coordinate-system)). They have to be computed in the **helicity frame** of the resonance, which means we have to boost into each of the three subsystems. For instance, for the $a_2$ resonance ($A^{12}$), this would be a boost into subsystem $p_1+p_2$, plus a rotation such that the $z$ axis points in the direction of $p_1+p_2$. The **helicity angles** $\\phi$ and $\\theta$ can then easily be computed from the [spherical coordinates](#spherical-coordinate-system) of the (boosted) $p_1'$ or $p_2'$." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - ":::{admonition} Reference frames after boost\n", - ":class: full-width\n", - "\n", - "| CM frame | Helicity frame |\n", - "|---|---|\n", - "| Before boosting | After boosting into the $p_1+p_2$ subsystem |\n", - "| | |\n", - "\n", - "The **production plane** (cyan) is defined by the momenta of the incoming particles that participate in the production of a particular state or particle. In the figure, we have $a+b \\to (R\\to 1+2)+3$ (subsystem $p_1+p_2$), meaning that the production plane is spanned by the momenta of $a$ (or $b$) and $1$.\n", - "\n", - "The **decay plane** (magenta) is defined by the momenta of the decay products of a particle. In the figure, the production plane is spanned by the momenta of $1$ and $2$.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "#### Numerical angle computation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To calculate the helicity angles $\\theta$ and $\\phi$, we define functions for boosting a combination of boost and rotation (around the $y$ and $z$ axis)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def theta_helicity(p_i: MomentumNumpy4D, p_ij: MomentumNumpy4D) -> np.ndarray:\n", - " return _boost_to_helicity_frame(p_i, p_ij).theta\n", - "\n", - "\n", - "def phi_helicity(p_i: MomentumNumpy4D, p_ij: MomentumNumpy4D) -> np.ndarray:\n", - " return _boost_to_helicity_frame(p_i, p_ij).phi\n", - "\n", - "\n", - "def _boost_to_helicity_frame(\n", - " p_i: MomentumNumpy4D, p_ij: MomentumNumpy4D\n", - ") -> MomentumNumpy4D:\n", - " return p_i.rotateZ(-p_ij.phi).rotateY(-p_ij.theta).boostZ(-p_ij.beta)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "theta1_phsp = theta_helicity(p1_phsp, p12_phsp)\n", - "theta2_phsp = theta_helicity(p2_phsp, p23_phsp)\n", - "theta3_phsp = theta_helicity(p3_phsp, p31_phsp)\n", - "phi1_phsp = phi_helicity(p1_phsp, p12_phsp)\n", - "phi2_phsp = phi_helicity(p2_phsp, p23_phsp)\n", - "phi3_phsp = phi_helicity(p3_phsp, p31_phsp)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The distribution of the phase space sample over the kinematic variables is shown later in [this section](#my_section), together the generated data and weighted phase space (model)." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "(parameter-values)=\n", - "### Toy model parameter values" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now that we have a numerical implementation for the amplitude model of Equation {eq}`model-dynamics-and-angular` and a means to compute the kinematic variables appearing in that expression, we are ready to visualize the model. First, however, we have to decide on some toy values for the parameters in the model. The toy model parameter values can be obtained from the data file from the [Lecture 11 in STRONG2020 HaSP School](https://indico.ific.uv.es/event/6803/contributions/21223/). In this tutorial, the values are modified to make the structures in the Dalitz plot more visible." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "toy_parameters = dict(\n", - " a_minus2=0,\n", - " a_minus1=0.5,\n", - " a_0=3.5,\n", - " a_plus1=4,\n", - " a_plus2=2.5,\n", - " b_minus1=-1.5,\n", - " b_0=4,\n", - " b_plus1=0.5,\n", - " c_0=2.5,\n", - " M12=1.32,\n", - " Gamma12=0.1,\n", - " M23=1.24 + 0.3,\n", - " Gamma23=0.1,\n", - " M31=1.57 + 0.3,\n", - " Gamma31=0.1,\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Note that the masses $(M)$ and widths $(\\Gamma)$ are properties of the three resonances.\n", - "\n", - "$$\n", - "\\begin{array}{rcl}\n", - "M_{12} &=& M_{\\eta\\pi^0} = M_{a_2} \\\\\n", - "M_{23} &=& M_{\\pi^0p} = M_{\\Delta^+} \\\\\n", - "M_{31} &=& M_{\\eta p} = M_{N^*} \\\\\n", - "\\Gamma_{12} &=& \\Gamma_{\\eta\\pi^0} = \\Gamma_{a_2} \\\\\n", - "\\Gamma_{23} &=& \\Gamma_{\\pi^0p} = \\Gamma_{\\Delta^+} \\\\\n", - "\\Gamma_{31} &=& \\Gamma_{\\eta p} = \\Gamma_{N^*} \\\\\n", - "\\end{array}\n", - "$$ (mass-and-width-parameters)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "(spherical-harmonics-visualization)=\n", - "### Spherical harmonics" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now have a look at the real part and imaginary part of $\\sum a_m Y_2^m (\\Omega_1)$ as well as $\\sum b_m Y_1^m (\\Omega_2)$ in Equation {eq}`model-dynamics-and-angular`. For this, we define a grid of values for $\\phi$ and $\\theta$ over which to visualize the amplitudes." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "PHI, THETA = np.meshgrid(\n", - " np.linspace(-np.pi, +np.pi, num=1_000),\n", - " np.linspace(0, np.pi, num=1_000),\n", - ")\n", - "Z12 = Ylm12(PHI, THETA, **toy_parameters)\n", - "Z23 = Ylm23(PHI, THETA, **toy_parameters)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "fig, axes = plt.subplots(figsize=(10, 4), ncols=2, sharey=True, dpi=120)\n", - "cmap_real = axes[0].pcolormesh(\n", - " np.degrees(PHI), np.degrees(THETA), Z12.real, cmap=plt.cm.coolwarm\n", - ")\n", - "cmap_imag = axes[1].pcolormesh(\n", - " np.degrees(PHI), np.degrees(THETA), Z12.imag, cmap=plt.cm.coolwarm\n", - ")\n", - "\n", - "axes[0].set_xlabel(R\"$\\phi$ [deg]\")\n", - "axes[0].set_ylabel(R\"$\\theta$ [deg]\")\n", - "axes[0].set_title(R\"Re($\\sum a_m Y_2^m (\\Omega_1)$)\")\n", - "axes[0].set_ylabel(R\"$\\theta$ [deg]\")\n", - "axes[1].set_xlabel(R\"$\\phi$ [deg]\")\n", - "axes[1].set_title(R\"Im($\\sum a_m Y_2^m (\\Omega_1)$)\")\n", - "\n", - "cbar_real = fig.colorbar(cmap_real, ax=axes[0])\n", - "cbar_imag = fig.colorbar(cmap_imag, ax=axes[1])\n", - "\n", - "fig.subplots_adjust(wspace=0.4, hspace=0.4)\n", - "fig.tight_layout()\n", - "plt.rcParams.update({\"font.size\": 10})\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "![](https://github.com/user-attachments/assets/95592061-ce3f-4e57-8c35-0c2202e0d662)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "fig, axes = plt.subplots(figsize=(10, 4), ncols=2, sharey=True, dpi=120)\n", - "cmap_real = axes[0].pcolormesh(\n", - " np.degrees(PHI), np.degrees(THETA), Z23.real, cmap=plt.cm.coolwarm\n", - ")\n", - "cmap_imag = axes[1].pcolormesh(\n", - " np.degrees(PHI), np.degrees(THETA), Z23.imag, cmap=plt.cm.coolwarm\n", - ")\n", - "\n", - "axes[0].set_xlabel(R\"$\\phi$ [deg]\")\n", - "axes[0].set_ylabel(R\"$\\theta$ [deg]\")\n", - "axes[0].set_title(R\"Re($\\sum b_m Y_1^m (\\Omega_2)$)\")\n", - "axes[0].set_ylabel(R\"$\\theta$ [deg]\")\n", - "axes[1].set_xlabel(R\"$\\phi$ [deg]\")\n", - "axes[1].set_title(R\"Im($\\sum b_m Y_1^m (\\Omega_2)$)\")\n", - "\n", - "cbar_real = fig.colorbar(cmap_real, ax=axes[0])\n", - "cbar_imag = fig.colorbar(cmap_imag, ax=axes[1])\n", - "\n", - "fig.subplots_adjust(wspace=0.4, hspace=0.4)\n", - "fig.tight_layout()\n", - "plt.rcParams.update({\"font.size\": 10})\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "![](https://github.com/user-attachments/assets/95acd332-6aa2-4839-81f9-fa433258ace2)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Dalitz Plots of (sub)models" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Dalitz plots are ideal for visualizing the model over the Mandelstam variables $s_{12}$, $s_{23}$, and $s_{31}$. In the following, we plot each of the 'sub-models' separately: Breit-Wigner (only), Spherical Harmonics (only), and Breit-Wigner $\\times$ Spherical Harmonics.\n", - "\n", - "To integrate out the model dependence on the helicity angles, we plot the Dalitz plots as a histogram over the phase space sample, taking the computed intensities for each event in the sample as weight." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "mystnb": { - "code_prompt_show": "Compute intensities from models for plotting histograms" - }, - "tags": [ - "hide-cell", - "hide-output", - "hide-input" - ] - }, - "outputs": [], - "source": [ - "BW_intensities = BW_model(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " s31_phsp,\n", - " **toy_parameters,\n", - ")\n", - "SH_intensities = SH_model(\n", - " phi1_phsp,\n", - " theta1_phsp,\n", - " phi2_phsp,\n", - " theta2_phsp,\n", - " **toy_parameters,\n", - ")\n", - "BW_SH_intensities = BW_SH_model(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " s31_phsp,\n", - " phi1_phsp,\n", - " theta1_phsp,\n", - " phi2_phsp,\n", - " theta2_phsp,\n", - " **toy_parameters,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "scroll-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "fig, axes = plt.subplots(figsize=(12, 4), ncols=3, sharey=True)\n", - "ax1, ax2, ax3 = axes\n", - "fig.suptitle(\"Dalitz Plots of sub-models\")\n", - "ax1.set_title(\"BW model only\")\n", - "ax2.set_title(\"SH model only\")\n", - "ax3.set_title(R\"BW $\\times$ SH model\")\n", - "for ax in axes:\n", - " ax.set_xlabel(R\"$m^2(\\eta \\pi^0)$\")\n", - " ax.set_ylabel(R\"$m^2(\\pi^0 p)$\")\n", - "\n", - "hist1 = ax1.hist2d(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " bins=100,\n", - " cmap=\"jet\",\n", - " cmin=1e-6,\n", - " density=True,\n", - " weights=BW_intensities,\n", - ")\n", - "hist2 = ax2.hist2d(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " bins=100,\n", - " cmap=\"jet\",\n", - " cmin=1e-6,\n", - " density=True,\n", - " weights=SH_intensities,\n", - ")\n", - "hist3 = ax3.hist2d(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " bins=100,\n", - " cmap=\"jet\",\n", - " cmin=1e-6,\n", - " density=True,\n", - " weights=BW_SH_intensities,\n", - ")\n", - "\n", - "fig.colorbar(hist1[3], ax=ax1)\n", - "fig.colorbar(hist2[3], ax=ax2)\n", - "fig.colorbar(hist3[3], ax=ax3)\n", - "fig.tight_layout()\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "![fig](https://github.com/user-attachments/assets/eac2eb28-4cd5-4115-a9eb-af7d3bd1ec86)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(generate-data)=\n", - "## Data Generation" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In an actual amplitude analysis, we are now ready to 'fit' the model we formulated to a measured data distribution. In this tutorial, however, we generate the data ourselves and then perform a test fit using a starting model with some 'guessed' parameter values." - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Hit and miss intensity sample" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The data can be generated with a similar hit-and-miss strategy as we saw in [](#phase-space-generation). In this case, the hit-and-miss is performed over the intensities computed from the model.\n", - "\n", - "The following function generates a data sample based on a model by using hit-and-miss filter applied to the phase space sample. The output is a {obj}`tuple` of five four-momenta $(p_1, p_2, p_3, p_a, p_b)$." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def generate_data(\n", - " model: callable, size: int, bunch_size: int = 500_000, seed: int | None = None\n", - ") -> tuple[MomentumNumpy4D, ...]:\n", - " rng = np.random.default_rng(seed=seed)\n", - " generated_seed = rng.integers(low=0, high=1_000_000)\n", - " progress_bar = tqdm(total=size)\n", - " phase_space = generate_phsp_all(bunch_size, generated_seed)\n", - " data_sample = _hit_and_miss_filter(model, phase_space, generated_seed)\n", - " progress_bar.n = min(get_size(data_sample), size)\n", - " progress_bar.update(n=0)\n", - " while get_size_data(data_sample) < size:\n", - " phase_space = generate_phsp_all(bunch_size, generated_seed)\n", - " bunch = _hit_and_miss_filter(model, phase_space)\n", - " data_sample = concatenate_data(data_sample, bunch)\n", - " progress_bar.n = min(get_size(data_sample), size)\n", - " progress_bar.update(n=0)\n", - " progress_bar.close()\n", - " return remove_overflow_data(data_sample, size)\n", - "\n", - "\n", - "def _hit_and_miss_filter(\n", - " model: callable, phase_space: tuple[MomentumNumpy4D, ...], seed: int | None = None\n", - ") -> tuple[MomentumNumpy4D, ...]:\n", - " p1, p2, p3, pa, pb = phase_space\n", - " p12 = p1 + p2\n", - " p23 = p2 + p3\n", - " p31 = p3 + p1\n", - "\n", - " intensities: np.ndarray = model(\n", - " s12=p12.m2,\n", - " s23=p23.m2,\n", - " s31=p31.m2,\n", - " phi1=phi_helicity(p1, p12),\n", - " theta1=theta_helicity(p1, p12),\n", - " phi2=phi_helicity(p2, p23),\n", - " theta2=theta_helicity(p2, p23),\n", - " **toy_parameters,\n", - " )\n", - " rng = np.random.default_rng(seed=seed) # FIX seed is used here for reproducibility\n", - " random_intensities = rng.uniform(0, intensities.max(), size=intensities.shape)\n", - " selector = intensities > random_intensities\n", - " return (\n", - " p1[selector],\n", - " p2[selector],\n", - " p3[selector],\n", - " pa[selector],\n", - " pb[selector],\n", - " )\n", - "\n", - "\n", - "def get_size_data(\n", - " data: tuple[\n", - " MomentumNumpy4D,\n", - " MomentumNumpy4D,\n", - " MomentumNumpy4D,\n", - " MomentumNumpy4D,\n", - " MomentumNumpy4D,\n", - " ],\n", - "):\n", - " return len(data[0])\n", - "\n", - "\n", - "def concatenate_data(\n", - " data1: tuple[MomentumNumpy4D, ...],\n", - " data2: tuple[MomentumNumpy4D, ...],\n", - ") -> tuple[MomentumNumpy4D, ...]:\n", - " return tuple(concatenate_vectors((pi1, pj2)) for pi1, pj2 in zip(data1, data2))\n", - "\n", - "\n", - "def concatenate_vectors(vectors: tuple[MomentumNumpy4D]) -> MomentumNumpy4D:\n", - " return vector.array({\n", - " \"px\": np.concatenate([p.px for p in vectors]),\n", - " \"py\": np.concatenate([p.py for p in vectors]),\n", - " \"pz\": np.concatenate([p.pz for p in vectors]),\n", - " \"E\": np.concatenate([p.e for p in vectors]),\n", - " })\n", - "\n", - "\n", - "def remove_overflow_data(\n", - " data: tuple[MomentumNumpy4D, ...], size: int\n", - ") -> tuple[MomentumNumpy4D, ...]:\n", - " return tuple(momentum[:size] for momentum in data)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{admonition} Hit-and-miss data generation\n", - ":class: dropdown\n", - "\n", - "A few functions are defined above for generating data sample from filtering phase space samples of particle decays and converting them into four-momentum vectors. The main ideas of them are:\n", - "\n", - "- `generate_data()`: The main function used to generate the data\n", - " - It simulates events until it gathers a specified number of valid data points that meet the model’s conditions.\n", - " - Parameters:\n", - " - model: A callable that computes the intensity of events based on given kinematic variables.\n", - " - size: The desired number of valid events to generate.\n", - " - bunch_size: The number of events to generate per iteration if the initial batch doesn't meet the required size. Defaults to size if not provided.\n", - " - seed: Seed for random number generation to ensure reproducibility.\n", - " - Process:\n", - " - Initialize a random number generator.\n", - " - Generate an initial set of phase space data (phase_space) and filter it using the model to obtain valid events (data_sample).\n", - " - If the filtered data (data_sample) is less than the desired size, additional data is generated in batches (bunch_size) and filtered until the size requirement is met.\n", - " - The function ensures the returned dataset size matches exactly the requested size.\n", - " \n", - "\n", - "- `_hit_and_miss_filter`: The other important function that filters the generated phase space data\n", - " - using the hit-and-miss Monte Carlo method based on the model intensities.\n", - " - Parameters:\n", - " - model: The model function to calculate intensities.\n", - " - phase_space: Tuple containing momentum vectors of particles involved in the event.\n", - " - seed: Seed for random number generation.\n", - " - Process:\n", - " - Calculate invariant masses and other kinematic variables required by the model.\n", - " - Call the model with these variables to compute intensities for each event.\n", - " - Generate random numbers and select events where the model intensity is greater than a random threshold, simulating a probability proportional to the intensity.\n", - " \n", - "There are other assisted functions that are used in the functions above:\n", - "\n", - "- `get_size_data`:\n", - "\n", - " - Simply returns the number of events in the data, used to check if more data needs to be generated.\n", - "\n", - "- `concatenate_data`:\n", - " - Merges two datasets into one, maintaining the structure of tuples of momentum vectors.\n", - "\n", - "- `concatenate_vectors`:\n", - " - Helper function to concatenate individual momentum vectors within the data tuples.\n", - "\n", - "- `remove_overflow_data`:\n", - " - Trims the dataset to the desired size if it exceeds the specified number due to batch generation.\n", - "\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we can use the function `generate_data()` to generate our data sample:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "data_events = 100_000\n", - "data = generate_data(BW_SH_model, data_events, seed=0, bunch_size=1_000_000)\n", - "p1_data, p2_data, p3_data, pa_data, pb_data = data" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Convert four-momenta to kinematic variables" - }, - "tags": [ - "hide-input" - ] - }, - "outputs": [], - "source": [ - "p12_data = p1_data + p2_data\n", - "p23_data = p2_data + p3_data\n", - "p31_data = p3_data + p1_data\n", - "\n", - "s12_data = p12_data.m2\n", - "s23_data = p23_data.m2\n", - "s31_data = p31_data.m2\n", - "\n", - "theta1_CM_data = p1_data.theta\n", - "theta2_CM_data = p2_data.theta\n", - "theta3_CM_data = p3_data.theta\n", - "phi1_CM_data = p1_data.phi\n", - "phi2_CM_data = p2_data.phi\n", - "phi3_CM_data = p3_data.phi\n", - "\n", - "theta1_data = theta_helicity(p1_data, p12_data)\n", - "theta2_data = theta_helicity(p2_data, p23_data)\n", - "theta3_data = theta_helicity(p3_data, p31_data)\n", - "phi1_data = phi_helicity(p1_data, p12_data)\n", - "phi2_data = phi_helicity(p2_data, p23_data)\n", - "phi3_data = phi_helicity(p3_data, p31_data)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "### Dalitz plot of data distribution" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "fig, ax = plt.subplots(figsize=(6, 5))\n", - "vmax = 0.15\n", - "fig.suptitle(f\"Dalitz plot of generated data (with cut at max={vmax})\")\n", - "hist = ax.hist2d(\n", - " s12_data,\n", - " s23_data,\n", - " bins=100,\n", - " cmin=1e-20,\n", - " density=True,\n", - " cmap=\"jet\",\n", - " vmax=vmax,\n", - ")\n", - "ax.set_xlabel(r\"$m^2_{\\eta \\pi^0}$\")\n", - "ax.set_ylabel(r\"$m^2_{\\pi^0 p}$\")\n", - "cbar = fig.colorbar(hist[3], ax=ax)\n", - "fig.tight_layout()\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://github.com/user-attachments/assets/5f22c929-3ea8-49c9-a9ce-5db175b829f1)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "(my_section)=\n", - "### 1D projections" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "The 1D projection distribution of CM angles, the helicity angles and invariant mass of the model, phase space, and data are shown in this section." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "thetaCM_subtitles = [\n", - " R\"$\\cos (\\theta_{1}^{{CM}}) \\equiv \\cos (\\theta_{\\eta}^{{CM}})$\",\n", - " R\"$\\cos (\\theta_{2}^{{CM}}) \\equiv \\cos (\\theta_{\\pi^0}^{{CM}})$\",\n", - " R\"$\\cos (\\theta_{3}^{{CM}}) \\equiv \\cos (\\theta_{p}^{{CM}})$\",\n", - "]\n", - "phiCM_subtitles = [\n", - " R\"$\\phi_1^{CM} \\equiv \\phi_{\\eta}^{CM}$\",\n", - " R\"$\\phi_2^{CM} \\equiv \\phi_{\\pi^0}^{CM}$\",\n", - " R\"$\\phi_3^{CM} \\equiv \\phi_{p}^{CM}$\",\n", - "]\n", - "theta_subtitles = [\n", - " R\"$\\cos (\\theta_{1}^{{12}}) \\equiv \\cos (\\theta_{\\eta}^{{\\eta \\pi^0}})$\",\n", - " R\"$\\cos (\\theta_{2}^{{23}}) \\equiv \\cos (\\theta_{\\pi^0}^{{\\pi^0 p}})$\",\n", - " R\"$\\cos (\\theta_{3}^{{31}}) \\equiv \\cos (\\theta_{p}^{{p \\eta}})$\",\n", - "]\n", - "phi_subtitles = [\n", - " R\"$\\phi_1^{12} \\equiv \\phi_{\\eta}^{{\\eta \\pi^0}}$\",\n", - " R\"$\\phi_2^{23} \\equiv \\phi_{\\pi^0}^{{\\pi^0 p}}$\",\n", - " R\"$\\phi_3^{31} \\equiv \\phi_{p}^{{p \\eta}}$\",\n", - "]\n", - "mass_subtitles = [\n", - " R\"$m_{12} \\equiv m_{\\eta \\pi^0}$\",\n", - " R\"$m_{23} \\equiv m_{\\pi^0 p}$\",\n", - " R\"$m_{31} \\equiv m_{p \\eta}$\",\n", - "]\n", - "\n", - "fig, (thetaCM_ax, phiCM_ax, theta_ax, phi_ax, mass_ax) = plt.subplots(\n", - " figsize=(13, 18), ncols=3, nrows=5\n", - ")\n", - "for i, ax1 in enumerate(thetaCM_ax, 1):\n", - " ax1.set_title(thetaCM_subtitles[i - 1])\n", - " ax1.set_xticks([-1, 0, 1])\n", - "for i, ax2 in enumerate(phiCM_ax, 1):\n", - " ax2.set_title(phiCM_subtitles[i - 1])\n", - " ax2.set_xticks([-np.pi, 0, np.pi])\n", - " ax2.set_xticklabels([R\"-$\\pi$\", 0, R\"$\\pi$\"])\n", - "\n", - "for i, ax3 in enumerate(theta_ax, 1):\n", - " ax3.set_title(theta_subtitles[i - 1])\n", - " ax3.set_xticks([-1, 0, 1])\n", - "\n", - "for i, ax4 in enumerate(phi_ax, 1):\n", - " ax4.set_title(phi_subtitles[i - 1])\n", - " ax4.set_xticks([-np.pi, 0, np.pi])\n", - " ax4.set_xticklabels([R\"-$\\pi$\", 0, R\"$\\pi$\"])\n", - "\n", - "for i, ax5 in enumerate(mass_ax, 1):\n", - " ax5.set_title(mass_subtitles[i - 1])\n", - "\n", - "thetaCM_ax[0].hist(\n", - " np.cos(theta1_CM_phsp),\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[1].hist(\n", - " np.cos(theta2_CM_phsp),\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[2].hist(\n", - " np.cos(theta3_CM_phsp),\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[0].hist(\n", - " np.cos(theta1_CM_phsp),\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[1].hist(\n", - " np.cos(theta2_CM_phsp),\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[2].hist(\n", - " np.cos(theta3_CM_phsp),\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[0].hist(\n", - " np.cos(theta1_CM_phsp),\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[1].hist(\n", - " np.cos(theta2_CM_phsp),\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[2].hist(\n", - " np.cos(theta3_CM_phsp),\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[0].hist(\n", - " np.cos(theta1_CM_phsp),\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[1].hist(\n", - " np.cos(theta2_CM_phsp),\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "thetaCM_ax[2].hist(\n", - " np.cos(theta3_CM_phsp),\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "\n", - "\n", - "phiCM_ax[0].hist(\n", - " phi1_CM_phsp,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[1].hist(\n", - " phi2_CM_phsp,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[2].hist(\n", - " phi3_CM_phsp,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[0].hist(\n", - " phi1_CM_phsp,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[1].hist(\n", - " phi2_CM_phsp,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[2].hist(\n", - " phi3_CM_phsp,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[0].hist(\n", - " phi1_CM_phsp,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[1].hist(\n", - " phi2_CM_phsp,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[2].hist(\n", - " phi3_CM_phsp,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[0].hist(\n", - " phi1_CM_phsp,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[1].hist(\n", - " phi2_CM_phsp,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "phiCM_ax[2].hist(\n", - " phi3_CM_phsp,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "\n", - "\n", - "theta_ax[0].hist(\n", - " np.cos(theta1_phsp),\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "theta_ax[1].hist(\n", - " np.cos(theta2_phsp),\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "theta_ax[2].hist(\n", - " np.cos(theta3_phsp),\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "theta_ax[0].hist(\n", - " np.cos(theta1_phsp),\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "theta_ax[1].hist(\n", - " np.cos(theta2_phsp),\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "theta_ax[2].hist(\n", - " np.cos(theta3_phsp),\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "theta_ax[0].hist(\n", - " np.cos(theta1_phsp),\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "theta_ax[1].hist(\n", - " np.cos(theta2_phsp),\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "theta_ax[2].hist(\n", - " np.cos(theta3_phsp),\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "theta_ax[0].hist(\n", - " np.cos(theta1_phsp),\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "theta_ax[1].hist(\n", - " np.cos(theta2_phsp),\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "theta_ax[2].hist(\n", - " np.cos(theta3_phsp),\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "\n", - "\n", - "phi_ax[0].hist(\n", - " phi1_phsp,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "phi_ax[1].hist(\n", - " phi2_phsp,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "phi_ax[2].hist(\n", - " phi3_phsp,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "phi_ax[0].hist(\n", - " phi1_phsp,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "phi_ax[1].hist(\n", - " phi2_phsp,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "phi_ax[2].hist(\n", - " phi3_phsp,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "phi_ax[0].hist(\n", - " phi1_phsp,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "phi_ax[1].hist(\n", - " phi2_phsp,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "phi_ax[2].hist(\n", - " phi3_phsp,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "phi_ax[0].hist(\n", - " phi1_phsp,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "phi_ax[1].hist(\n", - " phi2_phsp,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "phi_ax[2].hist(\n", - " phi3_phsp,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "\n", - "\n", - "mass_ax[0].hist(\n", - " p12_phsp.m,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "mass_ax[1].hist(\n", - " p23_phsp.m,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "mass_ax[2].hist(\n", - " p31_phsp.m,\n", - " bins=100,\n", - " color=\"black\",\n", - " histtype=\"step\",\n", - " label=\"phsp\",\n", - " density=True,\n", - ")\n", - "mass_ax[0].hist(\n", - " p12_phsp.m,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "mass_ax[1].hist(\n", - " p23_phsp.m,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "mass_ax[2].hist(\n", - " p31_phsp.m,\n", - " bins=100,\n", - " weights=BW_intensities,\n", - " color=\"orange\",\n", - " histtype=\"step\",\n", - " label=\"only BW\",\n", - " density=True,\n", - ")\n", - "mass_ax[0].hist(\n", - " p12_phsp.m,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "mass_ax[1].hist(\n", - " p23_phsp.m,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "mass_ax[2].hist(\n", - " p31_phsp.m,\n", - " bins=100,\n", - " weights=SH_intensities,\n", - " color=\"green\",\n", - " histtype=\"step\",\n", - " label=\"only SH\",\n", - " density=True,\n", - ")\n", - "mass_ax[0].hist(\n", - " p12_phsp.m,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "mass_ax[1].hist(\n", - " p23_phsp.m,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "mass_ax[2].hist(\n", - " p31_phsp.m,\n", - " bins=100,\n", - " weights=BW_SH_intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - ")\n", - "\n", - "\n", - "thetaCM_ax[0].hist(\n", - " np.cos(theta1_CM_data),\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - " alpha=0.5,\n", - ")\n", - "thetaCM_ax[1].hist(\n", - " np.cos(theta2_CM_data),\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - " alpha=0.5,\n", - ")\n", - "thetaCM_ax[2].hist(\n", - " np.cos(theta3_CM_data),\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - " alpha=0.5,\n", - ")\n", - "\n", - "phiCM_ax[0].hist(\n", - " phi1_CM_data,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - " alpha=0.5,\n", - ")\n", - "phiCM_ax[1].hist(\n", - " phi2_CM_data,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - " alpha=0.5,\n", - ")\n", - "phiCM_ax[2].hist(\n", - " phi3_CM_data,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - " alpha=0.5,\n", - ")\n", - "\n", - "theta_ax[0].hist(\n", - " np.cos(theta1_data),\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "theta_ax[1].hist(\n", - " np.cos(theta2_data),\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "theta_ax[2].hist(\n", - " np.cos(theta3_data),\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "\n", - "phi_ax[0].hist(\n", - " phi1_data,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "phi_ax[1].hist(\n", - " phi2_data,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "phi_ax[2].hist(\n", - " phi3_data,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "\n", - "mass_ax[0].hist(\n", - " p12_data.m,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "mass_ax[1].hist(\n", - " p23_data.m,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "mass_ax[2].hist(\n", - " p31_data.m,\n", - " bins=100,\n", - " label=\"data\",\n", - " density=True,\n", - ")\n", - "\n", - "thetaCM_ax[0].legend()\n", - "phiCM_ax[0].legend()\n", - "thetaCM_ax[1].legend()\n", - "phiCM_ax[1].legend()\n", - "thetaCM_ax[2].legend()\n", - "phiCM_ax[2].legend()\n", - "\n", - "theta_ax[0].legend()\n", - "phi_ax[0].legend()\n", - "theta_ax[1].legend()\n", - "phi_ax[1].legend()\n", - "theta_ax[2].legend()\n", - "phi_ax[2].legend()\n", - "\n", - "mass_ax[0].legend()\n", - "mass_ax[1].legend()\n", - "mass_ax[2].legend()\n", - "\n", - "fig.tight_layout()\n", - "plt.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "![](https://github.com/user-attachments/assets/ced99aea-2fb4-4198-9e41-6b3377c5af23)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "(fit-model)=\n", - "## Fitting " - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "We now pretend that we do not know the parameter values for the data distribution. In this particular situation, the parameters (to be fitted) are mass $M$ ($M_{12}$, $M_{23}$, and $M_{31}$) and decay width $\\Gamma$ ($\\Gamma_{12}$, $\\Gamma_{23}$, and $\\Gamma_{31}$). We also guess the coefficient values, but note that we keep one coefficient 'fixed'. The reason is that the model is normalized through the likelihood estimator, so what matters for the fits are the _ratios_ between the parameters." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Guessed parameter values" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "guessed_parameters = dict(\n", - " M12=1.4,\n", - " Gamma12=0.15,\n", - " M23=1.59,\n", - " Gamma23=0.05,\n", - " M31=1.82,\n", - " Gamma31=0.2,\n", - " # a_minus2 is set to fixed and not changed during fitting\n", - " a_minus1=0.51,\n", - " a_0=3.49,\n", - " a_plus1=4.08,\n", - " a_plus2=2.6,\n", - " b_minus1=-1.2,\n", - " b_0=4.04,\n", - " b_plus1=0.41,\n", - " c_0=2.66,\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "mystnb": { - "code_prompt_show": "Function for plotting model" - }, - "tags": [ - "hide-input", - "scroll-input" - ] - }, - "outputs": [], - "source": [ - "def compare_model_to_data(parameters: dict):\n", - " fig, axes = plt.subplots(figsize=(13, 12), ncols=3, nrows=5)\n", - " fig.suptitle(\n", - " R\"CM angles, Helicity angles and invariant mass (Before fitting: starting\"\n", - " R\" parameters for BW $\\times$ SH model)\",\n", - " y=1,\n", - " )\n", - "\n", - " thetaCM_ax, phiCM_ax, theta_ax, phi_ax, mass_ax = axes\n", - " for i, ax in enumerate(thetaCM_ax, 1):\n", - " ax.set_title(thetaCM_subtitles[i - 1])\n", - " ax.set_xticks([-1, 0, 1])\n", - " for i, ax in enumerate(phiCM_ax, 1):\n", - " ax.set_title(phiCM_subtitles[i - 1])\n", - " ax.set_xticks([-np.pi, 0, np.pi])\n", - " ax.set_xticklabels([R\"-$\\pi$\", 0, R\"$\\pi$\"])\n", - " for i, ax in enumerate(theta_ax, 1):\n", - " ax.set_title(theta_subtitles[i - 1])\n", - " ax.set_xticks([-1, 0, 1])\n", - " for i, ax in enumerate(phi_ax, 1):\n", - " ax.set_title(phi_subtitles[i - 1])\n", - " ax.set_xticks([-np.pi, 0, np.pi])\n", - " ax.set_xticklabels([R\"-$\\pi$\", 0, R\"$\\pi$\"])\n", - " for i, ax in enumerate(mass_ax, 1):\n", - " ax.set_title(mass_subtitles[i - 1])\n", - "\n", - " data_kwargs = dict(bins=100, label=\"data\", density=True)\n", - " data_phi_kwargs = dict(bins=50, label=\"data\", density=True)\n", - "\n", - " thetaCM_ax[0].hist(np.cos(theta1_CM_data), **data_kwargs, alpha=0.5)\n", - " thetaCM_ax[1].hist(np.cos(theta2_CM_data), **data_kwargs, alpha=0.5)\n", - " thetaCM_ax[2].hist(np.cos(theta3_CM_data), **data_kwargs, alpha=0.5)\n", - " theta_ax[0].hist(np.cos(theta1_data), **data_kwargs)\n", - " theta_ax[1].hist(np.cos(theta2_data), **data_kwargs)\n", - " theta_ax[2].hist(np.cos(theta3_data), **data_kwargs)\n", - " phiCM_ax[0].hist(phi1_CM_data, **data_phi_kwargs, alpha=0.5)\n", - " phiCM_ax[1].hist(phi2_CM_data, **data_phi_kwargs, alpha=0.5)\n", - " phiCM_ax[2].hist(phi3_CM_data, **data_phi_kwargs, alpha=0.5)\n", - " phi_ax[0].hist(phi1_data, **data_phi_kwargs)\n", - " phi_ax[1].hist(phi2_data, **data_phi_kwargs)\n", - " phi_ax[2].hist(phi3_data, **data_phi_kwargs)\n", - " mass_ax[0].hist(p12_data.m, **data_kwargs)\n", - " mass_ax[1].hist(p23_data.m, **data_kwargs)\n", - " mass_ax[2].hist(p31_data.m, **data_kwargs)\n", - "\n", - " intensities = BW_SH_model(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " s31_phsp,\n", - " phi1_phsp,\n", - " theta1_phsp,\n", - " phi2_phsp,\n", - " theta2_phsp,\n", - " **{**toy_parameters, **parameters},\n", - " )\n", - " model_kwargs = dict(\n", - " bins=100,\n", - " weights=intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - " )\n", - " model_phi_kwargs = dict(\n", - " bins=50,\n", - " weights=intensities,\n", - " color=\"red\",\n", - " histtype=\"step\",\n", - " label=\"BW&SH\",\n", - " density=True,\n", - " )\n", - " thetaCM_ax[0].hist(np.cos(theta1_CM_phsp), **model_kwargs)\n", - " thetaCM_ax[1].hist(np.cos(theta2_CM_phsp), **model_kwargs)\n", - " thetaCM_ax[2].hist(np.cos(theta3_CM_phsp), **model_kwargs)\n", - " phiCM_ax[0].hist(phi1_CM_phsp, **model_phi_kwargs)\n", - " phiCM_ax[1].hist(phi2_CM_phsp, **model_phi_kwargs)\n", - " phiCM_ax[2].hist(phi3_CM_phsp, **model_phi_kwargs)\n", - " theta_ax[0].hist(np.cos(theta1_phsp), **model_kwargs)\n", - " theta_ax[1].hist(np.cos(theta2_phsp), **model_kwargs)\n", - " theta_ax[2].hist(np.cos(theta3_phsp), **model_kwargs)\n", - " phi_ax[0].hist(phi1_phsp, **model_phi_kwargs)\n", - " phi_ax[1].hist(phi2_phsp, **model_phi_kwargs)\n", - " phi_ax[2].hist(phi3_phsp, **model_phi_kwargs)\n", - " mass_ax[0].hist(p12_phsp.m, **model_kwargs)\n", - " mass_ax[1].hist(p23_phsp.m, **model_kwargs)\n", - " mass_ax[2].hist(p31_phsp.m, **model_kwargs)\n", - "\n", - " for ax in axes.flatten():\n", - " ax.legend()\n", - " fig.tight_layout()\n", - " plt.show()" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "compare_model_to_data(guessed_parameters)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "tags": [ - "full-width" - ] - }, - "source": [ - "![](https://github.com/user-attachments/assets/cfe00d5b-4cc1-451c-850e-ac298a548ba7)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Estimator: likelihood function" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To optimize the model parameters so that the model 'fits' the data distribution, we need to define a measure for the 'distance' between the model to the data. The common choice for complicated, multidimensional models like an amplitude model is the **unbinned negative log likelihood function** (NLL)." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "def unbinned_nll(**parameters: float) -> float:\n", - " phsp = (\n", - " s12_phsp,\n", - " s23_phsp,\n", - " s31_phsp,\n", - " phi1_phsp,\n", - " theta1_phsp,\n", - " phi2_phsp,\n", - " theta2_phsp,\n", - " )\n", - " data = (\n", - " s12_data,\n", - " s23_data,\n", - " s31_data,\n", - " phi1_data,\n", - " theta1_data,\n", - " phi2_data,\n", - " theta2_data,\n", - " )\n", - " model_integral = BW_SH_model(*phsp, **parameters).mean()\n", - " data_intensities = BW_SH_model(*data, **parameters)\n", - " likelihoods = data_intensities / model_integral\n", - " log_likelihood = np.log(likelihoods).sum()\n", - " return -log_likelihood" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip}\n", - "The phase space (phsp) is used for the normalization calculation for the unbinned log likelihood function\n", - ":::" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Optimizer: Minuit" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "As a final step, we find the parameter values that minimize the unbinned Log likelihood function, i.e., we optimize the model with respect to the data distribution. In Python, this can be performed by e.g. [`iminuit`](https://scikit-hep.org/iminuit) package." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "arg_names = tuple(guessed_parameters)\n", - "progress_bar = None\n", - "\n", - "\n", - "def wrapped_estimator(*args: float) -> float:\n", - " # minuit migrad expects positional-argument functions\n", - " global progress_bar\n", - " if progress_bar is None:\n", - " progress_bar = tqdm(desc=\"Running fit\")\n", - " new_kwargs = {\n", - " **toy_parameters,\n", - " **dict(zip(arg_names, args)),\n", - " }\n", - " estimator_value = unbinned_nll(**new_kwargs)\n", - " progress_bar.set_postfix({\"estimator\": estimator_value})\n", - " progress_bar.update()\n", - " return estimator_value" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "optimizer = Minuit(\n", - " wrapped_estimator,\n", - " **guessed_parameters,\n", - " name=guessed_parameters,\n", - ")\n", - "optimizer.errordef = Minuit.LIKELIHOOD" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "e1f140d4a81e41ec956d036a306f86b6", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Running fit: 0it [00:00, ?it/s]" - ] - }, - "metadata": {}, - "output_type": "display_data" - } - ], - "source": [ - "optimizer.migrad()\n", - "progress_bar.close()\n", - "progress_bar = None # reset progress bar" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "### Final fit result" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "scroll-input", - "hide-input", - "keep_output" - ] - }, - "outputs": [ - { - "data": { - "text/html": [ - "
\n", - "\n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "
ParameterToy valueStarting valueOptimized valueDifference
0a_minus20.00NaNNaN0.0%
1a_minus10.500.510.5183163.7%
2a_03.503.493.6431714.1%
3a_plus14.004.084.1723944.3%
4a_plus22.502.602.6714726.9%
5b_minus1-1.50-1.20-1.600444-6.7%
6b_04.004.044.1207373.0%
7b_plus10.500.410.57866015.7%
8c_02.502.662.5925613.7%
9M121.321.401.3203590.0%
10Gamma120.100.150.1005160.5%
11M231.541.591.5399560.0%
12Gamma230.100.050.0989871.0%
13M311.871.821.8699460.0%
14Gamma310.100.200.1003070.3%
\n", - "
" - ], - "text/plain": [ - " Parameter Toy value Starting value Optimized value Difference\n", - "0 a_minus2 0.00 NaN NaN 0.0%\n", - "1 a_minus1 0.50 0.51 0.518316 3.7%\n", - "2 a_0 3.50 3.49 3.643171 4.1%\n", - "3 a_plus1 4.00 4.08 4.172394 4.3%\n", - "4 a_plus2 2.50 2.60 2.671472 6.9%\n", - "5 b_minus1 -1.50 -1.20 -1.600444 -6.7%\n", - "6 b_0 4.00 4.04 4.120737 3.0%\n", - "7 b_plus1 0.50 0.41 0.578660 15.7%\n", - "8 c_0 2.50 2.66 2.592561 3.7%\n", - "9 M12 1.32 1.40 1.320359 0.0%\n", - "10 Gamma12 0.10 0.15 0.100516 0.5%\n", - "11 M23 1.54 1.59 1.539956 0.0%\n", - "12 Gamma23 0.10 0.05 0.098987 1.0%\n", - "13 M31 1.87 1.82 1.869946 0.0%\n", - "14 Gamma31 0.10 0.20 0.100307 0.3%" - ] - }, - "execution_count": null, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "optimized_parameters = {p.name: p.value for p in optimizer.params}\n", - "pd.DataFrame({\n", - " \"Parameter\": list(toy_parameters.keys()),\n", - " \"Toy value\": list(toy_parameters.values()),\n", - " \"Starting value\": [guessed_parameters.get(name) for name in toy_parameters],\n", - " \"Optimized value\": [optimized_parameters.get(name) for name in toy_parameters],\n", - " \"Difference\": [\n", - " f\"{100 * abs(value - optimized_parameters.get(name, value)) / value if value else 0:.1f}%\"\n", - " for name, value in toy_parameters.items()\n", - " ],\n", - "})" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "Now we see the fit results of CM angles, Helicity angles and invariant mass as a comparison of data in 1D projection plots" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "tags": [ - "hide-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['svg']\n", - "compare_model_to_data(optimized_parameters)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://github.com/user-attachments/assets/f34d1cbc-81c0-4535-bf94-5346bc4df00e)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "#### Dalitz plots of phsp, data, and fit result" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "To conclude our analysis in this document, we examine the Dalitz plots for the following cases: phase space, generated data, default parameters of the model, and fitted parameters of the model.\n", - "\n", - "- Phase Space Dalitz Plot:\n", - " - The plot displays a flat distribution, representing the kinematically allowed region for the reaction. This flat distribution serves as a baseline, showing no dynamical effects or resonant structures.\n", - "- Generated Data Dalitz Plot:\n", - " - This plot is derived from synthetic data that simulates the actual experimental conditions. It includes potential resonances (and background processes). The generated data reflects the physical phenomena expected in the real experimental scenario.\n", - "- Model with Default Parameters Dalitz Plot:\n", - " - Here, the Dalitz plot illustrates the model's predictions using the initial set of parameters. This provides a theoretical reference for comparison with actual data and the fitted model.\n", - "- Fitted Parameters Dalitz Plot:\n", - " - After performing the fitting procedure, this plot shows the model's predictions with optimized parameters. The resonances and their characteristics are clearly visible, indicating the physical information extracted from the data.\n", - " \n", - "By comparing these Dalitz plots, we can observe how the model evolves from theoretical predictions to a fitted solution that aligns with the experimental or generated data. The physical information about possible resonances becomes more apparent, allowing us to extract meaningful insights into the underlying processes of the reaction.\n", - "\n", - "This comprehensive analysis showcases the interplay between kinematics and dynamics in understanding particle interactions, highlighting the power of Partial Wave Analysis in uncovering the nuances of hadronic reactions." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "hide-input", - "scroll-input", - "remove-output" - ] - }, - "outputs": [], - "source": [ - "%config InlineBackend.figure_formats = ['png']\n", - "fig, (ax1, ax2) = plt.subplots(figsize=(12, 5), ncols=2, sharey=True)\n", - "fig.suptitle(R\"Dalitz Plots of parameter defaults and fitted parameters\")\n", - "\n", - "hist1 = ax1.hist2d(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " bins=100,\n", - " cmap=\"jet\",\n", - " cmin=1e-6,\n", - " density=True,\n", - " vmax=0.15,\n", - " weights=BW_SH_intensities,\n", - ")\n", - "ax1.set_title(R\"Parameters defaults BW $\\times$ SH model (with cut at max=0.15)\")\n", - "ax1.set_xlabel(R\"$m^2(\\eta \\pi)$\")\n", - "ax1.set_ylabel(R\"$m^2(\\pi p)$\")\n", - "\n", - "BW_SH_intensities_fitted = BW_SH_model(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " s31_phsp,\n", - " phi1_phsp,\n", - " theta1_phsp,\n", - " phi2_phsp,\n", - " theta2_phsp,\n", - " **{**toy_parameters, **optimized_parameters},\n", - ")\n", - "hist2 = ax2.hist2d(\n", - " s12_phsp,\n", - " s23_phsp,\n", - " bins=100,\n", - " cmap=\"jet\",\n", - " cmin=1e-6,\n", - " density=True,\n", - " vmax=0.15,\n", - " weights=BW_SH_intensities_fitted,\n", - ")\n", - "ax2.set_title(R\"Fitted parameters BW $\\times$ SH model (with cut at max=0.15)\")\n", - "ax2.set_xlabel(R\"$m^2(\\eta \\pi)$\")\n", - "ax2.set_ylabel(R\"$m^2(\\pi p)$\")\n", - "\n", - "cbar1 = fig.colorbar(hist1[3], ax=ax1)\n", - "cbar2 = fig.colorbar(hist2[3], ax=ax2)\n", - "cbar3 = fig.colorbar(hist3[3], ax=ax3)\n", - "\n", - "fig.tight_layout()\n", - "fig.show()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "![](https://github.com/user-attachments/assets/9e002900-9e84-414f-bf42-003cc2cf5968)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "jp-MarkdownHeadingCollapsed": true - }, - "source": [ - "## Summary" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "In this document, we have covered the general workflow of PWA, which consists of three major parts:\n", - "\n", - "As a summary, what have we done in this document consists of three major parts:\n", - "- Amplitude Model formulation\n", - " - Introduced the theoretical model for the amplitude of the reaction.\n", - " - Defined the mathematical expressions representing the partial waves and their respective contributions.\n", - " - Incorporated the relevant physical parameters, such as resonance properties. \n", - "- Phase space and data sample generation (or data analysis)\n", - " - Generated phase space distributions for the three-body decay processes.\n", - " - Created synthetic data samples or analyzed experimental data.\n", - " - Ensured that the data accurately reflects the kinematic constraints and dynamics of the reaction. \n", - "- Perform fitting, analyze and evaluate the results\n", - " - Performed a fit of the amplitude model to the generated or experimental data.\n", - " - Used statistical techniques to optimize the model parameters.\n", - " - Analyzed the fit results to extract physical parameters.\n", - " - Evaluated the quality of the fit and the consistency of the model with the data.\n", - " - Interpreted the physical significance of the extracted parameters and identified potential resonances and their properties." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "This structured approach provides a comprehensive understanding of both the reaction dynamics and kinematics, and helps in extracting meaningful physical insights from the data." - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - ":::{tip} Extra\n", - "have a look for another example of PWA for $J/\\psi \\to \\gamma \\pi^0 \\pi^0$ by using `ComPWA` at [here](https://tensorwaves.readthedocs.io/stable/amplitude-analysis/) .\n", - ":::" - ] - } - ], - "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.10.14" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} diff --git a/docs/reports.md b/docs/reports.md deleted file mode 100644 index 21fe9b3d..00000000 --- a/docs/reports.md +++ /dev/null @@ -1,66 +0,0 @@ -# Technical reports - -These pages are a collection of findings while working on ComPWA packages such as -{mod}`ampform`, {mod}`qrules`, and {mod}`tensorwaves`. Most of these findings were not -implemented, but may become relevant later on or could be useful to other frameworks as -well. - -```{include} report-inventory.md - -``` - - - - - - -```{toctree} -:glob: -:hidden: -report/* -``` - -````{dropdown} Execution times -```{nb-exec-table} -``` -```` diff --git a/docs/utils/.gitignore b/docs/utils/.gitignore deleted file mode 100644 index d392f0e8..00000000 --- a/docs/utils/.gitignore +++ /dev/null @@ -1 +0,0 @@ -*.jar diff --git a/environment.yml b/environment.yml index 34dd3d47..7e2206e6 100644 --- a/environment.yml +++ b/environment.yml @@ -2,8 +2,7 @@ name: compwa.github.io channels: - conda-forge dependencies: - - graphviz # for binder - - juliaup + - graphviz - pip - python==3.12.* - pip: diff --git a/pyproject.toml b/pyproject.toml index 4032d28c..1a5a8db1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -59,6 +59,7 @@ doc = [ "sphinx-hep-pdgref", "sphinx-pybtex-etal-style", "sphinx-remove-toctrees", + "sphinx-reredirects", "sphinx-thebe", "sphinx-togglebutton", "sphinxcontrib-bibtex >=2", @@ -104,9 +105,6 @@ local_scheme = "no-local-version" version_scheme = "post-release" write_to = "version.py" -[tool.isort] -profile = "black" - [tool.pixi.project] channels = ["conda-forge"] platforms = ["linux-64"] @@ -117,7 +115,6 @@ PYTHONHASHSEED = "0" [tool.pixi.dependencies] graphviz = "*" -juliaup = "*" python = "3.12.*" [tool.pixi.environments] @@ -149,21 +146,13 @@ cmd = """ sphinx-autobuild \ --open-browser \ --re-ignore '.*/.ipynb_checkpoints/.*' \ + --re-ignore '.*/.virtual_documents/.*' \ --re-ignore '.*/__pycache__/.*' \ - --re-ignore '.*\\.gitignore' \ - --re-ignore 'docs/.*\\.csv' \ - --re-ignore 'docs/.*\\.gif' \ - --re-ignore 'docs/.*\\.gv' \ - --re-ignore 'docs/.*\\.inv' \ - --re-ignore 'docs/.*\\.json' \ - --re-ignore 'docs/.*\\.pickle' \ --re-ignore 'docs/.*\\.png' \ --re-ignore 'docs/.*\\.svg' \ --re-ignore 'docs/.*\\.yaml' \ --re-ignore 'docs/.*\\.yml' \ --re-ignore 'docs/_build/.*' \ - --re-ignore docs/_static/exported_intensity_model.py \ - --re-ignore docs/report/data/* \ --watch docs \ docs/ docs/_build/html """ @@ -246,7 +235,6 @@ addopts = [ "--ignore=docs/adr/001/sympy.ipynb", "--ignore=docs/conf.py", "--nbmake-timeout=1200", - "-k not docs/report/019.ipynb", ] filterwarnings = [ "error", @@ -384,24 +372,6 @@ split-on-trailing-comma = false "PLC2701", "PLW2901", ] -"docs/report/002.ipynb" = ["F821"] -"docs/report/003.ipynb" = ["F821", "S404", "S603"] -"docs/report/004.ipynb" = ["S404", "S603"] -"docs/report/005.ipynb" = ["F821", "S404", "S603"] -"docs/report/006.ipynb" = ["S404", "S603"] -"docs/report/009.ipynb" = ["S404", "S603"] -"docs/report/010.ipynb" = ["S404", "S603"] -"docs/report/011.ipynb" = ["F821"] -"docs/report/017.ipynb" = ["S404", "S603"] -"docs/report/020.ipynb" = ["F821"] -"docs/report/021.ipynb" = ["FIX004", "S404", "S603"] -"docs/report/022.ipynb" = ["F821"] -"docs/report/024.ipynb" = ["F821", "S102"] -"docs/report/025.ipynb" = ["S404", "S603"] -"docs/report/027.ipynb" = ["PLW1510", "S404", "S603"] -"docs/report/028.ipynb" = ["PERF203"] -"pin_nb_requirements.py" = ["INP001", "PLW2901"] -"setup.py" = ["D100"] [tool.ruff.lint.pydocstyle] convention = "google" diff --git a/tox.ini b/tox.ini index 8919e5c7..a9c5d852 100644 --- a/tox.ini +++ b/tox.ini @@ -30,21 +30,13 @@ commands = sphinx-autobuild \ --open-browser \ --re-ignore .*/.ipynb_checkpoints/.* \ + --re-ignore .*/.virtual_documents/.* \ --re-ignore .*/__pycache__/.* \ - --re-ignore .*\.gitignore \ - --re-ignore docs/.*\.csv \ - --re-ignore docs/.*\.gif \ - --re-ignore docs/.*\.gv \ - --re-ignore docs/.*\.inv \ - --re-ignore docs/.*\.json \ - --re-ignore docs/.*\.pickle \ --re-ignore docs/.*\.png \ --re-ignore docs/.*\.svg \ --re-ignore docs/.*\.yaml \ --re-ignore docs/.*\.yml \ --re-ignore docs/_build/.* \ - --re-ignore docs/_static/exported_intensity_model.py \ - --re-ignore docs/report/data/* \ --watch docs \ docs/ docs/_build/html description = @@ -81,21 +73,13 @@ commands = sphinx-autobuild \ --open-browser \ --re-ignore .*/.ipynb_checkpoints/.* \ + --re-ignore .*/.virtual_documents/.* \ --re-ignore .*/__pycache__/.* \ - --re-ignore .*\.gitignore \ - --re-ignore docs/.*\.csv \ - --re-ignore docs/.*\.gif \ - --re-ignore docs/.*\.gv \ - --re-ignore docs/.*\.inv \ - --re-ignore docs/.*\.json \ - --re-ignore docs/.*\.pickle \ --re-ignore docs/.*\.png \ --re-ignore docs/.*\.svg \ --re-ignore docs/.*\.yaml \ --re-ignore docs/.*\.yml \ --re-ignore docs/_build/.* \ - --re-ignore docs/_static/exported_intensity_model.py \ - --re-ignore docs/report/data/* \ --watch docs \ docs/ docs/_build/html description =