diff --git a/.constraints/py3.7.txt b/.constraints/py3.7.txt index 627300ef..3ebc709e 100644 --- a/.constraints/py3.7.txt +++ b/.constraints/py3.7.txt @@ -10,17 +10,17 @@ aquirdturtle-collapsible-headings==3.1.0 argon2-cffi==21.3.0 argon2-cffi-bindings==21.2.0 attrs==21.4.0 -babel==2.10.1 +babel==2.10.3 backcall==0.2.0 beautifulsoup4==4.11.1 black==22.3.0 bleach==5.0.0 -certifi==2022.5.18.1 +certifi==2022.6.15 cffi==1.15.0 cfgv==3.3.1 charset-normalizer==2.0.12 click==8.1.3 -colorama==0.4.4 +colorama==0.4.5 cycler==0.11.0 debugpy==1.6.0 decorator==5.1.1 @@ -30,20 +30,20 @@ docutils==0.17.1 entrypoints==0.4 execnet==1.9.0 fastjsonschema==2.15.3 -filelock==3.7.0 +filelock==3.7.1 fonttools==4.33.3 graphviz==0.20 greenlet==1.1.2 identify==2.5.1 idna==3.3 imagesize==1.3.0 -importlib-metadata==4.11.4 +importlib-metadata==4.12.0 iniconfig==1.1.1 -ipykernel==6.13.0 +ipykernel==6.15.0 ipympl==0.9.1 -ipython==7.33.0 +ipython==7.34.0 ipython-genutils==0.2.0 -ipywidgets==7.7.0 +ipywidgets==7.7.1 isort==5.10.1 jedi==0.18.1 jinja2==3.1.2 @@ -51,18 +51,18 @@ json5==0.9.8 jsonschema==3.2.0 jupyter==1.0.0 jupyter-cache==0.5.0 -jupyter-client==7.3.1 -jupyter-console==6.4.3 +jupyter-client==7.3.4 +jupyter-console==6.4.4 jupyter-core==4.10.0 -jupyter-server==1.17.0 -jupyterlab==3.4.2 +jupyter-server==1.18.0 +jupyterlab==3.4.3 jupyterlab-code-formatter==1.4.11 jupyterlab-markup==1.1.0 jupyterlab-myst==0.1.6 jupyterlab-pygments==0.2.2 jupyterlab-server==2.14.0 -jupyterlab-widgets==1.1.0 -kiwisolver==1.4.2 +jupyterlab-widgets==1.1.1 +kiwisolver==1.4.3 latexcodec==2.0.1 livereload==2.6.3 markdown-it-py==2.1.0 @@ -72,18 +72,18 @@ matplotlib-inline==0.1.3 mdit-py-plugins==0.3.0 mdurl==0.1.1 mistune==0.8.4 -mpl-interactions==0.21.0 +mpl-interactions==0.22.0 mypy-extensions==0.4.3 -myst-nb==0.15.0 -myst-parser==0.17.2 +myst-nb==0.16.0 +myst-parser==0.18.0 nbclassic==0.3.7 nbclient==0.5.13 nbconvert==6.5.0 nbformat==5.4.0 nbmake==1.3.0 nest-asyncio==1.5.5 -nodeenv==1.6.0 -notebook==6.4.11 +nodeenv==1.7.0 +notebook==6.4.12 notebook-shim==0.1.0 numpy==1.21.6 packaging==21.3 @@ -102,7 +102,7 @@ psutil==5.9.1 ptyprocess==0.7.0 py==1.11.0 pybtex==0.24.0 -pybtex-docutils==1.0.1 +pybtex-docutils==1.0.2 pycparser==2.21 pydantic==1.9.1 pydata-sphinx-theme==0.8.1 @@ -115,10 +115,10 @@ pytest-xdist==2.5.0 python-dateutil==2.8.2 pytz==2022.1 pyyaml==6.0 -pyzmq==23.0.0 -qtconsole==5.3.0 +pyzmq==23.2.0 +qtconsole==5.3.1 qtpy==2.1.0 -requests==2.27.1 +requests==2.28.0 requests-file==1.5.1 send2trash==1.8.0 six==1.16.0 @@ -128,7 +128,7 @@ soupsieve==2.3.2.post1 sphinx==4.3.2 ; python_version < "3.8.0" sphinx-autobuild==2021.3.14 sphinx-book-theme==0.3.2 -sphinx-codeautolink==0.10.0 +sphinx-codeautolink==0.11.0 sphinx-comments==0.0.3 sphinx-copybutton==0.5.0 sphinx-data-viewer==0.1.2 @@ -146,24 +146,24 @@ sphinxcontrib-needs==0.7.9 sphinxcontrib-plantuml==0.23 sphinxcontrib-qthelp==1.0.3 sphinxcontrib-serializinghtml==1.1.5 -sqlalchemy==1.4.36 -tabulate==0.8.9 +sqlalchemy==1.4.39 +tabulate==0.8.10 terminado==0.15.0 tinycss2==1.1.1 toml==0.10.2 tomli==2.0.1 tornado==6.1 tox==3.25.0 -traitlets==5.2.1.post0 +traitlets==5.3.0 typed-ast==1.5.4 typing-extensions==4.2.0 urllib3==1.26.9 -virtualenv==20.14.1 +virtualenv==20.15.0 wcwidth==0.2.5 webencodings==0.5.1 -websocket-client==1.3.2 +websocket-client==1.3.3 wheel==0.37.1 -widgetsnbextension==3.6.0 +widgetsnbextension==3.6.1 zipp==3.8.0 # The following packages are considered to be unsafe in a requirements file: diff --git a/.constraints/py3.8.txt b/.constraints/py3.8.txt index fd3f2acc..6038e678 100644 --- a/.constraints/py3.8.txt +++ b/.constraints/py3.8.txt @@ -11,17 +11,17 @@ argon2-cffi==21.3.0 argon2-cffi-bindings==21.2.0 asttokens==2.0.5 attrs==21.4.0 -babel==2.10.1 +babel==2.10.3 backcall==0.2.0 beautifulsoup4==4.11.1 black==22.3.0 bleach==5.0.0 -certifi==2022.5.18.1 +certifi==2022.6.15 cffi==1.15.0 cfgv==3.3.1 charset-normalizer==2.0.12 click==8.1.3 -colorama==0.4.4 +colorama==0.4.5 cycler==0.11.0 debugpy==1.6.0 decorator==5.1.1 @@ -32,20 +32,20 @@ entrypoints==0.4 execnet==1.9.0 executing==0.8.3 fastjsonschema==2.15.3 -filelock==3.7.0 +filelock==3.7.1 fonttools==4.33.3 graphviz==0.20 greenlet==1.1.2 identify==2.5.1 idna==3.3 imagesize==1.3.0 -importlib-metadata==4.11.4 +importlib-metadata==4.12.0 iniconfig==1.1.1 -ipykernel==6.13.0 +ipykernel==6.15.0 ipympl==0.9.1 -ipython==8.3.0 +ipython==8.4.0 ipython-genutils==0.2.0 -ipywidgets==7.7.0 +ipywidgets==7.7.1 isort==5.10.1 jedi==0.18.1 jinja2==3.1.2 @@ -53,18 +53,18 @@ json5==0.9.8 jsonschema==3.2.0 jupyter==1.0.0 jupyter-cache==0.5.0 -jupyter-client==7.3.1 -jupyter-console==6.4.3 +jupyter-client==7.3.4 +jupyter-console==6.4.4 jupyter-core==4.10.0 -jupyter-server==1.17.0 -jupyterlab==3.4.2 +jupyter-server==1.18.0 +jupyterlab==3.4.3 jupyterlab-code-formatter==1.4.11 jupyterlab-markup==1.1.0 jupyterlab-myst==0.1.6 jupyterlab-pygments==0.2.2 jupyterlab-server==2.14.0 -jupyterlab-widgets==1.1.0 -kiwisolver==1.4.2 +jupyterlab-widgets==1.1.1 +kiwisolver==1.4.3 latexcodec==2.0.1 livereload==2.6.3 markdown-it-py==2.1.0 @@ -74,20 +74,20 @@ matplotlib-inline==0.1.3 mdit-py-plugins==0.3.0 mdurl==0.1.1 mistune==0.8.4 -mpl-interactions==0.21.0 +mpl-interactions==0.22.0 mypy-extensions==0.4.3 -myst-nb==0.15.0 -myst-parser==0.17.2 +myst-nb==0.16.0 +myst-parser==0.18.0 nbclassic==0.3.7 nbclient==0.5.13 nbconvert==6.5.0 nbformat==5.4.0 nbmake==1.3.0 nest-asyncio==1.5.5 -nodeenv==1.6.0 -notebook==6.4.11 +nodeenv==1.7.0 +notebook==6.4.12 notebook-shim==0.1.0 -numpy==1.22.4 +numpy==1.23.0 packaging==21.3 pandocfilters==1.5.0 parso==0.8.3 @@ -105,7 +105,7 @@ ptyprocess==0.7.0 pure-eval==0.2.2 py==1.11.0 pybtex==0.24.0 -pybtex-docutils==1.0.1 +pybtex-docutils==1.0.2 pycparser==2.21 pydantic==1.9.1 pydata-sphinx-theme==0.8.1 @@ -118,10 +118,10 @@ pytest-xdist==2.5.0 python-dateutil==2.8.2 pytz==2022.1 pyyaml==6.0 -pyzmq==23.0.0 -qtconsole==5.3.0 +pyzmq==23.2.0 +qtconsole==5.3.1 qtpy==2.1.0 -requests==2.27.1 +requests==2.28.0 requests-file==1.5.1 send2trash==1.8.0 six==1.16.0 @@ -131,7 +131,7 @@ soupsieve==2.3.2.post1 sphinx==4.5.0 sphinx-autobuild==2021.3.14 sphinx-book-theme==0.3.2 -sphinx-codeautolink==0.10.0 +sphinx-codeautolink==0.11.0 sphinx-comments==0.0.3 sphinx-copybutton==0.5.0 sphinx-data-viewer==0.1.2 @@ -149,24 +149,24 @@ sphinxcontrib-needs==0.7.9 sphinxcontrib-plantuml==0.23 sphinxcontrib-qthelp==1.0.3 sphinxcontrib-serializinghtml==1.1.5 -sqlalchemy==1.4.36 -stack-data==0.2.0 -tabulate==0.8.9 +sqlalchemy==1.4.39 +stack-data==0.3.0 +tabulate==0.8.10 terminado==0.15.0 tinycss2==1.1.1 toml==0.10.2 tomli==2.0.1 tornado==6.1 tox==3.25.0 -traitlets==5.2.1.post0 +traitlets==5.3.0 typing-extensions==4.2.0 urllib3==1.26.9 -virtualenv==20.14.1 +virtualenv==20.15.0 wcwidth==0.2.5 webencodings==0.5.1 -websocket-client==1.3.2 +websocket-client==1.3.3 wheel==0.37.1 -widgetsnbextension==3.6.0 +widgetsnbextension==3.6.1 zipp==3.8.0 # The following packages are considered to be unsafe in a requirements file: diff --git a/.constraints/py3.9.txt b/.constraints/py3.9.txt index d5e4f298..3e99481c 100644 --- a/.constraints/py3.9.txt +++ b/.constraints/py3.9.txt @@ -11,17 +11,17 @@ argon2-cffi==21.3.0 argon2-cffi-bindings==21.2.0 asttokens==2.0.5 attrs==21.4.0 -babel==2.10.1 +babel==2.10.3 backcall==0.2.0 beautifulsoup4==4.11.1 black==22.3.0 bleach==5.0.0 -certifi==2022.5.18.1 +certifi==2022.6.15 cffi==1.15.0 cfgv==3.3.1 charset-normalizer==2.0.12 click==8.1.3 -colorama==0.4.4 +colorama==0.4.5 cycler==0.11.0 debugpy==1.6.0 decorator==5.1.1 @@ -32,20 +32,20 @@ entrypoints==0.4 execnet==1.9.0 executing==0.8.3 fastjsonschema==2.15.3 -filelock==3.7.0 +filelock==3.7.1 fonttools==4.33.3 graphviz==0.20 greenlet==1.1.2 identify==2.5.1 idna==3.3 imagesize==1.3.0 -importlib-metadata==4.11.4 +importlib-metadata==4.12.0 iniconfig==1.1.1 -ipykernel==6.13.0 +ipykernel==6.15.0 ipympl==0.9.1 -ipython==8.3.0 +ipython==8.4.0 ipython-genutils==0.2.0 -ipywidgets==7.7.0 +ipywidgets==7.7.1 isort==5.10.1 jedi==0.18.1 jinja2==3.1.2 @@ -53,18 +53,18 @@ json5==0.9.8 jsonschema==3.2.0 jupyter==1.0.0 jupyter-cache==0.5.0 -jupyter-client==7.3.1 -jupyter-console==6.4.3 +jupyter-client==7.3.4 +jupyter-console==6.4.4 jupyter-core==4.10.0 -jupyter-server==1.17.0 -jupyterlab==3.4.2 +jupyter-server==1.18.0 +jupyterlab==3.4.3 jupyterlab-code-formatter==1.4.11 jupyterlab-markup==1.1.0 jupyterlab-myst==0.1.6 jupyterlab-pygments==0.2.2 jupyterlab-server==2.14.0 -jupyterlab-widgets==1.1.0 -kiwisolver==1.4.2 +jupyterlab-widgets==1.1.1 +kiwisolver==1.4.3 latexcodec==2.0.1 livereload==2.6.3 markdown-it-py==2.1.0 @@ -74,20 +74,20 @@ matplotlib-inline==0.1.3 mdit-py-plugins==0.3.0 mdurl==0.1.1 mistune==0.8.4 -mpl-interactions==0.21.0 +mpl-interactions==0.22.0 mypy-extensions==0.4.3 -myst-nb==0.15.0 -myst-parser==0.17.2 +myst-nb==0.16.0 +myst-parser==0.18.0 nbclassic==0.3.7 nbclient==0.5.13 nbconvert==6.5.0 nbformat==5.4.0 nbmake==1.3.0 nest-asyncio==1.5.5 -nodeenv==1.6.0 -notebook==6.4.11 +nodeenv==1.7.0 +notebook==6.4.12 notebook-shim==0.1.0 -numpy==1.22.4 +numpy==1.23.0 packaging==21.3 pandocfilters==1.5.0 parso==0.8.3 @@ -105,7 +105,7 @@ ptyprocess==0.7.0 pure-eval==0.2.2 py==1.11.0 pybtex==0.24.0 -pybtex-docutils==1.0.1 +pybtex-docutils==1.0.2 pycparser==2.21 pydantic==1.9.1 pydata-sphinx-theme==0.8.1 @@ -118,10 +118,10 @@ pytest-xdist==2.5.0 python-dateutil==2.8.2 pytz==2022.1 pyyaml==6.0 -pyzmq==23.0.0 -qtconsole==5.3.0 +pyzmq==23.2.0 +qtconsole==5.3.1 qtpy==2.1.0 -requests==2.27.1 +requests==2.28.0 requests-file==1.5.1 send2trash==1.8.0 six==1.16.0 @@ -131,7 +131,7 @@ soupsieve==2.3.2.post1 sphinx==4.5.0 sphinx-autobuild==2021.3.14 sphinx-book-theme==0.3.2 -sphinx-codeautolink==0.10.0 +sphinx-codeautolink==0.11.0 sphinx-comments==0.0.3 sphinx-copybutton==0.5.0 sphinx-data-viewer==0.1.2 @@ -149,24 +149,24 @@ sphinxcontrib-needs==0.7.9 sphinxcontrib-plantuml==0.23 sphinxcontrib-qthelp==1.0.3 sphinxcontrib-serializinghtml==1.1.5 -sqlalchemy==1.4.36 -stack-data==0.2.0 -tabulate==0.8.9 +sqlalchemy==1.4.39 +stack-data==0.3.0 +tabulate==0.8.10 terminado==0.15.0 tinycss2==1.1.1 toml==0.10.2 tomli==2.0.1 tornado==6.1 tox==3.25.0 -traitlets==5.2.1.post0 +traitlets==5.3.0 typing-extensions==4.2.0 urllib3==1.26.9 -virtualenv==20.14.1 +virtualenv==20.15.0 wcwidth==0.2.5 webencodings==0.5.1 -websocket-client==1.3.2 +websocket-client==1.3.3 wheel==0.37.1 -widgetsnbextension==3.6.0 +widgetsnbextension==3.6.1 zipp==3.8.0 # The following packages are considered to be unsafe in a requirements file: diff --git a/.editorconfig b/.editorconfig index 9db907f1..d0ac7e38 100644 --- a/.editorconfig +++ b/.editorconfig @@ -7,10 +7,7 @@ indent_style = space insert_final_newline = true trim_trailing_whitespace = true -[*.ipynb] -indent_size = 1 - -[*.{py,pyi,toml}] +[*.{ipynb,py,toml}] indent_size = 4 # when adding words through vscode, this is the resulting output format diff --git a/.github/workflows/cd.yml b/.github/workflows/cd.yml index ad9c3e8a..17f41c94 100644 --- a/.github/workflows/cd.yml +++ b/.github/workflows/cd.yml @@ -38,7 +38,7 @@ jobs: steps: - uses: actions/checkout@v3 - name: Set up Python - uses: actions/setup-python@v3 + uses: actions/setup-python@v4 with: python-version: "3.8" - name: Install dependencies diff --git a/.github/workflows/ci-docs.yml b/.github/workflows/ci-docs.yml index 72602906..ef1a2115 100644 --- a/.github/workflows/ci-docs.yml +++ b/.github/workflows/ci-docs.yml @@ -20,7 +20,7 @@ jobs: steps: - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v3 + uses: actions/setup-python@v4 with: python-version: "3.8" - name: Install dependencies diff --git a/.github/workflows/ci-tests.yml b/.github/workflows/ci-tests.yml index 993ab610..f5e904e5 100644 --- a/.github/workflows/ci-tests.yml +++ b/.github/workflows/ci-tests.yml @@ -31,4 +31,4 @@ jobs: pip install -c .constraints/py3.8.txt .[test] - name: Run notebooks with pytest (diff only) if: steps.diff.outputs.modified-notebooks != '' - run: pytest --nbmake --nbmake-timeout=900 $(cat diff.txt) + run: pytest --nbmake --nbmake-timeout=1200 $(cat diff.txt) diff --git a/.github/workflows/linkcheck.yml b/.github/workflows/linkcheck.yml index f4b9f712..67a4376e 100644 --- a/.github/workflows/linkcheck.yml +++ b/.github/workflows/linkcheck.yml @@ -4,10 +4,12 @@ on: push: branches: - main + - epic/* - "[0-9]+.[0-9]+.x" pull_request: branches: - main + - epic/* - "[0-9]+.[0-9]+.x" workflow_dispatch: @@ -18,7 +20,7 @@ jobs: steps: - uses: actions/checkout@v3 - name: Set up Python 3.8 - uses: actions/setup-python@v3 + uses: actions/setup-python@v4 with: python-version: "3.8" - name: Install dependencies diff --git a/.github/workflows/notebooks.yml b/.github/workflows/notebooks.yml index c836495d..640c6265 100644 --- a/.github/workflows/notebooks.yml +++ b/.github/workflows/notebooks.yml @@ -22,4 +22,4 @@ jobs: sudo apt-get -y install graphviz pip install -c .constraints/py3.8.txt .[test] - run: | - pytest --nbmake --nbmake-timeout=900 ${{ github.event.inputs.notebook-selector }} + pytest --nbmake --nbmake-timeout=1200 ${{ github.event.inputs.notebook-selector }} diff --git a/.github/workflows/requirements-pr.yml b/.github/workflows/requirements-pr.yml index 7ab4bbee..0962efb8 100644 --- a/.github/workflows/requirements-pr.yml +++ b/.github/workflows/requirements-pr.yml @@ -5,6 +5,7 @@ on: branches: - main - epic/* + - "[0-9]+.[0-9]+.x" jobs: diff: diff --git a/.markdownlint.json b/.markdownlint.json index 453eed2e..23bb7fa6 100644 --- a/.markdownlint.json +++ b/.markdownlint.json @@ -1,4 +1,5 @@ { + "MD013": { "line_length": 88 }, "MD026": { "punctuation": ".,;:。，；：！" }, "MD033": { "allowed_elements": ["br"] diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 26f0795b..23d5b7df 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,7 +13,7 @@ repos: - id: check-useless-excludes - repo: https://github.com/pre-commit/pre-commit-hooks - rev: v4.2.0 + rev: v4.3.0 hooks: - id: check-ast - id: check-case-conflict @@ -48,7 +48,7 @@ repos: )$ - repo: https://github.com/ComPWA/repo-maintenance - rev: 0.0.128 + rev: 0.0.132 hooks: - id: check-dev-files args: @@ -69,7 +69,7 @@ repos: - id: blacken-docs - repo: https://github.com/streetsidesoftware/cspell-cli - rev: v5.9.1 + rev: v6.0.0 hooks: - id: cspell @@ -79,6 +79,7 @@ repos: - id: editorconfig-checker exclude: > (?x)^( + .*\.ipynb| .*\.md| .*\.py| LICENSE @@ -130,17 +131,17 @@ repos: metadata.varInspector - repo: https://github.com/pre-commit/mirrors-prettier - rev: v2.6.2 + rev: v2.7.1 hooks: - id: prettier - repo: https://github.com/ComPWA/mirrors-pyright - rev: v1.1.247 + rev: v1.1.255 hooks: - id: pyright - repo: https://github.com/asottile/pyupgrade - rev: v2.32.1 + rev: v2.34.0 hooks: - id: pyupgrade args: diff --git a/.prettierrc b/.prettierrc index 164d0b8b..ebfaf680 100644 --- a/.prettierrc +++ b/.prettierrc @@ -1,2 +1,2 @@ -printWidth: 79 +printWidth: 88 proseWrap: always diff --git a/.vscode/settings.json b/.vscode/settings.json index 2e5fed8c..dfaa1c43 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -19,7 +19,7 @@ }, "cSpell.enabled": true, "editor.formatOnSave": true, - "editor.rulers": [80], + "editor.rulers": [88], "files.watcherExclude": { "**/*_cache/**": true, "**/.eggs/**": true, @@ -42,7 +42,7 @@ "python.linting.pylamaEnabled": false, "python.linting.pylintEnabled": false, "python.testing.pytestEnabled": false, - "rewrap.wrappingColumn": 79, + "rewrap.wrappingColumn": 88, "search.exclude": { "**/tests/**/__init__.py": true, "*/.pydocstyle": true, diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 23b3180d..ebe740b6 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -3,7 +3,6 @@ [![Open in Visual Studio Code](https://open.vscode.dev/badges/open-in-vscode.svg)](https://open.vscode.dev/ComPWA/compwa-org) [![GitPod](https://img.shields.io/badge/gitpod-open-blue?logo=gitpod)](https://gitpod.io/#https://github.com/ComPWA/compwa-org) -This repository is part of the -[ComPWA Organization](https://github.com/ComPWA). For more information about -how to contribute to the packages, go to +This repository is part of the [ComPWA Organization](https://github.com/ComPWA). For +more information about how to contribute to the packages, go to [pwa.rtfd.io/develop.html](https://pwa.readthedocs.io/develop.html)! diff --git a/docs/_templates/genindex.html b/docs/_templates/genindex.html index 21849ec3..abea27ff 100644 --- a/docs/_templates/genindex.html +++ b/docs/_templates/genindex.html @@ -7,12 +7,11 @@ {%- for ismain, link in links[1:] -%} , {% if ismain %}{% endif -%} [{{ loop.index }}] {%- if ismain %}{% endif -%} + >{% if ismain %}{% endif -%} [{{ loop.index }}] {%- if ismain %}{% + endif -%} -{%- endfor %} {%- else %} {{ firstname|e }} {%- endif %} {% endmacro %} {%- -extends "layout.html" %} {% set title = _('Index') %} {% block body %} +{%- endfor %} {%- else %} {{ firstname|e }} {%- endif %} {% endmacro %} {%- extends +"layout.html" %} {% set title = _('Index') %} {% block body %}

{{ _('Index') }}

diff --git a/docs/about.md b/docs/about.md index 99375693..be3bff2c 100644 --- a/docs/about.md +++ b/docs/about.md @@ -1,22 +1,20 @@ # About -The ["Common Partial Wave Analysis"](https://github.com/ComPWA) organization -(ComPWA) aims to make amplitude analysis accessible through transparent and -interactive documentation, modern software development tools, and -collaboration-independent frameworks. **Contact details** can be found -[here](https://github.com/ComPWA). +The ["Common Partial Wave Analysis"](https://github.com/ComPWA) organization (ComPWA) +aims to make amplitude analysis accessible through transparent and interactive +documentation, modern software development tools, and collaboration-independent +frameworks. **Contact details** can be found [here](https://github.com/ComPWA). ## Main projects -ComPWA maintains **three main Python packages** with which you can do a full -partial wave analysis. The packages are designed as **libraries**, so that they -can be used separately by other projects. +ComPWA maintains **three main Python packages** with which you can do a full partial +wave analysis. The packages are designed as **libraries**, so that they can be used +separately by other projects. -Each of these libraries come with **interactive and interlinked documentation** -that is intended to bring theory and code closer together. The PWA Pages takes -that one step further: it is an independent and easy-to-maintain documentation -project that can serve as a central place to gather links to PWA theory and -software. +Each of these libraries come with **interactive and interlinked documentation** that is +intended to bring theory and code closer together. The PWA Pages takes that one step +further: it is an independent and easy-to-maintain documentation project that can serve +as a central place to gather links to PWA theory and software. :::{panels} @@ -76,162 +74,151 @@ A central knowledge-base for Partial Wave Analysis theory and software ::: Finally, the {doc}`technical reports ` on these pages -([compwa-org.rtfd.io](https://compwa-org.readthedocs.io)) describe more general -tips and tricks, some of which can be of interest to general Python users as -well! +([compwa-org.rtfd.io](https://compwa-org.readthedocs.io)) describe more general tips and +tricks, some of which can be of interest to general Python users as well! :::{dropdown} Deprecated projects -The main packages listed above originate from the following, deprecated -projects: +The main packages listed above originate from the following, deprecated projects: -- [ComPWA](https://compwa.readthedocs.io): a single framework for Partial Wave - Analysis written in C++. -- [pycompwa](https://compwa.github.io): the Python interface of ComPWA, which - also hosted a first version of the PWA Expert System. -- {doc}`PWA Expert System ` (split into - {doc}`QRules ` and {doc}`AmpForm `). +- [ComPWA](https://compwa.readthedocs.io): a single framework for Partial Wave Analysis + written in C++. +- [pycompwa](https://compwa.github.io): the Python interface of ComPWA, which also + hosted a first version of the PWA Expert System. +- {doc}`PWA Expert System ` (split into {doc}`QRules ` + and {doc}`AmpForm `). ::: ## Long-term development -Partial Wave Analysis is a complicated research discipline, where several -aspects of quantum field theory, experimental physics, statistics, regression -analysis, and high-performance computing come together. This has led to -{ref}`a large number of PWA frameworks ` that -taylor to the need of each collaboration. +Partial Wave Analysis is a complicated research discipline, where several aspects of +quantum field theory, experimental physics, statistics, regression analysis, and +high-performance computing come together. This has led to +{ref}`a large number of PWA frameworks ` that taylor to +the need of each collaboration. -This state of affairs is only natural: research requires a flexible and -specialized approach. If, say, some background component shows up in an ongoing -analysis, one may need to implement some formalism that can handle it or add -some alignment parameters that were not yet supplied by the framework. It's -therefore hard to facilitate ongoing research, while at the same time -developing a general, long-term PWA tool. +This state of affairs is only natural: research requires a flexible and specialized +approach. If, say, some background component shows up in an ongoing analysis, one may +need to implement some formalism that can handle it or add some alignment parameters +that were not yet supplied by the framework. It's therefore hard to facilitate ongoing +research, while at the same time developing a general, long-term PWA tool. This situation however has a few major disadvantages: -- Collaborations usually start developing a PWA framework from scratch. It - therefore takes years for packages to move beyond the basic PWA formalisms. -- Development is slow, because expertise is splintered: every group is working - on their own package. -- Results become unreliable: the fewer people use a package, the more bugs will - remain unnoticed. -- Once developers leave, the framework collapses. Sadly, this happens more - often than not, as developers are usually scientists on short-term contracts. +- Collaborations usually start developing a PWA framework from scratch. It therefore + takes years for packages to move beyond the basic PWA formalisms. +- Development is slow, because expertise is splintered: every group is working on their + own package. +- Results become unreliable: the fewer people use a package, the more bugs will remain + unnoticed. +- Once developers leave, the framework collapses. Sadly, this happens more often than + not, as developers are usually scientists on short-term contracts. ComPWA attempts to break this with the following ideals: - Prioritize {ref}`about:Developer Experience` over functionality. - Have everything {ref}`open source ` and remain collaboration-independent -- Offer functionality in the form of modules (libraries) that can serve as - building blocks in more specialized scripts and frameworks - ({ref}`about:Design`). +- Offer functionality in the form of modules (libraries) that can serve as building + blocks in more specialized scripts and frameworks ({ref}`about:Design`). - Record and share acquired knowledge through {ref}`modern, interlinked documentation `. - Follow industry standards and tools (e.g. {doc}`ticket system, CI, ADR, etc... `). -The first point is crucial. ComPWA rather sacrifices functionality for design -and developer experience related developments. Many other frameworks have -started with the same ideal of having a good software design etc., but soon -begin to drop those ideals for the understandable reasons described above. We -believe that only by sticking to those ideals, it is possible to maintain a -long-term and collaboration-independent common tool. +The first point is crucial. ComPWA rather sacrifices functionality for design and +developer experience related developments. Many other frameworks have started with the +same ideal of having a good software design etc., but soon begin to drop those ideals +for the understandable reasons described above. We believe that only by sticking to +those ideals, it is possible to maintain a long-term and collaboration-independent +common tool. ## Developer Experience -PWA is performed by researchers and this field of study is but relatively -small. In practice, users of PWA frameworks are therefore close to the -development of the framework itself. This means that, ideally, the gap between -the developer and the user should remain as small as possible. - -In order to close this gap as best as possible, ComPWA follows the following -ideas: - -1. Frameworks are built up as **modular libraries with an accessible and - well-documented [API](https://en.wikipedia.org/wiki/API)** (see for example - [here](https://ampform.readthedocs.io/en/stable/api/ampform.html)). Users - preferably set up their analysis through Jupyter notebooks or scripts that - use these libraries. This allows the user to adapt to the specific - challenges that their research challenges pose, while at the same time - thinking along or improving the with the design and features offered by the - library. -2. **New features are added to the library with care**, as to not let the - library grow over time with features that are specific to certain analyses. - Procedures are kept small and general enough, so that they can be used by - different user scripts. This relates to 1., because the API and underlying - code base should remain understandable. -3. It should be **easy to make the step from using the library to contributing - to the source code**. The main starting points are the - {ref}`example pages provided by each library ` - (accessible without any software but a browser). From there, it should be - just a matter of a running few command lines to start modifying and trying - out changes to the code-bases (see e.g. {ref}`develop:Local set-up`) in a +PWA is performed by researchers and this field of study is but relatively small. In +practice, users of PWA frameworks are therefore close to the development of the +framework itself. This means that, ideally, the gap between the developer and the user +should remain as small as possible. + +In order to close this gap as best as possible, ComPWA follows the following ideas: + +1. Frameworks are built up as **modular libraries with an accessible and well-documented + [API](https://en.wikipedia.org/wiki/API)** (see for example + [here](https://ampform.readthedocs.io/en/stable/api/ampform.html)). Users preferably + set up their analysis through Jupyter notebooks or scripts that use these libraries. + This allows the user to adapt to the specific challenges that their research + challenges pose, while at the same time thinking along or improving the with the + design and features offered by the library. +2. **New features are added to the library with care**, as to not let the library grow + over time with features that are specific to certain analyses. Procedures are kept + small and general enough, so that they can be used by different user scripts. This + relates to 1., because the API and underlying code base should remain understandable. +3. It should be **easy to make the step from using the library to contributing to the + source code**. The main starting points are the + {ref}`example pages provided by each library ` (accessible + without any software but a browser). From there, it should be just a matter of a + running few command lines to start modifying and trying out changes to the code-bases + (see e.g. {ref}`develop:Local set-up`) in a {ref}`standardized and automated environment `. -4. **Developments industry techniques and software development tools are - followed closely.** Despite the specific nature of PWA, many aspects can be - performed or supported by existing technologies or packages. Think of the - regression process that forms the basis of Machine Learning, but also the - growing number of tools that are popular in data science, like - [Pandas](https://pandas.pydata.org) and +4. **Developments industry techniques and software development tools are followed + closely.** Despite the specific nature of PWA, many aspects can be performed or + supported by existing technologies or packages. Think of the regression process that + forms the basis of Machine Learning, but also the growing number of tools that are + popular in data science, like [Pandas](https://pandas.pydata.org) and [Jupyter notebooks](https://jupyter.org). -5. As frameworks are {ref}`open source ` and users and - developers come and go, **decision making is recorded as thoroughly as - possible**. Fortunately, Git (commit history) and GitHub - ({ref}`issues, PRs, and discussions `) make this - extremely easy. In addition, larger decisions are recorded in the form of - {doc}`ADRs ` and explorations of challenges posed by physics and - software are hosted in the form of {doc}`/reports`. +5. As frameworks are {ref}`open source ` and users and developers + come and go, **decision making is recorded as thoroughly as possible**. Fortunately, + Git (commit history) and GitHub + ({ref}`issues, PRs, and discussions `) make this extremely + easy. In addition, larger decisions are recorded in the form of {doc}`ADRs ` + and explorations of challenges posed by physics and software are hosted in the form + of {doc}`/reports`. ### Design - Code modularity and transparency. For example, separation of {mod}`qrules`, - {mod}`ampform`, and {mod}`tensorwaves`. The former two include all of the - physics, while {mod}`tensorwaves` can use these amplitude models and perform - fits, but aims to keep physics logic contained upstream. -- Keep the code simple by sticking to the core responsibility: construct - amplitude models and fitting them to data. Avoid + {mod}`ampform`, and {mod}`tensorwaves`. The former two include all of the physics, + while {mod}`tensorwaves` can use these amplitude models and perform fits, but aims to + keep physics logic contained upstream. +- Keep the code simple by sticking to the core responsibility: construct amplitude + models and fitting them to data. Avoid ["feature creep"](https://en.wikipedia.org/wiki/Feature_creep)! -- Accommodate both stable development and flexibility for ongoing analyses (see - e.g. {ref}`develop:Branching model`). +- Accommodate both stable development and flexibility for ongoing analyses (see e.g. + {ref}`develop:Branching model`). - ComPWA values the [open-closed principle](https://en.wikipedia.org/wiki/Open%E2%80%93closed_principle). - Where possible, libraries are intended to give users the flexibility to - insert custom behavior (like custom dynamics) without having to introduce new - updates to the library. + Where possible, libraries are intended to give users the flexibility to insert custom + behavior (like custom dynamics) without having to introduce new updates to the + library. ### Open source -ComPWA repositories are intended to be collaboration-independent. As such, they -are always public and open source and free to try out and re-use under the +ComPWA repositories are intended to be collaboration-independent. As such, they are +always public and open source and free to try out and re-use under the [GPLv3+ license](https://www.gnu.org/licenses/gpl-3.0-standalone.html). -At the same time, open source projects come with many challenges: it is crucial -to +At the same time, open source projects come with many challenges: it is crucial to {ref}`maintain strict standards for the code-base ` from the start, when anyone is allowed to contribute. ## User Experience -PWA is a difficult field to get into and to navigate around. ComPWA therefore -puts most effort into maintaining **easily navigable and interactive -documentation** that **narrows the gap between code and theory**. - -- All libraries provide example pages (see e.g. {doc}`here `). - These pages are written as Jupyter notebook and provide a close link between - code (how to use the library in a script, with links to the API) and theory - (explanations of the basics being performed). -- Texts on the web-pages are thoroughly interlinked, so that the reader can - easily to navigate to literature or external resources for more information. - The intention is to make PWA a more accessible field for newcomers and to - provide reference to the literature that was consulted for the - implementation. -- Almost all ComPWA libraries are written in Python. This makes it easy for - analysis users to install and use. In addition, the Python community has - developed excellent tools that make it easy to document and maintain a clean - codebase, so that is easy to make the step to become a - {ref}`developer `. +PWA is a difficult field to get into and to navigate around. ComPWA therefore puts most +effort into maintaining **easily navigable and interactive documentation** that +**narrows the gap between code and theory**. + +- All libraries provide example pages (see e.g. {doc}`here `). These + pages are written as Jupyter notebook and provide a close link between code (how to + use the library in a script, with links to the API) and theory (explanations of the + basics being performed). +- Texts on the web-pages are thoroughly interlinked, so that the reader can easily to + navigate to literature or external resources for more information. The intention is to + make PWA a more accessible field for newcomers and to provide reference to the + literature that was consulted for the implementation. +- Almost all ComPWA libraries are written in Python. This makes it easy for analysis + users to install and use. In addition, the Python community has developed excellent + tools that make it easy to document and maintain a clean codebase, so that is easy to + make the step to become a {ref}`developer `. diff --git a/docs/adr.md b/docs/adr.md index dfcf2ecc..c8d855d9 100644 --- a/docs/adr.md +++ b/docs/adr.md @@ -13,7 +13,6 @@ maxdepth: 1 adr/* ``` -For new ADRs, please use {download}`adr/template.md` as basis. This template -was inspired by [MADR](https://adr.github.io/madr). General information about -architectural decision records is available at -[adr.github.io](https://adr.github.io). +For new ADRs, please use {download}`adr/template.md` as basis. This template was +inspired by [MADR](https://adr.github.io/madr). General information about architectural +decision records is available at [adr.github.io](https://adr.github.io). diff --git a/docs/adr/000.md b/docs/adr/000.md index bde8f845..a944a1dc 100644 --- a/docs/adr/000.md +++ b/docs/adr/000.md @@ -6,21 +6,20 @@ Status: **accepted** Deciders: @redeboer, @spflueger -Technical story: A large number of issues in the {mod}`expertsystem` are -correlated (e.g. {issue}`ComPWA/expertsystem#40`, -{issue}`ComPWA/expertsystem#44`, {issue}`ComPWA/expertsystem#22`) so that -resulting PRs (in this case, {pr}`ComPWA/expertsystem#42`) lacked direction. -This led us to consider ADRs. +Technical story: A large number of issues in the {mod}`expertsystem` are correlated +(e.g. {issue}`ComPWA/expertsystem#40`, {issue}`ComPWA/expertsystem#44`, +{issue}`ComPWA/expertsystem#22`) so that resulting PRs (in this case, +{pr}`ComPWA/expertsystem#42`) lacked direction. This led us to consider ADRs. ## Context and Problem Statement -We want to record architectural decisions made in this project. Which format -and structure should these records follow? +We want to record architectural decisions made in this project. Which format and +structure should these records follow? ## Considered Options -- [MADR](https://adr.github.io/madr/) 2.1.2 – The Markdown Architectural - Decision Records +- [MADR](https://adr.github.io/madr/) 2.1.2 – The Markdown Architectural Decision + Records - [Michael Nygard's template](https://cognitect.com/blog/2011/11/15/documenting-architecture-decisions) – The first incarnation of the term "ADR" - [Sustainable Architectural Decisions](https://www.infoq.com/articles/sustainable-architectural-design-decisions) @@ -33,8 +32,8 @@ and structure should these records follow? Chosen option: "MADR 2.1.2", because -- Implicit assumptions should be made explicit. Design documentation is - important to enable people understanding the decisions later on. See also +- Implicit assumptions should be made explicit. Design documentation is important to + enable people understanding the decisions later on. See also [A rational design process: How and why to fake it](https://ieeexplore.ieee.org/document/6312940/). - The MADR format is lean and fits our development style. - The MADR structure is comprehensible and facilitates usage & maintenance. diff --git a/docs/adr/001.md b/docs/adr/001.md index fa793cd7..886ea3cc 100644 --- a/docs/adr/001.md +++ b/docs/adr/001.md @@ -9,61 +9,53 @@ ## Context and problem statement From the perspective of a PWA fitter package, the responsibility of the -{mod}`expertsystem` is to construct a `AmplitudeModel` that serves as blueprint -for a function that can be evaluated. Such a **function** has the following -requirements: +{mod}`expertsystem` is to construct a `AmplitudeModel` that serves as blueprint for a +function that can be evaluated. Such a **function** has the following requirements: -1. It should be able to compute a list of real-valued intensities - $\mathbb{R}^m$ from a dataset of four-momenta - $\mathbb{R}^{m\times n\times4}$, where $m$ is the number of events and $n$ - is the number of final state particles. -2. It should contain **parameters** that can be tweaked, so that they can be - optimized with regard to a certain estimator. +1. It should be able to compute a list of real-valued intensities $\mathbb{R}^m$ from a + dataset of four-momenta $\mathbb{R}^{m\times n\times4}$, where $m$ is the number of + events and $n$ is the number of final state particles. +2. It should contain **parameters** that can be tweaked, so that they can be optimized + with regard to a certain estimator. ### Technical story -- {issue}`ComPWA/ampform#5`: Coupling parameters in the `AmplitudeModel` is - difficult (has to be done through the place where they are used in the - `dynamics` or `intensity` section) and counter-intuitive (cannot be done - through the `parameters` section) -- {issue}`ComPWA/expertsystem#440`: when overwriting existing dynamics, old - parameters are not cleaned up from the `parameters` section -- {issue}`ComPWA/expertsystem#441`: parameters contain a name that can be - changed, but that results in a mismatch between the key that is used in the - `parameters` section and the name of the parameter to which that entry - points. -- {pr}`ComPWA/ComPWA#226`: Use a math language for the blueprint of the - function. This was also discussed early to mid 2020, but dropped in favor of - custom python code + `amplitf`. The reasoning was that the effort of writing - some new math language plus generators converting a mathematical expression - into a function (using various back-ends) requires too much manpower. +- {issue}`ComPWA/ampform#5`: Coupling parameters in the `AmplitudeModel` is difficult + (has to be done through the place where they are used in the `dynamics` or `intensity` + section) and counter-intuitive (cannot be done through the `parameters` section) +- {issue}`ComPWA/expertsystem#440`: when overwriting existing dynamics, old parameters + are not cleaned up from the `parameters` section +- {issue}`ComPWA/expertsystem#441`: parameters contain a name that can be changed, but + that results in a mismatch between the key that is used in the `parameters` section + and the name of the parameter to which that entry points. +- {pr}`ComPWA/ComPWA#226`: Use a math language for the blueprint of the function. This + was also discussed early to mid 2020, but dropped in favor of custom python code + + `amplitf`. The reasoning was that the effort of writing some new math language plus + generators converting a mathematical expression into a function (using various + back-ends) requires too much manpower. ## Decision drivers ### Solution requirements -1. The `AmplitudeModel` has to be convertible to a function which can be - **evaluated using various computation back-ends** (numpy, tensorflow, - theano, jax, ...) -2. Ideally, the model should be complete in the sense that it contains all - information to construct the complete model. This means that some "common" - functions like a Breit-Wigner and Blatt-Weisskopf form factors should also - be contained inside the `AmplitudeModel`. This guarantees - **reproducibility**! +1. The `AmplitudeModel` has to be convertible to a function which can be **evaluated + using various computation back-ends** (numpy, tensorflow, theano, jax, ...) +2. Ideally, the model should be complete in the sense that it contains all information + to construct the complete model. This means that some "common" functions like a + Breit-Wigner and Blatt-Weisskopf form factors should also be contained inside the + `AmplitudeModel`. This guarantees **reproducibility**! 3. Adding new operators/models should not trigger many code modifications ([open-closed principle](https://en.wikipedia.org/wiki/Open%E2%80%93closed_principle)), for instance adding new dynamics or formalisms. 4. Extendible: - - Add or replace current parts of an existing model. For example replace the - dynamics part of some decay. - - Change a function plus a dataset to an estimator function. This is a - subtle but important point. The function should hide its details (which - backend and its mathematical expression) and yet be extendable to an - estimator. -5. Definition and easy extraction of components. Components are certain - sub-parts of the complete mathematical expression. This is at least needed - for the calculation of fit fractions, or plotting individual parts of the - intensity. + - Add or replace current parts of an existing model. For example replace the dynamics + part of some decay. + - Change a function plus a dataset to an estimator function. This is a subtle but + important point. The function should hide its details (which backend and its + mathematical expression) and yet be extendable to an estimator. +5. Definition and easy extraction of components. Components are certain sub-parts of the + complete mathematical expression. This is at least needed for the calculation of fit + fractions, or plotting individual parts of the intensity. ## Considered solutions @@ -75,20 +67,18 @@ contains five **sections** (instances of specific classes): - `kinematics`: defines initial and final state - `particles`: particle definitions (spin, etc.) -- `dynamics`: a mapping that defines which dynamics type to apply to which - particle -- `intensity`: the actual amplitude model that is to be converted by a fitter - package into a function as described above -- `parameters`: an inventory of parameters that are used in `intensity` and - `dynamics` +- `dynamics`: a mapping that defines which dynamics type to apply to which particle +- `intensity`: the actual amplitude model that is to be converted by a fitter package + into a function as described above +- `parameters`: an inventory of parameters that are used in `intensity` and `dynamics` This structure can be represented in YAML, see an example [here](https://github.com/ComPWA/expertsystem/blob/f4f1c55/tests/unit/io/expected_recipe.yml). -A fitter package converts `intensity` together with `dynamics` into a function. -Any references to parameters that `intensity` or `dynamics` contain are -converted into a parameter of the function. The parameters are initialized with -the value as listed in the `parameters` section of the `AmplitudeModel`. +A fitter package converts `intensity` together with `dynamics` into a function. Any +references to parameters that `intensity` or `dynamics` contain are converted into a +parameter of the function. The parameters are initialized with the value as listed in +the `parameters` section of the `AmplitudeModel`. ### Alternative solutions @@ -113,30 +103,29 @@ maxdepth: 1 - Not open-closed to new models - Conversion to various back-ends not DRY - Function replacement or extension feature becomes very difficult to handle. - - Model is not complete, since no complete mathematical description is used. - For example Breit-Wigner functions are referred to directly and their - implementations is not defined in the amplitude model. + - Model is not complete, since no complete mathematical description is used. For + example Breit-Wigner functions are referred to directly and their implementations is + not defined in the amplitude model. #### {doc}`SymPy <001/sympy>` - **Positive** - Easy to render amplitude model as LaTeX - - Model description is complete! Absolutely all information about the model - is included. (reproducibility) - - Follows open-closed principle. New models and formalism can be added - without any changes to other interfacing components (here: `tensorwaves`) + - Model description is complete! Absolutely all information about the model is + included. (reproducibility) + - Follows open-closed principle. New models and formalism can be added without any + changes to other interfacing components (here: `tensorwaves`) - Use [`lambdify`](https://docs.sympy.org/latest/tutorial/basic_operations.html#lambdify) to convert the expression to any back-end - Use [`Expr.subs`](https://docs.sympy.org/latest/modules/core.html#sympy.core.basic.Basic.subs) - (substitute) to couple parameters or replace components of the model, for - instance to set custom dynamics + (substitute) to couple parameters or replace components of the model, for instance + to set custom dynamics - **Negative** - - `lambdify` becomes a core dependency while its behavior cannot be modified, - but is defined by `sympy`. - - Need to keep track of components in the expression tree with symbol - mappings + - `lambdify` becomes a core dependency while its behavior cannot be modified, but is + defined by `sympy`. + - Need to keep track of components in the expression tree with symbol mappings #### {doc}`Python's operator library <001/operators>` @@ -149,12 +138,10 @@ maxdepth: 1 ## Decision outcome -Use {doc}`001/sympy`. Initially, we leave the existing amplitude builders -(modules +Use {doc}`001/sympy`. Initially, we leave the existing amplitude builders (modules [`helicity_decay`](https://expertsystem.readthedocs.io/en/0.6.8/api/expertsystem.amplitude.helicity_decay.html) and [`canonical_decay`](https://expertsystem.readthedocs.io/en/0.6.8/api/expertsystem.amplitude.canonical_decay.html)) -alongside a SymPy implementation, so that it's possible to compare the results. -Once it turns out the this set-up results in the same results and a comparable -performance, we replace the old amplitude builders with the new SymPy -implementation. +alongside a SymPy implementation, so that it's possible to compare the results. Once it +turns out the this set-up results in the same results and a comparable performance, we +replace the old amplitude builders with the new SymPy implementation. diff --git a/docs/adr/001/sympy.ipynb b/docs/adr/001/sympy.ipynb index de8b1b0e..4a806578 100644 --- a/docs/adr/001/sympy.ipynb +++ b/docs/adr/001/sympy.ipynb @@ -129,9 +129,7 @@ " sigma: 1 / (2 * i),\n", " }\n", " )\n", - " gauss = sp.exp(-((x - mu) ** 2) / (sigma**2)) / (\n", - " sigma * sp.sqrt(2 * sp.pi)\n", - " )\n", + " gauss = sp.exp(-((x - mu) ** 2) / (sigma**2)) / (sigma * sp.sqrt(2 * sp.pi))\n", " dyn_symbol = sp.Symbol(Rf\"\\mathrm{{dyn}}_{i}\")\n", " model.dynamics[dyn_symbol] = gauss\n", "\n", @@ -288,25 +286,19 @@ " line_color=\"black\",\n", " )\n", " p1 = sp.plotting.plot(\n", - " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_1\")].subs(\n", - " model.initial_values\n", - " ),\n", + " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_1\")].subs(model.initial_values),\n", " (x, 0, 4),\n", " line_color=\"red\",\n", " show=False,\n", " )\n", " p2 = sp.plotting.plot(\n", - " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_2\")].subs(\n", - " model.initial_values\n", - " ),\n", + " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_2\")].subs(model.initial_values),\n", " (x, 0, 4),\n", " line_color=\"blue\",\n", " show=False,\n", " )\n", " p3 = sp.plotting.plot(\n", - " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_3\")].subs(\n", - " model.initial_values\n", - " ),\n", + " model.dynamics[sp.Symbol(R\"\\mathrm{dyn}_3\")].subs(model.initial_values),\n", " (x, 0, 4),\n", " line_color=\"green\",\n", " show=False,\n", @@ -522,14 +514,10 @@ " \"x, A1, mu1, sigma1, A2, mu2, sigma2\"\n", ")\n", "gaussian1 = (\n", - " A1\n", - " / (sigma1 * sp.sqrt(2.0 * math.pi))\n", - " * sp.exp(-((x - mu1) ** 2) / (2 * sigma1))\n", + " A1 / (sigma1 * sp.sqrt(2.0 * math.pi)) * sp.exp(-((x - mu1) ** 2) / (2 * sigma1))\n", ")\n", "gaussian2 = (\n", - " A2\n", - " / (sigma2 * sp.sqrt(2.0 * math.pi))\n", - " * sp.exp(-((x - mu2) ** 2) / (2 * sigma2))\n", + " A2 / (sigma2 * sp.sqrt(2.0 * math.pi)) * sp.exp(-((x - mu2) ** 2) / (2 * sigma2))\n", ")\n", "\n", "gauss_sum = gaussian1 + gaussian2\n", diff --git a/docs/adr/002.md b/docs/adr/002.md index b7575c43..f838862a 100644 --- a/docs/adr/002.md +++ b/docs/adr/002.md @@ -9,16 +9,15 @@ ## Context and problem statement -Physics models usually include assumptions that simplify the structure of the -model. For example, splitting a model into a product of independent parts, in -which every part contains a certain responsibility. In case of partial wave -amplitude models, we can make a separation into a spin part and a dynamical -part. While the spin part controls the probability w.r.t angular kinematic -variables, the dynamics controls the probability on variable like the invariant -mass of states. +Physics models usually include assumptions that simplify the structure of the model. For +example, splitting a model into a product of independent parts, in which every part +contains a certain responsibility. In case of partial wave amplitude models, we can make +a separation into a spin part and a dynamical part. While the spin part controls the +probability w.r.t angular kinematic variables, the dynamics controls the probability on +variable like the invariant mass of states. -Generally, a dynamics part is simply a function, which is defined in complex -space, and consists of: +Generally, a dynamics part is simply a function, which is defined in complex space, and +consists of: - a mathematical **expression** (`sympy.Expr`) - a set of **parameters** in that expression that can be tweaked (optimized) @@ -26,65 +25,59 @@ space, and consists of: ### Technical story -- {issue}`ComPWA/expertsystem#124`: see - {ref}`adr/002:Issues with existing set-up`. -- {issue}`ComPWA/expertsystem#440`: no way to supply custom dynamics. Or at - least, {mod}`tensorwaves` does not understand those custom dynamics. +- {issue}`ComPWA/expertsystem#124`: see {ref}`adr/002:Issues with existing set-up`. +- {issue}`ComPWA/expertsystem#440`: no way to supply custom dynamics. Or at least, + {mod}`tensorwaves` does not understand those custom dynamics. - [ADR-001](./001.md): parameters _and_ variables are to be expressed as `sympy.Symbol`s. -- {pr}`ComPWA/expertsystem#454`: dynamics are specified as a mapping of - `sympy.Function` to a `sympy.Expr`, but now there is no way to supply those - `sympy.Expr`s with expected `sympy.Symbol`s (parameters and variables). +- {pr}`ComPWA/expertsystem#454`: dynamics are specified as a mapping of `sympy.Function` + to a `sympy.Expr`, but now there is no way to supply those `sympy.Expr`s with expected + `sympy.Symbol`s (parameters and variables). ### Issues with existing set-up -- There is no clear way to apply dynamics functions to a specific decaying - particle, that is, to a specific edge of the `StateTransitionGraph`s (`STG`). - Currently, we just work with a mapping of `Particle`s to some dynamics - expression, but this is not feasible when there there are identical particles - on the edges. -- The set of variables to which a dynamics expression applies, is determined by - the position within the `STG` that it is applied to. For instance, a - relativistic Breit-Wigner that is applied to the resonance in some 1-to-3 - isobar decay (described by an `STG` with final state edges 2, 3, 4 and - intermediate edge 1) would work on the invariant mass of edge 3 and 4 - (`mSq_3+4`). -- Just like variables, parameters need to be identifiable from their position - within the `STG` (take a relativistic Breit-Wigner _with form factors_, which - would require break-up momentum as a parameter), but also require some - suggested starting value (e.g. expected pole position). These starting values - are usually taken from the edge and node properties within the `STG`. +- There is no clear way to apply dynamics functions to a specific decaying particle, + that is, to a specific edge of the `StateTransitionGraph`s (`STG`). Currently, we just + work with a mapping of `Particle`s to some dynamics expression, but this is not + feasible when there there are identical particles on the edges. +- The set of variables to which a dynamics expression applies, is determined by the + position within the `STG` that it is applied to. For instance, a relativistic + Breit-Wigner that is applied to the resonance in some 1-to-3 isobar decay (described + by an `STG` with final state edges 2, 3, 4 and intermediate edge 1) would work on the + invariant mass of edge 3 and 4 (`mSq_3+4`). +- Just like variables, parameters need to be identifiable from their position within the + `STG` (take a relativistic Breit-Wigner _with form factors_, which would require + break-up momentum as a parameter), but also require some suggested starting value + (e.g. expected pole position). These starting values are usually taken from the edge + and node properties within the `STG`. ## Decision drivers The following points are nice to have or can influence the decision but are not essential and can be part of the users responsibility. -1. The parameters that a dynamics functions requires, are registered - automatically and linked together. -2. Kinematic variables used in dynamics functions are also linked - appropriately. +1. The parameters that a dynamics functions requires, are registered automatically and + linked together. +2. Kinematic variables used in dynamics functions are also linked appropriately. 3. It is easy to define custom dynamics (no boilerplate code). ### Solution requirements -1. It is easy to apply dynamics to specific components of the `STG`s. Note: - it's important that the dynamics can be applied to resonances of some - **selected** graphs and not generally all graphs in which the resonance - appears. +1. It is easy to apply dynamics to specific components of the `STG`s. Note: it's + important that the dynamics can be applied to resonances of some **selected** graphs + and not generally all graphs in which the resonance appears. 2. Where possible, suggested (initial) parameter values are provided as well. -3. It is possible to use and inspect the dynamics expression itself - independently from the `expertsystem`. -4. Follow open-closed principle. Probably the most important decision driver. - The solution should be flexible enough to handle any possible scenario, - without having to change the interface defined in requirement 1! +3. It is possible to use and inspect the dynamics expression itself independently from + the `expertsystem`. +4. Follow open-closed principle. Probably the most important decision driver. The + solution should be flexible enough to handle any possible scenario, without having to + change the interface defined in requirement 1! ## Considered solutions ### Group 1: expression builder -To satisfy [requirement 1](#solution-requirements), we propose the following -syntax: +To satisfy [requirement 1](#solution-requirements), we propose the following syntax: ```python # model: ModelInfo @@ -95,8 +88,8 @@ model.set_dynamics(graph, edge_id=1, expression_builder) :::{toggle} Another style would be to have `ModelInfo` contain a reference to the list of -`StateTransitionGraph`s. The user then needs some other way to express which -edges to apply the dynamics function to: +`StateTransitionGraph`s. The user then needs some other way to express which edges to +apply the dynamics function to: ```python model.set_dynamics( @@ -108,16 +101,14 @@ model.set_dynamics( ::: -Here, `expression_builder` is some function or method that can create a -dynamics expression. It can also be a class that contains both the -implementation of the expression and a static method to build itself from a -`StateTransitionGraph`. +Here, `expression_builder` is some function or method that can create a dynamics +expression. It can also be a class that contains both the implementation of the +expression and a static method to build itself from a `StateTransitionGraph`. The dynamics expression needs to be formulated in such a way that it satisfies -[the rest of the requirements](#solution-requirements). The following options -illustrate three different ways of formulating a dynamics expression, each -taking a relativistic Breit-Wigner and a relativistic Breit-Wigner _with form -factor_ as example. +[the rest of the requirements](#solution-requirements). The following options illustrate +three different ways of formulating a dynamics expression, each taking a relativistic +Breit-Wigner and a relativistic Breit-Wigner _with form factor_ as example. ```{toctree} --- @@ -138,42 +129,38 @@ A second branch of solutions would propose the following interface: model.set_dynamics(graph, edge_id=1, expression) ``` -The key difference is the usage of general sympy expression `sympy.Expr` as an -argument instead of constructing this through some builder object. +The key difference is the usage of general sympy expression `sympy.Expr` as an argument +instead of constructing this through some builder object. ## Solution evaluation ### 1: Expression builder -All of the solutions have the drawback arising from the choice of interface -using a `expression_builder`. This enforces the logic of correctly coupling -variables and parameters into these builders. This is extremely hard to get -right, since the code has to be able to handle arbitrarily complex models. And -always knowing what the user would like to do is more or less impossible. -Therefore it is much better to use a already built expression that is assumed -to be correctly built (see [solution group 2](#group-2-expression-only)). +All of the solutions have the drawback arising from the choice of interface using a +`expression_builder`. This enforces the logic of correctly coupling variables and +parameters into these builders. This is extremely hard to get right, since the code has +to be able to handle arbitrarily complex models. And always knowing what the user would +like to do is more or less impossible. Therefore it is much better to use a already +built expression that is assumed to be correctly built (see +[solution group 2](#group-2-expression-only)). -All of the solutions in this group also have the following additional -drawbacks. These are however more related to the correct building of the -dynamics expression: +All of the solutions in this group also have the following additional drawbacks. These +are however more related to the correct building of the dynamics expression: - There is an implicit assumption on the signature of the expression: the first - arguments are assumed to be the (kinematic) `variables` and the remaining - arguments are `parameters`. In addition, the arguments cannot be keywords, - but have to be positional. -- The number of `variables` and `parameters` is only verified at runtime (no - static typing, other than a check that each of the elements is - `sympy.Symbol`). - -Composition is the cleanest design, but is less in tune with the design of -{mod}`sympy`. {doc}`002/function` and {doc}`002/expr` follow {mod}`sympy` -implementations, but result in an obscure inheritance hierarchy with implicit -conventions. This can result in some nasty bugs, for instance if one were to -`__call__` method in either the `sympy.Function` or `sympy.Expr` -implementation. - -Pros and Cons that are specific to each of the implementations are listed -below. + arguments are assumed to be the (kinematic) `variables` and the remaining arguments + are `parameters`. In addition, the arguments cannot be keywords, but have to be + positional. +- The number of `variables` and `parameters` is only verified at runtime (no static + typing, other than a check that each of the elements is `sympy.Symbol`). + +Composition is the cleanest design, but is less in tune with the design of {mod}`sympy`. +{doc}`002/function` and {doc}`002/expr` follow {mod}`sympy` implementations, but result +in an obscure inheritance hierarchy with implicit conventions. This can result in some +nasty bugs, for instance if one were to `__call__` method in either the `sympy.Function` +or `sympy.Expr` implementation. + +Pros and Cons that are specific to each of the implementations are listed below. #### {doc}`002/composition` @@ -181,68 +168,63 @@ below. - Implementation of the expression is transparent - **Negative** - {ref}`adr/002/composition:Alternative signature`. - - The only way to see that `relativistic_breit_wigner_from_graph` is the - builder for `relativistic_breit_wigner`, is from its name. This makes it - implementing custom dynamics inconvenient and error-prone. - - Signature of the builder can only be checked with a - {class}`~typing.Protocol`, see {ref}`adr/002/composition:Type checking`. + - The only way to see that `relativistic_breit_wigner_from_graph` is the builder for + `relativistic_breit_wigner`, is from its name. This makes it implementing custom + dynamics inconvenient and error-prone. + - Signature of the builder can only be checked with a {class}`~typing.Protocol`, see + {ref}`adr/002/composition:Type checking`. #### {doc}`002/function` - **Positive** - `DynamicsFunction` behaves as a {class}`~sympy.core.function.Function` - - Implementation of the builder is kept together with the implementation of - the expression. + - Implementation of the builder is kept together with the implementation of the + expression. - **Negative** - It's not possible to identify variables and parameters #### {doc}`002/expr` - **Positive** - - When - {doc}`recursing through the amplitude model `, - it is still possible to identify instances of `DynamicsExpr` (before - `doit()` has been called). - - Additional properties and methods can be added and carried around by the - class. + - When {doc}`recursing through the amplitude model `, it + is still possible to identify instances of `DynamicsExpr` (before `doit()` has been + called). + - Additional properties and methods can be added and carried around by the class. - **Negative** - Boilerplate code required when implementing custom dynamics - {ref}`adr/002/expr:Issue with lambdify` ### 2: Expression-only -**Positive**: This choice of interface follows the principle of SOLID more than -solution group 1. By handing a complete expression of the dynamics to the -setter, its sole responsibility is to insert this expression at the correct -place in the full model expression. +**Positive**: This choice of interface follows the principle of SOLID more than solution +group 1. By handing a complete expression of the dynamics to the setter, its sole +responsibility is to insert this expression at the correct place in the full model +expression. -**Negative**: There are no direct negative aspects to this solution, as it just -splits up responsibilities. The construction of the expression with the correct -linking of parameters and initial values etc has to be performed by some other -code. This code is subject to the same issues mentioned in the individual -solutions of group 1. +**Negative**: There are no direct negative aspects to this solution, as it just splits +up responsibilities. The construction of the expression with the correct linking of +parameters and initial values etc has to be performed by some other code. This code is +subject to the same issues mentioned in the individual solutions of group 1. ## Decision outcome Use a composition based solution from group 2. -Important is the definition of the interface following solution group 2. This -ensures to be open-closed and keep the responsibilities separated. +Important is the definition of the interface following solution group 2. This ensures to +be open-closed and keep the responsibilities separated. -The `expertsystem` favors **composition over inheritance**: we intend to use -inheritance only to define interfaces, not to insert behavior. As such, the -design of the `expertsystem` is fundamentally different than that of SymPy. -That's another reason to favor composition here: the interfaces are not -determined by the dependency and instead remain -{doc}`contained within the dynamics class <002/composition>`. +The `expertsystem` favors **composition over inheritance**: we intend to use inheritance +only to define interfaces, not to insert behavior. As such, the design of the +`expertsystem` is fundamentally different than that of SymPy. That's another reason to +favor composition here: the interfaces are not determined by the dependency and instead +remain {doc}`contained within the dynamics class <002/composition>`. We decide to keep responsibilities as separated as possible. This means that: -1. The only responsibility of `set_dynamics` method is to attribute some - expression (`sympy.Expr`) the correct symbol within the complete amplitude - model. For now, this position is specified using some `StateTransitionGraph` - and an edge ID, but this syntax may be improved later (see - {pr}`ComPWA/expertsystem#458`): +1. The only responsibility of `set_dynamics` method is to attribute some expression + (`sympy.Expr`) the correct symbol within the complete amplitude model. For now, this + position is specified using some `StateTransitionGraph` and an edge ID, but this + syntax may be improved later (see {pr}`ComPWA/expertsystem#458`): ```python def set_dynamics( @@ -258,19 +240,18 @@ We decide to keep responsibilities as separated as possible. This means that: It is assumed that the `expression` is correct. -2. The user has the responsibility of formulating the `expression` with the - correct symbols. To aid the user in the construction of such expressions - some building code can handle some of the common tasks, such as +2. The user has the responsibility of formulating the `expression` with the correct + symbols. To aid the user in the construction of such expressions some building code + can handle some of the common tasks, such as - - A `VariablePool` can facilitate finding the correct symbol names (to avoid - typos). + - A `VariablePool` can facilitate finding the correct symbol names (to avoid typos). ```python mass = variable_pool.get_invariant_mass(graph, edge_id) ``` - - A `dynamics` module provides descriptions of common line-shapes as well as - some helper functions. An example would be: + - A `dynamics` module provides descriptions of common line-shapes as well as some + helper functions. An example would be: ```python inv_mass, mass0, gamma0 = build_relativistic_breit_wigner(graph, edge_id, particle) @@ -278,9 +259,8 @@ We decide to keep responsibilities as separated as possible. This means that: model.set_dynamics(graph, edge, rel_bw) ``` -3. The `SympyModel` has the responsibility of defining a the full model in - terms of an expression and keeping track of variables and parameters, for - instance: +3. The `SympyModel` has the responsibility of defining a the full model in terms of an + expression and keeping track of variables and parameters, for instance: ```python from __future__ import annotations diff --git a/docs/adr/002/composition.ipynb b/docs/adr/002/composition.ipynb index 21a052e9..934d966f 100644 --- a/docs/adr/002/composition.ipynb +++ b/docs/adr/002/composition.ipynb @@ -134,10 +134,7 @@ " width = gamma0 * (mass0 / mass) * (ff2 / ff02)\n", " width = width * sp.sqrt(q_squared / q0_squared)\n", " return (\n", - " relativistic_breit_wigner(mass, mass0, width)\n", - " * mass0\n", - " * gamma0\n", - " * sp.sqrt(ff2)\n", + " relativistic_breit_wigner(mass, mass0, width) * mass0 * gamma0 * sp.sqrt(ff2)\n", " )" ] }, diff --git a/docs/adr/002/expr.ipynb b/docs/adr/002/expr.ipynb index 7edaf650..6c03a073 100644 --- a/docs/adr/002/expr.ipynb +++ b/docs/adr/002/expr.ipynb @@ -154,11 +154,7 @@ " return (\n", " self.gamma0\n", " * self.mass0\n", - " / (\n", - " self.mass0**2\n", - " - self.mass**2\n", - " - self.gamma0 * self.mass0 * sp.I\n", - " )\n", + " / (self.mass0**2 - self.mass**2 - self.gamma0 * self.mass0 * sp.I)\n", " )\n", "\n", " @classmethod\n", @@ -198,9 +194,7 @@ " return sp.Expr.__new__(cls, *args)\n", "\n", " def doit(self, **hints) -> sp.Expr:\n", - " return RelativisticBreitWignerWithFF(\n", - " *self.args, **hints, evaluate=True\n", - " )\n", + " return RelativisticBreitWignerWithFF(*self.args, **hints, evaluate=True)\n", "\n", " @property\n", " def mass(self) -> sp.Symbol:\n", @@ -234,12 +228,8 @@ " # Computed variables\n", " q_squared = two_body_momentum_squared(self.mass, self.m_a, self.m_b)\n", " q0_squared = two_body_momentum_squared(self.mass0, self.m_a, self.m_b)\n", - " ff2 = blatt_weisskopf(\n", - " q_squared, self.meson_radius, self.angular_momentum\n", - " )\n", - " ff02 = blatt_weisskopf(\n", - " q0_squared, self.meson_radius, self.angular_momentum\n", - " )\n", + " ff2 = blatt_weisskopf(q_squared, self.meson_radius, self.angular_momentum)\n", + " ff02 = blatt_weisskopf(q0_squared, self.meson_radius, self.angular_momentum)\n", " width = self.gamma0 * (self.mass0 / self.mass) * (ff2 / ff02)\n", " width = width * sp.sqrt(q_squared / q0_squared)\n", " # Expression\n", @@ -406,9 +396,7 @@ " def decorator(decorated_class: type) -> type[DynamicsExpr]:\n", " def __new__(cls, *args: sp.Symbol, **hints) -> sp.Expr:\n", " if len(args) != n_args:\n", - " raise ValueError(\n", - " f\"{n_args} parameters expected, got {len(args)}\"\n", - " )\n", + " raise ValueError(f\"{n_args} parameters expected, got {len(args)}\")\n", " args = sp.sympify(args)\n", " evaluate = hints.get(\"evaluate\", False)\n", " if evaluate:\n", diff --git a/docs/adr/002/function.ipynb b/docs/adr/002/function.ipynb index 9bcf1b3d..bddf7d1d 100644 --- a/docs/adr/002/function.ipynb +++ b/docs/adr/002/function.ipynb @@ -105,9 +105,7 @@ " mass = args[0]\n", " mass0 = args[1]\n", " gamma0 = args[2]\n", - " return (\n", - " gamma0 * mass0 / (mass0**2 - mass**2 - gamma0 * mass0 * sp.I)\n", - " )\n", + " return gamma0 * mass0 / (mass0**2 - mass**2 - gamma0 * mass0 * sp.I)\n", "\n", " @classmethod\n", " def from_graph(\n", diff --git a/docs/adr/002/helpers.py b/docs/adr/002/helpers.py index add7daf2..86e1a44b 100644 --- a/docs/adr/002/helpers.py +++ b/docs/adr/002/helpers.py @@ -46,9 +46,7 @@ def two_body_momentum_squared( m_d: sp.Symbol, m_a: sp.Symbol, m_b: sp.Symbol ) -> sp.Expr: return ( - (m_d**2 - (m_a + m_b) ** 2) - * (m_d**2 - (m_a - m_b) ** 2) - / (4 * m_d**2) + (m_d**2 - (m_a + m_b) ** 2) * (m_d**2 - (m_a - m_b) ** 2) / (4 * m_d**2) ) diff --git a/docs/conf.py b/docs/conf.py index 706a763b..8a99cbe0 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -200,8 +200,7 @@ def get_minor_version(package_name: str) -> str: matches = re.match(r"^([0-9]+\.[0-9]+).*$", installed_version) if matches is None: raise ValueError( - "Could not find documentation for" - f" {package_name} v{installed_version}" + f"Could not find documentation for {package_name} v{installed_version}" ) return matches[1] @@ -297,7 +296,9 @@ def get_minor_version(package_name: str) -> str: "substitution", ] myst_heading_anchors = 4 -BINDER_LINK = f"https://mybinder.org/v2/gh/ComPWA/{REPO_NAME}/{BRANCH}?filepath=docs/usage" +BINDER_LINK = ( + f"https://mybinder.org/v2/gh/ComPWA/{REPO_NAME}/{BRANCH}?filepath=docs/usage" +) myst_substitutions = { "branch": BRANCH, "run_interactive": f""" @@ -360,9 +361,7 @@ def get_minor_version(package_name: str) -> str: } ON_RTD = os.environ.get("READTHEDOCS") is not None -PLANTUML_PATH = os.path.join( - os.path.dirname(__file__), "utils", "plantuml.jar" -) +PLANTUML_PATH = os.path.join(os.path.dirname(__file__), "utils", "plantuml.jar") if not os.path.exists(PLANTUML_PATH): print("\033[93;1mDowloading plantuml\033[0m") online_content = requests.get( @@ -427,9 +426,7 @@ def __init__(self): super().__init__(abbreviate_names=True) def format_names(self, role, as_sentence=True) -> Node: - formatted_names = names( - role, sep=", ", sep2=" and ", last_sep=", and " - ) + formatted_names = names(role, sep=", ", sep2=" and ", last_sep=", and ") if as_sentence: return sentence[formatted_names] else: diff --git a/docs/develop.md b/docs/develop.md index 3f7925d2..50bea100 100644 --- a/docs/develop.md +++ b/docs/develop.md @@ -4,14 +4,12 @@ [![Open in Visual Studio Code](https://open.vscode.dev/badges/open-in-vscode.svg)](https://open.vscode.dev/ComPWA/compwa-org) This page describes some of the tools and conventions followed by -[Common Partial Wave Analysis](https://github.com/ComPWA). Where possible, we -use the -[source code of the AmpForm repository](https://github.com/ComPWA/AmpForm) as -example, because its file structure is comparable to that of other ComPWA -repositories. +[Common Partial Wave Analysis](https://github.com/ComPWA). Where possible, we use the +[source code of the AmpForm repository](https://github.com/ComPWA/AmpForm) as example, +because its file structure is comparable to that of other ComPWA repositories. -::::{tip} To start developing, simply run the following from a cloned -repository on your machine: +::::{tip} To start developing, simply run the following from a cloned repository on your +machine: :::{tabbed} Conda @@ -45,65 +43,63 @@ See {ref}`develop:Virtual environment` for more info. ### Virtual environment When developing source code, it is safest to work within a -[virtual environment](https://realpython.com/python-virtual-environments-a-primer), -so that all package dependencies and developer tools are safely contained. This -is helpful in case something goes wrong with the dependencies: just trash the -environment and recreate it. In addition, you can easily install other versions -of the dependencies, without affecting other packages you may be working on. - -Two common tools to manage virtual environments are -[Conda](https://www.anaconda.com) and -[Python's built-in `venv`](https://docs.python.org/3/tutorial/venv.html). In -either case, you have to activate the environment whenever you want to run the -framework or use the developer tools. +[virtual environment](https://realpython.com/python-virtual-environments-a-primer), so +that all package dependencies and developer tools are safely contained. This is helpful +in case something goes wrong with the dependencies: just trash the environment and +recreate it. In addition, you can easily install other versions of the dependencies, +without affecting other packages you may be working on. + +Two common tools to manage virtual environments are [Conda](https://www.anaconda.com) +and [Python's built-in `venv`](https://docs.python.org/3/tutorial/venv.html). In either +case, you have to activate the environment whenever you want to run the framework or use +the developer tools. :::{tabbed} Conda environment -[Conda/Anaconda](https://www.anaconda.com) can be installed without -administrator rights, see instructions on -[this page](https://conda.io/projects/conda/en/latest/user-guide/install/index.html). -In addition, Conda can install more than just Python packages. +[Conda/Anaconda](https://www.anaconda.com) can be installed without administrator +rights, see instructions on +[this page](https://conda.io/projects/conda/en/latest/user-guide/install/index.html). In +addition, Conda can install more than just Python packages. -All packages {ref}`maintained by the ComPWA organization ` -provide a +All packages {ref}`maintained by the ComPWA organization ` provide +a [Conda environment file](https://docs.conda.io/projects/conda/en/latest/user-guide/tasks/manage-environments.html) -([`environment.yml`](https://github.com/ComPWA/ampform/blob/main/environment.yml)) -that defines all requirements when working on the source code of that -repository. To create an environment specific for this repository, simply -navigate to the main folder of the source code and run: +([`environment.yml`](https://github.com/ComPWA/ampform/blob/main/environment.yml)) that +defines all requirements when working on the source code of that repository. To create +an environment specific for this repository, simply navigate to the main folder of the +source code and run: ```shell conda env create ``` Conda now creates an environment with a name that is defined in the -[`environment.yml`](https://github.com/ComPWA/ampform/blob/main/environment.yml) -file. In addition, it will install the framework itself in -["editable" mode](#editable-installation), so that you can start developing -right away. +[`environment.yml`](https://github.com/ComPWA/ampform/blob/main/environment.yml) file. +In addition, it will install the framework itself in +["editable" mode](#editable-installation), so that you can start developing right away. ::: :::{tabbed} Python venv -If you have [Python's `venv`](https://docs.python.org/3/library/venv.html), -available on your system, you can create a virtual environment with it. -Navigate to some convenient folder and run: +If you have [Python's `venv`](https://docs.python.org/3/library/venv.html), available on +your system, you can create a virtual environment with it. Navigate to some convenient +folder and run: ```shell python3 -m venv ./venv ``` -This creates a folder called {file}`venv` where all Python packages will be -contained. To activate the environment, run: +This creates a folder called {file}`venv` where all Python packages will be contained. +To activate the environment, run: ```shell source ./venv/bin/activate ``` Now you can safely install the package you want to work on (see -["editable" mode](#editable-installation)), as well as any additional required -packages (see [optional dependencies](#optional-dependencies)): +["editable" mode](#editable-installation)), as well as any additional required packages +(see [optional dependencies](#optional-dependencies)): ```shell pip install -e . @@ -115,13 +111,12 @@ pip install -e . When developing a package, it is most convenient if you install it in ["editable" mode](https://pip.pypa.io/en/stable/cli/pip_install/#editable-installs). -This allows you to tweak the source code and try out new ideas immediately, -because the source code is considered the 'installation'. +This allows you to tweak the source code and try out new ideas immediately, because the +source code is considered the 'installation'. -With [`pip install`](https://pip.pypa.io/en/stable/cli/pip_install), a package -can be installed in "editable" mode with the -[`-e` flag](https://pip.pypa.io/en/stable/cli/pip_install/#install-editable). -Simply +With [`pip install`](https://pip.pypa.io/en/stable/cli/pip_install), a package can be +installed in "editable" mode with the +[`-e` flag](https://pip.pypa.io/en/stable/cli/pip_install/#install-editable). Simply [clone](https://git-scm.com/book/en/v2/Git-Basics-Getting-a-Git-Repository#_git_cloning) the repository you want to work on, navigate into it, and run: @@ -155,21 +150,21 @@ pip install .[test] # local directory, not editable pip install -e .[dev] # editable + all dev requirements ``` -Developers require several additional tools besides the dependencies required -to run the package itself (see {ref}`develop:Automated coding conventions`). -All those additional requirements can be installed with the last example. +Developers require several additional tools besides the dependencies required to run the +package itself (see {ref}`develop:Automated coding conventions`). All those additional +requirements can be installed with the last example. ### Pinning dependency versions -To ensure that developers use exactly the same versions of the package -dependencies and developer requirements, some of the repositories provide -[constraint files](https://pip.pypa.io/en/stable/user_guide/#constraints-files). -These files can be used to 'pin' all versions of installed packages as follows: +To ensure that developers use exactly the same versions of the package dependencies and +developer requirements, some of the repositories provide +[constraint files](https://pip.pypa.io/en/stable/user_guide/#constraints-files). These +files can be used to 'pin' all versions of installed packages as follows: :::{margin} -Requirements may differ per Python version, so there is one constraint file for -each version of Python that the package supports. +Requirements may differ per Python version, so there is one constraint file for each +version of Python that the package supports. ::: @@ -194,20 +189,18 @@ and :::{note} -Constraint files ensure that the framework is _deterministic and reproducible_ -(up to testing) _for all commits and versions_, which is vital for both users -(doing analysis) and for developers (for instance with -{ref}`continuous integration `). In other words, it -provides a way out of -["dependency hell"](https://en.wikipedia.org/wiki/Dependency_hell). +Constraint files ensure that the framework is _deterministic and reproducible_ (up to +testing) _for all commits and versions_, which is vital for both users (doing analysis) +and for developers (for instance with +{ref}`continuous integration `). In other words, it provides a +way out of ["dependency hell"](https://en.wikipedia.org/wiki/Dependency_hell). ::: ### Updating -It may be that new commits in the repository modify the dependencies. In that -case, you have to rerun this command after pulling new commits from the -repository: +It may be that new commits in the repository modify the dependencies. In that case, you +have to rerun this command after pulling new commits from the repository: ```shell git checkout main @@ -215,15 +208,15 @@ git pull pip install -c .constraints/py3.8.txt -e .[dev] ``` -If you still have problems, it may be that certain dependencies have become -redundant. In that case, trash the virtual environment and +If you still have problems, it may be that certain dependencies have become redundant. +In that case, trash the virtual environment and {ref}`create a new one `. ## Automated coding conventions -Where possible, we define and enforce our coding conventions through automated -tools, instead of describing them in documentation. These tools perform their -checks when you commit files locally (see {ref}`develop:Pre-commit`), when +Where possible, we define and enforce our coding conventions through automated tools, +instead of describing them in documentation. These tools perform their checks when you +commit files locally (see {ref}`develop:Pre-commit`), when {ref}`running tox `, and when you make a {ref}`pull request `. @@ -232,21 +225,18 @@ checks when you commit files locally (see {ref}`develop:Pre-commit`), when The tools are configured through files such as [`pyproject.toml`](https://github.com/ComPWA/ampform/blob/main/pyproject.toml), [`.pylintrc`](https://github.com/ComPWA/ampform/blob/main/.pylintrc), and -[`tox.ini`](https://github.com/ComPWA/ampform/blob/main/tox.ini), and the -workflow files under -[`.github`](https://github.com/ComPWA/ampform/blob/main/.github). If you run -into persistent linting errors, this may mean we need to further specify our -conventions. In that case, it's best to -{ref}`create an issue ` or a -{ref}`pull request ` and propose a policy change that -can be formulated through those config files. +[`tox.ini`](https://github.com/ComPWA/ampform/blob/main/tox.ini), and the workflow files +under [`.github`](https://github.com/ComPWA/ampform/blob/main/.github). If you run into +persistent linting errors, this may mean we need to further specify our conventions. In +that case, it's best to {ref}`create an issue ` or a +{ref}`pull request ` and propose a policy change that can be +formulated through those config files. ### Pre-commit -All {ref}`style checks ` are enforced through a tool -called [{command}`pre-commit`](https://pre-commit.com). It's best to activate -this tool locally as well. This has to be done only once, after you clone the -repository: +All {ref}`style checks ` are enforced through a tool called +[{command}`pre-commit`](https://pre-commit.com). It's best to activate this tool locally +as well. This has to be done only once, after you clone the repository: ```shell pre-commit install @@ -254,27 +244,26 @@ pre-commit install :::{margin} Initializing pre-commit -The first time you run {command}`pre-commit` after installing or updating its -checks, it may take some time to initialize. +The first time you run {command}`pre-commit` after installing or updating its checks, it +may take some time to initialize. ::: -Upon committing, {command}`pre-commit` runs a set of checks as defined in the -file +Upon committing, {command}`pre-commit` runs a set of checks as defined in the file [{file}`.pre-commit-config.yaml`](https://github.com/ComPWA/ampform/blob/main/.pre-commit-config.yaml) -over all staged files. You can also quickly run all checks over _all_ indexed -files in the repository with the command: +over all staged files. You can also quickly run all checks over _all_ indexed files in +the repository with the command: ```shell pre-commit run -a ``` -Whenever you {ref}`submit a pull request `, this command -is automatically run +Whenever you {ref}`submit a pull request `, this command is +automatically run [on GitHub actions](https://github.com/ComPWA/ampform/actions/workflows/ci-style.yml) -and [on pre-commit.ci](https://results.pre-commit.ci/install/github/18435973) , -ensuring that all files in the repository follow the same conventions as set in -the config files of these tools. +and [on pre-commit.ci](https://results.pre-commit.ci/install/github/18435973) , ensuring +that all files in the repository follow the same conventions as set in the config files +of these tools. ### Tox @@ -282,9 +271,9 @@ More thorough checks can be run in one go with the following command: :::{margin} Running jobs in parallel -The {code}`-p` flag lets the jobs run in parallel. It also provides a nicer -overview of the progress. See -[`--parallel`](https://tox.wiki/en/latest/config.html#cmdoption-tox-p) flag. +The {code}`-p` flag lets the jobs run in parallel. It also provides a nicer overview of +the progress. See [`--parallel`](https://tox.wiki/en/latest/config.html#cmdoption-tox-p) +flag. ::: @@ -294,16 +283,16 @@ tox -p This command will [run `pytest`](#testing), perform all {ref}`style checks `, -{ref}`build the documentation `, and verify -cross-references in the documentation and the API. It's especially recommended -to **run tox before submitting a pull request!** +{ref}`build the documentation `, and verify cross-references in +the documentation and the API. It's especially recommended to **run tox before +submitting a pull request!** More specialized {command}`tox` job are defined in the -[`tox.ini`](https://github.com/ComPWA/ampform/blob/main/tox.ini) config file, -under each {code}`testenv` section. You can list all environments, along with a -description of what they do, by running: +[`tox.ini`](https://github.com/ComPWA/ampform/blob/main/tox.ini) config file, under each +{code}`testenv` section. You can list all environments, along with a description of what +they do, by running: ```shell tox -av @@ -313,71 +302,66 @@ tox -av All {ref}`style checks `, testing of the {ref}`documentation and links `, and -{ref}`unit tests ` are performed upon each pull request -through [GitHub Actions](https://docs.github.com/en/actions) (see status -overview [here](https://github.com/ComPWA/ampform/actions)). The checks are -defined under the -[`.github`](https://github.com/ComPWA/ampform/blob/main/.github) folder. All -checks performed for each PR have to pass before the PR can be merged. +{ref}`unit tests ` are performed upon each pull request through +[GitHub Actions](https://docs.github.com/en/actions) (see status overview +[here](https://github.com/ComPWA/ampform/actions)). The checks are defined under the +[`.github`](https://github.com/ComPWA/ampform/blob/main/.github) folder. All checks +performed for each PR have to pass before the PR can be merged. ## Style checks ### Formatting -Formatters are tools that automatically format source code, or some document. -Naturally, this speeds up your own programming, but these tools are -particularly important when {ref}`collaborating `, -because a standardized format avoids line conflicts in Git and makes diffs in -code review easier to read. +Formatters are tools that automatically format source code, or some document. Naturally, +this speeds up your own programming, but these tools are particularly important when +{ref}`collaborating `, because a standardized format avoids line +conflicts in Git and makes diffs in code review easier to read. For the Python source code, we use [`black`](https://black.readthedocs.io) and [`isort`](https://isort.readthedocs.io). For other code, we use -[Prettier](https://prettier.io). All of these formatters are "opinionated -formatters": they offer only limited configuration options, as to make -formatting as conform as possible. +[Prettier](https://prettier.io). All of these formatters are "opinionated formatters": +they offer only limited configuration options, as to make formatting as conform as +possible. -{ref}`develop:Pre-commit` performs some additional formatting jobs. For -instance, it formats Jupyter notebooks with -[nbQA](https://github.com/nbQA-dev/nbQA) and strips them of any output cells -with [`nbstripout`](https://github.com/kynan/nbstripout). +{ref}`develop:Pre-commit` performs some additional formatting jobs. For instance, it +formats Jupyter notebooks with [nbQA](https://github.com/nbQA-dev/nbQA) and strips them +of any output cells with [`nbstripout`](https://github.com/kynan/nbstripout). ### Linting -Linters point out when certain style conventions are not correctly followed. -Unlike with {ref}`formatters `, you have to fix the errors -yourself. As mentioned in {ref}`develop:Automated coding conventions`, style -conventions are formulated in config files. +Linters point out when certain style conventions are not correctly followed. Unlike with +{ref}`formatters `, you have to fix the errors yourself. As +mentioned in {ref}`develop:Automated coding conventions`, style conventions are +formulated in config files. ### Spelling Throughout this repository, we follow American English -([en-us](https://www.andiamo.co.uk/resources/iso-language-codes)) spelling -conventions. As a tool, we use +([en-us](https://www.andiamo.co.uk/resources/iso-language-codes)) spelling conventions. +As a tool, we use [cSpell](https://github.com/streetsidesoftware/cspell/blob/master/packages/cspell/README.md), -because it allows to check variable names in camel case and snake case. This -way, a spelling checker helps you avoid mistakes in the code as well! cSpell is -enforced through pre-commit. +because it allows to check variable names in camel case and snake case. This way, a +spelling checker helps you avoid mistakes in the code as well! cSpell is enforced +through pre-commit. Accepted words are tracked through the -[`.cspell.json`](https://github.com/ComPWA/ampform/blob/main/.cspell.json) -file. As with the other config files, -[`.cspell.json`](https://github.com/ComPWA/ampform/blob/main/.cspell.json) -formulates our conventions with regard to spelling and can be continuously -updated while our code base develops. In the file, the `words` section lists -words that you want to see as suggested corrections, while `ignoreWords` are -just the words that won't be flagged. Try to be sparse in adding words: if some -word is just specific to one file, you can -[ignore it inline](https://www.npmjs.com/package/cspell#ignore), or you can add -the file to the `ignorePaths` section if you want to ignore it completely. +[`.cspell.json`](https://github.com/ComPWA/ampform/blob/main/.cspell.json) file. As with +the other config files, +[`.cspell.json`](https://github.com/ComPWA/ampform/blob/main/.cspell.json) formulates +our conventions with regard to spelling and can be continuously updated while our code +base develops. In the file, the `words` section lists words that you want to see as +suggested corrections, while `ignoreWords` are just the words that won't be flagged. Try +to be sparse in adding words: if some word is just specific to one file, you can +[ignore it inline](https://www.npmjs.com/package/cspell#ignore), or you can add the file +to the `ignorePaths` section if you want to ignore it completely. It is easiest to use cSpell in {ref}`develop:Visual Studio Code`, through the [Code Spell Checker](https://marketplace.visualstudio.com/items?itemName=streetsidesoftware.code-spell-checker) -extension: it provides linting, suggests corrections from the {code}`words` -section, and enables you to quickly add or ignore words through the -[`.cspell.json`](https://github.com/ComPWA/ampform/blob/main/.cspell.json) -file. +extension: it provides linting, suggests corrections from the {code}`words` section, and +enables you to quickly add or ignore words through the +[`.cspell.json`](https://github.com/ComPWA/ampform/blob/main/.cspell.json) file. ## Testing @@ -396,8 +380,7 @@ The flag {command}`-n auto` causes {code}`pytest` to In VScode, you can visualize test coverage are covered with [Coverage Gutters](https://marketplace.visualstudio.com/items?itemName=ryanluker.vscode-coverage-gutters). -For this you need to run {command}`pytest` with the flag -{command}`--cov-report=xml`). +For this you need to run {command}`pytest` with the flag {command}`--cov-report=xml`). ::: @@ -430,21 +413,18 @@ Jupyter notebooks can also be used as tests. See more info ## Documentation -The documentation that you find on ComPWA pages like -[pwa.rtfd.io](https://pwa.rtfd.io) is built with -[Sphinx](https://www.sphinx-doc.org). Sphinx also -[builds the API page](https://www.sphinx-doc.org/en/master/man/sphinx-apidoc.html) -of the packages and therefore checks whether the -[docstrings](https://www.python.org/dev/peps/pep-0257) in the Python source -code are valid and correctly interlinked. - -We make use of [Markedly Structured Text](https://myst-parser.readthedocs.io) -(MyST), so you can write the documentation in both -[Markdown](https://www.markdownguide.org) and +The documentation that you find on ComPWA pages like [pwa.rtfd.io](https://pwa.rtfd.io) +is built with [Sphinx](https://www.sphinx-doc.org). Sphinx also +[builds the API page](https://www.sphinx-doc.org/en/master/man/sphinx-apidoc.html) of +the packages and therefore checks whether the +[docstrings](https://www.python.org/dev/peps/pep-0257) in the Python source code are +valid and correctly interlinked. + +We make use of [Markedly Structured Text](https://myst-parser.readthedocs.io) (MyST), so +you can write the documentation in both [Markdown](https://www.markdownguide.org) and [reStructuredText](https://www.sphinx-doc.org/en/master/usage/restructuredtext/basics.html). -In addition, it's easy to write (interactive) code examples in Jupyter -notebooks and host them on the website (see -[MyST-NB](https://myst-nb.readthedocs.io))! +In addition, it's easy to write (interactive) code examples in Jupyter notebooks and +host them on the website (see [MyST-NB](https://myst-nb.readthedocs.io))! ### Documentation preview @@ -457,8 +437,8 @@ tox -e doc If you are doing a lot of work on the documentation, -[`sphinx-autobuild`](https://pypi.org/project/sphinx-autobuild) is a nice tool -to use. Just run: +[`sphinx-autobuild`](https://pypi.org/project/sphinx-autobuild) is a nice tool to use. +Just run: @@ -466,13 +446,12 @@ to use. Just run: tox -e doclive ``` -This will start a server [http://127.0.0.1:8000](http://127.0.0.1:8000) where -you can continuously preview the changes you make to the documentation. +This will start a server [http://127.0.0.1:8000](http://127.0.0.1:8000) where you can +continuously preview the changes you make to the documentation. -Finally, a nice feature of [Read the Docs](https://readthedocs.org), where we -host our documentation, is that documentation is built for each pull request as -well. This means that you can view the documentation for your changes as well. -For more info, see +Finally, a nice feature of [Read the Docs](https://readthedocs.org), where we host our +documentation, is that documentation is built for each pull request as well. This means +that you can view the documentation for your changes as well. For more info, see [here](https://docs.readthedocs.io/en/stable/guides/autobuild-docs-for-pull-requests.html), or just click "details" under the RTD check once you submit your PR. @@ -483,30 +462,28 @@ or just click "details" under the RTD check once you submit your PR. :::{tip} Sometimes it happens that your Jupyter installation does not recognize your -{ref}`virtual environment `. In that case, have a -look at +{ref}`virtual environment `. In that case, have a look at [these instructions](https://ipython.readthedocs.io/en/stable/install/kernel_install.html#kernels-for-different-environments). ::: :::: -The [docs](https://github.com/ComPWA/ampform/tree/main/docs) folder can also -contain Jupyter notebooks. These notebooks are rendered as HTML by -[MyST-NB](https://myst-nb.readthedocs.io). The notebooks are also run and -tested whenever you make a {ref}`pull request `, so they -also serve as **integration tests**. +The [docs](https://github.com/ComPWA/ampform/tree/main/docs) folder can also contain +Jupyter notebooks. These notebooks are rendered as HTML by +[MyST-NB](https://myst-nb.readthedocs.io). The notebooks are also run and tested +whenever you make a {ref}`pull request `, so they also serve as +**integration tests**. If you want to improve those notebooks, we recommend working with [Jupyter Lab](https://jupyterlab.readthedocs.io/en/stable), which is -{ref}`installed with the dev requirements `. -Jupyter Lab offers a nicer developer experience than the default Jupyter -notebook editor does. A few useful Jupyter Lab plugins are also installed -through the {ref}`optional dependencies `. +{ref}`installed with the dev requirements `. Jupyter Lab +offers a nicer developer experience than the default Jupyter notebook editor does. A few +useful Jupyter Lab plugins are also installed through the +{ref}`optional dependencies `. -Now, if you want to test all notebooks in the documentation folder and check -what their output cells will look like in the {ref}`develop:Documentation`, you -can do this with: +Now, if you want to test all notebooks in the documentation folder and check what their +output cells will look like in the {ref}`develop:Documentation`, you can do this with: @@ -514,10 +491,9 @@ can do this with: tox -e docnb ``` -This command takes more time than `tox -e doc`, but it is good practice to do -this before you submit a pull request. It's also possible to continuously -generate the HTML pages _including cell output_ while you work on the notebooks -with: +This command takes more time than `tox -e doc`, but it is good practice to do this +before you submit a pull request. It's also possible to continuously generate the HTML +pages _including cell output_ while you work on the notebooks with: ```shell EXECUTE_NB= tox -e doclive @@ -525,10 +501,9 @@ EXECUTE_NB= tox -e doclive :::{tip} -Notebooks are automatically formatted through -{ref}`pre-commit ` (see {ref}`develop:Formatting`). If you -want to format the notebooks automatically as you're working, you can do so -with +Notebooks are automatically formatted through {ref}`pre-commit ` +(see {ref}`develop:Formatting`). If you want to format the notebooks automatically as +you're working, you can do so with [`jupyterlab-code-formatter`](https://ryantam626.github.io/jupyterlab_code_formatter/index.html), which is automatically {ref}`installed with the dev requirements `. @@ -537,44 +512,42 @@ which is automatically ## Collaboration -The source code of all ComPWA repositories is maintained with -[Git](https://git-scm.com) and [GitHub](https://github.com). We keep track of -issues with the code, documentation, and developer set-up with GitHub issues -(see for instance [here](https://github.com/ComPWA/ampform/issues)). This is -also the place where you can +The source code of all ComPWA repositories is maintained with [Git](https://git-scm.com) +and [GitHub](https://github.com). We keep track of issues with the code, documentation, +and developer set-up with GitHub issues (see for instance +[here](https://github.com/ComPWA/ampform/issues)). This is also the place where you can [report bugs](https://github.com/ComPWA/ampform/issues/new/choose). ### Issue management -We keep track of issue dependencies, time estimates, planning, pipeline -statuses, et cetera with GitHub project boards -([GitHub Issues](https://github.com/features/issues)). The main project boards -are: +We keep track of issue dependencies, time estimates, planning, pipeline statuses, et +cetera with GitHub project boards ([GitHub Issues](https://github.com/features/issues)). +The main project boards are: - [Public roadmap](https://github.com/orgs/ComPWA/projects/4) - [Development planning](https://github.com/orgs/ComPWA/projects/5) -Some issues are not public. To get access, you can request to become member of -the [ComPWA](https://github.com/ComPWA) GitHub organization. Other information -that is publicly available are: +Some issues are not public. To get access, you can request to become member of the +[ComPWA](https://github.com/ComPWA) GitHub organization. Other information that is +publicly available are: -- [Issue labels](https://github.com/ComPWA/ampform/labels): help to categorize - issues by type (maintenance, enhancement, bug, etc.). The labels are also - used to in the sub-sections of the release notes. +- [Issue labels](https://github.com/ComPWA/ampform/labels): help to categorize issues by + type (maintenance, enhancement, bug, etc.). The labels are also used to in the + sub-sections of the release notes. - [Milestones](https://github.com/ComPWA/ampform/milestones?direction=asc&sort=title&state=open): way to bundle issues and PRs for upcoming releases. - [Releases](https://github.com/ComPWA/ampform/releases). -All of these are important for the {ref}`develop:Release flow` and therefore -also serve as a way to document the framework. +All of these are important for the {ref}`develop:Release flow` and therefore also serve +as a way to document the framework. ### Branching model While our aim is to maintain long-term, stable projects, -{ref}`PWA software projects ` are academic -projects that are subject to change and often require swift modifications or -new features for ongoing analyses. For this reason, we work in different layers -of development. These layers are represented by Git branches. +{ref}`PWA software projects ` are academic projects +that are subject to change and often require swift modifications or new features for +ongoing analyses. For this reason, we work in different layers of development. These +layers are represented by Git branches. ![Epic branches](./_static/epic-branches.svg) @@ -582,9 +555,9 @@ of development. These layers are represented by Git branches. #### `stable` branch -Represents the _latest release_ of the package that can be found on both the -GitHub release page and on PyPI (see {ref}`develop:Release flow`). The -documentation of the `stable` branch is also the default view +Represents the _latest release_ of the package that can be found on both the GitHub +release page and on PyPI (see {ref}`develop:Release flow`). The documentation of the +`stable` branch is also the default view {ref}`you see on Read the Docs ` (RTD). See e.g. [ampform.rtfd.io/en/stable](https://ampform.rtfd.io/en/stable). @@ -592,37 +565,35 @@ documentation of the `stable` branch is also the default view #### `main` branch -Represents the upcoming release of the package. This branch is not guaranteed -to be stable, but has high {ref}`CI standards ` and can -only be updated through reviewed pull requests. The documentation of the `main` -branch can be found on RTD under "latest", see e.g. +Represents the upcoming release of the package. This branch is not guaranteed to be +stable, but has high {ref}`CI standards ` and can only be +updated through reviewed pull requests. The documentation of the `main` branch can be +found on RTD under "latest", see e.g. [ampform.rtfd.io/en/latest](https://ampform.rtfd.io/en/latest). #### Epic branches :::{margin} -The word ["epic"](https://www.atlassian.com/agile/project-management/epics) is -used in +The word ["epic"](https://www.atlassian.com/agile/project-management/epics) is used in [agile software development](https://en.wikipedia.org/wiki/Agile_software_development). ::: -When working on a feature or larger refactoring that may take a longer time -(think of implementing a new PWA formalism), we isolate its development under -an 'epic branch', separate from the `main` branch. Eventually, this epic branch -is to be merged back into the `main`, until then it is available for discussion -and testing. +When working on a feature or larger refactoring that may take a longer time (think of +implementing a new PWA formalism), we isolate its development under an 'epic branch', +separate from the `main` branch. Eventually, this epic branch is to be merged back into +the `main`, until then it is available for discussion and testing. -Pull requests to an epic branch require no code review and the CI checks are -less strict. This allows for faster development, while still offering the -possibility to discuss new implementations and +Pull requests to an epic branch require no code review and the CI checks are less +strict. This allows for faster development, while still offering the possibility to +discuss new implementations and {ref}`keeping track of related issues `. -Epic branches can be installed through PyPI as well. Say that a certain epic is -located under the branch `epic/some-title` and that the source code is located -under [https://github.com/ComPWA/ampform](https://github.com/compwa/ampform), -it can be installed as follows: +Epic branches can be installed through PyPI as well. Say that a certain epic is located +under the branch `epic/some-title` and that the source code is located under +[https://github.com/ComPWA/ampform](https://github.com/compwa/ampform), it can be +installed as follows: ```shell python3 -m pip install git+https://github.com/ComPWA/ampform@epic/some-title @@ -630,22 +601,19 @@ python3 -m pip install git+https://github.com/ComPWA/ampform@epic/some-title #### Feature branches -The {ref}`main-branch` and {ref}`develop:epic branches` can be updated through -pull requests. It is best to create such a pull request from a separate branch, -which does not have any CI or code review restrictions. We call this a "feature -branch". +The {ref}`main-branch` and {ref}`develop:epic branches` can be updated through pull +requests. It is best to create such a pull request from a separate branch, which does +not have any CI or code review restrictions. We call this a "feature branch". ### Commit conventions -- Please use - [conventional commit messages](https://www.conventionalcommits.org): start +- Please use [conventional commit messages](https://www.conventionalcommits.org): start the commit with a semantic keyword (see e.g. - [Angular](https://github.com/angular/angular/blob/master/CONTRIBUTING.md#type) - or - [these examples](https://seesparkbox.com/foundry/semantic_commit_messages)), - followed by [a column](https://git-scm.com/docs/git-interpret-trailers), then - the message. The message itself should be in imperative mood — just imagine - the commit to give a command to the code framework. So for instance: + [Angular](https://github.com/angular/angular/blob/master/CONTRIBUTING.md#type) or + [these examples](https://seesparkbox.com/foundry/semantic_commit_messages)), followed + by [a column](https://git-scm.com/docs/git-interpret-trailers), then the message. The + message itself should be in imperative mood — just imagine the commit to give a + command to the code framework. So for instance: ```none ci: implement coverage report tools @@ -654,76 +622,71 @@ branch". docs!: rewrite welcome pages ``` -- Keep pull requests small. If the issue you try to address is too big, discuss - in the team whether the issue can be converted into an - [Epic](https://blog.zenhub.com/working-with-epics-in-github) and split up - into smaller tasks. +- Keep pull requests small. If the issue you try to address is too big, discuss in the + team whether the issue can be converted into an + [Epic](https://blog.zenhub.com/working-with-epics-in-github) and split up into smaller + tasks. - Before creating a pull request, run {ref}`develop:Tox`. -- Also use a [conventional commit message](https://www.conventionalcommits.org) - style for the PR title. This is because we follow a +- Also use a [conventional commit message](https://www.conventionalcommits.org) style + for the PR title. This is because we follow a [linear commit history](https://docs.github.com/en/github/administering-a-repository/requiring-a-linear-commit-history) - and the PR title will become the eventual commit message. A linear commit - history is important for the {ref}`develop:Release flow` and it is easier to - navigate through changes once something goes wrong. In fact, in a linear - commit history, commits that a have been merged into the {ref}`main-branch` - become more like small intermediate patches between the minor and major - releases. + and the PR title will become the eventual commit message. A linear commit history is + important for the {ref}`develop:Release flow` and it is easier to navigate through + changes once something goes wrong. In fact, in a linear commit history, commits that a + have been merged into the {ref}`main-branch` become more like small intermediate + patches between the minor and major releases. Note that a conventional commit message style is - {ref}`enforced through GitHub Actions `, as well as a - check on {ref}`PR labels ` (see overview + {ref}`enforced through GitHub Actions `, as well as a check on + {ref}`PR labels ` (see overview [here](https://github.com/ComPWA/ampform/actions?query=workflow%3A%22PR+linting%22)). The corresponding configuration file is [`commitlint.config.js`](https://github.com/ComPWA/ampform/blob/main/commitlint.config.js). -- PRs can only be merged through 'squash and merge'. There, you will see a - summary based on the separate commits that constitute this PR. Leave the - relevant commits in as bullet points. See the - [commit history](https://github.com/ComPWA/ampform/commits/main) for - examples. This comes in especially handy when - {ref}`drafting a release `! +- PRs can only be merged through 'squash and merge'. There, you will see a summary based + on the separate commits that constitute this PR. Leave the relevant commits in as + bullet points. See the + [commit history](https://github.com/ComPWA/ampform/commits/main) for examples. This + comes in especially handy when {ref}`drafting a release `! ### Release flow Releases are managed with the [GitHub release page](https://docs.github.com/en/free-pro-team@latest/github/administering-a-repository/managing-releases-in-a-repository), -see for instance -[the one for AmpForm](https://github.com/ComPWA/ampform/releases). The release -notes there are +see for instance [the one for AmpForm](https://github.com/ComPWA/ampform/releases). The +release notes there are [automatically generated from the PRs](https://github.com/release-drafter/release-drafter) -that were merged into the main branch since the previous tag and can be viewed -and edited as a release draft if you are a member of the ComPWA organization. -Each of the entries are generated from the PR titles, categorized by issue -label (see configuration in +that were merged into the main branch since the previous tag and can be viewed and +edited as a release draft if you are a member of the ComPWA organization. Each of the +entries are generated from the PR titles, categorized by issue label (see configuration +in [`.github/release-drafter.yml`](https://github.com/ComPWA/ampform/blob/main/.github/release-drafter.yml)). -Once a release is made on GitHub for a repository with source code for a Python -package, a new version is automatically published on [PyPI](https://pypi.org) -and the [`stable` branch](#stable-branch) is updated to this latest tag. The -package version is taken from the Git tag associated with the release on GitHub -(see [setuptools-scm](https://pypi.org/project/setuptools-scm)). This way, the -release notes on GitHub serve as a changelog as well! +Once a release is made on GitHub for a repository with source code for a Python package, +a new version is automatically published on [PyPI](https://pypi.org) and the +[`stable` branch](#stable-branch) is updated to this latest tag. The package version is +taken from the Git tag associated with the release on GitHub (see +[setuptools-scm](https://pypi.org/project/setuptools-scm)). This way, the release notes +on GitHub serve as a changelog as well! -Release tags **have to follow the -[Semantic Versioning scheme](https://semver.org)**! This ensures that the tag -can be used by [setuptools-scm](https://pypi.org/project/setuptools-scm) (in -case the repository is a Python package). In addition, -{ref}`milestones ` with the same name as the release -tag are automatically closed. +Release tags **have to follow the [Semantic Versioning scheme](https://semver.org)**! +This ensures that the tag can be used by +[setuptools-scm](https://pypi.org/project/setuptools-scm) (in case the repository is a +Python package). In addition, {ref}`milestones ` with the same +name as the release tag are automatically closed. ## Code editors -Even though we try to standardize the developer set-up of the repositories, we -encourage you to use the code editors that you feel comfortable with. Where -possible, we therefore define settings of linters, formatters, etc in config -files that are specific to those tools (such as `.pylintrc`), not in the -configuration files of the editors. +Even though we try to standardize the developer set-up of the repositories, we encourage +you to use the code editors that you feel comfortable with. Where possible, we therefore +define settings of linters, formatters, etc in config files that are specific to those +tools (such as `.pylintrc`), not in the configuration files of the editors. -Still, where code editor settings can be shared through configuration files in -the repository, we provide recommended settings for the code editor as well. -This is especially the case for [VSCode](#visual-studio-code). +Still, where code editor settings can be shared through configuration files in the +repository, we provide recommended settings for the code editor as well. This is +especially the case for [VSCode](#visual-studio-code). :::{tip} @@ -741,17 +704,16 @@ So we'll gladly integrate your editor settings where possible as you ### Visual Studio code -We recommend using [Visual Studio Code](https://code.visualstudio.com) as it's -free, regularly updated, and very flexible through it's wide offer of user -extensions. +We recommend using [Visual Studio Code](https://code.visualstudio.com) as it's free, +regularly updated, and very flexible through it's wide offer of user extensions. If you add or open this repository as a -[VSCode workspace](https://code.visualstudio.com/docs/editor/multi-root-workspaces), -the file +[VSCode workspace](https://code.visualstudio.com/docs/editor/multi-root-workspaces), the +file [`.vscode/settings.json`](https://github.com/ComPWA/ampform/blob/main/.vscode/settings.json) -will ensure that you have the right developer settings for this repository. In -addition, VSCode will automatically recommend you to install a number of -extensions that we use when working on this code base. +will ensure that you have the right developer settings for this repository. In addition, +VSCode will automatically recommend you to install a number of extensions that we use +when working on this code base. [They are defined](https://code.visualstudio.com/updates/v1_6#_workspace-extension-recommendations) in the [`.vscode/extensions.json`](https://github.com/ComPWA/ampform/blob/main/.vscode/extensions.json) @@ -763,9 +725,8 @@ as the VSCode settings that come with this are folder settings. :::{dropdown} Conda and VSCode -ComPWA projects are best developed -{ref}`with Conda ` and VSCode. The complete -developer install procedure then becomes: +ComPWA projects are best developed {ref}`with Conda ` and +VSCode. The complete developer install procedure then becomes: ```shell git clone https://github.com/ComPWA/ampform.git # or some other repo @@ -779,11 +740,11 @@ code . # open folder in VSCode ## Writing durable software -ComPWA strives to follow best practices from software development in industry. -Following these standards not only makes the code easier to maintain and the -software more reliable, it also provides you with the opportunity to learn -about these practices while developing the code-base. Below you can find some -resources we highly recommend you to be familiar with. +ComPWA strives to follow best practices from software development in industry. Following +these standards not only makes the code easier to maintain and the software more +reliable, it also provides you with the opportunity to learn about these practices while +developing the code-base. Below you can find some resources we highly recommend you to +be familiar with. :::{margin} @@ -799,8 +760,8 @@ Do you have other recommendations? Edit this page ``` - [Complete Python Development Guide ― testdriven.io](https://testdriven.io/guides/complete-python) -- _Effective Python: 90 Specific Ways to Write Better Python_ (2019) by Brett - Slatkin {cite}`slatkinEffectivePython902019`. See also +- _Effective Python: 90 Specific Ways to Write Better Python_ (2019) by Brett Slatkin + {cite}`slatkinEffectivePython902019`. See also [its website](https://effectivepython.com). - [Scikit-HEP Developer Information](https://scikit-hep.org/developer) @@ -809,10 +770,10 @@ Do you have other recommendations? Edit this page ``` - [Clean Code lesson series by "Uncle Bob" on YouTube](https://youtu.be/7EmboKQH8lM) -- _Clean Code: A Handbook of Agile Software Craftsmanship_ (2009) by Robert - Martin ("Uncle Bob") {cite}`martinCleanCodeHandbook2009` -- [This gist](https://gist.github.com/wojteklu/73c6914cc446146b8b533c0988cf8d29) - with a comprehensive summary of the core principles of Martin's _Clean Code_ +- _Clean Code: A Handbook of Agile Software Craftsmanship_ (2009) by Robert Martin + ("Uncle Bob") {cite}`martinCleanCodeHandbook2009` +- [This gist](https://gist.github.com/wojteklu/73c6914cc446146b8b533c0988cf8d29) with a + comprehensive summary of the core principles of Martin's _Clean Code_ ```{rubric} Test-Driven Development @@ -834,13 +795,13 @@ Do you have other recommendations? Edit this page [Design Patterns](https://sourcemaking.com/design_patterns) and [AntiPatterns](https://sourcemaking.com/antipatterns) - Composition over inheritance: - [Subclassing in Python Redux](https://hynek.me/articles/python-subclassing-redux) - by Hynek Schlawack. A comprehensive article on the topic with illustrative - examples in Python an several references to other important articles. + [Subclassing in Python Redux](https://hynek.me/articles/python-subclassing-redux) by + Hynek Schlawack. A comprehensive article on the topic with illustrative examples in + Python an several references to other important articles. - The classic - ["Gang of Four" (GOF) book](https://en.wikipedia.org/wiki/Design_Patterns): - _Design Patterns_ (1994) by Erich Gamma, Richard Helm, Ralph Johnson, and - John Vlissides {cite}`gammaDesignPatternsElements1995` + ["Gang of Four" (GOF) book](https://en.wikipedia.org/wiki/Design_Patterns): _Design + Patterns_ (1994) by Erich Gamma, Richard Helm, Ralph Johnson, and John Vlissides + {cite}`gammaDesignPatternsElements1995` - [C++ Core Guidelines](https://isocpp.github.io/CppCoreGuidelines/CppCoreGuidelines): while this document provides intended for C++ developers, it is an excellent, up-to-date set of guidelines that apply to any programming language. @@ -851,7 +812,7 @@ Do you have other recommendations? Edit this page - _Algorithms_ (2011) by Robert Sedgewick and Kevin Wayne {cite}`sedgewickAlgorithms2011`. See also - [their website](https://algs4.cs.princeton.edu), with several links to other - resources like their Coursera course. -- [LeetCode](https://leetcode.com/problemset): practice algorithms through - coding problems + [their website](https://algs4.cs.princeton.edu), with several links to other resources + like their Coursera course. +- [LeetCode](https://leetcode.com/problemset): practice algorithms through coding + problems diff --git a/docs/index.md b/docs/index.md index c1c106d7..6ccac7cd 100644 --- a/docs/index.md +++ b/docs/index.md @@ -13,10 +13,9 @@ This page is combines documentation on projects provided by the -[ComPWA organization on GitHub](https://github.com/ComPWA). It is more -technical than the [PWA Pages](https://pwa.rtfd.io) and focuses on the ComPWA -organization only. Read more about our ideals and ongoing projects on the -{doc}`/about` page. +[ComPWA organization on GitHub](https://github.com/ComPWA). It is more technical than +the [PWA Pages](https://pwa.rtfd.io) and focuses on the ComPWA organization only. Read +more about our ideals and ongoing projects on the {doc}`/about` page. ```{rubric} Table of contents diff --git a/docs/report/000.ipynb b/docs/report/000.ipynb index 835762b5..d053abf9 100644 --- a/docs/report/000.ipynb +++ b/docs/report/000.ipynb @@ -80,7 +80,7 @@ }, "outputs": [], "source": [ - "%pip install -q black==21.5b2 jax==0.2.13 jaxlib==0.1.67 numpy==1.19.5 sympy==1.8" + "%pip install -q black==21.5b2 jax==0.2.13 jaxlib==0.1.67 numpy==1.23 sympy==1.8" ] }, { @@ -328,7 +328,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Numpy also offers a special function that evaluates negative values even if the input values are real: {func}`numpy.lib.scimath.sqrt`:" + "Numpy also offers a special function that evaluates negative values even if the input values are real: {func}`numpy.emath.sqrt`:" ] }, { @@ -348,16 +348,14 @@ } ], "source": [ - "from numpy.lib.scimath import sqrt as csqrt\n", - "\n", - "csqrt(-1)" + "np.emath.sqrt(-1)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Unfortunately, the {mod}`jax.numpy` API does not interface to {mod}`numpy.lib.scimath`. It is possible to decorate {func}`numpy.lib.scimath.sqrt` be decorated with {func}`jax.jit`, but that **only works with static, hashable arguments**:" + "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**:" ] }, { @@ -366,7 +364,8 @@ "metadata": { "tags": [ "raises-exception", - "hide-output" + "hide-output", + "scroll-output" ] }, "outputs": [ @@ -377,38 +376,55 @@ "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mUnfilteredStackTrace\u001b[0m Traceback (most recent call last)", - "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0mjax_csqrt_error\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcsqrt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"cpu\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mjax_csqrt_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/traceback_util.py\u001b[0m in \u001b[0;36mreraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 143\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 144\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mcache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 425\u001b[0m \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_fun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_tree\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 426\u001b[0;31m out_flat = xla.xla_call(\n\u001b[0m\u001b[1;32m 427\u001b[0m \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mbind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1564\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mbind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1565\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mcall_bind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1566\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1555\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mmaybe_new_sublevel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1556\u001b[0;31m \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1557\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_lower\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mapply_todos\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv_trace_todo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m 1567\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mprocess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1568\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1569\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m 608\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 609\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimpl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 610\u001b[0m \u001b[0mprocess_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py\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[0;32mdef\u001b[0m \u001b[0m_xla_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mWrappedFun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdonated_invars\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 578\u001b[0;31m compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001b[0m\u001b[1;32m 579\u001b[0m *unsafe_map(arg_spec, args))\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mmemoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 263\u001b[0m \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mans\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py\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 \u001b[0mabstract_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0munzip2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg_specs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 652\u001b[0;31m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_avals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_to_jaxpr_final\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mabstract_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtransform_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"jit\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 653\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTracer\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py\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 \u001b[0mmain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjaxpr_stack\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# type: ignore\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1209\u001b[0;31m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_avals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrace_to_subjaxpr_dynamic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_avals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1210\u001b[0m \u001b[0;32mdel\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py\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 \u001b[0min_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew_arg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_avals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1188\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_wrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0min_tracers\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1189\u001b[0m \u001b[0mout_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull_raise\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mans\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 166\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 167\u001b[0m \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m<__array_function__ internals>\u001b[0m in \u001b[0;36msqrt\u001b[0;34m(*args, **kwargs)\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py\u001b[0m in \u001b[0;36msqrt\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 222\u001b[0m \"\"\"\n\u001b[0;32m--> 223\u001b[0;31m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_fix_real_lt_zero\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 224\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mnx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py\u001b[0m in \u001b[0;36m_fix_real_lt_zero\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 118\u001b[0m \"\"\"\n\u001b[0;32m--> 119\u001b[0;31m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0masarray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 120\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misreal\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m&\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/core/_asarray.py\u001b[0m in \u001b[0;36masarray\u001b[0;34m(a, dtype, order)\u001b[0m\n\u001b[1;32m 82\u001b[0m \"\"\"\n\u001b[0;32m---> 83\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morder\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0morder\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 84\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36m__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__array__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkw\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 472\u001b[0;31m \u001b[0;32mraise\u001b[0m \u001b[0mTracerArrayConversionError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 473\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", + "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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)", - "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0mjax_csqrt_error\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcsqrt\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"cpu\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mjax_csqrt_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m<__array_function__ internals>\u001b[0m in \u001b[0;36msqrt\u001b[0;34m(*args, **kwargs)\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py\u001b[0m in \u001b[0;36msqrt\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 221\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 222\u001b[0m \"\"\"\n\u001b[0;32m--> 223\u001b[0;31m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_fix_real_lt_zero\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 224\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mnx\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 225\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/lib/scimath.py\u001b[0m in \u001b[0;36m_fix_real_lt_zero\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 117\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 118\u001b[0m \"\"\"\n\u001b[0;32m--> 119\u001b[0;31m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0masarray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 120\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misreal\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;34m&\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 121\u001b[0m \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0m_tocomplex\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/core/_asarray.py\u001b[0m in \u001b[0;36masarray\u001b[0;34m(a, dtype, order)\u001b[0m\n\u001b[1;32m 81\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 82\u001b[0m \"\"\"\n\u001b[0;32m---> 83\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0marray\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdtype\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcopy\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;32mFalse\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0morder\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0morder\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 84\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 85\u001b[0m \u001b[0;34m\u001b[0m\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<__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-org/lib/python3.8/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-org/lib/python3.8/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(csqrt, backend=\"cpu\")\n", + "jax_csqrt_error = jax.jit(np.emath.sqrt, backend=\"cpu\")\n", "jax_csqrt_error(-1)" ] }, @@ -429,7 +445,7 @@ } ], "source": [ - "jax_csqrt = jax.jit(csqrt, backend=\"cpu\", static_argnums=0)\n", + "jax_csqrt = jax.jit(np.emath.sqrt, backend=\"cpu\", static_argnums=0)\n", "jax_csqrt(-1)" ] }, @@ -449,7 +465,7 @@ "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_csqrt\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", + "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" ] } @@ -617,27 +633,46 @@ "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mUnfilteredStackTrace\u001b[0m Traceback (most recent call last)", - "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0mjax_complex_sqrt_no_select\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp_complex_sqrt_no_select\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mjax_complex_sqrt_no_select\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/traceback_util.py\u001b[0m in \u001b[0;36mreraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 143\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 144\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mcache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 425\u001b[0m \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_fun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_tree\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 426\u001b[0;31m out_flat = xla.xla_call(\n\u001b[0m\u001b[1;32m 427\u001b[0m \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mbind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1564\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mbind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1565\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mcall_bind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1566\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1555\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mmaybe_new_sublevel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1556\u001b[0;31m \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1557\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_lower\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mapply_todos\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv_trace_todo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m 1567\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mprocess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1568\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1569\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m 608\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 609\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimpl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 610\u001b[0m \u001b[0mprocess_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py\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[0;32mdef\u001b[0m \u001b[0m_xla_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mWrappedFun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdonated_invars\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 578\u001b[0;31m compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001b[0m\u001b[1;32m 579\u001b[0m *unsafe_map(arg_spec, args))\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mmemoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 263\u001b[0m \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mans\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py\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 \u001b[0mabstract_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0munzip2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg_specs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 652\u001b[0;31m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_avals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_to_jaxpr_final\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mabstract_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtransform_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"jit\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 653\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTracer\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py\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 \u001b[0mmain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjaxpr_stack\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# type: ignore\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1209\u001b[0;31m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_avals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrace_to_subjaxpr_dynamic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_avals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1210\u001b[0m \u001b[0;32mdel\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py\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 \u001b[0min_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew_arg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_avals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1188\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_wrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0min_tracers\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1189\u001b[0m \u001b[0mout_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull_raise\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mans\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 166\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 167\u001b[0m \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/__init__.py\u001b[0m in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_lambdifygenerated\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1j\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36m__bool__\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 537\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__nonzero__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_nonzero\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 538\u001b[0;31m \u001b[0;32mdef\u001b[0m \u001b[0m__bool__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_bool\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 539\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__int__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0maval\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_int\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36merror\u001b[0;34m(self, arg)\u001b[0m\n\u001b[1;32m 959\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0merror\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 960\u001b[0;31m \u001b[0;32mraise\u001b[0m \u001b[0mConcretizationTypeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfname_context\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 961\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0merror\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", + "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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)", - "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0mjax_complex_sqrt_no_select\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mjax\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnp_complex_sqrt_no_select\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0mjax_complex_sqrt_no_select\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/__init__.py\u001b[0m in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_lambdifygenerated\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m1j\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m \u001b[0;34m<\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32melse\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", + "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)" ] } @@ -788,26 +823,45 @@ "traceback": [ "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", "\u001b[0;31mUnfilteredStackTrace\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_complex_sqrt_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/traceback_util.py\u001b[0m in \u001b[0;36mreraise_with_filtered_traceback\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 143\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 144\u001b[0m \u001b[0;32mexcept\u001b[0m \u001b[0mException\u001b[0m \u001b[0;32mas\u001b[0m \u001b[0me\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/_src/api.py\u001b[0m in \u001b[0;36mcache_miss\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m 425\u001b[0m \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_tree\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mflatten_fun\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_tree\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 426\u001b[0;31m out_flat = xla.xla_call(\n\u001b[0m\u001b[1;32m 427\u001b[0m \u001b[0mflat_fun\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mbind\u001b[0;34m(self, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1564\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mbind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1565\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mcall_bind\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1566\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mcall_bind\u001b[0;34m(primitive, fun, *args, **params)\u001b[0m\n\u001b[1;32m 1555\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mmaybe_new_sublevel\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1556\u001b[0;31m \u001b[0mouts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtop_trace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1557\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfull_lower\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mapply_todos\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0menv_trace_todo\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mouts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess\u001b[0;34m(self, trace, fun, tracers, params)\u001b[0m\n\u001b[1;32m 1567\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mprocess\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1568\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1569\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36mprocess_call\u001b[0;34m(self, primitive, f, tracers, params)\u001b[0m\n\u001b[1;32m 608\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 609\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0mprimitive\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mimpl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0mtracers\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 610\u001b[0m \u001b[0mprocess_map\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mprocess_call\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py\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[0;32mdef\u001b[0m \u001b[0m_xla_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mlu\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mWrappedFun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdevice\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbackend\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdonated_invars\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 578\u001b[0;31m compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,\n\u001b[0m\u001b[1;32m 579\u001b[0m *unsafe_map(arg_spec, args))\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mmemoized_fun\u001b[0;34m(fun, *args)\u001b[0m\n\u001b[1;32m 261\u001b[0m \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 262\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mcall\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 263\u001b[0m \u001b[0mcache\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mkey\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mans\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstores\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/xla.py\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 \u001b[0mabstract_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0m_\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0munzip2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg_specs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 652\u001b[0;31m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_avals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mpe\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtrace_to_jaxpr_final\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mabstract_args\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtransform_name\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"jit\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 653\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0many\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0misinstance\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mc\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcore\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mTracer\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mc\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mconsts\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py\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 \u001b[0mmain\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mjaxpr_stack\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# type: ignore\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1209\u001b[0;31m \u001b[0mjaxpr\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mout_avals\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mconsts\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrace_to_subjaxpr_dynamic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_avals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1210\u001b[0m \u001b[0;32mdel\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmain\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/interpreters/partial_eval.py\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 \u001b[0min_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnew_arg\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0min_avals\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1188\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfun\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcall_wrapped\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0min_tracers\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1189\u001b[0m \u001b[0mout_tracers\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmap\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mtrace\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mfull_raise\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mans\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/linear_util.py\u001b[0m in \u001b[0;36mcall_wrapped\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m 165\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 166\u001b[0;31m \u001b[0mans\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mdict\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mparams\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 167\u001b[0m \u001b[0;32mexcept\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/__init__.py\u001b[0m in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_lambdifygenerated\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mselect\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mless\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m1j\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdefault\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/jax/core.py\u001b[0m in \u001b[0;36m__array__\u001b[0;34m(self, *args, **kw)\u001b[0m\n\u001b[1;32m 471\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m__array__\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkw\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 472\u001b[0;31m \u001b[0;32mraise\u001b[0m \u001b[0mTracerArrayConversionError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 473\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", + "File \u001b[0;32m~/miniconda3/envs/compwa-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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-org/lib/python3.8/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)", - "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mjax_complex_sqrt_error\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0;36m1\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", - "\u001b[0;32m~/anaconda3/envs/compwa-org/lib/python3.8/site-packages/numpy/__init__.py\u001b[0m in \u001b[0;36m_lambdifygenerated\u001b[0;34m(x)\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0m_lambdifygenerated\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m \u001b[0;32mreturn\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mselect\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mless\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m0\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;32mTrue\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m[\u001b[0m\u001b[0;36m1j\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m-\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0msqrt\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mdefault\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mnan\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m", + "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)" ] } diff --git a/docs/report/001.ipynb b/docs/report/001.ipynb index d5413ca5..ce2da436 100644 --- a/docs/report/001.ipynb +++ b/docs/report/001.ipynb @@ -364,7 +364,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Just as noted in {ref}`report/000:Complex square root` though, {mod}`numpy.lib.scimath` is not provided by the NumPy API of {doc}`JAX `. As discussed there, we can at most decorate the {mod}`numpy.lib.scimath` version with {func}`jax.jit` and work with static arguments only:" + "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:" ] }, { @@ -436,7 +436,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "As concluded in {ref}`report/000:Conditional square root`, the alternative to lambdify to {mod}`numpy.lib.scimath` 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", + "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:" ] diff --git a/docs/report/002.ipynb b/docs/report/002.ipynb index 3320f6ff..a579e8c8 100644 --- a/docs/report/002.ipynb +++ b/docs/report/002.ipynb @@ -161,9 +161,7 @@ " allowed_interaction_types=[\"strong\", \"EM\"],\n", " formalism_type=\"canonical-helicity\",\n", ")\n", - "dot = qrules.io.asdot(\n", - " result, collapse_graphs=True, render_final_state_id=False\n", - ")\n", + "dot = qrules.io.asdot(result, collapse_graphs=True, render_final_state_id=False)\n", "graphviz.Source(dot)" ] }, @@ -281,9 +279,7 @@ ], "source": [ "amplitudes = {\n", - " name: expr\n", - " for name, expr in model.components.items()\n", - " if name.startswith(\"A\")\n", + " name: expr for name, expr in model.components.items() if name.startswith(\"A\")\n", "}\n", "sorted(amplitudes)" ] @@ -295,9 +291,7 @@ "outputs": [], "source": [ "intensities = {\n", - " name: expr\n", - " for name, expr in model.components.items()\n", - " if name.startswith(\"I\")\n", + " name: expr for name, expr in model.components.items() if name.startswith(\"I\")\n", "}" ] }, @@ -740,9 +734,7 @@ "outputs": [], "source": [ "sorted_amplitude_symbols = sorted(np_amplitudes, key=lambda s: s.name)\n", - "np_amplitude_expr = sp.lambdify(\n", - " sorted_amplitude_symbols, amplitude_expr, \"numpy\"\n", - ")" + "np_amplitude_expr = sp.lambdify(sorted_amplitude_symbols, amplitude_expr, \"numpy\")" ] }, { @@ -1622,9 +1614,7 @@ ], "source": [ "%time\n", - "top_expression, sub_expressions = split_expression(\n", - " expression, max_complexity=100\n", - ")" + "top_expression, sub_expressions = split_expression(expression, max_complexity=100)" ] }, { @@ -1716,8 +1706,7 @@ " 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)\n", - " for symbol in top_symbols\n", + " sp.lambdify(args, definitions[symbol], modules) for symbol in top_symbols\n", " ]\n", "\n", " def recombined_function(*args):\n", @@ -1872,9 +1861,7 @@ "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(\n", - " result, collapse_graphs=True, render_final_state_id=False\n", - ")\n", + "dot = qrules.io.asdot(result, collapse_graphs=True, render_final_state_id=False)\n", "graphviz.Source(dot)" ] }, diff --git a/docs/report/003.ipynb b/docs/report/003.ipynb index a58d30aa..4b5fd4c0 100644 --- a/docs/report/003.ipynb +++ b/docs/report/003.ipynb @@ -71,23 +71,7 @@ }, "outputs": [], "source": [ - "%pip install -q ampform==0.13.3 black==22.3.0 matplotlib==3.5.1 mpl-interactions==0.20.2 numpy==1.22.3 qrules==0.9.7 sympy==1.10.1" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "jupyter": { - "source_hidden": true - }, - "tags": [ - "remove-cell" - ] - }, - "outputs": [], - "source": [ - "%pip install ndim==0.1.7 quadpy==0.16.15 scipy==1.8.0" + "%pip install -q ampform==0.13.3 black==22.3.0 matplotlib==3.5.1 mpl-interactions==0.20.2 ndim==0.1.7 numpy==1.22.3 qrules==0.9.7 quadpy==0.16.15 scipy==1.8.0 sympy==1.10.1 x21==0.2.16" ] }, { @@ -234,9 +218,7 @@ " sp.log((m1**2 + m2**2 - s + 2 * sp.sqrt(s) * q) / (2 * m1 * m2)),\n", " evaluate=False,\n", " )\n", - " right_term = (\n", - " (m1**2 - m2**2) * (1 / s - 1 / (m1 + m2) ** 2) * sp.log(m1 / m2)\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", @@ -597,9 +579,9 @@ "q0 = sp.Symbol(\"q0\", real=True)\n", "L = sp.Symbol(\"L\", integer=True, positive=True)\n", "s_thr = (m1 + m2) ** 2\n", - "integrand = (\n", - " PhaseSpaceFactor(s_prime, m1, m2) * na2(s_prime, m1, m2, L, q0)\n", - ") / ((s_prime - s_thr) * (s_prime - s - epsilon * sp.I))\n", + "integrand = (PhaseSpaceFactor(s_prime, m1, m2) * na2(s_prime, m1, m2, L, q0)) / (\n", + " (s_prime - s_thr) * (s_prime - s - epsilon * sp.I)\n", + ")\n", "integrand" ] }, @@ -856,9 +838,7 @@ " def _numpycode(self, printer, *args):\n", " integration_vars, limits = _unpack_integral_limits(self)\n", " if len(limits) != 1:\n", - " raise ValueError(\n", - " f\"Cannot handle {len(limits)}-dimensional integrals\"\n", - " )\n", + " raise ValueError(f\"Cannot handle {len(limits)}-dimensional integrals\")\n", " integrate = \"quadpy_quad\"\n", " printer.module_imports[\"quadpy\"].update({f\"quad as {integrate}\"})\n", " limit_str = \"%s, %s\" % tuple(map(printer._print, limits[0]))\n", diff --git a/docs/report/004.ipynb b/docs/report/004.ipynb index cb2f5d5c..1fb494a7 100644 --- a/docs/report/004.ipynb +++ b/docs/report/004.ipynb @@ -107,10 +107,7 @@ "import matplotlib.pyplot as plt\n", "import numpy as np\n", "import sympy as sp\n", - "from ampform.dynamics import (\n", - " PhaseSpaceFactor,\n", - " relativistic_breit_wigner_with_ff,\n", - ")\n", + "from ampform.dynamics import PhaseSpaceFactor, relativistic_breit_wigner_with_ff\n", "from IPython.display import Math\n", "from ipywidgets import widgets\n", "from matplotlib import cm\n", @@ -458,9 +455,7 @@ " tight_layout=True,\n", ")\n", "ax1, ax2 = axes\n", - "fig.suptitle(\n", - " \"$S$-wave Breit-Wigner ($L=0$) plotted over the complex $s$-plane\"\n", - ")\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", diff --git a/docs/report/005.ipynb b/docs/report/005.ipynb index 6ebe0629..c91dffa4 100644 --- a/docs/report/005.ipynb +++ b/docs/report/005.ipynb @@ -809,9 +809,7 @@ " 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 = {\n", - " symbol: arg for arg, symbol in sliders._arg_to_symbol.items()\n", - " }\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", @@ -984,9 +982,7 @@ " buttons.append(sliders_copy.pop(dummy_name))\n", " h_box = ipywidgets.HBox(buttons)\n", " h_boxes.append(h_box)\n", - " remaining_sliders = sorted(\n", - " sliders_copy.values(), key=lambda s: s.description\n", - " )\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", diff --git a/docs/report/009.ipynb b/docs/report/009.ipynb index 970b4986..482fb3ed 100644 --- a/docs/report/009.ipynb +++ b/docs/report/009.ipynb @@ -301,9 +301,7 @@ "outputs": [], "source": [ "n_channels = 2\n", - "n_resonances, i, R, L = sp.symbols(\n", - " \"n_R, i, R, L\", integer=True, negative=False\n", - ")\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", @@ -447,8 +445,7 @@ ") -> sp.Expr:\n", " def residue_function(i):\n", " return gamma[R, i] * sp.sqrt(\n", - " M[R]\n", - " * CoupledWidth(m**2, M, Gamma, m_a, m_b, angular_momentum, R, i)\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", @@ -705,9 +702,7 @@ " 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 = {\n", - " symbol: arg for arg, symbol in sliders._arg_to_symbol.items()\n", - " }\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", @@ -953,12 +948,8 @@ " for R in range(n_resonances):\n", " buttons = [sliders_copy.pop(f\"m{R}\")]\n", " if n_channels == 1:\n", - " buttons.append(\n", - " sliders_copy.pop(symbol_to_arg[Rf\"\\Gamma_{{{R},0}}\"])\n", - " )\n", - " buttons.append(\n", - " sliders_copy.pop(symbol_to_arg[Rf\"\\gamma_{{{R},0}}\"])\n", - " )\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", diff --git a/docs/report/010.ipynb b/docs/report/010.ipynb index 64059d1b..605292bd 100644 --- a/docs/report/010.ipynb +++ b/docs/report/010.ipynb @@ -302,9 +302,9 @@ ") -> 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 = (\n", - " beta[R] * gamma[R, i] * M[R] * Gamma[R, i] * sp.sqrt(ff2)\n", - " ) / (M[R] ** 2 - m**2)\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", @@ -323,8 +323,7 @@ ") -> sp.Expr:\n", " def residue_function(i):\n", " return gamma[R, i] * sp.sqrt(\n", - " M[R]\n", - " * CoupledWidth(m**2, M, Gamma, m_a, m_b, angular_momentum, R, i)\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", @@ -474,9 +473,7 @@ " 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 = {\n", - " symbol: arg for arg, symbol in sliders._arg_to_symbol.items()\n", - " }\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", @@ -539,9 +536,7 @@ " for ax in axes:\n", " ax.set_xlim(x_min, x_max)\n", " if not title:\n", - " title = (\n", - " f\"{n_channels}-channel $F$-vector with {n_resonances} resonances\"\n", - " )\n", + " title = f\"{n_channels}-channel $F$-vector with {n_resonances} resonances\"\n", " fig.suptitle(title)\n", "\n", " # 2D plot\n", @@ -638,10 +633,7 @@ " )\n", " m_diff = m_a - m_b\n", " x_offset = (x_max - x_min) * 0.015\n", - " if (\n", - " m_diff > x_offset + 0.01\n", - " and s_thr - abs(m_diff) > x_offset\n", - " ):\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", @@ -694,12 +686,8 @@ " for R in range(n_resonances):\n", " buttons = [sliders_copy.pop(f\"m{R}\")]\n", " if n_channels == 1:\n", - " buttons.append(\n", - " sliders_copy.pop(symbol_to_arg[Rf\"\\Gamma_{{{R},0}}\"])\n", - " )\n", - " buttons.append(\n", - " sliders_copy.pop(symbol_to_arg[Rf\"\\gamma_{{{R},0}}\"])\n", - " )\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", diff --git a/docs/report/011.ipynb b/docs/report/011.ipynb index 6e10a748..b017a1a5 100644 --- a/docs/report/011.ipynb +++ b/docs/report/011.ipynb @@ -105,10 +105,7 @@ "from qrules.topology import Topology, create_isobar_topologies\n", "from sympy.printing.numpy import NumPyPrinter\n", "from sympy.printing.printer import Printer\n", - "from sympy.tensor.array.expressions.array_expressions import (\n", - " ArraySlice,\n", - " ArraySymbol,\n", - ")" + "from sympy.tensor.array.expressions.array_expressions import ArraySlice, ArraySymbol" ] }, { @@ -708,9 +705,7 @@ "outputs": [], "source": [ "def _print_RotationY(self: NumPyPrinter, expr: RotationY) -> str:\n", - " self.module_imports[\"numpy\"].update(\n", - " {\"array\", \"cos\", \"ones\", \"zeros\", \"sin\"}\n", - " )\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", @@ -743,9 +738,7 @@ "outputs": [], "source": [ "def _print_RotationZ(self: NumPyPrinter, expr: RotationZ) -> str:\n", - " self.module_imports[\"numpy\"].update(\n", - " {\"array\", \"cos\", \"ones\", \"zeros\", \"sin\"}\n", - " )\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", @@ -1663,9 +1656,7 @@ " return create_expression(cls, momentum, **hints)\n", "\n", " def evaluate(self) -> ArraySlice:\n", - " three_momentum = ArraySlice(\n", - " self.momentum, (slice(None), slice(1, None))\n", - " )\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", @@ -1958,26 +1949,18 @@ " events: dict[int, ArraySymbol], node_id: int\n", " ) -> dict[str, sp.Expr]:\n", " helicity_angles: dict[str, sp.Expr] = {}\n", - " child_state_ids = sorted(\n", - " topology.get_edge_ids_outgoing_from_node(node_id)\n", - " )\n", - " if all(\n", - " topology.edges[i].ending_node_id is None for i in child_state_ids\n", - " ):\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(\n", - " topology, state_id\n", - " )\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(\n", - " topology, state_id\n", - " )\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", @@ -2017,9 +2000,7 @@ " 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(\n", - " events, initial_state_edge.ending_node_id\n", - " )" + " return __recursive_helicity_angles(events, initial_state_edge.ending_node_id)" ] }, { diff --git a/docs/report/013.ipynb b/docs/report/013.ipynb index 08fd8c1d..85afd8ad 100644 --- a/docs/report/013.ipynb +++ b/docs/report/013.ipynb @@ -132,10 +132,7 @@ }, "outputs": [], "source": [ - "from tensorwaves.data import (\n", - " TFPhaseSpaceGenerator,\n", - " TFUniformRealNumberGenerator,\n", - ")\n", + "from tensorwaves.data import TFPhaseSpaceGenerator, TFUniformRealNumberGenerator\n", "\n", "phsp_generator = TFPhaseSpaceGenerator(\n", " initial_state_mass=PDG[\"Lambda(c)+\"].mass,\n", @@ -486,9 +483,7 @@ " 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(\n", - " model, \"K^*\", phsp, full_expression\n", - " )\n", + " intensities_k = compute_sub_intensities(model, \"K^*\", phsp, full_expression)\n", " intensities_delta = compute_sub_intensities(\n", " model, \"Delta^*\", phsp, full_expression\n", " )\n", diff --git a/docs/report/014.ipynb b/docs/report/014.ipynb index 8cec10bc..0c88f0c8 100644 --- a/docs/report/014.ipynb +++ b/docs/report/014.ipynb @@ -688,9 +688,7 @@ " 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)\n", - " - initial_state_ids\n", - " - final_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", @@ -908,9 +906,7 @@ }, "outputs": [], "source": [ - "ampform_expr = (\n", - " ampform.get_builder(reaction_two_resonances).formulate().expression\n", - ")\n", + "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()" @@ -1316,8 +1312,7 @@ " for state_id in transition.final_states\n", " ]\n", " resonances = [\n", - " Str(s.particle.latex)\n", - " for s in transition.intermediate_states.values()\n", + " Str(s.particle.latex) for s in transition.intermediate_states.values()\n", " ]\n", " indices = [*final_state_helicities, *resonances]\n", " coefficient = C[indices]\n", @@ -1522,9 +1517,7 @@ " 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(\n", - " transition\n", - " )\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", @@ -1537,8 +1530,7 @@ " dynamics, _ = dynamics_builder(parent_particle, variables)\n", " expr *= dynamics\n", " resonances = [\n", - " Str(s.particle.latex)\n", - " for s in transition.intermediate_states.values()\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", diff --git a/docs/report/015.ipynb b/docs/report/015.ipynb index a5cdff81..a8e61fa0 100644 --- a/docs/report/015.ipynb +++ b/docs/report/015.ipynb @@ -555,9 +555,7 @@ " s = sp.Rational(state.particle.spin)\n", " m = sp.Rational(state.spin_projection)\n", " display(\n", - " Math(\n", - " Rf\"|s_{i},m_{i}\\rangle=|{s},{m}\\rangle \\quad ({particle_name})\"\n", - " )\n", + " Math(Rf\"|s_{i},m_{i}\\rangle=|{s},{m}\\rangle \\quad ({particle_name})\")\n", " )\n", " if functor is formulate_rotation_chain:\n", " args = (transition, i)\n", @@ -663,9 +661,7 @@ ], "source": [ "transition_r = full_reaction.transitions[-1]\n", - "show_all_spin_matrices(\n", - " transition_r, formulate_helicity_rotation_chain, cleanup=True\n", - ")" + "show_all_spin_matrices(transition_r, formulate_helicity_rotation_chain, cleanup=True)" ] }, { @@ -1165,10 +1161,7 @@ } ], "source": [ - "from ampform.kinematics import (\n", - " compute_boost_chain,\n", - " create_four_momentum_symbols,\n", - ")\n", + "from ampform.kinematics import compute_boost_chain, create_four_momentum_symbols\n", "\n", "momenta = create_four_momentum_symbols(topology)\n", "for state_id in topology.outgoing_edge_ids:\n", diff --git a/docs/report/016.ipynb b/docs/report/016.ipynb index 94e22ffe..601d680c 100644 --- a/docs/report/016.ipynb +++ b/docs/report/016.ipynb @@ -79,7 +79,7 @@ } ], "source": [ - "%pip install -q ndim==0.1.7 quadpy==0.16.15 scipy==1.8.0 sympy==1.10.1" + "%pip install -q ndim==0.1.7 quadpy==0.16.15 scipy==1.8.0 sympy==1.10.1 x21==0.2.16" ] }, { @@ -292,7 +292,7 @@ "Input \u001b[0;32mIn [7]\u001b[0m, in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 5\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m s_prime \u001b[38;5;241m*\u001b[39m (s_prime \u001b[38;5;241m+\u001b[39m s \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39mj)\n\u001b[1;32m 8\u001b[0m s_array \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39mlinspace(\u001b[38;5;241m-\u001b[39m\u001b[38;5;241m1\u001b[39m, \u001b[38;5;241m1\u001b[39m, num\u001b[38;5;241m=\u001b[39m\u001b[38;5;241m10\u001b[39m)\n\u001b[0;32m----> 9\u001b[0m \u001b[43mquadpy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mquad\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 10\u001b[0m \u001b[43m \u001b[49m\u001b[43mpartial\u001b[49m\u001b[43m(\u001b[49m\u001b[43mparametrized_func\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43ms\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43ms_array\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m 11\u001b[0m \u001b[43m \u001b[49m\u001b[43ma\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0.0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 12\u001b[0m \u001b[43m \u001b[49m\u001b[43mb\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2.0\u001b[39;49m\u001b[43m,\u001b[49m\n\u001b[1;32m 13\u001b[0m \u001b[43m)\u001b[49m\n", "File \u001b[0;32m:55\u001b[0m, in \u001b[0;36mquad\u001b[0;34m(f, a, b, args, epsabs, epsrel, limit)\u001b[0m\n", "File \u001b[0;32m:45\u001b[0m, in \u001b[0;36mintegrate_adaptive\u001b[0;34m(f, intervals, eps_abs, eps_rel, criteria_connection, kronrod_degree, minimum_interval_length, max_num_subintervals, dot, domain_shape, range_shape)\u001b[0m\n", - "File \u001b[0;32m:128\u001b[0m, in \u001b[0;36m_gauss_kronrod_integrate\u001b[0;34m(k, f, intervals, dot, domain_shape, range_shape)\u001b[0m\n", + "File \u001b[0;32m:129\u001b[0m, in \u001b[0;36m_gauss_kronrod_integrate\u001b[0;34m(k, f, intervals, dot, domain_shape, range_shape)\u001b[0m\n", "File \u001b[0;32m:49\u001b[0m, in \u001b[0;36mg\u001b[0;34m(x)\u001b[0m\n", "Input \u001b[0;32mIn [7]\u001b[0m, in \u001b[0;36mparametrized_func\u001b[0;34m(s_prime, s)\u001b[0m\n\u001b[1;32m 4\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mparametrized_func\u001b[39m(s_prime, s):\n\u001b[0;32m----> 5\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m s_prime \u001b[38;5;241m*\u001b[39m (\u001b[43ms_prime\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m+\u001b[39;49m\u001b[43m \u001b[49m\u001b[43ms\u001b[49m \u001b[38;5;241m+\u001b[39m \u001b[38;5;241m1\u001b[39mj)\n", "\u001b[0;31mValueError\u001b[0m: operands could not be broadcast together with shapes (21,) (10,) " @@ -402,9 +402,7 @@ " def _numpycode(self, printer, *args):\n", " integration_vars, limits = _unpack_integral_limits(self)\n", " if len(limits) != 1:\n", - " raise ValueError(\n", - " f\"Cannot handle {len(limits)}-dimensional integrals\"\n", - " )\n", + " raise ValueError(f\"Cannot handle {len(limits)}-dimensional integrals\")\n", " integrate = \"quadpy_quad\"\n", " printer.module_imports[\"quadpy\"].update({f\"quad as {integrate}\"})\n", " limit_str = \"%s, %s\" % tuple(map(printer._print, limits[0]))\n", diff --git a/docs/report/017.ipynb b/docs/report/017.ipynb index 468d5e9a..b39c530e 100644 --- a/docs/report/017.ipynb +++ b/docs/report/017.ipynb @@ -196,9 +196,7 @@ "@make_commutative\n", "@implement_doit_method\n", "class Kibble(UnevaluatedExpression):\n", - " def __new__(\n", - " cls, sigma1, sigma2, sigma3, m0, m1, m2, m3, **hints\n", - " ) -> Kibble:\n", + " def __new__(cls, sigma1, sigma2, sigma3, m0, m1, m2, m3, **hints) -> Kibble:\n", " return create_expression(\n", " cls, sigma1, sigma2, sigma3, m0, m1, m2, m3, **hints\n", " )\n", diff --git a/docs/reports.md b/docs/reports.md index c66a9850..477052e3 100644 --- a/docs/reports.md +++ b/docs/reports.md @@ -1,9 +1,9 @@ # 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. +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. diff --git a/pyproject.toml b/pyproject.toml index 5efe584c..6ef311ac 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,6 @@ exclude = ''' )/ ''' include = '\.pyi?$' -line-length = 79 preview = true target-version = [ "py37", @@ -32,10 +31,12 @@ target-version = [ ] [tool.isort] -line_length = 79 profile = "black" [tool.nbqa.addopts] +black = [ + "--line-length=85", +] flake8 = [ "--extend-ignore=E402,F404,F821", ] diff --git a/tox.ini b/tox.ini index ed66fff5..553fc04a 100644 --- a/tox.ini +++ b/tox.ini @@ -107,7 +107,7 @@ allowlist_externals = passenv = EXECUTE_NB commands = - pytest --nbmake --nbmake-timeout=900 {posargs:docs} + pytest --nbmake --nbmake-timeout=1200 {posargs:docs} [testenv:sty] description =