pax_global_header00006660000000000000000000000064150354612710014516gustar00rootroot0000000000000052 comment=ffd5a6905c9e9f9dc1845326adbdc78f7ef67aae vector-1.6.3/000077500000000000000000000000001503546127100130275ustar00rootroot00000000000000vector-1.6.3/.all-contributorsrc000066400000000000000000000051521503546127100166630ustar00rootroot00000000000000{ "projectName": "vector", "projectOwner": "scikit-hep", "repoType": "github", "repoHost": "https://github.com", "files": [ "README.md" ], "imageSize": 100, "commit": false, "commitConvention": "angular", "contributors": [ { "login": "jpivarski", "name": "Jim Pivarski", "avatar_url": "https://avatars.githubusercontent.com/u/1852447?v=4", "profile": "https://github.com/jpivarski", "contributions": [ "maintenance", "code", "doc" ] }, { "login": "henryiii", "name": "Henry Schreiner", "avatar_url": "https://avatars.githubusercontent.com/u/4616906?v=4", "profile": "https://github.com/henryiii", "contributions": [ "maintenance", "code", "doc" ] }, { "login": "eduardo-rodrigues", "name": "Eduardo Rodrigues", "avatar_url": "https://avatars.githubusercontent.com/u/5013581?v=4", "profile": "https://github.com/eduardo-rodrigues", "contributions": [ "maintenance", "code", "doc" ] }, { "login": "LovelyBuggies", "name": "N!no", "avatar_url": "https://avatars.githubusercontent.com/u/29083689?v=4", "profile": "http://lovelybuggies.com.cn/", "contributions": [ "doc" ] }, { "login": "pfackeldey", "name": "Peter Fackeldey", "avatar_url": "https://avatars.githubusercontent.com/u/18463582?v=4", "profile": "https://github.com/pfackeldey", "contributions": [ "doc" ] }, { "login": "kreczko", "name": "Luke Kreczko", "avatar_url": "https://avatars.githubusercontent.com/u/1213276?v=4", "profile": "https://github.com/kreczko", "contributions": [ "code" ] }, { "login": "nsmith-", "name": "Nicholas Smith", "avatar_url": "https://avatars.githubusercontent.com/u/6587412?v=4", "profile": "https://github.com/nsmith-", "contributions": [ "ideas" ] }, { "login": "mayou36", "name": "Jonas Eschle", "avatar_url": "https://avatars.githubusercontent.com/u/17454848?v=4", "profile": "https://github.com/mayou36", "contributions": [ "ideas" ] }, { "login": "Saransh-cpp", "name": "Saransh Chopra", "avatar_url": "https://avatars.githubusercontent.com/u/74055102?v=4", "profile": "https://saransh-cpp.github.io/", "contributions": [ "maintenance", "code", "doc" ] } ], "contributorsPerLine": 7, "commitType": "docs" } vector-1.6.3/.git_archival.txt000066400000000000000000000002011503546127100162730ustar00rootroot00000000000000node: ffd5a6905c9e9f9dc1845326adbdc78f7ef67aae node-date: 2025-07-15T15:16:25+01:00 describe-name: v1.6.3 ref-names: tag: v1.6.3 vector-1.6.3/.gitattributes000066400000000000000000000000371503546127100157220ustar00rootroot00000000000000.git_archival.txt export-subst vector-1.6.3/.github/000077500000000000000000000000001503546127100143675ustar00rootroot00000000000000vector-1.6.3/.github/ISSUE_TEMPLATE/000077500000000000000000000000001503546127100165525ustar00rootroot00000000000000vector-1.6.3/.github/ISSUE_TEMPLATE/bug_report.yml000066400000000000000000000026551503546127100214550ustar00rootroot00000000000000name: Bug Report description: Create a bug report labels: ["bug (unverified)"] body: - type: markdown attributes: value: | Thank you for helping us by filing out this bug report! Please verify that this bug has not been reported before through the [issue tracker](https://github.com/scikit-hep/vector/issues). - type: input id: vector-version attributes: label: Vector Version description: What version of vector are you using? placeholder: python -m pip show vector validations: required: true - type: input id: python-version attributes: label: Python Version description: What version of Python are you running? placeholder: $ python --version validations: required: true - type: textarea id: os attributes: label: OS / Environment description: What Operating System or environment are you using? placeholder: Describe your environment. validations: required: true - type: textarea id: describe attributes: label: Describe the bug description: | - A clear and concise description of the bug. - Describe how to reproduce this bug using a [Minimum Workable Example](https://stackoverflow.com/help/minimal-reproducible-example). validations: required: true - type: textarea id: log-output attributes: label: Any additional but relevant log output vector-1.6.3/.github/ISSUE_TEMPLATE/config.yml000066400000000000000000000006431503546127100205450ustar00rootroot00000000000000blank_issues_enabled: true contact_links: - name: GitHub Discussions url: https://github.com/scikit-hep/vector/discussions about: | How do I do 'X' with vector and the Scikit-HEP ecosystem? - name: Gitter url: https://gitter.im/Scikit-HEP/vector?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge about: | How do I do 'X' with vector and the Scikit-HEP ecosystem? vector-1.6.3/.github/ISSUE_TEMPLATE/feature_request.yml000066400000000000000000000023121503546127100224760ustar00rootroot00000000000000name: Feature Request description: Suggest an idea for vector labels: ["feature"] body: - type: markdown attributes: value: | Thank you for helping us by suggesting a feature! Please verify that a similar feature has not been suggested before through the [issue tracker](https://github.com/scikit-hep/vector/issues). - type: textarea id: description attributes: label: Describe the potential feature description: A detailed description of the suggested feature. validations: required: true - type: textarea id: motivation attributes: label: Motivation description: Why is this feature needed? How will it help the HEP community? placeholder: | X would help users to perform Y with a better accuracy / speed. X is a novel feature required to perform W. ... - type: textarea id: implementation attributes: label: Possible Implementation description: If possible, add a possible implementation for the suggested feature. placeholder: | X can be implemented by referring to the Y paper or the Z blog post. I have implemented X locally using the following code - ... vector-1.6.3/.github/PULL_REQUEST_TEMPLATE.md000066400000000000000000000016221503546127100201710ustar00rootroot00000000000000## Description **Kindly take a look at [CONTRIBUTING.md](https://github.com/scikit-hep/vector/blob/main/.github/CONTRIBUTING.md).** _Please describe the purpose of this pull request. Reference and link to any relevant issues or pull requests._ ## Checklist - [ ] Have you followed the guidelines in our Contributing document? - [ ] Have you checked to ensure there aren't any other open Pull Requests for the required change? - [ ] Does your submission pass pre-commit? (`$ pre-commit run --all-files` or `$ nox -s lint`) - [ ] Does your submission pass tests? (`$ pytest` or `$ nox -s tests`) - [ ] Does the documentation build with your changes? (`$ cd docs; make clean; make html` or `$ nox -s docs`) - [ ] Does your submission pass the doctests? (`$ pytest --doctest-plus src/vector/` or `$ nox -s doctests`) ## Before Merging - [ ] Summarize the commit messages into a brief review of the Pull request. vector-1.6.3/.github/codecov.yml000066400000000000000000000000511503546127100165300ustar00rootroot00000000000000codecov: notify: after_n_builds: 6 vector-1.6.3/.github/dependabot.yml000066400000000000000000000003401503546127100172140ustar00rootroot00000000000000version: 2 updates: # Maintain dependencies for GitHub Actions - package-ecosystem: "github-actions" directory: "/" schedule: interval: "weekly" groups: actions: patterns: - "*" vector-1.6.3/.github/matchers/000077500000000000000000000000001503546127100161755ustar00rootroot00000000000000vector-1.6.3/.github/matchers/pylint.json000066400000000000000000000012341503546127100204070ustar00rootroot00000000000000{ "problemMatcher": [ { "severity": "warning", "pattern": [ { "regexp": "^([^:]+):(\\d+):(\\d+): ([A-DF-Z]\\d+): \\033\\[[\\d;]+m([^\\033]+).*$", "file": 1, "line": 2, "column": 3, "code": 4, "message": 5 } ], "owner": "pylint-warning" }, { "severity": "error", "pattern": [ { "regexp": "^([^:]+):(\\d+):(\\d+): (E\\d+): \\033\\[[\\d;]+m([^\\033]+).*$", "file": 1, "line": 2, "column": 3, "code": 4, "message": 5 } ], "owner": "pylint-error" } ] } vector-1.6.3/.github/workflows/000077500000000000000000000000001503546127100164245ustar00rootroot00000000000000vector-1.6.3/.github/workflows/cd.yml000066400000000000000000000026161503546127100175420ustar00rootroot00000000000000name: CD on: workflow_dispatch: release: types: - published jobs: dist: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 with: fetch-depth: 0 - uses: hynek/build-and-inspect-python-package@v2 publish: needs: [dist] runs-on: ubuntu-latest if: github.event_name == 'release' && github.event.action == 'published' environment: name: pypi url: https://pypi.org/p/vector permissions: id-token: write attestations: write steps: - uses: actions/download-artifact@v4 with: name: Packages path: dist - name: List distributions to be deployed run: ls -l dist/ - name: Generate artifact attestation for sdist and wheel uses: actions/attest-build-provenance@e8998f949152b193b063cb0ec769d69d929409be # v2.4.0 with: subject-path: "dist/vector-*" - name: Verify sdist artifact attestation env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: gh attestation verify dist/vector-*.tar.gz --repo ${{ github.repository }} - name: Verify wheel artifact attestation env: GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} run: gh attestation verify dist/vector-*.whl --repo ${{ github.repository }} - uses: pypa/gh-action-pypi-publish@release/v1 with: attestations: true vector-1.6.3/.github/workflows/ci.yml000066400000000000000000000061731503546127100175510ustar00rootroot00000000000000name: CI on: pull_request: push: branches: - main - develop workflow_dispatch: concurrency: # Skip intermediate builds: always. # Cancel intermediate builds: only if it is a pull request build. group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }} env: # The "FORCE_COLOR" variable, when set to 1, # tells Nox to colorize itself. FORCE_COLOR: "1" jobs: pre-commit: runs-on: ubuntu-latest name: Check SDist steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: python-version: 3.x - uses: astral-sh/setup-uv@v6 - name: PyLint run: uvx nox -s pylint -- --output-format=github check-lite: runs-on: ubuntu-latest strategy: fail-fast: false matrix: python-version: - "3.8" - "3.9" - "3.10" - "3.11" - "3.12" - "3.13" name: Python ${{ matrix.python-version }} - Lite steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} allow-prereleases: true - uses: astral-sh/setup-uv@v6 - name: Test lite package run: uvx nox -s lite-${{ matrix.python-version }} --verbose check-full: needs: [check-lite] runs-on: ubuntu-latest strategy: fail-fast: false matrix: python-version: - "3.8" - "3.9" - "3.10" - "3.11" - "3.12" - "3.13" name: Python ${{ matrix.python-version }} - Full steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} allow-prereleases: true - uses: astral-sh/setup-uv@v6 - name: Install nox run: uv tool install nox - name: Run doctests on Python 3.11 if: matrix.python-version == 3.11 run: nox -s doctests-${{ matrix.python-version }} --verbose - name: Test package and generate coverage report run: nox -s coverage-${{ matrix.python-version }} --verbose - name: Upload coverage report uses: codecov/codecov-action@v5.4.3 with: token: ${{ secrets.CODECOV_TOKEN }} discheck: runs-on: ubuntu-latest name: Disassemble check steps: - uses: actions/checkout@v4 - uses: actions/setup-python@v5 with: python-version: 3.8.13 - uses: astral-sh/setup-uv@v6 - name: Test compute features run: uvx nox -s disassemble --verbose pass: needs: [pre-commit, check-lite, check-full, discheck] runs-on: ubuntu-latest steps: - run: echo "All jobs passed" # root: # runs-on: ubuntu-latest # steps: # - uses: actions/checkout@v4 # - name: Get Conda # uses: conda-incubator/setup-miniconda@v2 # with: # environment-file: environment.yml # activate-environment: vector # - name: Run tests # shell: "bash -l {0}" # run: python -m pytest tests/root_tests -ra vector-1.6.3/.github/workflows/notebooks.yml000066400000000000000000000012171503546127100211530ustar00rootroot00000000000000name: Notebooks on: pull_request: branches: - main - develop workflow_dispatch: concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true env: # The "FORCE_COLOR" variable, when set to 1, # tells Nox to colorize itself. FORCE_COLOR: "1" jobs: test: runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 - name: Set up Python uses: actions/setup-python@v5 with: python-version: "3.10" - name: List installed Python packages run: python -m pip list - name: Test example notebooks run: pipx run nox -s notebooks vector-1.6.3/.gitignore000066400000000000000000000037171503546127100150270ustar00rootroot00000000000000# Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] *$py.class # C extensions *.so # Distribution / packaging .Python build/ develop-eggs/ dist/ downloads/ eggs/ .eggs/ lib/ lib64/ parts/ sdist/ var/ wheels/ pip-wheel-metadata/ share/python-wheels/ *.egg-info/ .installed.cfg *.egg MANIFEST # PyInstaller # Usually these files are written by a python script from a template # before PyInstaller builds the exe, so as to inject date/other infos into it. *.manifest *.spec # Installer logs pip-log.txt pip-delete-this-directory.txt # Unit test / coverage reports htmlcov/ .tox/ .nox/ .coverage .coverage.* .cache nosetests.xml coverage.xml *.cover *.py,cover .hypothesis/ .pytest_cache/ # Translations *.mo *.pot # Django stuff: *.log local_settings.py db.sqlite3 db.sqlite3-journal # Flask stuff: instance/ .webassets-cache # Scrapy stuff: .scrapy # Sphinx documentation docs/_build/ # PyBuilder target/ # Jupyter Notebook .ipynb_checkpoints # IPython profile_default/ ipython_config.py # pyenv .python-version # pipenv # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. # However, in case of collaboration, if having platform-specific dependencies or dependencies # having no cross-platform support, pipenv may install dependencies that don't work, or not # install all needed dependencies. #Pipfile.lock # PEP 582; used by e.g. github.com/David-OConnor/pyflow __pypackages__/ # Celery stuff celerybeat-schedule celerybeat.pid # SageMath parsed files *.sage.py # Environments .env .venv env/ venv/ ENV/ env.bak/ venv.bak/ # Spyder project settings .spyderproject .spyproject # Rope project settings .ropeproject # mkdocs documentation /site # mypy .mypy_cache/ .dmypy.json dmypy.json # Pyre type checker .pyre/ # Other editor files .vscode /.idea/* # Version file /src/vector/_version.py .pyre_configuration .watchmanconfig # All contributors stuff /node_modules /package.json /yarn.lock # MacOS .DS_Store vector-1.6.3/.pre-commit-config.yaml000066400000000000000000000033541503546127100173150ustar00rootroot00000000000000ci: autoupdate_commit_msg: "chore: update pre-commit hooks" autofix_commit_msg: "style: pre-commit fixes" repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v5.0.0 hooks: - id: check-added-large-files - id: check-case-conflict - id: check-merge-conflict - id: check-symlinks - id: check-yaml - id: debug-statements - id: end-of-file-fixer exclude: ^docs - id: mixed-line-ending - id: requirements-txt-fixer - id: trailing-whitespace - repo: https://github.com/astral-sh/ruff-pre-commit rev: "v0.12.3" hooks: - id: ruff-check args: ["--fix", "--show-fixes"] - id: ruff-format - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.16.1 hooks: - id: mypy files: src args: [] additional_dependencies: - numpy - packaging - sympy - repo: https://github.com/codespell-project/codespell rev: v2.4.1 hooks: - id: codespell args: ["-LHEP"] exclude: ^docs/usage/intro.ipynb$ - repo: https://github.com/rbubley/mirrors-prettier rev: "v3.6.2" hooks: - id: prettier types_or: [yaml, markdown, html, css, scss, javascript, json] exclude: assets/js/webapp\.js - repo: https://github.com/adamchainz/blacken-docs rev: 1.19.1 hooks: - id: blacken-docs args: ["-E"] additional_dependencies: [black~=24.0] - repo: https://github.com/pre-commit/pygrep-hooks rev: v1.10.0 hooks: - id: python-check-blanket-type-ignore exclude: ^src/vector/backends/_numba_object.py$ - id: rst-backticks - id: rst-directive-colons - id: rst-inline-touching-normal vector-1.6.3/.readthedocs.yml000066400000000000000000000006561503546127100161240ustar00rootroot00000000000000# .readthedocs.yml # Read the Docs configuration file # See https://docs.readthedocs.io/en/stable/config-file/v2.html for details # Required version: 2 build: os: ubuntu-22.04 tools: python: "3" # Build documentation in the docs/ directory with Sphinx sphinx: configuration: docs/conf.py # Include PDF and ePub formats: all python: install: - method: pip path: . extra_requirements: - docs vector-1.6.3/CITATION.cff000077500000000000000000000032701503546127100147260ustar00rootroot00000000000000cff-version: "1.2.0" authors: - family-names: Chopra given-names: Saransh orcid: "https://orcid.org/0000-0003-3046-7675" - family-names: Schreiner given-names: Henry orcid: "https://orcid.org/0000-0002-7833-783X" - family-names: Rodrigues given-names: Eduardo orcid: "https://orcid.org/0000-0003-2846-7625" - family-names: Eschle given-names: Jonas orcid: "https://orcid.org/0000-0002-7312-3699" - family-names: Pivarski given-names: Jim orcid: "https://orcid.org/0000-0002-6649-343X" contact: - family-names: Pivarski given-names: Jim orcid: "https://orcid.org/0000-0002-6649-343X" doi: 10.5281/zenodo.15263860 message: If you use this software, please cite our article in the Journal of Open Source Software. preferred-citation: authors: - family-names: Chopra given-names: Saransh orcid: "https://orcid.org/0000-0003-3046-7675" - family-names: Schreiner given-names: Henry orcid: "https://orcid.org/0000-0002-7833-783X" - family-names: Rodrigues given-names: Eduardo orcid: "https://orcid.org/0000-0003-2846-7625" - family-names: Eschle given-names: Jonas orcid: "https://orcid.org/0000-0002-7312-3699" - family-names: Pivarski given-names: Jim orcid: "https://orcid.org/0000-0002-6649-343X" date-published: 2025-05-23 doi: 10.21105/joss.07791 issn: 2475-9066 issue: 109 journal: Journal of Open Source Software publisher: name: Open Journals start: 7791 title: "Vector: JIT-compilable mathematical manipulations of ragged Lorentz vectors" type: article url: "https://joss.theoj.org/papers/10.21105/joss.07791" volume: 10 title: "Vector: JIT-compilable mathematical manipulations of ragged Lorentz vectors" vector-1.6.3/CONTRIBUTING.md000066400000000000000000000174551503546127100152740ustar00rootroot00000000000000# Contributing to vector If you are planning to develop `vector`, or want to use the latest commit of `vector` on your local machine, you might want to install it from the source. This installation is not recommended for users who want to use the stable version of `vector`. The steps below describe the installation process of `vector`'s latest commit. This page also describes how to test `vector`'s codebase and build `vector`'s documentation. **Note**: [Scikit-HEP's developer information](https://scikit-hep.org/developer) is a general and much more detailed collection of documentation available for developing `Scikit-HEP` packages. This developer guide is specific to `vector`. ## Installing vector We recommend using a virtual environment to install `vector`. This would isolate the library from your global `Python` environment, which would be beneficial for reproducing bugs, and the overall development of `vector`. The first step would be to clone `vector` - ```bash # fork scikit-hep/vector git clone https://github.com//vector.git ``` and then we can change the current working directory and enter `vector` - ```bash cd vector ``` ### Creating a virtual environment A virtual environment can be set up and activated using `venv` in both `UNIX` and `Windows` systems (or use your favorite environment tool). **UNIX**: ```bash python3 -m venv .env . .env/bin/activate ``` **Windows**: ``` python -m venv .env .env\bin\activate ``` ### Installation The developer installation of `vector` comes with several options - - `awkward`: installs [awkward](https://github.com/scikit-hep/awkward) for the `awkward` backend - `numba`: installs [numba](https://github.com/numba/numba) for JIT compilation - `sympy`: installs [sympy](https://github.com/sympy/sympy) for the `sympy` backend - `test`: the test dependencies - `test-extras`: extra dependencies to run disassemble and `dask_awkward` tests - `docs`: extra dependencies to build and develop `vector`'s documentation - `dev`: installs the `awkward`, `test`, `numba`, and `sympy` options These options can be used with `pip` with the editable (`-e`) mode of installation in the following way - ```bash pip install -e ".[dev, test]" ``` For example, if you want to install the `docs` dependencies along with the dependencies included above, use - ```bash pip install -e ".[dev, test, docs]" ``` Furthermore, `vector` can also be installed using `conda`. This installation also requires using a virtual environment - ```bash conda env create conda activate vector conda config --env --add channels conda-forge # optional ``` ### Adding vector for notebooks `vector` can be added to the notebooks using the following commands - ``` python -m ipykernel install --user --name vector ``` ## Activating pre-commit `vector` uses a set of `pre-commit` hooks and the `pre-commit` bot to format, type-check, and prettify the codebase. The hooks can be installed locally using - ```bash pre-commit install ``` This would run the checks every time a commit is created locally. The checks will only run on the files modified by that commit, but the checks can be triggered for all the files using - ```bash pre-commit run --all-files ``` If you would like to skip the failing checks and push the code for further discussion, use the `--no-verify` option with `git commit`. ## Testing vector `vector` is tested using `pytest` and `pytest-doctestplus`. `pytest` is responsible for testing the code, whose configuration is available in [pyproject.toml](https://github.com/scikit-hep/vector/blob/main/pyproject.toml).`pytest-doctestplus` is responsible for testing the examples available in every docstring, which prevents them from going stale. Additionally, `vector` also uses `pytest-cov` to calculate the coverage of these unit tests. ### Running tests locally The tests can be executed using the `test` and `test-extras` dependencies of `vector` in the following way - ```bash python -m pytest ``` To skip the notebook tests, use `--ignore=tests/test_notebooks.py` ### Running tests with coverage locally The coverage value can be obtained while running the tests using `pytest-cov` in the following way - ```bash python -m pytest --cov=vector tests/ ``` ### Running doctests The doctests can be executed using the `test` dependencies of `vector` in the following way - ```bash python -m pytest --doctest-plus src/vector/ ``` or, one can run the doctests along with the unit tests in the following way - ```bash python -m pytest --doctest-plus . ``` A much more detailed guide on testing with `pytest` for `Scikit-HEP` packages is available [here](https://scikit-hep.org/developer/pytest). ### Running notebook tests The Notebook tests can be executed individually in the following way - ```bash pytest tests/test_notebooks.py ``` ## Documenting vector `vector`'s documentation is mainly written in the form of [docstrings](https://peps.python.org/pep-0257) and [reStructurredText](https://docutils.sourceforge.io/docs/user/rst/quickref.html). The docstrings include the description, arguments, examples, return values, and attributes of a class or a function, and the `.rst` files enable us to render this documentation on `vector`'s documentation website. `vector` primarily uses [Sphinx](https://www.sphinx-doc.org/en/master/) for rendering documentation on its website. The configuration file (`conf.py`) for `sphinx` can be found [here](https://github.com/scikit-hep/vector/blob/main/docs/conf.py). The documentation is deployed on [https://readthedocs.io]() [here](https://vector.readthedocs.io/en/latest/). Ideally, with the addition of every new feature to `vector`, documentation should be added using comments, docstrings, and `.rst` files. ### Building documentation locally The documentation is located in the `docs` folder of the main repository. This documentation can be generated using the `docs` dependencies of `vector` in the following way - ```bash cd docs/ make clean make html ``` The commands executed above will clean any existing documentation build and create a new build under the `docs/_build` folder. You can view this build in any browser by opening the `index.html` file. ## Nox `vector` supports running various critical commands using [nox](https://github.com/wntrblm/nox) to make them less intimidating for new developers. All of these commands (or sessions in the language of `nox`) - `lint`, `tests`, `notebooks`, `doctests`, `docs`, and `build` - are defined in [noxfile.py](https://github.com/scikit-hep/vector/blob/main/noxfile.py). `nox` can be installed via `pip` using - ```bash pip install nox ``` The default sessions (`lint`, `lite`, `tests`, `doctests`, and `disassemble`) can be executed using - ```bash nox ``` ### Running pre-commit with nox The `pre-commit` hooks and `pylint` can be run with `nox` in the following way - ``` nox -s lint # run pre-commit on the default Python version nox -s pylint # run pylint on the default Python version ``` ### Building vector with nox `vector` can be built with `nox` in the following way - ``` nox -s build ``` ### Running tests with nox Tests can be run with `nox` in the following way - ``` nox -s lite # test only with the required dependencies on all available Python versions (use lite-3.11 for a specific Python version) nox -s tests # test with all extra dependencies on all available Python versions nox -s tests # run doctests on all available Python versions nox -s coverage # run the tests session and generate a coverage report on all available Python versions nox -s notebooks # test notebooks on the default Python version nox -s disassemble # check compute functions on Python 3.8 ``` ### Building documentation with nox Docs can be built with `nox` in the following way - ``` nox -s docs ``` Use the following command if you want to deploy the docs on `localhost` - ``` nox -s docs -- serve ``` vector-1.6.3/LICENSE000066400000000000000000000030551503546127100140370ustar00rootroot00000000000000BSD 3-Clause License Copyright (c) 2019-2024, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. All rights reserved. Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: * Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. * Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. * Neither the name of the vector package developers nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission. THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. vector-1.6.3/README.md000066400000000000000000000325151503546127100143140ustar00rootroot00000000000000Vector logo # Vector: arrays of 2D, 3D, and Lorentz vectors [![DOI][zenodo-badge]][zenodo-link] [![DOI][joss-badge]][joss-link] [![Scikit-HEP][sk-badge]][sk-link] [![Actions Status][actions-badge]][actions-link] [![Documentation Status][rtd-badge]][rtd-link] [![pre-commit.ci status][pre-commit-badge]][pre-commit-link] [![codecov percentage][codecov-badge]][codecov-link] [![PyPI platforms][pypi-platforms]][pypi-link] [![PyPI version][pypi-version]][pypi-link] [![Conda latest release][conda-version]][conda-link] [![LICENSE][license-badge]][license-link] [![GitHub Discussion][github-discussions-badge]][github-discussions-link] [![Gitter][gitter-badge]][gitter-link] ## Installation You can install Vector with [pip](https://pypi.org/project/vector/) and [conda](https://anaconda.org/conda-forge/vector). ```bash pip install vector ``` ## Introduction Vector is a Python library for 2D and 3D spatial vectors, as well as 4D space-time vectors. It is especially intended for performing geometric calculations on _arrays of vectors_, rather than one vector at a time in a Python for loop. Vector is part of the [Scikit-HEP project](https://scikit-hep.org/), High Energy Physics (HEP) tools in Python. ### Coordinate systems Vectors may be expressed in any of these coordinate systems: - the azimuthal plane may be Cartesian `x` `y` or polar `rho` ($\rho$) `phi` ($\phi$) - the longitudinal axis may be Cartesian `z`, polar `theta` ($\theta$), or pseudorapidity `eta` ($\eta$) - the temporal component for space-time vectors may be Cartesian `t` or proper time `tau` ($\tau$) in any combination. (That is, 4D vectors have 2×3×2 = 12 distinct coordinate systems.) Diagram of coordinate systems ### Backends Vectors may be included in any of these data types: - [vector.obj](https://vector.readthedocs.io/en/latest/src/make_object.html) objects (pure Python) - [NumPy structured arrays](https://numpy.org/doc/stable/user/basics.rec.html) of vectors - [Awkward Arrays](https://awkward-array.org/) of vectors (possibly within variable-length lists or nested record structures) - [SymPy expressions](https://www.sympy.org/en/index.html) for symbolic (non-numeric) manipulations - In [Numba-compiled functions](https://numba.pydata.org/), with [vector.obj](https://vector.readthedocs.io/en/latest/src/make_object.html) objects or Awkward Arrays Each of these "backends" provides the same suite of properties and methods, through a common "compute" library. ### Geometric versus momentum Finally, vectors come in two flavors: - geometric: only one name for each property or method - momentum: same property or method can be accessed with several synonyms, such as `pt` ($p_T$, transverse momentum) for the azimuthal magnitude `rho` ($\rho$) and `energy` and `mass` for the Cartesian time `t` and proper time `tau` ($\tau$). ### Familiar conventions Names and coordinate conventions were chosen to align with [ROOT](https://root.cern/)'s [TLorentzVector](https://root.cern.ch/doc/master/classTLorentzVector.html) and [Math::LorentzVector](https://root.cern.ch/doc/master/classROOT_1_1Math_1_1LorentzVector.html), as well as [scikit-hep/math](https://github.com/scikit-hep/scikit-hep/tree/master/skhep/math), [uproot-methods TLorentzVector](https://github.com/scikit-hep/uproot3-methods/blob/master/uproot3_methods/classes/TLorentzVector.py), [henryiii/hepvector](https://github.com/henryiii/hepvector), and [coffea.nanoevents.methods.vector](https://coffea-hep.readthedocs.io/en/latest/modules/coffea.nanoevents.methods.vector.html). ## Getting help - Source code on GitHub: [scikit-hep/vector](https://github.com/scikit-hep/vector) - Report bugs and request features on the [GitHub Issues page](https://github.com/scikit-hep/vector/issues) - Ask questions on the [GitHub Discussions page](https://github.com/scikit-hep/vector/discussions) - Real-time chat on Gitter: [Scikit-HEP/Vector](https://gitter.im/Scikit-HEP/vector) ## Contributing to Vector If you want to contribute to Vector, [pull requests](https://github.com/scikit-hep/vector/pulls) are welcome! Please install the latest version of the `main` branch from source or a fork: ```bash git clone https://github.com/scikit-hep/vector.git cd vector pip install -e . ``` Refer to [CONTRIBUTING.md](https://github.com/scikit-hep/vector/blob/main/.github/CONTRIBUTING.md) for more. ## Citing Vector To cite Vector, please use [![DOI][joss-badge]][joss-link] ```bib @article{Chopra2025, doi = {10.21105/joss.07791}, url = {https://doi.org/10.21105/joss.07791}, year = {2025}, publisher = {The Open Journal}, volume = {10}, number = {109}, pages = {7791}, author = {Saransh Chopra and Henry Schreiner and Eduardo Rodrigues and Jonas Eschle and Jim Pivarski}, title = {Vector: JIT-compilable mathematical manipulations of ragged Lorentz vectors}, journal = {Journal of Open Source Software} } ``` ## Documentation ### Tutorials - [Vector objects](https://vector.readthedocs.io/en/latest/src/object.html) - [NumPy arrays of vectors](https://vector.readthedocs.io/en/latest/src/numpy.html) - [Awkward Arrays of vectors](https://vector.readthedocs.io/en/latest/src/awkward.html) - [Compiling functions on vectors with Numba](https://vector.readthedocs.io/en/latest/src/numba.html) - [Vector expressions with SymPy](https://vector.readthedocs.io/en/latest/src/sympy.html) ### Vector constructors - [Making vector objects](https://vector.readthedocs.io/en/latest/src/make_object.html) - [Making NumPy arrays of vectors](https://vector.readthedocs.io/en/latest/src/make_numpy.html) - [Making Awkward Arrays of vectors](https://vector.readthedocs.io/en/latest/src/make_awkward.html) - [Making SymPy vector expressions](https://vector.readthedocs.io/en/latest/src/make_sympy.html) ### Vector functions - [Interface for all vectors](https://vector.readthedocs.io/en/latest/src/common.html) - [Interface for 2D vectors](https://vector.readthedocs.io/en/latest/src/vector2d.html) - [Interface for 3D vectors](https://vector.readthedocs.io/en/latest/src/vector3d.html) - [Interface for 4D vectors](https://vector.readthedocs.io/en/latest/src/vector4d.html) - [Interface for 2D momentum](https://vector.readthedocs.io/en/latest/src/momentum2d.html) - [Interface for 3D momentum](https://vector.readthedocs.io/en/latest/src/momentum3d.html) - [Interface for 4D momentum](https://vector.readthedocs.io/en/latest/src/momentum4d.html) ### More ways to learn - [Papers and talks](https://vector.readthedocs.io/en/latest/src/talks.html) ## Contributors Thanks goes to these wonderful people ([emoji key](https://allcontributors.org/docs/en/emoji-key)):
Jim Pivarski
Jim Pivarski
🚧 💻 📖		 Henry Schreiner
Henry Schreiner
🚧 💻 📖		 Eduardo Rodrigues
Eduardo Rodrigues
🚧 💻 📖		 N!no
N!no
📖		 Peter Fackeldey
Peter Fackeldey
📖		 Luke Kreczko
Luke Kreczko
💻		 Nicholas Smith
Nicholas Smith
🤔
Jonas Eschle
Jonas Eschle
🤔		 Saransh Chopra
Saransh Chopra
🚧 💻 📖
This project follows the [all-contributors](https://github.com/all-contributors/all-contributors) specification. Contributions of any kind welcome! See [CONTRIBUTING.md](./.github/CONTRIBUTING.md) for information on setting up a development environment. ## Acknowledgements This library was primarily developed by Saransh Chopra, Henry Schreiner, Jim Pivarski, Eduardo Rodrigues, and Jonas Eschle. Support for this work was provided by the National Science Foundation cooperative agreement [OAC-1836650](https://www.nsf.gov/awardsearch/showAward?AWD_ID=1836650) and [PHY-2323298](https://www.nsf.gov/awardsearch/showAward?AWD_ID=2323298) (IRIS-HEP) and [OAC-1450377](https://www.nsf.gov/awardsearch/showAward?AWD_ID=1450377) (DIANA/HEP). Any opinions, findings, conclusions or recommendations expressed in this material are those of the authors and do not necessarily reflect the views of the National Science Foundation. [actions-badge]: https://github.com/scikit-hep/vector/actions/workflows/ci.yml/badge.svg [actions-link]: https://github.com/scikit-hep/vector/actions [codecov-badge]: https://codecov.io/gh/scikit-hep/vector/branch/main/graph/badge.svg?token=YBv60ueORQ [codecov-link]: https://codecov.io/gh/scikit-hep/vector [conda-version]: https://img.shields.io/conda/vn/conda-forge/vector.svg [conda-link]: https://github.com/conda-forge/vector-feedstock [github-discussions-badge]: https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github [github-discussions-link]: https://github.com/scikit-hep/vector/discussions [gitter-badge]: https://badges.gitter.im/Scikit-HEP/vector.svg [gitter-link]: https://gitter.im/Scikit-HEP/vector?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge [joss-badge]: https://joss.theoj.org/papers/10.21105/joss.07791/status.svg [joss-link]: https://doi.org/10.21105/joss.07791 [license-badge]: https://img.shields.io/badge/License-BSD_3--Clause-blue.svg [license-link]: https://opensource.org/licenses/BSD-3-Clause [pre-commit-badge]: https://results.pre-commit.ci/badge/github/scikit-hep/vector/main.svg [pre-commit-link]: https://results.pre-commit.ci/repo/github/scikit-hep/vector [pypi-link]: https://pypi.org/project/vector/ [pypi-platforms]: https://img.shields.io/pypi/pyversions/vector [pypi-version]: https://badge.fury.io/py/vector.svg [rtd-badge]: https://readthedocs.org/projects/vector/badge/?version=latest [rtd-link]: https://vector.readthedocs.io/en/latest/?badge=latest [sk-badge]: https://scikit-hep.org/assets/images/Scikit--HEP-Project-blue.svg [sk-link]: https://scikit-hep.org/ [zenodo-badge]: https://zenodo.org/badge/DOI/10.5281/zenodo.15263860.svg [zenodo-link]: https://zenodo.org/records/15263860 vector-1.6.3/docs/000077500000000000000000000000001503546127100137575ustar00rootroot00000000000000vector-1.6.3/docs/Makefile000066400000000000000000000011721503546127100154200ustar00rootroot00000000000000# Minimal makefile for Sphinx documentation # # You can set these variables from the command line, and also # from the environment for the first two. SPHINXOPTS ?= SPHINXBUILD ?= sphinx-build SOURCEDIR = . BUILDDIR = _build # Put it first so that "make" without argument is like "make help". help: @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) .PHONY: help Makefile # Catch-all target: route all unknown targets to Sphinx using the new # "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). %: Makefile @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) vector-1.6.3/docs/_images/000077500000000000000000000000001503546127100153635ustar00rootroot00000000000000vector-1.6.3/docs/_images/LogoSrc.svg000066400000000000000000000145221503546127100174600ustar00rootroot00000000000000 image/svg+xml vector-1.6.3/docs/_images/coordinate-systems.png000066400000000000000000002771001503546127100217340ustar00rootroot00000000000000PNG  IHDRf Hq pHYs +tEXtSoftwarewww.inkscape.org< IDATxy|ս76$-=웈 Jm+mRj[XZmkomګZTKE+* XV5$$cldL|IVqjH3nF[LDDDDD–%J{%f""""",Xv,@SEDDDD$ٗ%5᜔|JDDDDD$Y=v<tq]"""""&l 7`۱twЈ431}2hiLDDDDDBM Lr;.vQ4 Oq. lqjA#f"""""bpcB23 f6^н2E """""*3Kn~ĻN0TY?""""" 3n4AaRJDDDDD$ u97("""""AQ}xwqMjE jLDDDDDEKȺ~dmI(1.bf5)My("""""zYr9P8FDDDDDSx?mPSIhLDDDDD:nePќ"`@S@#f"""""f6 _2u p%elVDDDDDNWcᄓr`8{k3iVT ] ĹQXz 4b&"""""'0< |Rݍ(CR#fo7݈[K@#f"""""̒UR QsRяz1 f6ߺs~c]''FDDDDDQs<nD=,qjO3n3nD8Ɠ}293oT w#xw%)3P8f<`9%qj"%f"""""!2/\ Hs8PSEDDDDB =РJ=I(149l+p`f~%KqluH+<E@`пC63=/I;.`;8^D Y"0_vf#ՠtơ34uIؙiP *ȡ; 'Vq,ȱH33SXWδ_iC93igyU)@18_.8$"ݍxVGt*`8+;z"%f""&f \ |Hr7ی/I{X8^0bfqdLW`q:DJDDUu%c_ݍKm^{,wxD$Y:~j ՠ-guVA*%f"",3 \\IG%ۺa73oQݬJ|6t H3?| x 5G׀%v9cf@ _mz4"q(4&Ӂa$W3OLDD0lpA >$qֹOafF&F+0 $tR7*xnkج.B1i&`bgoˢLD$™Y.#::7KrxB SM;8jP>YXSK45̙Z /ݍHَ*1Pu OoK9xB..N%|~ʀulq7$i9D|nϜs83SF @DO|4Ϊ.,W%c܍SVˁ7m|-=zt5qwɕD3 _qG!X 8Nt3\\ s9`9/I{ߨ'JԂ3FT9Q:3\U'Wb&"M-0p$k]R_eω. }xq.mY_{ q7"*tcf6  ]OwݐZgfi.WN#xY _.WS$v@H7dfD_YoYg{Txq.,_юU;%y]]GH7Rw7f|=܍F>|]Xlq:KKNW#`{W%jK9rlc5:v 7t"">Z :3Kn~B @EEȑ#ٰa?Ǐ_~Ӈ|bcc;%fW;`w3 13%e!-333g۴i^7o޼ `ܹJDDBntLD$\` (r8r?ѱ#?f%""]Ít!3;^Ƹgرcȑ#m:|5k]tE ""Ҳ5㜸zP(1fgf/g߉9_ڦs臈HXV*!"҉, -x^kK~ߟ;w%1H?qѱFDD:0_㛗]u*X[4WCIHHVJDD:b*-~ Hv9$7tSc'[Ģ|蔸DDmwc%f""tB,Ñ.4~xN?c'S\r 999tQqչ362Sl9JݎG'SotLE?DDBke""'̆ | }~F rrr(((kHSEC ;LD$48nu3VY\ϵ^p" >l@Ro}KIHfJDDafIfc` m e7x#QQ:+rbя.MDD:Ε̢ |#dR]IB95U^z)EDٹ3*-~ )""aL$DD3 X$0diGNN۷o2v8չFD$Y=EIhΝpa~p JDDB_˛Sv1%f"\> ΝKll`-M3gNB+us]H1|^#p$ egg3s̀c/K,iT_2} fx"">nvLD"M2Pao1eee|_oVH""1nv"Y/̀H3f 6lhe˖F{HHJr̭BD-33ۀwQR&Ѱ믿^IHtsHdfӁQn"ׯ%%%c׮]deeC8Znu,",2IKMM媫j/\IHxq57Rb&"݂EׁkF7n!ȑHEDžMLu;<O>TWW 6͛783+"FqdFD$lY=|2q3<臒2j0df_ݎE"1rH6o?ݻt12iusHX1u=L\+$e_Wd7;qs0+WB}w]"NfJD$Y6=fH_z+ԩS:UED”3ԍ=t;zVkmt шH'Qb&"r"3u;lO<طo6l`ٲeY@s9.ȥ(ED*3*%0zh6nb3x E%""]`8qsUeaf3{x%e&X`2ΕHH0)n"r2Yf ለHǹZYkDUf|s9F~_r&&&Ŕ)SHOOw12di3tcq\SEu>2 Zc&"̢߁4",xM&"""HtZm4b&"AafU()ps%f"z=`۱JD˘Y= f"%lv,""""'[+1NgfSpHձ2H2()sNCQb&"n^܎E$ت;CDD:FM܏>W$B;,\rCs.DC,X|XDܴxbJJJXbۡHiLD3܎EmK,`ѢE.G"""nfl2LDݎEm{aoLL:={r9Jнٓ={p_^{oya=y,\M^2܎E$,Y3`l۶!CtGEEE=zv(AWYYѣG_CNJPUU*A***Z|Oa*XN5b&"mbfSQR&w(YF m۶qyqysϹ7Hh3smv,"Ѻgyo٥D'裏0i$yyyxgMDs239)fvHt;P{5*fHLԯHg4h_at~u&\I4QDZef@IH#MM[b.D#ҴǏwڹ)++9;){,))ݯMFJDEfv%,XDBQsUUQTTTį~+&NH\\iii0rHn6<_0k,nw_aÆ@RR̙3;vTL?<ӧO'!!Crr2{.?xט5kfO?m?|?^XXȬYя~?s?oof֬Yy-]YY?#}-[oVb٤BLL &L#խ/~䐐@=eСC'{7v""YHg1x`Ón(//Cjx_ng}֧OϷ#G6{={'|칪ꫯn1[n~aWX< ,?/رs?222Ηky-}˭s6?i_wu6owm>6̙3x|͓eee~m|a̒}NZc&"M2GXDBŋeO}vl°aÂDm۶qRTT_ܹsɡy1o<^/3gdٲe-d̙|g\}\veՋ?wy#Gp5'x2ߓO>?eڴiİzj~pi=ՋG}={__̜9 /0]BBMxo~@=OBBV7 O>$7n>.]榛nbҤIs}q1.\?l̈f\xᅌ;LJJJؾ}; .䥗^\~]N?9֙mr;pfv5}MD4?n(Lk#f_?~{m? 6̪1,66^yFmkjdm 8ۨMqqM:5w-#fNMjMg8p<j֭kɓcGG:u=zQUVGFڽ{}g-Gy&t33_ѳ3 `f3' kftB$jjjx7ZlufLy7_9ѣx\zlٲ^z?/nt<&&&`}[oը#7a[lik@޽BSFgbϜ9sm7vXMi{6|7[nmw?gu6mp50CKD|Gw;x^***`Рϫٮ^xz\P"Ɍl2. D$?}gvڴiDGGSSS%''3tfOKK#99F%KJJ)Bttt1vi- E߿cǒb۳:+VtJǏoni+‡~Ν;˯Y[[ˡC"olK+1l,aR~h|*++ KKK4^/w񐞞NJJ;AKt/^LtFRb߾}4obb"%%%)=~ާOVs2mBMii)%%%' yѣ;ilC=ښ!rݡ3gf׀ƓCHMM 6hjkksf! IDAThxqq (((`РATWWxx<-NiʁONE8iƶھVKYUUU:&KΓ94|2U2d[?!1ccǎ%==t?ϟD:`wL$Y&v;>|_f8''KOOgڴiLi8a¶o>S?LJJJSn$'}c6ö#8*t#(5556ѣomͦ[N lߖY|4A?~_e,[,[*2,q-mbef)`۱RS$cbbHKK #9l0Ə3&&Ɵgf͚ S(E5j&pM͛[l{aƮZ3NFF֭kɴLqqq;vvMU4O޽>kC]gѾg&e/)1@f v,'+..t233INN>ĐNNNN1c0h ѵzjVX˕E.]ڦרlԩS_-m3uqM6iS{z^j9_^^jtbwyO=T ڵŶ}e߾}[ZXQQe˂Q j%f"~n3zqz:&''K^Ǒ V_wy?8`QuͿH|nk~H[4/.WzZW]]ͽW^e15,뭷6{C?y@Oi=z4ַf͚&TUU̚5 }_O?m5`h8͵Qz0M:3cf?v&))\Ǝ럖5={$..?V[[KII TTTb V^X _K,ik***Z݌Zu]޽{={vj{̝;kp0eʔ.iƌ-+ |ﾛC5js3P3 f68ڣjkk HRmjzz:^7륶鐻wСCx<| mI7u.hDs=X 6h"5k999;w 0h bccٶmUUU59spmӜ9s{(,,dϿ----`30a| yyy\~dffϿ/ ۸t&$$OpPYY5\3oK%f"ҹ̬ OjFj*[ӯ_?GYYY"2G~D>1+++c߾}x<RSSֵ^}&$$$ ??72jԨ.Scǎ5J2ի-bŊ,\uqQ<#Gdƌ\x-Qv뭷rW}GyyyI^߾}yٿ?}"++SN9װpFvv]\&Nȣ> IM׮JHH੧{e8pbC aݺu<,]}˜1c3g&Mo^^^K/eӦM<䓼лwo&O̜9s2d[nkرΑ3l|7ؔSO=߶fbCzwҷo_?| RSSILL mjiӦۂuYw `f6ˡtHmm-{!**PgfQZZ꿘:vBFFcƌ| 0!C}jjjx_WUUvZV\իm ;}QFYYY曌5є03fѢEي+kHn:nᆀϩzdܹBw`(Ғm25f"ݔ 4^6o_S?IJJ")))`1h Ǩ(={4XPPg}F||<hZi;òeZ]Oqg.2*++y7˻:\QVVc=c=ˣGfIhz!nlFGJD!3 Ui{\8111ZEEE@єkR^^aJuÛoyҋ=\qeQZZʢEIDIKKcl߾=JbCƍ8\PEA[cL1(`0X:[}BV[[XQ,x<&x)--S33^/555TUUIVVVf<={lC裏IىٚЗ%^z%.233|pmƎ;XjСCőԩS ?NZZ+CW:Rb& ˜1c03Cx%77YUU񔖖y^(++#**ʟ#VXum+mLeqƑp|[omƄ HMMзo_=z4wqj?wwO=T̛7{ x&__[[_O/ԯ_?ѣӦM㮻b׮]'ݑ5f":3o%7|pvr=z$++ Kuu][GB>30`@O9r|ƎKLL )))r)-X+s7ndƍ .d߾}|+_aڵnٓ3}t***={6q@f͚Eii)̛7޽{}'e'O[naуÇ{nVZ6WBJjM:30ef_vTQQAee%555$%%i~ѣǏCuuu"֭[!---$ ih*H~OZ9orSNa֬Y|'??'xڞ~7vn?o7pwp9}՟/77˗l뒛ɓկ~WJOhH2LQ'pÑjXqСy6^w׽^/ǎc߿#G}v j'`=^z1nܸm;N4pV3r-= 6[o/`޽{5jJҁbw:vѣ;F^;qN?t6l5kZl'g7 @޽7o^*[2` x8|pHKʻ%f"a̮vnHNN&33ha!it_~Q\\hZǷmFUUQQQL6(jjj8v)))i*Hm۶{nVtm"c߾},\UVg=|)++#11dddsq9PVVٳ;gy>}M>?{2rHoq3uG%ģLD2\ݎ-yU—8p|x^hi 7.**o)vEEhHg̬7?jkk~}wӒKJJMgc? ۉ~rJV\Iaa!w}7w}7ɜuY\ve\ve&!/8՝(1 uzV$5EGGm Gܹѣ8㐑#(//'77M3,++}k呞Nzz?{Ȑ!m)Xg}ZJJJXd K,{r w}>GWmM:N?"MeeeR[[KJJ n$!.66}6:Bvvve7rVXr%|( 6K^^lDv>gŋ~GjXa5h[nO?QF9Ά)((`kرJ^I/.(23 q۷sNUeٳ'yyyL24V YYY555SVVo{a>^k"Xv-Gl۶z[ih6)ؾ}I3??o|<DEEQRRٳ~;w\q<#lݺYSTTi}IP)1-&@jj*z*b7%;Xe֭U3cͤ(H"66UK!%%CEݸqcS~_S^z;+W gΜ /K~m~mq9vݒ&3HWW.qqqۗ>}2KLLL@6bM\Пo_RRغuk+޽{)** HDB]YYWfǎc999{߇;M[{hŋ{4ik 7o7֭[kQQ^z)K.mvǧ~_ ȑ#ԇ:ok 4b&71"⎘] fggT|kX@aa!v| 9r؀"5kPYYI~~>qqqVOg,_72tPf͚9r/fÆ 0m45y{rbfx<ϟߨly駙1cEEE̞={ {fڴi9Cqq1k׮eɒ%TVVpwG%;"ynJٰaNq D233 ..KzzxsbgřgItt4zE:S\\ڵ:a̘1UVqQ}рǍ /&__SSW_Maa!>(Æ k%\~>V^͏~FMVVMKMM;wIWBN_VOH2>=n= h*PK^իWУG&kg㉎իWE߾}:t(R&]ܿFjDEE릲Yr% ,`ɒ%瓘H袋⋉cΜ9y@Çs7s7E]bo803ǡ? 9))^xhYJJ `ƍ,[͛7~>Lll,Og̙JP3jS (CϞ=13 ] ]رc'%%!ChtÇar֍"&&ٳg3{fL0 &4:^_d}Z;רQ:4D*308B8L:@%tqqq 4Mn]SS?8<O;"mqq֯_e:\C$VGJDBt;Pv"[k***(--*bVTTxLJJ $$$0n8bbbs;4PL vk.㕘IHTr 222'$$ lOjjjS\\LJJ G?gf*Y;x ~)L08K$UL2tn.WUFi())=<κLg&CҦiREXRq~ ˺/**\-"+_QĕEliiӦI4iɜ23!I9䙙ܯ"<ݘLܟf###˥ESJr\.|/fwU@vRz˚"ڽiZT3!@Aĺ B)))dJ6M_1y|l6iiik`C)eppm۶a6e0KiFX3!"L)~F/&''ikkC)Eff!nR]v199䤞ȹ\.B0w8'O7òv=bɚ_?9qN,23=LZI(b||@JEԐR7NجV+;wdbbbΪbCO>Qt:IMMsZlݺ***$)by|:$fBD֭@A(& >''))iѡP(QJƎ;7˹/TTTiYYY޽{N!Ē$\.K_ab644đ#GHII8Wh͌v "7䒘 9Fo4M[ӍB,2s4mh&|&&&欼ͬy<9?_---x<v;LLLfrrRF2͚$1"RFBPTXd!Ql6X,x$Z~<7n$;;?Nzz>Ml$fBDMP^{ .kܽ{ !DUTTPQQA l6e3% 3efqhh!}ƚ7 cl"̔R Ǎ#^L&gMF"R,eXزe ~ہgn211!eBLiSF!+fBא5S#{̈́djjqxx ebbT?uꔾv%i@Ánj?6:LRJ}8Ν;AU)%%%LMMp8hhhл;&&&=[nebbI}!%3㴵gnRJ&4-`t a0CUD^r*++)++-&))$rrrٲe .kNյCl6S__OFFFYBD}Ob$fBϧ-F8N~?J)rssXz{{9v0ǶIbf3y<##2&&&HKKI؂&Di4:  8֋vbK.KyR(evcө'fᖞNzzdzf#III9yDӴb6"i $ k8+)e+~4M"r+$11g=kt YKBR82YďVJ)^wndL*rc9FGG(..jBiBNNEEEFL/n2txk!bLicFĪX!+c=IDԑRFTSSSPQQAff&&ZMNN&990XiAbR& ٕQEE~t:iiiE ?OB^8Y1benb*(( ##$eBe6֕Vá7(..ֻE #͚\$1bRFDZYVIDT lRzz>Lee%iiiF 0]z)DGӴCFR("zHX z-=ѱX,444DmRrAĴRɊˠ8ֻ!:;; BP\\ltHBqN4778t۷o.EEEz=BD!o5M; <FA&'')IVBRRRHLL$`X0F$&I2>iZA, DJJvҕQD)eg3>mɡ$#[P(D   G7M~nt%+fB,ݗ+rrrسgOmB_tuu}v}֭[kYN>MGG7nB,gX Y1b R6Lh%cǎzinnk!"o^4с 4_V)J(Lbb!uNJB233ILLc,%vc!g4M;bt+B!UF!599ɾ}X+ B/`BJJ ۶mc``M dOhA$fBR2Ef1Ftq}RJ۷o'!!͆f3:4!qFZLbtb.ٌn'333&f>VWBDx$}rrÁg2l@VKV̄XR*8|ܹ0DBx^=I|ttt055Ejj*555  @ii)8q>4McdggM  ))$NQQ`0H^^^$샃:u Z H94AbqX(HHHg۷˅lK./ gL#;# @^^0=j߾}NYYقǏqnSSSL&f3iiilݺO~~ѡ LӴF. T1:cp8L&\ݜ>}GCCT0󑜜!4qhfHiZ|>| v;㤤HBFJKK)**vӧOKyyN8?ci=us1}g-R 8ʙL&`b!\ecܦ8x q)++;Dgl6ϙ'5S HLLUqq1$e fYE}FӴ#FaIz0 CYb Riii!!!* pF_ىHgJJJJ"~?dggi ǎcpp;vȾFI4GFaIĺ pq+<MMMz̊Bvv6IIIG'`MBz83pPQQAEEAFKss3[ə$eIFYL&b-)RX`0abbgt8qkhhI@rr2EEEQ3%2^bSNvQJ&"%%ݎne)L+dL7F!VoxxVʨ28r0zsBUUޭN+W9y$>O ioob'6lvٶmۜV%b |ZӴ60$fbPJm4:ʝrݸn9Bzz:Q[[+%blzc144͛ݘ ӉR L&~i!iO"Hb&<$D2331ͤJpB!}خje۶mQ*S$mllC}v,Oss:~233uD,^ Dv#|>;@jfLsOBBvѣv222$rip٩vtt2tuu1<< y'1Fwo{JA"A6l@[[呝-aq$99y^Brq)--*&̕hF]]("4M;itDڊ!`}mX@ݼCFFV¶m0뮌srss QMӞ5:h#kJ5:~@cǎ! ;v^`zM(gw[O(--*[)+#]]]tttIMM5+$CEwD4ް"n)ҁ ʸz4@l1LTWWcJJJ餬Ψ^;@ @__~_oz"0inFb&ٷa6d2 2d"77R)..f``@:ЙKФ4zfrr[X,8x TV㙞[YY)b54 qI)>oCDd0bF  555>|478J q,9 BtttPZZJRRرC`xx1$1++{F$1qG)eb{ar"Aa#Ѣݎ)M fddoNbb$ea"]En4- ֘$f"} 2:ymmmLMMdZBL&222(//Nn`!Kϔ7Ngg'TWW{Ѣcbb8999G?jjjTY-+b:)6:h'+J6:6FFF|LLSSS9rŢl*//'//TZ VDahooféEjj* 477aILnވUiD,L T&H 1Bh.Pzq3<٣'"33)++q 6n܈bfqˌ2!"VHb&ɿF!֭[QJa6%)d2QZZ ӝ)G+//\gaoL&6oLNNΜ>rss zijjzihh$;wέb&Oh<3 Jn4:V!D ӧD4|ؘ^vI233IOO78jKf'e;vnmFrr|tHMM]׫b4 bl1O)UWJJш(i?6:X%j"&)r҂k۸q#7n4:5viiiA)ӧ$;;{>\.0]&]#g=2I)ɓ'잞'fyyyAٺut_\cMMMn.RY>etL~DQJ2I;NOgWGYY&i۶mf͛rSN֦^XXȶmo4 q4atLVD,BeM$&&RD8rJ(**:Eb9egB[bb" L&}5e||Ʉf[N8AwBuu5 IQ|$9OjA:IDLQJ}W&ɓ'񦽽ӧO+dx<&''SEz.et:9t}e}e5FF i"Hb&bRWcvIS040ș}p8|vttp8dƍY_)Ǐ;8qm۶-u3 niic !VMƌ"Hb&bR*x6:=V+渚a6::Jff&eee8Nl6|.8biF}}= !!aI-`Ç/] QXXht8"nݺ5)))EAAA6\866 /dllqɡ=\s5Kq0'b̻Y hnnt222b$Nr'Nc2$19"Hb&R[CHNZjx`6)))0Rݻ}sO<|ӟ_n199ɡCHII~NuDZZN$7hctDꈨ x8Dtx vm|;_{?XZ5S~~>$%%X,2!=z`0ls"֝DLD<\ht,"%&&RTTdtKp8zVUU>VXXHAArkX(e|gxWظ|K_-cxxx pal6uVٹsg 466sNmft(8!#$1QA)-1::(.p eeedffBCCü!+ߏ.K{y}\qEs)cGGǜ6[n]IOOuԻl>}M0dggwӉgƍE)y?Gb]jt$fJ:d;FrQRRL!ONN""qMMMp8ɑ2qV3gY1$f3ÔRz8 l6rww7kzXb5M&""B3Ro]t =y!15QK#pN>M{{;tpl䤞 Jӏk ֹx<?y,!!+"LQK4lH㙗\. x㍰S( z%?D)B4 M>b aZz%++ Irrضw^%B!L&_Wq\svWSXj|<y줹YѣGkY3N'rܱ, ߯Ǯv8j!;FIb&F)eXQPP@AAA؎ p:j͛inn2j 樉E.7~;RSS} ̿fg>0CSJxСC477+KKK sxxϣ>:XN+2lт Ǐoŵ4AgPJYkErݴ299Icc#iii$$$sNC`jjIvt2XopUWSO-ѣGyԜ_wuzrixؽ{7}iss3MMMttt_PPK/Dcc؞={Lf_ +IWue_f$1*^vXB)Evv:Ý>}Z/jmmeklFFF8r䈾D>ƏcRlݺo|#lܸTxxGx#++kH矧'N͉'=(_~9k5lvrVhX,+))hD6:N3*J/"c>W,))iىRJ \]]@͛7GURXJ)_&V+|+t,v_Wa-#imm߿GΝ;կ~E~~PVV,vsUW/FI2͔4-2ŒIb&VL) aK;vrrVl6u9M?IZZ>9wbK_n{٫^999o䡇"7we)_}U~a<ɓ'Wtx'yeϞ=L`UW]ų>UW]0כl6)O4AI2)rPDRRR:͆je||Ʉb]G>¯~+~_o222EEEreq7/g<3 >o}kIe 3[kgꫯ^MQ/p8W%S$1K < T+rz`2﬜N'JFFƊQ9X8# b6kkQTqqg}?Oaiƽ߿l޼<|u]O?5HCbL7pQ v;LDI~mOGGǢ9r @rr2Qri1:'kf*;;;bIٙG^u:::xW=b/+)dK.Yk~?7p/²^Kٹs';wdÆ F#UMӆBLLR*M) nb B311]9J)&pjvCj`$"^(hjjb||;fh<s=[ݻwj璚 /'>Űr3?sF;ɄnnJ}|wIb&N+댎E;raQ[[Kbb"l޼yMc J(//_V6bQ44 lxcP(ij>˖-[[333oˇ?ǰ knF}[2ƨ"{Ă)]|Y%Q.55.4'Iɓl߾Nbb";wQ1RĖ177M6JZZ!LMMӟx`ūv999$[lo^4);W RTT4ߵf.hů|c?&P@ @ hMaW@%)1s0SXXXHKKWfk&&&x78pO6:X+et:tuurnϓO>Imm-7|>#LW\q?0~c׿&//. \|<䓫%hbb`6G$fQFJJ)3oO',P?hFRR^x!:\.׊DǃbۍGā5Sy' MЇ>}Ǯ]ԧ>*?TKKI|MN:bxyiiiO|cf[vJB%+fR=$)1( ~_n")wtj!cnkVxGn[0)3L|f߾}d׮]LMM-چ8=b5;)Z ݻwldd{dJs!N>?_OFF'yVi 2lSJ-pXY*!l)<6s饗r9O,,ZK9qļע$mrrzJn6zzz=d2qu/KvءW^axxx<NMMqI|ϵTIIIwym}bn~iHoo/_qI!TRR%bRJe2صkRRRbpta2HIIa":8N=_{J)R|;⡇;s3 Ykk+<̂~9k׿uxnNHyy9wy93<æM<盓$&&ZdffӥI.ytM\s5s7sĊ;HIq|cPJ%w" ijj~?;w? ]팍Jcc#z⒛w---z%`[wRƱ1}Uuu5V^}|7ȽKMM9/t:Ϊ?<{R~/⋱X,z͛஻ӟ?7x .`U[W_}kvc{ 0l֝;i+#VGvcz5-Vq2w066R4Mq4i[-BR;v>6mڤ7)((yy衇%gJLLnzǞ)c祗^lVϕR+JHm6۶m?0so+\KSSSⳍX.ʖ-[v<Xc OKK˼$l4M{hzikkJRR[l!99Y6{G0f(3Y;FzzzDn #o[_:SRR7|3{/6> IDAT:lذa1z<zzz[{^:/1뮻ϟyN:Eyy!A@,,^N)U0P~r::tÇٿL%''cqM7E R1tfΝyrssIOO֦Z7:qPXTT'eo|<øcZ[>GQQъK/QYYɯ WtϞ=|&99.LkRݻ/}K+5y<0}({Ś8itba!ǁvSoo/oMMM9rÇsȑ\,]v!Rq6VSSS,k'F!%011Aff&Q[[KaaaX/7w<ςϙI_uyWVL&ӊ[ٳMسgϜ$l6O}sc?ILMGG0OmtbaTjcY+^>]p ۳!v;::::]$ 288b!!!ajX~?͸nnJVV&)l|_? f[n|غ~.ͥoiӦ9e3nV>=矟S(D,JIb'R9FDzVSO=DW\>HJJJϵ^ qQf3v"!!4o?{1!H<5b)Y2SozsQ?:uo~<裋6v>p=}%{v5MbǹKL̲kyꩧǾebfZٸq E)IbR*qcYkۿ_Wygxӟ+d(,,$33SE($8q7o^kWY(?6d҅xWMM nIlٲɓ|[b޽|ww37 顢b7nE;~s^{nݺx;7Ac SJ]< ,q)[hmmG?Y9^/zbL&yyyQ\\L^^RTT7XoB& MӨnSUUcNXV3T6Y,[-r HKKl6S__[FGG> ?]g뮻b꽪'NӭӲ}_袋hllw;ݻ~xe 67hQ$f1a o||__GW_갖B|IIIdee-7{%4n|sIIJJj}vl6۲ʙ d2]򒔔$Uc-J9tJ)}+]omm+_ ?OMrssg>w;999zb600@UUUXǛ~SO=Wո*yԩSD0*$f1F)Ucf\wu\wuSOwEDKvv6.>>hmm=s㡔K_qlmmeh`sZq q'Nill\q>̃>O~EW. NIIPogNtMs=>籱1^x>я9xO;n"H&pH\e߾}$H^.`lWVGsVNz~鰞cjji n~!5!箻ⳟl\DҭC+79o͗%^zE;e)))+/^9eeeh $1JrQłFGGygw{{1n%?3e``rsVV|״tn[CRCCCtvv ())xUNj&߿.#Z0޽{?~}EuzU"MԨ` &Id71Yu7bVwF7l, FذPI:Sqp~s{#!!AVTTɉJ]z7⧟~jtR@ 0hҙ &y<`kbؾ}h믿m۶ugB`dnAhuf_ܸŜaɒ%Xd ӱcl,[ xg:}(PUULǏQWW-LWK) | Ç.4iud `|EEEƍg؆/(ծ; :J3 B@|p&;;wƮ]PUUCEZZ7TJ)IpttT9Riu)J;vҥKsEbb^Ȁ M~ZZ6oaBUVYLBVff&R)зo_())A~Tsi]]ўkbݺujӧO`8qӦMS}a(..ʕ+qsX}Y1 wSz0L APYOL 86mڵkqAܹgϞeˁnݺŋ#11[ftR՝ĕ(ZJQ"קOQ=!=B߾}uPaڵ8}t ǪU`0Z2h8qۿ^٭[qFz ۷owߏ|*m6*:cʿ(֚&&f& +Ku-^}Ux![|wx1~vf4666񁏏O%+lݽo])YgJ%ӡT*!tnN6RlBTUUaС0)KKKG}3gδAE'dN!M4^{5,!!۶mrV}QF/rׯ۷ĎBĬ3,̄B8@'qFlذmi& 0 u;;;#FA pVJY[[Z(ikFK<mlB.}7H$^Q"T|ᇸzj2d?^3ZN*1h׸8ۣ PSST̘1 jky<Kf֭[qIv"=zVUxzzHHB=0M\B!p:J;<шFYY  uhinnÇ@Am=&ږR()R|( zm1dܾ}j7„`Æ 6'zg`ܹLH-s21srr̙3U߿_,))I%I&apݷ Ĝ9s~vLT⧟~ʕ+u>/#:JM8Ba zCy{{w: N544޽{JJ2dA>}>GWK)Nڊ*vǔRͭ1FMK̐L&CAABBBQF|[uֱi2f/ALL M:@LLBaW^yE%1;r}tEX`*Y{jK*Z1@M0hbfhbB=fs EښDH,mh1w۱r Ϗ&i&Ro߆H$D"ϭRѣGfddd{qaժU 44666pv6|6IjllѣGߥ\tImL@˞k è|}jV[Byyy9SĬBBu,+LRտ"<<ܤpssp&Uk)e{{µRJ///Zhb1b1wX (JXf M2sw}ǎ]N̞wvvaTg\]]Bvpr=!,MLM̺]OFZdggÆ ǃ'<<<8- x]*TK)khJS)8S(իW3q}Qb4suuExx8rss1x`!11=47n֯_ɓ'wcL޽ Z;v RmZD"HR1ooo|QI>>>P*].ɤLm6ghbM!l]OF` }JIbVV cƌQI %%E,))Iml؎'NLjJ`xf@ߘϝ q\BQh-%Z 7%ԙP q46mK)E"P^^ ٳ@T(jЄ}􁏏<==/O@_WF``At&|D"~"""weǎ?~c{9lذAe~w\#`߾}jڔWRfE̤!hYSF2i}Ghh(rrrOOOC@֖p \!!! m$fw祔 (//GffVXYYaXne|||ЫW/ez:%f&MBXXfغu+V\S,{Q9F\\\ǕB(aB0.C&fz"e2!JyyyD:t(zd f^Zg[W)D75$\Ç񈋋3z {{{ӧO?Oӧ!Jaccq b 6`ɓ'wzlSS"##2޷o_$%%[*8qӦMSyϾ}`@NNfϞ2tÇ(..DDDe7SGuT :c#B=CIC ØE BA.c&ך>AEEgsgee~!77{Q3;;;iՕ(++7ߨLذIkP(pE\x8 DwYn!7ט@ ɓqav_*1޽{1vXcƌoxsB!Ο?CaǎlV<w֪,1 71dJ Ôrґ"B~P(PPPR ݻ73h%)JH$rlaccv4^bP*Ņ 7nt۶Xx\+<==1h `ǎj_k˗/X,FEEp[Bpu|VVV]|F믝"~?F+#P!={9s={D"CzB /h}D T*O?r!C#ҝR$0 ٺuVP(ƍd͚5$22RW[[[MnJJKKRT*I>}T>ǏyrssuNƏO233|3gΐ3gΐ.O/Ι _`@hsk]Zֺ+<<*ʴBR '''^+==3gd׶uŋsN6~~_VVVWP`ҥصk8Ǯ]xbΛ" s IDATdٳKkaذayӵ/d_\7o\W\O?'ORbgggaɘ;w.&L\,?fm 0W’џZ HMʨ@.#33bBel2իWr#oNn6,04DxlѸq\o[lCPXXu֩%v7|W,'NTyݕ Ԅ3f%eHII1XR4YhPYY$,Y^^^ߘ0aE!11"`;XFkҤI*322 sݐ?1j(5&eoX:iKùaeeIfkkkr7SѶ[XkIs0 \\\ ooo}7nT8m4>}Zc;y'''|G8p GohSSSu>WWaԩ8}ʸ;N:h]ؾ};JJJpZ몪w^̛7Ŏ;t*G$l_cc`r\~ס(=EB:ˁByuzp477sUjhh n"wܡPTTDlmmUӇTWWkuӝg7ł|>mgyFF~G>|HnJƍGѪ'ƍ#6m"4~SURRB*++)..NsqF_~Q|۷ &oMM )!.u1!/\AQ)--EFFQXXȎuӜ3`РAFلRlٲDeO?v}ڬ"ׯ2ӹUXXc"33Se<00.\PwEZZ˱{ntuSիW#,, Xz5ҌV7njF#33Z{zb:}e41k!`~2i f[G2yg1ooo̟?_s;jݽ{qũ>xLNNƎ|p!88hօ'-Zd )) .t-PVV6oތ &_~Xt)9ݼܔ稦Ĭ'$mcԥ##}`R&q0|p&cU^^ QYY Tj?555*c*t}:tHf͚… :u͛xgc#F&?kooXٳUUUp/_??+**Ž;0k,稯MÓ'Op-!//1b𭮮Fnnch񁏏JY`AX"˼!z㠨BPXX '0`J K#Xφy֔Əm۶wŦM0nܸg=D"RRRxbՋ=GIII7F >:rtÐ!CTƮ^j7ΎpKznD3;(56622ӧBBB*1~nj:wkzϞ=+KII /Κ5 ǎ3ΦXjPPP۷#&&ui+V"""vZwc<ϨU#GTyakQl:KCB^)=e BHHH[=+|}}UZ|ɓ'j3Z׹w\pB_|89s樕kU3Xd Q^^,\u!** xwnܸѣ1a 6(x7o4u 7h4#%mmB:ZJ=zDcZwɣ̛D"8]QF͛7?Ntƍ2ܹs:8uT:u}={l֮m߾˖-Sa}wm6[)HdOZcʔ)9sbXn޼3zРV-00`#ݻǖFFFϯ)0 m hb!ok(A۷Q[[  7QQB(MMM߿ލ);*c|>cTݻHJJR󑗗~u\7n7rUVaӦM:h.J%222D;b񈉉yz,BR)b1lmm J&Ye)77!!!!ݽ{UUU(8::rŁYuv }㠨V è4F3ܿwիdxhjjNNjMAUUUNs ((}P(駟vڀoٲ&ek޽{x!n݊gg% RSSb "** k׮Evvv7FǏʕ+!B{"""TLu}aDDYRZyBtf-Z:Rg!(++co}}}AܐRԠ suuǃA6ݱo_rkumx<,[y׮](++xR~׿T|>v܉{O]֧NBYYvލgOJ%ӱn:#88TݤX멆֭[FzLS=࿄ARFf=e@٨GPPƛZcW^eg'Nh*e2z6q ٝ;w0h U?Q-zZss3.\f=^z%cdbHII#GPQQq{3fȑjG$!77666psscߪpa_I (\]!,:ʲUTT><FRyzKuN$>O颗^z 4u._ѣGu7?d_!''݇RsEJJʸ:^x /_FJJ >ku=&OXšO>F[*_sߊ hhhZH:arČ:?Ou;R aؙTWWc=nC[BzDM6 'N0`dHd233QFч^ _ 1 QŸy&ٱPDEEѤZkc*.5|ۍu₿*c?.]2VSS)S%e^^^8{,Mʌؽ{7bbb:,mz CDDV^4m(P"}U)V(xCQF=]of=cF}Cu+Lk׮A.hٌ&cBMM r9Fi: :gHHHP[h^<ի9rl%yر˗/cϞ=0a&OBcSNmҏbccˮKKLLQRRqEEEرcvL4 񈋋kim;677%1 WILy;jG8}ٴXl)#!\AY&Š vWWW#3ְm=q!!!puuIZ4o 1~xwyFRKʢpy>c۶m(..ݻwfDFFv,ŋw9:Jշo_F[ÇF!( rrJi-f0'}0 -2emd2bfڶ1 W}MmۆQF3O<'OTsȑ#pv[U֍׮]GɓHNNƉ'ܬ\/ŋXbXDFFv9޽{at})ۮ565of3f6)8'BH0_s93oqe9d2Bԑ2(R;wkjc>HR= _ñrv>66ǎIYԯ_?,Y(//GBB.\kVV֭[(vcv;.Ĭjǿ !XdbF ,qPucz°a/7#JhllD]]d2tLBRҔ).//}*<<\udС,X6Ǟ={PYYSNa𸂂?ooo,ZlsMB!233Ph򺼼\mT NNNwk!B a))z1_80  <( ɩeeeHOOǭ[̝S0O_}^yYe˖aϞ= "::]v YĿ{ży۶mCYYd2jkkQ__TjE|ѣGF HDFF'#!\ғY\bFx8(USS@T"??}2K;[_gk̝Ǝ02Fkq.}:7T\Z _|EFq!22k׮eۺu+4%H+VM֮]lcI@G !WXo:}(=O>mmm9r888^^^Y[|>_RuffRy-ˑSjkk! ! hJQ:xOIIJ{P 0 0`TWWΝ;gV_|Ǐ###CelĉBիU|>Oyݻ7-ZB$%%aɒ%샳/ɓE!99eO'}OT޽{wU"P:{ z"J! 9\AB޽<@^^^GGYB߿jA$uCWݻon^|.CBE?ʸ@ NkPؾ};JKKqZJm? Bݻf͂;bcccTVVǏ}Nc])2 zJu`FeO}7#QTTBx<@ MMM]>DzeT^/Z|W*cX`AciK*b޼yطoʸ=0w\OYM7mڄܽ{6m¸q:L7EUYYD|~:qygNVOFtt4zРAsαS:`:0t,-1 `qP=L&CNNjjjCڨg6(-~VǏwXf_^e:o,][[Sƍ*㞞8qy)=xb1>|s!)) ZiӦ!&&/BWdgg[oGGGv<##ׯÇUpi|Jezr9ݻeVnp\.ǝ;w OOOͧB*bزe {s֑!C`=zN׻~:^xvV>,-<?f;-_ckjjpi$''hllҵn݊e˖T|}}=J%J%zu8T`0\b,"1_Nܹ2iwEUU0p@K) \p舨(#2pQdffe 00Ǐ1}t[j9sj͜9㠨Δ //___#Hd$&&jt^0c bڴiprre~foKIb$qeR)صcr9997=!ΝвQ84(JTTTw{/ S[%%%ajr vmR3yP(R'>RDFFeeeiuƍ̛7nB äq)l2mɁ\.СC,2 5Md2:$5ogo/B2L2UHNNt(Nx6lbbb0n$==rÆ(F toXJb`&qPl -Ecc#t=zI%9VXvcj*lڴ(TUUرcHIIcJycرF{qE477c„ Fe0 sYKI*xrezr9{P(pM|֖(}( ؙ۷oC&! :_W1ai&\(KWRRDkkktkX,FjM IDATj*RRRpTTThuMxL:N]њYYYajvK Ì:ScI\AP4JLv5455&NH?P}a !xeq> oGQTˆѣGsuS디:t{{{L<={6<=.P(@sd)]D3 s L߱Bf8uiݻwE2F;ׯ_gK}Y>)r"5 K,~2.w^̛7(U۟!&L0{!%%ɸtVZ|=xj8|駴$'ulE! ## puuExx8aQp455 MiKGdb1x<z9;wZx"228pڸmۦ( Fa>:S`w^ >> A{9̚5 *'OƑ#G ReXJHNNFBB:SLAll,f͚W^ygϞg};wkVS[%2.\nnnZQ*t}EC.C,t3ؔC"Ϟ=S&ZƆ>m۶ٳbg;ʮ눽=>#ݶ //Om?J7Jm566bj&MIgas!:xb1233!H E7p !zk̴TWW'''o~,^r\e|ѢEصkEQ=[SS]}3 m닐#2 vM7gw>@RRR;w3G*crJw **JQ |x7 Q0 c}ƌ>B! u]RNA2=J hllTOX|aʔ)jI;RSSiRFhu\1 777Xp!ynMÇ~ߩI$hjjb_ϛ7O-)?>M Culg!Rt'q H$8::bРAZ,H$z9g)((@ii)ԩSqmz{{ĉߘZJۛٳg7@&hP[[T"燯3f09% F7Į]^;447nܠ]< a/q1 @ qP]SSS>>lhh@EEGQ477QS*iaa!Ne_~HMMEpp0aR d2 zmV__n lL{nllD\\AN1bͳ_}U*cJ)\df3#0(ٴ۩B  REEEGDd>u:w7N-)1b.\@2lҥKy&\zeΝ~";3R sձf)eSS퍠 ZFDTo}{{{9z#66VmIp899qE!.]Bss3q`eeD6D"F, Fzz:mcx39u\3;_wڧT*k׮ZoCBBhRFWwQL6M-)űchRFBwww8;;Ӥ١o߾nɒ%4)3<+\a ׵P\CII .6T ;eD7`Μ9욒V<[[[Rf رc1h 19bbd|8uƿ۰ae0 v6gΏ<j_@@B!xࢢꫯ;%^mRlH$ :(sAŸv= qDwUYݿ ͛+Wή#3k'TY(PQ*e[`2dBBBMEqhZ 9[PQQzp)EH$;wP^^\u8&c{jCXrewfTČ2:D7oYYY$pw})\.T*eר-{BB׫P[[*ǰT 4覦իWcj7`[lMвS@(999(..:?{wg]/zLfOI&4I4MK[v ^(A*W8Tr=h eEni$m}Mf&?Gnj)m&yfidf2|Knn.df&v"!֭\?Ζ-[=<bjUEVdDXw$Ѩ>MLL066f;vd2ywշիW.@AAv@pLVT*//ӉlV{q\fyp{@1mEFx n喨|͕ۏ0:#-:FRSSW? /5o[*++ٰa_Wn:8 ,X0&_@,I3!D (k`RRRhhhblٲL&&&8z(z^mulj3 \dw}Wp8xIi@ ]vZRfXشiӴ,BKʄI)n-[%KPRR"'Y~}Ԙ(ukZ999DbPQQ"330ԋ`07HMM<ݻQRjjjHdj$MZv-?|xFF=\r #RMuddMm9.@AAVG 8`|}98NniWU֬Y}( /ٽ{7Ʒnʷm| 9VQA%wZd-) /ђHb65lCdffbX#;UJJ ?<BFݎ狚E$999TTT1LbkcxNN/%D\.###PTT$W"??_Nܜ磵J_"QcV?F6wyQBz!997Czkr+#p$z@@qq1}}}|llǣgddP]]MSS-/''l6ۜ3V%b+VDX,l6%*t:|>qK ###\xӒB~N`r @CCLʤp88p|===z*s~DYV6mڄl>uY<QcBnV|gE٧w&ʄ 0c6on7mmm92Άh+bbb۷x0͜}䜀lb͊+ 38BPVV}}}\|߿?j_]ۖ;SQpX@H1oǎ\.`r~?Q;\={8묳zQ7nӟqWUUJ^}ըO|lٲEXhTM@bM=sz@6'={3c###`41 ڋ"`rƽ/`f%%%i+??_{3ɄbIXKOO. ͖M@;:: H[ZSPBIgg'Qt$)֤ BPWW'_#rFj&K/_"w _Q.26nܨ%GAQ.9<(]v j$ؠwcB `0zjRRRBڊf#///i.ؼ^/ `m_~ssVqQh#Z/d߷o^zjrJ^yY{߿,YD^J1O&oklld֭3{袋 ^}U.Y7tOx$pHy*p/ 55lE!??PKIIѶiF^CBdfj0!HOOOu.E_[n墋.*E+zl6kͼK`T2&_kOֆ  FW/ ^ugt_UUUa~mmLkM7ݔMOm;XPWBb&3qj;Giood2  [WW Yi6QUU+?88ˣIe:jW_ͦMdkg?>-IСCR\\LeeGwfttf sţv?|j#<7|J'?֭[/iE$AQY#$ZkIٜ tuu!RQQAFF466el3JG΢<<'e2(((HF ]Vk|'|r L}YT $i& }}}JOO%%%21KB .{Vh4+}&('y_%f;D*XVinn&==KQQQTgi-ZDAAڙUU$ ʕ+rDUU(--;y_ui #~ӟ-!mmm( fYKSU|)(( ##CBVV+WACQzZٿlV !hooUUmvVb lQU8 --M{C7 6M@//{l0~?0t:eb&bP__ZUHd} ߊadd<@&s)}rmK;#QQQ!|+CSe% 6~1EQx衇{5EQhhhYkD-Iz kLy |I EB$2BT;Dv1B!vڅ jkkg\H=`?Zj*Y#!PU51 !yGƍF#?8֭)2׍$>֭[lTWWlFFFT%%%TWWgE@E}6!{ett VXAMM |yLgp<v]K|>}}}_)IY8[o婧OII駟'ֆbaX7ŋ)YUUUh8EL_ ! I@@}t:p\tuuQYYsdRDjj*K,abb"ŋΪ%+!PP(hL-~kGgdd.LȦ ZoOoo/ .9*)Yc4Bm@QmO}fNJ|v~?^ P\\U]b<1HMMM"!B ###!}}}\tEܹ3j~80ughhp8M۷ogxxl: ` `RݕN p(=z/g6EzEapp!G$e@#G1LjoBvh4RRRBQQQB$eDd&j4F#yyy&lhh6IMMSiNoEQw(n3!D`'!{n&&&p:*JQ"MW^/uuuÁ瓫d²a~GH ˓[i"+h rUY/㌌+3 ໊$ QRRBMMM^=zۍ,Z(fW]|k_?qԸh _NgjRSSuHEmmmx<IIIA@NN555*HSD"3!U:C-@ uvH;8NL&v.1>|!&&&()),s뭷cEf6n7߬SdX,vn7K,Vn7&)atJ'#GX>/l&;;* H*"|b&$tuuT:iyyy(㡩IKf&&&n6pz7.O~:j<--{kVbdH{|>vq|fOgg'!PUr#IsA$^BQ8{&+-\P爤x5aǎSWW+Ov BF1cgO}^{-j<;;^zs=WbΝ;fEmGOOmmmjAA555:GW^^ILw+ށ$BOdRv"/n4+o>TUeppiLeGGG+5nyW83t,-XPWWs4^0UUU|>@J&|YQH2)#E6>ijT K`0( fbj)scǎ lBccNž4v;6Mkrػw/V5涪Jsۍc1(zE?zlb] Xw,G#BȤLU. @{Sv\p˜xBct\z6Zx ***t,>ʎ;x< 6t:ڞ$<~`Ŋ<ۙgE٤w (V2>LN ---,\RIn,i `޽~nV:֨{.RƗ-[kn)z B*Igdd{*.˗Ȥlf*\?)ށ$B ;xhmmEUUڢΘI\fP`NDVL&STc;iIYgŖ-[dRv222>Faa!/E۫\= ؘXYQ&e+*! =@v= `V455EXikkt`+gzzꫵ]t?|(WPm۶ZQ=Ф$`x<YmX]]]tuuۮ uH#ZEQ;d3{@ޱ шnbP^^7e̥NIIvپ};`9A\{x>Β%KHtP:k###B!齜Їvc0*J \(ʟD+W;x(eeeTUUCВE!''EJJVFo߿_l6kSh϶{/}Kڹ|+<#rgjC`r5p8fcɒ%1[I4 RSvyMIFTVV47FE٪w $q B$IQ.\ȢEXt~^nqi(BA.i![neǎ=ztV;{'?IÁappPg2WIzz:˄ȑ#GVV~,#hrQe ,vzYւ p:Q "}NgԹArrrY#|gy&jb뮻n~ i>>g6MZ0d2$my<-*..ǃ۷:S"EQDhZ[[9t*0J1d2QXXHnn.999p̸j X8IMM ݻ󑟟}]w7oL6o'?yY%;;OMM<طoG`0!Y줥шjV3!0 yPI@EQ //wH1PQz"߼'fB5cXsX6>>ΡCɕ2I1oj5UUٿ?GВXa0laٴ>FGGɘ#m:::ꪫxw'''^{O|H3FQQ9N:;; BSVVS+ntt6`=DJJ v F#=dffb۵R\jAQNl^3!D2 7πd`0r(//xKww70yQVPP0o=fBQcNz|>X, p%Fݦ7|+VW,|nZp8%55U!L&Ftn7C0$%%E[%rQf^&WD:9j%c&HU$ٷL7B룡!f*yUU1 9r/X[X` `l6;S###deei۷ovpgh=rhmmrQZZJm۾ž}cɒ%p ўI{|ieGށH'o^ ! Ȥl\.l{R\R:VXOW+ o'u;@KK w޴w}Պ``(Nvv6]]]  ػw/`P]RJ3vɕ2řoX cll;vk.휋$ųDi W\;CII C"xZm|rGettt~?|Dbzz:X,jjjM*..^ 6bԉ3Y 'nt(ҫw Zb3ڊ!H cttɊn:[d~?\s ;}w˦MZ!ͦ5+.//`0*---\.eeeF!N0LZLjVUUUs˅dJhStvv*jVUsHaa!4~3LdggcdEpɕNN͜ǯI)l6c2(--;I5P[[;gL6l@[[ ؾK\veG_s5+Q+FNCCridʀ[}K, ժ%Pn:::صkIXAd۵1ݮ%^`P#E="<۷o{b[NN%%%Ȫɽ,aBd1Y~# v>Ѫ͕,Y;~r7O;rM7Oq(bhhŢ%N2 hXJZ痢@KaAyKK YYYTTT` sF<FQصkN.Ӄj[*sN|>&HIN8+e;D:=s&=؇LNm"A&eRšz!ɑ#GzҊ|ԊrwN[;O:>K&ivP((!Σ577)((`…S),,$-->}èz> ׯV->@ tjۃ\.VlllLHx^y4)j=d2 !$c4)..`bbvVfr{$%\*++9rȜ5Qz|_b&ߣn( ?я}2QebZq:Qgo$55ڸ 'A}}=B&&& ^p8 lxx|>F~ŋFV!̤ۊ4LN Dc~_[A|l&55-1LKKjWOO̙fN?R҅P4{f51BO1efʏg_͂ ͍g>8y"3swOD}h4Op-J cZ93Q"e}>_JCoo/5"(,==: 'jjPUuschh@ B!ϝH"e||T~6YzӺ룭..]*q EQ;ivZbsetia-RҙE.O H{{1! g?˦Mf6ٸq#_i=LMBhUV0k" _TRRBNN^wZOL0J"g֦&kS?z-55 RRRc…M;v:IY(&A~?CCC21S!=Z2<LEQDH ! Y24۷~zzz29(`2Nle͚5ZUQXXGVZuJ7_RRRɡT hkkChEQPUݻw X,3'I'KUU>ClL&F\)sRQ_^rEQ4wNkKJf'5>>>/R3LTTThv޽[fw26l0G)//wߍq&38 JJJdvxxIGGmmm$qٳG>lCʕO] fEQ4Nk*KuOxC鄄088Hff&z$I1j%O8NcccTb2())!334mr||rUѪ^/^5nmmPPP@QQ\INd޽B!$͊E[ wBeb,IMQl6Q$Iǐ(!&wuI%e0 넷kjjIo,áBp88Nm(0$HuP-2I]ہB !L/xB! & Ijh"O /K/4k____Voǀ^☓Us\ܹS;6166Fww7Fe5rʟٲsQ%w 8Bn(ɫ^rssY`idxf^/K,*1 IDAT׳nݺ B066(Zv= L6&˲펒k.BYYY,]TH@Z:KQiόeA,Ikxxp8z#I1odd#?i'e|[nIdDQiI`d2 zR__'ɛX+" jjYas=.3WR!xfX,r$#GhkX|90C}nY5DVǵqǃ&/#@1l6[­&J'VPPlСC466_l6v]"Az!ͿM!nsbIjp'1JIܹT-TW]u)-+(({+_ܪw&&&8|0Ôj:t… ))]S $\(I !l@ OK*sSԌf/}{3aJ8*6ܹq/_S"==<`:::"//}Qܦw>fȤlVa>LNNVUI Ljp8җtߛw룪J+ǖ188 !X,|++ vk-%FFF@稤Ӄ`0P^^.3ν>^sR_ ~܆|9z(G%??FCr{&Ǔɝw}ݧHncBhАV, ''Gfhhrsse233ihh2%meD&fsK]y{Rf>tjظqqou7UmP:}SEapp0A__t 6ddd"UIDVeW/)AH:ab&dbI:fƦ$u]Zg6gx9c2(**VU`r;f&/oNNNv]NNŠ`0Ν;ͥEQd$S]]M(" ղ;IMZy%餥&I)zgy嗧fnfx3Yg hgiGGGRGky` PU={zz4yDQ#z!D+fk3#d3\$I3zƌF#_~9)((cl4jriO=taaX-No c2$I5`AHD%$tQz{{ZwHWsh>>ATUZJHcA-1B088ѣQGngX,fOR#`0Xtl17~(7G&fBVc,Ir|$ {C=``ڵ_^j p\FtV:Byyy0l6Қ}{<mFZZeeerf)++ZabbB0B[1[?oQ$@ },WˤD()%%L|>|d<٤DSSEʕ+zAm|ppܾ6;Y%pvڊKMMr/Isk(nb13!D#p<ǒT8 G6sY뮻g}1zzzXn L*#EJIIMn%i83bٲe'w:ٳP(YbC%%%D3|Z@ l&CaIMM-xT[[KZZ,^XIL&9y<ǦP lZ144DGGTUUQUUSAUUNV=rO&es(L s+d9 ˗wH4+Ӻ{ި 9Bee%IতP__OWW-[bj69D%d2QPP S~wRlJ̄:Œ4X,Q"$)߿k ׮]>Yviii!;;Ӻ/)6eggOۅ0<(^z~9L)~ ~gpt dffhnnߧ]ϐKۡ( 555dddpꫯbX<#m۶{ӎk-1TTTDII z;z!6@auuu@YY555UUսc?--- /^ qڦ}K/_onűcVZߺ$ňwRl~?I4z{{+%e\r >ƒISS999_]tG&##n455|ۏ?Ak˖-Ӓ [[[o~[oow-Ugg'BF#znReAH- YX2v#Yf /Rxii)_(--)I}}}*VyS p& pOjT_ 6zjv;77ѣӾsO= xxg}M6[oEXj?0{<~~mX,V^w8^Qb[d+b]H~_8r\˿LK233yuO`LP2Chjj9")^L&JKK)--p8sss;::BPPP0+^V^{M~yy~ݻwi&~_Ln)>6/| | _`pp 6}v dĄS}ڼ@AH$Pټl6~!%GΏL&~LpBFFFȭX4SS/fkkkZ cـ[B!9쳵-jY .cn,_?OaX'?9x<<#Q=КXtSQQmM7t ?%KJҮ(JՐb (d9yb4̡4~o|cp饗ѱfV\)[XHsbP^^Nyy6t:\-)v g1񸓀۶m&+/lڴ}ktww—eDߧ.>,L-j#! S̓P($lJIc˖-||_NQ}IeppPhdYgł ())^#G0<<zdYmm6IEÉwz 217HYIJD\yDZ_W1Vk|z4/_tj\.6Q Jg0vb׿ /ddd}{+ ++똱ʓO>&iii|_{=6wy~ncҥKyyq:rrM7m:Bh9N|IټX,ddd N$Ţ6niʪx|euu5S]]w8RȈ IMME.SRRHMM[rw/?}OEEW^>z_/k+k^{-?Y` Moi_w8lٲ\im.\O|gyFkgSWeeU6wR|P$I񭣣ﾛ^x!j<;;{F"@ l; If||Ţm߾ ҥKo@g_϶m``` * `WV-ZO~.첓3 QRRᠦiyǹ۵ϗ,YwImm-===lڴ?Hss3UUU3ONs&8NlEQ>;)y!Ii3LӒ2ﻪj^8F2!===Fe$I]]]C0DQ-) ߿_]h4rUWE%esAgeeokFoXe00ȯ>Ǐc8z(..Z>iڤ#Y?M$o ye0JRRUUU1LMqƍlܸ뮛`;wzIII@^H1wQw_ߙL}rB AA@Yu=jUwkWںݮbO* DApw&3|͐HIǃ＃d~ޟJjkk)(( 77ׯGcc^lkk+$y՗%%%yv gnuuuܹ0wó>˹sصk |_W_7 z)}5Q\:Y15Tޅ?qQ=jt(B\2ѣG)**-~t:78"!i|mv;h#&&SbX:PZZʞ={hkk#99>JAOSSW]uW_}5/Q+8H~ƍ#<2 lDVV}7o:+Vk%$$.~ӟrw!.d׮]ڵ#GNs&@ {7 D+&GJGGDŽmjeΜ9QXXHxx8FGG3sL.]v\dׅp1q\TVVoo,^w6,__WC۶m>l2F_cP֮];\^^xѴ4JJJx뭷xy())Gl޼W_}u󉑑JBĈhJdPQQAoo/&<⒴꫿g&::zt19pfbXF|f33gb$&&RPP}lذ 6pAfY4}Ggg555dggxXz5oq:uJѣGV}N8Xn]cl6檫 tvV+]v+5M;`t" Ǎta||}AJJ wǩSذaׯm۶p8 PAA7nܨW ֪g0)odBB---TUU xoڵk`ݺu\wuC>dgXXpaL!FĄVB@ww7һ9/)ٷo9sp$s1v;?~|T=k,̙CJJ IIIwuF;c=Ƒ#G8qO=7ηeہ._cpn!As\8ءuYrss|ޑ~; CAOBf, Z\+ӧ'U9ndd܌fU31nB)5* E_&d2ͦ'zɑ#GHLL$--7PJ?i@*++c TTTDww7M6-[Xt)k$xqylܸnWbLpc(DpAVMuu5PIDP1r ']# Bff&mmmL6M21?>Ns\fuEEEtRZ744r'22RPw8i>(YYY{ڰa~c >/҈}w˗/^{M_Q|vs׳|r֮];`׷M74;oy"tꍝ?5)FDqۤݷӑ(4M#55AO9AYYYdgg$hll$??4ÓCBBy˅ijbzzzHNNfΜ9<ӃWNg>!c:䱾VX1L&wuאzعs';wpd e͚5CcݺuXVJKKteUl磔!"(444p N>Bm2{aKCCF$&6JJJxtu,޾ IDATuҢw{t:l6bccߝ6Xzӧ9p=s~M7ݤimme֭\{keee|[oex$$$ y>á7D ɹo`sĈM5:Bj,26V l6'N@)ٳgbfǃj%>>c łl68O?vG[[9uaaa撒B\\UUUڵkk2W^yti&fΜɂ СC:tH\LL UUUz3-[1h|bȶ3y6eKb4 "pXV}Cbb^);d[}WVV((( ...#f\uU~+Ⱦ6===X,gy^_bի<ַeZ[=WHHwyzwij}_Ò/l2{9Ν;7h|SL)dx6id ѣGE)W\at8BҢdggArrr@_SRSSC\\9PWWԩފaKNQJp8&!.'Ox}Çp- ,lcu]74:v @aa!)))ocϞ=ٳz?vww͛TC{zzOV5hJ0T$ <1$_bVjhnz<L&y"ǩSazYp1lfѢEU^&_ii) ?YFFz&I Tkk^ iZ6mo6l`ڵC>Px}k˫ ٳ;wncM~eڵ477i&֯_Ν;7oެ 2FuV%1tHb&{lݼ-)E ihhn뛔"hii&/^\\>nMZ*++ill$44[n?ö) x{ｧ?k<;Ckk+0p숈dbb"_Wٶm~ ̅ dR51*TJ y(&sb2}ƹsr nll$11qBDBBf͢jyWgժU{ـ,RN,,YŒ3ub0."N'555z]0O~Bjj*RKq`yrss),,$33dF:::8}W;&&FOn/~Qo|_*b9_II8 )-[^x?OtM,\͛7_&O>o]]]۷}q)Ù(3:&f{ b{GvMFbb"fsdX$# [ 6l~s=dz>֭[l#>_[[o6rssrOá'`QQQnÇ  X~L:OӬ\t~/\~1m1===7!??qַm۶ {??<3pvvF [ڛ]E1IcZ a1q8$$$is̡; 7jjj`ɒ%R(CXXAں~k}??>hW_}^`[,Yv2_2FLL^h Geܹ<:tz~m.]ʵ^O~`_?!d2裏zﵸ.}}w}|DGGUW]Œ%Kl2LӴCơCo?Ң~-[ӇVtY 555~z1; |k_GGСC9}]^}U{9 xf=ɸ MVVmmmDGG޲vWWMtt4u]GUUfٯϹs}is6W^yE_챧GO.fIYY۶meÆ ; |YGVFj$12SJ ޶H\Ckk+NiJv9|0ny]T0'))I.`ĀN8 //.ή]k`ttt0{lΞ=KZZn׿{կJv IMM;>RZZdb~:+dCXjfL2[lc?S?XYYɏ~#~ߢb֬Y9rdw3ejjj|x<`6?bTj6P=1i_[6"vsQJJJ2:1566 xv^~,X@1*.bP䐘8`s`ѷ Ə~2}K_Ç<ηcX;F-ΈYt*o?ۛUSSfpM||<___8}tcywԧ>5dR+XÕq/;;_R~z~Hyy\ý^o>#]G,\/1xvb,^g+sb,qq9BWWп}J)9r䈞tڪܹsvqF;[_ϟ?_|ĉ~_gǎp $&&^3L~{$''Ovݔs1;`r}P~@^^\s_#M6{FRꫯɑIG!+d }tsv4І7ֿ1dpP]]R EMMގYjΜ9x̚5Q<ǎ<Ӂݹs'~ @_~eHѐDRRJ)EssGll~:?)7o~ß'N>MRRS_2 ;C[[0ϿuV^{5N:EDD3gk.}cO>4,YW_=h{@~~> ~G멫3g2uT233yꩧxꩧ8p_9`]];w+IIIC0T#>< ^{ｗիWSSSЯcϞ={KOttߪU)hс/1m!Ѐxɷ;""Pϟb`MDDw?effMCCA\.^~e㑔x|;wNʕ+L(`ݺư7΂ ʒM8C:u4jjj]l[]lʔ)}PRRŽ;hnnfƍlܸ(}ErU -c+$f`D1F}}=a 1eeeR__  `5\j1nf3&|Ӊ1:K}veNSNWNK/{Kx ǏuV}/̺uFN4Jq8ꢫN.\xKW-[Fee%{0zK߇z~rs[vxAp!^}U6l@YYnsy<}o`vv6/߁ᠭ LDD̽}*4M6$ P{z~ll$fÍ---L:਄㡱![kb.2DHCu|9bsOc_Ydxx8zXlـ{m}dž . ﵴoSUUU֒… :>oYﾫ}[7d2pB.\ȿ˿PPP@qq1VuĬ_ ,٨wF"`E d軁p(˜7oIRZ[[ٻw/'N~nS\\餸X/N' & @xS|k_zUU{oևBԍ7Ν;? o~=k\uU"mhh fopW{޷ p{8@qq1rˀ+}(F ˀi =αL pWJSq\.%%%8N[Vh"a8VUoР?WHH2jZmnݪ //_z۹;anÆ z+MBCCyo9[o5qJ)>~_ZFFF~r-@W~!|x6lp@;C%@%}rWZdee巗NBԵF!`M֌W EHHȘjGoo/EEE>_ㅆ&55uDDpX,\~9s 5lxܹ<Ӄjr-_g۶m444e] 7ෂ4c bbbKη~??C__w1 q8s8~2vO}S|oժUX,N'?Yjg;ro}KOcbb馛kP m)CNa+93zzz mov e"={nOO&55(L&yyyF4Aϧ@yo{m:s `_p뭷{E;999}{̘1C/lSߤrժU<K9s/^}>K]b`Zl~egΜܹsG:u*}(,,[n!11b?>H[ouYcRnUJjvzCŤSbL:uJjpDjooW{QmmmF"Fv+etb۷;w&|ImٲE]wud2$%%'|RlOǾ~m_S_?G-YDfkJNN/w.ۭ2334MSSL-ZF/<6My<R===jǎjWn[)R[g}VEFFoXX?~=oooJHHP.8vUZZJKKUccEb{=iZR ƀoIwwo|OWWc2aR"v;555444.'ۭrPCWZŪUɓDDDpe ;_tttһkpo/륂|~wRӧn:Jٹs'?8x ;rdee@llwLyy9dff?޹s'f-[ƃ>]wx\x1eݜ;w<[/)~iZсc nIF]dd$& *NbX=zӉ墠$IYp\={zI&Laa!ǎ#**tCuv7xv<HOLL>k(>|ޝ_w~yY5k7;\./+[|M6励{qzݻw}vaʴiӘ6m_F>0fmww7k֬ ?_UU/3TK~pgDkD1L2)SHR&Fto===\.$=&8N,aK`1L ?&k޽{=\?iD  y7'v-ccco;2?Ν wĬ bYoxrq-[ YYYֲi&͛}Ν;_󐑈=\61*iwM\D󩬬/,̚5ht:9zr=Wzx饗/|> ]):}Y{{; x-JzN0egg;ɓ'ٵk 1|/^|ɉKtt4kmmՓϸ8=)kool٢g?{Q1ָn(&Q%Ix!:u8--mcN<ɪU5k_yx衇X|9?MMMC>υJHH`@[[eee>GcDD˖-d2qwj\bL=htXíUTߔ' @FP0̚5 `Acc#IIItww#wcve]C=ҥKZ<쳜>IWJ,UaѢEC>޽{۶m~A^^.2:;;yx饗x7Fn$4M#..߰OC] >xZ[[11n߾]x= kpǒ%Kv IDAT1֬Yfd2O#mmmw'nNF+NAJWW6 ͆П?)) uQ3\.v2o1LDcRF0_)>>Trss [ZZL2EOMMMnhhhz;|"f8pGy%K0}t͛ǣ>ʮ]-J1׍رc}>|կx'nuV, <󏄦i$$$cܹۡsٳb:;;/yNиRjn^Cʅb#'RR,;33h eԩz\vJJJx<:t k.… , _N'0pt:yW1cC>ONEnݺ={F(>8ŋ/|__ځ5%%$ZZZVsx<4553d󚌌17 X|`t b njKBLRJ)?N]]TVJYD5sL%⨨(i֮]Çyǹ;}B*V8?yx zzz<8s~Qd2HJJox<^9ݻ9y)Fdc "6vE҂&.}fiXV}?Dn->a Kaa!̚5Ky 8 .l#Gk֬[r\lݺ.wtt//m=Ê+{ӥoL&̙wM) с$>>ɄW'-ESQO3uvvRWW?lx}byWرcMMMsWs]wFQQl/tw>0PWWj@!ùsӧIIIٳ~+B۶mcŊ?{ashjj\.fMp\ٳǃlfҥ={V3c=iڿ_íKT\\s!66VO [n{~ŋՅbUI$))3g1l6o;&Y2L&ey'<|ѢE,Zgݻ׳qF}>)~fMˣłl&,,H$1dM̔R@82X,*w7d͚5\?Σ>ի$ɢE4:1)..&++óX6mڤ?IȌ34 ^r>MXt)K.Ǯ]X~=+WNJlfԩL:UoJKK G7i]2iP1Q z D)u%cRl6f\fF;,,>HWNhNpN[׿{nK\s v"&&'O,zzzhll$**Jf@Ӵbgp-IݻwsΞ=kt8Wq'{#nSSSR 8>LQQ-IgXA~ ׮];)WȐ8Pc@{{шKe˖1=^L J)<ȩSvҤx(**VʌIUVow}B^uB9F P?[Xdd$ᄇ WSSsAWUUQ$"iF||>KM2d21m4JJJ00jgn4VXy@4Xr%F,&I&} V@)%&9spKSS+ сd$D_?![lp,22z7$FqM$9D[1 Iƍ$eCaa%fcdIII$%%ƤUYY`̙hFLL!xb6oLEE7ofϞ=:u6\.7x$eb,L3`+fOIWJ%" ӧ撙itHB1ӣ+MIȣ-55ӧϗLLj%)\Οc;51nI$fBLl===ݻOlx~$%%.ۧA:? &1bccX,DFFkt8B1' 6 ipDCgg'梨B$eB t:CwٕW^Ihhѡ!IJJ NXiJ1 zzz(**vSQQAHHF%CKoJ E!$J)N:ٳgZ,XA !DStM*DVWcijjf! !DP43f`ZILLLJR6}ǒ455*Èફpp_'OCjj*SN58"!$@j @$&& D1illat8Agr *++G!.Aw,Ǩ $HIIBLtuu-/Rc)BL4M>LB}2DG!YWW.(uRt:5k!ĥp5M+7:14?`z{{"!!P"`vzzz?BӴ!K!S=itl Pkp,N`/*IӴJpY(&b|#Q]]MUUN mtHBqQ#(BLr. 4ؐdoY_*v %"1,ǃ|#B6_\\UUESSs՛~!8jWӴncCѐ+f>J8-`؆#F>222Q x~~>F4&9v,\pRn ! S 瘦ig I!W|4MkSJ6WmHb8,^X 1i;">>^2!he9 >DFh֡Zw"#7)s\L ]$dKX,˄jfΜit8B U@%P दiՆE&%Q)c_J) o_CMM Νc餧B rQYYINN>@!sxK}ƛ|i1*@!E'VJ5 G\ŋi+ !R"Z[[@ʱ@7q"iݠ80܈Kϧi+J"%c!(\.")l8N4M#..FEGGޭv^ĪN|d&4M2(N!%/(__pąjoobit(BQt:ٽ{7,\FO{{;̘1Bxuq'+[nKy!Eh qQJR!&O^]..V+ŋ~F!Ępz{6̨k x-QA !ڻiPJ}lfytwwS]]Moo/F#HlN>Lll,sl6KR&ĥsM|]?~kqNӴVB\Q^nsy<ݫ؄cdo]j֮)>d2Ayy9& &P!!!QRRBxx8999F$D j>sOMӔ !ߘVWJ= z r 33S: jkkq8(6mZPPJFxx8F%x!MŽ%BRF-'Fg!caꫯ}Yeee̘1#J!FY؍7۩iZ! !ݘkM)u H deeIIAd2=\jΝ;@oo/s58"!M)$l7(0!ą۰+=8~8འ7o !.f>Zvҏ &&HsvI65MJ G1;iH?1btIOO9"gwNNtttet8B 0ۀ-ub J$-~ٳ8N222n"8ZZZ:uѡ1iQBLlVJBJKKn7SJ1Q AӴ}F#pѱ;;}P6 ajl6ᤥKj;w7`P,4@b4|i%ؙfRSS0Ll6#BPWWGEEn/455qiN'G{̟SK$Ib" 4J)x0O0ɡl)["o t&**.RRRu@4mA!XRƾR?6:rҦ[۫h5Euu5999HZ`ctb w`l2:nSZZ yM rQZZӥ&jkBWJ`?0Xa߾}{;;J:G%[mْ%ǻed74)ӖK!B:eZ I(tRL`hi- IS vHv,Km-ɖ,v#\>3k{|<'CyN%R^sy}Q9ǥKD"a5G(0`fiADD|=3[ֆmP(+g[Ddy^|EfffX[[#ٕ3"a2)u%qFs?h6H$hooW1 X,zjkky0]]]Z-bpC!DDW6o;G;tPA@$H|8r<ٯERSOpmy]$-"Ϙmk.G|g)v+++SO=EccH"e#JQUUcgVVV$s)ėb25 ""V2+ff|7ݹs'jll,o۪DBЎ2d287nؑ+K>R&"墤Ůjjj8zFY"iIIDAT`fffϗ'O:xBlhwR϶s_wbŨe"`zz1 {scJrᜣ"["9GMMͶdfBDd܊ɾ&>WRccc{Ddgmhhh׻z*/^2CLC즒,ff6Q۴¹s瘞:o&tuuҲb699KKK\p v(Ȯs""eٟ9~,, d'///F5Qdں#_*4"E"""$Ϙm末 񝥘r N8#۬D$(333>&'|m%_s'sZfs92i޽KSSuX"`bb^I1KfvwVmf}(v,HpEFGGuLd,,,Dy&3 .]brrAI]$EH*bQ … ,..Hl?qqX,@<ϝ )2i!"KYle; ;s¾Lpe8~8mmm# Fb1X_#066ɓ'܉O{|񥬊s;?MNNFme2*++=&ylpF|񥜶2`f|wbՕ+eApuT"gyy9w 1?? ###{Tʤ}^LD]9 nmma($RTI&466rGH$rwfggilla8 Vv+ff ٭ d27En~~ωDG42AQWWWsQ̌fQd7BDķ;csݹs'wtKK8"E#1119r.AȎohh  ;H>̞BDķr/f5d/>;KKӹ˦s$ jjj<)A h#SQQCJ>޽lfSC1{Lιg|kb1jjjrOh{N&155E__-kTjDf8cfEDuЊc2myQ__+eKKK 388ʊd"&ܯ'''{nBLJUDD^?m/9^z%({s*r x&nݺ@ee%o~ ÞӉg{|)F*fk.}g)ED/RQQ3g4B_ ^4套^9[;H>E'uLDY6Y9no玫ٳA+eD"=ɕJB$I8@ww7wav#%UDDVvsQ8ϟ'JO>st:͞={W)r})Ezgf󝣜DEE'Nm-SActtJihh'̝u)a03meιW7R.s@"`hh>ԗmgjj J$nߠDDG+fRw㾳 3˕2/_&144ĝ;w4ī 8|V}; |8;K[]]͝1K& ===0g w9`v7ػ Ye5WL[DD ^.fv9}㔵̓?^J:vnaaAŬSQQA&W*.00jfv;皀̦G?މ'>ef1aDD䵴bV`sg"FvKKK<pZZZrJpgO3;;˱cTʊ#[v3jBI3@>/φ/cf=GbVso hEb_;w3^q,//3::Juu5*b'C!y39k8d [?٢O~vE%gBgfs<_Y8Zd>=gӣ\%`9O"""A+f9wKyI&\~9y6Q2_ S$ahh:uxa[~e3;H>8\?d'IIӹꋋ !O?.$.0LNNr̙wnd\ ΗO]K#""O:TQlx0;lH9RFI&|ҷW266F4W\U #;-fv3J1+fsYOȎhKkk+[F_|x>ιsrTUU سgO?4T)?me,f[~w-:OfiiH$Bcc#q?RZ:&LRWW-neb4qg/mP1+fv9c~N ù"&hn#d8nXd2U͗2Rܰ+XTU ;)rCuu5d H&ajkk TWWNֶ? EADDDJY14 ܧ;%Ŭ}xD"AUU@477N l@3\В|[H"U2<ѳ2;4鍡|)LȖZ]vxVDDDCLp~ xP9 1gJ&""?*fk~x?psى_DDDԩk8突![ОGDv#{OYwrb&;^8"3{싾3y$9#{ہ~H|Q3DDDܨcqΝޱxP7l>3}ADDDʕloή?H1"o """LO%/sْv hKDr"w|39Fu{vK8hLDDpIAr΅Ȗemն칶fb7oo\2!U$ky_,}k{/d/zu?W-y="ɮ """H qUjAm&[lۀMh'rY#/]%DDDD^KLD6^t?#d y(""RTD䡜s{M@R7 +fFDDDKLDsFd'OYw*f"RxpsX0"""TD$8 Ɏ<4wgmf~㈈O*f"RRsd'9n\bF_+3)iι>%m@.WDDD^TD8:7? [֪Q` _y*f"" [pj֟Q 14K+""".'GTIENDB`vector-1.6.3/docs/_images/coordinate-systems.svg000066400000000000000000000405221503546127100217430ustar00rootroot00000000000000 x y z ϕ ρ θ azimuthalplane longitudinalaxis vector-1.6.3/docs/_images/vector-logo.png000066400000000000000000000237071503546127100203420ustar00rootroot00000000000000PNG  IHDRg>, pHYs ? ?.{stEXtSoftwarewww.inkscape.org< IDATxk]U'ڧ.)"!!mo7tKcm~z588x~/ GbAD ;@,0滍M6]z\?17g c/|!"Z)xeccck֬hǫ|=g2<8WsWƘUaZo7X9Zvgڟ9z@(yOnrR+WDO{ZkGD(h(lV q$<DSQ_\b<""e}Rφnnllk[n=S/瓚)~|;==q8ѣ<;D$d8 ;k'OLL"t3sNNN,">J!yo߅7oٲf3%AhIR&tO]uU{7hJ8ȮS:?u<?;rL"gfu>Jd{`DQ V-J??x#e5FDh ,G1Ƹ&L]*LgRyWWdtpt>ϧ*aqK-}G]v ^bMROnrVs^ED.I_Oq=1-"OIlmZ MDcZNVu{#'/(&j@Tp>Jr)Q+r'54{)ZmxmwYk9?cg0^~z?\TgpPwYΧ5fI+_Xk?ϼID^9`pAQ&L.TMZ(?͛7Zr>z~ϣ$.JL꽡H|q=}}}@R9'~ϣ$jZ U!c8,i(Mβn 5.t,НC<ԘTSD_-EɁG#j dwwwkIۇA??q*T7ZazTE+/my҇,5GADr\p!*L$NH|"s5jzz:X,py0G`P-s188 ác88OjPR:\s NKZs>- 7YP2 l%d3 sB"g& P\Oe4Rf _qPR:ԴF0{{{c6uJ"JYPHd2DGqt8dDuiUiWq`j]-A?CisMej[.c^ `w@ d6]j78c̛#L@M˼_k59Ν;?uKx"= ׃w"3``Z Zeэ7:=14 _+ʌe˖S9jsy qPD$MxKy*httk7mt6-PSQ1NɄLDnikk{!e _bܠ[en1Zcmv,tlܲ K}Bkkks}e`b8&? `Ϟ=NTjP1s{l?^O 1%ɿp|mSO=+8p`}t9Pܸqnc`iNN1\wjP(j8c.zRV)jI6ޓc{.ِ Wiwxx ;7RL;w,8ߓGP*A'{ŕߏ"ktϧ)GOC$o Hdgzv=g q2¾0giZ-D $""]#RIMl:,"dg.c~fzz\XU*^#,n.8_ߓS]+hcmh 4fq DL/Y$p^'kQc"]گPj:yJd((++DS_HsKS{R[y8F8v ;֮] ЖI`Æ U*f}EVS c>0tZk 8T.yeq /BG G1|5ű_>3j)"VD JP1_KqMaz1fLýP oDx/-vZqDž1"""" R\l8Cr1r|} =E\ hcLf#L}K.}cP8Ɂ, """y%>`In̶0KD|? B&wҥK_0%?dLڵkDclR*7J]]]nR K$bgyB&DDDDT9D cɴ`LDDDD *_"y8=`""""GD<_cY7|7 """SDD`8LDDDDx1'JWc>o*"Jr,C c"7z͛7?Z)Qb``iJr1-$ d:'"\:|ֳ>00qVT^k}N8G:ND.C#4c8~OD4xmYZʬϛ17Y9Z!|$U\^8oDNq%tf>:zXk?P[kq$n\ ޣz|G9qR}7'p]:7sQzc.ӡ k1P\\vT ҺO76::Ză.Q_f1 W>Y.q,72 hR&jN Ϫmmm/QV'4 Q?4M{ڊ:;T5ВLڮ.,-1!T*̳e˖á]3,'Bn1FKRժ@g =xh.EI&\b>Zɻׯ_miM k퐈:8)dp V] }L~Ӵp{*8piIP 9+"pB8Pl =c8m155NOPɫ1_2F\O4)"tn-#?-iXȀ&j@!P?ժwcS/ߌoF'&&`iNxXntt`e%g{]_ZkD:;; e'Yq1&ͮjhZT*k |3Z @ 5u_hZi'; @s4ഌ?|UWiYʐ2iZ[`E+@\b}M;Dȥܨy@LE%1|JDLcVY*icv۾}{={R.QE 9 =+@bÀh+"ɵtt$."'(fi[B˅  q](;;;S?tuYn&:'5EoذA5::ZRJ5nqdK;Rg o;EDҼ&OEэƘo 8k 0`QuLh)(dî]pL8;C0ۮ]X:\2GE#Qr<(ps1󺷐׀Z1FUX='G3xxxkݺuCfMȠW_'OJujB15 #MdͧWCǑ(d E UKYk]p^~O7#drL111q2'gܹ"!Ƙ媯۷+!"*vTJK\APvժ%B@@Od8ZN|HD^:Dd-{)j)@EO)j6 hEAk "Aτ{Du *-ׁ%KT^"x kt8NӖN%9rE8LEdW э7ـ:'&sP%j]k׮ DZ###8ZXU"sBǑp*شi3ג\l߾CQ|v4ϡ -7>i'5#-ׁB^v޽@G8P(ZC#N4FٳbV:['CcEQXuttpxnj@,5FD:0LkK%]r^%Njq&<0G/pb@@j*Do{{{ bjJ (b\ D#JɴѭUі@7ncr)Yepa鳽C;Bɫ)wN]Ss3̩1jF/EK`y&Z--{~::i)PuO#ti (+V<8"IY%ޑU`YChP`*EoF'{zzZՕA<'`cLˏc\ ",KlA] (SFk\DQfD'">OCkIP) 뭷/٠-[sU_"-E1"ZB-c"ro zK7l0 L٠VOr~}cѲy-_%E "/tֲwA[n=cqm 䄖/v>'Ye pXO.& [BU84?&8m""AOd[;jM"@vZREEM4ָr(lL=cJs Ø{7-XwSN9uKN*%7o4 0Y7Z*mcinr_`X\y7k{0HIDAT׮enc͚53PZ ȓ8^K 力[ 1!BnXrjKy"tur-^u*4FE)'"r0F9j[rK6V򇄖]֭ 5k]14T%"ng"2aE;ݴKO ђ| ;"c׃mjT*q᮷^`cLK@DQ4y-7z&bYJ&ª10c)Gml2Gs@]o m`40rVQ`ea [k?F\Z(z3Cۦ:3.Ko߾sj/%%Y-4j55{(*t,"_Gqݝ!s-Znܸqw&V)ZoX͍Esq2Q򇄖+T5ư8K\CW(߄EMgC)AK'oIƘzo8ZB mիkhvYc;E$Mj[[D:,G" p|$p Zn|-&áշ@rz>hQ30m (Vhv pww^"g7ÁcrkESݻܓU0%]h f7zzzڶ@XA%D&ְlC[C6u<;d -ng#[ '3Mt؏&<+BZU{!"Uc̻qW"^gge(Z@"Z&q%@sA<LNN`B~%"x kpiMRqX\1-` d0jO0!O1nBPM1d\D 8`:Iu`<:F3G\$"W5CItw 1_T*m<-)׽"]#!ZS&"5ư?T]e||dScXfc<'t<`UA@ύUeځ`CS p -ׁMC#Ƙ73+6?c^չd86܃|rEQtI)1$ZѸ@SE8Yq$ C6w""uVKLb6Ee6c&ぁbhhdy{}#"6\ dK9:(rCS/|Hm}Ƙu]]]gW=%6 kퟄ!qO 1`c̵c~\.;89.,֪7aIƘ˻~zsş |qn#gz7"e+Ö6"r+[&&&n5Y7E W[k "'M _h؀+.J l[l鞙9;t`iMVf2:Vfmu& پ}{={^ .Dm C c~hrKL}xjJ?op;$~4,cg.¹_.333]8oZ1QA_;ENDB`vector-1.6.3/docs/_static/000077500000000000000000000000001503546127100154055ustar00rootroot00000000000000vector-1.6.3/docs/_static/css/000077500000000000000000000000001503546127100161755ustar00rootroot00000000000000vector-1.6.3/docs/_static/css/awkward_output.css000066400000000000000000000003611503546127100217670ustar00rootroot00000000000000/* to render Awkward Array reprs correctly */ html[data-theme="dark"] div.output_area.rendered_html pre { background-color: #000000df; padding: 15px; } html[data-theme="light"] div.output_area.rendered_html pre { background: unset; } vector-1.6.3/docs/conf.py000066400000000000000000000060551503546127100152640ustar00rootroot00000000000000# Configuration file for the Sphinx documentation builder. # # This file only contains a selection of the most common options. For a full # list see the documentation: # https://www.sphinx-doc.org/en/master/usage/configuration.html from __future__ import annotations import importlib.metadata from datetime import datetime # -- Path setup -------------------------------------------------------------- # If extensions (or modules to document with autodoc) are in another directory, # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. # # import os # import sys # sys.path.insert(0, os.path.abspath('.')) # -- Project information ----------------------------------------------------- project = "Vector" copyright = f"2019\u2012{datetime.now().year}" author = ( "Saransh Chopra, Henry Schreiner, Jim Pivarski, Eduardo Rodrigues, and Jonas Eschle" ) version = importlib.metadata.version("vector") # -- General configuration --------------------------------------------------- # Add any Sphinx extension module names here, as strings. They can be # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ "nbsphinx", "myst_parser", "sphinx.ext.autodoc", "sphinx.ext.mathjax", "sphinx.ext.napoleon", "sphinx_copybutton", "sphinx_math_dollar", ] source_suffix = [".rst", ".md"] mathjax3_config = { "tex2jax": { "inlineMath": [["\\(", "\\)"]], "displayMath": [["\\[", "\\]"]], }, } # Add any paths that contain templates here, relative to this directory. templates_path = ["_templates"] # List of patterns, relative to source directory, that match files and # directories to ignore when looking for source files. # This pattern also affects html_static_path and html_extra_path. exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] autodoc_mock_imports = ["numba"] # -- Options for HTML output ------------------------------------------------- # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. # html_theme = "sphinx_book_theme" # Config for the Sphinx book html_baseurl = "https://vector.readthedocs.io/en/latest/" html_logo = "_images/vector-logo.png" html_title = f"Vector {version}" html_theme_options = { "home_page_in_toc": True, "repository_url": "https://github.com/scikit-hep/vector", "use_repository_button": True, "use_issues_button": True, "use_edit_page_button": True, } # Add any paths that contain custom static files (such as style sheets) here, # relative to this directory. They are copied after the builtin static files, # so a file named "default.css" will overwrite the builtin "default.css". html_static_path = ["_static"] html_css_files = ["css/awkward_output.css"] # -- Options for notebooks -------------------------------------------------- highlight_language = "python3" nbsphinx_execute_arguments = [ "--InlineBackend.figure_formats={'png2x'}", "--InlineBackend.rc={'figure.dpi': 96}", ] vector-1.6.3/docs/index.md000066400000000000000000000175031503546127100154160ustar00rootroot00000000000000![](_images/LogoSrc.svg) # Overview [![DOI][zenodo-badge]][zenodo-link] [![DOI][joss-badge]][joss-link] [![Scikit-HEP][sk-badge]][sk-link] [![Actions Status][actions-badge]][actions-link] [![Documentation Status][rtd-badge]][rtd-link] [![pre-commit.ci status][pre-commit-badge]][pre-commit-link] [![codecov percentage][codecov-badge]][codecov-link] [![PyPI platforms][pypi-platforms]][pypi-link] [![PyPI version][pypi-version]][pypi-link] [![Conda latest release][conda-version]][conda-link] [![LICENSE][license-badge]][license-link] [![GitHub Discussion][github-discussions-badge]][github-discussions-link] [![Gitter][gitter-badge]][gitter-link] ## Installation You can install Vector with [pip](https://pypi.org/project/vector/) and [conda](https://anaconda.org/conda-forge/vector). ```bash pip install vector ``` ## Introduction Vector is a Python library for 2D and 3D spatial vectors, as well as 4D space-time vectors. It is especially intended for performing geometric calculations on _arrays of vectors_, rather than one vector at a time in a Python for loop. Vector is part of the [Scikit-HEP project](https://scikit-hep.org/), High Energy Physics (HEP) tools in Python. ### Coordinate systems Vectors may be expressed in any of these coordinate systems: - the azimuthal plane may be Cartesian `x` `y` or polar `rho` ($\rho$) `phi` ($\phi$) - the longitudinal axis may be Cartesian `z`, polar `theta` ($\theta$), or pseudorapidity `eta` ($\eta$) - the temporal component for space-time vectors may be Cartesian `t` or proper time `tau` ($\tau$) in any combination. (That is, 4D vectors have 2×3×2 = 12 distinct coordinate systems.) ![](_images/coordinate-systems.svg) ### Backends Vectors may be included in any of these data types: - [vector.obj](src/make_object.md) objects (pure Python) - [NumPy structured arrays](https://numpy.org/doc/stable/user/basics.rec.html) of vectors - [Awkward Arrays](https://awkward-array.org/) of vectors (possibly within variable-length lists or nested record structures) - [SymPy expressions](https://www.sympy.org/en/index.html) for symbolic (non-numeric) manipulations - In [Numba-compiled functions](https://numba.pydata.org/), with [vector.obj](src/make_object.md) objects or Awkward Arrays Each of these "backends" provides the same suite of properties and methods, through a common "compute" library. ### Geometric versus momentum Finally, vectors come in two flavors: - geometric: only one name for each property or method - momentum: same property or method can be accessed with several synonyms, such as `pt` ($p_T$, transverse momentum) for the azimuthal magnitude `rho` ($\rho$) and `energy` and `mass` for the Cartesian time `t` and proper time `tau` ($\tau$). ### Familiar conventions Names and coordinate conventions were chosen to align with [ROOT](https://root.cern/)'s [TLorentzVector](https://root.cern.ch/doc/master/classTLorentzVector.html) and [Math::LorentzVector](https://root.cern.ch/doc/master/classROOT_1_1Math_1_1LorentzVector.html), as well as [scikit-hep/math](https://github.com/scikit-hep/scikit-hep/tree/master/skhep/math), [uproot-methods TLorentzVector](https://github.com/scikit-hep/uproot3-methods/blob/master/uproot3_methods/classes/TLorentzVector.py), [henryiii/hepvector](https://github.com/henryiii/hepvector), and [coffea.nanoevents.methods.vector](https://coffea-hep.readthedocs.io/en/latest/modules/coffea.nanoevents.methods.vector.html). Vector follows the $(-, -, -, +)$ (`x, y, z, t`) metric convention for Lorentz vectors. The $(-,-,-,+)$ metric convention for Lorentz vectors corresponds to the Minkowski metric: ```{math} g_{\mu\nu} = \text{diag}(-1, -1, -1, +1) ``` For a Lorentz vector $p^\mu = (p_x, p_y, p_z, E)$, the squared norm (or invariant mass squared) is given by: ```{math} p^\mu p_\mu = g_{\mu\nu} p^\mu p^\nu = E^2 - p_x^2 - p_y^2 - p_z^2 ``` This convention is widely used in high-energy physics (HEP), including frameworks such as ROOT. Further, the transformations and rotations (including boosts) are active, and the Euler angle conventions align with the formalisations in the [GenVector package](https://root.cern/topical/GenVector.pdf). More precisely, the implementation of `rotate_euler` uses the matrices defined in the [wikipedia article](https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix). ## Getting help - Source code on GitHub: [scikit-hep/vector](https://github.com/scikit-hep/vector) - Report bugs and request features on the [GitHub Issues page](https://github.com/scikit-hep/vector/issues) - Ask questions on the [GitHub Discussions page](https://github.com/scikit-hep/vector/discussions) - Real-time chat on Gitter: [Scikit-HEP/Vector](https://gitter.im/Scikit-HEP/vector) ## Contributing to Vector If you want to contribute to Vector, [pull requests](https://github.com/scikit-hep/vector/pulls) are welcome! Please install the latest version of the `main` branch from source or a fork: ```bash git clone https://github.com/scikit-hep/vector.git cd vector pip install -e . ``` Refer to [CONTRIBUTING.md](https://github.com/scikit-hep/vector/blob/main/.github/CONTRIBUTING.md) for more. ## Citing Vector To cite Vector, please use [![DOI][joss-badge]][joss-link] ```bib @article{Chopra2025, doi = {10.21105/joss.07791}, url = {https://doi.org/10.21105/joss.07791}, year = {2025}, publisher = {The Open Journal}, volume = {10}, number = {109}, pages = {7791}, author = {Saransh Chopra and Henry Schreiner and Eduardo Rodrigues and Jonas Eschle and Jim Pivarski}, title = {Vector: JIT-compilable mathematical manipulations of ragged Lorentz vectors}, journal = {Journal of Open Source Software} } ``` ## Documentation ```{toctree} :maxdepth: 1 :caption: Tutorials src/object.ipynb src/numpy.ipynb src/awkward.ipynb src/numba.ipynb src/sympy.ipynb ``` ```{toctree} :maxdepth: 1 :caption: Vector constructors src/make_object.md src/make_numpy.md src/make_awkward.md src/make_sympy.md ``` ```{toctree} :maxdepth: 1 :caption: Vector functions src/common.md src/vector2d.md src/vector3d.md src/vector4d.md src/momentum2d.md src/momentum3d.md src/momentum4d.md ``` ```{toctree} :maxdepth: 1 :caption: More ways to learn src/talks.md ``` [actions-badge]: https://github.com/scikit-hep/vector/actions/workflows/ci.yml/badge.svg [actions-link]: https://github.com/scikit-hep/vector/actions [codecov-badge]: https://codecov.io/gh/scikit-hep/vector/branch/main/graph/badge.svg?token=YBv60ueORQ [codecov-link]: https://codecov.io/gh/scikit-hep/vector [conda-version]: https://img.shields.io/conda/vn/conda-forge/vector.svg [conda-link]: https://github.com/conda-forge/vector-feedstock [github-discussions-badge]: https://img.shields.io/static/v1?label=Discussions&message=Ask&color=blue&logo=github [github-discussions-link]: https://github.com/scikit-hep/vector/discussions [gitter-badge]: https://badges.gitter.im/Scikit-HEP/vector.svg [gitter-link]: https://gitter.im/Scikit-HEP/vector?utm_source=badge&utm_medium=badge&utm_campaign=pr-badge&utm_content=badge [joss-badge]: https://joss.theoj.org/papers/10.21105/joss.07791/status.svg [joss-link]: https://doi.org/10.21105/joss.07791 [license-badge]: https://img.shields.io/badge/License-BSD_3--Clause-blue.svg [license-link]: https://opensource.org/licenses/BSD-3-Clause [pre-commit-badge]: https://results.pre-commit.ci/badge/github/scikit-hep/vector/main.svg [pre-commit-link]: https://results.pre-commit.ci/repo/github/scikit-hep/vector [pypi-link]: https://pypi.org/project/vector/ [pypi-platforms]: https://img.shields.io/pypi/pyversions/vector [pypi-version]: https://badge.fury.io/py/vector.svg [rtd-badge]: https://readthedocs.org/projects/vector/badge/?version=latest [rtd-link]: https://vector.readthedocs.io/en/latest/?badge=latest [sk-badge]: https://scikit-hep.org/assets/images/Scikit--HEP-Project-blue.svg [sk-link]: https://scikit-hep.org/ [zenodo-badge]: https://zenodo.org/badge/DOI/10.5281/zenodo.15263860.svg [zenodo-link]: https://zenodo.org/records/15263860 vector-1.6.3/docs/make.bat000066400000000000000000000014331503546127100153650ustar00rootroot00000000000000@ECHO OFF pushd %~dp0 REM Command file for Sphinx documentation if "%SPHINXBUILD%" == "" ( set SPHINXBUILD=sphinx-build ) set SOURCEDIR=. set BUILDDIR=_build if "%1" == "" goto help %SPHINXBUILD% >NUL 2>NUL if errorlevel 9009 ( echo. echo.The 'sphinx-build' command was not found. Make sure you have Sphinx echo.installed, then set the SPHINXBUILD environment variable to point echo.to the full path of the 'sphinx-build' executable. Alternatively you echo.may add the Sphinx directory to PATH. echo. echo.If you don't have Sphinx installed, grab it from echo.http://sphinx-doc.org/ exit /b 1 ) %SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% goto end :help %SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% :end popd vector-1.6.3/docs/src/000077500000000000000000000000001503546127100145465ustar00rootroot00000000000000vector-1.6.3/docs/src/awkward.ipynb000066400000000000000000000774011503546127100172620ustar00rootroot00000000000000{ "cells": [ { "cell_type": "markdown", "id": "13ca882a-504d-409f-b6df-ac91f4e8f803", "metadata": {}, "source": [ "# Awkward Arrays of vectors" ] }, { "cell_type": "markdown", "id": "26894013-8d55-45b5-9592-6a82c6e7440a", "metadata": {}, "source": [ "First, [install](../index.md#installation) and import Vector and [Awkward Array](https://awkward-array.org/)." ] }, { "cell_type": "code", "execution_count": 1, "id": "8f504534-1fa5-4dfb-b0c9-f7a2b2d12eae", "metadata": {}, "outputs": [], "source": [ "import vector\n", "import awkward as ak" ] }, { "cell_type": "markdown", "id": "b119a69a-19e6-4568-bdde-fca92ec23f2f", "metadata": {}, "source": [ "## Making an Awkward Array of vectors" ] }, { "cell_type": "markdown", "id": "bb1e0f46-627e-463a-840b-bac3bf0588eb", "metadata": {}, "source": [ "Awkward Arrays are arrays with more complex data structures than NumPy allows, such as variable-length lists, nested records, missing and even heterogeneous data (different data types in the same array).\n", "\n", "Vectors can be included among those data structures. In this context, vectors are Awkward \"records,\" objects with named fields, that can be nested inside of other structures. The vector properties and methods are implemented through Awkward Array's [behavior](https://awkward-array.org/doc/main/reference/ak.behavior.html) mechanism. Unlike [vector objects](object.md) and [NumPy subclasses](numpy.md), the vectors can't be ordinary Python classes because they might be nested within other data structures, such as variable-length lists, and these lists are implemented in a columnar way that isn't open to Python's introspection.\n", "\n", "Let's start with an example. Below, we create an Awkward Array using its [ak.Array](https://awkward-array.org/doc/main/reference/generated/ak.Array.html) constructor, but include `with_name` and `behavior` arguments:" ] }, { "cell_type": "code", "execution_count": 2, "id": "42b68d0a-d867-4fed-84ad-59d4732e112b", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{x: 1.1, y: 2.2}, {x: 3.3, y: 4.4}],\n",
       " [],\n",
       " [{x: 5.5, y: 6.6}]]\n",
       "----------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 80 B\n",
       "type: 3 * var * Vector2D[\n",
       "    x: float64,\n",
       "    y: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "arr = ak.Array(\n", " [\n", " [{\"x\": 1.1, \"y\": 2.2}, {\"x\": 3.3, \"y\": 4.4}],\n", " [],\n", " [{\"x\": 5.5, \"y\": 6.6}],\n", " ],\n", " with_name=\"Vector2D\",\n", " behavior=vector.backends.awkward.behavior,\n", ")\n", "arr" ] }, { "cell_type": "markdown", "id": "e674ffdd-a42e-4707-911d-def1a0010858", "metadata": {}, "source": [ "The above array contains 3 lists, the first has length 2, the second has length 0, and the third has length 1. The lists contain records with field names `\"x\"` and `\"y\"`, and the record type is named `\"Vector2D\"`. In addition, this array has `behavior` from `vector.backends.awkward.behavior`, which is a large dict containing classes and functions to implement vector operations.\n", "\n", "For instance, we can compute `rho` and `phi` coordinates in the same way as with the [NumPy subclasses](numpy.md), an array at a time:" ] }, { "cell_type": "code", "execution_count": 3, "id": "1d5484f9-07e1-4192-bf88-3b32f1028bf5", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[2.46, 5.5],\n",
       " [],\n",
       " [8.59]]\n",
       "-----------------------\n",
       "backend: cpu\n",
       "nbytes: 56 B\n",
       "type: 3 * var * float64
" ], "text/plain": [ "" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "arr.rho" ] }, { "cell_type": "code", "execution_count": 4, "id": "61e61f29-015d-4f53-94cd-07444811c9ec", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[1.11, 0.927],\n",
       " [],\n",
       " [0.876]]\n",
       "-----------------------\n",
       "backend: cpu\n",
       "nbytes: 56 B\n",
       "type: 3 * var * float64
" ], "text/plain": [ "" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "arr.phi" ] }, { "cell_type": "markdown", "id": "66d1c600-568a-41cc-b31a-f7eaedef1c68", "metadata": {}, "source": [ "As with NumPy, performing operations an array at a time is usually much faster than writing Python for loops. What Awkward Array provides on top of that is the ability to do these operations _through_ variable-length lists and other structures.\n", "\n", "An Awkward Array needs all of the following for its records to be interpreted as vectors:\n", "\n", "1. the record name, which can be assigned using [ak.with_name](https://awkward-array.org/doc/main/reference/generated/ak.with_name.html) or as a constructor argument, must be one of `\"Vector2D\"`, `\"Momentum2D\"`, `\"Vector3D\"`, `\"Momentum3D\"`, `\"Vector4D\"`, and `\"Momentum4D\"`\n", "2. the field names must be recognized coordinate names, following the same conventions as [vector objects](object.md)\n", "3. the array must have `vector.backends.awkward.behavior` as its `behavior`.\n", "\n", "When Awkward Arrays are saved in files, such as with [ak.to_parquet](https://awkward-array.org/doc/main/reference/generated/ak.to_parquet.html), they retain their record names and field names, so conditions 1 and 2 above are persistent. They don't preserve condition 3, the behaviors, since these are Python classes and functions.\n", "\n", "To make sure that Vector behaviors are always available, you can call [vector.register_awkward](make_awkward.md#vector.register_awkward) at the beginning of every script, like this:\n", "\n", "```python\n", "import awkward as ak\n", "import vector\n", "vector.register_awkward()\n", "```\n", "\n", "This function copies Vector's behaviors into Awkward's global [ak.behavior](https://awkward-array.org/doc/main/reference/ak.behavior.html) so that any array with the right record and field names (such as one read from a file) automatically have Vector behaviors.\n", "\n", "Vector also has a [vector.Array](make_awkward.md#vector.Array) constructor, which works like [ak.Array](https://awkward-array.org/doc/main/reference/generated/ak.Array.html) but sets `with_name` automatically, as well as [vector.zip](make_awkward.md#vector.zip), which works like [ak.zip](https://awkward-array.org/doc/main/reference/generated/ak.zip.html) and sets `with_name` automatically. However, these functions still require you to set field names appropriately and if you need to do something complex, it's easier to use Awkward Array's own functions and assign the record name after the array is built, using [ak.with_name](https://awkward-array.org/doc/main/reference/generated/ak.with_name.html)." ] }, { "cell_type": "markdown", "id": "467f6d6d-dd6a-4c6e-8ff2-cc7b61e6b949", "metadata": {}, "source": [ "## Using an Awkward array of vectors" ] }, { "cell_type": "markdown", "id": "8677dc3e-1482-42c4-9157-75c0cea61f45", "metadata": {}, "source": [ "First, let's make some arrays to use in examples:" ] }, { "cell_type": "code", "execution_count": 5, "id": "025fcf28-c946-4c52-9c06-333a217ead52", "metadata": {}, "outputs": [], "source": [ "import numpy as np\n", "import awkward as ak\n", "import vector\n", "\n", "vector.register_awkward()" ] }, { "cell_type": "code", "execution_count": 6, "id": "15050418-eab0-4da4-a1bd-c8ceef2a47c2", "metadata": {}, "outputs": [], "source": [ "def array_of_momentum3d(num_vectors):\n", " return ak.zip(\n", " {\n", " \"px\": np.random.normal(0, 1, num_vectors),\n", " \"py\": np.random.normal(0, 1, num_vectors),\n", " \"pz\": np.random.normal(0, 1, num_vectors),\n", " },\n", " with_name=\"Momentum3D\",\n", " )\n", "\n", "\n", "def array_of_lists_of_momentum3d(mean_num_per_list, num_lists):\n", " num_per_list = np.random.poisson(mean_num_per_list, num_lists)\n", " return ak.unflatten(\n", " array_of_momentum3d(np.sum(num_per_list)),\n", " num_per_list,\n", " )\n", "\n", "\n", "a = array_of_momentum3d(10)\n", "b = array_of_lists_of_momentum3d(1.5, 10)" ] }, { "cell_type": "code", "execution_count": 7, "id": "43fbd738-8243-45bf-af35-2b09ef196bff", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[{px: -0.541, py: 0.714, pz: -0.983},\n",
       " {px: 0.61, py: -0.48, pz: -0.0845},\n",
       " {px: 0.929, py: 0.718, pz: 0.0935},\n",
       " {px: 1.52, py: -1.39, pz: -0.0187},\n",
       " {px: 0.675, py: 0.0216, pz: -1.12},\n",
       " {px: 0.542, py: -0.524, pz: 0.586},\n",
       " {px: 0.927, py: 0.476, pz: 0.602},\n",
       " {px: 0.615, py: 1.23, pz: -0.59},\n",
       " {px: -0.626, py: 0.072, pz: 0.136},\n",
       " {px: 2.23, py: 1.05, pz: -0.355}]\n",
       "--------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 240 B\n",
       "type: 10 * Momentum3D[\n",
       "    px: float64,\n",
       "    py: float64,\n",
       "    pz: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a" ] }, { "cell_type": "code", "execution_count": 8, "id": "c347edb1-78b3-4963-a503-b877597c2ee5", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{px: 0.73, py: 0.862, pz: -0.129}, {px: 0.746, py: 0.614, pz: ..., ...}],\n",
       " [],\n",
       " [{px: 0.493, py: -0.633, pz: -0.413}, {px: 0.836, py: -1.77, ...}],\n",
       " [{px: -0.58, py: -0.586, pz: -0.513}, {...}, ..., {px: -0.201, py: 1.52, ...}],\n",
       " [{px: -0.474, py: 1.6, pz: -0.1}, {px: 0.218, py: 0.155, pz: 0.834}],\n",
       " [{px: -0.881, py: -0.183, pz: -0.725}, {px: -0.108, py: -0.889, ...}],\n",
       " [{px: -1.77, py: 2.04, pz: 0.896}],\n",
       " [{px: -1.08, py: -0.591, pz: -0.0674}, {px: -0.183, py: -0.289, ...}],\n",
       " [{px: -1.08, py: -1.15, pz: 2.6}],\n",
       " [{px: -0.419, py: -0.363, pz: -0.563}, {px: -2.48, py: -0.228, ...}]]\n",
       "--------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 544 B\n",
       "type: 10 * var * Momentum3D[\n",
       "    px: float64,\n",
       "    py: float64,\n",
       "    pz: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "b" ] }, { "cell_type": "markdown", "id": "7a6e73c4-cc2c-4cb5-9807-13c1bc970f78", "metadata": {}, "source": [ "Awkward Array uses array-at-a-time functions like NumPy, so if we want to compute dot products of each vector in `a` with every vector of each list in `b`, we'd say:" ] }, { "cell_type": "code", "execution_count": 9, "id": "144ee86e-2b47-4526-9d6b-3d490be69ba8", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[0.347, -1.32],\n",
       " [],\n",
       " [-0.035, -0.385],\n",
       " [-0.0551, 0.194, 2.28, -2.86, -2.39],\n",
       " [-0.173, -0.782],\n",
       " [-0.807, 0.704],\n",
       " [-0.128],\n",
       " [-1.35, -0.446],\n",
       " [0.945],\n",
       " [-1.11, -5.58]]\n",
       "--------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 240 B\n",
       "type: 10 * var * float64
" ], "text/plain": [ "" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a.dot(b)" ] }, { "cell_type": "markdown", "id": "1088c2a6-e744-4308-bd74-4ed9471d3bd1", "metadata": {}, "source": [ "Note that `a` and `b` have different numbers of vectors, but the same array lengths. The operation above [broadcasts](https://awkward-array.org/doc/main/user-guide/how-to-math-broadcasting.html) array `a` into `b`, like the following code:" ] }, { "cell_type": "code", "execution_count": 10, "id": "9edab13c-5b8e-47f2-afac-c0dd7241a8ca", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "[0.3470024410105361 -1.32262374856204 ]\n", "[]\n", "[-0.0350058354646221 -0.3845727581312185 ]\n", "[-0.05513156034267678 0.1939866533163302 2.2806023784977056 -2.861051860111839 -2.3860250251022475 ]\n", "[-0.17338156209593328 -0.7816799910864897 ]\n", "[-0.8071770038569903 0.7044860439866077 ]\n", "[-0.1279745100114199 ]\n", "[-1.3483450919978617 -0.44626327125953613 ]\n", "[0.9449043237160631 ]\n", "[-1.1124522691465186 -5.579496184145224 ]\n" ] } ], "source": [ "for i in range(len(a)):\n", " print(\"[\", end=\"\")\n", "\n", " for j in range(len(b[i])):\n", " out = a[i].dot(b[i, j])\n", "\n", " print(out, end=\" \")\n", "\n", " print(\"]\")" ] }, { "cell_type": "markdown", "id": "0e17784d-a838-4c87-a708-f530bd0fbe59", "metadata": {}, "source": [ "Like NumPy, the array-at-a-time expression is more concise and faster:" ] }, { "cell_type": "code", "execution_count": 11, "id": "1e3ffd74-469d-4ccd-984d-706413442700", "metadata": {}, "outputs": [], "source": [ "a = array_of_momentum3d(10000)\n", "b = array_of_lists_of_momentum3d(1.5, 10000)" ] }, { "cell_type": "code", "execution_count": 12, "id": "135e44f1-7e25-49ef-8805-8638b3d0e9be", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "9.08 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)\n" ] } ], "source": [ "%%timeit -n1 -r1\n", "\n", "out = np.zeros(10000)\n", "\n", "for i in range(len(a)):\n", " for j in range(len(b[i])):\n", " out[i] += a[i].dot(b[i, j])" ] }, { "cell_type": "code", "execution_count": 13, "id": "6b4d2c04-9034-4558-905a-af4539be1d40", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "2.44 ms ± 20.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n" ] } ], "source": [ "%%timeit\n", "\n", "out = np.sum(a.dot(b), axis=1)" ] }, { "cell_type": "markdown", "id": "cd509b13-8e87-40b0-b478-bb40154dd5c1", "metadata": {}, "source": [ "(Note the units.)\n", "\n", "Just as with NumPy, all of the coordinate transformations and vector operations are implemented for Awkward Arrays of vectors." ] }, { "cell_type": "markdown", "id": "716cd51c-88ec-42e7-89ae-719cef7e4368", "metadata": {}, "source": [ "## Some troubleshooting hints" ] }, { "cell_type": "markdown", "id": "904561c9-b998-4e4a-8ce1-e799931d9285", "metadata": {}, "source": [ "Make sure that the Vector behaviors are actually installed and applied to your data. In the data type, the record type should appear as `\"Vector2D\"`, `\"Momentum2D\"`, `\"Vector3D\"`, `\"Momentum3D\"`, `\"Vector4D\"`, or `\"Momentum4D\"`, rather than the generic curly brackets `{` and `}`, and if you extract one record from the array, can you perform a vector operation on it?\n", "\n", "Make sure that your arrays broadcast the way that you want them to. If the vector behaviors are clouding the picture, make simpler arrays with numbers in place of records. Can you add them with `+`? (Addition uses the same broadcasting rules as all other operations.)\n", "\n", "If your code runs but doesn't give the results you expect, try slicing the arrays to just the first two items with `arr[:2]`. Step through the calculation on just two elements, observing the results of each operation. Are they what you expect?" ] }, { "cell_type": "markdown", "id": "60db1ccb-8d0a-4bc8-9050-6f68b03e2ee4", "metadata": {}, "source": [ "## Advanced: subclassing Awkward-Vector behaviors" ] }, { "cell_type": "markdown", "id": "00d4f231-46f3-41a0-8573-9564aad8b29b", "metadata": {}, "source": [ "It is possible to write subclasses for Awkward-Vector behaviors as mixins to extend the vector functionalities. For instance, the `MomentumAwkward` classes can be extended in the following way:" ] }, { "cell_type": "code", "execution_count": 14, "id": "967444cc-a4f1-4e2b-8289-3224333218f1", "metadata": {}, "outputs": [], "source": [ "behavior = vector.backends.awkward.behavior\n", "\n", "\n", "@ak.mixin_class(behavior)\n", "class TwoVector(vector.backends.awkward.MomentumAwkward2D):\n", " pass\n", "\n", "\n", "@ak.mixin_class(behavior)\n", "class ThreeVector(vector.backends.awkward.MomentumAwkward3D):\n", " pass\n", "\n", "\n", "# required for transforming vectors\n", "# the class names must always end with \"Array\"\n", "TwoVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821\n", "TwoVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821\n", "TwoVectorArray.MomentumClass = TwoVectorArray # noqa: F821\n", "\n", "ThreeVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821\n", "ThreeVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821\n", "ThreeVectorArray.MomentumClass = ThreeVectorArray # noqa: F821" ] }, { "cell_type": "code", "execution_count": 15, "id": "42ea1b53-6003-4294-ac52-3e014eab92b7", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{pt: 1, phi: 1.2}, {pt: 2, phi: 1.4}],\n",
       " [],\n",
       " [{pt: 3, phi: 1.6}],\n",
       " [{pt: 4, phi: 3.4}]]\n",
       "------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 104 B\n",
       "type: 4 * var * TwoVector[\n",
       "    pt: int64,\n",
       "    phi: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec = ak.zip(\n", " {\n", " \"pt\": [[1, 2], [], [3], [4]],\n", " \"phi\": [[1.2, 1.4], [], [1.6], [3.4]],\n", " },\n", " with_name=\"TwoVector\",\n", " behavior=behavior,\n", ")\n", "vec" ] }, { "cell_type": "markdown", "id": "c2889f3f-5fb3-42ea-bd5d-8c0135aa7c81", "metadata": {}, "source": [ "The binary operators are not automatically registered by Awkward, but Vector methods can be used to perform operations on subclassed vectors." ] }, { "cell_type": "code", "execution_count": 16, "id": "b4ecc08c-253b-46b7-a78a-1aca4558d863", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 2, phi: 1.2}, {rho: 4, phi: 1.4}],\n",
       " [],\n",
       " [{rho: 6, phi: 1.6}],\n",
       " [{rho: 8, phi: -2.88}]]\n",
       "---------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 104 B\n",
       "type: 4 * var * TwoVector[\n",
       "    rho: float64,\n",
       "    phi: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec.add(vec)" ] }, { "cell_type": "markdown", "id": "41426373-4492-4ff6-9496-3e5d88ef630e", "metadata": {}, "source": [ "Similarly, other vector methods can be used by the new methods internally." ] }, { "cell_type": "code", "execution_count": 17, "id": "2c05135c-ba70-4b24-9d9f-2b04f72bf1ce", "metadata": {}, "outputs": [], "source": [ "import numbers" ] }, { "cell_type": "code", "execution_count": 18, "id": "8e24d10e-a1ab-4614-9e8e-62a6378c0864", "metadata": {}, "outputs": [], "source": [ "@ak.mixin_class(behavior)\n", "class LorentzVector(vector.backends.awkward.MomentumAwkward4D):\n", " @ak.mixin_class_method(np.divide, {numbers.Number})\n", " def divide(self, factor):\n", " return self.scale(1 / factor)\n", "\n", "\n", "# required for transforming vectors\n", "# the class names must always end with \"Array\"\n", "LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821\n", "LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821\n", "LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821\n", "LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821" ] }, { "cell_type": "code", "execution_count": 19, "id": "203112e0-217e-4dc5-8145-881d48c9cf08", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{pt: 1, eta: 1.2, phi: 0.3, energy: 50}, {pt: 2, eta: 1.4, ...}],\n",
       " [],\n",
       " [{pt: 3, eta: 1.6, phi: 0.5, energy: 52}],\n",
       " [{pt: 4, eta: 3.4, phi: 0.6, energy: 60}]]\n",
       "-----------------------------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 168 B\n",
       "type: 4 * var * LorentzVector[\n",
       "    pt: int64,\n",
       "    eta: float64,\n",
       "    phi: float64,\n",
       "    energy: int64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec = ak.zip(\n", " {\n", " \"pt\": [[1, 2], [], [3], [4]],\n", " \"eta\": [[1.2, 1.4], [], [1.6], [3.4]],\n", " \"phi\": [[0.3, 0.4], [], [0.5], [0.6]],\n", " \"energy\": [[50, 51], [], [52], [60]],\n", " },\n", " with_name=\"LorentzVector\",\n", " behavior=behavior,\n", ")\n", "vec" ] }, { "cell_type": "code", "execution_count": 20, "id": "6bf279a8-fca7-4c16-aca8-90565818a646", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 0.5, phi: 0.3, eta: 1.2, t: 25}, {rho: 1, phi: 0.4, ...}],\n",
       " [],\n",
       " [{rho: 1.5, phi: 0.5, eta: 1.6, t: 26}],\n",
       " [{rho: 2, phi: 0.6, eta: 3.4, t: 30}]]\n",
       "-----------------------------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 168 B\n",
       "type: 4 * var * LorentzVector[\n",
       "    rho: float64,\n",
       "    phi: float64,\n",
       "    eta: float64,\n",
       "    t: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec / 2" ] }, { "cell_type": "code", "execution_count": 21, "id": "6648d398-60e4-4a7d-bd88-a7dd4a35e8d0", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 1, phi: 0.3}, {rho: 2, phi: 0.4}],\n",
       " [],\n",
       " [{rho: 3, phi: 0.5}],\n",
       " [{rho: 4, phi: 0.6}]]\n",
       "-------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 104 B\n",
       "type: 4 * var * TwoVector[\n",
       "    rho: int64,\n",
       "    phi: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec.like(vector.obj(x=1, y=2))" ] }, { "cell_type": "code", "execution_count": 22, "id": "09854ba6-f571-42f8-95a9-b0ac52553cda", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 1, phi: 0.3, eta: 1.2}, {rho: 2, phi: 0.4, eta: 1.4}],\n",
       " [],\n",
       " [{rho: 3, phi: 0.5, eta: 1.6}],\n",
       " [{rho: 4, phi: 0.6, eta: 3.4}]]\n",
       "---------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 136 B\n",
       "type: 4 * var * ThreeVector[\n",
       "    rho: int64,\n",
       "    phi: float64,\n",
       "    eta: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 22, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec.like(vector.obj(x=1, y=2, z=3))" ] }, { "cell_type": "markdown", "id": "59cf3641-e227-4f0b-a1fa-99b0d196f02f", "metadata": {}, "source": [ "It is also possible to manually add binary operations in vector's behavior dict to enable binary operations." ] }, { "cell_type": "code", "execution_count": 23, "id": "11778dd0-4726-43cc-817b-3326528dd144", "metadata": {}, "outputs": [], "source": [ "_binary_dispatch_cls = {\n", " \"TwoVector\": TwoVector,\n", " \"ThreeVector\": ThreeVector,\n", " \"LorentzVector\": LorentzVector,\n", "}\n", "_rank = [TwoVector, ThreeVector, LorentzVector]\n", "\n", "for lhs, lhs_to in _binary_dispatch_cls.items():\n", " for rhs, rhs_to in _binary_dispatch_cls.items():\n", " out_to = min(lhs_to, rhs_to, key=_rank.index)\n", " behavior[(np.add, lhs, rhs)] = out_to.add\n", " behavior[(np.subtract, lhs, rhs)] = out_to.subtract" ] }, { "cell_type": "code", "execution_count": 24, "id": "3088485c-70e4-461d-86b3-5db81148389f", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 2, phi: 0.3, eta: 1.2, t: 100}, {rho: 4, phi: 0.4, eta: 1.4, ...}],\n",
       " [],\n",
       " [{rho: 6, phi: 0.5, eta: 1.6, t: 104}],\n",
       " [{rho: 8, phi: 0.6, eta: 3.4, t: 120}]]\n",
       "---------------------------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 168 B\n",
       "type: 4 * var * LorentzVector[\n",
       "    rho: float64,\n",
       "    phi: float64,\n",
       "    eta: float64,\n",
       "    t: int64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec + vec" ] }, { "cell_type": "code", "execution_count": 25, "id": "040e47cc-013d-423e-b5df-b0982a2279aa", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 2, phi: 0.3}, {rho: 4, phi: 0.4}],\n",
       " [],\n",
       " [{rho: 6, phi: 0.5}],\n",
       " [{rho: 8, phi: 0.6}]]\n",
       "---------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 104 B\n",
       "type: 4 * var * TwoVector[\n",
       "    rho: float64,\n",
       "    phi: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 25, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec.to_2D() + vec.to_2D()" ] }, { "cell_type": "markdown", "id": "e2259fbe-3eea-4fde-bedb-b5631f639ca3", "metadata": {}, "source": [ "Finally, instead of manually registering the superclass ufuncs, one can use the utility `copy_behaviors` function to copy behavior items for a new subclass -" ] }, { "cell_type": "code", "execution_count": 26, "id": "4af071da-8612-49a2-9a56-7e49487cf866", "metadata": {}, "outputs": [], "source": [ "behavior.update(ak._util.copy_behaviors(\"Vector2D\", \"TwoVector\", behavior))\n", "behavior.update(ak._util.copy_behaviors(\"Vector3D\", \"ThreeVector\", behavior))\n", "behavior.update(ak._util.copy_behaviors(\"Momentum4D\", \"LorentzVector\", behavior))" ] }, { "cell_type": "code", "execution_count": 27, "id": "1c434751-a932-4694-92fc-9e48b51905b0", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 2, phi: 0.3, eta: 1.2, t: 100}, {rho: 4, phi: 0.4, eta: 1.4, ...}],\n",
       " [],\n",
       " [{rho: 6, phi: 0.5, eta: 1.6, t: 104}],\n",
       " [{rho: 8, phi: 0.6, eta: 3.4, t: 120}]]\n",
       "------------------------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 168 B\n",
       "type: 4 * var * Momentum4D[\n",
       "    rho: float64,\n",
       "    phi: float64,\n",
       "    eta: float64,\n",
       "    t: int64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 27, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec + vec" ] }, { "cell_type": "code", "execution_count": 28, "id": "09e2bc41-340d-4ad6-82d8-f08cf869644a", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{rho: 2, phi: 0.3}, {rho: 4, phi: 0.4}],\n",
       " [],\n",
       " [{rho: 6, phi: 0.5}],\n",
       " [{rho: 8, phi: 0.6}]]\n",
       "--------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 104 B\n",
       "type: 4 * var * Vector2D[\n",
       "    rho: float64,\n",
       "    phi: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 28, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vec.to_2D() + vec.to_2D()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.10" } }, "nbformat": 4, "nbformat_minor": 5 } vector-1.6.3/docs/src/common.md000066400000000000000000000013151503546127100163600ustar00rootroot00000000000000# Interface for all vectors Vectors of all backends, 2D/3D/4D, geometric and momentum, have the following attributes, properties, and methods. For interfaces specialized to 2D/3D/4D vectors or momentum vectors, see - [Interface for 2D vectors](vector2d.md) - [Interface for 3D vectors](vector3d.md) - [Interface for 4D vectors](vector4d.md) - [Interface for 2D momentum](momentum2d.md) - [Interface for 3D momentum](momentum3d.md) - [Interface for 4D momentum](momentum4d.md) ```{eval-rst} .. autoclass:: vector._methods.VectorProtocol :members: :inherited-members: :member-order: bysource :exclude-members: lib,ProjectionClass2D,ProjectionClass3D,ProjectionClass4D,GenericClass,MomentumClass ``` vector-1.6.3/docs/src/make_awkward.md000066400000000000000000000047241503546127100175340ustar00rootroot00000000000000# Making Awkward Arrays of vectors An Awkward Array of vectors is an Awkward Array containing appropriately named records, appropriately named fields, and the Vector [behaviors](https://awkward-array.org/doc/main/reference/ak.behavior.html) registered in the array. Here's a complete example for illustration: ```python >>> import awkward as ak >>> import vector >>> vector.register_awkward() >>> >>> vec = ak.Array([ ... [{"x": 1.1, "y": 2.2}, {"x": 3.3, "y": 4.4}], ... [], ... [{"x": 5.5, "y": 6.6}], ... ], with_name="Vector2D") >>> >>> abs(vec) ``` In the above, 1. `vector.register_awkward()` loads Vector's `vector.backends.awkward.behavior` dict of functionality into the global `ak.behavior` 2. the Awkward Array contains records (inside variable-length lists) with field names `"x"` and `"y"` 3. those records are labeled with type name `"Vector2D"` and thus the `abs` function computes the magnitude of each record as `sqrt(x**2 + y**2)`, through the variable-length lists. It is not necessary to install Vector's behaviors globally. They could be installed in the `vec` array only by passing `behavior=vector.backends.awkward.behavior` to the [ak.Array](https://awkward-array.org/doc/main/reference/generated/ak.Array.html) constructor. The records can contain more fields than those that specify coordinates, which can be useful for specifying properties of a particle other than its momentum. Only the coordinate names are considered when performing vector calculations. Coordinates must be numbers (not, for instance, lists of numbers). Be careful about field names that coincide with coordinates, such as `rho` (azimuthal magnitude) and `tau` (proper time). The `vector.Array` function (`vector.awk` is a synonym) is an alternative to the [ak.Array](https://awkward-array.org/doc/main/reference/generated/ak.Array.html) constructor, which installs Vector's behavior in the new array (not globally in `ak.behavior`). The `vector.zip` function is an alternative to the [ak.zip](https://awkward-array.org/doc/main/reference/generated/ak.zip.html) function, which installs Vector's behavior in the new array (not globally in `ak.behavior`). Awkward Arrays can be used in [Numba-compiled functions](https://numba.pydata.org/), including those that contain vectors. ```{eval-rst} .. autofunction:: vector.register_awkward ``` ```{eval-rst} .. autofunction:: vector.Array ``` ```{eval-rst} .. autofunction:: vector.zip ``` vector-1.6.3/docs/src/make_numpy.md000066400000000000000000000030401503546127100172320ustar00rootroot00000000000000# Making NumPy arrays of vectors A NumPy array of vectors is a subclass of [np.ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html) with vector properties and methods. The [dtype](https://numpy.org/doc/stable/reference/arrays.dtypes.html) of this array is [structured](https://numpy.org/doc/stable/user/basics.rec.html) to specify the coordinate names; an array with fields `x` and `y` (Cartesian) performs computations differently from an array with fields `rho` and `phi` (polar). To create a NumPy array of vectors, 1. use the `vector.array` function (`vector.arr` is a synonym) 2. use the `vector.VectorNumpy` class constructor 3. or cast a structured NumPy array as the appropriate class, which can avoid copying data. ## General constructor ```{eval-rst} .. autofunction:: vector.array ``` ```{eval-rst} .. autoclass:: vector.VectorNumpy ``` ## Casting structured arrays [NumPy structured arrays](https://numpy.org/doc/stable/user/basics.rec.html) with appropriately named fields (see above) can be _cast_ as arrays of vectors using [np.ndarray.view](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.view.html). Use the NumPy array subclass with the appropriate dimension below. ```{eval-rst} .. autoclass:: vector.VectorNumpy2D ``` ```{eval-rst} .. autoclass:: vector.MomentumNumpy2D ``` ```{eval-rst} .. autoclass:: vector.VectorNumpy3D ``` ```{eval-rst} .. autoclass:: vector.MomentumNumpy3D ``` ```{eval-rst} .. autoclass:: vector.VectorNumpy4D ``` ```{eval-rst} .. autoclass:: vector.MomentumNumpy4D ``` vector-1.6.3/docs/src/make_object.md000066400000000000000000000030111503546127100173260ustar00rootroot00000000000000# Making vector objects A vector object represents a single vector, rather than an array of vectors. Lists of vector objects are slower to compute and have more memory overhead than arrays of vectors, _unless_ those computations are performed in [Numba-compiled functions](https://numba.pydata.org/). To create a vector object, use the `vector.obj` function with appropriate arguments for 2D/3D/4D and geometric versus momentum. ## General constructor ```{eval-rst} .. autofunction:: vector.obj ``` ## 2D constructors ```{eval-rst} .. autoclass:: vector.VectorObject2D :members: from_rhophi,from_xy ``` ```{eval-rst} .. autoclass:: vector.MomentumObject2D :members: from_rhophi,from_xy ``` ## 3D constructors ```{eval-rst} .. autoclass:: vector.VectorObject3D :members: from_rhophieta,from_rhophitheta,from_rhophiz,from_xyeta,from_xytheta,from_xyz ``` ```{eval-rst} .. autoclass:: vector.MomentumObject3D :members: from_rhophieta,from_rhophitheta,from_rhophiz,from_xyeta,from_xytheta,from_xyz ``` ## 4D constructors ```{eval-rst} .. autoclass:: vector.VectorObject4D :members: from_rhophietat,from_rhophietatau,from_rhophithetat,from_rhophithetatau,from_rhophizt,from_rhophiztau,from_xyetat,from_xyetatau,from_xythetat,from_xythetatau,from_xyzt,from_xyztau ``` ```{eval-rst} .. autoclass:: vector.MomentumObject4D :members: from_rhophietat,from_rhophietatau,from_rhophithetat,from_rhophithetatau,from_rhophizt,from_rhophiztau,from_xyetat,from_xyetatau,from_xythetat,from_xythetatau,from_xyzt,from_xyztau ``` vector-1.6.3/docs/src/make_sympy.md000066400000000000000000000022451503546127100172510ustar00rootroot00000000000000# Making SymPy vector expressions SymPy expressions are not numerical, they're purely algebraic. However, the same Vector computations can be performed on them. To construct a symbolic vector, first create symbols for its components and ensure that they are real-valued (not complex), ```python >>> import sympy >>> x, y, z, t, px, py, pz, eta, tau = sympy.symbols( ... "x y z t px py pz eta tau", real=True ... ) ``` then use one of Vector's SymPy constructors (geometric or momentum), ```python >>> vector.VectorSympy2D(x=x, y=y) VectorSympy2D(x=x, y=y) >>> >>> vector.MomentumSympy3D(px=px, py=py, pz=pz) MomentumSympy3D(px=px, py=py, pz=pz) >>> >>> vector.VectorSympy4D(x=x, y=y, eta=eta, tau=tau) vector.VectorSympy4D(x=x, y=y, eta=eta, tau=tau) ``` which are documented below. ## 2D constructors ```{eval-rst} .. autoclass:: vector.VectorSympy2D ``` ```{eval-rst} .. autoclass:: vector.MomentumSympy2D ``` ## 3D constructors ```{eval-rst} .. autoclass:: vector.VectorSympy3D ``` ```{eval-rst} .. autoclass:: vector.MomentumSympy3D ``` ## 4D constructors ```{eval-rst} .. autoclass:: vector.VectorSympy4D ``` ```{eval-rst} .. autoclass:: vector.MomentumSympy4D ``` vector-1.6.3/docs/src/momentum2d.md000066400000000000000000000005611503546127100171610ustar00rootroot00000000000000# Interface for 2D momentum 2D momentum vectors of all backends have the following attributes, properties, and methods. This includes momentum synonyms. For purely geometric vectors, see [Interface for 2D vectors](vector2d.md). ```{eval-rst} .. autoclass:: vector._methods.MomentumProtocolPlanar :members: :inherited-members: :member-order: bysource ``` vector-1.6.3/docs/src/momentum3d.md000066400000000000000000000005621503546127100171630ustar00rootroot00000000000000# Interface for 3D momentum 3D momentum vectors of all backends have the following attributes, properties, and methods. This includes momentum synonyms. For purely geometric vectors, see [Interface for 3D vectors](vector3d.md). ```{eval-rst} .. autoclass:: vector._methods.MomentumProtocolSpatial :members: :inherited-members: :member-order: bysource ``` vector-1.6.3/docs/src/momentum4d.md000066400000000000000000000005621503546127100171640ustar00rootroot00000000000000# Interface for 4D momentum 4D momentum vectors of all backends have the following attributes, properties, and methods. This includes momentum synonyms. For purely geometric vectors, see [Interface for 4D vectors](vector4d.md). ```{eval-rst} .. autoclass:: vector._methods.MomentumProtocolLorentz :members: :inherited-members: :member-order: bysource ``` vector-1.6.3/docs/src/numba.ipynb000066400000000000000000000240561503546127100167220ustar00rootroot00000000000000{ "cells": [ { "cell_type": "markdown", "id": "1af51709-6860-4e5c-96c9-6a437eaed606", "metadata": {}, "source": [ "# Compiling functions on vectors with Numba" ] }, { "cell_type": "markdown", "id": "03064225-ea6b-4f6b-a5bf-08a47b62350e", "metadata": {}, "source": [ "First, [install](../index.md#installation) and import Vector and [Numba](https://numba.pydata.org/)." ] }, { "cell_type": "code", "execution_count": 1, "id": "df8bea35-cf3c-4f87-8e3d-8ffe06264811", "metadata": {}, "outputs": [], "source": [ "import vector\n", "import numba as nb\n", "import numpy as np" ] }, { "cell_type": "markdown", "id": "7fff9620-227b-4e5e-a211-4daf1fb96298", "metadata": {}, "source": [ "Numba is a just-in-time (JIT) compiler for a mathematically relevant subset of NumPy and Python. It allows you to write fast code without leaving the Python environment. The drawback of Numba is that it can only compile code blocks involving objects and functions that it recognizes.\n", "\n", "The Vector library includes extensions to inform Numba about [vector objects](object.md), [arrays of vectors](numpy.md), and [arrays of Awkward Arrays](awkward.md). At the time of writing, the implementation of vector NumPy arrays is incomplete (see issue [#43](https://github.com/scikit-hep/vector/issues/43)).\n", "\n", "Consider the following function:" ] }, { "cell_type": "code", "execution_count": 2, "id": "0e7f7c5c-2ce3-4f55-9159-363910fac6b9", "metadata": {}, "outputs": [], "source": [ "def compute_mass(v1, v2):\n", " return (v1 + v2).mass" ] }, { "cell_type": "code", "execution_count": 3, "id": "16d0fe04-fc8d-4d1f-bb68-0f869c1d2bdb", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "18.3 μs ± 144 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)\n" ] } ], "source": [ "%timeit compute_mass(vector.obj(px=1, py=2, pz=3, E=4), vector.obj(px=-1, py=-2, pz=-3, E=4))" ] }, { "cell_type": "code", "execution_count": 4, "id": "22516614-48eb-47de-96a9-12781fdeb66f", "metadata": {}, "outputs": [], "source": [ "@nb.njit\n", "def compute_mass(v1, v2):\n", " return (v1 + v2).mass" ] }, { "cell_type": "code", "execution_count": 5, "id": "293e83fa-d059-4f14-a70f-3c4d91de8f9b", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "The slowest run took 27658.64 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "93 ms ± 227 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%timeit -n 1 compute_mass(vector.obj(px=1, py=2, pz=3, E=4), vector.obj(px=-1, py=-2, pz=-3, E=4))" ] }, { "cell_type": "markdown", "id": "69c36f6d-3b54-4db8-9890-9245d64fefe7", "metadata": {}, "source": [ "When the two `MomentumObject4D` objects are passed as arguments, Numba recognizes them and replaces the Python objects with low-level structs. When it compiles the function, it recognizes `+` as the 4D `add` function and recognizes `.mass` as the `tau` component of the result.\n", "\n", "Although this demonstrates that Numba can manipulate vector objects, there is no performance advantage (and a likely disadvantage) to compiling a calculation on just a few vectors. The `@nb.njit` result showcases this behavior, where the run (the actual just-in-time compilation + the run) takes much longer to run that the non-JIT version.\n", "\n", "Once the function has been JIT-compiled, the subsequent runs are comparable to the non-JIT version, but is no performance advantage on just a few vectors." ] }, { "cell_type": "code", "execution_count": 6, "id": "d0fb0559-b944-43d1-b44e-1b959c1bbbd4", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "24.1 μs ± 116 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n" ] } ], "source": [ "%timeit compute_mass(vector.obj(px=1, py=2, pz=3, E=4), vector.obj(px=-1, py=-2, pz=-3, E=4))" ] }, { "cell_type": "markdown", "id": "0bb7c5d1-2310-4a9a-8fdb-176179e3f84f", "metadata": {}, "source": [ "However, the real advantage comes when many vectors are involved, in arrays." ] }, { "cell_type": "code", "execution_count": 7, "id": "63f614a2-969d-465d-8d74-a44f6f77eab4", "metadata": {}, "outputs": [ { "data": { "text/html": [ "
[[{x: -1.47, y: -0.254, z: 1.44, t: 11}, ..., {x: -0.96, y: -0.5, z: ..., ...}],\n",
       " [{x: 0.46, y: -1.97, z: -0.815, t: 9.54}, ..., {x: 0.399, y: 1.36, ...}],\n",
       " [],\n",
       " [{x: -0.484, y: 1.29, z: -0.875, t: 9.92}, {x: 1.05, y: -1.39, ...}],\n",
       " [{x: -0.441, y: -0.38, z: -1.99, t: 10.2}],\n",
       " [{x: -0.974, y: -0.324, z: 0.0935, t: 8.76}, {x: 0.989, y: 0.577, ...}],\n",
       " [{x: -2.34, y: 0.919, z: 2.36, t: 10.6}],\n",
       " [{x: 0.109, y: 0.589, z: -0.555, t: 11.3}, ..., {x: 0.681, y: -0.827, ...}],\n",
       " [{x: 0.716, y: 0.122, z: 0.21, t: 10}, ..., {x: -1.41, y: -1.53, z: ..., ...}],\n",
       " [],\n",
       " ...,\n",
       " [{x: 1.51, y: -1.67, z: 0.00979, t: 10.5}, ..., {x: -0.347, y: -0.968, ...}],\n",
       " [{x: 0.765, y: -0.107, z: 0.0998, t: 10.2}],\n",
       " [{x: 1.96, y: -0.0284, z: -0.204, t: 8.2}, {x: 0.411, y: -1.67, ...}],\n",
       " [{x: 0.398, y: 0.242, z: -0.609, t: 8.86}],\n",
       " [{x: -0.429, y: 1.23, z: 1.44, t: 10.1}],\n",
       " [{x: -0.52, y: -0.614, z: 0.586, t: 10.6}, ..., {x: -0.732, y: 0.394, ...}],\n",
       " [{x: -0.21, y: -0.225, z: 0.503, t: 11.6}],\n",
       " [],\n",
       " [{x: -1.63, y: 0.324, z: 0.701, t: 9.05}, {x: -0.000179, y: -0.59, ...}]]\n",
       "---------------------------------------------------------------------------------------------\n",
       "backend: cpu\n",
       "nbytes: 3.0 kB\n",
       "type: 50 * var * Momentum4D[\n",
       "    x: float64,\n",
       "    y: float64,\n",
       "    z: float64,\n",
       "    t: float64\n",
       "]
" ], "text/plain": [ "" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "# This is still not a large number. You want millions.\n", "array = vector.Array(\n", " [\n", " [\n", " dict(\n", " {x: np.random.normal(0, 1) for x in (\"px\", \"py\", \"pz\")},\n", " E=np.random.normal(10, 1),\n", " )\n", " for inner in range(np.random.poisson(1.5))\n", " ]\n", " for outer in range(50)\n", " ]\n", ")\n", "array" ] }, { "cell_type": "code", "execution_count": 8, "id": "f6edd6c9-a423-4662-bcae-c685686633c8", "metadata": {}, "outputs": [], "source": [ "def compute_masses(array):\n", " out = np.empty(len(array), np.float64)\n", " for i, event in enumerate(array):\n", " total = vector.obj(px=0.0, py=0.0, pz=0.0, E=0.0)\n", " for vec in event:\n", " total = total + vec\n", " out[i] = total.mass\n", " return out" ] }, { "cell_type": "code", "execution_count": 9, "id": "fa854373-aeba-4274-8ab8-c1c85d5a931b", "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "98.5 ms ± 305 μs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" ] } ], "source": [ "%timeit compute_masses(array)" ] }, { "cell_type": "code", "execution_count": 10, "id": "a6ca1a20-3f5f-485b-8cc0-663cdbf70ec6", "metadata": {}, "outputs": [], "source": [ "@nb.njit\n", "def compute_masses(array):\n", " out = np.empty(len(array), np.float64)\n", " for i, event in enumerate(array):\n", " total = vector.obj(px=0.0, py=0.0, pz=0.0, E=0.0)\n", " for vec in event:\n", " total = total + vec\n", " out[i] = total.mass\n", " return out" ] }, { "cell_type": "code", "execution_count": 11, "id": "c9bbfba4-1000-4d86-9804-340f8494076d", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "The slowest run took 2514.01 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "69.1 ms ± 169 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%timeit -n 1 compute_masses(array)" ] }, { "cell_type": "markdown", "id": "48d33ad4-a1ad-47f5-ac19-d8d2359265cb", "metadata": {}, "source": [ "This time, given that the function operates on a large array, the subsequent runs are much faster (by a considerable factor) than the non-JIT version runs." ] }, { "cell_type": "code", "execution_count": 12, "id": "56f6f10d-ee09-47c8-81da-a639a27a9645", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "194 μs ± 2.57 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n" ] } ], "source": [ "%timeit compute_masses(array)" ] }, { "cell_type": "code", "execution_count": null, "id": "950f1c48-956b-436b-8b8e-9a2b65fe45d8", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.11" } }, "nbformat": 4, "nbformat_minor": 5 } vector-1.6.3/docs/src/numpy.ipynb000066400000000000000000000611751503546127100167730ustar00rootroot00000000000000{ "cells": [ { "cell_type": "markdown", "id": "20258df9-4dbb-4457-945c-c9b6abe58cf9", "metadata": {}, "source": [ "# NumPy arrays of vectors" ] }, { "cell_type": "markdown", "id": "1144533f-0290-43f2-a9f2-adc9d042ed1a", "metadata": {}, "source": [ "First, [install](../index.md#installation) and import Vector and [NumPy](https://numpy.org/). (Vector requires NumPy, so if you can use Vector, you've already installed NumPy.)" ] }, { "cell_type": "code", "execution_count": 1, "id": "04ce97c4-d7db-458a-9976-68af8eeb9b58", "metadata": {}, "outputs": [], "source": [ "import vector\n", "import numpy as np" ] }, { "cell_type": "markdown", "id": "c9492633-fa30-458d-ba40-3c1b12a067b6", "metadata": {}, "source": [ "## Making an array of vectors" ] }, { "cell_type": "markdown", "id": "f84dbc67-e195-4ee7-a11f-cd8040c09fa8", "metadata": {}, "source": [ "If you want to do calculations with large numbers of vectors or performance is important, it's better to make arrays of vectors than lists of [vector objects](object.md).\n", "\n", "The [vector.array](make_numpy.md#general-constructor) function is a general-purpose constructor. It works like [np.array](https://numpy.org/doc/stable/reference/generated/numpy.array.html) and expects the [dtype](https://numpy.org/doc/stable/reference/arrays.dtypes.html) to correspond to a [structured array](https://numpy.org/doc/stable/user/basics.rec.html) in which the field names are recognized coordinate names:" ] }, { "cell_type": "code", "execution_count": 2, "id": "de83bfb0-37ec-4d1a-a6c4-31de7ee94998", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorNumpy2D([(1.1, 2.2), (3.3, 4.4), (5.5, 6.6)],\n", " dtype=[('x', ' abs(b), a, b)" ] }, { "cell_type": "code", "execution_count": 22, "id": "b1a1236b-81f2-40a4-b5c5-5c46a7cb5e57", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorNumpy3D((496.97860925, -22.8454872, -148.7924227),\n", " dtype=[('x', '\n", "\n", "This one constructor, [vector.obj](make_object.md), can output a variety of data types. If you want to control the type more explicitly, you can use [vector.VectorObject2D](make_object.md#vector.VectorObject2D), [vector.MomentumObject2D](make_object.md#vector.MomentumObject2D), [vector.VectorObject3D](make_object.md#vector.VectorObject3D), [vector.MomentumObject3D](make_object.md#vector.MomentumObject3D), [vector.VectorObject4D](make_object.md#vector.VectorObject4D), and [vector.MomentumObject4D](make_object.md#vector.MomentumObject4D) to construct or check the type explicitly. These classes also have `from_*` methods to construct vectors from positional arguments, rather than keyword arguments." ] }, { "cell_type": "markdown", "id": "d9a23b9b-a19f-49a9-b1ae-7b86c51a17f1", "metadata": {}, "source": [ "## Using a vector" ] }, { "cell_type": "markdown", "id": "959d7c4e-f7a6-4220-877a-dfedc6a873b7", "metadata": {}, "source": [ "Vector objects have a suite of properties and methods appropriate to their type (2D/3D/4D, geometric or momentum). For example, to compute the cross-product of two vectors, you would use [cross](vector3d.md#vector._methods.VectorProtocolSpatial.cross):" ] }, { "cell_type": "code", "execution_count": 5, "id": "51d9350f-66e7-447c-840c-ed6bd899bd37", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject3D(x=6, y=0, z=-3)" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a = vector.obj(x=2, y=3, z=4)\n", "b = vector.obj(x=1, y=0, z=2)\n", "\n", "a.cross(b)" ] }, { "cell_type": "markdown", "id": "98530a7c-5413-4268-9312-91e1e8054c50", "metadata": {}, "source": [ "or to compute the angle between them, you would use [deltaangle](vector3d.md#vector._methods.VectorProtocolSpatial.deltaangle):" ] }, { "cell_type": "code", "execution_count": 6, "id": "eee16a4d-b15b-434b-868b-ef7a11d8e101", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "0.590872750145419" ] }, "execution_count": 6, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a.deltaangle(b)" ] }, { "cell_type": "markdown", "id": "ac749138-ed71-4a07-92a8-af33785eb5ea", "metadata": {}, "source": [ "or to compute their sum, you would use `+`:" ] }, { "cell_type": "code", "execution_count": 7, "id": "5686fc89-5cb0-4a5f-b521-41dd7d0a3b10", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject3D(x=3, y=3, z=6)" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a + b" ] }, { "cell_type": "markdown", "id": "eeaf4807-b338-4f17-ba35-a7a3f7f43d97", "metadata": {}, "source": [ "In this last example, the `+` operator overloads the [add](common.md#vector._methods.VectorProtocol.add) method. Similarly, multiplication between a vector and a scalar number overloads [scale](common.md#vector._methods.VectorProtocol.scale), etc. Since they overload standard operators, vectors can be used in Python built-in functions like [sum](https://docs.python.org/3/library/functions.html#sum), as long as you provide a `start`:" ] }, { "cell_type": "code", "execution_count": 8, "id": "10b5f3e3-db5d-421e-89a4-5929fe720af1", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject2D(x=45, y=4.5)" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vs = [vector.obj(x=x, y=x / 10) for x in range(10)]\n", "\n", "sum(vs, start=vector.obj(x=0, y=0))" ] }, { "cell_type": "markdown", "id": "32c06273-06ae-4cfa-8e79-f55586148283", "metadata": {}, "source": [ "The same applies to [abs](https://docs.python.org/3/library/functions.html#abs) for the vector's magnitude, but note that this depends on the number of dimensions:" ] }, { "cell_type": "code", "execution_count": 9, "id": "e821c6d7-db8c-47cd-93a0-5eb7bf3e023f", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "5.0" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "abs(vector.obj(x=3, y=4)) # sqrt(3**2 + 4**2)" ] }, { "cell_type": "code", "execution_count": 10, "id": "bf89db7e-6326-46e4-8add-fa5e7e7df799", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "3.0" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "abs(vector.obj(x=1, y=2, z=2)) # sqrt(1**2 + 2**2 + 2**2)" ] }, { "cell_type": "code", "execution_count": 11, "id": "9f6d19e6-ca20-4f8b-a8d9-57950c24655b", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "3.0" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "abs(vector.obj(x=3, y=3, z=3, t=6)) # sqrt(6**2 - 3**2 - 3**2 - 3**2)" ] }, { "cell_type": "markdown", "id": "2d9407af-0641-403c-822c-d89344dd95db", "metadata": {}, "source": [ "Equality ([equal](common.md#vector._methods.VectorProtocol.equal)) and inequality ([not_equal](common.md#vector._methods.VectorProtocol.not_equal)) are defined:" ] }, { "cell_type": "code", "execution_count": 12, "id": "41ca1487-4a33-493f-9d9b-dd452255ae73", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vector.obj(x=3, y=4) == vector.obj(x=3, y=4)" ] }, { "cell_type": "markdown", "id": "bf85dc4f-88de-4e07-9303-c1f364a0a227", "metadata": {}, "source": [ "But you'll probably want to use [isclose](common.md#vector._methods.VectorProtocol.isclose) (and possibly specify tolerances):" ] }, { "cell_type": "code", "execution_count": 13, "id": "7604a176-efd9-4cbd-947c-90c8bbb20685", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "False" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vector.obj(x=3, y=4) == vector.obj(rho=5, phi=0.9272952180016122)" ] }, { "cell_type": "code", "execution_count": 14, "id": "0a3d1aa1-eeda-486a-84b2-e9d43e102117", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vector.obj(x=3, y=4).isclose(vector.obj(rho=5, phi=0.9272952180016122))" ] }, { "cell_type": "markdown", "id": "70a3c4b3-8a81-4bb2-b6f7-1ebb57889c0a", "metadata": {}, "source": [ "The full set of properties and methods available to each type of vector (2D/3D/4D, geometric or momentum) is described in\n", "\n", "* [Interface for all vectors](common.md)\n", "* [Interface for 2D vectors](vector2d.md)\n", "* [Interface for 3D vectors](vector3d.md)\n", "* [Interface for 4D vectors](vector4d.md)\n", "* [Interface for 2D momentum](momentum2d.md)\n", "* [Interface for 3D momentum](momentum3d.md)\n", "* [Interface for 4D momentum](momentum4d.md)" ] }, { "cell_type": "markdown", "id": "6d1f1795-1051-4b32-b9b8-222a2f2228f9", "metadata": {}, "source": [ "## Using coordinate systems" ] }, { "cell_type": "markdown", "id": "0e80cf47-801f-421a-818a-e70db3356a5a", "metadata": {}, "source": [ "A vector can be constructed using any combination of coordinate systems and computations will be performed using whatever coordinate system it has. Thus, after creating vectors, you can write code that does not depend on the coordinate system—it becomes a hidden implementation detail." ] }, { "cell_type": "code", "execution_count": 15, "id": "a96ddaad-7c6b-487f-8d23-fdd2fde159bc", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject2D(x=5.0, y=1.0)" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a = vector.obj(x=2, y=1)\n", "b = vector.obj(rho=3, phi=0)\n", "\n", "a + b" ] }, { "cell_type": "markdown", "id": "4f4552b3-bd21-4760-90fc-92002a2f14ad", "metadata": {}, "source": [ "Some of the properties of a vector are coordinates, so you can use Vector to convert coordinates." ] }, { "cell_type": "code", "execution_count": 16, "id": "d0ce4192-f634-40f3-900a-4d143b0050d6", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(2.23606797749979, 0.4636476090008061)" ] }, "execution_count": 16, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a.rho, a.phi" ] }, { "cell_type": "markdown", "id": "633a6970-6986-40e3-a3aa-1fde0edc0a0a", "metadata": {}, "source": [ "Since the way that you access the original coordinates is the same as the way that you access converted coordinates," ] }, { "cell_type": "code", "execution_count": 17, "id": "38006968-210d-4568-a94a-e13d31b048a7", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "(2, 1)" ] }, "execution_count": 17, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a.x, a.y" ] }, { "cell_type": "markdown", "id": "31f7e65a-ed4e-48a8-a761-4c70cf81eca4", "metadata": {}, "source": [ "these conversions are part of the coordinate-abstraction.\n", "\n", "For reasons of numerical precision, you might want to open this black box and explicitly change the coordinate system. These methods start with `to_*`:" ] }, { "cell_type": "code", "execution_count": 18, "id": "83c22d5b-a251-42c2-9527-87695c9d5dac", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject2D(rho=2.23606797749979, phi=0.4636476090008061)" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "a.to_rhophi()" ] }, { "cell_type": "code", "execution_count": 19, "id": "01754e13-3c0c-4b1e-be0e-0e0d01585b63", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject2D(x=3.0, y=0.0)" ] }, "execution_count": 19, "metadata": {}, "output_type": "execute_result" } ], "source": [ "b.to_xy()" ] }, { "cell_type": "markdown", "id": "3c4a7d35-74a6-4d90-b195-68fc6a77cbec", "metadata": {}, "source": [ "## Geometric versus momentum vectors" ] }, { "cell_type": "markdown", "id": "459d089b-8569-4eff-a316-6ff40cc1b145", "metadata": {}, "source": [ "Vectors come in two flavors:\n", "\n", "* geometric: only one name for each property or method\n", "* momentum: same property or method can be accessed with several synonyms (which assume that the vector is a [momentum](https://en.wikipedia.org/wiki/Momentum) vector)." ] }, { "cell_type": "code", "execution_count": 20, "id": "e86873c1-f8aa-447d-b265-2a6b42b9e1ba", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorObject3D(x=1, y=2, z=3)" ] }, "execution_count": 20, "metadata": {}, "output_type": "execute_result" } ], "source": [ "v = vector.obj(x=1, y=2, z=3)\n", "v" ] }, { "cell_type": "code", "execution_count": 21, "id": "0d4fab89-cf8c-43f9-a7d9-34083ac70eb7", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "MomentumObject3D(px=1, py=2, pz=3)" ] }, "execution_count": 21, "metadata": {}, "output_type": "execute_result" } ], "source": [ "p = vector.obj(px=1, py=2, pz=3)\n", "p" ] }, { "cell_type": "markdown", "id": "1f19524b-e98a-461e-a098-6d5b803dc693", "metadata": {}, "source": [ "Calculations are the same in both cases:" ] }, { "cell_type": "code", "execution_count": 22, "id": "0eaf2b80-9e94-41bb-b251-1b713df0e829", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "3.7416573867739413" ] }, "execution_count": 22, "metadata": {}, "output_type": "execute_result" } ], "source": [ "abs(v)" ] }, { "cell_type": "code", "execution_count": 23, "id": "f1c35d52-9e8d-44d1-b1e8-a4c0f3405370", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "3.7416573867739413" ] }, "execution_count": 23, "metadata": {}, "output_type": "execute_result" } ], "source": [ "abs(p)" ] }, { "cell_type": "markdown", "id": "dec60082-29ca-47b8-b9c4-3d68ff967a06", "metadata": {}, "source": [ "but there are more ways to express some operations:" ] }, { "cell_type": "code", "execution_count": 24, "id": "1a3540ed-c4f1-4df0-af75-0a0c75b83c89", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "2.23606797749979" ] }, "execution_count": 24, "metadata": {}, "output_type": "execute_result" } ], "source": [ "v.rho" ] }, { "cell_type": "code", "execution_count": 25, "id": "061c00d9-2134-41bb-9275-a81eccb4a71c", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "2.23606797749979" ] }, "execution_count": 25, "metadata": {}, "output_type": "execute_result" } ], "source": [ "p.rho" ] }, { "cell_type": "code", "execution_count": 26, "id": "0fb9656b-97d7-49f0-b61e-c99bccf2e31d", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "2.23606797749979" ] }, "execution_count": 26, "metadata": {}, "output_type": "execute_result" } ], "source": [ "p.pt" ] }, { "cell_type": "markdown", "id": "375ac938-9ef0-493b-b94a-52cab0c1b359", "metadata": {}, "source": [ "The geometric vector satisfies the [Zen of Python](https://en.wikipedia.org/wiki/Zen_of_Python) stipulation that\n", "\n", "> There should be one-- and preferably only one --obvious way to do it.\n", "\n", "and code that uses, for example, \"`pt`\" to specify \"distance from the beamline\" is obfuscated code. However, the most common use for these vectors in High Energy Physics (HEP) is to represent the momentum of particles. For that purpose, using \"`rho`\" for $p_T$ is not self-documenting.\n", "\n", "Momentum vectors have all of the same properties and methods as geometric vectors _as well as_ momentum synonyms. In some cases, there are multiple momentum synonyms for adherence to different conventions. For example, energy and mass (the [temporal component of momentum](https://en.wikipedia.org/wiki/Four-momentum), as Cartesian and proper time, respectively) have four different spellings:\n", "\n", "| energy spelling | mass spelling | rationale |\n", "|:--:|:--:|:--|\n", "| `t` | `tau` | geometric coordinates; $\\tau$ for proper time is conventional |\n", "| `energy` | `mass` | full names are more self-documenting in the code |\n", "| `e` | `m` | all other coordinates are lower-case single letters (sometimes Greek letters) |\n", "| `E` | `M` | capital E and M (only!) are used in other HEP vector libraries |\n", "\n", "If any momentum components are used to construct a vector (or if [vector.MomentumObject2D](make_object.md#vector.MomentumObject2D), [vector.MomentumObject3D](make_object.md#vector.MomentumObject3D), or [vector.MomentumObject4D](make_object.md#vector.MomentumObject4D) are used explicitly), then the vector is momentum and all synonyms become available." ] }, { "cell_type": "markdown", "id": "0560e95f-48e9-454d-9db1-f85386d5e9ed", "metadata": {}, "source": [ "## Numeric data types and numerical error" ] }, { "cell_type": "markdown", "id": "c3ba958e-9c1b-4721-8f40-b1112ea6a068", "metadata": {}, "source": [ "Vector does not require any specific numeric data type, such as `np.float32` or `np.float64`, it only requires that vector components are some kind of number, including integers." ] }, { "cell_type": "code", "execution_count": 27, "id": "15ca1301-574c-43dd-bb83-f85c71a5f8dc", "metadata": {}, "outputs": [], "source": [ "v = vector.obj(x=1, y=2.2)" ] }, { "cell_type": "code", "execution_count": 28, "id": "c1883cf8-231f-4996-840a-65c4f30eced4", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "int" ] }, "execution_count": 28, "metadata": {}, "output_type": "execute_result" } ], "source": [ "type(v.x)" ] }, { "cell_type": "code", "execution_count": 29, "id": "18f515ee-1c9d-480f-b4c8-ba3033afcff0", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "float" ] }, "execution_count": 29, "metadata": {}, "output_type": "execute_result" } ], "source": [ "type(v.y)" ] }, { "cell_type": "code", "execution_count": 30, "id": "b1a2e4d9-8381-4bc2-9332-52fa475d6c4b", "metadata": {}, "outputs": [], "source": [ "import numpy as np" ] }, { "cell_type": "code", "execution_count": 31, "id": "e35f37d0-b6bd-45d0-ad78-7a17629a4e3a", "metadata": {}, "outputs": [], "source": [ "v = vector.obj(x=np.float32(1.1), y=np.float64(2.2))" ] }, { "cell_type": "code", "execution_count": 32, "id": "6bd30abb-4566-4551-9fb7-c7bc8643d487", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "numpy.float32" ] }, "execution_count": 32, "metadata": {}, "output_type": "execute_result" } ], "source": [ "type(v.x)" ] }, { "cell_type": "code", "execution_count": 33, "id": "e4fabf24-2799-4b0a-9e29-c106986c3330", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "numpy.float64" ] }, "execution_count": 33, "metadata": {}, "output_type": "execute_result" } ], "source": [ "type(v.y)" ] }, { "cell_type": "markdown", "id": "11c3e336-4ced-4704-a1c8-54363687c20f", "metadata": {}, "source": [ "The same formulas are applied, regardless of the numeric type, so if the numerical error is larger than you expect it to be, check your types (and coordinate systems)." ] }, { "cell_type": "markdown", "id": "0dab118c-e262-45ec-b121-4c0172603b4a", "metadata": {}, "source": [ "## Application to other backends" ] }, { "cell_type": "markdown", "id": "75959cc7-a64b-48ac-9b48-1244887d7d9a", "metadata": {}, "source": [ "Everything stated above about vector objects (except their methods of construction) apply equally to all other backends. Arrays of vectors and symbolic vector expressions in SymPy have the same properties and methods as vector objects, they can hide choice of coordinate system as an abstraction, and the set of synonyms can be minimal for geometric vectors and maximal for momentum vectors. Therefore, it can be convenient to use vector objects as a quick way to debug issues in large arrays. However, note that different backends can use different libraries for computations, and results might differ in numerical error.\n", "\n", "In particular, note that SymPy vector expressions have a different sign convention for operations on space-like and negative time-like 4D vectors. For all other backends, Vector's conventions were chosen to agree with popular HEP libraries, particularly [ROOT](https://root.cern), but for the SymPy backend, those conventions would insert piecewise if-then branches, which would complicate symbolic expressions." ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.10" } }, "nbformat": 4, "nbformat_minor": 5 } vector-1.6.3/docs/src/sympy.ipynb000066400000000000000000000303651503546127100170010ustar00rootroot00000000000000{ "cells": [ { "cell_type": "markdown", "id": "dd79e619-a48d-415b-9d7a-7dd5b3970475", "metadata": {}, "source": [ "# Vector expressions with SymPy" ] }, { "cell_type": "markdown", "id": "3f40d8aa-c1d4-4106-859a-04d8452673a0", "metadata": {}, "source": [ "First, [install](../index.md#installation) and import Vector and [SymPy](https://www.sympy.org/)." ] }, { "cell_type": "code", "execution_count": 1, "id": "405dce26-762d-438b-8b8c-f693d0333a32", "metadata": {}, "outputs": [], "source": [ "import vector\n", "import sympy" ] }, { "cell_type": "markdown", "id": "00b96b54-33e4-459b-8fa6-bbebf1df0319", "metadata": {}, "source": [ "## How the SymPy backend differs from the others" ] }, { "cell_type": "markdown", "id": "07e49af1-a2e4-43e9-b265-38d917456ce0", "metadata": {}, "source": [ "SymPy is a computer algebra system like Mathematica and Maple. It primarily deals with algebraic expressions, rather than concrete numbers. However, all of the coordinate transformations and vector manipulations can be applied symbolically through Vector's SymPy backend.\n", "\n", "When comparing SymPy to the other backends, note that SymPy vector expressions have a different sign convention for operations on space-like and negative time-like 4D vectors. For all other backends, Vector's conventions were chosen to agree with popular HEP libraries, particularly [ROOT](https://root.cern), but for the SymPy backend, those conventions would insert piecewise if-then branches, which would complicate symbolic expressions.\n", "\n", "When vector expressions are evaluated numerically, you can expect agreement in 2D and 3D vector operations, as well as 4D vector operations if all of the vectors have a positive time-like part (which is [necessary for real momentum vectors and causal relationships between events](https://en.wikipedia.org/wiki/Light_cone))." ] }, { "cell_type": "markdown", "id": "afce6222-03bb-40b7-8789-4e98782ebede", "metadata": {}, "source": [ "## Making a vector expression" ] }, { "cell_type": "markdown", "id": "10257bc9-a064-4ede-bcdb-d906a25598e2", "metadata": {}, "source": [ "Before making a vector expression, we need symbols for each of the components, so use the [sympy.symbols](https://docs.sympy.org/latest/modules/core.html#sympy.core.symbol.symbols) function. Be sure to [tell SymPy to assume](https://docs.sympy.org/latest/guides/assumptions.html) that they are all real-valued, not complex numbers." ] }, { "cell_type": "code", "execution_count": 2, "id": "e74eaaea-c607-47c0-a228-9b3530863bf7", "metadata": {}, "outputs": [], "source": [ "x, y, z, t, px, py, pz, eta, tau = sympy.symbols(\"x y z t px py pz eta tau\", real=True)" ] }, { "cell_type": "markdown", "id": "8216735e-eb82-4ba9-92c3-12f45a9400a8", "metadata": {}, "source": [ "Now we can make vectors [just as we did with objects](object.md), though these lack concrete numerical values." ] }, { "cell_type": "code", "execution_count": 3, "id": "0beb59b0-7cea-4d06-b437-da0597ba027e", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorSympy2D(x=x, y=y)" ] }, "execution_count": 3, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vector.VectorSympy2D(x=x, y=y)" ] }, { "cell_type": "code", "execution_count": 4, "id": "8dc8eb68-89e3-4794-b533-e9f5d521770b", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "MomentumSympy3D(px=px, py=py, pz=pz)" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vector.MomentumSympy3D(px=px, py=py, pz=pz)" ] }, { "cell_type": "code", "execution_count": 5, "id": "921724ac-2d20-4a82-8e0b-3cc571d7d909", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorSympy4D(x=x, y=y, eta=eta, tau=tau)" ] }, "execution_count": 5, "metadata": {}, "output_type": "execute_result" } ], "source": [ "vector.VectorSympy4D(x=x, y=y, eta=eta, tau=tau)" ] }, { "cell_type": "markdown", "id": "6ac26f0b-8fb6-410c-865f-6fefaeb2302b", "metadata": {}, "source": [ "## Using a vector expression" ] }, { "cell_type": "markdown", "id": "59c644ed-32af-42f1-bac7-a27c7ec55035", "metadata": {}, "source": [ "All of the vector operations performed on these expressions return symbolic results." ] }, { "cell_type": "code", "execution_count": 6, "id": "9dd7d004-9635-427c-8394-f0e65906f564", "metadata": {}, "outputs": [], "source": [ "v = vector.VectorSympy2D(x=x, y=y)" ] }, { "cell_type": "code", "execution_count": 7, "id": "3b3af4ee-b0a7-4379-ab8e-ac674037db18", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\sqrt{x^{2} + y^{2}}$" ], "text/plain": [ "sqrt(x**2 + y**2)" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "v.rho" ] }, { "cell_type": "code", "execution_count": 8, "id": "9430dd6e-ceef-4977-ab1f-5d78ff7d6d83", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\text{True}$" ], "text/plain": [ "True" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "sympy.Eq(v.rho, abs(v))" ] }, { "cell_type": "code", "execution_count": 9, "id": "7c674475-b882-44fc-b01a-f04800eb0205", "metadata": {}, "outputs": [], "source": [ "v = vector.VectorSympy4D(x=x, y=y, z=z, t=t)" ] }, { "cell_type": "code", "execution_count": 10, "id": "c25ea976-b913-4717-9858-3e2c926a99c2", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\sqrt{\\left|{- t^{2} + x^{2} + y^{2} + z^{2}}\\right|}$" ], "text/plain": [ "sqrt(Abs(-t**2 + x**2 + y**2 + z**2))" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "v.tau" ] }, { "cell_type": "code", "execution_count": 11, "id": "3d258221-4cd4-43fb-8710-d9bc96827132", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle t^{2} - x^{2} - y^{2} - z^{2} > 0$" ], "text/plain": [ "t**2 - x**2 - y**2 - z**2 > 0" ] }, "execution_count": 11, "metadata": {}, "output_type": "execute_result" } ], "source": [ "v.is_timelike()" ] }, { "cell_type": "code", "execution_count": 12, "id": "86dcae0d-22ba-4621-82b8-397a94fecfc6", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "VectorSympy4D(x=x*(1 + x**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))) + x/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x*y**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + x*z**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)), y=y*(1 + y**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))) + y/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2*y/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y*z**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)), z=z*(1 + z**2/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))) + z/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2*z/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y**2*z/(t**2*(1 + 1/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))*(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)), t=t/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + z**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)))" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "boosted = v.boost(v.to_beta3())\n", "boosted" ] }, { "cell_type": "code", "execution_count": 13, "id": "1e5a3bf0-1784-4d29-8b7c-1b8904de8a20", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\frac{t}{\\sqrt{1 - \\frac{x^{2}}{t^{2}} - \\frac{y^{2}}{t^{2}} - \\frac{z^{2}}{t^{2}}}} + \\frac{x^{2}}{t \\sqrt{1 - \\frac{x^{2}}{t^{2}} - \\frac{y^{2}}{t^{2}} - \\frac{z^{2}}{t^{2}}}} + \\frac{y^{2}}{t \\sqrt{1 - \\frac{x^{2}}{t^{2}} - \\frac{y^{2}}{t^{2}} - \\frac{z^{2}}{t^{2}}}} + \\frac{z^{2}}{t \\sqrt{1 - \\frac{x^{2}}{t^{2}} - \\frac{y^{2}}{t^{2}} - \\frac{z^{2}}{t^{2}}}}$" ], "text/plain": [ "t/sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2) + x**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + y**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2)) + z**2/(t*sqrt(1 - x**2/t**2 - y**2/t**2 - z**2/t**2))" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "boosted.t" ] }, { "cell_type": "markdown", "id": "d7e6de0b-917e-4d87-bade-74bd62db93f2", "metadata": {}, "source": [ "They can be [simplified](https://docs.sympy.org/latest/modules/simplify/simplify.html):" ] }, { "cell_type": "code", "execution_count": 14, "id": "7bcac010-14a5-4d81-9e8a-c4cc9d84067a", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\frac{t \\sqrt{\\frac{t^{2} - x^{2} - y^{2} - z^{2}}{t^{2}}} \\left(t^{2} + x^{2} + y^{2} + z^{2}\\right)}{t^{2} - x^{2} - y^{2} - z^{2}}$" ], "text/plain": [ "t*sqrt((t**2 - x**2 - y**2 - z**2)/t**2)*(t**2 + x**2 + y**2 + z**2)/(t**2 - x**2 - y**2 - z**2)" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "boosted.t.simplify()" ] }, { "cell_type": "markdown", "id": "8e29d643-52e8-4b7d-bbd3-2b864d5efdb6", "metadata": {}, "source": [ "And the symbols can be [replaced with numerical values](https://docs.sympy.org/latest/modules/core.html#sympy.core.basic.Basic.subs):" ] }, { "cell_type": "code", "execution_count": 15, "id": "0b9b401d-43a1-411f-9d88-f7c91fa6c5e7", "metadata": {}, "outputs": [ { "data": { "text/latex": [ "$\\displaystyle \\frac{57 \\sqrt{86}}{43}$" ], "text/plain": [ "57*sqrt(86)/43" ] }, "execution_count": 15, "metadata": {}, "output_type": "execute_result" } ], "source": [ "boosted.t.subs({x: 3, y: 2, z: 1, t: 10})" ] }, { "cell_type": "markdown", "id": "55befc25-e5a1-465c-9dfa-850f0dc06cf5", "metadata": {}, "source": [ "Or [converted into code](https://docs.sympy.org/latest/modules/printing.html#prettyprinter-class) for a programming language:" ] }, { "cell_type": "code", "execution_count": 16, "id": "57fa1123-2d1a-47ac-ad8e-31a0ceee8747", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ " t*sqrt((t**2 - x**2 - y**2 - z**2)/t**2)*(t**2 + x**2 + y**2 + z**\n", " @ 2)/(t**2 - x**2 - y**2 - z**2)\n" ] } ], "source": [ "import sympy.printing.fortran\n", "\n", "print(sympy.printing.fortran.fcode(boosted.t.simplify()))" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.10" } }, "nbformat": 4, "nbformat_minor": 5 } vector-1.6.3/docs/src/talks.md000066400000000000000000000035171503546127100162140ustar00rootroot00000000000000# Papers and talks ## Papers - [CHEP 2024](https://indico.cern.ch/event/1338689) (to be published): [A new SymPy backend for vector: uniting experimental and theoretical physicists](https://indi.to/zTs5b) - [JOSS](https://joss.theoj.org) (2025-05-23): [Vector: JIT-compilable mathematical manipulations of ragged Lorentz vectors](https://doi.org/10.21105/joss.07791) ## Talks - [PyHEP 2024](https://indico.cern.ch/event/1384010) (2024-07-03): [🎥](https://www.youtube.com/watch?v=aVPhP_APhqw) [A new SymPy backend for vector: uniting experimental and theoretical physicists](https://indi.to/pfTC6) - [PyHEP 2023](https://indico.cern.ch/event/1252095) (2023-10-09): [🎥](https://www.youtube.com/watch?v=JHEAb2R3xzE&list=PLKZ9c4ONm-VlAorAG8kR09ZqhMfHiH2LJ&index=10) [What’s new with Vector? First major release is out!](https://indi.to/35ym5) - [PyHEP 2022](https://indico.cern.ch/event/1150631) (2022-09-13): [🎥](https://www.youtube.com/watch?v=4iveMzrbe7s&list=PLKZ9c4ONm-VkohKG-skzEG_gklMaSgaO7&index=15) [Constructing HEP vectors and analyzing HEP data using Vector](https://indi.to/bPmMc) - [DANCE/CoDaS@Snowmass 2022](https://indico.cern.ch/event/1151329) (2022-07-20): [Analysis Grand Challenge / HEP Scientific Python Ecosystem](https://indico.cern.ch/event/1151329/timetable/#3-analysis-grand-challenge-hep) - [IRIS-HEP AGC Tools 2022](https://indico.cern.ch/event/1126109) (2022-04-25): [🎥](https://www.youtube.com/watch?v=O9KvsDMKOmY) [Foundation libraries (Uproot, Awkward, hist, mplhep)](https://indico.cern.ch/event/1126109/contributions/4780138) - [IRIS-HEP AGC Tools 2021](https://indico.cern.ch/event/1076231) (2021-11-03): [🎥](https://indico.cern.ch/event/1076231/contributions/4560398/attachments/2338579/4017718/agc_uproot_awk.mp4) [Data handling: Uproot, Awkward & Vector](https://indico.cern.ch/event/1076231/contributions/4560398) vector-1.6.3/docs/src/vector2d.md000066400000000000000000000005041503546127100166170ustar00rootroot00000000000000# Interface for 2D vectors 2D vectors of all backends have the following attributes, properties, and methods. For the momentum synonyms, see [Interface for 2D momentum](momentum2d.md). ```{eval-rst} .. autoclass:: vector._methods.VectorProtocolPlanar :members: :inherited-members: :member-order: bysource ``` vector-1.6.3/docs/src/vector3d.md000066400000000000000000000005051503546127100166210ustar00rootroot00000000000000# Interface for 3D vectors 3D vectors of all backends have the following attributes, properties, and methods. For the momentum synonyms, see [Interface for 3D momentum](momentum3d.md). ```{eval-rst} .. autoclass:: vector._methods.VectorProtocolSpatial :members: :inherited-members: :member-order: bysource ``` vector-1.6.3/docs/src/vector4d.md000066400000000000000000000005051503546127100166220ustar00rootroot00000000000000# Interface for 4D vectors 4D vectors of all backends have the following attributes, properties, and methods. For the momentum synonyms, see [Interface for 4D momentum](momentum4d.md). ```{eval-rst} .. autoclass:: vector._methods.VectorProtocolLorentz :members: :inherited-members: :member-order: bysource ``` vector-1.6.3/environment.yml000066400000000000000000000004361503546127100161210ustar00rootroot00000000000000name: vector channels: - conda-forge dependencies: - jupyterlab >=1.2 - nb_conda_kernels - pip >=18 - pytest >=6 - numba >=0.57 - numpy >=1.13.3 - root >=6.18.04 - pip: - "awkward>=2.0.0" - "uproot>=5.0.0" - "scikit-hep-testdata>=0.2.0" - -e . vector-1.6.3/noxfile.py000066400000000000000000000060321503546127100150460ustar00rootroot00000000000000from __future__ import annotations from pathlib import Path import nox nox.options.sessions = ["lint", "lite", "tests", "doctests", "disassemble"] nox.needs_version = ">=2024.4.15" nox.options.default_venv_backend = "uv|virtualenv" ALL_PYTHON = ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"] DIR = Path(__file__).parent.resolve() @nox.session(reuse_venv=True) def lint(session: nox.Session) -> None: """Run the linter.""" session.install("pre-commit") session.run("pre-commit", "run", "--all-files", *session.posargs) @nox.session(reuse_venv=True) def pylint(session: nox.Session) -> None: """Run pylint.""" session.install("pylint") session.install("-e", ".") session.run("pylint", "src/vector/", *session.posargs) @nox.session(reuse_venv=True, python=ALL_PYTHON) def lite(session: nox.Session) -> None: """Run lightweight tests.""" session.install("-e", ".[test]") session.run("pytest", "--ignore", "tests/test_notebooks.py", *session.posargs) @nox.session(reuse_venv=True, python=ALL_PYTHON) def tests(session: nox.Session) -> None: """Run the unit and regular tests.""" session.install("-e", ".[awkward,numba,test,test-extras,sympy]") session.run( "pytest", "--ignore", "tests/test_notebooks.py", *session.posargs, ) @nox.session(reuse_venv=True, python=ALL_PYTHON, default=False) def coverage(session: nox.Session) -> None: """Run tests and compute coverage.""" session.posargs.extend(["--cov=vector", "--cov-report=xml"]) tests(session) @nox.session(reuse_venv=True, python=ALL_PYTHON) def doctests(session: nox.Session) -> None: """Run the doctests.""" session.install("-e", ".[awkward,numba,test,test-extras,sympy]") session.run("pytest", "--doctest-plus", "src/vector/", *session.posargs) @nox.session(reuse_venv=True, default=False) def notebooks(session: nox.Session) -> None: """Run the notebook tests""" session.install("-e", ".[awkward,numba,test,sympy]") session.install("jupyter") session.run("pytest", "tests/test_notebooks.py", *session.posargs) @nox.session(reuse_venv=True, default=False) def docs(session: nox.Session) -> None: """Build the docs. Pass "serve" to serve.""" session.install("-e", ".[docs]") session.chdir("docs") session.run("sphinx-build", "-M", "html", ".", "_build") if session.posargs: if "serve" in session.posargs: print("Launching docs at http://localhost:8001/ - use Ctrl-C to quit") session.run("python", "-m", "http.server", "8001", "-d", "_build/html") else: print("Unsupported argument to docs") @nox.session(reuse_venv=True, default=False) def build(session: nox.Session) -> None: """Build an SDist and wheel.""" session.install("build") session.run("python", "-m", "build") @nox.session(reuse_venv=True, python="3.8", default=False) def disassemble(session: nox.Session) -> None: """Disassemble run.""" session.install("-e", ".[test-extras]") session.posargs.extend(["-m", "dis"]) lite(session) vector-1.6.3/pyproject.toml000066400000000000000000000136031503546127100157460ustar00rootroot00000000000000[build-system] build-backend = "hatchling.build" requires = [ "hatch-vcs", "hatchling", ] [project] name = "vector" description = "Vector classes and utilities" readme = { file = "README.md", content-type = "text/markdown" } keywords = [ "vector", ] maintainers = [ { name = "The Scikit-HEP admins", email = "scikit-hep-admins@googlegroups.com" }, ] authors = [ { name = "Jim Pivarski, Henry Schreiner, Eduardo Rodrigues", email = "eduardo.rodrigues@cern.ch" }, ] requires-python = ">=3.8" classifiers = [ "Development Status :: 5 - Production/Stable", "Intended Audience :: Developers", "Intended Audience :: Science/Research", "License :: OSI Approved :: BSD License", "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3 :: Only", "Programming Language :: Python :: 3.8", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering", "Topic :: Scientific/Engineering :: Information Analysis", "Topic :: Scientific/Engineering :: Mathematics", "Topic :: Scientific/Engineering :: Physics", "Typing :: Typed", ] dynamic = [ "version", ] dependencies = [ "numpy>=1.13.3", "packaging>=19", ] optional-dependencies.awkward = [ "awkward>=2", ] optional-dependencies.dev = [ "awkward>=2", "dask-awkward", "nox", "numba>=0.57; python_version<'3.14'", "papermill>=2.4", "pytest>=6", "pytest-cov>=3", "pytest-doctestplus", "sympy", ] optional-dependencies.docs = [ "awkward>=2", "sympy", "ipykernel", "myst-parser>0.13", "nbsphinx", "sphinx>=4", "sphinx-book-theme>=0.0.42", "sphinx-copybutton", "sphinx-math-dollar", ] optional-dependencies.numba = [ "numba>=0.57; python_version<'3.14'", ] optional-dependencies.sympy = [ "sympy", ] optional-dependencies.test = [ "nox", "papermill>=2.4", "pytest>=6", "pytest-cov>=3", "pytest-doctestplus", ] optional-dependencies.test-extras = [ "dask_awkward", "spark-parser", 'uncompyle6; python_version == "3.8"', ] urls."Bug Tracker" = "https://github.com/scikit-hep/vector/issues" urls.Changelog = "https://vector.readthedocs.io/en/latest/changelog.html" urls.Discussions = "https://github.com/scikit-hep/vector/discussions" urls.Documentation = "https://vector.readthedocs.io/" urls.Homepage = "https://github.com/scikit-hep/vector" entry-points.numba_extensions.init = "vector:register_numba" [tool.hatch] version.source = "vcs" build.hooks.vcs.version-file = "src/vector/_version.py" [tool.ruff.lint] extend-select = [ "B", # flake8-bugbear "I", # isort "C4", # flake8-comprehensions "ISC", # flake8-implicit-str-concat "PGH", # pygrep-hooks "PIE", # flake8-pie "PL", # pylint "PT", # flake8-pytest-style "RUF", # Ruff-specific "SIM", # flake8-simplify "T20", # flake8-print "UP", # pyupgrade "YTT", # flake8-2020 ] ignore = [ "PLR09", # Too many X "PLR2004", # Magic values "ISC001", # Conflicts with formatter ] typing-modules = [ "vector._typeutils", ] isort.required-imports = [ "from __future__ import annotations", ] [tool.ruff.lint.per-file-ignores] "noxfile.py" = [ "T20", ] "tests/*" = [ "T20", ] "docs/src/*.ipynb" = [ "T20", "I001", "I002", "F811" ] "src/vector/__init__.py" = [ "PLC0415", ] "src/vector/_methods.py" = [ "PLC0415", ] "src/vector/backends/_numba_object.py" = [ "PGH003", ] "src/vector/backends/awkward.py" = [ "PLC0415", ] "src/vector/backends/awkward_constructors.py" = [ "PLC0415", ] "src/vector/backends/object.py" = [ "PLC0415", ] "tests/backends/test_operators.py" = [ "SIM201", "SIM202", "PLR0124", ] [tool.pylint] master.py-version = "3.8" reports.output-format = "colorized" similarities.ignore-imports = "yes" master.jobs = "0" messages_control.enable = [ "useless-suppression", ] messages_control.disable = [ "fixme", "invalid-name", "line-too-long", "missing-class-docstring", "missing-function-docstring", "missing-module-docstring", "protected-access", "too-few-public-methods", "too-many-arguments", "too-many-branches", "too-many-lines", "too-many-locals", "too-many-return-statements", "too-many-statements", "too-many-positional-arguments", "wrong-import-position", "duplicate-code", "import-outside-toplevel", "import-error", "disallowed-name", "cyclic-import", "no-self-use", "unnecessary-lambda", "unnecessary-pass", "too-many-ancestors", "redefined-outer-name", "unused-argument", "no-else-return", "redefined-builtin", "too-many-boolean-expressions", "too-many-public-methods", "inconsistent-return-statements", "reimported", "cast_python_value", "unused-variable", "unsubscriptable-object", "abstract-method", "no-member", "consider-using-enumerate", "global-statement", "function-redefined", "ungrouped-imports", "unused-import", "pointless-string-statement", "useless-option-value", "cast_python_value", "unknown-option-value", "no-else-raise", "unidiomatic-typecheck", "possibly-used-before-assignment", # TODO: Fix this "used-before-assignment", # TODO: Fix this ] [tool.pytest.ini_options] minversion = "6.0" xfail_strict = true addopts = [ "-ra", "--strict-markers", "--strict-config", ] testpaths = [ "tests", ] markers = [ "slow", "numba", "awkward", "dis", "sympy", ] log_cli_level = "DEBUG" filterwarnings = [ "error", "ignore::DeprecationWarning", "ignore::UserWarning", ] [tool.mypy] files = [ "src/vector", ] python_version = "3.8" strict = true warn_return_any = false enable_error_code = [ "ignore-without-code", "truthy-bool", "redundant-expr", ] warn_unreachable = false [[tool.mypy.overrides]] module = [ "vector._compute.*.*", "numba.*", "awkward.*", ] ignore_missing_imports = true disallow_untyped_defs = false disallow_untyped_calls = false vector-1.6.3/src/000077500000000000000000000000001503546127100136165ustar00rootroot00000000000000vector-1.6.3/src/vector/000077500000000000000000000000001503546127100151205ustar00rootroot00000000000000vector-1.6.3/src/vector/__init__.py000066400000000000000000000107731503546127100172410ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import importlib.metadata import typing import packaging.version from vector._methods import ( Azimuthal, AzimuthalRhoPhi, AzimuthalXY, Coordinates, Longitudinal, LongitudinalEta, LongitudinalTheta, LongitudinalZ, Lorentz, Momentum, Planar, Spatial, Temporal, TemporalT, TemporalTau, Vector, Vector2D, Vector3D, Vector4D, dim, ) from vector._version import version as __version__ from vector.backends.awkward_constructors import Array, zip from vector.backends.awkward_constructors import Array as awk from vector.backends.numpy import ( MomentumNumpy2D, MomentumNumpy3D, MomentumNumpy4D, VectorNumpy, VectorNumpy2D, VectorNumpy3D, VectorNumpy4D, array, ) from vector.backends.numpy import array as arr from vector.backends.object import ( MomentumObject2D, MomentumObject3D, MomentumObject4D, VectorObject, VectorObject2D, VectorObject3D, VectorObject4D, obj, ) def _import_awkward() -> None: awk_version = packaging.version.Version(importlib.metadata.version("awkward")) if awk_version < packaging.version.Version("2.0.0"): # the only context users will see this message is if they're trying to use vector.awk # VectorAwkward is still set to None msg = f"awkward {awk_version} is too old; please upgrade to 2.0.0 or later" raise ImportError(msg) try: import awkward _import_awkward() except ImportError: awkward = None if not typing.TYPE_CHECKING: VectorAwkward = None else: from vector.backends.awkward import VectorAwkward, awkward_transform try: import sympy # type: ignore[import-untyped] except ImportError: sympy = None if not typing.TYPE_CHECKING: VectorSympy = None else: from vector.backends.sympy import ( MomentumSympy2D, MomentumSympy3D, MomentumSympy4D, VectorSympy, VectorSympy2D, VectorSympy3D, VectorSympy4D, ) __all__: tuple[str, ...] = ( "Array", "Azimuthal", "AzimuthalRhoPhi", "AzimuthalXY", "Coordinates", "Longitudinal", "LongitudinalEta", "LongitudinalTheta", "LongitudinalZ", "Lorentz", "Momentum", "MomentumNumpy2D", "MomentumNumpy3D", "MomentumNumpy4D", "MomentumObject2D", "MomentumObject3D", "MomentumObject4D", "MomentumSympy2D", "MomentumSympy3D", "MomentumSympy4D", "Planar", "Spatial", "Temporal", "TemporalT", "TemporalTau", "Vector", "Vector2D", "Vector3D", "Vector4D", "VectorAwkward", "VectorNumpy", "VectorNumpy2D", "VectorNumpy3D", "VectorNumpy4D", "VectorObject", "VectorObject2D", "VectorObject3D", "VectorObject4D", "VectorSympy", "VectorSympy2D", "VectorSympy3D", "VectorSympy4D", "__version__", "arr", "array", "awk", "awkward_transform", "dim", "obj", "register_awkward", "register_numba", "zip", ) def __dir__() -> tuple[str, ...]: return ( tuple(s for s in __all__ if s != "VectorAwkward") if awkward is None else __all__ ) def register_numba() -> None: """ Make Vector types known to Numba's compiler, so that JIT-compilations with vector arguments do not fail. This usually isn't necessary, as it is passed to Numba's ``entry_point`` and is therefore executed as soon as Numba is imported. """ import vector.backends._numba_object import vector.backends.numba_numpy # noqa: F401 _awkward_registered = False def register_awkward() -> None: """ Make Vector behaviors known to Awkward Array's ``ak.behavior`` mechanism. If you call this function, any records named ``Vector2D``, ``Vector3D``, ``Vector4D``, ``Momentum2D``, ``Momentum3D``, and ``Momentum4D`` will have vector properties and methods. If you do not call this function, only arrays created with :func:`vector.Array` (or derived from such an array) have vector properties and methods. """ import awkward import vector.backends.awkward global _awkward_registered # noqa: PLW0603 awkward.behavior.update(vector.backends.awkward.behavior) _awkward_registered = True vector-1.6.3/src/vector/_compute/000077500000000000000000000000001503546127100167335ustar00rootroot00000000000000vector-1.6.3/src/vector/_compute/__init__.py000066400000000000000000000003641503546127100210470ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/src/vector/_compute/lorentz/000077500000000000000000000000001503546127100204305ustar00rootroot00000000000000vector-1.6.3/src/vector/_compute/lorentz/Et.py000066400000000000000000000065101503546127100213540ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.Et(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import Et2, t from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return lib.sqrt(Et2.xy_z_t(lib, x, y, z, t)) def xy_z_tau(lib, x, y, z, tau): return xy_z_t(lib, x, y, z, t.xy_z_tau(lib, x, y, z, tau)) def xy_theta_t(lib, x, y, theta, t): return t * lib.sin(theta) def xy_theta_tau(lib, x, y, theta, tau): return xy_theta_t(lib, x, y, theta, t.xy_theta_tau(lib, x, y, theta, tau)) def xy_eta_t(lib, x, y, eta, t): expmeta = lib.exp(-eta) return t * (2 / (expmeta + 1 / expmeta)) def xy_eta_tau(lib, x, y, eta, tau): return xy_eta_t(lib, x, y, eta, t.xy_eta_tau(lib, x, y, eta, tau)) def rhophi_z_t(lib, rho, phi, z, t): return t * rho / lib.sqrt(rho**2 + z**2) def rhophi_z_tau(lib, rho, phi, z, tau): return rhophi_z_t(lib, rho, phi, z, t.rhophi_z_tau(lib, rho, phi, z, tau)) def rhophi_theta_t(lib, rho, phi, theta, t): return t * lib.sin(theta) def rhophi_theta_tau(lib, rho, phi, theta, tau): return rhophi_theta_t( lib, rho, phi, theta, t.rhophi_theta_tau(lib, rho, phi, theta, tau) ) def rhophi_eta_t(lib, rho, phi, eta, t): expmeta = lib.exp(-eta) return t * (2 / (expmeta + 1 / expmeta)) def rhophi_eta_tau(lib, rho, phi, eta, tau): return rhophi_eta_t(lib, rho, phi, eta, t.rhophi_eta_tau(lib, rho, phi, eta, tau)) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/Et2.py000066400000000000000000000065631503546127100214460ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.Et2(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): pt2 = x**2 + y**2 return t**2 * pt2 / (pt2 + z**2) def xy_z_tau(lib, x, y, z, tau): return xy_z_t(lib, x, y, z, t.xy_z_tau(lib, x, y, z, tau)) def xy_theta_t(lib, x, y, theta, t): return (t * lib.sin(theta)) ** 2 def xy_theta_tau(lib, x, y, theta, tau): return xy_theta_t(lib, x, y, theta, t.xy_theta_tau(lib, x, y, theta, tau)) def xy_eta_t(lib, x, y, eta, t): expmeta = lib.exp(-eta) return (t * (2 / (expmeta + 1 / expmeta))) ** 2 def xy_eta_tau(lib, x, y, eta, tau): return xy_eta_t(lib, x, y, eta, t.xy_eta_tau(lib, x, y, eta, tau)) def rhophi_z_t(lib, rho, phi, z, t): pt2 = rho**2 return t**2 * pt2 / (pt2 + z**2) def rhophi_z_tau(lib, rho, phi, z, tau): return rhophi_z_t(lib, rho, phi, z, t.rhophi_z_tau(lib, rho, phi, z, tau)) def rhophi_theta_t(lib, rho, phi, theta, t): return (t * lib.sin(theta)) ** 2 def rhophi_theta_tau(lib, rho, phi, theta, tau): return rhophi_theta_t( lib, rho, phi, theta, t.rhophi_theta_tau(lib, rho, phi, theta, tau) ) def rhophi_eta_t(lib, rho, phi, eta, t): expmeta = lib.exp(-eta) return (t * (2 / (expmeta + 1 / expmeta))) ** 2 def rhophi_eta_tau(lib, rho, phi, eta, tau): return rhophi_eta_t(lib, rho, phi, eta, t.rhophi_eta_tau(lib, rho, phi, eta, tau)) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/Mt.py000066400000000000000000000063651503546127100213740ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.Mt(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import Mt2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return lib.sqrt(Mt2.xy_z_t(lib, x, y, z, t)) def xy_z_tau(lib, x, y, z, tau): return lib.sqrt(Mt2.xy_z_tau(lib, x, y, z, tau)) def xy_theta_t(lib, x, y, theta, t): return lib.sqrt(Mt2.xy_theta_t(lib, x, y, theta, t)) def xy_theta_tau(lib, x, y, theta, tau): return lib.sqrt(Mt2.xy_theta_tau(lib, x, y, theta, tau)) def xy_eta_t(lib, x, y, eta, t): return lib.sqrt(Mt2.xy_eta_t(lib, x, y, eta, t)) def xy_eta_tau(lib, x, y, eta, tau): return lib.sqrt(Mt2.xy_eta_tau(lib, x, y, eta, tau)) def rhophi_z_t(lib, rho, phi, z, t): return lib.sqrt(Mt2.rhophi_z_t(lib, rho, phi, z, t)) def rhophi_z_tau(lib, rho, phi, z, tau): return lib.sqrt(Mt2.rhophi_z_tau(lib, rho, phi, z, tau)) def rhophi_theta_t(lib, rho, phi, theta, t): return lib.sqrt(Mt2.rhophi_theta_t(lib, rho, phi, theta, t)) def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.sqrt(Mt2.rhophi_theta_tau(lib, rho, phi, theta, tau)) def rhophi_eta_t(lib, rho, phi, eta, t): return lib.sqrt(Mt2.rhophi_eta_t(lib, rho, phi, eta, t)) def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.sqrt(Mt2.rhophi_eta_tau(lib, rho, phi, eta, tau)) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/Mt2.py000066400000000000000000000065641503546127100214570ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.Mt2(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import tau2 from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return t**2 - z**2 def xy_z_tau(lib, x, y, z, tau): return lib.maximum(tau2.xy_z_tau(lib, x, y, z, tau) + x**2 + y**2, 0) def xy_theta_t(lib, x, y, theta, t): return xy_z_t(lib, x, y, z.xy_theta(lib, x, y, theta), t) def xy_theta_tau(lib, x, y, theta, tau): return lib.maximum(tau2.xy_theta_tau(lib, x, y, theta, tau) + x**2 + y**2, 0) def xy_eta_t(lib, x, y, eta, t): return xy_z_t(lib, x, y, z.xy_eta(lib, x, y, eta), t) def xy_eta_tau(lib, x, y, eta, tau): return lib.maximum(tau2.xy_eta_tau(lib, x, y, eta, tau) + x**2 + y**2, 0) def rhophi_z_t(lib, rho, phi, z, t): return t**2 - z**2 def rhophi_z_tau(lib, rho, phi, z, tau): return lib.maximum(tau2.rhophi_z_tau(lib, rho, phi, z, tau) + rho**2, 0) def rhophi_theta_t(lib, rho, phi, theta, t): return rhophi_z_t(lib, rho, phi, z.rhophi_theta(lib, rho, phi, theta), t) def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.maximum(tau2.rhophi_theta_tau(lib, rho, phi, theta, tau) + rho**2, 0) def rhophi_eta_t(lib, rho, phi, eta, t): return rhophi_z_t(lib, rho, phi, z.rhophi_eta(lib, rho, phi, eta), t) def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.maximum(tau2.rhophi_eta_tau(lib, rho, phi, eta, tau) + rho**2, 0) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/__init__.py000066400000000000000000000046631503546127100225520ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Compute functions for lorentz vectors, which is to say 4D. Each function is a module with variants for each coordinate system (or combination of coordinate systems) as functions within the module. Each module has a ``dispatch_map`` (dict) that maps coordinate types to the appropriate function and its return type(s), and a ``dispatch`` (function) uses this information to call the right function and return the right type. The compute functions themselves are restricted to a minimum of Python features: no statements other than assignments and one return, no assumptions about data types. In particular, if statements and loops are not allowed. The tests/test_compute_features.py suite ensures that these rules are followed (though that set of allowed features can be expanded if it doesn't prevent the addition of new backends). """ from __future__ import annotations import vector._compute.lorentz.add import vector._compute.lorentz.beta import vector._compute.lorentz.boost_beta3 import vector._compute.lorentz.boost_p4 import vector._compute.lorentz.boostX_beta import vector._compute.lorentz.boostX_gamma import vector._compute.lorentz.boostY_beta import vector._compute.lorentz.boostY_gamma import vector._compute.lorentz.boostZ_beta import vector._compute.lorentz.boostZ_gamma import vector._compute.lorentz.deltaRapidityPhi import vector._compute.lorentz.deltaRapidityPhi2 import vector._compute.lorentz.dot import vector._compute.lorentz.equal import vector._compute.lorentz.Et import vector._compute.lorentz.Et2 import vector._compute.lorentz.gamma import vector._compute.lorentz.is_lightlike import vector._compute.lorentz.is_spacelike import vector._compute.lorentz.is_timelike import vector._compute.lorentz.isclose import vector._compute.lorentz.Mt import vector._compute.lorentz.Mt2 import vector._compute.lorentz.not_equal import vector._compute.lorentz.rapidity import vector._compute.lorentz.scale import vector._compute.lorentz.subtract import vector._compute.lorentz.t import vector._compute.lorentz.t2 import vector._compute.lorentz.tau import vector._compute.lorentz.tau2 import vector._compute.lorentz.to_beta3 import vector._compute.lorentz.transform4D import vector._compute.lorentz.unit # noqa: F401 vector-1.6.3/src/vector/_compute/lorentz/add.py000066400000000000000000000162101503546127100215320ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.add(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t, tau from vector._compute.spatial import add from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): spatial_add, azimuthal, longitudinal = add.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] if azimuthal1 is AzimuthalXY: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.xy_z_t elif temporal1 is TemporalTau: to_t1 = t.xy_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.xy_theta_t elif temporal1 is TemporalTau: to_t1 = t.xy_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.xy_eta_t elif temporal1 is TemporalTau: to_t1 = t.xy_eta_tau elif azimuthal1 is AzimuthalRhoPhi: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.rhophi_z_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.rhophi_theta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.rhophi_eta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_eta_tau if azimuthal2 is AzimuthalXY: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.xy_z_t elif temporal2 is TemporalTau: to_t2 = t.xy_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.xy_theta_t elif temporal2 is TemporalTau: to_t2 = t.xy_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.xy_eta_t elif temporal2 is TemporalTau: to_t2 = t.xy_eta_tau elif azimuthal2 is AzimuthalRhoPhi: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.rhophi_z_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.rhophi_theta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.rhophi_eta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_eta_tau if temporal1 is TemporalT or temporal2 is TemporalT: def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): coord1, coord2, coord3 = spatial_add( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) t1 = to_t1(lib, coord11, coord12, coord13, coord14) t2 = to_t2(lib, coord21, coord22, coord23, coord24) return (coord1, coord2, coord3, t1 + t2) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, azimuthal, longitudinal, TemporalT) else: if azimuthal is AzimuthalXY: if longitudinal is LongitudinalZ: to_tau = tau.xy_z_t elif longitudinal is LongitudinalTheta: to_tau = tau.xy_theta_t elif longitudinal is LongitudinalEta: to_tau = tau.xy_eta_t elif azimuthal is AzimuthalRhoPhi: if longitudinal is LongitudinalZ: to_tau = tau.rhophi_z_t elif longitudinal is LongitudinalTheta: to_tau = tau.rhophi_theta_t elif longitudinal is LongitudinalEta: to_tau = tau.rhophi_eta_t def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): coord1, coord2, coord3 = spatial_add( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) t1 = to_t1(lib, coord11, coord12, coord13, coord14) t2 = to_t2(lib, coord21, coord22, coord23, coord24) return ( coord1, coord2, coord3, to_tau(lib, coord1, coord2, coord3, t1 + t2), ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = ( f, azimuthal, longitudinal, TemporalTau, ) # type: ignore[assignment] for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/beta.py000066400000000000000000000103461503546127100217210ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.beta(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.lorentz import t from vector._compute.spatial import mag from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return lib.nan_to_num(mag.xy_z(lib, x, y, z) / t, nan=0, posinf=inf, neginf=-inf) def xy_z_tau(lib, x, y, z, tau): return lib.nan_to_num( mag.xy_z(lib, x, y, z) / t.xy_z_tau(lib, x, y, z, tau), nan=0, posinf=inf, neginf=-inf, ) def xy_theta_t(lib, x, y, theta, t): return lib.nan_to_num( mag.xy_theta(lib, x, y, theta) / t, nan=0, posinf=inf, neginf=-inf ) def xy_theta_tau(lib, x, y, theta, tau): return lib.nan_to_num( mag.xy_theta(lib, x, y, theta) / t.xy_theta_tau(lib, x, y, theta, tau), nan=0, posinf=inf, neginf=-inf, ) def xy_eta_t(lib, x, y, eta, t): return lib.nan_to_num( mag.xy_eta(lib, x, y, eta) / t, nan=0, posinf=inf, neginf=-inf ) def xy_eta_tau(lib, x, y, eta, tau): return lib.nan_to_num( mag.xy_eta(lib, x, y, eta) / t.xy_eta_tau(lib, x, y, eta, tau), nan=0, posinf=inf, neginf=-inf, ) def rhophi_z_t(lib, rho, phi, z, t): return lib.nan_to_num( mag.rhophi_z(lib, rho, phi, z) / t, nan=0, posinf=inf, neginf=-inf ) def rhophi_z_tau(lib, rho, phi, z, tau): return lib.nan_to_num( mag.rhophi_z(lib, rho, phi, z) / t.rhophi_z_tau(lib, rho, phi, z, tau), nan=0, posinf=inf, neginf=-inf, ) def rhophi_theta_t(lib, rho, phi, theta, t): return lib.nan_to_num( mag.rhophi_theta(lib, rho, phi, theta) / t, nan=0, posinf=inf, neginf=-inf ) def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.nan_to_num( mag.rhophi_theta(lib, rho, phi, theta) / t.rhophi_theta_tau(lib, rho, phi, theta, tau), nan=0, posinf=inf, neginf=-inf, ) def rhophi_eta_t(lib, rho, phi, eta, t): return lib.nan_to_num( mag.rhophi_eta(lib, rho, phi, eta) / t, nan=0, posinf=inf, neginf=-inf ) def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.nan_to_num( mag.rhophi_eta(lib, rho, phi, eta) / t.rhophi_eta_tau(lib, rho, phi, eta, tau), nan=0, posinf=inf, neginf=-inf, ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boostX_beta.py000066400000000000000000000142631503546127100232610ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boostX(self, beta=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, beta, x, y, z, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def xy_z_tau(lib, beta, x, y, z, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z tee = t.xy_z_tau(lib, x, y, z, tau) return (gam * exx + bgam * tee, why, zee, tau) def xy_theta_t(lib, beta, x, y, theta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def xy_theta_tau(lib, beta, x, y, theta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t.xy_theta_tau(lib, x, y, theta, tau) return (gam * exx + bgam * tee, why, zee, tau) def xy_eta_t(lib, beta, x, y, eta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def xy_eta_tau(lib, beta, x, y, eta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t.xy_eta_tau(lib, x, y, eta, tau) return (gam * exx + bgam * tee, why, zee, tau) def rhophi_z_t(lib, beta, rho, phi, z, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def rhophi_z_tau(lib, beta, rho, phi, z, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t.rhophi_z_tau(lib, rho, phi, z, tau) return (gam * exx + bgam * tee, why, zee, tau) def rhophi_theta_t(lib, beta, rho, phi, theta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def rhophi_theta_tau(lib, beta, rho, phi, theta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t.rhophi_theta_tau(lib, rho, phi, theta, tau) return (gam * exx + bgam * tee, why, zee, tau) def rhophi_eta_t(lib, beta, rho, phi, eta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def rhophi_eta_tau(lib, beta, rho, phi, eta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t.rhophi_eta_tau(lib, rho, phi, eta, tau) return (gam * exx + bgam * tee, why, zee, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), } def dispatch(beta: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, beta, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boostX_gamma.py000066400000000000000000000150331503546127100234240ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boostX(self, gamma=) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, gamma, x, y, z, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def xy_z_tau(lib, gamma, x, y, z, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z tee = t.xy_z_tau(lib, x, y, z, tau) return (gam * exx + bgam * tee, why, zee, tau) def xy_theta_t(lib, gamma, x, y, theta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def xy_theta_tau(lib, gamma, x, y, theta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t.xy_theta_tau(lib, x, y, theta, tau) return (gam * exx + bgam * tee, why, zee, tau) def xy_eta_t(lib, gamma, x, y, eta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def xy_eta_tau(lib, gamma, x, y, eta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t.xy_eta_tau(lib, x, y, eta, tau) return (gam * exx + bgam * tee, why, zee, tau) def rhophi_z_t(lib, gamma, rho, phi, z, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def rhophi_z_tau(lib, gamma, rho, phi, z, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t.rhophi_z_tau(lib, rho, phi, z, tau) return (gam * exx + bgam * tee, why, zee, tau) def rhophi_theta_t(lib, gamma, rho, phi, theta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def rhophi_theta_tau(lib, gamma, rho, phi, theta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t.rhophi_theta_tau(lib, rho, phi, theta, tau) return (gam * exx + bgam * tee, why, zee, tau) def rhophi_eta_t(lib, gamma, rho, phi, eta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t return (gam * exx + bgam * tee, why, zee, bgam * exx + gam * tee) def rhophi_eta_tau(lib, gamma, rho, phi, eta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t.rhophi_eta_tau(lib, rho, phi, eta, tau) return (gam * exx + bgam * tee, why, zee, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), } def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, gamma, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boostY_beta.py000066400000000000000000000142631503546127100232620ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boostY(self, beta=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, beta, x, y, z, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def xy_z_tau(lib, beta, x, y, z, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z tee = t.xy_z_tau(lib, x, y, z, tau) return (exx, gam * why + bgam * tee, zee, tau) def xy_theta_t(lib, beta, x, y, theta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def xy_theta_tau(lib, beta, x, y, theta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t.xy_theta_tau(lib, x, y, theta, tau) return (exx, gam * why + bgam * tee, zee, tau) def xy_eta_t(lib, beta, x, y, eta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def xy_eta_tau(lib, beta, x, y, eta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t.xy_eta_tau(lib, x, y, eta, tau) return (exx, gam * why + bgam * tee, zee, tau) def rhophi_z_t(lib, beta, rho, phi, z, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def rhophi_z_tau(lib, beta, rho, phi, z, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t.rhophi_z_tau(lib, rho, phi, z, tau) return (exx, gam * why + bgam * tee, zee, tau) def rhophi_theta_t(lib, beta, rho, phi, theta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def rhophi_theta_tau(lib, beta, rho, phi, theta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t.rhophi_theta_tau(lib, rho, phi, theta, tau) return (exx, gam * why + bgam * tee, zee, tau) def rhophi_eta_t(lib, beta, rho, phi, eta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def rhophi_eta_tau(lib, beta, rho, phi, eta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t.rhophi_eta_tau(lib, rho, phi, eta, tau) return (exx, gam * why + bgam * tee, zee, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), } def dispatch(beta: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, beta, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boostY_gamma.py000066400000000000000000000150361503546127100234300ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boostY(self, gamma=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, gamma, x, y, z, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def xy_z_tau(lib, gamma, x, y, z, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z tee = t.xy_z_tau(lib, x, y, z, tau) return (exx, gam * why + bgam * tee, zee, tau) def xy_theta_t(lib, gamma, x, y, theta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def xy_theta_tau(lib, gamma, x, y, theta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_theta(lib, x, y, theta) tee = t.xy_theta_tau(lib, x, y, theta, tau) return (exx, gam * why + bgam * tee, zee, tau) def xy_eta_t(lib, gamma, x, y, eta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def xy_eta_tau(lib, gamma, x, y, eta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x why = y zee = z.xy_eta(lib, x, y, eta) tee = t.xy_eta_tau(lib, x, y, eta, tau) return (exx, gam * why + bgam * tee, zee, tau) def rhophi_z_t(lib, gamma, rho, phi, z, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def rhophi_z_tau(lib, gamma, rho, phi, z, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z tee = t.rhophi_z_tau(lib, rho, phi, z, tau) return (exx, gam * why + bgam * tee, zee, tau) def rhophi_theta_t(lib, gamma, rho, phi, theta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def rhophi_theta_tau(lib, gamma, rho, phi, theta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t.rhophi_theta_tau(lib, rho, phi, theta, tau) return (exx, gam * why + bgam * tee, zee, tau) def rhophi_eta_t(lib, gamma, rho, phi, eta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t return (exx, gam * why + bgam * tee, zee, bgam * why + gam * tee) def rhophi_eta_tau(lib, gamma, rho, phi, eta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) exx = x.rhophi(lib, rho, phi) why = y.rhophi(lib, rho, phi) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t.rhophi_eta_tau(lib, rho, phi, eta, tau) return (exx, gam * why + bgam * tee, zee, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), } def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, gamma, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boostZ_beta.py000066400000000000000000000131431503546127100232570ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boostZ(self, beta=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, beta, x, y, z, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z tee = t return (x, y, gam * zee + bgam * tee, bgam * zee + gam * tee) def xy_z_tau(lib, beta, x, y, z, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z tee = t.xy_z_tau(lib, x, y, z, tau) return (x, y, gam * zee + bgam * tee, tau) def xy_theta_t(lib, beta, x, y, theta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.xy_theta(lib, x, y, theta) tee = t return (x, y, gam * zee + bgam * tee, bgam * zee + gam * tee) def xy_theta_tau(lib, beta, x, y, theta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.xy_theta(lib, x, y, theta) tee = t.xy_theta_tau(lib, x, y, theta, tau) return (x, y, gam * zee + bgam * tee, tau) def xy_eta_t(lib, beta, x, y, eta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.xy_eta(lib, x, y, eta) tee = t return (x, y, gam * zee + bgam * tee, bgam * zee + gam * tee) def xy_eta_tau(lib, beta, x, y, eta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.xy_eta(lib, x, y, eta) tee = t.xy_eta_tau(lib, x, y, eta, tau) return (x, y, gam * zee + bgam * tee, tau) def rhophi_z_t(lib, beta, rho, phi, z, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z tee = t return (rho, phi, gam * zee + bgam * tee, bgam * zee + gam * tee) def rhophi_z_tau(lib, beta, rho, phi, z, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z tee = t.rhophi_z_tau(lib, rho, phi, z, tau) return (rho, phi, gam * zee + bgam * tee, tau) def rhophi_theta_t(lib, beta, rho, phi, theta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.rhophi_theta(lib, rho, phi, theta) tee = t return (rho, phi, gam * zee + bgam * tee, bgam * zee + gam * tee) def rhophi_theta_tau(lib, beta, rho, phi, theta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.rhophi_theta(lib, rho, phi, theta) tee = t.rhophi_theta_tau(lib, rho, phi, theta, tau) return (rho, phi, gam * zee + bgam * tee, tau) def rhophi_eta_t(lib, beta, rho, phi, eta, t): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.rhophi_eta(lib, rho, phi, eta) tee = t return (rho, phi, gam * zee + bgam * tee, bgam * zee + gam * tee) def rhophi_eta_tau(lib, beta, rho, phi, eta, tau): gam = (1 - beta**2) ** -0.5 bgam = beta * gam zee = z.rhophi_eta(lib, rho, phi, eta) tee = t.rhophi_eta_tau(lib, rho, phi, eta, tau) return (rho, phi, gam * zee + bgam * tee, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), } def dispatch(beta: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, beta, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boostZ_gamma.py000066400000000000000000000137161503546127100234340ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boostZ(self, gamma=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, gamma, x, y, z, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z tee = t return (x, y, gam * zee + bgam * tee, bgam * zee + gam * tee) def xy_z_tau(lib, gamma, x, y, z, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z tee = t.xy_z_tau(lib, x, y, z, tau) return (x, y, gam * zee + bgam * tee, tau) def xy_theta_t(lib, gamma, x, y, theta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.xy_theta(lib, x, y, theta) tee = t return (x, y, gam * zee + bgam * tee, bgam * zee + gam * tee) def xy_theta_tau(lib, gamma, x, y, theta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.xy_theta(lib, x, y, theta) tee = t.xy_theta_tau(lib, x, y, theta, tau) return (x, y, gam * zee + bgam * tee, tau) def xy_eta_t(lib, gamma, x, y, eta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.xy_eta(lib, x, y, eta) tee = t return (x, y, gam * zee + bgam * tee, bgam * zee + gam * tee) def xy_eta_tau(lib, gamma, x, y, eta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.xy_eta(lib, x, y, eta) tee = t.xy_eta_tau(lib, x, y, eta, tau) return (x, y, gam * zee + bgam * tee, tau) def rhophi_z_t(lib, gamma, rho, phi, z, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z tee = t return (rho, phi, gam * zee + bgam * tee, bgam * zee + gam * tee) def rhophi_z_tau(lib, gamma, rho, phi, z, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z tee = t.rhophi_z_tau(lib, rho, phi, z, tau) return (rho, phi, gam * zee + bgam * tee, tau) def rhophi_theta_t(lib, gamma, rho, phi, theta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t return (rho, phi, gam * zee + bgam * tee, bgam * zee + gam * tee) def rhophi_theta_tau(lib, gamma, rho, phi, theta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.rhophi_theta(lib, rho, phi, theta) tee = t.rhophi_theta_tau(lib, rho, phi, theta, tau) return (rho, phi, gam * zee + bgam * tee, tau) def rhophi_eta_t(lib, gamma, rho, phi, eta, t): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t return (rho, phi, gam * zee + bgam * tee, bgam * zee + gam * tee) def rhophi_eta_tau(lib, gamma, rho, phi, eta, tau): gam = lib.absolute(gamma) bgam = lib.copysign(lib.sqrt(gam**2 - 1), gamma) zee = z.rhophi_eta(lib, rho, phi, eta) tee = t.rhophi_eta_tau(lib, rho, phi, eta, tau) return (rho, phi, gam * zee + bgam * tee, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), } def dispatch(gamma: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, gamma, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boost_beta3.py000066400000000000000000000236521503546127100232160ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boost_beta3(self, beta3) or .. code-block:: python Lorentz.boost(self, beta3=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import transform4D from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) def cartesian_t(lib, x1, y1, z1, t1, betax, betay, betaz): bp2 = betax**2 + betay**2 + betaz**2 gamma = 1 / lib.sqrt(1 - bp2) bgam = gamma**2 / (1 + gamma) xx = 1 + bgam * betax * betax yy = 1 + bgam * betay * betay zz = 1 + bgam * betaz * betaz xy = bgam * betax * betay xz = bgam * betax * betaz yz = bgam * betay * betaz yx = xy zx = xz zy = yz xt = gamma * betax yt = gamma * betay zt = gamma * betaz tx = xt ty = yt tz = zt tt = gamma # fmt: off return transform4D.cartesian_t( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, tx, ty, tz, tt, x1, y1, z1, t1, ) # fmt: on def cartesian_tau(lib, x1, y1, z1, tau1, betax, betay, betaz): bp2 = betax**2 + betay**2 + betaz**2 gamma = 1 / lib.sqrt(1 - bp2) bgam = gamma**2 / (1 + gamma) xx = 1 + bgam * betax * betax yy = 1 + bgam * betay * betay zz = 1 + bgam * betaz * betaz xy = bgam * betax * betay xz = bgam * betax * betaz yz = bgam * betay * betaz yx = xy zx = xz zy = yz xt = gamma * betax yt = gamma * betay zt = gamma * betaz # fmt: off return transform4D.cartesian_tau( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, x1, y1, z1, tau1, ) # fmt: on def cartesian_t_xy_z(lib, x1, y1, z1, t1, x2, y2, z2): return cartesian_t(lib, x1, y1, z1, t1, x2, y2, z2) def cartesian_tau_xy_z(lib, x1, y1, z1, tau1, x2, y2, z2): return cartesian_tau(lib, x1, y1, z1, tau1, x2, y2, z2) def cartesian_t_xy_theta(lib, x1, y1, z1, t1, x2, y2, theta2): return cartesian_t(lib, x1, y1, z1, t1, x2, y2, z.xy_theta(lib, x2, y2, theta2)) def cartesian_tau_xy_theta(lib, x1, y1, z1, tau1, x2, y2, theta2): return cartesian_tau(lib, x1, y1, z1, tau1, x2, y2, z.xy_theta(lib, x2, y2, theta2)) def cartesian_t_xy_eta(lib, x1, y1, z1, t1, x2, y2, eta2): return cartesian_t(lib, x1, y1, z1, t1, x2, y2, z.xy_eta(lib, x2, y2, eta2)) def cartesian_tau_xy_eta(lib, x1, y1, z1, tau1, x2, y2, eta2): return cartesian_tau(lib, x1, y1, z1, tau1, x2, y2, z.xy_eta(lib, x2, y2, eta2)) def cartesian_t_rhophi_z(lib, x1, y1, z1, t1, rho2, phi2, z2): return cartesian_t( lib, x1, y1, z1, t1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2 ) def cartesian_tau_rhophi_z(lib, x1, y1, z1, tau1, rho2, phi2, z2): return cartesian_tau( lib, x1, y1, z1, tau1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2 ) def cartesian_t_rhophi_theta(lib, x1, y1, z1, t1, rho2, phi2, theta2): return cartesian_t( lib, x1, y1, z1, t1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def cartesian_tau_rhophi_theta(lib, x1, y1, z1, tau1, rho2, phi2, theta2): return cartesian_tau( lib, x1, y1, z1, tau1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def cartesian_t_rhophi_eta(lib, x1, y1, z1, t1, rho2, phi2, eta2): return cartesian_t( lib, x1, y1, z1, t1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def cartesian_tau_rhophi_eta(lib, x1, y1, z1, tau1, rho2, phi2, eta2): return cartesian_tau( lib, x1, y1, z1, tau1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalZ): ( cartesian_t_xy_z, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalZ): ( cartesian_tau_xy_z, AzimuthalXY, LongitudinalZ, TemporalTau, ), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalTheta, ): ( cartesian_t_xy_theta, AzimuthalXY, LongitudinalZ, TemporalT, ), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalTheta, ): ( cartesian_tau_xy_theta, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalEta): ( cartesian_t_xy_eta, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalEta): ( cartesian_tau_xy_eta, AzimuthalXY, LongitudinalZ, TemporalTau, ), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalZ, ): ( cartesian_t_rhophi_z, AzimuthalXY, LongitudinalZ, TemporalT, ), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalZ, ): ( cartesian_tau_rhophi_z, AzimuthalXY, LongitudinalZ, TemporalTau, ), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalTheta, ): ( cartesian_t_rhophi_theta, AzimuthalXY, LongitudinalZ, TemporalT, ), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalTheta, ): ( cartesian_tau_rhophi_theta, AzimuthalXY, LongitudinalZ, TemporalTau, ), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalEta, ): ( cartesian_t_rhophi_eta, AzimuthalXY, LongitudinalZ, TemporalT, ), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalEta, ): ( cartesian_tau_rhophi_eta, AzimuthalXY, LongitudinalZ, TemporalTau, ), } def make_conversion(azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2): if (azimuthal1, longitudinal1, temporal1) != ( AzimuthalXY, LongitudinalZ, TemporalT, ): if azimuthal1 is AzimuthalXY: to_x = x.xy to_y = y.xy if longitudinal1 is LongitudinalZ: to_z = z.xy_z elif longitudinal1 is LongitudinalTheta: to_z = z.xy_theta elif longitudinal1 is LongitudinalEta: to_z = z.xy_eta elif azimuthal1 is AzimuthalRhoPhi: to_x = x.rhophi to_y = y.rhophi if longitudinal1 is LongitudinalZ: to_z = z.rhophi_z elif longitudinal1 is LongitudinalTheta: to_z = z.rhophi_theta elif longitudinal1 is LongitudinalEta: to_z = z.rhophi_eta cartesian, azout, lout, tout = dispatch_map[ AzimuthalXY, LongitudinalZ, temporal1, azimuthal2, longitudinal2 ] def f(lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23): return cartesian( lib, to_x(lib, coord11, coord12), to_y(lib, coord11, coord12), to_z(lib, coord11, coord12, coord13), coord14, coord21, coord22, coord23, ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2 ] = (f, azout, lout, tout) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/boost_p4.py000066400000000000000000000514431503546127100225420ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.boost_p4(self, p4) or .. code-block:: python Lorentz.boost(self, p4=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import transform4D from vector._compute.planar import x, y from vector._compute.spatial import mag2, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) def cartesian_t(lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z2): gamma = energy / mass mass2_gamma_1 = mass2 * (gamma + 1) gbetax = x2 / mass gbetay = y2 / mass gbetaz = z2 / mass xx = 1 + (x2 * x2) / mass2_gamma_1 yy = 1 + (y2 * y2) / mass2_gamma_1 zz = 1 + (z2 * z2) / mass2_gamma_1 xy = (x2 * y2) / mass2_gamma_1 xz = (x2 * z2) / mass2_gamma_1 yz = (y2 * z2) / mass2_gamma_1 yx = xy zx = xz zy = yz xt = gbetax yt = gbetay zt = gbetaz tx = xt ty = yt tz = zt tt = gamma # fmt: off return transform4D.cartesian_t( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, tx, ty, tz, tt, x1, y1, z1, t1, ) # fmt: on def cartesian_tau(lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z2): gamma = energy / mass mass2_gamma_1 = mass2 * (gamma + 1) gbetax = x2 / mass gbetay = y2 / mass gbetaz = z2 / mass xx = 1 + (x2 * x2) / mass2_gamma_1 yy = 1 + (y2 * y2) / mass2_gamma_1 zz = 1 + (z2 * z2) / mass2_gamma_1 xy = (x2 * y2) / mass2_gamma_1 xz = (x2 * z2) / mass2_gamma_1 yz = (y2 * z2) / mass2_gamma_1 yx = xy zx = xz zy = yz xt = gbetax yt = gbetay zt = gbetaz # fmt: off return transform4D.cartesian_tau( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, x1, y1, z1, tau1, ) # fmt: on def cartesian_t_xy_z_t(lib, x1, y1, z1, t1, x2, y2, z2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.xy_z(lib, x2, y2, z2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_t(lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z2) def cartesian_tau_xy_z_t(lib, x1, y1, z1, tau1, x2, y2, z2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.xy_z(lib, x2, y2, z2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau(lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z2) def cartesian_t_xy_z_tau(lib, x1, y1, z1, t1, x2, y2, z2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.xy_z(lib, x2, y2, z2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_t(lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z2) def cartesian_tau_xy_z_tau(lib, x1, y1, z1, tau1, x2, y2, z2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.xy_z(lib, x2, y2, z2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau(lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z2) def cartesian_t_xy_theta_t(lib, x1, y1, z1, t1, x2, y2, theta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.xy_theta(lib, x2, y2, theta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def cartesian_tau_xy_theta_t(lib, x1, y1, z1, tau1, x2, y2, theta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.xy_theta(lib, x2, y2, theta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def cartesian_t_xy_theta_tau(lib, x1, y1, z1, t1, x2, y2, theta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.xy_theta(lib, x2, y2, theta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def cartesian_tau_xy_theta_tau(lib, x1, y1, z1, tau1, x2, y2, theta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.xy_theta(lib, x2, y2, theta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def cartesian_t_xy_eta_t(lib, x1, y1, z1, t1, x2, y2, eta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.xy_eta(lib, x2, y2, eta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z.xy_eta(lib, x2, y2, eta2) ) def cartesian_tau_xy_eta_t(lib, x1, y1, z1, tau1, x2, y2, eta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.xy_eta(lib, x2, y2, eta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z.xy_eta(lib, x2, y2, eta2) ) def cartesian_t_xy_eta_tau(lib, x1, y1, z1, t1, x2, y2, eta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.xy_eta(lib, x2, y2, eta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x2, y2, z.xy_eta(lib, x2, y2, eta2) ) def cartesian_tau_xy_eta_tau(lib, x1, y1, z1, tau1, x2, y2, eta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.xy_eta(lib, x2, y2, eta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x2, y2, z.xy_eta(lib, x2, y2, eta2) ) def cartesian_t_rhophi_z_t(lib, x1, y1, z1, t1, rho2, phi2, z2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.rhophi_z(lib, rho2, phi2, z2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def cartesian_tau_rhophi_z_t(lib, x1, y1, z1, tau1, rho2, phi2, z2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.rhophi_z(lib, rho2, phi2, z2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def cartesian_t_rhophi_z_tau(lib, x1, y1, z1, t1, rho2, phi2, z2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.rhophi_z(lib, rho2, phi2, z2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def cartesian_tau_rhophi_z_tau(lib, x1, y1, z1, tau1, rho2, phi2, z2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.rhophi_z(lib, rho2, phi2, z2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def cartesian_t_rhophi_theta_t(lib, x1, y1, z1, t1, rho2, phi2, theta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.rhophi_theta(lib, rho2, phi2, theta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def cartesian_tau_rhophi_theta_t(lib, x1, y1, z1, tau1, rho2, phi2, theta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.rhophi_theta(lib, rho2, phi2, theta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def cartesian_t_rhophi_theta_tau(lib, x1, y1, z1, t1, rho2, phi2, theta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.rhophi_theta(lib, rho2, phi2, theta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def cartesian_tau_rhophi_theta_tau(lib, x1, y1, z1, tau1, rho2, phi2, theta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.rhophi_theta(lib, rho2, phi2, theta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def cartesian_t_rhophi_eta_t(lib, x1, y1, z1, t1, rho2, phi2, eta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.rhophi_eta(lib, rho2, phi2, eta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def cartesian_tau_rhophi_eta_t(lib, x1, y1, z1, tau1, rho2, phi2, eta2, t2): energy = t2 energy2 = energy**2 mass2 = energy2 - mag2.rhophi_eta(lib, rho2, phi2, eta2) mass = lib.sqrt(mass2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def cartesian_t_rhophi_eta_tau(lib, x1, y1, z1, t1, rho2, phi2, eta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.rhophi_eta(lib, rho2, phi2, eta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_t( lib, x1, y1, z1, t1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def cartesian_tau_rhophi_eta_tau(lib, x1, y1, z1, tau1, rho2, phi2, eta2, tau2): mass = tau2 mass2 = mass**2 energy2 = mass2 + mag2.rhophi_eta(lib, rho2, phi2, eta2) energy = lib.sqrt(energy2) # NaN for spacelike boosts propagates everywhere! return cartesian_tau( lib, x1, y1, z1, tau1, energy, mass, mass2, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalZ, TemporalT): ( cartesian_t_xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalZ, TemporalT): ( cartesian_tau_xy_z_t, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalZ, TemporalTau): ( cartesian_t_xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalT, ), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalZ, TemporalTau, ): ( cartesian_tau_xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalTheta, TemporalT, ): (cartesian_t_xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalTheta, TemporalT, ): (cartesian_tau_xy_theta_t, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalTheta, TemporalTau, ): (cartesian_t_xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalTheta, TemporalTau, ): (cartesian_tau_xy_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau), (AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalEta, TemporalT): ( cartesian_t_xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalT, ), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalEta, TemporalT, ): ( cartesian_tau_xy_eta_t, AzimuthalXY, LongitudinalZ, TemporalTau, ), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalXY, LongitudinalEta, TemporalTau, ): (cartesian_t_xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalXY, LongitudinalEta, TemporalTau, ): (cartesian_tau_xy_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ): (cartesian_t_rhophi_z_t, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ): (cartesian_tau_rhophi_z_t, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ): (cartesian_t_rhophi_z_tau, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ): (cartesian_tau_rhophi_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalTheta, TemporalT, ): (cartesian_t_rhophi_theta_t, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalTheta, TemporalT, ): (cartesian_tau_rhophi_theta_t, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalTheta, TemporalTau, ): (cartesian_t_rhophi_theta_tau, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalTheta, TemporalTau, ): (cartesian_tau_rhophi_theta_tau, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalEta, TemporalT, ): (cartesian_t_rhophi_eta_t, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalEta, TemporalT, ): (cartesian_tau_rhophi_eta_t, AzimuthalXY, LongitudinalZ, TemporalTau), ( AzimuthalXY, LongitudinalZ, TemporalT, AzimuthalRhoPhi, LongitudinalEta, TemporalTau, ): (cartesian_t_rhophi_eta_tau, AzimuthalXY, LongitudinalZ, TemporalT), ( AzimuthalXY, LongitudinalZ, TemporalTau, AzimuthalRhoPhi, LongitudinalEta, TemporalTau, ): (cartesian_tau_rhophi_eta_tau, AzimuthalXY, LongitudinalZ, TemporalTau), } def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): if (azimuthal1, longitudinal1, temporal1) != ( AzimuthalXY, LongitudinalZ, TemporalT, ): if azimuthal1 is AzimuthalXY: to_x = x.xy to_y = y.xy if longitudinal1 is LongitudinalZ: to_z = z.xy_z elif longitudinal1 is LongitudinalTheta: to_z = z.xy_theta elif longitudinal1 is LongitudinalEta: to_z = z.xy_eta elif azimuthal1 is AzimuthalRhoPhi: to_x = x.rhophi to_y = y.rhophi if longitudinal1 is LongitudinalZ: to_z = z.rhophi_z elif longitudinal1 is LongitudinalTheta: to_z = z.rhophi_theta elif longitudinal1 is LongitudinalEta: to_z = z.rhophi_eta cartesian, azout, lout, tout = dispatch_map[ AzimuthalXY, LongitudinalZ, temporal1, azimuthal2, longitudinal2, temporal2 ] def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): return cartesian( lib, to_x(lib, coord11, coord12), to_y(lib, coord11, coord12), to_z(lib, coord11, coord12, coord13), coord14, coord21, coord22, coord23, coord24, ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, azout, lout, tout) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/deltaRapidityPhi.py000066400000000000000000000062511503546127100242460ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.deltaRapidityPhi(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import deltaRapidityPhi2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): lorentz_deltaRapidityPhi2, _ = deltaRapidityPhi2.dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24, ): return lib.sqrt( lorentz_deltaRapidityPhi2( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24, ) ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, float) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch( v1: typing.Any, v2: typing.Any, ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/deltaRapidityPhi2.py000066400000000000000000000064351503546127100243340ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.deltaRapidityPhi2(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import rapidity from vector._compute.planar import deltaphi from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): planar_deltaphi, _ = deltaphi.dispatch_map[azimuthal1, azimuthal2] lorentz_rapidity1, _ = rapidity.dispatch_map[azimuthal1, longitudinal1, temporal1] lorentz_rapidity2, _ = rapidity.dispatch_map[azimuthal2, longitudinal2, temporal2] def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24, ): return ( planar_deltaphi(lib, coord11, coord12, coord21, coord22) ** 2 + ( lorentz_rapidity1(lib, coord11, coord12, coord13, coord14) - lorentz_rapidity2(lib, coord21, coord22, coord23, coord24) ) ** 2 ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, float) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch( v1: typing.Any, v2: typing.Any, ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/dot.py000066400000000000000000000127711503546127100216000ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.dot(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import dot from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): spatial_dot, _ = dot.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] if azimuthal1 is AzimuthalXY: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.xy_z_t elif temporal1 is TemporalTau: to_t1 = t.xy_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.xy_theta_t elif temporal1 is TemporalTau: to_t1 = t.xy_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.xy_eta_t elif temporal1 is TemporalTau: to_t1 = t.xy_eta_tau elif azimuthal1 is AzimuthalRhoPhi: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.rhophi_z_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.rhophi_theta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.rhophi_eta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_eta_tau if azimuthal2 is AzimuthalXY: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.xy_z_t elif temporal2 is TemporalTau: to_t2 = t.xy_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.xy_theta_t elif temporal2 is TemporalTau: to_t2 = t.xy_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.xy_eta_t elif temporal2 is TemporalTau: to_t2 = t.xy_eta_tau elif azimuthal2 is AzimuthalRhoPhi: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.rhophi_z_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.rhophi_theta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.rhophi_eta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_eta_tau def f(lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24): return ( to_t1(lib, coord11, coord12, coord13, coord14) * to_t2(lib, coord21, coord22, coord23, coord24) ) - spatial_dot(lib, coord11, coord12, coord13, coord21, coord22, coord23) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, float) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/equal.py000066400000000000000000000135341503546127100221170ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.equal(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import equal from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): spatial_equal, _ = equal.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] if azimuthal1 is AzimuthalXY: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.xy_z_t elif temporal1 is TemporalTau: to_t1 = t.xy_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.xy_theta_t elif temporal1 is TemporalTau: to_t1 = t.xy_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.xy_eta_t elif temporal1 is TemporalTau: to_t1 = t.xy_eta_tau elif azimuthal1 is AzimuthalRhoPhi: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.rhophi_z_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.rhophi_theta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.rhophi_eta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_eta_tau if azimuthal2 is AzimuthalXY: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.xy_z_t elif temporal2 is TemporalTau: to_t2 = t.xy_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.xy_theta_t elif temporal2 is TemporalTau: to_t2 = t.xy_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.xy_eta_t elif temporal2 is TemporalTau: to_t2 = t.xy_eta_tau elif azimuthal2 is AzimuthalRhoPhi: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.rhophi_z_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.rhophi_theta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.rhophi_eta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_eta_tau if temporal1 == temporal2: def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): return (coord14 == coord24) & spatial_equal( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) else: def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): return ( to_t1(lib, coord11, coord12, coord13, coord14) == to_t2(lib, coord21, coord22, coord23, coord24) ) & spatial_equal(lib, coord11, coord12, coord13, coord21, coord22, coord23) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/gamma.py000066400000000000000000000102161503546127100220640ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.gamma(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.lorentz import t, tau from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return lib.nan_to_num( t / tau.xy_z_t(lib, x, y, z, t), nan=lib.inf, posinf=inf, neginf=-inf ) def xy_z_tau(lib, x, y, z, tau): return lib.nan_to_num( t.xy_z_tau(lib, x, y, z, tau) / tau, nan=lib.inf, posinf=inf, neginf=-inf ) def xy_theta_t(lib, x, y, theta, t): return lib.nan_to_num( t / tau.xy_theta_t(lib, x, y, theta, t), nan=lib.inf, posinf=inf, neginf=-inf ) def xy_theta_tau(lib, x, y, theta, tau): return lib.nan_to_num( t.xy_theta_tau(lib, x, y, theta, tau) / tau, nan=lib.inf, posinf=inf, neginf=-inf, ) def xy_eta_t(lib, x, y, eta, t): return lib.nan_to_num( t / tau.xy_eta_t(lib, x, y, eta, t), nan=lib.inf, posinf=inf, neginf=-inf ) def xy_eta_tau(lib, x, y, eta, tau): return lib.nan_to_num( t.xy_eta_tau(lib, x, y, eta, tau) / tau, nan=lib.inf, posinf=inf, neginf=-inf ) def rhophi_z_t(lib, rho, phi, z, t): return lib.nan_to_num( t / tau.rhophi_z_t(lib, rho, phi, z, t), nan=lib.inf, posinf=inf, neginf=-inf ) def rhophi_z_tau(lib, rho, phi, z, tau): return lib.nan_to_num( t.rhophi_z_tau(lib, rho, phi, z, tau) / tau, nan=lib.inf, posinf=inf, neginf=-inf, ) def rhophi_theta_t(lib, rho, phi, theta, t): return lib.nan_to_num( t / tau.rhophi_theta_t(lib, rho, phi, theta, t), nan=lib.inf, posinf=inf, neginf=-inf, ) def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.nan_to_num( t.rhophi_theta_tau(lib, rho, phi, theta, tau) / tau, nan=lib.inf, posinf=inf, neginf=-inf, ) def rhophi_eta_t(lib, rho, phi, eta, t): return lib.nan_to_num( t / tau.rhophi_eta_t(lib, rho, phi, eta, t), nan=lib.inf, posinf=inf, neginf=-inf, ) def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.nan_to_num( t.rhophi_eta_tau(lib, rho, phi, eta, tau) / tau, nan=lib.inf, posinf=inf, neginf=-inf, ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/is_lightlike.py000066400000000000000000000037741503546127100234640ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.is_lightlike(self, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import dot from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) dispatch_map = {} def make_function(azimuthal, longitudinal, temporal): dot_function, _ = dot.dispatch_map[ azimuthal, longitudinal, temporal, azimuthal, longitudinal, temporal ] def f(lib, tolerance, coord1, coord2, coord3, coord4): return lib.absolute( dot_function( lib, coord1, coord2, coord3, coord4, coord1, coord2, coord3, coord4 ) ) < lib.absolute(tolerance) dispatch_map[azimuthal, longitudinal, temporal] = (f, bool) for azimuthal in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal in (TemporalT, TemporalTau): make_function(azimuthal, longitudinal, temporal) def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, tolerance, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/is_spacelike.py000066400000000000000000000037201503546127100234370ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.is_spacelike(self, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import dot from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) dispatch_map = {} def make_function(azimuthal, longitudinal, temporal): dot_function, _ = dot.dispatch_map[ azimuthal, longitudinal, temporal, azimuthal, longitudinal, temporal ] def f(lib, tolerance, coord1, coord2, coord3, coord4): return dot_function( lib, coord1, coord2, coord3, coord4, coord1, coord2, coord3, coord4 ) < lib.absolute(tolerance) dispatch_map[azimuthal, longitudinal, temporal] = (f, bool) for azimuthal in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal in (TemporalT, TemporalTau): make_function(azimuthal, longitudinal, temporal) def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, tolerance, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/is_timelike.py000066400000000000000000000037171503546127100233100ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.is_timelike(self, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import dot from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) dispatch_map = {} def make_function(azimuthal, longitudinal, temporal): dot_function, _ = dot.dispatch_map[ azimuthal, longitudinal, temporal, azimuthal, longitudinal, temporal ] def f(lib, tolerance, coord1, coord2, coord3, coord4): return dot_function( lib, coord1, coord2, coord3, coord4, coord1, coord2, coord3, coord4 ) > lib.absolute(tolerance) dispatch_map[azimuthal, longitudinal, temporal] = (f, bool) for azimuthal in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal in (TemporalT, TemporalTau): make_function(azimuthal, longitudinal, temporal) def dispatch(tolerance: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, tolerance, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/isclose.py000066400000000000000000000155221503546127100224500ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.isclose(self, rtol=..., atol=..., equal_nan=...) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import isclose from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): spatial_isclose, _ = isclose.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] if azimuthal1 is AzimuthalXY: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.xy_z_t elif temporal1 is TemporalTau: to_t1 = t.xy_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.xy_theta_t elif temporal1 is TemporalTau: to_t1 = t.xy_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.xy_eta_t elif temporal1 is TemporalTau: to_t1 = t.xy_eta_tau elif azimuthal1 is AzimuthalRhoPhi: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.rhophi_z_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.rhophi_theta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.rhophi_eta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_eta_tau if azimuthal2 is AzimuthalXY: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.xy_z_t elif temporal2 is TemporalTau: to_t2 = t.xy_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.xy_theta_t elif temporal2 is TemporalTau: to_t2 = t.xy_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.xy_eta_t elif temporal2 is TemporalTau: to_t2 = t.xy_eta_tau elif azimuthal2 is AzimuthalRhoPhi: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.rhophi_z_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.rhophi_theta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.rhophi_eta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_eta_tau if temporal1 == temporal2: def f( lib, rtol, atol, equal_nan, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24, ): return lib.isclose( coord14, coord24, rtol, atol, equal_nan ) & spatial_isclose( lib, rtol, atol, equal_nan, coord11, coord12, coord13, coord21, coord22, coord23, ) else: def f( lib, rtol, atol, equal_nan, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24, ): return lib.isclose( to_t1(lib, coord11, coord12, coord13, coord14), to_t2(lib, coord21, coord22, coord23, coord24), rtol, atol, equal_nan, ) & spatial_isclose( lib, rtol, atol, equal_nan, coord11, coord12, coord13, coord21, coord22, coord23, ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch( rtol: typing.Any, atol: typing.Any, equal_nan: typing.Any, v1: typing.Any, v2: typing.Any, ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), rtol, atol, equal_nan, *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/not_equal.py000066400000000000000000000135401503546127100227740ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.not_equal(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import equal from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): spatial_equal, _ = equal.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] if azimuthal1 is AzimuthalXY: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.xy_z_t elif temporal1 is TemporalTau: to_t1 = t.xy_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.xy_theta_t elif temporal1 is TemporalTau: to_t1 = t.xy_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.xy_eta_t elif temporal1 is TemporalTau: to_t1 = t.xy_eta_tau elif azimuthal1 is AzimuthalRhoPhi: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.rhophi_z_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.rhophi_theta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.rhophi_eta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_eta_tau if azimuthal2 is AzimuthalXY: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.xy_z_t elif temporal2 is TemporalTau: to_t2 = t.xy_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.xy_theta_t elif temporal2 is TemporalTau: to_t2 = t.xy_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.xy_eta_t elif temporal2 is TemporalTau: to_t2 = t.xy_eta_tau elif azimuthal2 is AzimuthalRhoPhi: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.rhophi_z_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.rhophi_theta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.rhophi_eta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_eta_tau if temporal1 == temporal2: def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): return (coord14 != coord24) & spatial_equal( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) else: def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): return ( to_t1(lib, coord11, coord12, coord13, coord14) != to_t2(lib, coord21, coord22, coord23, coord24) ) & spatial_equal(lib, coord11, coord12, coord13, coord21, coord22, coord23) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/rapidity.py000066400000000000000000000071631503546127100226360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.rapidity(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return 0.5 * lib.log((t + z) / (t - z)) def xy_z_tau(lib, x, y, z, tau): return xy_z_t(lib, x, y, z, t.xy_z_tau(lib, x, y, z, tau)) def xy_theta_t(lib, x, y, theta, t): return xy_z_t(lib, x, y, z.xy_theta(lib, x, y, theta), t) def xy_theta_tau(lib, x, y, theta, tau): return xy_z_t( lib, x, y, z.xy_theta(lib, x, y, theta), t.xy_theta_tau(lib, x, y, theta, tau) ) def xy_eta_t(lib, x, y, eta, t): return xy_z_t(lib, x, y, z.xy_eta(lib, x, y, eta), t) def xy_eta_tau(lib, x, y, eta, tau): return xy_z_t( lib, x, y, z.xy_eta(lib, x, y, eta), t.xy_eta_tau(lib, x, y, eta, tau) ) def rhophi_z_t(lib, rho, phi, z, t): return 0.5 * lib.log((t + z) / (t - z)) def rhophi_z_tau(lib, rho, phi, z, tau): return rhophi_z_t(lib, rho, phi, z, t.rhophi_z_tau(lib, rho, phi, z, tau)) def rhophi_theta_t(lib, rho, phi, theta, t): return rhophi_z_t(lib, rho, phi, z.rhophi_theta(lib, rho, phi, theta), t) def rhophi_theta_tau(lib, rho, phi, theta, tau): return rhophi_z_t( lib, rho, phi, z.rhophi_theta(lib, rho, phi, theta), t.rhophi_theta_tau(lib, rho, phi, theta, tau), ) def rhophi_eta_t(lib, rho, phi, eta, t): return rhophi_z_t(lib, rho, phi, z.rhophi_eta(lib, rho, phi, eta), t) def rhophi_eta_tau(lib, rho, phi, eta, tau): return rhophi_z_t( lib, rho, phi, z.rhophi_eta(lib, rho, phi, eta), t.rhophi_eta_tau(lib, rho, phi, eta, tau), ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/scale.py000066400000000000000000000115031503546127100220710ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.scale(self, factor) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import scale as scale3d from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, factor, x, y, z, t): sx, sy, sz = scale3d.xy_z(lib, factor, x, y, z) return (sx, sy, sz, t * factor) def xy_z_tau(lib, factor, x, y, z, tau): sx, sy, sz = scale3d.xy_z(lib, factor, x, y, z) return (sx, sy, sz, tau * factor) def xy_theta_t(lib, factor, x, y, theta, t): sx, sy, stheta = scale3d.xy_theta(lib, factor, x, y, theta) return (sx, sy, stheta, t * factor) def xy_theta_tau(lib, factor, x, y, theta, tau): sx, sy, stheta = scale3d.xy_theta(lib, factor, x, y, theta) return (sx, sy, stheta, tau * factor) def xy_eta_t(lib, factor, x, y, eta, t): sx, sy, seta = scale3d.xy_eta(lib, factor, x, y, eta) return (sx, sy, seta, t * factor) def xy_eta_tau(lib, factor, x, y, eta, tau): sx, sy, seta = scale3d.xy_eta(lib, factor, x, y, eta) return (sx, sy, seta, tau * factor) def rhophi_z_t(lib, factor, rho, phi, z, t): srho, sphi, sz = scale3d.rhophi_z(lib, factor, rho, phi, z) return (srho, sphi, sz, t * factor) def rhophi_z_tau(lib, factor, rho, phi, z, tau): srho, sphi, sz = scale3d.rhophi_z(lib, factor, rho, phi, z) return (srho, sphi, sz, tau * factor) def rhophi_theta_t(lib, factor, rho, phi, theta, t): srho, sphi, stheta = scale3d.rhophi_theta(lib, factor, rho, phi, theta) return (srho, sphi, stheta, t * factor) def rhophi_theta_tau(lib, factor, rho, phi, theta, tau): srho, sphi, stheta = scale3d.rhophi_theta(lib, factor, rho, phi, theta) return (srho, sphi, stheta, tau * factor) def rhophi_eta_t(lib, factor, rho, phi, eta, t): srho, sphi, seta = scale3d.rhophi_eta(lib, factor, rho, phi, eta) return (srho, sphi, seta, t * factor) def rhophi_eta_tau(lib, factor, rho, phi, eta, tau): srho, sphi, seta = scale3d.rhophi_eta(lib, factor, rho, phi, eta) return (srho, sphi, seta, tau * factor) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalTheta, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalTheta, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalEta, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalEta, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalRhoPhi, LongitudinalTheta, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalRhoPhi, LongitudinalTheta, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalRhoPhi, LongitudinalEta, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalRhoPhi, LongitudinalEta, TemporalTau, ), } def dispatch(factor: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, factor, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/subtract.py000066400000000000000000000162461503546127100226420ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.subtract(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t, tau from vector._compute.spatial import subtract from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, _ttype, ) dispatch_map = {} def make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ): spatial_subtract, azimuthal, longitudinal = subtract.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] if azimuthal1 is AzimuthalXY: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.xy_z_t elif temporal1 is TemporalTau: to_t1 = t.xy_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.xy_theta_t elif temporal1 is TemporalTau: to_t1 = t.xy_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.xy_eta_t elif temporal1 is TemporalTau: to_t1 = t.xy_eta_tau elif azimuthal1 is AzimuthalRhoPhi: if longitudinal1 is LongitudinalZ: if temporal1 is TemporalT: to_t1 = t.rhophi_z_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_z_tau elif longitudinal1 is LongitudinalTheta: if temporal1 is TemporalT: to_t1 = t.rhophi_theta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_theta_tau elif longitudinal1 is LongitudinalEta: if temporal1 is TemporalT: to_t1 = t.rhophi_eta_t elif temporal1 is TemporalTau: to_t1 = t.rhophi_eta_tau if azimuthal2 is AzimuthalXY: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.xy_z_t elif temporal2 is TemporalTau: to_t2 = t.xy_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.xy_theta_t elif temporal2 is TemporalTau: to_t2 = t.xy_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.xy_eta_t elif temporal2 is TemporalTau: to_t2 = t.xy_eta_tau elif azimuthal2 is AzimuthalRhoPhi: if longitudinal2 is LongitudinalZ: if temporal2 is TemporalT: to_t2 = t.rhophi_z_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_z_tau elif longitudinal2 is LongitudinalTheta: if temporal2 is TemporalT: to_t2 = t.rhophi_theta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_theta_tau elif longitudinal2 is LongitudinalEta: if temporal2 is TemporalT: to_t2 = t.rhophi_eta_t elif temporal2 is TemporalTau: to_t2 = t.rhophi_eta_tau if temporal1 is TemporalT or temporal2 is TemporalT: def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): coord1, coord2, coord3 = spatial_subtract( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) t1 = to_t1(lib, coord11, coord12, coord13, coord14) t2 = to_t2(lib, coord21, coord22, coord23, coord24) return (coord1, coord2, coord3, t1 - t2) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = (f, azimuthal, longitudinal, TemporalT) else: if azimuthal is AzimuthalXY: if longitudinal is LongitudinalZ: to_tau = tau.xy_z_t elif longitudinal is LongitudinalTheta: to_tau = tau.xy_theta_t elif longitudinal is LongitudinalEta: to_tau = tau.xy_eta_t elif azimuthal is AzimuthalRhoPhi: if longitudinal is LongitudinalZ: to_tau = tau.rhophi_z_t elif longitudinal is LongitudinalTheta: to_tau = tau.rhophi_theta_t elif longitudinal is LongitudinalEta: to_tau = tau.rhophi_eta_t def f( lib, coord11, coord12, coord13, coord14, coord21, coord22, coord23, coord24 ): coord1, coord2, coord3 = spatial_subtract( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) t1 = to_t1(lib, coord11, coord12, coord13, coord14) t2 = to_t2(lib, coord21, coord22, coord23, coord24) return ( coord1, coord2, coord3, to_tau(lib, coord1, coord2, coord3, t1 - t2), ) dispatch_map[ azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2 ] = ( f, azimuthal, longitudinal, TemporalTau, ) # type: ignore[assignment] for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal1 in (TemporalT, TemporalTau): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in ( LongitudinalZ, LongitudinalTheta, LongitudinalEta, ): for temporal2 in (TemporalT, TemporalTau): make_conversion( azimuthal1, longitudinal1, temporal1, azimuthal2, longitudinal2, temporal2, ) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _ttype(v1), _aztype(v2), _ltype(v2), _ttype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v1.temporal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, *v2.temporal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/lorentz/t.py000066400000000000000000000067051503546127100212550ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.t(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return t xy_z_t.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_z_tau(lib, x, y, z, tau): return lib.sqrt(t2.xy_z_tau(lib, x, y, z, tau)) def xy_theta_t(lib, x, y, theta, t): return t xy_theta_t.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_theta_tau(lib, x, y, theta, tau): return lib.sqrt(t2.xy_theta_tau(lib, x, y, theta, tau)) def xy_eta_t(lib, x, y, eta, t): return t xy_eta_t.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_eta_tau(lib, x, y, eta, tau): return lib.sqrt(t2.xy_eta_tau(lib, x, y, eta, tau)) def rhophi_z_t(lib, rho, phi, z, t): return t rhophi_z_t.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_z_tau(lib, rho, phi, z, tau): return lib.sqrt(t2.rhophi_z_tau(lib, rho, phi, z, tau)) def rhophi_theta_t(lib, rho, phi, theta, t): return t rhophi_theta_t.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.sqrt(t2.rhophi_theta_tau(lib, rho, phi, theta, tau)) def rhophi_eta_t(lib, rho, phi, eta, t): return t rhophi_eta_t.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.sqrt(t2.rhophi_eta_tau(lib, rho, phi, eta, tau)) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/t2.py000066400000000000000000000066071503546127100213400ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.t2(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import tau2 from vector._compute.spatial import mag2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return t**2 def xy_z_tau(lib, x, y, z, tau): return lib.maximum(tau2.xy_z_tau(lib, x, y, z, tau) + mag2.xy_z(lib, x, y, z), 0) def xy_theta_t(lib, x, y, theta, t): return t**2 def xy_theta_tau(lib, x, y, theta, tau): return lib.maximum( tau2.xy_theta_tau(lib, x, y, theta, tau) + mag2.xy_theta(lib, x, y, theta), 0 ) def xy_eta_t(lib, x, y, eta, t): return t**2 def xy_eta_tau(lib, x, y, eta, tau): return lib.maximum( tau2.xy_eta_tau(lib, x, y, eta, tau) + mag2.xy_eta(lib, x, y, eta), 0 ) def rhophi_z_t(lib, rho, phi, z, t): return t**2 def rhophi_z_tau(lib, rho, phi, z, tau): return lib.maximum( tau2.rhophi_z_tau(lib, rho, phi, z, tau) + mag2.rhophi_z(lib, rho, phi, z), 0 ) def rhophi_theta_t(lib, rho, phi, theta, t): return t**2 def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.maximum( tau2.rhophi_theta_tau(lib, rho, phi, theta, tau) + mag2.rhophi_theta(lib, rho, phi, theta), 0, ) def rhophi_eta_t(lib, rho, phi, eta, t): return t**2 def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.maximum( tau2.rhophi_eta_tau(lib, rho, phi, eta, tau) + mag2.rhophi_eta(lib, rho, phi, eta), 0, ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/tau.py000066400000000000000000000074671503546127100216110ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.tau(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import tau2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): squared = tau2.xy_z_t(lib, x, y, z, t) return lib.copysign(lib.sqrt(lib.absolute(squared)), squared) def xy_z_tau(lib, x, y, z, tau): return tau xy_z_tau.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_theta_t(lib, x, y, theta, t): squared = tau2.xy_theta_t(lib, x, y, theta, t) return lib.copysign(lib.sqrt(lib.absolute(squared)), squared) def xy_theta_tau(lib, x, y, theta, tau): return tau xy_theta_tau.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_eta_t(lib, x, y, eta, t): squared = tau2.xy_eta_t(lib, x, y, eta, t) return lib.copysign(lib.sqrt(lib.absolute(squared)), squared) def xy_eta_tau(lib, x, y, eta, tau): return tau xy_eta_tau.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_z_t(lib, rho, phi, z, t): squared = tau2.rhophi_z_t(lib, rho, phi, z, t) return lib.copysign(lib.sqrt(lib.absolute(squared)), squared) def rhophi_z_tau(lib, rho, phi, z, tau): return tau rhophi_z_tau.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_theta_t(lib, rho, phi, theta, t): squared = tau2.rhophi_theta_t(lib, rho, phi, theta, t) return lib.copysign(lib.sqrt(lib.absolute(squared)), squared) def rhophi_theta_tau(lib, rho, phi, theta, tau): return tau rhophi_theta_tau.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_eta_t(lib, rho, phi, eta, t): squared = tau2.rhophi_eta_t(lib, rho, phi, eta, t) return lib.copysign(lib.sqrt(lib.absolute(squared)), squared) def rhophi_eta_tau(lib, rho, phi, eta, tau): return tau rhophi_eta_tau.__awkward_transform_allowed__ = False # type:ignore[attr-defined] dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/tau2.py000066400000000000000000000060761503546127100216660ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Lorentz.tau2(self) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import mag2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return t**2 - mag2.xy_z(lib, x, y, z) def xy_z_tau(lib, x, y, z, tau): return lib.copysign(tau**2, tau) def xy_theta_t(lib, x, y, theta, t): return t**2 - mag2.xy_theta(lib, x, y, theta) def xy_theta_tau(lib, x, y, theta, tau): return lib.copysign(tau**2, tau) def xy_eta_t(lib, x, y, eta, t): return t**2 - mag2.xy_eta(lib, x, y, eta) def xy_eta_tau(lib, x, y, eta, tau): return lib.copysign(tau**2, tau) def rhophi_z_t(lib, rho, phi, z, t): return t**2 - mag2.rhophi_z(lib, rho, phi, z) def rhophi_z_tau(lib, rho, phi, z, tau): return lib.copysign(tau**2, tau) def rhophi_theta_t(lib, rho, phi, theta, t): return t**2 - mag2.rhophi_theta(lib, rho, phi, theta) def rhophi_theta_tau(lib, rho, phi, theta, tau): return lib.copysign(tau**2, tau) def rhophi_eta_t(lib, rho, phi, eta, t): return t**2 - mag2.rhophi_eta(lib, rho, phi, eta) def rhophi_eta_tau(lib, rho, phi, eta, tau): return lib.copysign(tau**2, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, float), (AzimuthalXY, LongitudinalZ, TemporalTau): (xy_z_tau, float), (AzimuthalXY, LongitudinalTheta, TemporalT): (xy_theta_t, float), (AzimuthalXY, LongitudinalTheta, TemporalTau): (xy_theta_tau, float), (AzimuthalXY, LongitudinalEta, TemporalT): (xy_eta_t, float), (AzimuthalXY, LongitudinalEta, TemporalTau): (xy_eta_tau, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): (rhophi_z_t, float), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): (rhophi_z_tau, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): (rhophi_theta_t, float), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): (rhophi_theta_tau, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): (rhophi_eta_t, float), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): (rhophi_eta_tau, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/to_beta3.py000066400000000000000000000077621503546127100225160ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.to_beta3(self) """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): return (x / t, y / t, z / t) def xy_z_tau(lib, x, y, z, tau): return xy_z_t(lib, x, y, z, t.xy_z_tau(lib, x, y, z, tau)) def xy_theta_t(lib, x, y, theta, t): return (x / t, y / t, theta) def xy_theta_tau(lib, x, y, theta, tau): return xy_theta_t(lib, x, y, theta, t.xy_theta_tau(lib, x, y, theta, tau)) def xy_eta_t(lib, x, y, eta, t): return (x / t, y / t, eta) def xy_eta_tau(lib, x, y, eta, tau): return xy_eta_t(lib, x, y, eta, t.xy_eta_tau(lib, x, y, eta, tau)) def rhophi_z_t(lib, rho, phi, z, t): return (rho / t, phi, z / t) def rhophi_z_tau(lib, rho, phi, z, tau): return rhophi_z_t(lib, rho, phi, z, t.rhophi_z_tau(lib, rho, phi, z, tau)) def rhophi_theta_t(lib, rho, phi, theta, t): return (rho / t, phi, theta) def rhophi_theta_tau(lib, rho, phi, theta, tau): return rhophi_theta_t( lib, rho, phi, theta, t.rhophi_theta_tau(lib, rho, phi, theta, tau) ) def rhophi_eta_t(lib, rho, phi, eta, t): return (rho / t, phi, eta) def rhophi_eta_tau(lib, rho, phi, eta, tau): return rhophi_eta_t(lib, rho, phi, eta, t.rhophi_eta_tau(lib, rho, phi, eta, tau)) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): (xy_z_t, AzimuthalXY, LongitudinalZ, None), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, None, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalTheta, None, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalTheta, None, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalEta, None, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalEta, None, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalRhoPhi, LongitudinalZ, None, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalRhoPhi, LongitudinalZ, None, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalRhoPhi, LongitudinalTheta, None, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalRhoPhi, LongitudinalTheta, None, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalRhoPhi, LongitudinalEta, None, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalRhoPhi, LongitudinalEta, None, ), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/transform4D.py000066400000000000000000000134061503546127100232110ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.transform4D(self, obj) where ``obj`` has ``obj["xx"]``, ``obj["xy"]``, etc. """ from __future__ import annotations import typing import numpy from vector._compute.lorentz import t from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. def cartesian_t( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, tx, ty, tz, tt, x, y, z, t ): xp = xx * x + xy * y + xz * z + xt * t yp = yx * x + yy * y + yz * z + yt * t zp = zx * x + zy * y + zz * z + zt * t tp = tx * x + ty * y + tz * z + tt * t return (xp, yp, zp, tp) def cartesian_tau(lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, x, y, z, tau): tee = t.xy_z_tau(lib, x, y, z, tau) xp = xx * x + xy * y + xz * z + xt * tee yp = yx * x + yy * y + yz * z + yt * tee zp = zx * x + zy * y + zz * z + zt * tee return (xp, yp, zp, tau) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( cartesian_t, AzimuthalXY, LongitudinalZ, TemporalT, ) } def make_conversion(azimuthal, longitudinal, temporal): if (azimuthal, longitudinal, temporal) != (AzimuthalXY, LongitudinalZ, TemporalT): if azimuthal is AzimuthalXY: to_x = x.xy to_y = y.xy if longitudinal is LongitudinalZ: to_z = z.xy_z if temporal is TemporalT: to_t = t.xy_z_t elif temporal is TemporalTau: to_t = t.xy_z_tau elif longitudinal is LongitudinalTheta: to_z = z.xy_theta if temporal is TemporalT: to_t = t.xy_theta_t elif temporal is TemporalTau: to_t = t.xy_theta_tau elif longitudinal is LongitudinalEta: to_z = z.xy_eta if temporal is TemporalT: to_t = t.xy_eta_t elif temporal is TemporalTau: to_t = t.xy_eta_tau elif azimuthal is AzimuthalRhoPhi: to_x = x.rhophi to_y = y.rhophi if longitudinal is LongitudinalZ: to_z = z.rhophi_z if temporal is TemporalT: to_t = t.rhophi_z_t elif temporal is TemporalTau: to_t = t.rhophi_z_tau elif longitudinal is LongitudinalTheta: to_z = z.rhophi_theta if temporal is TemporalT: to_t = t.rhophi_theta_t elif temporal is TemporalTau: to_t = t.rhophi_theta_tau elif longitudinal is LongitudinalEta: to_z = z.rhophi_eta if temporal is TemporalT: to_t = t.rhophi_eta_t elif temporal is TemporalTau: to_t = t.rhophi_eta_tau cartesian, azout, lout, tout = dispatch_map[ AzimuthalXY, LongitudinalZ, TemporalT ] def f( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, tx, ty, tz, tt, coord1, coord2, coord3, coord4, ): return cartesian( lib, xx, xy, xz, xt, yx, yy, yz, yt, zx, zy, zz, zt, tx, ty, tz, tt, to_x(lib, coord1, coord2), to_y(lib, coord1, coord2), to_z(lib, coord1, coord2, coord3), to_t(lib, coord1, coord2, coord3, coord4), ) dispatch_map[azimuthal, longitudinal, temporal] = (f, azout, lout, tout) for azimuthal in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for temporal in (TemporalT, TemporalTau): make_conversion(azimuthal, longitudinal, temporal) def dispatch(obj: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, obj["xx"], obj["xy"], obj["xz"], obj["xt"], obj["yx"], obj["yy"], obj["yz"], obj["yt"], obj["zx"], obj["zy"], obj["zz"], obj["zt"], obj["tx"], obj["ty"], obj["tz"], obj["tt"], *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/lorentz/unit.py000066400000000000000000000151071503546127100217650ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Lorentz.unit(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.lorentz import tau2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, TemporalT, TemporalTau, _aztype, _flavor_of, _from_signature, _ltype, _ttype, ) def xy_z_t(lib, x, y, z, t): squared = tau2.xy_z_t(lib, x, y, z, t) norm = lib.sqrt(lib.absolute(squared)) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(z / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(t / norm, nan=0, posinf=inf, neginf=-inf), ) def xy_z_tau(lib, x, y, z, tau): norm = lib.absolute(tau) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(z / norm, nan=0, posinf=inf, neginf=-inf), lib.copysign(1, tau), ) def xy_theta_t(lib, x, y, theta, t): squared = tau2.xy_theta_t(lib, x, y, theta, t) norm = lib.sqrt(lib.absolute(squared)) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), theta, lib.nan_to_num(t / norm, nan=0, posinf=inf, neginf=-inf), ) def xy_theta_tau(lib, x, y, theta, tau): norm = lib.absolute(tau) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), theta, lib.copysign(1, tau), ) def xy_eta_t(lib, x, y, eta, t): squared = tau2.xy_eta_t(lib, x, y, eta, t) norm = lib.sqrt(lib.absolute(squared)) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), eta, lib.nan_to_num(t / norm, nan=0, posinf=inf, neginf=-inf), ) def xy_eta_tau(lib, x, y, eta, tau): norm = lib.absolute(tau) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), eta, lib.copysign(1, tau), ) def rhophi_z_t(lib, rho, phi, z, t): squared = tau2.rhophi_z_t(lib, rho, phi, z, t) norm = lib.sqrt(lib.absolute(squared)) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, lib.nan_to_num(z / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(t / norm, nan=0, posinf=inf, neginf=-inf), ) def rhophi_z_tau(lib, rho, phi, z, tau): norm = lib.absolute(tau) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, lib.nan_to_num(z / norm, nan=0, posinf=inf, neginf=-inf), lib.copysign(1, tau), ) def rhophi_theta_t(lib, rho, phi, theta, t): squared = tau2.rhophi_theta_t(lib, rho, phi, theta, t) norm = lib.sqrt(lib.absolute(squared)) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, theta, lib.nan_to_num(t / norm, nan=0, posinf=inf, neginf=-inf), ) def rhophi_theta_tau(lib, rho, phi, theta, tau): norm = lib.absolute(tau) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, theta, lib.copysign(1, tau), ) def rhophi_eta_t(lib, rho, phi, eta, t): squared = tau2.rhophi_eta_t(lib, rho, phi, eta, t) norm = lib.sqrt(lib.absolute(squared)) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, eta, lib.nan_to_num(t / norm, nan=0, posinf=inf, neginf=-inf), ) def rhophi_eta_tau(lib, rho, phi, eta, tau): norm = lib.absolute(tau) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, eta, lib.copysign(1, tau), ) dispatch_map = { (AzimuthalXY, LongitudinalZ, TemporalT): ( xy_z_t, AzimuthalXY, LongitudinalZ, TemporalT, ), (AzimuthalXY, LongitudinalZ, TemporalTau): ( xy_z_tau, AzimuthalXY, LongitudinalZ, TemporalTau, ), (AzimuthalXY, LongitudinalTheta, TemporalT): ( xy_theta_t, AzimuthalXY, LongitudinalTheta, TemporalT, ), (AzimuthalXY, LongitudinalTheta, TemporalTau): ( xy_theta_tau, AzimuthalXY, LongitudinalTheta, TemporalTau, ), (AzimuthalXY, LongitudinalEta, TemporalT): ( xy_eta_t, AzimuthalXY, LongitudinalEta, TemporalT, ), (AzimuthalXY, LongitudinalEta, TemporalTau): ( xy_eta_tau, AzimuthalXY, LongitudinalEta, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalT): ( rhophi_z_t, AzimuthalRhoPhi, LongitudinalZ, TemporalT, ), (AzimuthalRhoPhi, LongitudinalZ, TemporalTau): ( rhophi_z_tau, AzimuthalRhoPhi, LongitudinalZ, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalT): ( rhophi_theta_t, AzimuthalRhoPhi, LongitudinalTheta, TemporalT, ), (AzimuthalRhoPhi, LongitudinalTheta, TemporalTau): ( rhophi_theta_tau, AzimuthalRhoPhi, LongitudinalTheta, TemporalTau, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalT): ( rhophi_eta_t, AzimuthalRhoPhi, LongitudinalEta, TemporalT, ), (AzimuthalRhoPhi, LongitudinalEta, TemporalTau): ( rhophi_eta_tau, AzimuthalRhoPhi, LongitudinalEta, TemporalTau, ), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), _ttype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements, *v.temporal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/000077500000000000000000000000001503546127100202105ustar00rootroot00000000000000vector-1.6.3/src/vector/_compute/planar/__init__.py000066400000000000000000000035071503546127100223260ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Compute functions for planar vectors, which is to say 2D, 3D, and 4D. Each function is a module with variants for each coordinate system (or combination of coordinate systems) as functions within the module. Each module has a ``dispatch_map`` (dict) that maps coordinate types to the appropriate function and its return type(s), and a ``dispatch`` (function) uses this information to call the right function and return the right type. The compute functions themselves are restricted to a minimum of Python features: no statements other than assignments and one return, no assumptions about data types. In particular, if statements and loops are not allowed. The tests/test_compute_features.py suite ensures that these rules are followed (though that set of allowed features can be expanded if it doesn't prevent the addition of new backends). """ from __future__ import annotations import vector._compute.planar.add import vector._compute.planar.deltaphi import vector._compute.planar.dot import vector._compute.planar.equal import vector._compute.planar.is_antiparallel import vector._compute.planar.is_parallel import vector._compute.planar.is_perpendicular import vector._compute.planar.isclose import vector._compute.planar.not_equal import vector._compute.planar.phi import vector._compute.planar.rho import vector._compute.planar.rho2 import vector._compute.planar.rotateZ import vector._compute.planar.scale import vector._compute.planar.subtract import vector._compute.planar.transform2D import vector._compute.planar.unit import vector._compute.planar.x import vector._compute.planar.y # noqa: F401 vector-1.6.3/src/vector/_compute/planar/add.py000066400000000000000000000040211503546127100213070ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.add(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) def rectify(lib, phi): return (phi + lib.pi) % (2 * lib.pi) - lib.pi # specialized def xy_xy(lib, x1, y1, x2, y2): return (x1 + x2, y1 + y2) def xy_rhophi(lib, x1, y1, rho2, phi2): return xy_xy(lib, x1, y1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2)) def rhophi_xy(lib, rho1, phi1, x2, y2): return xy_xy(lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), x2, y2) # specialized def rhophi_rhophi(lib, rho1, phi1, rho2, phi2): diff = phi2 - phi1 u = rho2 * lib.cos(diff) v = rho2 * lib.sin(diff) return ( lib.sqrt((rho1 + u) ** 2 + v**2), rectify(lib, phi1 + lib.arctan2(v, rho1 + u)), ) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, AzimuthalXY), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, AzimuthalXY), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, AzimuthalXY), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, AzimuthalRhoPhi), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/deltaphi.py000066400000000000000000000034411503546127100223560ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.deltaphi(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import phi from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) def rectify(lib, phi): return (phi + lib.pi) % (2 * lib.pi) - lib.pi def xy_xy(lib, x1, y1, x2, y2): return rectify(lib, phi.xy(lib, x1, y1) - phi.xy(lib, x2, y2)) def xy_rhophi(lib, x1, y1, rho2, phi2): return rectify(lib, phi.xy(lib, x1, y1) - phi2) def rhophi_xy(lib, rho1, phi1, x2, y2): return rectify(lib, phi1 - phi.xy(lib, x2, y2)) def rhophi_rhophi(lib, rho1, phi1, rho2, phi2): return rectify(lib, phi1 - phi2) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, float), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, float), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, float), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, float), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/dot.py000066400000000000000000000033441503546127100213540ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.dot(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) def xy_xy(lib, x1, y1, x2, y2): return x1 * x2 + y1 * y2 def xy_rhophi(lib, x1, y1, rho2, phi2): return x1 * x.rhophi(lib, rho2, phi2) + y1 * y.rhophi(lib, rho2, phi2) def rhophi_xy(lib, rho1, phi1, x2, y2): return x.rhophi(lib, rho1, phi1) * x2 + y.rhophi(lib, rho1, phi1) * y2 def rhophi_rhophi(lib, rho1, phi1, rho2, phi2): return rho1 * rho2 * lib.cos(phi1 - phi2) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, float), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, float), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, float), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, float), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/equal.py000066400000000000000000000035751503546127100217030ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.equal(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) # Policy: turn (rho, phi) into (x, y) # (if not already the same) # same types def xy_xy(lib, x1, y1, x2, y2): return (x1 == x2) & (y1 == y2) def xy_rhophi(lib, x1, y1, rho2, phi2): return xy_xy(lib, x1, y1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2)) def rhophi_xy(lib, rho1, phi1, x2, y2): return xy_xy(lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), x2, y2) # same types def rhophi_rhophi(lib, rho1, phi1, rho2, phi2): return (rho1 == rho2) & (phi1 == phi2) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, bool), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, bool), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, bool), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, bool), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/is_antiparallel.py000066400000000000000000000036301503546127100237270ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.is_antiparallel(self, other, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.planar import dot, rho from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) dispatch_map = {} def make_function(azimuthal1, azimuthal2): dot_function, _ = dot.dispatch_map[azimuthal1, azimuthal2] rho1_function, _ = rho.dispatch_map[azimuthal1,] rho2_function, _ = rho.dispatch_map[azimuthal2,] def f(lib, tolerance, coord11, coord12, coord21, coord22): return dot_function(lib, coord11, coord12, coord21, coord22) < ( lib.absolute(tolerance) - 1 ) * rho1_function(lib, coord11, coord12) * rho2_function(lib, coord21, coord22) dispatch_map[azimuthal1, azimuthal2] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): make_function(azimuthal1, azimuthal2) def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), tolerance, *v1.azimuthal.elements, *v2.azimuthal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/is_parallel.py000066400000000000000000000036241503546127100230560ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.is_parallel(self, other, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.planar import dot, rho from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) dispatch_map = {} def make_function(azimuthal1, azimuthal2): dot_function, _ = dot.dispatch_map[azimuthal1, azimuthal2] rho1_function, _ = rho.dispatch_map[azimuthal1,] rho2_function, _ = rho.dispatch_map[azimuthal2,] def f(lib, tolerance, coord11, coord12, coord21, coord22): return dot_function(lib, coord11, coord12, coord21, coord22) > ( 1 - lib.absolute(tolerance) ) * rho1_function(lib, coord11, coord12) * rho2_function(lib, coord21, coord22) dispatch_map[azimuthal1, azimuthal2] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): make_function(azimuthal1, azimuthal2) def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), tolerance, *v1.azimuthal.elements, *v2.azimuthal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/is_perpendicular.py000066400000000000000000000036231503546127100241160ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.is_perpendicular(self, other, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.planar import dot, rho from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) dispatch_map = {} def make_function(azimuthal1, azimuthal2): dot_function, _ = dot.dispatch_map[azimuthal1, azimuthal2] rho1_function, _ = rho.dispatch_map[azimuthal1,] rho2_function, _ = rho.dispatch_map[azimuthal2,] def f(lib, tolerance, coord11, coord12, coord21, coord22): return dot_function(lib, coord11, coord12, coord21, coord22) < lib.absolute( tolerance ) * rho1_function(lib, coord11, coord12) * rho2_function(lib, coord21, coord22) dispatch_map[azimuthal1, azimuthal2] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): make_function(azimuthal1, azimuthal2) def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), tolerance, *v1.azimuthal.elements, *v2.azimuthal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/isclose.py000066400000000000000000000047571503546127100222400ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.isclose(self, other, rtol=..., atol=..., equal_nan=...) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) # Policy: turn (rho, phi) into (x, y) # (if not already the same) # same types def xy_xy(lib, rtol, atol, equal_nan, x1, y1, x2, y2): return lib.isclose(x1, x2, rtol, atol, equal_nan) & lib.isclose( y1, y2, rtol, atol, equal_nan ) def xy_rhophi(lib, rtol, atol, equal_nan, x1, y1, rho2, phi2): return xy_xy( lib, rtol, atol, equal_nan, x1, y1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), ) def rhophi_xy(lib, rtol, atol, equal_nan, rho1, phi1, x2, y2): return xy_xy( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), x2, y2, ) # same types def rhophi_rhophi(lib, rtol, atol, equal_nan, rho1, phi1, rho2, phi2): return lib.isclose(rho1, rho2, rtol, atol, equal_nan) & lib.isclose( phi1, phi2, rtol, atol, equal_nan ) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, bool), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, bool), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, bool), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, bool), } def dispatch( rtol: typing.Any, atol: typing.Any, equal_nan: typing.Any, v1: typing.Any, v2: typing.Any, ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), rtol, atol, equal_nan, *v1.azimuthal.elements, *v2.azimuthal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/not_equal.py000066400000000000000000000036011503546127100225510ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.not_equal(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) # Policy: turn (rho, phi) into (x, y) # (if not already the same) # same types def xy_xy(lib, x1, y1, x2, y2): return (x1 != x2) & (y1 != y2) def xy_rhophi(lib, x1, y1, rho2, phi2): return xy_xy(lib, x1, y1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2)) def rhophi_xy(lib, rho1, phi1, x2, y2): return xy_xy(lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), x2, y2) # same types def rhophi_rhophi(lib, rho1, phi1, rho2, phi2): return (rho1 != rho2) & (phi1 != phi2) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, bool), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, bool), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, bool), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, bool), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/phi.py000066400000000000000000000021301503546127100213360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Planar.phi(self) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def xy(lib, x, y): return lib.arctan2(y, x) def rhophi(lib, rho, phi): return phi rhophi.__awkward_transform_allowed__ = False # type:ignore[attr-defined] dispatch_map = { (AzimuthalXY,): (xy, float), (AzimuthalRhoPhi,): (rhophi, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/rho.py000066400000000000000000000022131503546127100213500ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Planar.rho(self) """ from __future__ import annotations import typing import numpy from vector._compute.planar import rho2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def xy(lib, x, y): return lib.sqrt(rho2.xy(lib, x, y)) def rhophi(lib, rho, phi): return rho rhophi.__awkward_transform_allowed__ = False # type:ignore[attr-defined] dispatch_map = { (AzimuthalXY,): (xy, float), (AzimuthalRhoPhi,): (rhophi, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/rho2.py000066400000000000000000000020121503546127100214270ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Planar.rho2(self) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def xy(lib, x, y): return x**2 + y**2 def rhophi(lib, rho, phi): return rho**2 dispatch_map = { (AzimuthalXY,): (xy, float), (AzimuthalRhoPhi,): (rhophi, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/rotateZ.py000066400000000000000000000023401503546127100222110ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.rotateZ(self, angle) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def rectify(lib, phi): return (phi + lib.pi) % (2 * lib.pi) - lib.pi def xy(lib, angle, x, y): s = lib.sin(angle) c = lib.cos(angle) return c * x - s * y, s * x + c * y def rhophi(lib, angle, rho, phi): return rho, rectify(lib, phi + angle) dispatch_map = { (AzimuthalXY,): (xy, AzimuthalXY), (AzimuthalRhoPhi,): (rhophi, AzimuthalRhoPhi), } def dispatch(angle: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, angle, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/scale.py000066400000000000000000000024751503546127100216610ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.scale(self, factor) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def rectify(lib, phi): return (phi + lib.pi) % (2 * lib.pi) - lib.pi def xy(lib, factor, x, y): return (x * factor, y * factor) def rhophi(lib, factor, rho, phi): absfactor = lib.absolute(factor) sign = lib.sign(factor) turn_if_negative = -0.5 * (sign - 1) * lib.pi return (rho * absfactor, rectify(lib, phi + turn_if_negative)) dispatch_map = { (AzimuthalXY,): (xy, AzimuthalXY), (AzimuthalRhoPhi,): (rhophi, AzimuthalRhoPhi), } def dispatch(factor: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, factor, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/subtract.py000066400000000000000000000040371503546127100224150ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.subtract(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, ) def rectify(lib, phi): return (phi + lib.pi) % (2 * lib.pi) - lib.pi # specialized def xy_xy(lib, x1, y1, x2, y2): return (x1 - x2, y1 - y2) def xy_rhophi(lib, x1, y1, rho2, phi2): return xy_xy(lib, x1, y1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2)) def rhophi_xy(lib, rho1, phi1, x2, y2): return xy_xy(lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), x2, y2) # specialized def rhophi_rhophi(lib, rho1, phi1, rho2, phi2): diff = phi2 - phi1 + lib.pi u = rho2 * lib.cos(diff) v = rho2 * lib.sin(diff) return ( lib.sqrt((rho1 + u) ** 2 + v**2), rectify(lib, phi1 + lib.arctan2(v, rho1 + u)), ) dispatch_map = { (AzimuthalXY, AzimuthalXY): (xy_xy, AzimuthalXY), (AzimuthalXY, AzimuthalRhoPhi): (xy_rhophi, AzimuthalXY), (AzimuthalRhoPhi, AzimuthalXY): (rhophi_xy, AzimuthalXY), (AzimuthalRhoPhi, AzimuthalRhoPhi): (rhophi_rhophi, AzimuthalRhoPhi), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _aztype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v2.azimuthal.elements ), returns, 2, ) vector-1.6.3/src/vector/_compute/planar/transform2D.py000066400000000000000000000026731503546127100227730ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.transform2D(self, obj) where ``obj`` has ``obj["xx"]``, ``obj["xy"]``, etc. """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. def cartesian(lib, xx, xy, yx, yy, x, y): return (xx * x + xy * y, yx * x + yy * y) def rhophi(lib, xx, xy, yx, yy, rho, phi): return cartesian( lib, xx, xy, yx, yy, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi) ) dispatch_map = { (AzimuthalXY,): (cartesian, AzimuthalXY), (AzimuthalRhoPhi,): (rhophi, AzimuthalXY), } def dispatch(obj: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, obj["xx"], obj["xy"], obj["yx"], obj["yy"], *v.azimuthal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/unit.py000066400000000000000000000024631503546127100215460ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Planar.unit(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.planar import rho from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def xy(lib, x, y): norm = rho.xy(lib, x, y) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), ) def rhophi(lib, rho, phi): return (1, phi) rhophi.__awkward_transform_allowed__ = False # type:ignore[attr-defined] dispatch_map = { (AzimuthalXY,): (xy, AzimuthalXY), (AzimuthalRhoPhi,): (rhophi, AzimuthalRhoPhi), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/x.py000066400000000000000000000021211503546127100210250ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Planar.x(self) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def xy(lib, x, y): return x xy.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi(lib, rho, phi): return rho * lib.cos(phi) dispatch_map = { (AzimuthalXY,): (xy, float), (AzimuthalRhoPhi,): (rhophi, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/planar/y.py000066400000000000000000000021211503546127100210260ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Planar.y(self) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, _aztype, _flavor_of, _from_signature, ) def xy(lib, x, y): return y xy.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi(lib, rho, phi): return rho * lib.sin(phi) dispatch_map = { (AzimuthalXY,): (xy, float), (AzimuthalRhoPhi,): (rhophi, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature(__name__, dispatch_map, (_aztype(v),)) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)(v.lib, *v.azimuthal.elements), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/000077500000000000000000000000001503546127100203705ustar00rootroot00000000000000vector-1.6.3/src/vector/_compute/spatial/__init__.py000066400000000000000000000043651503546127100225110ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Compute functions for spatial vectors, which is to say 3D and 4D. Each function is a module with variants for each coordinate system (or combination of coordinate systems) as functions within the module. Each module has a ``dispatch_map`` (dict) that maps coordinate types to the appropriate function and its return type(s), and a ``dispatch`` (function) uses this information to call the right function and return the right type. The compute functions themselves are restricted to a minimum of Python features: no statements other than assignments and one return, no assumptions about data types. In particular, if statements and loops are not allowed. The tests/test_compute_features.py suite ensures that these rules are followed (though that set of allowed features can be expanded if it doesn't prevent the addition of new backends). """ from __future__ import annotations import vector._compute.spatial.add import vector._compute.spatial.costheta import vector._compute.spatial.cottheta import vector._compute.spatial.cross import vector._compute.spatial.deltaangle import vector._compute.spatial.deltaeta import vector._compute.spatial.deltaR import vector._compute.spatial.deltaR2 import vector._compute.spatial.dot import vector._compute.spatial.equal import vector._compute.spatial.eta import vector._compute.spatial.is_antiparallel import vector._compute.spatial.is_parallel import vector._compute.spatial.is_perpendicular import vector._compute.spatial.isclose import vector._compute.spatial.mag import vector._compute.spatial.mag2 import vector._compute.spatial.not_equal import vector._compute.spatial.rotate_axis import vector._compute.spatial.rotate_euler import vector._compute.spatial.rotate_quaternion import vector._compute.spatial.rotateX import vector._compute.spatial.rotateY import vector._compute.spatial.scale import vector._compute.spatial.subtract import vector._compute.spatial.theta import vector._compute.spatial.transform3D import vector._compute.spatial.unit import vector._compute.spatial.z # noqa: F401 vector-1.6.3/src/vector/_compute/spatial/add.py000066400000000000000000000366771503546127100215150ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.add(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import add, x, y from vector._compute.spatial import eta, theta, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # keep them in xy_z (default, anyway) def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return (x1 + x2, y1 + y2, z1 + z2) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_theta(lib, x2, y2, theta2)) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_eta(lib, x2, y2, eta2)) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2 ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z2) # keep them in xy_theta def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): x = x1 + x2 y = y1 + y2 z1 = z.xy_theta(lib, x1, y1, theta1) z2 = z.xy_theta(lib, x2, y2, theta2) return (x, y, theta.xy_z(lib, x, y, z1 + z2)) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z.xy_eta(lib, x2, y2, eta2), ) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z2) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z.xy_theta(lib, x2, y2, theta2), ) # keep them in xy_eta def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): x = x1 + x2 y = y1 + y2 z1 = z.xy_eta(lib, x1, y1, eta1) z2 = z.xy_eta(lib, x2, y2, eta2) return (x, y, eta.xy_z(lib, x, y, z1 + z2)) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z2 ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_eta(lib, x2, y2, eta2), ) # keep them in rhophi_z def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): rho, phi = add.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) return (rho, phi, z1 + z2) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z2, ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z.xy_eta(lib, x2, y2, eta2), ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) # keep them in rhophi_theta def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): rho, phi = add.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) z1 = z.rhophi_theta(lib, rho1, phi1, theta1) z2 = z.rhophi_theta(lib, rho2, phi2, theta2) return (rho, phi, theta.rhophi_z(lib, rho, phi, z1 + z2)) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z2, ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z.xy_eta(lib, x2, y2, eta2), ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) # keep them in rhophi_eta def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): rho, phi = add.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) z1 = z.rhophi_eta(lib, rho1, phi1, eta1) z2 = z.rhophi_eta(lib, rho2, phi2, eta2) return (rho, phi, eta.rhophi_z(lib, rho, phi, z1 + z2)) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( xy_z_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( xy_z_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, AzimuthalXY, LongitudinalTheta, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( xy_eta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, AzimuthalXY, LongitudinalEta, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, AzimuthalRhoPhi, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, AzimuthalRhoPhi, LongitudinalTheta, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, AzimuthalRhoPhi, LongitudinalEta, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/costheta.py000066400000000000000000000040031503546127100225510ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.costheta(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.spatial import mag, theta from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): return lib.nan_to_num(z / mag.xy_z(lib, x, y, z), nan=1.0, posinf=inf, neginf=-inf) def xy_theta(lib, x, y, theta): return lib.cos(theta) def xy_eta(lib, x, y, eta): return lib.cos(theta.xy_eta(lib, x, y, eta)) def rhophi_z(lib, rho, phi, z): return lib.nan_to_num( z / mag.rhophi_z(lib, rho, phi, z), nan=1.0, posinf=inf, neginf=-inf ) def rhophi_theta(lib, rho, phi, theta): return lib.cos(theta) def rhophi_eta(lib, rho, phi, eta): return lib.cos(theta.rhophi_eta(lib, rho, phi, eta)) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/cottheta.py000066400000000000000000000040171503546127100225570ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.cottheta(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.planar import rho from vector._compute.spatial import theta from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): return lib.nan_to_num(z / rho.xy(lib, x, y), nan=lib.inf, posinf=inf, neginf=-inf) def xy_theta(lib, x, y, theta): return 1 / lib.tan(theta) def xy_eta(lib, x, y, eta): return 1 / lib.tan(theta.xy_eta(lib, x, y, eta)) def rhophi_z(lib, rho, phi, z): return lib.nan_to_num(z / rho, nan=lib.inf, posinf=inf, neginf=-inf) def rhophi_theta(lib, rho, phi, theta): return 1 / lib.tan(theta) def rhophi_eta(lib, rho, phi, eta): return 1 / lib.tan(theta.rhophi_eta(lib, rho, phi, eta)) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/cross.py000066400000000000000000000110311503546127100220670ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.cross(self, other) Note that this returns a 3D vector even for 4D inputs. The ``None`` at the end of the return signature indicates termination. """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # Cross-product is only computed in Cartesian coordinates; the rest are conversions. def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return (y1 * z2 - z1 * y2, z1 * x2 - x1 * z2, x1 * y2 - y1 * x2) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( xy_z_xy_z, AzimuthalXY, LongitudinalZ, None, ), } def make_conversion(azimuthal1, longitudinal1, azimuthal2, longitudinal2): if (azimuthal1, longitudinal1, azimuthal2, longitudinal2) != ( AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ, ): if azimuthal1 is AzimuthalXY: to_x1 = x.xy to_y1 = y.xy if longitudinal1 is LongitudinalZ: to_z1 = z.xy_z elif longitudinal1 is LongitudinalTheta: to_z1 = z.xy_theta elif longitudinal1 is LongitudinalEta: to_z1 = z.xy_eta elif azimuthal1 is AzimuthalRhoPhi: to_x1 = x.rhophi to_y1 = y.rhophi if longitudinal1 is LongitudinalZ: to_z1 = z.rhophi_z elif longitudinal1 is LongitudinalTheta: to_z1 = z.rhophi_theta elif longitudinal1 is LongitudinalEta: to_z1 = z.rhophi_eta if azimuthal2 is AzimuthalXY: to_x2 = x.xy to_y2 = y.xy if longitudinal2 is LongitudinalZ: to_z2 = z.xy_z elif longitudinal2 is LongitudinalTheta: to_z2 = z.xy_theta elif longitudinal2 is LongitudinalEta: to_z2 = z.xy_eta elif azimuthal2 is AzimuthalRhoPhi: to_x2 = x.rhophi to_y2 = y.rhophi if longitudinal2 is LongitudinalZ: to_z2 = z.rhophi_z elif longitudinal2 is LongitudinalTheta: to_z2 = z.rhophi_theta elif longitudinal2 is LongitudinalEta: to_z2 = z.rhophi_eta cartesian, azout, lout, tout = dispatch_map[ AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ ] def f(lib, coord11, coord12, coord13, coord21, coord22, coord23): return cartesian( lib, to_x1(lib, coord11, coord12), to_y1(lib, coord11, coord12), to_z1(lib, coord11, coord12, coord13), to_x2(lib, coord21, coord22), to_y2(lib, coord21, coord22), to_z2(lib, coord21, coord22, coord23), ) dispatch_map[azimuthal1, longitudinal1, azimuthal2, longitudinal2] = ( f, azout, lout, tout, ) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): make_conversion(azimuthal1, longitudinal1, azimuthal2, longitudinal2) def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/deltaR.py000066400000000000000000000254541503546127100221670ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.deltaR(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import deltaR2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return lib.sqrt(deltaR2.xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2)) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return lib.sqrt(deltaR2.xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2)) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return lib.sqrt(deltaR2.xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2)) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return lib.sqrt(deltaR2.xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2)) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return lib.sqrt(deltaR2.xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2)) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return lib.sqrt(deltaR2.xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2)) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return lib.sqrt(deltaR2.xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2)) def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): return lib.sqrt(deltaR2.xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2)) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return lib.sqrt(deltaR2.xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2)) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return lib.sqrt(deltaR2.xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2)) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return lib.sqrt( deltaR2.xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2) ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return lib.sqrt(deltaR2.xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2)) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return lib.sqrt(deltaR2.xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2)) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return lib.sqrt(deltaR2.xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2)) def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): return lib.sqrt(deltaR2.xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2)) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return lib.sqrt(deltaR2.xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2)) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return lib.sqrt(deltaR2.xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2)) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return lib.sqrt(deltaR2.xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2)) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return lib.sqrt(deltaR2.rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2)) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return lib.sqrt(deltaR2.rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2)) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return lib.sqrt(deltaR2.rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2)) def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): return lib.sqrt(deltaR2.rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2)) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return lib.sqrt( deltaR2.rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2) ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return lib.sqrt(deltaR2.rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2)) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return lib.sqrt(deltaR2.rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2)) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return lib.sqrt( deltaR2.rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2) ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return lib.sqrt(deltaR2.rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2)) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return lib.sqrt( deltaR2.rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2) ) def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): return lib.sqrt( deltaR2.rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2) ) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return lib.sqrt( deltaR2.rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2) ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return lib.sqrt(deltaR2.rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2)) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return lib.sqrt(deltaR2.rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2)) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return lib.sqrt(deltaR2.rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2)) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return lib.sqrt(deltaR2.rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2)) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return lib.sqrt( deltaR2.rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2) ) def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): return lib.sqrt( deltaR2.rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2) ) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, float), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, float), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, float), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, float, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/deltaR2.py000066400000000000000000000324241503546127100222440ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.deltaR2(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import deltaphi from vector._compute.spatial import deltaeta from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2) ** 2 ) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2) ** 2 ) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2) ** 2 ) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2) ** 2 ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2) ** 2 ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2) ** 2 ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2) ** 2 ) def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2) ** 2 ) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2) ** 2 ) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2) ** 2 ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2) ** 2 ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2) ** 2 ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2) ** 2 ) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2) ** 2 ) def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): return ( deltaphi.xy_xy(lib, x1, y1, x2, y2) ** 2 + deltaeta.xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2) ** 2 ) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2) ** 2 ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2) ** 2 ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return ( deltaphi.xy_rhophi(lib, x1, y1, rho2, phi2) ** 2 + deltaeta.xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2) ** 2 ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2) ** 2 ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2) ** 2 ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2) ** 2 ) def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2) ** 2 ) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2) ** 2 ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2) ** 2 ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2) ** 2 ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2) ** 2 ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2) ** 2 ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2) ** 2 ) def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_theta_rhophi_theta( lib, rho1, phi1, theta1, rho2, phi2, theta2 ) ** 2 ) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2) ** 2 ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2) ** 2 ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2) ** 2 ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return ( deltaphi.rhophi_xy(lib, rho1, phi1, x2, y2) ** 2 + deltaeta.rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2) ** 2 ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2) ** 2 ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2) ** 2 ) def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): return ( deltaphi.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) ** 2 + deltaeta.rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2) ** 2 ) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, float), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, float), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, float), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, float, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/deltaangle.py000066400000000000000000000454761503546127100230620ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.deltaangle(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import dot, mag from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): v1m = mag.xy_z(lib, x1, y1, z1) v2m = mag.xy_z(lib, x2, y2, z2) return lib.arccos( lib.maximum( -1, lib.minimum(1, dot.xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2) / v1m / v2m) ) ) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): v1m = mag.xy_z(lib, x1, y1, z1) v2m = mag.xy_theta(lib, x2, y2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2) / v1m / v2m ), ) ) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): v1m = mag.xy_z(lib, x1, y1, z1) v2m = mag.xy_eta(lib, x2, y2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum(1, dot.xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2) / v1m / v2m), ) ) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): v1m = mag.xy_z(lib, x1, y1, z1) v2m = mag.rhophi_z(lib, rho2, phi2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2) / v1m / v2m ), ) ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): v1m = mag.xy_z(lib, x1, y1, z1) v2m = mag.rhophi_theta(lib, rho2, phi2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2) / v1m / v2m, ), ) ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): v1m = mag.xy_z(lib, x1, y1, z1) v2m = mag.rhophi_eta(lib, rho2, phi2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2) / v1m / v2m ), ) ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): v1m = mag.xy_theta(lib, x1, y1, theta1) v2m = mag.xy_z(lib, x2, y2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2) / v1m / v2m ), ) ) def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): v1m = mag.xy_theta(lib, x1, y1, theta1) v2m = mag.xy_theta(lib, x2, y2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2) / v1m / v2m, ), ) ) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): v1m = mag.xy_theta(lib, x1, y1, theta1) v2m = mag.xy_eta(lib, x2, y2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2) / v1m / v2m ), ) ) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): v1m = mag.xy_theta(lib, x1, y1, theta1) v2m = mag.rhophi_z(lib, rho2, phi2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2) / v1m / v2m, ), ) ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): v1m = mag.xy_theta(lib, x1, y1, theta1) v2m = mag.rhophi_theta(lib, rho2, phi2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2) / v1m / v2m, ), ) ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): v1m = mag.xy_theta(lib, x1, y1, theta1) v2m = mag.rhophi_eta(lib, rho2, phi2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2) / v1m / v2m, ), ) ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): v1m = mag.xy_eta(lib, x1, y1, eta1) v2m = mag.xy_z(lib, x2, y2, z2) return lib.arccos( lib.maximum( -1, lib.minimum(1, dot.xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2) / v1m / v2m), ) ) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): v1m = mag.xy_eta(lib, x1, y1, eta1) v2m = mag.xy_theta(lib, x2, y2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2) / v1m / v2m ), ) ) def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): v1m = mag.xy_eta(lib, x1, y1, eta1) v2m = mag.xy_eta(lib, x2, y2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2) / v1m / v2m ), ) ) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): v1m = mag.xy_eta(lib, x1, y1, eta1) v2m = mag.rhophi_z(lib, rho2, phi2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2) / v1m / v2m ), ) ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): v1m = mag.xy_eta(lib, x1, y1, eta1) v2m = mag.rhophi_theta(lib, rho2, phi2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2) / v1m / v2m, ), ) ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): v1m = mag.xy_eta(lib, x1, y1, eta1) v2m = mag.rhophi_eta(lib, rho2, phi2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2) / v1m / v2m, ), ) ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): v1m = mag.rhophi_z(lib, rho1, phi1, z1) v2m = mag.xy_z(lib, x2, y2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2) / v1m / v2m ), ) ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): v1m = mag.rhophi_z(lib, rho1, phi1, z1) v2m = mag.xy_theta(lib, x2, y2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2) / v1m / v2m, ), ) ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): v1m = mag.rhophi_z(lib, rho1, phi1, z1) v2m = mag.xy_eta(lib, x2, y2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2) / v1m / v2m ), ) ) def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): v1m = mag.rhophi_z(lib, rho1, phi1, z1) v2m = mag.rhophi_z(lib, rho2, phi2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2) / v1m / v2m, ), ) ) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): v1m = mag.rhophi_z(lib, rho1, phi1, z1) v2m = mag.rhophi_theta(lib, rho2, phi2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2) / v1m / v2m, ), ) ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): v1m = mag.rhophi_z(lib, rho1, phi1, z1) v2m = mag.rhophi_eta(lib, rho2, phi2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2) / v1m / v2m, ), ) ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): v1m = mag.rhophi_theta(lib, rho1, phi1, theta1) v2m = mag.xy_z(lib, x2, y2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2) / v1m / v2m, ), ) ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): v1m = mag.rhophi_theta(lib, rho1, phi1, theta1) v2m = mag.xy_theta(lib, x2, y2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2) / v1m / v2m, ), ) ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): v1m = mag.rhophi_theta(lib, rho1, phi1, theta1) v2m = mag.xy_eta(lib, x2, y2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2) / v1m / v2m, ), ) ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): v1m = mag.rhophi_theta(lib, rho1, phi1, theta1) v2m = mag.rhophi_z(lib, rho2, phi2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2) / v1m / v2m, ), ) ) def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): v1m = mag.rhophi_theta(lib, rho1, phi1, theta1) v2m = mag.rhophi_theta(lib, rho2, phi2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_theta_rhophi_theta( lib, rho1, phi1, theta1, rho2, phi2, theta2 ) / v1m / v2m, ), ) ) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): v1m = mag.rhophi_theta(lib, rho1, phi1, theta1) v2m = mag.rhophi_eta(lib, rho2, phi2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2) / v1m / v2m, ), ) ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): v1m = mag.rhophi_eta(lib, rho1, phi1, eta1) v2m = mag.xy_z(lib, x2, y2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2) / v1m / v2m ), ) ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): v1m = mag.rhophi_eta(lib, rho1, phi1, eta1) v2m = mag.xy_theta(lib, x2, y2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2) / v1m / v2m, ), ) ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): v1m = mag.rhophi_eta(lib, rho1, phi1, eta1) v2m = mag.xy_eta(lib, x2, y2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2) / v1m / v2m, ), ) ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): v1m = mag.rhophi_eta(lib, rho1, phi1, eta1) v2m = mag.rhophi_z(lib, rho2, phi2, z2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2) / v1m / v2m, ), ) ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): v1m = mag.rhophi_eta(lib, rho1, phi1, eta1) v2m = mag.rhophi_theta(lib, rho2, phi2, theta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2) / v1m / v2m, ), ) ) def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): v1m = mag.rhophi_eta(lib, rho1, phi1, eta1) v2m = mag.rhophi_eta(lib, rho2, phi2, eta2) return lib.arccos( lib.maximum( -1, lib.minimum( 1, dot.rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2) / v1m / v2m, ), ) ) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, float), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, float), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, float), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, float, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/deltaeta.py000066400000000000000000000236441503546127100225360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.deltaeta(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import eta from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return eta.xy_z(lib, x1, y1, z1) - eta.xy_z(lib, x2, y2, z2) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return eta.xy_z(lib, x1, y1, z1) - eta.xy_theta(lib, x2, y2, theta2) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return eta.xy_z(lib, x1, y1, z1) - eta2 def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return eta.xy_z(lib, x1, y1, z1) - eta.rhophi_z(lib, rho2, phi2, z2) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return eta.xy_z(lib, x1, y1, z1) - eta.rhophi_theta(lib, rho2, phi2, theta2) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return eta.xy_z(lib, x1, y1, z1) - eta2 def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return eta.xy_theta(lib, x1, y1, theta1) - eta.xy_z(lib, x2, y2, z2) def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): return eta.xy_theta(lib, x1, y1, theta1) - eta.xy_theta(lib, x2, y2, theta2) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return eta.xy_theta(lib, x1, y1, theta1) - eta2 def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return eta.xy_theta(lib, x1, y1, theta1) - eta.rhophi_z(lib, rho2, phi2, z2) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return eta.xy_theta(lib, x1, y1, theta1) - eta.rhophi_theta(lib, rho2, phi2, theta2) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return eta.xy_theta(lib, x1, y1, theta1) - eta2 def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return eta1 - eta.xy_z(lib, x2, y2, z2) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return eta1 - eta.xy_theta(lib, x2, y2, theta2) def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): return eta1 - eta2 def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return eta1 - eta.rhophi_z(lib, rho2, phi2, z2) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return eta1 - eta.rhophi_theta(lib, rho2, phi2, theta2) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return eta1 - eta2 def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return eta.rhophi_z(lib, rho1, phi1, z1) - eta.xy_z(lib, x2, y2, z2) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return eta.rhophi_z(lib, rho1, phi1, z1) - eta.xy_theta(lib, x2, y2, theta2) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return eta.rhophi_z(lib, rho1, phi1, z1) - eta2 def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): return eta.rhophi_z(lib, rho1, phi1, z1) - eta.rhophi_z(lib, rho2, phi2, z2) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return eta.rhophi_z(lib, rho1, phi1, z1) - eta.rhophi_theta(lib, rho2, phi2, theta2) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return eta.rhophi_z(lib, rho1, phi1, z1) - eta2 def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return eta.rhophi_theta(lib, rho1, phi1, theta1) - eta.xy_z(lib, x2, y2, z2) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return eta.rhophi_theta(lib, rho1, phi1, theta1) - eta.xy_theta(lib, x2, y2, theta2) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return eta.rhophi_theta(lib, rho1, phi1, theta1) - eta2 def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return eta.rhophi_theta(lib, rho1, phi1, theta1) - eta.rhophi_z(lib, rho2, phi2, z2) def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): return eta.rhophi_theta(lib, rho1, phi1, theta1) - eta.rhophi_theta( lib, rho2, phi2, theta2 ) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return eta.rhophi_theta(lib, rho1, phi1, theta1) - eta2 def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return eta1 - eta.xy_z(lib, x2, y2, z2) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return eta1 - eta.xy_theta(lib, x2, y2, theta2) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return eta1 - eta2 def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return eta1 - eta.rhophi_z(lib, rho2, phi2, z2) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return eta1 - eta.rhophi_theta(lib, rho2, phi2, theta2) def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): return eta1 - eta2 dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, float), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, float), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, float), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, float, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/dot.py000066400000000000000000000362231503546127100215360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.dot(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import theta, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # specialized def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return x1 * x2 + y1 * y2 + z1 * z2 def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_z(lib, x1, y1, z1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_theta(lib, x2, y2, theta2), ) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_z(lib, x1, y1, z1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_eta(lib, x2, y2, eta2), ) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_z(lib, x1, y1, z1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_z(lib, rho2, phi2, z2), ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_z(lib, x1, y1, z1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_z(lib, x1, y1, z1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_theta(lib, x1, y1, theta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_z(lib, x2, y2, z2), ) def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_theta(lib, x1, y1, theta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_theta(lib, x2, y2, theta2), ) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_theta(lib, x1, y1, theta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_eta(lib, x2, y2, eta2), ) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_z(lib, rho2, phi2, z2), ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_eta(lib, x1, y1, eta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_z(lib, x2, y2, z2), ) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_eta(lib, x1, y1, eta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_theta(lib, x2, y2, theta2), ) def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_eta(lib, x1, y1, eta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_eta(lib, x2, y2, eta2), ) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_z(lib, rho2, phi2, z2), ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.xy(lib, x1, y1), y.xy(lib, x1, y1), z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_z(lib, rho1, phi1, z1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_z(lib, x2, y2, z2), ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_z(lib, rho1, phi1, z1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_theta(lib, x2, y2, theta2), ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_z(lib, rho1, phi1, z1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_eta(lib, x2, y2, eta2), ) # specialized def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): return rho1 * rho2 * lib.cos(phi1 - phi2) + z1 * z2 def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return rhophi_z_rhophi_z( lib, rho1, phi1, z1, rho2, phi2, z.rhophi_theta(lib, rho2, phi2, theta2) ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return rhophi_z_rhophi_z( lib, rho1, phi1, z1, rho2, phi2, z.rhophi_eta(lib, rho2, phi2, eta2) ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_z(lib, x2, y2, z2), ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_theta(lib, x2, y2, theta2), ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_eta(lib, x2, y2, eta2), ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return rhophi_z_rhophi_z( lib, rho1, phi1, z.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, z2 ) # specialized def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): return ( rho1 * rho2 * (lib.cos(phi1 - phi2) + 1 / (lib.tan(theta1) * lib.tan(theta2))) ) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return rhophi_theta_rhophi_theta( lib, rho1, phi1, theta1, rho2, phi2, theta.rhophi_eta(lib, rho2, phi2, eta2) ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_z(lib, x2, y2, z2), ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_theta(lib, x2, y2, theta2), ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.xy(lib, x2, y2), y.xy(lib, x2, y2), z.xy_eta(lib, x2, y2, eta2), ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return rhophi_theta_rhophi_theta( lib, rho1, phi1, theta.rhophi_eta(lib, rho1, phi1, eta1), rho2, phi2, theta.rhophi_z(lib, rho2, phi2, z2), ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return rhophi_theta_rhophi_theta( lib, rho1, phi1, theta.rhophi_eta(lib, rho1, phi1, eta1), rho2, phi2, theta2 ) # specialized def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): expmeta1 = lib.exp(-eta1) expmeta2 = lib.exp(-eta2) invtantheta1 = 0.5 * (1 - expmeta1**2) / expmeta1 invtantheta2 = 0.5 * (1 - expmeta2**2) / expmeta2 return rho1 * rho2 * (lib.cos(phi1 - phi2) + invtantheta1 * invtantheta2) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, float), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, float), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, float, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, float), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, float, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, float, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, float, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/equal.py000066400000000000000000000307551503546127100220630ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.equal(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import eta, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # Policy: turn (rho, phi) into (x, y) # turn theta into z, eta into z, theta into eta # (if not already the same) # same types def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return (x1 == x2) & (y1 == y2) & (z1 == z2) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_theta(lib, x2, y2, theta2)) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_eta(lib, x2, y2, eta2)) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2 ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z2) # same types def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): return (x1 == x2) & (y1 == y2) & (theta1 == theta2) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return xy_eta_xy_eta(lib, x1, y1, eta.xy_theta(lib, x1, y1, theta1), x2, y2, eta2) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return xy_theta_xy_theta( lib, x1, y1, theta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), theta2, ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return xy_eta_xy_eta( lib, x1, y1, eta.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta2, ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z2) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta.xy_theta(lib, x2, y2, theta2)) # same types def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): return (x1 == x2) & (y1 == y2) & (eta1 == eta2) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return xy_eta_xy_eta( lib, x1, y1, eta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return xy_eta_xy_eta( lib, x1, y1, eta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta2 ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z2 ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_eta(lib, x2, y2, eta2), ) # same types def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): return (rho1 == rho2) & (phi1 == phi2) & (z1 == z2) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return rhophi_z_rhophi_z( lib, rho1, phi1, z1, rho2, phi2, z.rhophi_theta(lib, rho2, phi2, theta2) ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return rhophi_z_rhophi_z( lib, rho1, phi1, z1, rho2, phi2, z.rhophi_eta(lib, rho2, phi2, eta2) ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z2, ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return xy_theta_xy_theta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), theta1, x2, y2, theta2, ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return xy_eta_xy_eta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, eta2, ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return rhophi_z_rhophi_z( lib, rho1, phi1, z.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, z2 ) # same types def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): return (rho1 == rho2) & (phi1 == phi2) & (theta1 == theta2) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return rhophi_eta_rhophi_eta( lib, rho1, phi1, eta.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, eta2 ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z2, ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return xy_eta_xy_eta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta1, x2, y2, eta.xy_theta(lib, x2, y2, theta2), ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return xy_eta_xy_eta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta1, x2, y2, eta2 ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return rhophi_z_rhophi_z( lib, rho1, phi1, z.rhophi_eta(lib, rho1, phi1, eta1), rho2, phi2, z2 ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return rhophi_eta_rhophi_eta( lib, rho1, phi1, eta1, rho2, phi2, eta.rhophi_theta(lib, rho2, phi2, theta2) ) # same types def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): return (rho1 == rho2) & (phi1 == phi2) & (eta1 == eta2) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, bool), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, bool), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, bool), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, bool, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/eta.py000066400000000000000000000051451503546127100215200ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.eta(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) # TODO: https://github.com/scikit-hep/vector/issues/615 # revert back to `nan_to_num` implementation once # https://github.com/cupy/cupy/issues/9143 is fixed. # `lib.where` works but there is no SymPy equivalent for the function. def xy_z(lib, x, y, z): return ( lib.where( z != 0, lib.arcsinh(lib.where(z != 0, z / lib.sqrt(x**2 + y**2), z)), z ) * 1 ) def xy_theta(lib, x, y, theta): return lib.nan_to_num( -lib.log(lib.tan(0.5 * theta)), nan=0.0, posinf=inf, neginf=-inf ) def xy_eta(lib, x, y, eta): return eta xy_eta.__awkward_transform_allowed__ = False # type:ignore[attr-defined] # TODO: https://github.com/scikit-hep/vector/issues/615 # revert back to `nan_to_num` implementation once # https://github.com/cupy/cupy/issues/9143 is fixed. # `lib.where` works but there is no SymPy equivalent for the function. def rhophi_z(lib, rho, phi, z): return lib.where(z != 0, lib.arcsinh(lib.where(z != 0, z / rho, z)), z) * 1 def rhophi_theta(lib, rho, phi, theta): return -lib.log(lib.tan(0.5 * theta)) def rhophi_eta(lib, rho, phi, eta): return eta rhophi_eta.__awkward_transform_allowed__ = False # type:ignore[attr-defined] dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/is_antiparallel.py000066400000000000000000000050111503546127100241020ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.is_antiparallel(self, other, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import dot, mag from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) dispatch_map = {} def make_function(azimuthal1, longitudinal1, azimuthal2, longitudinal2): dot_function, _ = dot.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] mag1_function, _ = mag.dispatch_map[azimuthal1, longitudinal1] mag2_function, _ = mag.dispatch_map[azimuthal2, longitudinal2] def f(lib, tolerance, coord11, coord12, coord13, coord21, coord22, coord23): return dot_function( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) < (lib.absolute(tolerance) - 1) * mag1_function( lib, coord11, coord12, coord13 ) * mag2_function(lib, coord21, coord22, coord23) dispatch_map[azimuthal1, longitudinal1, azimuthal2, longitudinal2] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): make_function(azimuthal1, longitudinal1, azimuthal2, longitudinal2) def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), tolerance, *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/is_parallel.py000066400000000000000000000050051503546127100232310ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.is_parallel(self, other, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import dot, mag from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) dispatch_map = {} def make_function(azimuthal1, longitudinal1, azimuthal2, longitudinal2): dot_function, _ = dot.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] mag1_function, _ = mag.dispatch_map[azimuthal1, longitudinal1] mag2_function, _ = mag.dispatch_map[azimuthal2, longitudinal2] def f(lib, tolerance, coord11, coord12, coord13, coord21, coord22, coord23): return dot_function( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) > (1 - lib.absolute(tolerance)) * mag1_function( lib, coord11, coord12, coord13 ) * mag2_function(lib, coord21, coord22, coord23) dispatch_map[azimuthal1, longitudinal1, azimuthal2, longitudinal2] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): make_function(azimuthal1, longitudinal1, azimuthal2, longitudinal2) def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), tolerance, *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/is_perpendicular.py000066400000000000000000000050041503546127100242710ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.is_perpendicular(self, other, tolerance=...) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import dot, mag from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) dispatch_map = {} def make_function(azimuthal1, longitudinal1, azimuthal2, longitudinal2): dot_function, _ = dot.dispatch_map[ azimuthal1, longitudinal1, azimuthal2, longitudinal2 ] mag1_function, _ = mag.dispatch_map[azimuthal1, longitudinal1] mag2_function, _ = mag.dispatch_map[azimuthal2, longitudinal2] def f(lib, tolerance, coord11, coord12, coord13, coord21, coord22, coord23): return dot_function( lib, coord11, coord12, coord13, coord21, coord22, coord23 ) < lib.absolute(tolerance) * mag1_function( lib, coord11, coord12, coord13 ) * mag2_function(lib, coord21, coord22, coord23) dispatch_map[azimuthal1, longitudinal1, azimuthal2, longitudinal2] = (f, bool) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): make_function(azimuthal1, longitudinal1, azimuthal2, longitudinal2) def dispatch(tolerance: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), tolerance, *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/isclose.py000066400000000000000000000403311503546127100224040ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.isclose(self, other, rtol=..., atol=..., equal_nan=...) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import eta, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # Policy: turn (rho, phi) into (x, y) # turn theta into z, eta into z, theta into eta # (if not already the same) # same types def xy_z_xy_z(lib, rtol, atol, equal_nan, x1, y1, z1, x2, y2, z2): return ( lib.isclose(x1, x2, rtol, atol, equal_nan) & lib.isclose(y1, y2, rtol, atol, equal_nan) & lib.isclose(z1, z2, rtol, atol, equal_nan) ) def xy_z_xy_theta(lib, rtol, atol, equal_nan, x1, y1, z1, x2, y2, theta2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z1, x2, y2, z.xy_theta(lib, x2, y2, theta2) ) def xy_z_xy_eta(lib, rtol, atol, equal_nan, x1, y1, z1, x2, y2, eta2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z1, x2, y2, z.xy_eta(lib, x2, y2, eta2) ) def xy_z_rhophi_z(lib, rtol, atol, equal_nan, x1, y1, z1, rho2, phi2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_z_rhophi_theta(lib, rtol, atol, equal_nan, x1, y1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_z_rhophi_eta(lib, rtol, atol, equal_nan, x1, y1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_theta_xy_z(lib, rtol, atol, equal_nan, x1, y1, theta1, x2, y2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z2 ) # same types def xy_theta_xy_theta(lib, rtol, atol, equal_nan, x1, y1, theta1, x2, y2, theta2): return ( lib.isclose(x1, x2, rtol, atol, equal_nan) & lib.isclose(y1, y2, rtol, atol, equal_nan) & lib.isclose(theta1, theta2, rtol, atol, equal_nan) ) def xy_theta_xy_eta(lib, rtol, atol, equal_nan, x1, y1, theta1, x2, y2, eta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x1, y1, eta.xy_theta(lib, x1, y1, theta1), x2, y2, eta2, ) def xy_theta_rhophi_z(lib, rtol, atol, equal_nan, x1, y1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_theta_rhophi_theta( lib, rtol, atol, equal_nan, x1, y1, theta1, rho2, phi2, theta2 ): return xy_theta_xy_theta( lib, rtol, atol, equal_nan, x1, y1, theta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), theta2, ) def xy_theta_rhophi_eta(lib, rtol, atol, equal_nan, x1, y1, theta1, rho2, phi2, eta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x1, y1, eta.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta2, ) def xy_eta_xy_z(lib, rtol, atol, equal_nan, x1, y1, eta1, x2, y2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z2 ) def xy_eta_xy_theta(lib, rtol, atol, equal_nan, x1, y1, eta1, x2, y2, theta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x1, y1, eta1, x2, y2, eta.xy_theta(lib, x2, y2, theta2), ) # same types def xy_eta_xy_eta(lib, rtol, atol, equal_nan, x1, y1, eta1, x2, y2, eta2): return ( lib.isclose(x1, x2, rtol, atol, equal_nan) & lib.isclose(y1, y2, rtol, atol, equal_nan) & lib.isclose(eta1, eta2, rtol, atol, equal_nan) ) def xy_eta_rhophi_z(lib, rtol, atol, equal_nan, x1, y1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_eta_rhophi_theta(lib, rtol, atol, equal_nan, x1, y1, eta1, rho2, phi2, theta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x1, y1, eta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_eta_rhophi_eta(lib, rtol, atol, equal_nan, x1, y1, eta1, rho2, phi2, eta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x1, y1, eta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta2, ) def rhophi_z_xy_z(lib, rtol, atol, equal_nan, rho1, phi1, z1, x2, y2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z2, ) def rhophi_z_xy_theta(lib, rtol, atol, equal_nan, rho1, phi1, z1, x2, y2, theta2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_z_xy_eta(lib, rtol, atol, equal_nan, rho1, phi1, z1, x2, y2, eta2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_eta(lib, x2, y2, eta2), ) # same types def rhophi_z_rhophi_z(lib, rtol, atol, equal_nan, rho1, phi1, z1, rho2, phi2, z2): return ( lib.isclose(rho1, rho2, rtol, atol, equal_nan) & lib.isclose(phi1, phi2, rtol, atol, equal_nan) & lib.isclose(z1, z2, rtol, atol, equal_nan) ) def rhophi_z_rhophi_theta( lib, rtol, atol, equal_nan, rho1, phi1, z1, rho2, phi2, theta2 ): return rhophi_z_rhophi_z( lib, rtol, atol, equal_nan, rho1, phi1, z1, rho2, phi2, z.rhophi_theta(lib, rho2, phi2, theta2), ) def rhophi_z_rhophi_eta(lib, rtol, atol, equal_nan, rho1, phi1, z1, rho2, phi2, eta2): return rhophi_z_rhophi_z( lib, rtol, atol, equal_nan, rho1, phi1, z1, rho2, phi2, z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_theta_xy_z(lib, rtol, atol, equal_nan, rho1, phi1, theta1, x2, y2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z2, ) def rhophi_theta_xy_theta( lib, rtol, atol, equal_nan, rho1, phi1, theta1, x2, y2, theta2 ): return xy_theta_xy_theta( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), theta1, x2, y2, theta2, ) def rhophi_theta_xy_eta(lib, rtol, atol, equal_nan, rho1, phi1, theta1, x2, y2, eta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, eta2, ) def rhophi_theta_rhophi_z( lib, rtol, atol, equal_nan, rho1, phi1, theta1, rho2, phi2, z2 ): return rhophi_z_rhophi_z( lib, rtol, atol, equal_nan, rho1, phi1, z.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, z2, ) # same types def rhophi_theta_rhophi_theta( lib, rtol, atol, equal_nan, rho1, phi1, theta1, rho2, phi2, theta2 ): return ( lib.isclose(rho1, rho2, rtol, atol, equal_nan) & lib.isclose(phi1, phi2, rtol, atol, equal_nan) & lib.isclose(theta1, theta2, rtol, atol, equal_nan) ) def rhophi_theta_rhophi_eta( lib, rtol, atol, equal_nan, rho1, phi1, theta1, rho2, phi2, eta2 ): return rhophi_eta_rhophi_eta( lib, rtol, atol, equal_nan, rho1, phi1, eta.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, eta2, ) def rhophi_eta_xy_z(lib, rtol, atol, equal_nan, rho1, phi1, eta1, x2, y2, z2): return xy_z_xy_z( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z2, ) def rhophi_eta_xy_theta(lib, rtol, atol, equal_nan, rho1, phi1, eta1, x2, y2, theta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta1, x2, y2, eta.xy_theta(lib, x2, y2, theta2), ) def rhophi_eta_xy_eta(lib, rtol, atol, equal_nan, rho1, phi1, eta1, x2, y2, eta2): return xy_eta_xy_eta( lib, rtol, atol, equal_nan, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta1, x2, y2, eta2, ) def rhophi_eta_rhophi_z(lib, rtol, atol, equal_nan, rho1, phi1, eta1, rho2, phi2, z2): return rhophi_z_rhophi_z( lib, rtol, atol, equal_nan, rho1, phi1, z.rhophi_eta(lib, rho1, phi1, eta1), rho2, phi2, z2, ) def rhophi_eta_rhophi_theta( lib, rtol, atol, equal_nan, rho1, phi1, eta1, rho2, phi2, theta2 ): return rhophi_eta_rhophi_eta( lib, rtol, atol, equal_nan, rho1, phi1, eta1, rho2, phi2, eta.rhophi_theta(lib, rho2, phi2, theta2), ) # same types def rhophi_eta_rhophi_eta( lib, rtol, atol, equal_nan, rho1, phi1, eta1, rho2, phi2, eta2 ): return ( lib.isclose(rho1, rho2, rtol, atol, equal_nan) & lib.isclose(phi1, phi2, rtol, atol, equal_nan) & lib.isclose(eta1, eta2, rtol, atol, equal_nan) ) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, bool), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, bool), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, bool), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, bool, ), } def dispatch( rtol: typing.Any, atol: typing.Any, equal_nan: typing.Any, v1: typing.Any, v2: typing.Any, ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), rtol, atol, equal_nan, *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/mag.py000066400000000000000000000040711503546127100215100ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.mag(self) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import mag2 from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): return lib.sqrt(mag2.xy_z(lib, x, y, z)) def xy_theta(lib, x, y, theta): return lib.sqrt(x**2 + y**2) / lib.absolute(lib.sin(theta)) def xy_eta(lib, x, y, eta): expmeta = lib.exp(-eta) invsintheta = 0.5 * (1 + expmeta**2) / expmeta return lib.sqrt(x**2 + y**2) * invsintheta def rhophi_z(lib, rho, phi, z): return lib.sqrt(mag2.rhophi_z(lib, rho, phi, z)) def rhophi_theta(lib, rho, phi, theta): return rho / lib.absolute(lib.sin(theta)) def rhophi_eta(lib, rho, phi, eta): expmeta = lib.exp(-eta) invsintheta = 0.5 * (1 + expmeta**2) / expmeta return rho * invsintheta dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/mag2.py000066400000000000000000000037201503546127100215720ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.mag2(self) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): return x**2 + y**2 + z**2 def xy_theta(lib, x, y, theta): return (x**2 + y**2) / lib.sin(theta) ** 2 def xy_eta(lib, x, y, eta): expmeta = lib.exp(-eta) invsintheta = 0.5 * (1 + expmeta**2) / expmeta return (x**2 + y**2) * invsintheta**2 def rhophi_z(lib, rho, phi, z): return rho**2 + z**2 def rhophi_theta(lib, rho, phi, theta): return rho**2 / lib.sin(theta) ** 2 def rhophi_eta(lib, rho, phi, eta): expmeta = lib.exp(-eta) invsintheta = 0.5 * (1 + expmeta**2) / expmeta return rho**2 * invsintheta**2 dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/not_equal.py000066400000000000000000000307611503546127100227400ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.not_equal(self, other) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import eta, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # Policy: turn (rho, phi) into (x, y) # turn theta into z, eta into z, theta into eta # (if not already the same) # same types def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return (x1 != x2) & (y1 != y2) & (z1 != z2) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_theta(lib, x2, y2, theta2)) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_eta(lib, x2, y2, eta2)) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2 ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z2) # same types def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): return (x1 != x2) & (y1 != y2) & (theta1 != theta2) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return xy_eta_xy_eta(lib, x1, y1, eta.xy_theta(lib, x1, y1, theta1), x2, y2, eta2) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return xy_theta_xy_theta( lib, x1, y1, theta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), theta2, ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return xy_eta_xy_eta( lib, x1, y1, eta.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta2, ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z2) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta.xy_theta(lib, x2, y2, theta2)) # same types def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): return (x1 != x2) & (y1 != y2) & (eta1 != eta2) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return xy_eta_xy_eta( lib, x1, y1, eta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return xy_eta_xy_eta( lib, x1, y1, eta1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), eta2 ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z2 ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_eta(lib, x2, y2, eta2), ) # same types def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): return (rho1 != rho2) & (phi1 != phi2) & (z1 != z2) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return rhophi_z_rhophi_z( lib, rho1, phi1, z1, rho2, phi2, z.rhophi_theta(lib, rho2, phi2, theta2) ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return rhophi_z_rhophi_z( lib, rho1, phi1, z1, rho2, phi2, z.rhophi_eta(lib, rho2, phi2, eta2) ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z2, ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return xy_theta_xy_theta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), theta1, x2, y2, theta2, ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return xy_eta_xy_eta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, eta2, ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return rhophi_z_rhophi_z( lib, rho1, phi1, z.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, z2 ) # same types def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): return (rho1 != rho2) & (phi1 != phi2) & (theta1 != theta2) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return rhophi_eta_rhophi_eta( lib, rho1, phi1, eta.rhophi_theta(lib, rho1, phi1, theta1), rho2, phi2, eta2 ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z2, ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return xy_eta_xy_eta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta1, x2, y2, eta.xy_theta(lib, x2, y2, theta2), ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return xy_eta_xy_eta( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), eta1, x2, y2, eta2 ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return rhophi_z_rhophi_z( lib, rho1, phi1, z.rhophi_eta(lib, rho1, phi1, eta1), rho2, phi2, z2 ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return rhophi_eta_rhophi_eta( lib, rho1, phi1, eta1, rho2, phi2, eta.rhophi_theta(lib, rho2, phi2, theta2) ) # same types def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): return (rho1 != rho2) & (phi1 != phi2) & (eta1 != eta2) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): (xy_z_xy_z, bool), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): (xy_z_xy_eta, bool), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): (xy_eta_xy_z, bool), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, bool, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, bool, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, bool, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/rotateX.py000066400000000000000000000050361503546127100223740ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.rotateX(self, angle) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. def xy_z(lib, angle, x, y, z): s = lib.sin(angle) c = lib.cos(angle) return (x, c * y - s * z, s * y + c * z) def xy_theta(lib, angle, x, y, theta): return xy_z(lib, angle, x, y, z.xy_theta(lib, x, y, theta)) def xy_eta(lib, angle, x, y, eta): return xy_z(lib, angle, x, y, z.xy_eta(lib, x, y, eta)) def rhophi_z(lib, angle, rho, phi, z): return xy_z(lib, angle, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z) def rhophi_theta(lib, angle, rho, phi, theta): return xy_z( lib, angle, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z.rhophi_theta(lib, rho, phi, theta), ) def rhophi_eta(lib, angle, rho, phi, eta): return xy_z( lib, angle, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z.rhophi_eta(lib, rho, phi, eta), ) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalTheta): (xy_theta, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalEta): (xy_eta, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, AzimuthalXY, LongitudinalZ), } def dispatch(angle: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, angle, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/rotateY.py000066400000000000000000000050371503546127100223760ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.rotateY(self, angle) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. def xy_z(lib, angle, x, y, z): s = lib.sin(angle) c = lib.cos(angle) return (c * x + s * z, y, -s * x + c * z) def xy_theta(lib, angle, x, y, theta): return xy_z(lib, angle, x, y, z.xy_theta(lib, x, y, theta)) def xy_eta(lib, angle, x, y, eta): return xy_z(lib, angle, x, y, z.xy_eta(lib, x, y, eta)) def rhophi_z(lib, angle, rho, phi, z): return xy_z(lib, angle, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z) def rhophi_theta(lib, angle, rho, phi, theta): return xy_z( lib, angle, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z.rhophi_theta(lib, rho, phi, theta), ) def rhophi_eta(lib, angle, rho, phi, eta): return xy_z( lib, angle, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z.rhophi_eta(lib, rho, phi, eta), ) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalTheta): (xy_theta, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalEta): (xy_eta, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, AzimuthalXY, LongitudinalZ), } def dispatch(angle: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, angle, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/rotate_axis.py000066400000000000000000000120761503546127100232720ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.rotate_axis(self, axis, angle) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. def cartesian(lib, angle, x1, y1, z1, x2, y2, z2): norm = lib.sqrt(x1**2 + y1**2 + z1**2) ux = x1 / norm uy = y1 / norm uz = z1 / norm c = lib.cos(angle) s = lib.sin(angle) c1 = 1 - c xp = ( (c + ux**2 * c1) * x2 + (ux * uy * c1 - uz * s) * y2 + (ux * uz * c1 + uy * s) * z2 ) yp = ( (ux * uy * c1 + uz * s) * x2 + (c + uy**2 * c1) * y2 + (uy * uz * c1 - ux * s) * z2 ) zp = ( (ux * uz * c1 - uy * s) * x2 + (uy * uz * c1 + ux * s) * y2 + (c + uz**2 * c1) * z2 ) return (xp, yp, zp) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( cartesian, AzimuthalXY, LongitudinalZ, ), } def make_conversion(azimuthal1, longitudinal1, azimuthal2, longitudinal2): if (azimuthal1, longitudinal1, azimuthal2, longitudinal2) != ( AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ, ): if azimuthal1 is AzimuthalXY: to_x1 = x.xy to_y1 = y.xy if longitudinal1 is LongitudinalZ: to_z1 = z.xy_z elif longitudinal1 is LongitudinalTheta: to_z1 = z.xy_theta elif longitudinal1 is LongitudinalEta: to_z1 = z.xy_eta elif azimuthal1 is AzimuthalRhoPhi: to_x1 = x.rhophi to_y1 = y.rhophi if longitudinal1 is LongitudinalZ: to_z1 = z.rhophi_z elif longitudinal1 is LongitudinalTheta: to_z1 = z.rhophi_theta elif longitudinal1 is LongitudinalEta: to_z1 = z.rhophi_eta if azimuthal2 is AzimuthalXY: to_x2 = x.xy to_y2 = y.xy if longitudinal2 is LongitudinalZ: to_z2 = z.xy_z elif longitudinal2 is LongitudinalTheta: to_z2 = z.xy_theta elif longitudinal2 is LongitudinalEta: to_z2 = z.xy_eta elif azimuthal2 is AzimuthalRhoPhi: to_x2 = x.rhophi to_y2 = y.rhophi if longitudinal2 is LongitudinalZ: to_z2 = z.rhophi_z elif longitudinal2 is LongitudinalTheta: to_z2 = z.rhophi_theta elif longitudinal2 is LongitudinalEta: to_z2 = z.rhophi_eta cartesian, azout, lout = dispatch_map[ AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ ] def f(lib, angle, coord11, coord12, coord13, coord21, coord22, coord23): return cartesian( lib, angle, to_x1(lib, coord11, coord12), to_y1(lib, coord11, coord12), to_z1(lib, coord11, coord12, coord13), to_x2(lib, coord21, coord22), to_y2(lib, coord21, coord22), to_z2(lib, coord21, coord22, coord23), ) dispatch_map[azimuthal1, longitudinal1, azimuthal2, longitudinal2] = ( f, azout, lout, ) for azimuthal1 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal1 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for azimuthal2 in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal2 in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): make_conversion(azimuthal1, longitudinal1, azimuthal2, longitudinal2) def dispatch(angle: typing.Any, v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), # v1 is the axis about which we're rotating _ltype(v1), _aztype(v2), # v2 is the primary vector, the one being rotated _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v2) return handler._wrap_result( # note: _handler_of(v2) _flavor_of(v2), # note: _flavor_of(v2) handler._wrap_dispatched_function(function)( _lib_of(v1, v2), angle, *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/rotate_euler.py000066400000000000000000000236571503546127100234510ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.rotate_euler(self, phi, theta, psi, order=...) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. # Matrices copied from https://en.wikipedia.org/wiki/Euler_angles#Rotation_matrix # # ROOT uses the same names, phi, theta, psi, but takes the arguments in the # opposite order (i.e. first psi, then theta, finally phi) and takes their # direction to be the opposite way (e.g. Wikipedia's psi is ROOT's -psi for all # three). I've left the matrices in terms of the same c1, c2, c3, s1, s2, s3, # but changed the definitions of c1, c2, c3, s1, s2, s3 to agree with ROOT. # # Also, ROOT's angle order convention is "zxz", so that's our default. # # https://github.com/root-project/root/blob/f8efb11a51cbe5b5152ebef19a4f7b78744ca2fa/math/genvector/src/3DConversions.cxx#L420-L436 def cartesian_xzx(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c2) * x + (-c3 * s2) * y + (s2 * s3) * z yp = (c1 * s2) * x + (c1 * c2 * c3 - s1 * s3) * y + (-c3 * s1 - c1 * c2 * s3) * z zp = (s1 * s2) * x + (c1 * s3 + c2 * c3 * s1) * y + (c1 * c3 - c2 * s1 * s3) * z return (xp, yp, zp) def cartesian_xyx(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c2) * x + (s2 * s3) * y + (c3 * s2) * z yp = (s1 * s2) * x + (c1 * c3 - c2 * s1 * s3) * y + (-c1 * s3 - c2 * c3 * s1) * z zp = (-c1 * s2) * x + (c3 * s1 + c1 * c2 * s3) * y + (c1 * c2 * c3 - s1 * s3) * z return (xp, yp, zp) def cartesian_yxy(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c3 - c2 * s1 * s3) * x + (s1 * s2) * y + (c1 * s3 + c2 * c3 * s1) * z yp = (s2 * s3) * x + (c2) * y + (-c3 * s2) * z zp = (-c3 * s1 - c1 * c2 * s3) * x + (c1 * s2) * y + (c1 * c2 * c3 - s1 * s3) * z return (xp, yp, zp) def cartesian_yzy(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c2 * c3 - s1 * s3) * x + (-c1 * s2) * y + (c3 * s1 + c1 * c2 * s3) * z yp = (c3 * s2) * x + (c2) * y + (s2 * s3) * z zp = (-c1 * s3 - c2 * c3 * s1) * x + (s1 * s2) * y + (c1 * c3 - c2 * s1 * s3) * z return (xp, yp, zp) def cartesian_zyz(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c2 * c3 - s1 * s3) * x + (-c3 * s1 - c1 * c2 * s3) * y + (c1 * s2) * z yp = (c1 * s3 + c2 * c3 * s1) * x + (c1 * c3 - c2 * s1 * s3) * y + (s1 * s2) * z zp = (-c3 * s2) * x + (s2 * s3) * y + (c2) * z return (xp, yp, zp) def cartesian_zxz(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c3 - c2 * s1 * s3) * x + (-c1 * s3 - c2 * c3 * s1) * y + (s1 * s2) * z yp = (c3 * s1 + c1 * c2 * s3) * x + (c1 * c2 * c3 - s1 * s3) * y + (-c1 * s2) * z zp = (s2 * s3) * x + (c3 * s2) * y + (c2) * z return (xp, yp, zp) def cartesian_xzy(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c2 * c3) * x + (-s2) * y + (c2 * s3) * z yp = (s1 * s3 + c1 * c3 * s2) * x + (c1 * c2) * y + (c1 * s2 * s3 - c3 * s1) * z zp = (c3 * s1 * s2 - c1 * s3) * x + (c2 * s1) * y + (c1 * c3 + s1 * s2 * s3) * z return (xp, yp, zp) def cartesian_xyz(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c2 * c3) * x + (-c2 * s3) * y + (s2) * z yp = (c1 * s3 + c3 * s1 * s2) * x + (c1 * c3 - s1 * s2 * s3) * y + (-c2 * s1) * z zp = (s1 * s3 - c1 * c3 * s2) * x + (c3 * s1 + c1 * s2 * s3) * y + (c1 * c2) * z return (xp, yp, zp) def cartesian_yxz(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c3 + s1 * s2 * s3) * x + (c3 * s1 * s2 - c1 * s3) * y + (c2 * s1) * z yp = (c2 * s3) * x + (c2 * c3) * y + (-s2) * z zp = (c1 * s2 * s3 - c3 * s1) * x + (c1 * c3 * s2 + s1 * s3) * y + (c1 * c2) * z return (xp, yp, zp) def cartesian_yzx(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c2) * x + (s1 * s3 - c1 * c3 * s2) * y + (c3 * s1 + c1 * s2 * s3) * z yp = (s2) * x + (c2 * c3) * y + (-c2 * s3) * z zp = (-c2 * s1) * x + (c1 * s3 + c3 * s1 * s2) * y + (c1 * c3 - s1 * s2 * s3) * z return (xp, yp, zp) def cartesian_zyx(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c2) * x + (c1 * s2 * s3 - c3 * s1) * y + (s1 * s3 + c1 * c3 * s2) * z yp = (c2 * s1) * x + (c1 * c3 + s1 * s2 * s3) * y + (c3 * s1 * s2 - c1 * s3) * z zp = (-s2) * x + (c2 * s3) * y + (c2 * c3) * z return (xp, yp, zp) def cartesian_zxy(lib, phi, theta, psi, x, y, z): c1 = lib.cos(psi) s1 = -lib.sin(psi) c2 = lib.cos(theta) s2 = -lib.sin(theta) c3 = lib.cos(phi) s3 = -lib.sin(phi) xp = (c1 * c3 - s1 * s2 * s3) * x + (-c2 * s1) * y + (c1 * s3 + c3 * s1 * s2) * z yp = (c3 * s1 + c1 * s2 * s3) * x + (c1 * c2) * y + (s1 * s3 - c1 * c3 * s2) * z zp = (-c2 * s3) * x + (s2) * y + (c2 * c3) * z return (xp, yp, zp) dispatch_map = { (AzimuthalXY, LongitudinalZ, "xzx"): (cartesian_xzx, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "xyx"): (cartesian_xyx, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "yxy"): (cartesian_yxy, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "yzy"): (cartesian_yzy, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "zyz"): (cartesian_zyz, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "zxz"): (cartesian_zxz, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "xzy"): (cartesian_xzy, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "xyz"): (cartesian_xyz, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "yxz"): (cartesian_yxz, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "yzx"): (cartesian_yzx, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "zyx"): (cartesian_zyx, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalZ, "zxy"): (cartesian_zxy, AzimuthalXY, LongitudinalZ), } def make_conversion(azimuthal, longitudinal, order): if (azimuthal, longitudinal) != (AzimuthalXY, LongitudinalZ): if azimuthal is AzimuthalXY: to_x = x.xy to_y = y.xy if longitudinal is LongitudinalZ: to_z = z.xy_z elif longitudinal is LongitudinalTheta: to_z = z.xy_theta elif longitudinal is LongitudinalEta: to_z = z.xy_eta elif azimuthal is AzimuthalRhoPhi: to_x = x.rhophi to_y = y.rhophi if longitudinal is LongitudinalZ: to_z = z.rhophi_z elif longitudinal is LongitudinalTheta: to_z = z.rhophi_theta elif longitudinal is LongitudinalEta: to_z = z.rhophi_eta cartesian, azout, lout = dispatch_map[AzimuthalXY, LongitudinalZ, order] def f(lib, phi, theta, psi, coord1, coord2, coord3): return cartesian( lib, phi, theta, psi, to_x(lib, coord1, coord2), to_y(lib, coord1, coord2), to_z(lib, coord1, coord2, coord3), ) dispatch_map[azimuthal, longitudinal, order] = (f, azout, lout) for azimuthal in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): for order in ( "xzx", "xyx", "yxy", "yzy", "zyz", "zxz", "xzy", "xyz", "yxz", "yzx", "zyx", "zxy", ): make_conversion(azimuthal, longitudinal, order) def dispatch( phi: typing.Any, theta: typing.Any, psi: typing.Any, order: typing.Any, v: typing.Any, ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), order, ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, phi, theta, psi, *v.azimuthal.elements, *v.longitudinal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/rotate_quaternion.py000066400000000000000000000071431503546127100245120ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.rotate_quaternion(self, u, i, j, k) """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. # Follows ROOT's conventions. # # https://github.com/root-project/root/blob/f8efb11a51cbe5b5152ebef19a4f7b78744ca2fa/math/genvector/src/3DConversions.cxx#L478-L502 # # I don't know how this relates to Wikipedia's representation: # # https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation#Quaternion-derived_rotation_matrix def cartesian(lib, u, i, j, k, x, y, z): q00 = u * u q01 = u * i q02 = u * j q03 = u * k q11 = i * i q12 = i * j q13 = i * k q22 = j * j q23 = j * k q33 = k * k xp = (q00 + q11 - q22 - q33) * x + (2 * (q12 - q03)) * y + (2 * (q02 + q13)) * z yp = (2 * (q12 + q03)) * x + (q00 - q11 + q22 - q33) * y + (2 * (q23 - q01)) * z zp = (2 * (q13 - q02)) * x + (2 * (q23 + q01)) * y + (q00 - q11 - q22 + q33) * z return (xp, yp, zp) dispatch_map = { (AzimuthalXY, LongitudinalZ): (cartesian, AzimuthalXY, LongitudinalZ), } def make_conversion(azimuthal, longitudinal): if (azimuthal, longitudinal) != (AzimuthalXY, LongitudinalZ): if azimuthal is AzimuthalXY: to_x = x.xy to_y = y.xy if longitudinal is LongitudinalZ: to_z = z.xy_z elif longitudinal is LongitudinalTheta: to_z = z.xy_theta elif longitudinal is LongitudinalEta: to_z = z.xy_eta elif azimuthal is AzimuthalRhoPhi: to_x = x.rhophi to_y = y.rhophi if longitudinal is LongitudinalZ: to_z = z.rhophi_z elif longitudinal is LongitudinalTheta: to_z = z.rhophi_theta elif longitudinal is LongitudinalEta: to_z = z.rhophi_eta cartesian, azout, lout = dispatch_map[AzimuthalXY, LongitudinalZ] def f(lib, u, i, j, k, coord1, coord2, coord3): return cartesian( lib, u, i, j, k, to_x(lib, coord1, coord2), to_y(lib, coord1, coord2), to_z(lib, coord1, coord2, coord3), ) dispatch_map[azimuthal, longitudinal] = (f, azout, lout) for azimuthal in (AzimuthalXY, AzimuthalRhoPhi): for longitudinal in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): make_conversion(azimuthal, longitudinal) def dispatch( u: typing.Any, i: typing.Any, j: typing.Any, k: typing.Any, vec: typing.Any ) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(vec), _ltype(vec), ), ) with numpy.errstate(all="ignore"): return vec._wrap_result( _flavor_of(vec), vec._wrap_dispatched_function(function)( vec.lib, u, i, j, k, *vec.azimuthal.elements, *vec.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/scale.py000066400000000000000000000054721503546127100220410ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.scale(self, factor) """ from __future__ import annotations import typing import numpy from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def rectify(lib, phi): return (phi + lib.pi) % (2 * lib.pi) - lib.pi def xy_z(lib, factor, x, y, z): return (x * factor, y * factor, z * factor) def xy_theta(lib, factor, x, y, theta): sign = lib.sign(factor) flip_if_negative = lib.absolute(theta + (0.5 * (sign - 1) * lib.pi)) return (x * factor, y * factor, flip_if_negative) def xy_eta(lib, factor, x, y, eta): return (x * factor, y * factor, eta * lib.sign(factor)) def rhophi_z(lib, factor, rho, phi, z): absfactor = lib.absolute(factor) sign = lib.sign(factor) turn_if_negative = -0.5 * (sign - 1) * lib.pi return (rho * absfactor, rectify(lib, phi + turn_if_negative), z * factor) def rhophi_theta(lib, factor, rho, phi, theta): absfactor = lib.absolute(factor) sign = lib.sign(factor) turn_if_negative = -0.5 * (sign - 1) * lib.pi flip_if_negative = lib.absolute(theta + (0.5 * (sign - 1) * lib.pi)) return (rho * absfactor, rectify(lib, phi + turn_if_negative), flip_if_negative) def rhophi_eta(lib, factor, rho, phi, eta): absfactor = lib.absolute(factor) sign = lib.sign(factor) turn_if_negative = -0.5 * (sign - 1) * lib.pi return (rho * absfactor, rectify(lib, phi + turn_if_negative), eta * sign) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalTheta): (xy_theta, AzimuthalXY, LongitudinalTheta), (AzimuthalXY, LongitudinalEta): (xy_eta, AzimuthalXY, LongitudinalEta), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, AzimuthalRhoPhi, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta, AzimuthalRhoPhi, LongitudinalTheta, ), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, AzimuthalRhoPhi, LongitudinalEta), } def dispatch(factor: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, factor, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/subtract.py000066400000000000000000000367301503546127100226020ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.subtract(self, angle) """ from __future__ import annotations import typing import numpy from vector._compute.planar import subtract, x, y from vector._compute.spatial import eta, theta, z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _handler_of, _lib_of, _ltype, ) # keep them in xy_z (default, anyway) def xy_z_xy_z(lib, x1, y1, z1, x2, y2, z2): return (x1 - x2, y1 - y2, z1 - z2) def xy_z_xy_theta(lib, x1, y1, z1, x2, y2, theta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_theta(lib, x2, y2, theta2)) def xy_z_xy_eta(lib, x1, y1, z1, x2, y2, eta2): return xy_z_xy_z(lib, x1, y1, z1, x2, y2, z.xy_eta(lib, x2, y2, eta2)) def xy_z_rhophi_z(lib, x1, y1, z1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2 ) def xy_z_rhophi_theta(lib, x1, y1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_z_rhophi_eta(lib, x1, y1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_theta_xy_z(lib, x1, y1, theta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z2) # keep them in xy_theta def xy_theta_xy_theta(lib, x1, y1, theta1, x2, y2, theta2): x = x1 - x2 y = y1 - y2 z1 = z.xy_theta(lib, x1, y1, theta1) z2 = z.xy_theta(lib, x2, y2, theta2) return (x, y, theta.xy_z(lib, x, y, z1 - z2)) def xy_theta_xy_eta(lib, x1, y1, theta1, x2, y2, eta2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x2, y2, z.xy_eta(lib, x2, y2, eta2), ) def xy_theta_rhophi_z(lib, x1, y1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_theta_rhophi_theta(lib, x1, y1, theta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_theta_rhophi_eta(lib, x1, y1, theta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z.xy_theta(lib, x1, y1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def xy_eta_xy_z(lib, x1, y1, eta1, x2, y2, z2): return xy_z_xy_z(lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z2) def xy_eta_xy_theta(lib, x1, y1, eta1, x2, y2, theta2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x2, y2, z.xy_theta(lib, x2, y2, theta2), ) # keep them in xy_eta def xy_eta_xy_eta(lib, x1, y1, eta1, x2, y2, eta2): x = x1 - x2 y = y1 - y2 z1 = z.xy_eta(lib, x1, y1, eta1) z2 = z.xy_eta(lib, x2, y2, eta2) return (x, y, eta.xy_z(lib, x, y, z1 - z2)) def xy_eta_rhophi_z(lib, x1, y1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def xy_eta_rhophi_theta(lib, x1, y1, eta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def xy_eta_rhophi_eta(lib, x1, y1, eta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x1, y1, z.xy_eta(lib, x1, y1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_z_xy_z(lib, rho1, phi1, z1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z2 ) def rhophi_z_xy_theta(lib, rho1, phi1, z1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_z_xy_eta(lib, rho1, phi1, z1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x2, y2, z.xy_eta(lib, x2, y2, eta2), ) # keep them in rhophi_z def rhophi_z_rhophi_z(lib, rho1, phi1, z1, rho2, phi2, z2): rho, phi = subtract.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) return (rho, phi, z1 - z2) def rhophi_z_rhophi_theta(lib, rho1, phi1, z1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) def rhophi_z_rhophi_eta(lib, rho1, phi1, z1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z1, x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_theta_xy_z(lib, rho1, phi1, theta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z2, ) def rhophi_theta_xy_theta(lib, rho1, phi1, theta1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_theta_xy_eta(lib, rho1, phi1, theta1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x2, y2, z.xy_eta(lib, x2, y2, eta2), ) def rhophi_theta_rhophi_z(lib, rho1, phi1, theta1, rho2, phi2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) # keep them in rhophi_theta def rhophi_theta_rhophi_theta(lib, rho1, phi1, theta1, rho2, phi2, theta2): rho, phi = subtract.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) z1 = z.rhophi_theta(lib, rho1, phi1, theta1) z2 = z.rhophi_theta(lib, rho2, phi2, theta2) return (rho, phi, theta.rhophi_z(lib, rho, phi, z1 - z2)) def rhophi_theta_rhophi_eta(lib, rho1, phi1, theta1, rho2, phi2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_theta(lib, rho1, phi1, theta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_eta(lib, rho2, phi2, eta2), ) def rhophi_eta_xy_z(lib, rho1, phi1, eta1, x2, y2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z2, ) def rhophi_eta_xy_theta(lib, rho1, phi1, eta1, x2, y2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z.xy_theta(lib, x2, y2, theta2), ) def rhophi_eta_xy_eta(lib, rho1, phi1, eta1, x2, y2, eta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x2, y2, z.xy_eta(lib, x2, y2, eta2), ) def rhophi_eta_rhophi_z(lib, rho1, phi1, eta1, rho2, phi2, z2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z2, ) def rhophi_eta_rhophi_theta(lib, rho1, phi1, eta1, rho2, phi2, theta2): return xy_z_xy_z( lib, x.rhophi(lib, rho1, phi1), y.rhophi(lib, rho1, phi1), z.rhophi_eta(lib, rho1, phi1, eta1), x.rhophi(lib, rho2, phi2), y.rhophi(lib, rho2, phi2), z.rhophi_theta(lib, rho2, phi2, theta2), ) # keep them in rhophi_eta def rhophi_eta_rhophi_eta(lib, rho1, phi1, eta1, rho2, phi2, eta2): rho, phi = subtract.rhophi_rhophi(lib, rho1, phi1, rho2, phi2) z1 = z.rhophi_eta(lib, rho1, phi1, eta1) z2 = z.rhophi_eta(lib, rho2, phi2, eta2) return (rho, phi, eta.rhophi_z(lib, rho, phi, z1 - z2)) dispatch_map = { (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( xy_z_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( xy_z_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( xy_z_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( xy_z_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( xy_z_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( xy_z_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( xy_theta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( xy_theta_xy_theta, AzimuthalXY, LongitudinalTheta, ), (AzimuthalXY, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( xy_theta_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( xy_theta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_theta_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( xy_theta_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( xy_eta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( xy_eta_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( xy_eta_xy_eta, AzimuthalXY, LongitudinalEta, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( xy_eta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( xy_eta_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalXY, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( xy_eta_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalZ): ( rhophi_z_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalTheta): ( rhophi_z_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalXY, LongitudinalEta): ( rhophi_z_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_z_rhophi_z, AzimuthalRhoPhi, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_z_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalZ, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_z_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalZ): ( rhophi_theta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalTheta): ( rhophi_theta_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalXY, LongitudinalEta): ( rhophi_theta_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_theta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta_rhophi_theta, AzimuthalRhoPhi, LongitudinalTheta, ), (AzimuthalRhoPhi, LongitudinalTheta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_theta_rhophi_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalZ): ( rhophi_eta_xy_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalTheta): ( rhophi_eta_xy_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalXY, LongitudinalEta): ( rhophi_eta_xy_eta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalZ): ( rhophi_eta_rhophi_z, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_eta_rhophi_theta, AzimuthalXY, LongitudinalZ, ), (AzimuthalRhoPhi, LongitudinalEta, AzimuthalRhoPhi, LongitudinalEta): ( rhophi_eta_rhophi_eta, AzimuthalRhoPhi, LongitudinalEta, ), } def dispatch(v1: typing.Any, v2: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v1), _ltype(v1), _aztype(v2), _ltype(v2), ), ) with numpy.errstate(all="ignore"): handler = _handler_of(v1, v2) return handler._wrap_result( _flavor_of(v1, v2), handler._wrap_dispatched_function(function)( _lib_of(v1, v2), *v1.azimuthal.elements, *v1.longitudinal.elements, *v2.azimuthal.elements, *v2.longitudinal.elements, ), returns, 2, ) vector-1.6.3/src/vector/_compute/spatial/theta.py000066400000000000000000000040131503546127100220450ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.theta(self) """ from __future__ import annotations import typing import numpy from vector._compute.spatial import costheta from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): return lib.arccos(costheta.xy_z(lib, x, y, z)) def xy_theta(lib, x, y, theta): return theta xy_theta.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_eta(lib, x, y, eta): return 2.0 * lib.arctan(lib.exp(-eta)) def rhophi_z(lib, rho, phi, z): return lib.arccos(costheta.rhophi_z(lib, rho, phi, z)) def rhophi_theta(lib, rho, phi, theta): return theta rhophi_theta.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_eta(lib, rho, phi, eta): return 2.0 * lib.arctan(lib.exp(-eta)) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/transform3D.py000066400000000000000000000067601503546127100231550ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.transform3D(self, obj) where ``obj` has ``obj["xx"]``, ``obj["xy"]``, etc. """ from __future__ import annotations import typing import numpy from vector._compute.planar import x, y from vector._compute.spatial import z from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) # Rotation is only computed in Cartesian coordinates; the rest are conversions. def cartesian(lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x, y, z): xp = xx * x + xy * y + xz * z yp = yx * x + yy * y + yz * z zp = zx * x + zy * y + zz * z return (xp, yp, zp) def xy_theta(lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x, y, theta): return cartesian( lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x, y, z.xy_theta(lib, x, y, theta) ) def xy_eta(lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x, y, eta): return cartesian( lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x, y, z.xy_eta(lib, x, y, eta) ) def rhophi_z(lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, rho, phi, z): return cartesian( lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z, ) def rhophi_theta(lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, rho, phi, theta): return cartesian( lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z.rhophi_theta(lib, rho, phi, theta), ) def rhophi_eta(lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, rho, phi, eta): return cartesian( lib, xx, xy, xz, yx, yy, yz, zx, zy, zz, x.rhophi(lib, rho, phi), y.rhophi(lib, rho, phi), z.rhophi_eta(lib, rho, phi, eta), ) dispatch_map = { (AzimuthalXY, LongitudinalZ): (cartesian, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalTheta): (xy_theta, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalEta): (xy_eta, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, AzimuthalXY, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, AzimuthalXY, LongitudinalZ), } def dispatch(obj: typing.Any, v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, obj["xx"], obj["xy"], obj["xz"], obj["yx"], obj["yy"], obj["yz"], obj["zx"], obj["zy"], obj["zz"], *v.azimuthal.elements, *v.longitudinal.elements, ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/unit.py000066400000000000000000000056461503546127100217340ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python Spatial.unit(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.spatial import mag from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): norm = mag.xy_z(lib, x, y, z) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(z / norm, nan=0, posinf=inf, neginf=-inf), ) def xy_theta(lib, x, y, theta): norm = mag.xy_theta(lib, x, y, theta) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), theta, ) def xy_eta(lib, x, y, eta): norm = mag.xy_eta(lib, x, y, eta) return ( lib.nan_to_num(x / norm, nan=0, posinf=inf, neginf=-inf), lib.nan_to_num(y / norm, nan=0, posinf=inf, neginf=-inf), eta, ) def rhophi_z(lib, rho, phi, z): norm = mag.rhophi_z(lib, rho, phi, z) return ( lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, lib.nan_to_num(z / norm, nan=0, posinf=inf, neginf=-inf), ) def rhophi_theta(lib, rho, phi, theta): norm = mag.rhophi_theta(lib, rho, phi, theta) return (lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, theta) def rhophi_eta(lib, rho, phi, eta): norm = mag.rhophi_eta(lib, rho, phi, eta) return (lib.nan_to_num(rho / norm, nan=0, posinf=inf, neginf=-inf), phi, eta) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, AzimuthalXY, LongitudinalZ), (AzimuthalXY, LongitudinalTheta): (xy_theta, AzimuthalXY, LongitudinalTheta), (AzimuthalXY, LongitudinalEta): (xy_eta, AzimuthalXY, LongitudinalEta), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, AzimuthalRhoPhi, LongitudinalZ), (AzimuthalRhoPhi, LongitudinalTheta): ( rhophi_theta, AzimuthalRhoPhi, LongitudinalTheta, ), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, AzimuthalRhoPhi, LongitudinalEta), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_compute/spatial/z.py000066400000000000000000000041161503546127100212150ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ .. code-block:: python @property Spatial.z(self) """ from __future__ import annotations import typing from math import inf import numpy from vector._compute.planar import rho from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, _aztype, _flavor_of, _from_signature, _ltype, ) def xy_z(lib, x, y, z): return z xy_z.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def xy_theta(lib, x, y, theta): return lib.nan_to_num( rho.xy(lib, x, y) / lib.tan(theta), nan=0.0, posinf=inf, neginf=-inf ) def xy_eta(lib, x, y, eta): return rho.xy(lib, x, y) * lib.sinh(eta) def rhophi_z(lib, rho, phi, z): return z rhophi_z.__awkward_transform_allowed__ = False # type:ignore[attr-defined] def rhophi_theta(lib, rho, phi, theta): return lib.nan_to_num(rho / lib.tan(theta), nan=0.0, posinf=inf, neginf=-inf) def rhophi_eta(lib, rho, phi, eta): return rho * lib.sinh(eta) dispatch_map = { (AzimuthalXY, LongitudinalZ): (xy_z, float), (AzimuthalXY, LongitudinalTheta): (xy_theta, float), (AzimuthalXY, LongitudinalEta): (xy_eta, float), (AzimuthalRhoPhi, LongitudinalZ): (rhophi_z, float), (AzimuthalRhoPhi, LongitudinalTheta): (rhophi_theta, float), (AzimuthalRhoPhi, LongitudinalEta): (rhophi_eta, float), } def dispatch(v: typing.Any) -> typing.Any: function, *returns = _from_signature( __name__, dispatch_map, ( _aztype(v), _ltype(v), ), ) with numpy.errstate(all="ignore"): return v._wrap_result( _flavor_of(v), v._wrap_dispatched_function(function)( v.lib, *v.azimuthal.elements, *v.longitudinal.elements ), returns, 1, ) vector-1.6.3/src/vector/_methods.py000066400000000000000000004251711503546127100173060ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import typing from contextlib import suppress import vector from vector._typeutils import ( BoolCollection, FloatArray, ScalarCollection, TransformProtocol2D, TransformProtocol3D, TransformProtocol4D, ) Module = typing.Any # returns a module, but we can't be specific about which one class Coordinates: pass class Azimuthal(Coordinates): @property def elements(self) -> tuple[ScalarCollection, ScalarCollection]: """ Azimuthal coordinates as a tuple. Each coordinate may be a scalar, a NumPy array, an Awkward Array, etc., but they are not vectors. """ raise AssertionError class Longitudinal(Coordinates): @property def elements(self) -> tuple[ScalarCollection]: """ Longitudinal coordinates as a tuple. Each coordinate may be a scalar, a NumPy array, an Awkward Array, etc., but they are not vectors. """ raise AssertionError class Temporal(Coordinates): @property def elements(self) -> tuple[ScalarCollection]: """ Temporal coordinates as a tuple. Each coordinate may be a scalar, a NumPy array, an Awkward Array, etc., but they are not vectors. """ raise AssertionError class AzimuthalXY(Azimuthal): """ Attributes: x (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $x$ coordinate(s). y (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $y$ coordinate(s). """ x: ScalarCollection y: ScalarCollection class AzimuthalRhoPhi(Azimuthal): r""" Attributes: rho (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\rho$ coordinate(s). phi (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\phi$ coordinate(s). """ rho: ScalarCollection phi: ScalarCollection class LongitudinalZ(Longitudinal): """ Attributes: z (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $z$ coordinate(s). """ z: ScalarCollection class LongitudinalTheta(Longitudinal): r""" Attributes: theta (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\theta$ coordinate(s). """ theta: ScalarCollection class LongitudinalEta(Longitudinal): r""" Attributes: eta (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\eta$ coordinate(s). """ eta: ScalarCollection class TemporalT(Temporal): """ Attributes: t (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $t$ coordinate(s). """ t: ScalarCollection class TemporalTau(Temporal): r""" Attributes: tau (scalar, ``np.ndarray``, ``ak.Array``, etc.): The $\tau$ coordinate(s). """ tau: ScalarCollection SameVectorType = typing.TypeVar("SameVectorType", bound="VectorProtocol") class VectorProtocol: """ Attributes: lib (module): The module used for functions used in compute functions (such as ``sqrt``, ``sin``, ``cos``). Usually ``numpy``. ProjectionClass2D (type): The class that would result from projecting this vector onto azimuthal coordinates only. ProjectionClass3D (type): The class that would result from projecting this vector onto azimuthal and longitudinal coordinates only. ProjectionClass4D (type): The class that would result from projecting this vector onto azimuthal, longitudinal, and temporal coordinates. GenericClass (type): The most generic concrete class for this type, for vectors without momentum-synonyms. MomentumClass (type): The momentum class for this type, for vectors with momentum-synonyms. """ @property def lib(self) -> Module: ... def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). Wraps the raw result of a compute function as a scalar, an array of scalars, a vector, or an array of vectors. """ raise AssertionError def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] raise AssertionError ProjectionClass2D: type[VectorProtocolPlanar] ProjectionClass3D: type[VectorProtocolSpatial] ProjectionClass4D: type[VectorProtocolLorentz] GenericClass: type[VectorProtocol] MomentumClass: type[VectorProtocol] def to_Vector2D(self) -> VectorProtocolPlanar: """Projects this vector/these vectors onto azimuthal coordinates only.""" raise AssertionError def to_Vector3D(self) -> VectorProtocolSpatial: """ Projects this vector/these vectors onto azimuthal and longitudinal coordinates only. If 2D, a default $z$ component of $0$ is imputed. The longitudinal coordinate can be passed as a named argument. """ raise AssertionError def to_Vector4D(self) -> VectorProtocolLorentz: """ Projects this vector/these vectors onto azimuthal, longitudinal, and temporal coordinates. If 3D, a default $t$ component of $0$ is imputed. If 2D, a $z$ component of $0$ is imputed along with a default $t$ component of $0$. The longitudinal and temporal coordinates can be passed as named arguments. """ raise AssertionError def to_2D(self) -> VectorProtocolPlanar: """ Projects this vector/these vectors onto azimuthal coordinates only. Alias for :meth:`vector._methods.VectorProtocol.to_Vector2D`. """ raise AssertionError def to_3D(self) -> VectorProtocolSpatial: """ Projects this vector/these vectors onto azimuthal and longitudinal coordinates only. If 2D, a default $z$ component of $0$ is imputed. The longitudinal coordinate can be passed as a named argument. Alias for :meth:`vector._methods.VectorProtocol.to_Vector3D`. """ raise AssertionError def to_4D(self) -> VectorProtocolLorentz: """ Projects this vector/these vectors onto azimuthal, longitudinal, and temporal coordinates. If 3D, a default $t$ component of $0$ is imputed. If 2D, a $z$ component of $0$ is imputed along with a default $t$ component of $0$. The longitudinal and temporal coordinates can be passed as named arguments. Alias for :meth:`vector._methods.VectorProtocol.to_Vector4D`. """ raise AssertionError def to_xy(self) -> VectorProtocolPlanar: """ Converts to $x$-$y$ coordinates, possibly eliminating dimensions with a projection. """ raise AssertionError def to_pxpy(self) -> VectorProtocolPlanar: """ Converts to $px$-$py$ coordinates, possibly eliminating dimensions with a projection. """ raise AssertionError def to_rhophi(self) -> VectorProtocolPlanar: r""" Converts to $\rho$-$\phi$ coordinates, possibly eliminating dimensions with a projection. """ raise AssertionError def to_ptphi(self) -> VectorProtocolPlanar: r""" Converts to $pt$-$\phi$ coordinates, possibly eliminating dimensions with a projection. """ raise AssertionError def to_xyz(self) -> VectorProtocolSpatial: """ Converts to $x$-$y$-$z$ coordinates, possibly eliminating or imputing dimensions with a projection. The $z$ coordinate can be passed as a named argument. """ raise AssertionError def to_xytheta(self) -> VectorProtocolSpatial: r""" Converts to $x$-$y$-$\theta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $theta$ coordinate can be passed as a named argument. """ raise AssertionError def to_xyeta(self) -> VectorProtocolSpatial: r""" Converts to $x$-$y$-$\eta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $eta$ coordinate can be passed as a named argument. """ raise AssertionError def to_pxpypz(self) -> VectorProtocolSpatial: """ Converts to $px$-$py$-$pz$ coordinates, possibly eliminating or imputing dimensions with a projection. The $pz$ coordinate can be passed as a named argument. """ raise AssertionError def to_pxpytheta(self) -> VectorProtocolSpatial: r""" Converts to $px$-$py$-$\theta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $theta$ coordinate can be passed as a named argument. """ raise AssertionError def to_pxpyeta(self) -> VectorProtocolSpatial: r""" Converts to $px$-$py$-$\eta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $eta$ coordinate can be passed as a named argument. """ raise AssertionError def to_rhophiz(self) -> VectorProtocolSpatial: r""" Converts to $\rho$-$\phi$-$z$ coordinates, possibly eliminating or imputing dimensions with a projection. The $z$ coordinate can be passed as a named argument. """ raise AssertionError def to_rhophitheta(self) -> VectorProtocolSpatial: r""" Converts to $\rho$-$\phi$-$\theta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $theta$ coordinate can be passed as a named argument. """ raise AssertionError def to_rhophieta(self) -> VectorProtocolSpatial: r""" Converts to $\rho$-$\phi$-$\eta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $eta$ coordinate can be passed as a named argument. """ raise AssertionError def to_ptphipz(self) -> VectorProtocolSpatial: r""" Converts to $pt$-$\phi$-$pz$ coordinates, possibly eliminating or imputing dimensions with a projection. The $pz$ coordinate can be passed as a named argument. """ raise AssertionError def to_ptphitheta(self) -> VectorProtocolSpatial: r""" Converts to $pt$-$\phi$-$\theta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $theta$ coordinate can be passed as a named argument. """ raise AssertionError def to_ptphieta(self) -> VectorProtocolSpatial: r""" Converts to $pt$-$\phi$-$\eta$ coordinates, possibly eliminating or imputing dimensions with a projection. The $eta$ coordinate can be passed as a named argument. """ raise AssertionError def to_xyzt(self) -> VectorProtocolLorentz: """ Converts to $x$-$y$-$z$-$t$ coordinates, possibly imputing dimensions with a projection. The $z$ and $t$ coordinates can be passed as a named argument. """ raise AssertionError def to_xyztau(self) -> VectorProtocolLorentz: r""" Converts to $x$-$y$-$z$-$\tau$ coordinates, possibly imputing dimensions with a projection. The $z$ and $tau$ coordinates can be passed as a named argument. """ raise AssertionError def to_xythetat(self) -> VectorProtocolLorentz: r""" Converts to $x$-$y$-$\theta$-$t$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $t$ coordinates can be passed as a named argument. """ raise AssertionError def to_xythetatau(self) -> VectorProtocolLorentz: r""" Converts to $x$-$y$-$\theta$-$\tau$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $tau$ coordinates can be passed as a named argument. """ raise AssertionError def to_xyetat(self) -> VectorProtocolLorentz: r""" Converts to $x$-$y$-$\eta$-$t$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $t$ coordinates can be passed as a named argument. """ raise AssertionError def to_xyetatau(self) -> VectorProtocolLorentz: r""" Converts to $x$-$y$-$\eta$-$\tau$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $tau$ coordinates can be passed as a named argument. """ raise AssertionError def to_pxpypzenergy(self) -> VectorProtocolLorentz: r""" Converts to $px$-$py$-$pz$-$energy$ coordinates, possibly imputing dimensions with a projection. The $pz$ and $energy$ coordinates can be passed as a named argument. """ raise AssertionError def to_pxpythetaenergy(self) -> VectorProtocolLorentz: r""" Converts to $px$-$py$-$\theta$-$energy$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $energy$ coordinates can be passed as a named argument. """ raise AssertionError def to_pxpyetaenergy(self) -> VectorProtocolLorentz: r""" Converts to $px$-$py$-$\eta$-$energy$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $energy$ coordinates can be passed as a named argument. """ raise AssertionError def to_pxpypzmass(self) -> VectorProtocolLorentz: r""" Converts to $px$-$py$-$pz$-$mass$ coordinates, possibly imputing dimensions with a projection. The $pz$ and $mass$ coordinates can be passed as a named argument. """ raise AssertionError def to_pxpythetamass(self) -> VectorProtocolLorentz: r""" Converts to $px$-$py$-$\theta$-$energy$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $mass$ coordinates can be passed as a named argument. """ raise AssertionError def to_pxpyetamass(self) -> VectorProtocolLorentz: r""" Converts to $px$-$py$-$\eta$-$mass$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $mass$ coordinates can be passed as a named argument. """ raise AssertionError def to_rhophizt(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$z$-$t$ coordinates, possibly imputing dimensions with a projection. The $z$ and $t$ coordinates can be passed as a named argument. """ raise AssertionError def to_rhophiztau(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$z$-$\tau$ coordinates, possibly imputing dimensions with a projection. The $z$ and $tau$ coordinates can be passed as a named argument. """ raise AssertionError def to_rhophithetat(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$\theta$-$t$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $t$ coordinates can be passed as a named argument. """ raise AssertionError def to_rhophithetatau(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$\theta$-$\tau$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $tau$ coordinates can be passed as a named argument. """ raise AssertionError def to_rhophietat(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$\eta$-$t$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $t$ coordinates can be passed as a named argument. """ raise AssertionError def to_rhophietatau(self) -> VectorProtocolLorentz: r""" Converts to $\rho$-$\phi$-$\eta$-$\tau$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $tau$ coordinates can be passed as a named argument. """ raise AssertionError def to_ptphipzenergy(self) -> VectorProtocolLorentz: r""" Converts to $pt$-$\phi$-$pz$-$energy$ coordinates, possibly imputing dimensions with a projection. The $pz$ and $energy$ coordinates can be passed as a named argument. """ raise AssertionError def to_ptphithetaenergy(self) -> VectorProtocolLorentz: r""" Converts to $pt$-$\phi$-$\theta$-$energy$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $energy$ coordinates can be passed as a named argument. """ raise AssertionError def to_ptphietaenergy(self) -> VectorProtocolLorentz: r""" Converts to $pt$-$\phi$-$\eta$-$energy$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $energy$ coordinates can be passed as a named argument. """ raise AssertionError def to_ptphipzmass(self) -> VectorProtocolLorentz: r""" Converts to $pt$-$\phi$-$pz$-$mass$ coordinates, possibly imputing dimensions with a projection. The $pz$ and $mass$ coordinates can be passed as a named argument. """ raise AssertionError def to_ptphithetamass(self) -> VectorProtocolLorentz: r""" Converts to $pt$-$\phi$-$\theta$-$mass$ coordinates, possibly imputing dimensions with a projection. The $theta$ and $mass$ coordinates can be passed as a named argument. """ raise AssertionError def to_ptphietamass(self) -> VectorProtocolLorentz: r""" Converts to $pt$-$\phi$-$\theta$-$mass$ coordinates, possibly imputing dimensions with a projection. The $eta$ and $mass$ coordinates can be passed as a named argument. """ raise AssertionError def unit(self: SameVectorType) -> SameVectorType: """ Returns vector(s) normalized to unit length, which is `rho == 1` for 2D vectors, `mag == 1` for 3D vectors, and `tau == 1` for 4D vectors. """ raise AssertionError def dot(self, other: VectorProtocol) -> ScalarCollection: """ Vector dot product of ``self`` with ``other``. This method is equivalent to the ``@`` operator. """ raise AssertionError def add(self, other: VectorProtocol) -> VectorProtocol: """ Sum of ``self`` and ``other``. This method is equivalent to the ``+`` operator. """ raise AssertionError def subtract(self, other: VectorProtocol) -> VectorProtocol: """ Difference of ``self`` minus ``other``. This method is equivalent to the ``-`` operator. """ raise AssertionError def scale(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: """ Returns vector(s) scaled by a ``factor``, changing the length(s) but not the direction(s). This method is equivalent to the ``*`` operator. """ raise AssertionError def equal(self, other: VectorProtocol) -> BoolCollection: """ Returns True if ``self`` is exactly equal to ``other`` (possibly for arrays of vectors), False otherwise. This method is equivalent to the ``==`` operator. Typically, you'll want to check :meth:`vector._methods.VectorProtocol.isclose` to allow for numerical errors. """ raise AssertionError def not_equal(self, other: VectorProtocol) -> BoolCollection: """ Returns False if ``self`` is exactly equal to ``other`` (possibly for arrays of vectors), True otherwise. This method is equivalent to the ``!=`` operator. Typically, you'll want to check :meth:`vector._methods.VectorProtocol.isclose` to allow for numerical errors. """ raise AssertionError def isclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: """ Returns True if ``self`` is approximately equal to ``other`` (possibly for arrays of vectors), False otherwise. The relative tolerance (``rtol``) and absolute tolerance (``atol``) are interpreted as in ``np.isclose``: .. code-block:: python close_enough = abs(self - other) <= atol + rtol * abs(other) """ raise AssertionError def like(self, other: VectorProtocol) -> VectorProtocol: """ Projects the vector into the geometric coordinates of the `other` vector. Value(s) of $0$ is/are imputed while transforming vector from a lower geometric dimension to a higher geometric dimension. .. code-block:: python vec_4d + vec_3d.like(vec_4d) For more flexibility (passing new coordinate values), see :meth:`vector._methods.Vector2D.to_Vector3D`, :meth:`vector._methods.Vector2D.to_Vector4D`, and :meth:`vector._methods.Vector3D.to_Vector4D`, which can be used as: .. code-block:: python vec_2d.to_Vector3D(z=3.0) vec_2d.to_Vector4D(z=3.0, t=4.0) vec_3d.to_Vector4D(t=4.0) """ raise AssertionError class VectorProtocolPlanar(VectorProtocol): @property def azimuthal(self) -> Azimuthal: """ Container of azimuthal coordinates, for use in dispatching to compute functions or to identify coordinate system with ``isinstance``. """ raise AssertionError @property def x(self) -> ScalarCollection: """The Cartesian $x$ coordinate of the vector or every vector in the array.""" raise AssertionError @property def y(self) -> ScalarCollection: """The Cartesian $y$ coordinate of the vector or every vector in the array.""" raise AssertionError @property def rho(self) -> ScalarCollection: r""" The polar $\rho$ coordinate of the vector or every vector in the array. This is also the magnitude of the 2D azimuthal part of the vector (not including any longitudinal or temporal parts). """ raise AssertionError @property def rho2(self) -> ScalarCollection: r"""The polar $\rho$ coordinate squared of the vector or every vector in the array.""" raise AssertionError @property def phi(self) -> ScalarCollection: r""" The polar $\phi$ coordinate of the vector or every vector in the array (in radians, always between $-\pi$ and $\pi$). """ raise AssertionError def scale2D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: """ Returns vector(s) with the 2D part scaled by a ``factor``, not affecting any longitudinal or temporal parts. """ raise AssertionError @property def neg2D(self: SameVectorType) -> SameVectorType: """ Returns vector(s) with the 2D part negated, not affecting any longitudinal or temporal parts. """ raise AssertionError def deltaphi(self, other: VectorProtocol) -> ScalarCollection: r"""Signed difference in $\phi$ of ``self`` minus ``other`` (in radians).""" raise AssertionError def rotateZ(self: SameVectorType, angle: ScalarCollection) -> SameVectorType: """ Rotates the vector(s) by a given ``angle`` (in radians) around the longitudinal axis. Note that the ``angle`` can be an array with the same length as the vectors, if the vectors are in an array. """ raise AssertionError def transform2D(self: SameVectorType, obj: TransformProtocol2D) -> SameVectorType: """ Arbitrarily transforms the vector(s) by .. code-block:: python obj["xx"] obj["xy"] obj["yx"] obj["yy"] leaving any longitudinal or temporal coordinates unchanged. There is no restriction on the type of ``obj``; it just has to provide those components (which can be arrays if the vectors are in an array). """ raise AssertionError def is_parallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: r""" Returns True if ``self`` and ``other`` are pointing in the same direction (i.e. not "antiparallel"; dot product is nearly ``abs(self) * abs(other)``). The ``tolerance`` is measured in units of $\cos(\Delta\alpha)$ where $\Delta\alpha$ is ``self.deltaangle(other)``. """ raise AssertionError def is_antiparallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: r""" Returns True if ``self`` and ``other`` are pointing in opposite directions (i.e. dot product is nearly ``-abs(self) * abs(other)``). The ``tolerance`` is measured in units of $\cos(\Delta\alpha)$ where $\Delta\alpha$ is ``self.deltaangle(other)``. """ raise AssertionError def is_perpendicular( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: r""" Returns True if ``self`` and ``other`` are pointing in perpendicular directions (i.e. dot product is nearly ``0``). The ``tolerance`` is measured in units of $\cos(\Delta\alpha)$ where $\Delta\alpha$ is ``self.deltaangle(other)``. """ raise AssertionError class VectorProtocolSpatial(VectorProtocolPlanar): @property def longitudinal(self) -> Longitudinal: """ Container of longitudinal coordinates, for use in dispatching to compute functions or to identify coordinate system with ``isinstance``. """ raise AssertionError @property def z(self) -> ScalarCollection: """The Cartesian $z$ coordinate of the vector or every vector in the array.""" raise AssertionError @property def theta(self) -> ScalarCollection: r""" The spherical $\theta$ coordinate (polar angle) of the vector or every vector in the array (in radians, always between $0$ ($+z$) and $\pi$ ($-z$)). """ raise AssertionError @property def eta(self) -> ScalarCollection: r""" The pseudorapidity $\eta$ coordinate of the vector or every vector in the array (in radians, always between $0$ ($+z$) and $\pi$ ($-z$)). """ raise AssertionError @property def costheta(self) -> ScalarCollection: r""" The $\cos\theta$ coordinate of the vector or every vector in the array (has the same sign as $z$). """ raise AssertionError @property def cottheta(self) -> ScalarCollection: r""" The $\cot\theta$ coordinate of the vector or every vector in the array (has the same sign as $z$). """ raise AssertionError @property def mag(self) -> ScalarCollection: """The magnitude of the vector(s) in 3D (not including any temporal parts).""" raise AssertionError @property def mag2(self) -> ScalarCollection: """The magnitude-squared of the vector(s) in 3D (not including any temporal parts).""" raise AssertionError def scale3D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: """ Returns vector(s) with the 3D part scaled by a ``factor``, not affecting any longitudinal or temporal parts. """ raise AssertionError @property def neg3D(self: SameVectorType) -> SameVectorType: """ Returns vector(s) with the 3D part negated, not affecting any longitudinal or temporal parts. """ raise AssertionError def cross(self, other: VectorProtocolSpatial) -> VectorProtocolSpatial: """ The 3D cross-product of ``self`` with ``other``. Even if ``self`` or ``other`` is 4D, the resulting vector(s) is/are 3D. """ raise AssertionError def deltaangle( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: r""" Angle in 3D space between ``self`` and ``other``, which is always positive, between $0$ and $\pi$. """ raise AssertionError def deltaeta( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: r"""Signed difference in $\eta$ of ``self`` minus ``other``.""" raise AssertionError def deltaR( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: r""" Sum in quadrature of :meth:`vector._methods.VectorProtocolPlanar.deltaphi` and :meth:`vector._methods.VectorProtocolSpatial.deltaeta`: $$\Delta R = \sqrt{\Delta\phi^2 + \Delta\eta^2}$$ """ raise AssertionError def deltaR2( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: r""" Square of the sum in quadrature of :meth:`vector._methods.VectorProtocolPlanar.deltaphi` and :meth:`vector._methods.VectorProtocolSpatial.deltaeta`: $$\Delta R^2 = \Delta\phi^2 + \Delta\eta^2$$ """ raise AssertionError def rotateX(self: SameVectorType, angle: ScalarCollection) -> SameVectorType: """ Rotates the vector(s) by a given ``angle`` (in radians) around the $x$ axis. Note that the ``angle`` can be an array with the same length as the vectors, if the vectors are in an array. """ raise AssertionError def rotateY(self: SameVectorType, angle: ScalarCollection) -> SameVectorType: """ Rotates the vector(s) by a given ``angle`` (in radians) around the $y$ axis. Note that the ``angle`` can be an array with the same length as the vectors, if the vectors are in an array. """ raise AssertionError def rotate_axis( self: SameVectorType, axis: VectorProtocolSpatial, angle: ScalarCollection ) -> SameVectorType: """ Rotates the vector(s) by a given ``angle`` (in radians) around the axis indicated by another vector, ``axis``. The magnitude of ``axis`` is ignored. Note that the ``axis`` and ``angle`` can be arrays with the same length as the vectors, if the vectors are in an array. """ raise AssertionError def rotate_euler( self: SameVectorType, phi: ScalarCollection, theta: ScalarCollection, psi: ScalarCollection, order: str = "zxz", ) -> SameVectorType: """ Rotates the vector(s) by three given angles: ``phi``, ``theta``, and ``psi`` (in radians). The ``order`` string determines which axis each rotation is applied around: - ``"zxz"``, ``"xyx"``, ``"yzy"``, ``"zyz"``, ``"xzx"``, and ``"yxy"`` are proper Euler angles - ``"zxz"``, ``"xyx"``, ``"yzy"``, ``"zyz"``, ``"xzx"``, and ``"yxy"`` are Tait-Bryan angles (see :meth:`vector._methods.VectorProtocolSpatial.rotate_nautical`) The names ``phi``, ``theta``, and ``psi`` agree with `Wikipedia's terminology `_, and both the names and order agree with `ROOT's Math::EulerAngles `_. The default ``order = "zxz"`` is also ROOT's convention. Note that the angles can be arrays with the same lengths as the vectors, if the vectors are in an array. """ raise AssertionError def rotate_nautical( self: SameVectorType, yaw: ScalarCollection, pitch: ScalarCollection, roll: ScalarCollection, ) -> SameVectorType: """ Rotates the vector(s) by three given angles: ``yaw``, ``pitch``, and ``roll`` (in radians). These are Tait-Bryan angles often used for boats and planes (see `this lesson `__ and `this lesson `__). This function is entirely equivalent to .. code-block:: python rotate_euler(roll, pitch, yaw, order="zyx") Note that the angles can be arrays with the same lengths as the vectors, if the vectors are in an array. """ raise AssertionError def rotate_quaternion( self: SameVectorType, u: ScalarCollection, i: ScalarCollection, j: ScalarCollection, k: ScalarCollection, ) -> SameVectorType: """ Rotates the vector(s) by four angles as quaternion coefficients (in radians). Four angles are sometimes preferred over three because the latter has a pathology known as "gimbal lock." This function follows the same conventions as `ROOT's Math::Quaternion `_. Note that the angles can be arrays with the same lengths as the vectors, if the vectors are in an array. """ raise AssertionError def transform3D(self: SameVectorType, obj: TransformProtocol3D) -> SameVectorType: """ Arbitrarily transforms the vector(s) by .. code-block:: python obj["xx"] obj["xy"] obj["xz"] obj["yx"] obj["yy"] obj["yz"] obj["zx"] obj["zy"] obj["zz"] leaving any temporal coordinate unchanged. There is no restriction on the type of ``obj``; it just has to provide those components (which can be arrays if the vectors are in an array). """ raise AssertionError def is_parallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: raise AssertionError def is_antiparallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: raise AssertionError def is_perpendicular( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: raise AssertionError class VectorProtocolLorentz(VectorProtocolSpatial): @property def temporal(self) -> Temporal: """ Container of temporal coordinates, for use in dispatching to compute functions or to identify coordinate system with ``isinstance``. """ raise AssertionError @property def t(self) -> ScalarCollection: r""" The Cartesian $t$ (time) coordinate of the vector or every vector in the array. If $t$ is derived from $\tau$, it is not allowed to be ``NaN``. .. code-block:: python t = sqrt(max(copysign(tau**2, tau) + mag**2, 0)) """ raise AssertionError @property def t2(self) -> ScalarCollection: r""" The Cartesian $t$ (time) coordinate squared of the vector or every vector in the array. If $t^2$ is derived from $\tau$, it is not allowed to be negative. .. code-block:: python t2 = max(copysign(tau**2, tau) + mag**2, 0) """ raise AssertionError @property def tau(self) -> ScalarCollection: r""" The Lorentz magnitude $\tau$ (proper time) of the vector or every vector in the array. If $\tau$ is derived from $t$, spacelike vectors are represented by negative proper times. .. code-block:: python tau = copysign(sqrt(abs(t**2 - mag**2)), t**2 - mag**2) """ raise AssertionError @property def tau2(self) -> ScalarCollection: r""" The Lorentz magnitude $\tau$ (proper time) squared of the vector or every vector in the array. .. code-block:: python tau2 = t**2 - mag**2 """ raise AssertionError @property def beta(self) -> ScalarCollection: """ The speed(s) of the Lorentz vector or array of vectors, in which lightlike vectors have ``beta == 1``. """ raise AssertionError @property def gamma(self) -> ScalarCollection: r""" The time dilation/length contraction factor(s) of the Lorentz vector or array of vectors: $t/\tau$. """ raise AssertionError @property def rapidity(self) -> ScalarCollection: """ The rapidity relative to the longitudinal axis of the Lorentz vector or array of vectors. .. code-block:: python 0.5 * log((t + z) / (t - z)) """ raise AssertionError def deltaRapidityPhi(self, other: VectorProtocolLorentz) -> ScalarCollection: r""" Sum in quadrature of :meth:`vector._methods.VectorProtocolPlanar.deltaphi` and the difference in :attr:`vector._methods.VectorProtocolLorentz.rapidity` of the two vectors: $$\Delta R_{\mbox{rapidity}} = \sqrt{\Delta\phi^2 + \Delta \mbox{rapidity}^2}$$ """ raise AssertionError def deltaRapidityPhi2(self, other: VectorProtocolLorentz) -> ScalarCollection: r""" Square of the sum in quadrature of :meth:`vector._methods.VectorProtocolPlanar.deltaphi` and the difference in :attr:`vector._methods.VectorProtocolLorentz.rapidity` of the two vectors: $$\Delta R_{\mbox{rapidity}} = \Delta\phi^2 + \Delta \mbox{rapidity}^2$$ """ raise AssertionError def scale4D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: """Same as ``scale``.""" raise AssertionError @property def neg4D(self: SameVectorType) -> SameVectorType: """Same as multiplying by -1.""" raise AssertionError def boost_p4(self: SameVectorType, p4: VectorProtocolLorentz) -> SameVectorType: """ Boosts the vector or array of vectors in a direction and magnitude given by the 4D vector or array of vectors ``p4``. This function is equivalent to but more numerically stable than .. code-block:: python boost_beta3(p4.to_beta3()) where :meth:`vector._methods.VectorProtocolLorentz.to_beta3` converts a 4D Lorentz vector into a 3D velocity (in which lightlike velocities have ``mag == 1``). Note that ``v.boost_p4(v)`` does not boost into the center-of-mass (CM) frame of ``v``; it boosts *away* from its CM frame. Neither does ``v.boost_p4(-v)``, since that negates the time component of ``v`` as well. To boost to the center-of-mass frame of a vector ``v``, use :meth:`vector._methods.VectorProtocolLorentz.boostCM_of_p4`. For instance, ``v.boostCM_of_p4(v)`` is guaranteed to have spatial components close to zero and a temporal component close to ``v.tau``. """ raise AssertionError def boost_beta3( self: SameVectorType, beta3: VectorProtocolSpatial ) -> SameVectorType: """ Boosts the vector or array of vectors in a direction and magnitude given by the 3D velocity or array of velocity vectors ``beta3``. Note that ``v.boost_beta3(v.to_beta3())`` does not boost into the center-of-mass (CM) frame of ``v``; it boosts *away* from its CM frame. Neither does ``v.boost_beta3((-v).to_beta3())``, since that negates the time component of ``v`` as well. On the other hand, ``v.boost_beta3(-(v.to_beta3()))`` *would* boost to the center-of-mass frame. However, there's a function for that: :meth:`vector._methods.VectorProtocolLorentz.boostCM_of_beta3` is explicit about boosting to a center-of-mass (CM) frame and it handles the negative sign for you: ``v.boostCM_of_beta3(v.to_beta3())`` is guaranteed to have spatial components close to zero and a temporal component close to ``v.tau``. """ raise AssertionError def boost( self: SameVectorType, booster: VectorProtocolSpatial | VectorProtocolLorentz ) -> SameVectorType: """ Boosts the vector or array of vectors using the 3D or 4D ``booster``. If ``booster`` is 3D, it is interpreted as a velocity (in which lightlike velocities have ``mag == 1``) and :meth:`vector._methods.VectorProtocolLorentz.boost_beta3` is called. If ``booster`` is 4D, it is interpreted as a Lorentz vector and :meth:`vector._methods.VectorProtocolLorentz.boost_p4` is called. Note that ``v.boost(v)`` does not boost into the center-of-mass (CM) frame of ``v``; it boosts *away* from its CM frame. Neither does ``v.boost(-v)``, since that negates the time component of ``v`` as well. To boost to the center-of-mass frame of a vector ``v``, use :meth:`vector._methods.VectorProtocolLorentz.boostCM_of`. For instance, ``v.boostCM_of(v)`` is guaranteed to have spatial components close to zero and a temporal component close to ``v.tau``. """ raise AssertionError def boostCM_of_p4( self: SameVectorType, p4: VectorProtocolLorentz ) -> SameVectorType: """ Boosts the vector or array of vectors to the center-of-mass (CM) frame of the 4D vector or array of vectors ``p4``. This function is equivalent to but more numerically stable than .. code-block:: python boostCM_of_beta3(p4.to_beta3()) Note that ``v.boostCM_of_p4(v)`` is guaranteed to have spatial components close to zero and a temporal component close to ``v.tau``. """ raise AssertionError def boostCM_of_beta3( self: SameVectorType, beta3: VectorProtocolSpatial ) -> SameVectorType: """ Boosts the vector or array of vectors to the center-of-mass (CM) frame of the 3D velocity or array of velocity vectors ``beta3``. Note that ``v.boostCM_of_beta3(v.to_beta3())`` is guaranteed to have spatial components close to zero and a temporal component close to ``v.tau``. """ raise AssertionError def boostCM_of( self: SameVectorType, booster: VectorProtocolSpatial | VectorProtocolLorentz ) -> SameVectorType: """ Boosts the vector or array of vectors to the center-of-mass (CM) frame of the 3D or 4D ``booster``. If ``booster`` is 3D, it is interpreted as a velocity (in which lightlike velocities have ``mag == 1``) and :meth:`vector._methods.VectorProtocolLorentz.boostCM_of_beta3` is called. If ``booster`` is 4D, it is interpreted as a Lorentz vector and :meth:`vector._methods.VectorProtocolLorentz.boostCM_of_p4` is called. Note that ``v.boostCM_of(v)`` is guaranteed to have spatial components close to zero and a temporal component close to ``v.tau``. """ raise AssertionError def boostX( self: SameVectorType, beta: ScalarCollection | None = None, gamma: ScalarCollection | None = None, ) -> SameVectorType: """ Boosts the vector or array of vectors in the $x$ direction by a speed ``beta`` (in which lightlike boosts have ``beta == 1``) or time dilation/length contraction factor ``gamma``. Either ``beta`` xor ``gamma`` must be specified, not both or neither. If ``beta`` or ``gamma`` is negative, it is taken as a boost in the $-x$ direction. """ raise AssertionError def boostY( self: SameVectorType, beta: ScalarCollection | None = None, gamma: ScalarCollection | None = None, ) -> SameVectorType: """ Boosts the vector or array of vectors in the $y$ direction by a speed ``beta`` (in which lightlike boosts have ``beta == 1``) or time dilation/length contraction factor ``gamma``. Either ``beta`` xor ``gamma`` must be specified, not both or neither. If ``beta`` or ``gamma`` is negative, it is taken as a boost in the $-y$ direction. """ raise AssertionError def boostZ( self: SameVectorType, beta: ScalarCollection | None = None, gamma: ScalarCollection | None = None, ) -> SameVectorType: """ Boosts the vector or array of vectors in the $z$ direction by a speed ``beta`` (in which lightlike boosts have ``beta == 1``) or time dilation/length contraction factor ``gamma``. Either ``beta`` xor ``gamma`` must be specified, not both or neither. If ``beta`` or ``gamma`` is negative, it is taken as a boost in the $-z$ direction. """ raise AssertionError def transform4D(self: SameVectorType, obj: TransformProtocol4D) -> SameVectorType: """ Arbitrarily transforms the vector(s) by .. code-block:: python obj["xx"] obj["xy"] obj["xz"] obj["xt"] obj["yx"] obj["yy"] obj["yz"] obj["yt"] obj["zx"] obj["zy"] obj["zz"] obj["zt"] obj["tx"] obj["ty"] obj["tz"] obj["tt"] There is no restriction on the type of ``obj``; it just has to provide those components (which can be arrays if the vectors are in an array). """ raise AssertionError def to_beta3(self) -> VectorProtocolSpatial: """ Converts the 4D Lorentz vector or array of vectors into a 3D velocity vector or array of vectors, in which lightlike velocities have ``mag == 1``. """ raise AssertionError def is_timelike(self, tolerance: ScalarCollection = 0) -> BoolCollection: """ Returns True if the vector or a vector in the array is pointing in a timelike direction, ``t**2 > mag**2``, False otherwise. The ``tolerance`` is in units of ``t`` and ``mag``. Note that - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_timelike` is ``0`` - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_spacelike` is ``0`` - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_lightlike` is ``1e-5`` If you want to use these methods to divide space-time into non-overlapping regions (the light-cone), use the same ``tolerance`` for each. """ raise AssertionError def is_spacelike(self, tolerance: ScalarCollection = 0) -> BoolCollection: """ Returns True if the vector or a vector in the array is pointing in a spacelike direction, ``t**2 < mag**2``, False otherwise. The ``tolerance`` is in units of ``t`` and ``mag``. Note that - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_timelike` is ``0`` - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_spacelike` is ``0`` - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_lightlike` is ``1e-5`` If you want to use these methods to divide space-time into non-overlapping regions (the light-cone), use the same ``tolerance`` for each. """ raise AssertionError def is_lightlike(self, tolerance: ScalarCollection = 1e-5) -> BoolCollection: """ Returns True if the vector or a vector in the array is pointing in a lightlike direction, ``t**2 == mag**2``, False otherwise. The ``tolerance`` is in units of ``t`` and ``mag``. Note that - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_timelike` is ``0`` - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_spacelike` is ``0`` - the default ``tolerance`` for :meth:`vector._methods.VectorProtocolLorentz.is_lightlike` is ``1e-5`` If you want to use these methods to divide space-time into non-overlapping regions (the light-cone), use the same ``tolerance`` for each. """ raise AssertionError class MomentumProtocolPlanar(VectorProtocolPlanar): @property def px(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolPlanar.x`.""" raise AssertionError @property def py(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolPlanar.y`.""" raise AssertionError @property def pt(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolPlanar.rho`.""" raise AssertionError @property def pt2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolPlanar.rho2`.""" raise AssertionError class MomentumProtocolSpatial(VectorProtocolSpatial, MomentumProtocolPlanar): @property def pz(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolSpatial.z`.""" raise AssertionError @property def pseudorapidity(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolSpatial.eta`.""" raise AssertionError @property def p(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolSpatial.mag`.""" raise AssertionError @property def p2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolSpatial.mag2`.""" raise AssertionError class MomentumProtocolLorentz(VectorProtocolLorentz, MomentumProtocolSpatial): @property def E(self) -> ScalarCollection: """Momentum-synonyor :attr:`vector._methods.VectorProtocolLorentz.t`.""" raise AssertionError @property def e(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.t`.""" raise AssertionError @property def energy(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.t`.""" raise AssertionError @property def E2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorent2`.""" raise AssertionError @property def e2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.t2`.""" raise AssertionError @property def energy2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.t2`.""" raise AssertionError @property def M(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.tau`.""" raise AssertionError @property def m(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.tau`.""" raise AssertionError @property def mass(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.tau`.""" raise AssertionError @property def M2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.tau2`.""" raise AssertionError @property def m2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.tau2`.""" raise AssertionError @property def mass2(self) -> ScalarCollection: """Momentum-synonym for :attr:`vector._methods.VectorProtocolLorentz.tau2`.""" raise AssertionError @property def Et(self) -> ScalarCollection: r""" Transverse energy of the four-momentum vector or array of vectors: $E_T = E \sin\theta$. """ raise AssertionError @property def et(self) -> ScalarCollection: r""" Transverse energy of the four-momentum vector or array of vectors: $E_T = E \sin\theta$. """ raise AssertionError @property def transverse_energy(self) -> ScalarCollection: """Synonym for :attr:`vector._methods.MomentumProtocolLorentz.Et`.""" raise AssertionError @property def Et2(self) -> ScalarCollection: r""" Transverse energy squared of the four-momentum vector or array of vectors: $E_T^2 = E^2 \sin^2\theta$. """ raise AssertionError @property def et2(self) -> ScalarCollection: r""" Transverse energy squared of the four-momentum vector or array of vectors: $E_T^2 = E^2 \sin^2\theta$. """ raise AssertionError @property def transverse_energy2(self) -> ScalarCollection: """Synonym for :attr:`vector._methods.MomentumProtocolLorentz.Et2`.""" raise AssertionError @property def Mt(self) -> ScalarCollection: r""" Transverse mass of the four-momentum vector or array of vectors: $M_T = \sqrt{t^2 - z^2}$. """ raise AssertionError @property def mt(self) -> ScalarCollection: r""" Transverse mass of the four-momentum vector or array of vectors: $M_T = \sqrt{t^2 - z^2}$. """ raise AssertionError @property def transverse_mass(self) -> ScalarCollection: """Synonym for :attr:`vector._methods.MomentumProtocolLorentz.Mt`.""" raise AssertionError @property def Mt2(self) -> ScalarCollection: r""" Transverse mass squared of the four-momentum vector or array of vectors: $M_T^2 = t^2 - z^2$. """ raise AssertionError @property def mt2(self) -> ScalarCollection: r""" Transverse mass squared of the four-momentum vector or array of vectors: $M_T^2 = t^2 - z^2$. """ raise AssertionError @property def transverse_mass2(self) -> ScalarCollection: """Synonym for :attr:`vector._methods.MomentumProtocolLorentz.Mt2`.""" raise AssertionError class Vector(VectorProtocol): @typing.overload def __new__(cls, *, x: float, y: float) -> vector.VectorObject2D: ... @typing.overload def __new__(cls, *, rho: float, phi: float) -> vector.VectorObject2D: ... @typing.overload def __new__(cls, *, x: float, y: float, z: float) -> vector.VectorObject3D: ... @typing.overload def __new__(cls, *, x: float, y: float, eta: float) -> vector.VectorObject3D: ... @typing.overload def __new__(cls, *, x: float, y: float, theta: float) -> vector.VectorObject3D: ... @typing.overload def __new__(cls, *, rho: float, phi: float, z: float) -> vector.VectorObject3D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float ) -> vector.VectorObject3D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float ) -> vector.VectorObject3D: ... @typing.overload def __new__(cls, *, px: float, py: float) -> vector.MomentumObject2D: ... @typing.overload def __new__(cls, *, x: float, py: float) -> vector.MomentumObject2D: ... @typing.overload def __new__(cls, *, px: float, y: float) -> vector.MomentumObject2D: ... @typing.overload def __new__(cls, *, pt: float, phi: float) -> vector.MomentumObject2D: ... @typing.overload def __new__(cls, *, x: float, y: float, pz: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, x: float, py: float, z: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, x: float, py: float, pz: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, px: float, y: float, z: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, px: float, y: float, pz: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, px: float, py: float, z: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, px: float, py: float, pz: float) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, pt: float, phi: float, z: float) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, pz: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, theta: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, x: float, py: float, eta: float) -> vector.MomentumObject3D: ... @typing.overload def __new__(cls, *, px: float, y: float, eta: float) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, eta: float ) -> vector.MomentumObject3D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, t: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, t: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, z: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, pz: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, t: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, t: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, theta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, t: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, t: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, eta: float, t: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, tau: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, tau: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, ptau: float, phi: float, z: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, ptau: float, phi: float, pz: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, tau: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, tau: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, ptau: float, phi: float, theta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, tau: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, tau: float ) -> vector.VectorObject4D: ... @typing.overload def __new__( cls, *, ptau: float, phi: float, eta: float, tau: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pE: float, phi: float, z: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pE: float, phi: float, pz: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pE: float, phi: float, theta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pE: float, phi: float, eta: float, E: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pe: float, phi: float, z: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pe: float, phi: float, pz: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pe: float, phi: float, theta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pe: float, phi: float, eta: float, e: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, z: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, pz: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, theta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, eta: float, energy: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pM: float, phi: float, z: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pM: float, phi: float, pz: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pM: float, phi: float, theta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pM: float, phi: float, eta: float, M: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pm: float, phi: float, z: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pm: float, phi: float, pz: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pm: float, phi: float, theta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pm: float, phi: float, eta: float, m: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, z: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, pz: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, z: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, pz: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, z: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, pz: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, z: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, pz: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, z: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, pz: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, z: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, pz: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, theta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, theta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, theta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, theta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, theta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, theta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, y: float, eta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, x: float, py: float, eta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, y: float, eta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, px: float, py: float, eta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, rho: float, phi: float, eta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__( cls, *, pt: float, phi: float, eta: float, mass: float ) -> vector.MomentumObject4D: ... @typing.overload def __new__(cls, __azimuthal: Azimuthal) -> Vector: ... @typing.overload def __new__( cls, __azimuthal: Azimuthal, __longitudinal: Longitudinal ) -> Vector: ... @typing.overload def __new__( cls, __azimuthal: Azimuthal, __longitudinal: Longitudinal, __temporal: Temporal, ) -> Vector: ... def __new__(cls, *args: typing.Any, **kwargs: float) -> Vector: if cls is not Vector: return super().__new__(cls) return vector.obj(*args, **kwargs) def to_xy(self) -> VectorProtocolPlanar: from vector._compute import planar return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self)), [AzimuthalXY, None], 1, ) def to_pxpy(self) -> VectorProtocolPlanar: return self.to_xy() def to_rhophi(self) -> VectorProtocolPlanar: from vector._compute import planar return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self)), [AzimuthalRhoPhi, None], 1, ) def to_ptphi(self) -> VectorProtocolPlanar: return self.to_rhophi() def to_xyz(self, *, z: float | FloatArray = 0.0) -> VectorProtocolSpatial: from vector._compute import planar, spatial lcoord = z if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.z.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord), [AzimuthalXY, LongitudinalZ, None], 1, ) def to_xytheta(self, *, theta: float | FloatArray = 0.0) -> VectorProtocolSpatial: from vector._compute import planar, spatial lcoord = theta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.theta.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord), [AzimuthalXY, LongitudinalTheta, None], 1, ) def to_xyeta(self, *, eta: float | FloatArray = 0.0) -> VectorProtocolSpatial: from vector._compute import planar, spatial lcoord = eta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.eta.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord), [AzimuthalXY, LongitudinalEta, None], 1, ) def to_pxpypz(self, *, pz: float | FloatArray = 0.0) -> VectorProtocolSpatial: return self.to_xyz(z=pz) def to_pxpytheta(self, *, theta: float | FloatArray = 0.0) -> VectorProtocolSpatial: return self.to_xytheta(theta=theta) def to_pxpyeta(self, *, eta: float | FloatArray = 0.0) -> VectorProtocolSpatial: return self.to_xyeta(eta=eta) def to_rhophiz(self, *, z: float | FloatArray = 0.0) -> VectorProtocolSpatial: from vector._compute import planar, spatial lcoord = z if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.z.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord), [AzimuthalRhoPhi, LongitudinalZ, None], 1, ) def to_rhophitheta( self, *, theta: float | FloatArray = 0.0 ) -> VectorProtocolSpatial: from vector._compute import planar, spatial lcoord = theta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.theta.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord), [AzimuthalRhoPhi, LongitudinalTheta, None], 1, ) def to_rhophieta(self, *, eta: float | FloatArray = 0.0) -> VectorProtocolSpatial: from vector._compute import planar, spatial lcoord = eta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.eta.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord), [AzimuthalRhoPhi, LongitudinalEta, None], 1, ) def to_ptphipz(self, *, pz: float | FloatArray = 0.0) -> VectorProtocolSpatial: return self.to_rhophiz(z=pz) def to_ptphitheta( self, *, theta: float | FloatArray = 0.0 ) -> VectorProtocolSpatial: return self.to_rhophitheta(theta=theta) def to_ptphieta(self, *, eta: float | FloatArray = 0.0) -> VectorProtocolSpatial: return self.to_rhophieta(eta=eta) def to_xyzt( self, *, z: float | FloatArray = 0.0, t: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = z if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.z.dispatch(self) tcoord = t if isinstance(self, Vector4D): tcoord = lorentz.t.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord, tcoord), [AzimuthalXY, LongitudinalZ, TemporalT], 1, ) def to_xyztau( self, *, z: float | FloatArray = 0.0, tau: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = z if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.z.dispatch(self) tcoord = tau if isinstance(self, Vector4D): tcoord = lorentz.tau.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord, tcoord), [AzimuthalXY, LongitudinalZ, TemporalTau], 1, ) def to_xythetat( self, *, theta: float | FloatArray = 0.0, t: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = theta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.theta.dispatch(self) tcoord = t if isinstance(self, Vector4D): tcoord = lorentz.t.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord, tcoord), [AzimuthalXY, LongitudinalTheta, TemporalT], 1, ) def to_xythetatau( self, *, theta: float | FloatArray = 0.0, tau: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = theta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.theta.dispatch(self) tcoord = tau if isinstance(self, Vector4D): tcoord = lorentz.tau.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord, tcoord), [AzimuthalXY, LongitudinalTheta, TemporalTau], 1, ) def to_xyetat( self, *, eta: float | FloatArray = 0.0, t: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = eta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.eta.dispatch(self) tcoord = t if isinstance(self, Vector4D): tcoord = lorentz.t.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord, tcoord), [AzimuthalXY, LongitudinalEta, TemporalT], 1, ) def to_xyetatau( self, *, eta: float | FloatArray = 0.0, tau: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = eta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.eta.dispatch(self) tcoord = tau if isinstance(self, Vector4D): tcoord = lorentz.tau.dispatch(self) return self._wrap_result( type(self), (planar.x.dispatch(self), planar.y.dispatch(self), lcoord, tcoord), [AzimuthalXY, LongitudinalEta, TemporalTau], 1, ) def to_pxpypzenergy( self, *, pz: float | FloatArray = 0.0, energy: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_xyzt(z=pz, t=energy) def to_pxpythetaenergy( self, *, theta: float | FloatArray = 0.0, energy: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_xythetat(theta=theta, t=energy) def to_pxpyetaenergy( self, *, eta: float | FloatArray = 0.0, energy: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_xyetat(eta=eta, t=energy) def to_pxpypzmass( self, *, pz: float | FloatArray = 0.0, mass: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_xyztau(z=pz, tau=mass) def to_pxpythetamass( self, *, theta: float | FloatArray = 0.0, mass: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_xythetatau(theta=theta, tau=mass) def to_pxpyetamass( self, *, eta: float | FloatArray = 0.0, mass: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_xyetatau(eta=eta, tau=mass) def to_rhophizt( self, *, z: float | FloatArray = 0.0, t: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = z if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.z.dispatch(self) tcoord = t if isinstance(self, Vector4D): tcoord = lorentz.t.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord, tcoord), [AzimuthalRhoPhi, LongitudinalZ, TemporalT], 1, ) def to_rhophiztau( self, *, z: float | FloatArray = 0.0, tau: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = z if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.z.dispatch(self) tcoord = tau if isinstance(self, Vector4D): tcoord = lorentz.tau.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord, tcoord), [AzimuthalRhoPhi, LongitudinalZ, TemporalTau], 1, ) def to_rhophithetat( self, *, theta: float | FloatArray = 0.0, t: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = theta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.theta.dispatch(self) tcoord = t if isinstance(self, Vector4D): tcoord = lorentz.t.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord, tcoord), [AzimuthalRhoPhi, LongitudinalTheta, TemporalT], 1, ) def to_rhophithetatau( self, *, theta: float | FloatArray = 0.0, tau: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = theta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.theta.dispatch(self) tcoord = tau if isinstance(self, Vector4D): tcoord = lorentz.tau.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord, tcoord), [AzimuthalRhoPhi, LongitudinalTheta, TemporalTau], 1, ) def to_rhophietat( self, *, eta: float | FloatArray = 0.0, t: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = eta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.eta.dispatch(self) tcoord = t if isinstance(self, Vector4D): tcoord = lorentz.t.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord, tcoord), [AzimuthalRhoPhi, LongitudinalEta, TemporalT], 1, ) def to_rhophietatau( self, *, eta: float | FloatArray = 0.0, tau: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: from vector._compute import lorentz, planar, spatial lcoord = eta if isinstance(self, (Vector3D, Vector4D)): lcoord = spatial.eta.dispatch(self) tcoord = tau if isinstance(self, Vector4D): tcoord = lorentz.tau.dispatch(self) return self._wrap_result( type(self), (planar.rho.dispatch(self), planar.phi.dispatch(self), lcoord, tcoord), [AzimuthalRhoPhi, LongitudinalEta, TemporalTau], 1, ) def to_ptphipzenergy( self, *, pz: float | FloatArray = 0.0, energy: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_rhophizt(z=pz, t=energy) def to_ptphithetaenergy( self, *, theta: float | FloatArray = 0.0, energy: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_rhophithetat(theta=theta, t=energy) def to_ptphietaenergy( self, *, eta: float | FloatArray = 0.0, energy: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_rhophietat(eta=eta, t=energy) def to_ptphipzmass( self, *, pz: float | FloatArray = 0.0, mass: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_rhophiztau(z=pz, tau=mass) def to_ptphithetamass( self, *, theta: float | FloatArray = 0.0, mass: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_rhophithetatau(theta=theta, tau=mass) def to_ptphietamass( self, *, eta: float | FloatArray = 0.0, mass: float | FloatArray = 0.0 ) -> VectorProtocolLorentz: return self.to_rhophietatau(eta=eta, tau=mass) def like(self, other: VectorProtocol) -> VectorProtocol: if isinstance(other, Vector2D): return self.to_Vector2D() elif isinstance(other, Vector3D): return self.to_Vector3D() else: return self.to_Vector4D() class Vector2D(Vector, VectorProtocolPlanar): def to_Vector2D(self) -> VectorProtocolPlanar: return self def to_Vector3D( self, *, z: float | FloatArray | None = None, pz: float | FloatArray | None = None, theta: float | FloatArray | None = None, eta: float | FloatArray | None = None, ) -> VectorProtocolSpatial: """ Converts a 2D vector to 3D vector. The scalar longitudinal coordinate is broadcasted for NumPy and Awkward vectors. Only a single longitudinal coordinate should be provided. Examples: >>> import vector >>> vec = vector.VectorObject2D(x=1, y=2) >>> vec.to_Vector3D(z=1) VectorObject3D(x=1, y=2, z=1) >>> vec = vector.MomentumObject2D(px=1, py=2) >>> vec.to_Vector3D(pz=4) MomentumObject3D(px=1, py=2, pz=4) """ if sum(x is not None for x in (z, pz, theta, eta)) > 1: raise TypeError( "At most one longitudinal coordinate (`z`/`pz`, `theta`, or `eta`) may be assigned (non-None)" ) l_value: float | FloatArray = 0.0 l_type: type[Longitudinal] = LongitudinalZ if any(coord is not None for coord in (z, pz)): l_value = next(coord for coord in (z, pz) if coord is not None) elif eta is not None: l_value = eta l_type = LongitudinalEta elif theta is not None: l_value = theta l_type = LongitudinalTheta return self._wrap_result( type(self), (*self.azimuthal.elements, l_value), [_aztype(self), l_type, None], 1, ) def to_Vector4D( self, *, z: float | FloatArray | None = None, pz: float | FloatArray | None = None, theta: float | FloatArray | None = None, eta: float | FloatArray | None = None, t: float | FloatArray | None = None, e: float | FloatArray | None = None, E: float | FloatArray | None = None, energy: float | FloatArray | None = None, tau: float | FloatArray | None = None, m: float | FloatArray | None = None, M: float | FloatArray | None = None, mass: float | FloatArray | None = None, ) -> VectorProtocolLorentz: """ Converts a 2D vector to 4D vector. The scalar longitudinal and temporal coordinates are broadcasted for NumPy and Awkward vectors. Only a single longitudinal and temporal coordinate should be provided. Examples: >>> import vector >>> vec = vector.VectorObject2D(x=1, y=2) >>> vec.to_Vector4D(z=3, t=4) VectorObject4D(x=1, y=2, z=3, t=4) >>> vec = vector.MomentumObject2D(px=1, py=2) >>> vec.to_Vector4D(pz=4, energy=4) MomentumObject4D(px=1, py=2, pz=4, E=4) """ if sum(x is not None for x in (z, pz, theta, eta)) > 1: raise TypeError( "At most one longitudinal coordinate (`z`/`pz`, `theta`, or `eta`) may be assigned (non-None)" ) elif sum(x is not None for x in (t, tau, m, M, mass, e, E, energy)) > 1: raise TypeError( "At most one longitudinal coordinate (`t`/`e`/`E`/`energy`, `tau`/`m`/`M`/`mass`) may be assigned (non-None)" ) t_value: float | FloatArray = 0.0 t_type: type[Temporal] = TemporalT if any(coord is not None for coord in (tau, m, M, mass)): t_type = TemporalTau t_value = next(coord for coord in (tau, m, M, mass) if coord is not None) elif any(coord is not None for coord in (t, e, E, energy)): t_value = next(coord for coord in (t, e, E, energy) if coord is not None) l_value: float | FloatArray = 0.0 l_type: type[Longitudinal] = LongitudinalZ if any(coord is not None for coord in (z, pz)): l_value = next(coord for coord in (z, pz) if coord is not None) elif eta is not None: l_value = eta l_type = LongitudinalEta elif theta is not None: l_value = theta l_type = LongitudinalTheta return self._wrap_result( type(self), (*self.azimuthal.elements, l_value, t_value), [_aztype(self), l_type, t_type], 1, ) def to_2D(self) -> VectorProtocolPlanar: """ Alias for :meth:`vector._methods.Vector2D.to_Vector2D`. """ return self.to_Vector2D() def to_3D(self, **kwargs: float | None) -> VectorProtocolSpatial: """ Alias for :meth:`vector._methods.Vector2D.to_Vector3D`. """ return self.to_Vector3D(**kwargs) def to_4D(self, **kwargs: float | None) -> VectorProtocolLorentz: """ Alias for :meth:`vector._methods.Vector2D.to_Vector4D`. """ return self.to_Vector4D(**kwargs) class Vector3D(Vector, VectorProtocolSpatial): def to_Vector2D(self) -> VectorProtocolPlanar: return self._wrap_result( type(self), self.azimuthal.elements, [_aztype(self), None], 1, ) def to_Vector3D(self) -> VectorProtocolSpatial: return self def to_Vector4D( self, *, t: float | FloatArray | None = None, e: float | FloatArray | None = None, E: float | FloatArray | None = None, energy: float | FloatArray | None = None, tau: float | FloatArray | None = None, m: float | FloatArray | None = None, M: float | FloatArray | None = None, mass: float | FloatArray | None = None, ) -> VectorProtocolLorentz: """ Converts a 3D vector to 4D vector. The scalar temporal coordinate are broadcasted for NumPy and Awkward vectors. Only a single temporal coordinate should be provided. Examples: >>> import vector >>> vec = vector.VectorObject3D(x=1, y=2, z=3) >>> vec.to_Vector4D(t=4) VectorObject4D(x=1, y=2, z=3, t=4) >>> vec = vector.MomentumObject3D(px=1, py=2, pz=3) >>> vec.to_Vector4D(M=4) MomentumObject4D(px=1, py=2, pz=3, mass=4) """ if sum(x is not None for x in (t, tau, m, M, mass, e, E, energy)) > 1: raise TypeError( "At most one longitudinal coordinate (`t`/`e`/`E`/`energy`, `tau`/`m`/`M`/`mass`) may be assigned (non-None)" ) t_value: float | FloatArray = 0.0 t_type: type[Temporal] = TemporalT if any(coord is not None for coord in (tau, m, M, mass)): t_type = TemporalTau t_value = next(coord for coord in (tau, m, M, mass) if coord is not None) elif any(coord is not None for coord in (t, e, E, energy)): t_value = next(coord for coord in (t, e, E, energy) if coord is not None) return self._wrap_result( type(self), (*self.azimuthal.elements, *self.longitudinal.elements, t_value), [_aztype(self), _ltype(self), t_type], 1, ) def to_2D(self) -> VectorProtocolPlanar: """ Alias for :meth:`vector._methods.Vector3D.to_Vector2D`. """ return self.to_Vector2D() def to_3D(self) -> VectorProtocolSpatial: """ Alias for :meth:`vector._methods.Vector3D.to_Vector3D`. """ return self.to_Vector3D() def to_4D(self, **kwargs: float | None) -> VectorProtocolLorentz: """ Alias for :meth:`vector._methods.Vector3D.to_Vector4D`. """ return self.to_Vector4D(**kwargs) class Vector4D(Vector, VectorProtocolLorentz): def to_Vector2D(self) -> VectorProtocolPlanar: return self._wrap_result( type(self), self.azimuthal.elements, [_aztype(self), None], 1, ) def to_Vector3D(self) -> VectorProtocolSpatial: return self._wrap_result( type(self), self.azimuthal.elements + self.longitudinal.elements, [_aztype(self), _ltype(self), None], 1, ) def to_Vector4D(self) -> VectorProtocolLorentz: return self def to_2D(self) -> VectorProtocolPlanar: """ Alias for :meth:`vector._methods.Vector4D.to_Vector2D`. """ return self.to_Vector2D() def to_3D(self) -> VectorProtocolSpatial: """ Alias for :meth:`vector._methods.Vector4D.to_Vector3D`. """ return self.to_Vector3D() def to_4D(self) -> VectorProtocolLorentz: """ Alias for :meth:`vector._methods.Vector4D.to_Vector4D`. """ return self.to_Vector4D() class Planar(VectorProtocolPlanar): @property def x(self) -> ScalarCollection: from vector._compute.planar import x return x.dispatch(self) @property def y(self) -> ScalarCollection: from vector._compute.planar import y return y.dispatch(self) @property def rho(self) -> ScalarCollection: from vector._compute.planar import rho return rho.dispatch(self) @property def rho2(self) -> ScalarCollection: from vector._compute.planar import rho2 return rho2.dispatch(self) @property def phi(self) -> ScalarCollection: from vector._compute.planar import phi return phi.dispatch(self) def deltaphi(self, other: VectorProtocol) -> ScalarCollection: from vector._compute.planar import deltaphi return deltaphi.dispatch(self, other) def rotateZ(self: SameVectorType, angle: ScalarCollection) -> SameVectorType: from vector._compute.planar import rotateZ return rotateZ.dispatch(angle, self) def transform2D(self: SameVectorType, obj: TransformProtocol2D) -> SameVectorType: from vector._compute.planar import transform2D return transform2D.dispatch(obj, self) def is_parallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: from vector._compute.planar import is_parallel _maybe_same_dimension_error(self, other, self.is_parallel.__name__) return is_parallel.dispatch(tolerance, self, other) def is_antiparallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: from vector._compute.planar import is_antiparallel _maybe_same_dimension_error(self, other, self.is_antiparallel.__name__) return is_antiparallel.dispatch(tolerance, self, other) def is_perpendicular( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: from vector._compute.planar import is_perpendicular _maybe_same_dimension_error(self, other, self.is_perpendicular.__name__) return is_perpendicular.dispatch(tolerance, self, other) def unit(self: SameVectorType) -> SameVectorType: from vector._compute.planar import unit return unit.dispatch(self) def dot(self, other: VectorProtocol) -> ScalarCollection: _maybe_same_dimension_error(self, other, self.dot.__name__) module = _compute_module_of(self, other) return module.dot.dispatch(self, other) def add(self, other: VectorProtocol) -> VectorProtocol: _maybe_same_dimension_error(self, other, self.add.__name__) module = _compute_module_of(self, other) return module.add.dispatch(self, other) def subtract(self, other: VectorProtocol) -> VectorProtocol: _maybe_same_dimension_error(self, other, self.subtract.__name__) module = _compute_module_of(self, other) return module.subtract.dispatch(self, other) @property def neg2D(self: SameVectorType) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(-1, self) def scale2D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(factor, self) def scale(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(factor, self) def equal(self, other: VectorProtocol) -> BoolCollection: from vector._compute.planar import equal _maybe_same_dimension_error(self, other, self.equal.__name__) return equal.dispatch(self, other) def not_equal(self, other: VectorProtocol) -> BoolCollection: from vector._compute.planar import not_equal _maybe_same_dimension_error(self, other, self.not_equal.__name__) return not_equal.dispatch(self, other) def isclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: from vector._compute.planar import isclose _maybe_same_dimension_error(self, other, self.isclose.__name__) return isclose.dispatch(rtol, atol, equal_nan, self, other) class Spatial(Planar, VectorProtocolSpatial): @property def z(self) -> ScalarCollection: from vector._compute.spatial import z return z.dispatch(self) @property def theta(self) -> ScalarCollection: from vector._compute.spatial import theta return theta.dispatch(self) @property def eta(self) -> ScalarCollection: from vector._compute.spatial import eta return eta.dispatch(self) @property def costheta(self) -> ScalarCollection: from vector._compute.spatial import costheta return costheta.dispatch(self) @property def cottheta(self) -> ScalarCollection: from vector._compute.spatial import cottheta return cottheta.dispatch(self) @property def mag(self) -> ScalarCollection: from vector._compute.spatial import mag return mag.dispatch(self) @property def mag2(self) -> ScalarCollection: from vector._compute.spatial import mag2 return mag2.dispatch(self) def cross(self, other: VectorProtocolSpatial) -> VectorProtocolSpatial: from vector._compute.spatial import cross if dim(self) != 3 or dim(other) != 3: raise TypeError("cross is only defined for 3D vectors") return cross.dispatch(self, other) def deltaangle( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: from vector._compute.spatial import deltaangle if dim(other) != 3 and dim(other) != 4: raise TypeError(f"{other!r} is not a 3D or a 4D vector") return deltaangle.dispatch(self, other) def deltaeta( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: from vector._compute.spatial import deltaeta if dim(other) != 3 and dim(other) != 4: raise TypeError(f"{other!r} is not a 3D or a 4D vector") return deltaeta.dispatch(self, other) def deltaR( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: from vector._compute.spatial import deltaR if dim(other) != 3 and dim(other) != 4: raise TypeError(f"{other!r} is not a 3D or a 4D vector") return deltaR.dispatch(self, other) def deltaR2( self, other: VectorProtocolSpatial | VectorProtocolLorentz ) -> ScalarCollection: from vector._compute.spatial import deltaR2 if dim(other) != 3 and dim(other) != 4: raise TypeError(f"{other!r} is not a 3D or a 4D vector") return deltaR2.dispatch(self, other) def rotateX(self: SameVectorType, angle: ScalarCollection) -> SameVectorType: from vector._compute.spatial import rotateX return rotateX.dispatch(angle, self) def rotateY(self: SameVectorType, angle: ScalarCollection) -> SameVectorType: from vector._compute.spatial import rotateY return rotateY.dispatch(angle, self) def rotate_axis( self: SameVectorType, axis: VectorProtocolSpatial, angle: ScalarCollection ) -> SameVectorType: from vector._compute.spatial import rotate_axis if dim(axis) != 3: raise TypeError(f"{axis!r} is not a 3D vector") return rotate_axis.dispatch(angle, axis, self) def rotate_euler( self: SameVectorType, phi: ScalarCollection, theta: ScalarCollection, psi: ScalarCollection, order: str = "zxz", ) -> SameVectorType: from vector._compute.spatial import rotate_euler return rotate_euler.dispatch(phi, theta, psi, order.lower(), self) def rotate_nautical( self: SameVectorType, yaw: ScalarCollection, pitch: ScalarCollection, roll: ScalarCollection, ) -> SameVectorType: # The order of arguments is reversed because rotate_euler # follows ROOT's argument order: phi, theta, psi. from vector._compute.spatial import rotate_euler return rotate_euler.dispatch(roll, pitch, yaw, "zyx", self) def rotate_quaternion( self: SameVectorType, u: ScalarCollection, i: ScalarCollection, j: ScalarCollection, k: ScalarCollection, ) -> SameVectorType: from vector._compute.spatial import rotate_quaternion return rotate_quaternion.dispatch(u, i, j, k, self) def transform3D(self: SameVectorType, obj: TransformProtocol3D) -> SameVectorType: from vector._compute.spatial import transform3D return transform3D.dispatch(obj, self) def is_parallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: from vector._compute.spatial import is_parallel _maybe_same_dimension_error(self, other, self.is_parallel.__name__) return is_parallel.dispatch(tolerance, self, other) def is_antiparallel( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: from vector._compute.spatial import is_antiparallel _maybe_same_dimension_error(self, other, self.is_antiparallel.__name__) return is_antiparallel.dispatch(tolerance, self, other) def is_perpendicular( self, other: VectorProtocol, tolerance: ScalarCollection = 1e-5 ) -> BoolCollection: from vector._compute.spatial import is_perpendicular _maybe_same_dimension_error(self, other, self.is_perpendicular.__name__) return is_perpendicular.dispatch(tolerance, self, other) def unit(self: SameVectorType) -> SameVectorType: from vector._compute.spatial import unit return unit.dispatch(self) def dot(self, other: VectorProtocol) -> ScalarCollection: _maybe_same_dimension_error(self, other, self.dot.__name__) module = _compute_module_of(self, other) return module.dot.dispatch(self, other) def add(self, other: VectorProtocol) -> VectorProtocol: _maybe_same_dimension_error(self, other, self.add.__name__) module = _compute_module_of(self, other) return module.add.dispatch(self, other) def subtract(self, other: VectorProtocol) -> VectorProtocol: _maybe_same_dimension_error(self, other, self.subtract.__name__) module = _compute_module_of(self, other) return module.subtract.dispatch(self, other) @property def neg2D(self: SameVectorType) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(-1, self) @property def neg3D(self: SameVectorType) -> SameVectorType: from vector._compute.spatial import scale return scale.dispatch(-1, self) def scale2D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(factor, self) def scale3D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.spatial import scale return scale.dispatch(factor, self) def scale(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.spatial import scale return scale.dispatch(factor, self) def equal(self, other: VectorProtocol) -> BoolCollection: from vector._compute.spatial import equal _maybe_same_dimension_error(self, other, self.equal.__name__) return equal.dispatch(self, other) def not_equal(self, other: VectorProtocol) -> BoolCollection: from vector._compute.spatial import not_equal _maybe_same_dimension_error(self, other, self.not_equal.__name__) return not_equal.dispatch(self, other) def isclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: from vector._compute.spatial import isclose _maybe_same_dimension_error(self, other, self.isclose.__name__) return isclose.dispatch(rtol, atol, equal_nan, self, other) class Lorentz(Spatial, VectorProtocolLorentz): @property def t(self) -> ScalarCollection: from vector._compute.lorentz import t return t.dispatch(self) @property def t2(self) -> ScalarCollection: from vector._compute.lorentz import t2 return t2.dispatch(self) @property def tau(self) -> ScalarCollection: from vector._compute.lorentz import tau return tau.dispatch(self) @property def tau2(self) -> ScalarCollection: from vector._compute.lorentz import tau2 return tau2.dispatch(self) @property def beta(self) -> ScalarCollection: from vector._compute.lorentz import beta return beta.dispatch(self) @property def gamma(self) -> ScalarCollection: from vector._compute.lorentz import gamma return gamma.dispatch(self) @property def rapidity(self) -> ScalarCollection: from vector._compute.lorentz import rapidity return rapidity.dispatch(self) def deltaRapidityPhi(self, other: VectorProtocolLorentz) -> ScalarCollection: from vector._compute.lorentz import deltaRapidityPhi if not dim(other) == 4: raise TypeError(f"{other!r} is not a 4D vector") return deltaRapidityPhi.dispatch(self, other) def deltaRapidityPhi2(self, other: VectorProtocolLorentz) -> ScalarCollection: from vector._compute.lorentz import deltaRapidityPhi2 if not dim(other) == 4: raise TypeError(f"{other!r} is not a 4D vector") return deltaRapidityPhi2.dispatch(self, other) def boost_p4(self: SameVectorType, p4: VectorProtocolLorentz) -> SameVectorType: from vector._compute.lorentz import boost_p4 if dim(p4) != 4: raise TypeError(f"{p4!r} is not a 4D vector") return boost_p4.dispatch(self, p4) def boost_beta3( self: SameVectorType, beta3: VectorProtocolSpatial ) -> SameVectorType: from vector._compute.lorentz import boost_beta3 if dim(beta3) != 3: raise TypeError(f"{beta3!r} is not a 3D vector") return boost_beta3.dispatch(self, beta3) def boost( self: SameVectorType, booster: VectorProtocolSpatial | VectorProtocolLorentz ) -> SameVectorType: from vector._compute.lorentz import boost_beta3, boost_p4 if isinstance(booster, Vector3D): return boost_beta3.dispatch(self, booster) elif isinstance(booster, Vector4D): return boost_p4.dispatch(self, booster) else: raise TypeError( "specify a Vector3D to boost by beta (velocity with c=1) or " "a Vector4D to boost by a momentum 4-vector" ) def boostCM_of_p4( self: SameVectorType, p4: VectorProtocolLorentz ) -> SameVectorType: from vector._compute.lorentz import boost_p4 if dim(p4) != 4: raise TypeError(f"{p4!r} is not a 4D momentum vector") return boost_p4.dispatch(self, p4.neg3D) def boostCM_of_beta3( self: SameVectorType, beta3: VectorProtocolSpatial ) -> SameVectorType: from vector._compute.lorentz import boost_beta3 if dim(beta3) != 3: raise TypeError(f"{beta3!r} is not a 3D momentum vector") return boost_beta3.dispatch(self, beta3.neg3D) def boostCM_of( self: SameVectorType, booster: VectorProtocolSpatial | VectorProtocolLorentz ) -> SameVectorType: from vector._compute.lorentz import boost_beta3, boost_p4 if isinstance(booster, Vector3D): return boost_beta3.dispatch(self, booster.neg3D) elif isinstance(booster, Vector4D): return boost_p4.dispatch(self, booster.neg3D) else: raise TypeError( "specify a Vector3D to boost by beta (velocity with c=1) or " "a Vector4D to boost by a momentum 4-vector" ) def boostX( self: SameVectorType, beta: ScalarCollection | None = None, gamma: ScalarCollection | None = None, ) -> SameVectorType: from vector._compute.lorentz import boostX_beta, boostX_gamma if beta is not None and gamma is None: return boostX_beta.dispatch(beta, self) elif beta is None and gamma is not None: return boostX_gamma.dispatch(gamma, self) else: raise TypeError("specify 'beta' xor 'gamma', not both or neither") def boostY( self: SameVectorType, beta: ScalarCollection | None = None, gamma: ScalarCollection | None = None, ) -> SameVectorType: from vector._compute.lorentz import boostY_beta, boostY_gamma if beta is not None and gamma is None: return boostY_beta.dispatch(beta, self) elif beta is None and gamma is not None: return boostY_gamma.dispatch(gamma, self) else: raise TypeError("specify 'beta' xor 'gamma', not both or neither") def boostZ( self: SameVectorType, beta: ScalarCollection | None = None, gamma: ScalarCollection | None = None, ) -> SameVectorType: from vector._compute.lorentz import boostZ_beta, boostZ_gamma if beta is not None and gamma is None: return boostZ_beta.dispatch(beta, self) elif beta is None and gamma is not None: return boostZ_gamma.dispatch(gamma, self) else: raise TypeError("specify 'beta' xor 'gamma', not both or neither") def transform4D(self: SameVectorType, obj: TransformProtocol4D) -> SameVectorType: from vector._compute.lorentz import transform4D return transform4D.dispatch(obj, self) def to_beta3(self) -> VectorProtocolSpatial: from vector._compute.lorentz import to_beta3 return to_beta3.dispatch(self) def is_timelike(self, tolerance: ScalarCollection = 0) -> BoolCollection: from vector._compute.lorentz import is_timelike return is_timelike.dispatch(tolerance, self) def is_spacelike(self, tolerance: ScalarCollection = 0) -> BoolCollection: from vector._compute.lorentz import is_spacelike return is_spacelike.dispatch(tolerance, self) def is_lightlike(self, tolerance: ScalarCollection = 1e-5) -> BoolCollection: from vector._compute.lorentz import is_lightlike return is_lightlike.dispatch(tolerance, self) def unit(self: SameVectorType) -> SameVectorType: from vector._compute.lorentz import unit return unit.dispatch(self) def dot(self, other: VectorProtocol) -> ScalarCollection: _maybe_same_dimension_error(self, other, self.dot.__name__) module = _compute_module_of(self, other) return module.dot.dispatch(self, other) def add(self, other: VectorProtocol) -> VectorProtocol: _maybe_same_dimension_error(self, other, self.add.__name__) module = _compute_module_of(self, other) return module.add.dispatch(self, other) def subtract(self, other: VectorProtocol) -> VectorProtocol: _maybe_same_dimension_error(self, other, self.subtract.__name__) module = _compute_module_of(self, other) return module.subtract.dispatch(self, other) @property def neg2D(self: SameVectorType) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(-1, self) @property def neg3D(self: SameVectorType) -> SameVectorType: from vector._compute.spatial import scale return scale.dispatch(-1, self) @property def neg4D(self: SameVectorType) -> SameVectorType: from vector._compute.lorentz import scale return scale.dispatch(-1, self) def scale2D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.planar import scale return scale.dispatch(factor, self) def scale3D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.spatial import scale return scale.dispatch(factor, self) def scale4D(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.lorentz import scale return scale.dispatch(factor, self) def scale(self: SameVectorType, factor: ScalarCollection) -> SameVectorType: from vector._compute.lorentz import scale return scale.dispatch(factor, self) def equal(self, other: VectorProtocol) -> BoolCollection: from vector._compute.lorentz import equal _maybe_same_dimension_error(self, other, self.equal.__name__) return equal.dispatch(self, other) def not_equal(self, other: VectorProtocol) -> BoolCollection: from vector._compute.lorentz import not_equal _maybe_same_dimension_error(self, other, self.not_equal.__name__) return not_equal.dispatch(self, other) def isclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: from vector._compute.lorentz import isclose _maybe_same_dimension_error(self, other, self.isclose.__name__) return isclose.dispatch(rtol, atol, equal_nan, self, other) class Momentum: pass class PlanarMomentum(Momentum, MomentumProtocolPlanar): @property def px(self) -> ScalarCollection: return self.x @property def py(self) -> ScalarCollection: return self.y @property def pt(self) -> ScalarCollection: return self.rho @property def pt2(self) -> ScalarCollection: return self.rho2 class SpatialMomentum(PlanarMomentum, MomentumProtocolSpatial): @property def pz(self) -> ScalarCollection: return self.z @property def pseudorapidity(self) -> ScalarCollection: return self.eta @property def p(self) -> ScalarCollection: return self.mag @property def p2(self) -> ScalarCollection: return self.mag2 class LorentzMomentum(SpatialMomentum, MomentumProtocolLorentz): @property def E(self) -> ScalarCollection: return self.t @property def e(self) -> ScalarCollection: return self.t @property def energy(self) -> ScalarCollection: return self.t @property def E2(self) -> ScalarCollection: return self.t2 @property def e2(self) -> ScalarCollection: return self.t2 @property def energy2(self) -> ScalarCollection: return self.t2 @property def M(self) -> ScalarCollection: return self.tau @property def m(self) -> ScalarCollection: return self.tau @property def mass(self) -> ScalarCollection: return self.tau @property def M2(self) -> ScalarCollection: return self.tau2 @property def m2(self) -> ScalarCollection: return self.tau2 @property def mass2(self) -> ScalarCollection: return self.tau2 @property def Et(self) -> ScalarCollection: from vector._compute.lorentz import Et return Et.dispatch(self) @property def et(self) -> ScalarCollection: from vector._compute.lorentz import Et return Et.dispatch(self) @property def transverse_energy(self) -> ScalarCollection: return self.Et @property def Et2(self) -> ScalarCollection: from vector._compute.lorentz import Et2 return Et2.dispatch(self) @property def et2(self) -> ScalarCollection: from vector._compute.lorentz import Et2 return Et2.dispatch(self) @property def transverse_energy2(self) -> ScalarCollection: return self.Et2 @property def Mt(self) -> ScalarCollection: from vector._compute.lorentz import Mt return Mt.dispatch(self) @property def mt(self) -> ScalarCollection: from vector._compute.lorentz import Mt return Mt.dispatch(self) @property def transverse_mass(self) -> ScalarCollection: return self.Mt @property def Mt2(self) -> ScalarCollection: from vector._compute.lorentz import Mt2 return Mt2.dispatch(self) @property def mt2(self) -> ScalarCollection: from vector._compute.lorentz import Mt2 return Mt2.dispatch(self) @property def transverse_mass2(self) -> ScalarCollection: return self.Mt2 def dim(v: VectorProtocol) -> int: """Returns the number of dimensions in a vector: 2, 3, or 4.""" if isinstance(v, Vector2D): return 2 elif isinstance(v, Vector3D): return 3 elif isinstance(v, Vector4D): return 4 else: raise TypeError(f"{v!r} is not a vector.Vector") def _maybe_same_dimension_error( v1: VectorProtocol, v2: VectorProtocol, operation: str ) -> None: """Raises an error if the vectors are not of the same dimension.""" if dim(v1) != dim(v2): raise TypeError( f"""{v1!r} and {v2!r} do not have the same dimension; use a.like(b).{operation}(b) or a.{operation}(b.like(a)) or the binary operation equivalent to project or embed one of the vectors to match the other's dimensionality """ ) def _compute_module_of( one: VectorProtocol, two: VectorProtocol, nontemporal: bool = False ) -> Module: """ Determines which compute module to use for functions of two vectors (the one with minimum dimension). If ``nontemporal``, use a spatial module even if both vectors are 4D. """ if not isinstance(one, Vector): raise TypeError(f"{one!r} is not a Vector") if not isinstance(two, Vector): raise TypeError(f"{two!r} is not a Vector") if isinstance(one, Vector2D): import vector._compute.planar return vector._compute.planar elif isinstance(one, Vector3D): if isinstance(two, Vector2D): import vector._compute.planar return vector._compute.planar else: import vector._compute.spatial return vector._compute.spatial elif isinstance(one, Vector4D): if isinstance(two, Vector2D): import vector._compute.planar return vector._compute.planar elif isinstance(two, Vector3D) or nontemporal: import vector._compute.spatial return vector._compute.spatial else: import vector._compute.lorentz return vector._compute.lorentz raise AssertionError(repr(one)) _coordinate_class_to_names = { AzimuthalXY: ("x", "y"), AzimuthalRhoPhi: ("rho", "phi"), LongitudinalZ: ("z",), LongitudinalTheta: ("theta",), LongitudinalEta: ("eta",), TemporalT: ("t",), TemporalTau: ("tau",), } _repr_generic_to_momentum = { "x": "px", "y": "py", "rho": "pt", "z": "pz", "t": "E", "tau": "mass", } _repr_momentum_to_generic = { "px": "x", "py": "y", "pt": "rho", "pz": "z", "E": "t", "e": "t", "energy": "t", "M": "tau", "m": "tau", "mass": "tau", } _coordinate_order = [ "x", "px", "y", "py", "rho", "pt", "phi", "z", "pz", "theta", "eta", "t", "E", "e", "energy", "tau", "M", "m", "mass", ] def _aztype(obj: VectorProtocolPlanar) -> type[Coordinates]: """ Determines the Azimuthal type of a vector for use in looking up a dispatched function. """ if hasattr(obj, "azimuthal"): for t in type(obj.azimuthal).__mro__: if t in (AzimuthalXY, AzimuthalRhoPhi): return t raise AssertionError(repr(obj)) def _ltype(obj: VectorProtocolSpatial) -> type[Coordinates]: """ Determines the Longitudinal type of a vector for use in looking up a dispatched function. """ if hasattr(obj, "longitudinal"): for t in type(obj.longitudinal).__mro__: if t in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): return t raise AssertionError(repr(obj)) def _ttype(obj: VectorProtocolLorentz) -> type[Coordinates]: """ Determines the Temporal type of a vector for use in looking up a dispatched function. """ if hasattr(obj, "temporal"): for t in type(obj.temporal).__mro__: if t in (TemporalT, TemporalTau): return t raise AssertionError(repr(obj)) def _lib_of(*objects: VectorProtocol) -> Module: # NumPy-like module """ Determines the ``lib`` of a vector or set of vectors, complaining if they're incompatible. """ lib: typing.Any | None = None for obj in objects: if isinstance(obj, Vector): if lib is None: lib = obj.lib elif lib != obj.lib: raise TypeError( f"cannot use {lib} and {obj.lib} in the same calculation" ) assert lib is not None return lib def _from_signature( name: str, dispatch_map: dict[typing.Any, typing.Any], signature: tuple[typing.Any, ...], ) -> tuple[typing.Any, ...]: """ Gets a function and its return type from a ``dispatch_map`` and the ``signature`` to search for (complaining if none is found). """ result = dispatch_map.get(signature) if result is None: raise TypeError( f"function {'.'.join(name.split('.')[-2:])!r} has no signature {signature}" ) return result _handler_priority = [ "vector.backends.object", "vector.backends.numpy", "vector.backends.sympy", "vector.backends.awkward", ] def _get_handler_index(obj: VectorProtocol) -> int: """Returns the index of the first valid handler checking the list of parent classes""" for cls in type(obj).__mro__: with suppress(ValueError): return _handler_priority.index(cls.__module__) raise AssertionError( f"Could not find a valid handler for {obj}! This should not happen." ) def _check_instance( any_or_all: typing.Callable[[typing.Iterable[bool]], bool], objects: tuple[VectorProtocol, ...], clas: type[VectorProtocol], ) -> bool: return any_or_all(isinstance(v, clas) for v in objects) def _handler_of(*objects: VectorProtocol) -> VectorProtocol: """ Determines which vector should wrap the output of a dispatched function. Awkward Arrays have higher priority than NumPy arrays, which have higher priority than Python objects, which has the effect of "promoting" Python objects to NumPy arrays to Awkward Arrays whenever two are used in the same formula. """ handler: VectorProtocol | None = None for obj in objects: if not isinstance(obj, Vector): continue if handler is None or _get_handler_index(obj) > _get_handler_index(handler): handler = obj assert handler is not None return handler def _flavor_of(*objects: VectorProtocol) -> type[VectorProtocol]: """ Determines the flavor of the output of a dispatched function, where "flavor" is generic vs momentum. """ from vector.backends.numpy import VectorNumpy from vector.backends.object import VectorObject handler: VectorProtocol | None = None is_momentum = any(isinstance(obj, Momentum) for obj in objects) for obj in objects: if isinstance(obj, Vector): if handler is None: handler = obj elif isinstance(obj, VectorObject): pass elif isinstance(obj, VectorNumpy): handler = obj assert handler is not None if is_momentum: return handler.MomentumClass else: return handler.GenericClass vector-1.6.3/src/vector/_typeutils.py000066400000000000000000000033641503546127100177010ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import typing from typing import Protocol, TypedDict __all__ = [ "BoolCollection", "FloatArray", "Protocol", "ScalarCollection", "TransformProtocol2D", "TransformProtocol3D", "TransformProtocol4D", ] def __dir__() -> list[str]: return __all__ # Represents a number, a NumPy array, an Awkward Array, etc., of non-vectors. ScalarCollection = typing.Any # Represents a bool, a NumPy array of bools, an Awkward Array of bools, etc. BoolCollection = typing.Any class TransformProtocol2D(TypedDict): xx: ScalarCollection xy: ScalarCollection yx: ScalarCollection yy: ScalarCollection class TransformProtocol3D(TypedDict): xx: ScalarCollection xy: ScalarCollection xz: ScalarCollection yx: ScalarCollection yy: ScalarCollection yz: ScalarCollection zx: ScalarCollection zy: ScalarCollection zz: ScalarCollection class TransformProtocol4D(TypedDict): xx: ScalarCollection xy: ScalarCollection xz: ScalarCollection xt: ScalarCollection yx: ScalarCollection yy: ScalarCollection yz: ScalarCollection yt: ScalarCollection zx: ScalarCollection zy: ScalarCollection zz: ScalarCollection zt: ScalarCollection tx: ScalarCollection ty: ScalarCollection tz: ScalarCollection tt: ScalarCollection if typing.TYPE_CHECKING: import numpy.typing FloatArray = numpy.typing.NDArray[numpy.float64] else: import numpy FloatArray = numpy.ndarray vector-1.6.3/src/vector/_version.pyi000066400000000000000000000001661503546127100174720ustar00rootroot00000000000000from __future__ import annotations version: str version_tuple: tuple[int, int, int] | tuple[int, int, int, str, str] vector-1.6.3/src/vector/backends/000077500000000000000000000000001503546127100166725ustar00rootroot00000000000000vector-1.6.3/src/vector/backends/__init__.py000066400000000000000000000011371503546127100210050ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Vector comes loaded with 3 + 2 backends; a pure Python object backend, a NumPy backend, an Awkward Array backend, an Object-Numba, and an Awkward-Numba backend to leverage JIT (Just In Time) compiled calculations on vectors. Other potential future vanilla backends include Tensorflow and JAX, and other possible future Numba-backends include Numba-NumPy. """ vector-1.6.3/src/vector/backends/_numba.py000066400000000000000000000033301503546127100205040ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import types import typing import numba import vector._compute.lorentz import vector._compute.planar import vector._compute.spatial names_and_modules = [ ("planar", vector._compute.planar), ("spatial", vector._compute.spatial), ("lorentz", vector._compute.lorentz), ] numba_modules: typing.Any = {} registered = set() for groupname, module in names_and_modules: numba_modules[groupname] = {} for modname, submodule in module.__dict__.items(): if isinstance(submodule, types.ModuleType) and submodule.__name__.startswith( "vector._compute." ): new_name = submodule.__name__.replace( "vector._compute.", "vector._compute._numba." ) numba_modules[groupname][modname] = {} for name, obj in submodule.__dict__.items(): if ( isinstance(obj, types.FunctionType) and name != "dispatch" and obj.__module__ == submodule.__name__ ): numba.extending.register_jitable(obj) registered.add(obj) for key, value in submodule.dispatch_map.items(): function, *returns = value if function not in registered: numba.extending.register_jitable(function) registered.add(function) numba_modules[groupname][modname][key] = (function, *returns) vector-1.6.3/src/vector/backends/_numba_object.py000066400000000000000000003413651503546127100220470ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. # type: ignore """ Implements VectorObjects in Numba. Every function should be made usable in Numba. """ from __future__ import annotations import operator import numba import numpy import vector from vector._methods import ( AzimuthalRhoPhi, AzimuthalXY, LongitudinalEta, LongitudinalTheta, LongitudinalZ, Momentum, TemporalT, TemporalTau, _from_signature, ) from vector.backends._numba import numba_modules from vector.backends.object import ( AzimuthalObject, AzimuthalObjectRhoPhi, AzimuthalObjectXY, LongitudinalObject, LongitudinalObjectEta, LongitudinalObjectTheta, LongitudinalObjectZ, MomentumObject2D, MomentumObject3D, MomentumObject4D, TemporalObject, TemporalObjectT, TemporalObjectTau, VectorObject2D, VectorObject3D, VectorObject4D, _coord_object_type, ) @numba.extending.overload(numpy.nan_to_num) # FIXME: This needs to go into Numba! def nan_to_num(x, copy=True, nan=0.0, posinf=None, neginf=None): if isinstance(x, numba.types.Array): def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None): out = numpy.copy(x).reshape(-1) if copy else x.reshape(-1) for i in range(len(out)): if numpy.isnan(out[i]): out[i] = nan if posinf is not None and numpy.isinf(out[i]) and out[i] > 0: out[i] = posinf if neginf is not None and numpy.isinf(out[i]) and out[i] < 0: out[i] = neginf return out.reshape(x.shape) else: def nan_to_num_impl(x, copy=True, nan=0.0, posinf=None, neginf=None): if numpy.isnan(x): return nan if posinf is not None and numpy.isinf(x) and x > 0: return posinf if neginf is not None and numpy.isinf(x) and x < 0: return neginf return x return nan_to_num_impl # Since CoordinateObjects are NamedTuples, we get their types wrapped for free. def is_azimuthaltype(typ): return isinstance( typ, (numba.types.NamedTuple, numba.types.NamedUniTuple) ) and issubclass(typ.instance_class, AzimuthalObject) def is_longitudinaltype(typ): return isinstance( typ, (numba.types.NamedTuple, numba.types.NamedUniTuple) ) and issubclass(typ.instance_class, LongitudinalObject) def is_temporaltype(typ): return isinstance( typ, (numba.types.NamedTuple, numba.types.NamedUniTuple) ) and issubclass(typ.instance_class, TemporalObject) def numba_aztype(typ): for t in typ.azimuthaltype.instance_class.__mro__: if t in (AzimuthalXY, AzimuthalRhoPhi): return t raise AssertionError def numba_ltype(typ): for t in typ.longitudinaltype.instance_class.__mro__: if t in (LongitudinalZ, LongitudinalTheta, LongitudinalEta): return t raise AssertionError def numba_ttype(typ): for t in typ.temporaltype.instance_class.__mro__: if t in (TemporalT, TemporalTau): return t raise AssertionError # VectorObject2D (and momentum) ############################################### class VectorObject2DType(numba.types.Type): instance_class = VectorObject2D def __init__(self, azimuthaltype): super().__init__(name=f"VectorObject2DType({azimuthaltype})") self.azimuthaltype = azimuthaltype class MomentumObject2DType(VectorObject2DType): instance_class = MomentumObject2D def __init__(self, azimuthaltype): super().__init__(azimuthaltype) self.name = f"MomentumObject2DType({azimuthaltype})" @numba.extending.register_model(VectorObject2DType) @numba.extending.register_model(MomentumObject2DType) class VectorObject2DModel(numba.extending.models.StructModel): def __init__(self, dmm, fe_type): members = [ ("azimuthal", fe_type.azimuthaltype), ] super().__init__(dmm, fe_type, members) @numba.extending.typeof_impl.register(VectorObject2D) def VectorObject2D_typeof(val, c): if isinstance(val, MomentumObject2D): return MomentumObject2DType(numba.typeof(val.azimuthal)) else: return VectorObject2DType(numba.typeof(val.azimuthal)) numba.extending.make_attribute_wrapper(VectorObject2DType, "azimuthal", "azimuthal") numba.extending.make_attribute_wrapper(MomentumObject2DType, "azimuthal", "azimuthal") @numba.extending.type_callable(VectorObject2D) def VectorObject2D_constructor_typer(context): def typer(azimuthaltype): if is_azimuthaltype(azimuthaltype): return VectorObject2DType(azimuthaltype) else: raise numba.TypingError( "VectorObject2D constructor requires an AzimuthalObject as its argument" ) return typer @numba.extending.type_callable(MomentumObject2D) def MomentumObject2D_constructor_typer(context): def typer(azimuthaltype): if is_azimuthaltype(azimuthaltype): return MomentumObject2DType(azimuthaltype) else: raise numba.TypingError( "MomentumObject2D constructor requires an AzimuthalObject as its argument" ) return typer @numba.extending.lower_builtin(VectorObject2D, numba.types.Type) @numba.extending.lower_builtin(MomentumObject2D, numba.types.Type) def VectorObject2D_constructor_impl(context, builder, sig, args): typ = sig.return_type proxyout = numba.core.cgutils.create_struct_proxy(typ)(context, builder) proxyout.azimuthal = args[0] return proxyout._getvalue() @numba.extending.unbox(VectorObject2DType) def VectorObject2D_unbox(typ, obj, c): azimuthal_obj = c.pyapi.object_getattr_string(obj, "azimuthal") proxyout = numba.core.cgutils.create_struct_proxy(typ)(c.context, c.builder) proxyout.azimuthal = c.pyapi.to_native_value(typ.azimuthaltype, azimuthal_obj).value c.pyapi.decref(azimuthal_obj) is_error = numba.core.cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) return numba.extending.NativeValue(proxyout._getvalue(), is_error) @numba.extending.box(VectorObject2DType) def VectorObject2D_box(typ, val, c): if isinstance(typ, MomentumObject2DType): cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(MomentumObject2D)) else: cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(VectorObject2D)) proxyin = numba.core.cgutils.create_struct_proxy(typ)( c.context, c.builder, value=val ) azimuthal_obj = c.pyapi.from_native_value(typ.azimuthaltype, proxyin.azimuthal) output_obj = c.pyapi.call_function_objargs(cls_obj, (azimuthal_obj,)) c.pyapi.decref(cls_obj) c.pyapi.decref(azimuthal_obj) return output_obj # VectorObject3D (and momentum) ############################################### class VectorObject3DType(numba.types.Type): instance_class = VectorObject3D def __init__(self, azimuthaltype, longitudinaltype): super().__init__( name=f"VectorObject3DType({azimuthaltype}, {longitudinaltype})" ) self.azimuthaltype = azimuthaltype self.longitudinaltype = longitudinaltype class MomentumObject3DType(VectorObject3DType): instance_class = MomentumObject3D def __init__(self, azimuthaltype, longitudinaltype): super().__init__(azimuthaltype, longitudinaltype) self.name = f"MomentumObject3DType({azimuthaltype}, {longitudinaltype})" @numba.extending.register_model(VectorObject3DType) @numba.extending.register_model(MomentumObject3DType) class VectorObject3DModel(numba.extending.models.StructModel): def __init__(self, dmm, fe_type): members = [ ("azimuthal", fe_type.azimuthaltype), ("longitudinal", fe_type.longitudinaltype), ] super().__init__(dmm, fe_type, members) @numba.extending.typeof_impl.register(VectorObject3D) def VectorObject3D_typeof(val, c): if isinstance(val, MomentumObject3D): return MomentumObject3DType( numba.typeof(val.azimuthal), numba.typeof(val.longitudinal) ) else: return VectorObject3DType( numba.typeof(val.azimuthal), numba.typeof(val.longitudinal) ) numba.extending.make_attribute_wrapper(VectorObject3DType, "azimuthal", "azimuthal") numba.extending.make_attribute_wrapper( VectorObject3DType, "longitudinal", "longitudinal" ) numba.extending.make_attribute_wrapper(MomentumObject3DType, "azimuthal", "azimuthal") numba.extending.make_attribute_wrapper( MomentumObject3DType, "longitudinal", "longitudinal" ) @numba.extending.type_callable(VectorObject3D) def VectorObject3D_constructor_typer(context): def typer(azimuthaltype, longitudinaltype): if is_azimuthaltype(azimuthaltype) and is_longitudinaltype(longitudinaltype): return VectorObject3DType(azimuthaltype, longitudinaltype) else: raise numba.TypingError( "VectorObject3D constructor requires an AzimuthalObject and a " "LongitudinalObject as its arguments" ) return typer @numba.extending.type_callable(MomentumObject3D) def MomentumObject3D_constructor_typer(context): def typer(azimuthaltype, longitudinaltype): if is_azimuthaltype(azimuthaltype) and is_longitudinaltype(longitudinaltype): return MomentumObject3DType(azimuthaltype, longitudinaltype) else: raise numba.TypingError( "MomentumObject3D constructor requires an AzimuthalObject and a " "LongitudinalObject as its arguments" ) return typer @numba.extending.lower_builtin(VectorObject3D, numba.types.Type, numba.types.Type) @numba.extending.lower_builtin(MomentumObject3D, numba.types.Type, numba.types.Type) def VectorObject3D_constructor_impl(context, builder, sig, args): typ = sig.return_type proxyout = numba.core.cgutils.create_struct_proxy(typ)(context, builder) proxyout.azimuthal = args[0] proxyout.longitudinal = args[1] return proxyout._getvalue() @numba.extending.unbox(VectorObject3DType) def VectorObject3D_unbox(typ, obj, c): azimuthal_obj = c.pyapi.object_getattr_string(obj, "azimuthal") longitudinal_obj = c.pyapi.object_getattr_string(obj, "longitudinal") proxyout = numba.core.cgutils.create_struct_proxy(typ)(c.context, c.builder) proxyout.azimuthal = c.pyapi.to_native_value(typ.azimuthaltype, azimuthal_obj).value proxyout.longitudinal = c.pyapi.to_native_value( typ.longitudinaltype, longitudinal_obj ).value c.pyapi.decref(azimuthal_obj) c.pyapi.decref(longitudinal_obj) is_error = numba.core.cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) return numba.extending.NativeValue(proxyout._getvalue(), is_error) @numba.extending.box(VectorObject3DType) def VectorObject3D_box(typ, val, c): if isinstance(typ, MomentumObject3DType): cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(MomentumObject3D)) else: cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(VectorObject3D)) proxyin = numba.core.cgutils.create_struct_proxy(typ)( c.context, c.builder, value=val ) azimuthal_obj = c.pyapi.from_native_value(typ.azimuthaltype, proxyin.azimuthal) longitudinal_obj = c.pyapi.from_native_value( typ.longitudinaltype, proxyin.longitudinal ) output_obj = c.pyapi.call_function_objargs( cls_obj, (azimuthal_obj, longitudinal_obj) ) c.pyapi.decref(cls_obj) c.pyapi.decref(azimuthal_obj) c.pyapi.decref(longitudinal_obj) return output_obj # VectorObject4D (and momentum) ############################################### class VectorObject4DType(numba.types.Type): instance_class = VectorObject4D def __init__(self, azimuthaltype, longitudinaltype, temporaltype): super().__init__( name=f"VectorObject4DType({azimuthaltype}, {longitudinaltype}, {temporaltype})" ) self.azimuthaltype = azimuthaltype self.longitudinaltype = longitudinaltype self.temporaltype = temporaltype class MomentumObject4DType(VectorObject4DType): instance_class = MomentumObject4D def __init__(self, azimuthaltype, longitudinaltype, temporaltype): super().__init__(azimuthaltype, longitudinaltype, temporaltype) self.name = ( f"MomentumObject4DType({azimuthaltype}, {longitudinaltype}, {temporaltype})" ) @numba.extending.register_model(VectorObject4DType) @numba.extending.register_model(MomentumObject4DType) class VectorObject4DModel(numba.extending.models.StructModel): def __init__(self, dmm, fe_type): members = [ ("azimuthal", fe_type.azimuthaltype), ("longitudinal", fe_type.longitudinaltype), ("temporal", fe_type.temporaltype), ] super().__init__(dmm, fe_type, members) @numba.extending.typeof_impl.register(VectorObject4D) def VectorObject4D_typeof(val, c): if isinstance(val, MomentumObject4D): return MomentumObject4DType( numba.typeof(val.azimuthal), numba.typeof(val.longitudinal), numba.typeof(val.temporal), ) else: return VectorObject4DType( numba.typeof(val.azimuthal), numba.typeof(val.longitudinal), numba.typeof(val.temporal), ) numba.extending.make_attribute_wrapper(VectorObject4DType, "azimuthal", "azimuthal") numba.extending.make_attribute_wrapper( VectorObject4DType, "longitudinal", "longitudinal" ) numba.extending.make_attribute_wrapper(VectorObject4DType, "temporal", "temporal") numba.extending.make_attribute_wrapper(MomentumObject4DType, "azimuthal", "azimuthal") numba.extending.make_attribute_wrapper( MomentumObject4DType, "longitudinal", "longitudinal" ) numba.extending.make_attribute_wrapper(MomentumObject4DType, "temporal", "temporal") @numba.extending.type_callable(VectorObject4D) def VectorObject4D_constructor_typer(context): def typer(azimuthaltype, longitudinaltype, temporaltype): if ( is_azimuthaltype(azimuthaltype) and is_temporaltype(temporaltype) and is_temporaltype(temporaltype) ): return VectorObject4DType(azimuthaltype, longitudinaltype, temporaltype) else: raise numba.TypingError( "VectorObject4D constructor requires an AzimuthalObject and a " "LongitudinalObject and a TemporalObject as its arguments" ) return typer @numba.extending.type_callable(MomentumObject4D) def MomentumObject4D_constructor_typer(context): def typer(azimuthaltype, longitudinaltype, temporaltype): if ( is_azimuthaltype(azimuthaltype) and is_temporaltype(temporaltype) and is_temporaltype(temporaltype) ): return MomentumObject4DType(azimuthaltype, longitudinaltype, temporaltype) else: raise numba.TypingError( "MomentumObject4D constructor requires an AzimuthalObject and a " "LongitudinalObject and a TemporalObject as its arguments" ) return typer @numba.extending.lower_builtin( VectorObject4D, numba.types.Type, numba.types.Type, numba.types.Type ) @numba.extending.lower_builtin( MomentumObject4D, numba.types.Type, numba.types.Type, numba.types.Type ) def VectorObject4D_constructor_impl(context, builder, sig, args): typ = sig.return_type proxyout = numba.core.cgutils.create_struct_proxy(typ)(context, builder) proxyout.azimuthal = args[0] proxyout.longitudinal = args[1] proxyout.temporal = args[2] return proxyout._getvalue() @numba.extending.unbox(VectorObject4DType) def VectorObject4D_unbox(typ, obj, c): azimuthal_obj = c.pyapi.object_getattr_string(obj, "azimuthal") longitudinal_obj = c.pyapi.object_getattr_string(obj, "longitudinal") temporal_obj = c.pyapi.object_getattr_string(obj, "temporal") proxyout = numba.core.cgutils.create_struct_proxy(typ)(c.context, c.builder) proxyout.azimuthal = c.pyapi.to_native_value(typ.azimuthaltype, azimuthal_obj).value proxyout.longitudinal = c.pyapi.to_native_value( typ.longitudinaltype, longitudinal_obj ).value proxyout.temporal = c.pyapi.to_native_value(typ.temporaltype, temporal_obj).value c.pyapi.decref(azimuthal_obj) c.pyapi.decref(longitudinal_obj) c.pyapi.decref(temporal_obj) is_error = numba.core.cgutils.is_not_null(c.builder, c.pyapi.err_occurred()) return numba.extending.NativeValue(proxyout._getvalue(), is_error) @numba.extending.box(VectorObject4DType) def VectorObject4D_box(typ, val, c): if isinstance(typ, MomentumObject4DType): cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(MomentumObject4D)) else: cls_obj = c.pyapi.unserialize(c.pyapi.serialize_object(VectorObject4D)) proxyin = numba.core.cgutils.create_struct_proxy(typ)( c.context, c.builder, value=val ) azimuthal_obj = c.pyapi.from_native_value(typ.azimuthaltype, proxyin.azimuthal) longitudinal_obj = c.pyapi.from_native_value( typ.longitudinaltype, proxyin.longitudinal ) temporal_obj = c.pyapi.from_native_value(typ.temporaltype, proxyin.temporal) output_obj = c.pyapi.call_function_objargs( cls_obj, (azimuthal_obj, longitudinal_obj, temporal_obj) ) c.pyapi.decref(cls_obj) c.pyapi.decref(azimuthal_obj) c.pyapi.decref(longitudinal_obj) c.pyapi.decref(temporal_obj) return output_obj # vector.obj factory function ################################################# @numba.jit(nopython=True) def vector_obj_Azimuthal_xy(x, px, y, py, rho, pt, phi): return AzimuthalObjectXY(x, y) @numba.jit(nopython=True) def vector_obj_Azimuthal_xpy(x, px, y, py, rho, pt, phi): return AzimuthalObjectXY(x, py) @numba.jit(nopython=True) def vector_obj_Azimuthal_pxy(x, px, y, py, rho, pt, phi): return AzimuthalObjectXY(px, y) @numba.jit(nopython=True) def vector_obj_Azimuthal_pxpy(x, px, y, py, rho, pt, phi): return AzimuthalObjectXY(px, py) @numba.jit(nopython=True) def vector_obj_Azimuthal_rhophi(x, px, y, py, rho, pt, phi): return AzimuthalObjectRhoPhi(rho, phi) @numba.jit(nopython=True) def vector_obj_Azimuthal_ptphi(x, px, y, py, rho, pt, phi): return AzimuthalObjectRhoPhi(pt, phi) @numba.jit(nopython=True) def vector_obj_Longitudinal_z(z, pz, theta, eta): return LongitudinalObjectZ(z) @numba.jit(nopython=True) def vector_obj_Longitudinal_pz(z, pz, theta, eta): return LongitudinalObjectZ(pz) @numba.jit(nopython=True) def vector_obj_Longitudinal_theta(z, pz, theta, eta): return LongitudinalObjectTheta(theta) @numba.jit(nopython=True) def vector_obj_Longitudinal_eta(z, pz, theta, eta): return LongitudinalObjectEta(eta) @numba.jit(nopython=True) def vector_obj_Temporal_t(t, E, e, energy, tau, M, m, mass): return TemporalObjectT(t) @numba.jit(nopython=True) def vector_obj_Temporal_E(t, E, e, energy, tau, M, m, mass): return TemporalObjectT(E) @numba.jit(nopython=True) def vector_obj_Temporal_e(t, E, e, energy, tau, M, m, mass): return TemporalObjectT(e) @numba.jit(nopython=True) def vector_obj_Temporal_energy(t, E, e, energy, tau, M, m, mass): return TemporalObjectT(energy) @numba.jit(nopython=True) def vector_obj_Temporal_tau(t, E, e, energy, tau, M, m, mass): return TemporalObjectTau(tau) @numba.jit(nopython=True) def vector_obj_Temporal_M(t, E, e, energy, tau, M, m, mass): return TemporalObjectTau(M) @numba.jit(nopython=True) def vector_obj_Temporal_m(t, E, e, energy, tau, M, m, mass): return TemporalObjectTau(m) @numba.jit(nopython=True) def vector_obj_Temporal_mass(t, E, e, energy, tau, M, m, mass): return TemporalObjectTau(mass) @numba.extending.overload(vector.obj) def vector_obj( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): if unrecognized_argument is not None: raise numba.TypingError( "only keyword arguments are allowed in vector.obj; no positional arguments" ) has_x = x is not None has_px = px is not None has_y = y is not None has_py = py is not None has_rho = rho is not None has_pt = pt is not None has_phi = phi is not None has_z = z is not None has_pz = pz is not None has_theta = theta is not None has_eta = eta is not None has_t = t is not None has_E = E is not None has_e = e is not None has_energy = energy is not None has_tau = tau is not None has_M = M is not None has_m = m is not None has_mass = mass is not None is_momentum = False azimuthal = None longitudinal = None temporal = None if (has_x and has_y) and not (has_rho or has_pt or has_phi): if has_px or has_py: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): x/px or y/py" ) azimuthal = vector_obj_Azimuthal_xy elif (has_x and has_py) and not (has_rho or has_pt or has_phi): is_momentum = True if has_px or has_y: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): x/px or y/py" ) azimuthal = vector_obj_Azimuthal_xpy elif (has_px and has_y) and not (has_rho or has_pt or has_phi): is_momentum = True if has_x or has_py: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): x/px or y/py" ) azimuthal = vector_obj_Azimuthal_pxy elif (has_px and has_py) and not (has_rho or has_pt or has_phi): is_momentum = True if has_x or has_y: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): x/px or y/py" ) azimuthal = vector_obj_Azimuthal_pxpy elif (has_rho and has_phi) and not (has_x or has_px or has_y or has_py): if has_pt: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): rho/pt" ) azimuthal = vector_obj_Azimuthal_rhophi elif (has_pt and has_phi) and not (has_x or has_px or has_y or has_py): is_momentum = True if has_rho: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): rho/pt" ) azimuthal = vector_obj_Azimuthal_ptphi if has_z and not (has_theta or has_eta): if has_pz: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): z/pz" ) longitudinal = vector_obj_Longitudinal_z elif has_pz and not (has_theta or has_eta): is_momentum = True if has_z: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): z/pz" ) longitudinal = vector_obj_Longitudinal_pz elif has_theta and not (has_z or has_eta): longitudinal = vector_obj_Longitudinal_theta elif has_eta and not (has_z or has_theta): longitudinal = vector_obj_Longitudinal_eta if has_t and not (has_tau or has_M or has_m or has_mass): if has_E or has_e or has_energy: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): t/E/e/energy" ) temporal = vector_obj_Temporal_t elif has_E and not (has_tau or has_M or has_m or has_mass): is_momentum = True if has_t or has_e or has_energy: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): t/E/e/energy" ) temporal = vector_obj_Temporal_E elif has_e and not (has_tau or has_M or has_m or has_mass): is_momentum = True if has_t or has_E or has_energy: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): t/E/e/energy" ) temporal = vector_obj_Temporal_e elif has_energy and not (has_tau or has_M or has_m or has_mass): is_momentum = True if has_t or has_E or has_e: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): t/E/e/energy" ) temporal = vector_obj_Temporal_energy elif has_tau and not (has_t or has_E or has_e or has_energy): if has_M or has_m or has_mass: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): tau/M/m/mass" ) temporal = vector_obj_Temporal_tau elif has_M and not (has_t or has_E or has_e or has_energy): is_momentum = True if has_tau or has_m or has_mass: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): tau/M/m/mass" ) temporal = vector_obj_Temporal_M elif has_m and not (has_t or has_E or has_e or has_energy): is_momentum = True if has_tau or has_M or has_mass: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): tau/M/m/mass" ) temporal = vector_obj_Temporal_m elif has_mass and not (has_t or has_E or has_e or has_energy): is_momentum = True if has_tau or has_M or has_m: raise numba.TypingError( "duplicate coordinates (through momentum-aliases): tau/M/m/mass" ) temporal = vector_obj_Temporal_mass if azimuthal is not None and longitudinal is not None and temporal is not None: if is_momentum: def vector_obj_impl( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): return MomentumObject4D( azimuthal(x, px, y, py, rho, pt, phi), longitudinal(z, pz, theta, eta), temporal(t, E, e, energy, tau, M, m, mass), ) else: def vector_obj_impl( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): return VectorObject4D( azimuthal(x, px, y, py, rho, pt, phi), longitudinal(z, pz, theta, eta), temporal(t, E, e, energy, tau, M, m, mass), ) elif azimuthal is not None and longitudinal is not None: if is_momentum: def vector_obj_impl( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): return MomentumObject3D( azimuthal(x, px, y, py, rho, pt, phi), longitudinal(z, pz, theta, eta), ) else: def vector_obj_impl( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): return VectorObject3D( azimuthal(x, px, y, py, rho, pt, phi), longitudinal(z, pz, theta, eta), ) elif azimuthal is not None: if is_momentum: def vector_obj_impl( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): return MomentumObject2D(azimuthal(x, px, y, py, rho, pt, phi)) else: def vector_obj_impl( unrecognized_argument=None, x=None, px=None, y=None, py=None, rho=None, pt=None, phi=None, z=None, pz=None, theta=None, eta=None, t=None, E=None, e=None, energy=None, tau=None, M=None, m=None, mass=None, ): return VectorObject2D(azimuthal(x, px, y, py, rho, pt, phi)) else: raise numba.TypingError( "unrecognized combination of coordinates, allowed combinations are:\n\n" " (2D) x= y=\n" " (2D) rho= phi=\n" " (3D) x= y= z=\n" " (3D) x= y= theta=\n" " (3D) x= y= eta=\n" " (3D) rho= phi= z=\n" " (3D) rho= phi= theta=\n" " (3D) rho= phi= eta=\n" " (4D) x= y= z= t=\n" " (4D) x= y= z= tau=\n" " (4D) x= y= theta= t=\n" " (4D) x= y= theta= tau=\n" " (4D) x= y= eta= t=\n" " (4D) x= y= eta= tau=\n" " (4D) rho= phi= z= t=\n" " (4D) rho= phi= z= tau=\n" " (4D) rho= phi= theta= t=\n" " (4D) rho= phi= theta= tau=\n" " (4D) rho= phi= eta= t=\n" " (4D) rho= phi= eta= tau=" ) return vector_obj_impl # properties and methods ###################################################### @numba.extending.overload_method(VectorObject2DType, "to_Vector2D") def VectorObject2D_to_Vector2D(v): def VectorObject2D_to_Vector2D_impl(v): return v return VectorObject2D_to_Vector2D_impl @numba.extending.overload_method(VectorObject2DType, "to_Vector3D") def VectorObject2D_to_Vector3D(v): if issubclass(v.instance_class, Momentum): def VectorObject2D_to_Vector3D_impl(v): return MomentumObject3D(v.azimuthal, LongitudinalObjectZ(0.0)) else: def VectorObject2D_to_Vector3D_impl(v): return VectorObject3D(v.azimuthal, LongitudinalObjectZ(0.0)) return VectorObject2D_to_Vector3D_impl @numba.extending.overload_method(VectorObject2DType, "to_Vector4D") def VectorObject2D_to_Vector4D(v): if issubclass(v.instance_class, Momentum): def VectorObject2D_to_Vector4D_impl(v): return MomentumObject4D( v.azimuthal, LongitudinalObjectZ(0.0), TemporalObjectT(0.0) ) else: def VectorObject2D_to_Vector4D_impl(v): return VectorObject4D( v.azimuthal, LongitudinalObjectZ(0.0), TemporalObjectT(0.0) ) return VectorObject2D_to_Vector4D_impl @numba.extending.overload_method(VectorObject3DType, "to_Vector2D") def VectorObject3D_to_Vector2D(v): if issubclass(v.instance_class, Momentum): def VectorObject3D_to_Vector2D_impl(v): return MomentumObject2D(v.azimuthal) else: def VectorObject3D_to_Vector2D_impl(v): return VectorObject2D(v.azimuthal) return VectorObject3D_to_Vector2D_impl @numba.extending.overload_method(VectorObject3DType, "to_Vector3D") def VectorObject3D_to_Vector3D(v): def VectorObject3D_to_Vector3D_impl(v): return v return VectorObject3D_to_Vector3D_impl @numba.extending.overload_method(VectorObject3DType, "to_Vector4D") def VectorObject3D_to_Vector4D(v): if issubclass(v.instance_class, Momentum): def VectorObject3D_to_Vector4D_impl(v): return MomentumObject4D(v.azimuthal, v.longitudinal, TemporalObjectT(0.0)) else: def VectorObject3D_to_Vector4D_impl(v): return VectorObject4D(v.azimuthal, v.longitudinal, TemporalObjectT(0.0)) return VectorObject3D_to_Vector4D_impl @numba.extending.overload_method(VectorObject4DType, "to_Vector2D") def VectorObject4D_to_Vector2D(v): if issubclass(v.instance_class, Momentum): def VectorObject4D_to_Vector2D_impl(v): return MomentumObject2D(v.azimuthal) else: def VectorObject4D_to_Vector2D_impl(v): return VectorObject2D(v.azimuthal) return VectorObject4D_to_Vector2D_impl @numba.extending.overload_method(VectorObject4DType, "to_Vector3D") def VectorObject4D_to_Vector3D(v): if issubclass(v.instance_class, Momentum): def VectorObject4D_to_Vector3D_impl(v): return MomentumObject3D(v.azimuthal, v.longitudinal) else: def VectorObject4D_to_Vector3D_impl(v): return VectorObject3D(v.azimuthal, v.longitudinal) return VectorObject4D_to_Vector3D_impl @numba.extending.overload_method(VectorObject4DType, "to_Vector4D") def VectorObject4D_to_Vector4D(v): def VectorObject4D_to_Vector4D_impl(v): return v return VectorObject4D_to_Vector4D_impl @numba.jit(nopython=True) def make_AzimuthalObjectXY(v): return AzimuthalObjectXY(v.x, v.y) @numba.jit(nopython=True) def make_AzimuthalObjectRhoPhi(v): return AzimuthalObjectRhoPhi(v.rho, v.phi) @numba.jit(nopython=True) def make_LongitudinalObjectZ(v): return LongitudinalObjectZ(v.z) @numba.jit(nopython=True) def make_LongitudinalObjectZ_zero(v): return LongitudinalObjectZ(0.0) @numba.jit(nopython=True) def make_LongitudinalObjectTheta(v): return LongitudinalObjectTheta(v.theta) @numba.jit(nopython=True) def make_LongitudinalObjectTheta_zero(v): return LongitudinalObjectTheta(0.0) @numba.jit(nopython=True) def make_LongitudinalObjectEta(v): return LongitudinalObjectEta(v.eta) @numba.jit(nopython=True) def make_LongitudinalObjectEta_zero(v): return LongitudinalObjectEta(0.0) @numba.jit(nopython=True) def make_TemporalObjectT(v): return TemporalObjectT(v.t) @numba.jit(nopython=True) def make_TemporalObjectT_zero(v): return TemporalObjectT(0.0) @numba.jit(nopython=True) def make_TemporalObjectTau(v): return TemporalObjectTau(v.tau) @numba.jit(nopython=True) def make_TemporalObjectTau_zero(v): return TemporalObjectTau(0.0) make_coordobject = { (VectorObject2DType, AzimuthalObjectXY): make_AzimuthalObjectXY, (VectorObject3DType, AzimuthalObjectXY): make_AzimuthalObjectXY, (VectorObject4DType, AzimuthalObjectXY): make_AzimuthalObjectXY, (VectorObject2DType, AzimuthalObjectRhoPhi): make_AzimuthalObjectRhoPhi, (VectorObject3DType, AzimuthalObjectRhoPhi): make_AzimuthalObjectRhoPhi, (VectorObject4DType, AzimuthalObjectRhoPhi): make_AzimuthalObjectRhoPhi, (VectorObject2DType, LongitudinalObjectZ): make_LongitudinalObjectZ_zero, (VectorObject3DType, LongitudinalObjectZ): make_LongitudinalObjectZ, (VectorObject4DType, LongitudinalObjectZ): make_LongitudinalObjectZ, (VectorObject2DType, LongitudinalObjectTheta): make_LongitudinalObjectTheta_zero, (VectorObject3DType, LongitudinalObjectTheta): make_LongitudinalObjectTheta, (VectorObject4DType, LongitudinalObjectTheta): make_LongitudinalObjectTheta, (VectorObject2DType, LongitudinalObjectEta): make_LongitudinalObjectEta_zero, (VectorObject3DType, LongitudinalObjectEta): make_LongitudinalObjectEta, (VectorObject4DType, LongitudinalObjectEta): make_LongitudinalObjectEta, (VectorObject2DType, TemporalObjectT): make_TemporalObjectT_zero, (VectorObject3DType, TemporalObjectT): make_TemporalObjectT_zero, (VectorObject4DType, TemporalObjectT): make_TemporalObjectT, (VectorObject2DType, TemporalObjectTau): make_TemporalObjectTau_zero, (VectorObject3DType, TemporalObjectTau): make_TemporalObjectTau_zero, (VectorObject4DType, TemporalObjectTau): make_TemporalObjectTau, } def add_coordinate_change(vectortype, azcoordtype, lcoordtype, tcoordtype): methodname = "to_" if azcoordtype is AzimuthalObjectXY: methodname += "xy" elif azcoordtype is AzimuthalObjectRhoPhi: methodname += "rhophi" if lcoordtype is LongitudinalObjectZ: methodname += "z" if lcoordtype is LongitudinalObjectTheta: methodname += "theta" if lcoordtype is LongitudinalObjectEta: methodname += "eta" if tcoordtype is TemporalObjectT: methodname += "t" if tcoordtype is TemporalObjectTau: methodname += "tau" @numba.extending.overload_method(vectortype, methodname) def overloader(v): if lcoordtype is None and tcoordtype is None: azcoords = make_coordobject[vectortype, azcoordtype] if issubclass(v.instance_class, Momentum): def overloader_impl(v): return MomentumObject2D(azcoords(v)) else: def overloader_impl(v): return VectorObject2D(azcoords(v)) elif tcoordtype is None: azcoords = make_coordobject[vectortype, azcoordtype] lcoords = make_coordobject[vectortype, lcoordtype] if issubclass(v.instance_class, Momentum): def overloader_impl(v): return MomentumObject3D(azcoords(v), lcoords(v)) else: def overloader_impl(v): return VectorObject3D(azcoords(v), lcoords(v)) else: azcoords = make_coordobject[vectortype, azcoordtype] lcoords = make_coordobject[vectortype, lcoordtype] tcoords = make_coordobject[vectortype, tcoordtype] if issubclass(v.instance_class, Momentum): def overloader_impl(v): return MomentumObject4D(azcoords(v), lcoords(v), tcoords(v)) else: def overloader_impl(v): return VectorObject4D(azcoords(v), lcoords(v), tcoords(v)) return overloader_impl for vectortype in (VectorObject2DType, VectorObject3DType, VectorObject4DType): for azcoordtype in (AzimuthalObjectXY, AzimuthalObjectRhoPhi): for lcoordtype in ( None, LongitudinalObjectZ, LongitudinalObjectTheta, LongitudinalObjectEta, ): for tcoordtype in (None, TemporalObjectT, TemporalObjectTau): only_tcoordtype = lcoordtype is None and tcoordtype is not None if not only_tcoordtype: add_coordinate_change( vectortype, azcoordtype, lcoordtype, tcoordtype ) @numba.jit(nopython=True) def azimuthalxy_coord1(v): return v.azimuthal.x @numba.jit(nopython=True) def azimuthalxy_coord2(v): return v.azimuthal.y @numba.jit(nopython=True) def azimuthalrhophi_coord1(v): return v.azimuthal.rho @numba.jit(nopython=True) def azimuthalrhophi_coord2(v): return v.azimuthal.phi @numba.jit(nopython=True) def longitudinalz_coord1(v): return v.longitudinal.z @numba.jit(nopython=True) def longitudinaltheta_coord1(v): return v.longitudinal.theta @numba.jit(nopython=True) def longitudinaleta_coord1(v): return v.longitudinal.eta @numba.jit(nopython=True) def temporalt_coord1(v): return v.temporal.t @numba.jit(nopython=True) def temporaltau_coord1(v): return v.temporal.tau getcoord1 = { AzimuthalXY: azimuthalxy_coord1, AzimuthalRhoPhi: azimuthalrhophi_coord1, LongitudinalZ: longitudinalz_coord1, LongitudinalTheta: longitudinaltheta_coord1, LongitudinalEta: longitudinaleta_coord1, TemporalT: temporalt_coord1, TemporalTau: temporaltau_coord1, } getcoord2 = { AzimuthalXY: azimuthalxy_coord2, AzimuthalRhoPhi: azimuthalrhophi_coord2, } planar_properties = ["x", "y", "rho", "rho2", "phi"] spatial_properties = ["z", "theta", "eta", "costheta", "cottheta", "mag", "mag2"] lorentz_properties = ["t", "t2", "tau", "tau2", "beta", "gamma", "rapidity"] lorentz_momentum_properties = ["Et", "Et2", "Mt", "Mt2"] def add_planar_property(vectortype, propertyname): @numba.extending.overload_attribute(vectortype, propertyname) def overloader(v): function, *returns = _from_signature( propertyname, numba_modules["planar"][propertyname], (numba_aztype(v),) ) coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] def overloader_impl(v): return function(numpy, coord1(v), coord2(v)) return overloader_impl def add_spatial_property(vectortype, propertyname): @numba.extending.overload_attribute(vectortype, propertyname) def overloader(v): function, *returns = _from_signature( propertyname, numba_modules["spatial"][propertyname], (numba_aztype(v), numba_ltype(v)), ) coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] def overloader_impl(v): return function(numpy, coord1(v), coord2(v), coord3(v)) return overloader_impl def add_lorentz_property(vectortype, propertyname): @numba.extending.overload_attribute(vectortype, propertyname) def overloader(v): function, *returns = _from_signature( propertyname, numba_modules["lorentz"][propertyname], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] def overloader_impl(v): return function(numpy, coord1(v), coord2(v), coord3(v), coord4(v)) return overloader_impl for propertyname in planar_properties: for vectortype in (VectorObject2DType, VectorObject3DType, VectorObject4DType): add_planar_property(vectortype, propertyname) for propertyname in spatial_properties: for vectortype in (VectorObject3DType, VectorObject4DType): add_spatial_property(vectortype, propertyname) for propertyname in lorentz_properties: for vectortype in (VectorObject4DType,): add_lorentz_property(vectortype, propertyname) for propertyname in lorentz_momentum_properties: for vectortype in (MomentumObject4DType,): add_lorentz_property(vectortype, propertyname) planar_binary_methods = ["deltaphi"] spatial_binary_methods = ["deltaangle", "deltaeta", "deltaR", "deltaR2"] lorentz_binary_methods = ["deltaRapidityPhi", "deltaRapidityPhi2"] general_binary_methods = ["dot", "add", "subtract", "equal", "not_equal"] def dimension_of(v): if isinstance(v, VectorObject2DType): return 2 elif isinstance(v, VectorObject3DType): return 3 elif isinstance(v, VectorObject4DType): return 4 def flavor_of(v1, v2): if issubclass(v1.instance_class, Momentum) and issubclass( v2.instance_class, Momentum ): return v1.instance_class else: return v1.instance_class.GenericClass def add_binary_method(vectortype, gn, methodname): @numba.extending.overload_method(vectortype, methodname) def overloader(v1, v2): groupname = gn min_dimension = min(dimension_of(v1), dimension_of(v2)) if (methodname in ("equal", "not_equal")) and dimension_of(v1) != dimension_of( v2 ): raise numba.TypingError( f"{type(v1).__name__} and {type(v2).__name__} do not have the same dimension" ) if min_dimension == 2: if groupname is None: groupname = "planar" coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] elif min_dimension == 3: if groupname is None: groupname = "spatial" coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord13 = getcoord1[numba_ltype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] coord23 = getcoord1[numba_ltype(v2)] elif min_dimension == 4: if groupname is None: groupname = "lorentz" coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord13 = getcoord1[numba_ltype(v1)] coord14 = getcoord1[numba_ttype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] coord23 = getcoord1[numba_ltype(v2)] coord24 = getcoord1[numba_ttype(v2)] if groupname == "planar": signature = (numba_aztype(v1), numba_aztype(v2)) elif groupname == "spatial": signature = ( numba_aztype(v1), numba_ltype(v1), numba_aztype(v2), numba_ltype(v2), ) elif groupname == "lorentz": signature = ( numba_aztype(v1), numba_ltype(v1), numba_ttype(v1), numba_aztype(v2), numba_ltype(v2), numba_ttype(v2), ) function, *returns = _from_signature( groupname + "." + methodname, numba_modules[groupname][methodname], signature, ) if returns in ([bool], [float]): if groupname == "planar": def overloader_impl(v1, v2): return function( numpy, coord11(v1), coord12(v1), coord21(v2), coord22(v2) ) elif groupname == "spatial": def overloader_impl(v1, v2): return function( numpy, coord11(v1), coord12(v1), coord13(v1), coord21(v2), coord22(v2), coord23(v2), ) elif groupname == "lorentz": def overloader_impl(v1, v2): return function( numpy, coord11(v1), coord12(v1), coord13(v1), coord14(v1), coord21(v2), coord22(v2), coord23(v2), coord24(v2), ) elif groupname == "planar": instance_class = flavor_of(v1, v2).ProjectionClass2D azcoords = _coord_object_type[returns[0]] def overloader_impl(v1, v2): out1, out2 = function( numpy, coord11(v1), coord12(v1), coord21(v2), coord22(v2) ) return instance_class(azcoords(out1, out2)) elif groupname == "spatial": instance_class = flavor_of(v1, v2).ProjectionClass3D azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] def overloader_impl(v1, v2): out1, out2, out3 = function( numpy, coord11(v1), coord12(v1), coord13(v1), coord21(v2), coord22(v2), coord23(v2), ) return instance_class(azcoords(out1, out2), lcoords(out3)) elif groupname == "lorentz": instance_class = flavor_of(v1, v2).ProjectionClass4D azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] def overloader_impl(v1, v2): out1, out2, out3, out4 = function( numpy, coord11(v1), coord12(v1), coord13(v1), coord14(v1), coord21(v2), coord22(v2), coord23(v2), coord24(v2), ) return instance_class( azcoords(out1, out2), lcoords(out3), tcoords(out4) ) return overloader_impl for methodname in planar_binary_methods: for vectortype in (VectorObject2DType, VectorObject3DType, VectorObject4DType): add_binary_method(vectortype, "planar", methodname) for methodname in spatial_binary_methods: for vectortype in (VectorObject3DType, VectorObject4DType): add_binary_method(vectortype, "spatial", methodname) for methodname in lorentz_binary_methods: for vectortype in (VectorObject4DType,): add_binary_method(vectortype, "lorentz", methodname) for methodname in general_binary_methods: add_binary_method(VectorObject2DType, None, methodname) for methodname in general_binary_methods: add_binary_method(VectorObject3DType, None, methodname) for methodname in general_binary_methods: add_binary_method(VectorObject4DType, None, methodname) tolerance_methods = ["is_parallel", "is_antiparallel", "is_perpendicular"] def add_tolerance_method(vectortype, methodname): @numba.extending.overload_method(vectortype, methodname) def overloader(v1, v2, tolerance=1e-5): if isinstance(v1, VectorObject2DType) and not isinstance( v2, VectorObject2DType ): def overloader_impl(v1, v2, tolerance=1e-5): return v1.to_Vector3D().is_parallel(v2, tolerance=tolerance) return overloader_impl if not isinstance(v1, VectorObject2DType) and isinstance( v2, VectorObject2DType ): def overloader_impl(v1, v2, tolerance=1e-5): return v1.is_parallel(v2.to_Vector3D(), tolerance=tolerance) return overloader_impl if issubclass(vectortype, VectorObject2DType): groupname = "planar" signature = (numba_aztype(v1), numba_aztype(v2)) coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] elif issubclass(vectortype, (VectorObject3DType, VectorObject4DType)): groupname = "spatial" signature = ( numba_aztype(v1), numba_ltype(v1), numba_aztype(v2), numba_ltype(v2), ) coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord13 = getcoord1[numba_ltype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] coord23 = getcoord1[numba_ltype(v2)] function, *returns = _from_signature( groupname + "." + methodname, numba_modules[groupname][methodname], signature, ) if issubclass(vectortype, VectorObject2DType): def overloader_impl(v1, v2, tolerance=1e-5): return function( numpy, tolerance, coord11(v1), coord12(v1), coord21(v2), coord22(v2), ) elif issubclass(vectortype, (VectorObject3DType, VectorObject4DType)): def overloader_impl(v1, v2, tolerance=1e-5): return function( numpy, tolerance, coord11(v1), coord12(v1), coord13(v1), coord21(v2), coord22(v2), coord23(v2), ) return overloader_impl for methodname in tolerance_methods: add_tolerance_method(VectorObject2DType, methodname) for methodname in tolerance_methods: add_tolerance_method(VectorObject3DType, methodname) for methodname in tolerance_methods: add_tolerance_method(VectorObject4DType, methodname) def add_isclose_method(vectortype): @numba.extending.overload_method(vectortype, "isclose") def overloader(v1, v2, rtol=1e-05, atol=1e-08, equal_nan=False): if isinstance(v1, VectorObject2DType) and isinstance(v2, VectorObject2DType): groupname = "planar" signature = (numba_aztype(v1), numba_aztype(v2)) coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] elif isinstance(v1, VectorObject3DType) and isinstance(v2, VectorObject3DType): groupname = "spatial" signature = ( numba_aztype(v1), numba_ltype(v1), numba_aztype(v2), numba_ltype(v2), ) coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord13 = getcoord1[numba_ltype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] coord23 = getcoord1[numba_ltype(v2)] elif isinstance(v1, VectorObject4DType) and isinstance(v2, VectorObject4DType): groupname = "lorentz" signature = ( numba_aztype(v1), numba_ltype(v1), numba_ttype(v1), numba_aztype(v2), numba_ltype(v2), numba_ttype(v2), ) coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord13 = getcoord1[numba_ltype(v1)] coord14 = getcoord1[numba_ttype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] coord23 = getcoord1[numba_ltype(v2)] coord24 = getcoord1[numba_ttype(v2)] else: raise numba.TypingError( f"{type(v1).__name__} and {type(v2).__name__} do not have the same dimension" ) function, *returns = _from_signature( groupname + ".isclose", numba_modules[groupname]["isclose"], signature, ) if isinstance(v1, VectorObject2DType) and isinstance(v2, VectorObject2DType): def overloader_impl(v1, v2, rtol=1e-05, atol=1e-08, equal_nan=False): return function( numpy, rtol, atol, equal_nan, coord11(v1), coord12(v1), coord21(v2), coord22(v2), ) elif isinstance(v1, VectorObject3DType) and isinstance(v2, VectorObject3DType): def overloader_impl(v1, v2, rtol=1e-05, atol=1e-08, equal_nan=False): return function( numpy, rtol, atol, equal_nan, coord11(v1), coord12(v1), coord13(v1), coord21(v2), coord22(v2), coord23(v2), ) elif isinstance(v1, VectorObject4DType) and isinstance(v2, VectorObject4DType): def overloader_impl(v1, v2, rtol=1e-05, atol=1e-08, equal_nan=False): return function( numpy, rtol, atol, equal_nan, coord11(v1), coord12(v1), coord13(v1), coord14(v1), coord21(v2), coord22(v2), coord23(v2), coord24(v2), ) return overloader_impl add_isclose_method(VectorObject2DType) add_isclose_method(VectorObject3DType) add_isclose_method(VectorObject4DType) # the rest are special in one way or another ################################## def add_rotateZ(vectortype): @numba.extending.overload_method(vectortype, "rotateZ") def overloader(v, angle): if isinstance(angle, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["planar"]["rotateZ"], (numba_aztype(v),) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] azcoords = _coord_object_type[returns[0]] if issubclass(vectortype, VectorObject2DType): def overloader_impl(v, angle): out1, out2 = function(numpy, angle, coord1(v), coord2(v)) return instance_class(azcoords(out1, out2)) elif issubclass(vectortype, VectorObject3DType): def overloader_impl(v, angle): out1, out2 = function(numpy, angle, coord1(v), coord2(v)) return instance_class(azcoords(out1, out2), v.longitudinal) elif issubclass(vectortype, VectorObject4DType): def overloader_impl(v, angle): out1, out2 = function(numpy, angle, coord1(v), coord2(v)) return instance_class( azcoords(out1, out2), v.longitudinal, v.temporal ) return overloader_impl else: raise numba.TypingError( "'angle' must be an integer or a floating-point number" ) for vectortype in (VectorObject2DType, VectorObject3DType, VectorObject4DType): add_rotateZ(vectortype) def add_transform2D(vectortype): @numba.extending.overload_method(vectortype, "transform2D") def overloader(v, obj): function, *returns = _from_signature( "", numba_modules["planar"]["transform2D"], (numba_aztype(v),) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] azcoords = _coord_object_type[returns[0]] if issubclass(vectortype, VectorObject2DType): def overloader_impl(v, obj): out1, out2 = function( numpy, obj["xx"], obj["xy"], obj["yx"], obj["yy"], coord1(v), coord2(v), ) return instance_class(azcoords(out1, out2)) elif issubclass(vectortype, VectorObject3DType): def overloader_impl(v, obj): out1, out2 = function( numpy, obj["xx"], obj["xy"], obj["yx"], obj["yy"], coord1(v), coord2(v), ) return instance_class(azcoords(out1, out2), v.longitudinal) elif issubclass(vectortype, VectorObject4DType): def overloader_impl(v, obj): out1, out2 = function( numpy, obj["xx"], obj["xy"], obj["yx"], obj["yy"], coord1(v), coord2(v), ) return instance_class(azcoords(out1, out2), v.longitudinal, v.temporal) return overloader_impl for vectortype in (VectorObject2DType, VectorObject3DType, VectorObject4DType): add_transform2D(vectortype) @numba.extending.overload_method(VectorObject2DType, "unit") def VectorObject2DType_unit(v): function, *returns = _from_signature( "", numba_modules["planar"]["unit"], (numba_aztype(v),) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] azcoords = _coord_object_type[returns[0]] def VectorObject2DType_unit_impl(v): out1, out2 = function(numpy, coord1(v), coord2(v)) return instance_class(azcoords(out1, out2)) return VectorObject2DType_unit_impl @numba.extending.overload_method(VectorObject3DType, "unit") def VectorObject3DType_unit(v): function, *returns = _from_signature( "", numba_modules["spatial"]["unit"], (numba_aztype(v), numba_ltype(v)) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] def VectorObject3DType_unit_impl(v): out1, out2, out3 = function(numpy, coord1(v), coord2(v), coord3(v)) return instance_class(azcoords(out1, out2), lcoords(out3)) return VectorObject3DType_unit_impl @numba.extending.overload_method(VectorObject4DType, "unit") def VectorObject4DType_unit(v): function, *returns = _from_signature( "", numba_modules["lorentz"]["unit"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] def VectorObject4DType_unit_impl(v): out1, out2, out3, out4 = function( numpy, coord1(v), coord2(v), coord3(v), coord4(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), tcoords(out4)) return VectorObject4DType_unit_impl @numba.extending.overload_method(VectorObject2DType, "scale") @numba.extending.overload_method(VectorObject2DType, "scale2D") def VectorObject2DType_scale(v, factor): if isinstance(factor, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["planar"]["scale"], (numba_aztype(v),) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] azcoords = _coord_object_type[returns[0]] def VectorObject2DType_scale_impl(v, factor): out1, out2 = function(numpy, factor, coord1(v), coord2(v)) return instance_class(azcoords(out1, out2)) return VectorObject2DType_scale_impl else: raise numba.TypingError( "'factor' must be an integer or a floating-point number" ) @numba.extending.overload_method(VectorObject3DType, "scale") @numba.extending.overload_method(VectorObject3DType, "scale3D") def VectorObject3DType_scale(v, factor): if isinstance(factor, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["spatial"]["scale"], (numba_aztype(v), numba_ltype(v)) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] def VectorObject3DType_scale_impl(v, factor): out1, out2, out3 = function(numpy, factor, coord1(v), coord2(v), coord3(v)) return instance_class(azcoords(out1, out2), lcoords(out3)) return VectorObject3DType_scale_impl else: raise numba.TypingError( "'factor' must be an integer or a floating-point number" ) @numba.extending.overload_method(VectorObject4DType, "scale") @numba.extending.overload_method(VectorObject4DType, "scale4D") def VectorObject4DType_scale(v, factor): if isinstance(factor, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["lorentz"]["scale"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] def VectorObject4DType_scale_impl(v, factor): out1, out2, out3, out4 = function( numpy, factor, coord1(v), coord2(v), coord3(v), coord4(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), tcoords(out4)) return VectorObject4DType_scale_impl else: raise numba.TypingError( "'factor' must be an integer or a floating-point number" ) @numba.extending.overload_method(VectorObject3DType, "scale2D") @numba.extending.overload_method(VectorObject4DType, "scale2D") def VectorObject34DType_scale2D(v, factor): if isinstance(factor, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["planar"]["scale"], (numba_aztype(v),) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] azcoords = _coord_object_type[returns[0]] if isinstance(v, VectorObject3DType): def VectorObject34DType_scale2D_impl(v, factor): out1, out2 = function(numpy, factor, coord1(v), coord2(v)) return instance_class(azcoords(out1, out2), v.longitudinal) else: def VectorObject34DType_scale2D_impl(v, factor): out1, out2 = function(numpy, factor, coord1(v), coord2(v)) return instance_class(azcoords(out1, out2), v.longitudinal, v.temporal) return VectorObject34DType_scale2D_impl else: numba.TypingError("'factor' must be an integer or a floating-point number") @numba.extending.overload_method(VectorObject4DType, "scale3D") def VectorObject4DType_scale3D(v, factor): if isinstance(factor, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["spatial"]["scale"], (numba_aztype(v), numba_ltype(v)) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] def VectorObject4DType_scale3D_impl(v, factor): out1, out2, out3 = function(numpy, factor, coord1(v), coord2(v), coord3(v)) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject4DType_scale3D_impl else: numba.TypingError("'factor' must be an integer or a floating-point number") @numba.extending.overload_method(VectorObject3DType, "cross") @numba.extending.overload_method(VectorObject4DType, "cross") def VectorObject34DType_cross(v1, v2): if isinstance(v1, VectorObject3DType) and isinstance(v2, VectorObject3DType): function, *returns = _from_signature( "", numba_modules["spatial"]["cross"], (numba_aztype(v1), numba_ltype(v1), numba_aztype(v2), numba_ltype(v2)), ) instance_class = flavor_of(v1, v2).ProjectionClass3D coord11 = getcoord1[numba_aztype(v1)] coord12 = getcoord2[numba_aztype(v1)] coord13 = getcoord1[numba_ltype(v1)] coord21 = getcoord1[numba_aztype(v2)] coord22 = getcoord2[numba_aztype(v2)] coord23 = getcoord1[numba_ltype(v2)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] def VectorObject34DType_cross_impl(v1, v2): out1, out2, out3 = function( numpy, coord11(v1), coord12(v1), coord13(v1), coord21(v2), coord22(v2), coord23(v2), ) return instance_class(azcoords(out1, out2), lcoords(out3)) elif isinstance(v1, VectorObject3DType) and isinstance(v2, VectorObject4DType): def VectorObject34DType_cross_impl(v1, v2): return v1.cross(v2.to_Vector3D()) elif isinstance(v1, VectorObject4DType) and isinstance(v2, VectorObject3DType): def VectorObject34DType_cross_impl(v1, v2): return v1.to_Vector3D().cross(v2) elif isinstance(v1, VectorObject4DType) and isinstance(v2, VectorObject4DType): def VectorObject34DType_cross_impl(v1, v2): return v1.to_Vector3D().cross(v2.to_Vector3D()) else: raise numba.TypingError( "both vectors in 'cross' must be 3D or 4D (result is 3D)" ) return VectorObject34DType_cross_impl @numba.extending.overload_method(VectorObject3DType, "rotateX") @numba.extending.overload_method(VectorObject4DType, "rotateX") def VectorObject34DType_rotateX(v, angle): if isinstance(angle, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["spatial"]["rotateX"], (numba_aztype(v), numba_ltype(v)) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_rotateX_impl(v, angle): out1, out2, out3 = function( numpy, angle, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_rotateX_impl(v, angle): out1, out2, out3 = function( numpy, angle, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_rotateX_impl else: raise numba.TypingError("'angle' must be an integer or a floating-point number") @numba.extending.overload_method(VectorObject3DType, "rotateY") @numba.extending.overload_method(VectorObject4DType, "rotateY") def VectorObject34DType_rotateY(v, angle): if isinstance(angle, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["spatial"]["rotateY"], (numba_aztype(v), numba_ltype(v)) ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_rotateY_impl(v, angle): out1, out2, out3 = function( numpy, angle, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_rotateY_impl(v, angle): out1, out2, out3 = function( numpy, angle, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_rotateY_impl else: raise numba.TypingError("'angle' must be an integer or a floating-point number") @numba.extending.overload_method(VectorObject3DType, "rotate_axis") @numba.extending.overload_method(VectorObject4DType, "rotate_axis") def VectorObject34DType_rotate_axis(v, axis, angle): if isinstance(angle, (numba.types.Integer, numba.types.Float)): function, *returns = _from_signature( "", numba_modules["spatial"]["rotate_axis"], (numba_aztype(axis), numba_ltype(axis), numba_aztype(v), numba_ltype(v)), ) instance_class = v.instance_class coord11 = getcoord1[numba_aztype(axis)] coord12 = getcoord2[numba_aztype(axis)] coord13 = getcoord1[numba_ltype(axis)] coord21 = getcoord1[numba_aztype(v)] coord22 = getcoord2[numba_aztype(v)] coord23 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_rotate_axis(v, axis, angle): out1, out2, out3 = function( numpy, angle, coord11(axis), coord12(axis), coord13(axis), coord21(v), coord22(v), coord23(v), ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_rotate_axis(v, axis, angle): out1, out2, out3 = function( numpy, angle, coord11(axis), coord12(axis), coord13(axis), coord21(v), coord22(v), coord23(v), ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_rotate_axis else: raise numba.TypingError("'angle' must be an integer or a floating-point number") @numba.extending.overload_method(VectorObject3DType, "rotate_euler") @numba.extending.overload_method(VectorObject4DType, "rotate_euler") def VectorObject34DType_rotate_euler(v, phi, theta, psi, order="zxz"): if ( isinstance(phi, (numba.types.Integer, numba.types.Float)) and isinstance(theta, (numba.types.Integer, numba.types.Float)) and isinstance(psi, (numba.types.Integer, numba.types.Float)) ): if isinstance(order, str): pass elif isinstance(order, numba.types.StringLiteral): order = order.literal_value elif isinstance(order, (numba.types.UnicodeType, numba.types.Bytes)): raise numba.TypingError( "'order' argument for 'rotate_euler' must be a compile-time string" ) function, *returns = _from_signature( "", numba_modules["spatial"]["rotate_euler"], (numba_aztype(v), numba_ltype(v), order), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_rotate_axis_impl(v, phi, theta, psi, order="zxz"): out1, out2, out3 = function( numpy, phi, theta, psi, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_rotate_axis_impl(v, phi, theta, psi, order="zxz"): out1, out2, out3 = function( numpy, phi, theta, psi, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_rotate_axis_impl else: raise numba.TypingError( "'phi', 'theta', and 'psi' must be integers or floating-point numbers" ) @numba.extending.overload_method(VectorObject3DType, "rotate_nautical") @numba.extending.overload_method(VectorObject4DType, "rotate_nautical") def VectorObject34DType_rotate_nautical(v, yaw, pitch, roll): if ( isinstance(yaw, (numba.types.Integer, numba.types.Float)) and isinstance(pitch, (numba.types.Integer, numba.types.Float)) and isinstance(roll, (numba.types.Integer, numba.types.Float)) ): function, *returns = _from_signature( "", numba_modules["spatial"]["rotate_euler"], (numba_aztype(v), numba_ltype(v), "zyx"), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_rotate_nautical_impl(v, yaw, pitch, roll): out1, out2, out3 = function( numpy, roll, pitch, yaw, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_rotate_nautical_impl(v, yaw, pitch, roll): out1, out2, out3 = function( numpy, roll, pitch, yaw, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_rotate_nautical_impl else: raise numba.TypingError( "'yaw', 'pitch', and 'roll' must be integers or floating-point numbers" ) @numba.extending.overload_method(VectorObject3DType, "rotate_quaternion") @numba.extending.overload_method(VectorObject4DType, "rotate_quaternion") def VectorObject34DType_rotate_quaternion(v, u, i, j, k): if ( isinstance(u, (numba.types.Integer, numba.types.Float)) and isinstance(i, (numba.types.Integer, numba.types.Float)) and isinstance(j, (numba.types.Integer, numba.types.Float)) and isinstance(k, (numba.types.Integer, numba.types.Float)) ): function, *returns = _from_signature( "", numba_modules["spatial"]["rotate_quaternion"], (numba_aztype(v), numba_ltype(v)), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_rotate_quaternion_impl(v, u, i, j, k): out1, out2, out3 = function( numpy, u, i, j, k, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_rotate_quaternion_impl(v, u, i, j, k): out1, out2, out3 = function( numpy, u, i, j, k, coord1(v), coord2(v), coord3(v) ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_rotate_quaternion_impl else: raise numba.TypingError( "'u', 'i', 'j', and 'k' must be integers or floating-point numbers" ) @numba.extending.overload_method(VectorObject3DType, "transform3D") @numba.extending.overload_method(VectorObject4DType, "transform3D") def VectorObject34DType_transform3D(v, obj): function, *returns = _from_signature( "", numba_modules["spatial"]["transform3D"], (numba_aztype(v), numba_ltype(v)), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] if isinstance(v, VectorObject3DType): def VectorObject34DType_transform3D_impl(v, obj): out1, out2, out3 = function( numpy, obj["xx"], obj["xy"], obj["xz"], obj["yx"], obj["yy"], obj["yz"], obj["zx"], obj["zy"], obj["zz"], coord1(v), coord2(v), coord3(v), ) return instance_class(azcoords(out1, out2), lcoords(out3)) else: def VectorObject34DType_transform3D_impl(v, obj): out1, out2, out3 = function( numpy, obj["xx"], obj["xy"], obj["xz"], obj["yx"], obj["yy"], obj["yz"], obj["zx"], obj["zy"], obj["zz"], coord1(v), coord2(v), coord3(v), ) return instance_class(azcoords(out1, out2), lcoords(out3), v.temporal) return VectorObject34DType_transform3D_impl @numba.extending.overload_method(VectorObject4DType, "boost_p4") def VectorObject4DType_boost_p4(v, p4): function, *returns = _from_signature( "", numba_modules["lorentz"]["boost_p4"], ( numba_aztype(v), numba_ltype(v), numba_ttype(v), numba_aztype(p4), numba_ltype(p4), numba_ttype(p4), ), ) instance_class = v.instance_class coord11 = getcoord1[numba_aztype(v)] coord12 = getcoord2[numba_aztype(v)] coord13 = getcoord1[numba_ltype(v)] coord14 = getcoord1[numba_ttype(v)] coord21 = getcoord1[numba_aztype(p4)] coord22 = getcoord2[numba_aztype(p4)] coord23 = getcoord1[numba_ltype(p4)] coord24 = getcoord1[numba_ttype(p4)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] def VectorObject4DType_boost_p4_impl(v, p4): out1, out2, out3, out4 = function( numpy, coord11(v), coord12(v), coord13(v), coord14(v), coord21(p4), coord22(p4), coord23(p4), coord24(p4), ) return instance_class(azcoords(out1, out2), lcoords(out3), tcoords(out4)) return VectorObject4DType_boost_p4_impl @numba.extending.overload_method(VectorObject4DType, "boost_beta3") def VectorObject4DType_boost_beta3(v, beta3): function, *returns = _from_signature( "", numba_modules["lorentz"]["boost_beta3"], ( numba_aztype(v), numba_ltype(v), numba_ttype(v), numba_aztype(beta3), numba_ltype(beta3), ), ) instance_class = v.instance_class coord11 = getcoord1[numba_aztype(v)] coord12 = getcoord2[numba_aztype(v)] coord13 = getcoord1[numba_ltype(v)] coord14 = getcoord1[numba_ttype(v)] coord21 = getcoord1[numba_aztype(beta3)] coord22 = getcoord2[numba_aztype(beta3)] coord23 = getcoord1[numba_ltype(beta3)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] def VectorObject4DType_boost_beta3_impl(v, beta3): out1, out2, out3, out4 = function( numpy, coord11(v), coord12(v), coord13(v), coord14(v), coord21(beta3), coord22(beta3), coord23(beta3), ) return instance_class(azcoords(out1, out2), lcoords(out3), tcoords(out4)) return VectorObject4DType_boost_beta3_impl @numba.extending.overload_method(VectorObject4DType, "boost") def VectorObject4DType_boost(v, booster): if isinstance(booster, VectorObject3DType): def VectorObject4DType_boost_impl(v, booster): return v.boost_beta3(booster) elif isinstance(booster, VectorObject4DType): def VectorObject4DType_boost_impl(v, booster): return v.boost_p4(booster) else: raise numba.TypingError( "specify a Vector3D to boost by beta (velocity with c=1) or " "a Vector4D to boost by a momentum 4-vector" ) return VectorObject4DType_boost_impl @numba.extending.overload_method(VectorObject4DType, "boostCM_of_p4") def VectorObject4DType_boostCM_of_p4(v, p4): def VectorObject4DType_boostCM_of_p4_impl(v, p4): return v.boost_p4(p4.neg3D) return VectorObject4DType_boostCM_of_p4_impl @numba.extending.overload_method(VectorObject4DType, "boostCM_of_beta3") def VectorObject4DType_boostCM_of_beta3(v, beta3): def VectorObject4DType_boostCM_of_beta3_impl(v, beta3): return v.boost_beta3(beta3.neg3D) return VectorObject4DType_boostCM_of_beta3_impl @numba.extending.overload_method(VectorObject4DType, "boostCM_of") def VectorObject4DType_boostCM_of(v, booster): if isinstance(booster, VectorObject3DType): def VectorObject4DType_boostCM_of_impl(v, booster): return v.boost_beta3(booster.neg3D) elif isinstance(booster, VectorObject4DType): def VectorObject4DType_boostCM_of_impl(v, booster): return v.boost_p4(booster.neg3D) else: raise numba.TypingError( "specify a Vector3D to boost to the CM of beta (velocity with c=1) or " "a Vector4D to boost to the CM of a momentum 4-vector" ) return VectorObject4DType_boostCM_of_impl def add_boostXYZ(methodname): @numba.extending.overload_method(VectorObject4DType, methodname) def VectorObject4DType_boostXYZ(v, beta=None, gamma=None): if not isinstance( beta, (type(None), numba.types.NoneType, numba.types.Omitted) ) and isinstance( gamma, (type(None), numba.types.NoneType, numba.types.Omitted) ): if not isinstance(beta, (numba.types.Integer, numba.types.Float)): raise numba.TypingError( "'beta' must be an integer or a floating-point number" ) function, *returns = _from_signature( "", numba_modules["lorentz"][methodname + "_beta"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) elif isinstance( beta, (type(None), numba.types.NoneType, numba.types.Omitted) ) and not isinstance( gamma, (type(None), numba.types.NoneType, numba.types.Omitted) ): if not isinstance(gamma, (numba.types.Integer, numba.types.Float)): raise numba.TypingError( "'gamma' must be an integer or a floating-point number" ) function, *returns = _from_signature( "", numba_modules["lorentz"][methodname + "_gamma"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) else: raise numba.TypingError("specify 'beta' xor 'gamma', not both or neither") instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] if not isinstance( beta, (type(None), numba.types.NoneType, numba.types.Omitted) ): def VectorObject4DType_boostXYZ_impl(v, beta=None, gamma=None): out1, out2, out3, out4 = function( numpy, beta, coord1(v), coord2(v), coord3(v), coord4(v), ) return instance_class( azcoords(out1, out2), lcoords(out3), tcoords(out4) ) else: def VectorObject4DType_boostXYZ_impl(v, beta=None, gamma=None): out1, out2, out3, out4 = function( numpy, gamma, coord1(v), coord2(v), coord3(v), coord4(v), ) return instance_class( azcoords(out1, out2), lcoords(out3), tcoords(out4) ) return VectorObject4DType_boostXYZ_impl for methodname in "boostX", "boostY", "boostZ": add_boostXYZ(methodname) @numba.extending.overload_method(VectorObject4DType, "to_beta3") def VectorObject4DType_to_beta3(v): function, *returns = _from_signature( "", numba_modules["lorentz"]["to_beta3"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) instance_class = v.instance_class.ProjectionClass3D coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] def VectorObject4DType_to_beta3_impl(v): out1, out2, out3 = function( numpy, coord1(v), coord2(v), coord3(v), coord4(v), ) return instance_class(azcoords(out1, out2), lcoords(out3)) return VectorObject4DType_to_beta3_impl @numba.extending.overload_method(VectorObject4DType, "transform4D") def VectorObject4DType_transform4D(v, obj): function, *returns = _from_signature( "", numba_modules["lorentz"]["transform4D"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) instance_class = v.instance_class coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] azcoords = _coord_object_type[returns[0]] lcoords = _coord_object_type[returns[1]] tcoords = _coord_object_type[returns[2]] def VectorObject4DType_transform4D_impl(v, obj): out1, out2, out3, out4 = function( numpy, obj["xx"], obj["xy"], obj["xz"], obj["xt"], obj["yx"], obj["yy"], obj["yz"], obj["yt"], obj["zx"], obj["zy"], obj["zz"], obj["zt"], obj["tx"], obj["ty"], obj["tz"], obj["tt"], coord1(v), coord2(v), coord3(v), coord4(v), ) return instance_class(azcoords(out1, out2), lcoords(out3), tcoords(out4)) return VectorObject4DType_transform4D_impl @numba.extending.overload_method(VectorObject4DType, "is_timelike") def VectorObject4DType_is_timelike(v, tolerance=0): function, *returns = _from_signature( "", numba_modules["lorentz"]["is_timelike"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] def VectorObject4DType_is_timelike_impl(v, tolerance=0): return function( numpy, tolerance, coord1(v), coord2(v), coord3(v), coord4(v), ) return VectorObject4DType_is_timelike_impl @numba.extending.overload_method(VectorObject4DType, "is_spacelike") def VectorObject4DType_is_spacelike(v, tolerance=0): function, *returns = _from_signature( "", numba_modules["lorentz"]["is_spacelike"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] def VectorObject4DType_is_spacelike_impl(v, tolerance=0): return function( numpy, tolerance, coord1(v), coord2(v), coord3(v), coord4(v), ) return VectorObject4DType_is_spacelike_impl @numba.extending.overload_method(VectorObject4DType, "is_lightlike") def VectorObject4DType_is_lightlike(v, tolerance=1e-5): function, *returns = _from_signature( "", numba_modules["lorentz"]["is_lightlike"], (numba_aztype(v), numba_ltype(v), numba_ttype(v)), ) coord1 = getcoord1[numba_aztype(v)] coord2 = getcoord2[numba_aztype(v)] coord3 = getcoord1[numba_ltype(v)] coord4 = getcoord1[numba_ttype(v)] def VectorObject4DType_is_lightlike_impl(v, tolerance=1e-5): return function( numpy, tolerance, coord1(v), coord2(v), coord3(v), coord4(v), ) return VectorObject4DType_is_lightlike_impl @numba.extending.overload_attribute(VectorObject2DType, "neg2D") @numba.extending.overload_attribute(VectorObject3DType, "neg2D") @numba.extending.overload_attribute(VectorObject4DType, "neg2D") def VectorObject234DType_neg2D(v): def VectorObject234DType_neg2D_impl(v): return v.scale2D(-1) return VectorObject234DType_neg2D_impl @numba.extending.overload_attribute(VectorObject3DType, "neg3D") @numba.extending.overload_attribute(VectorObject4DType, "neg3D") def VectorObject34DType_neg3D(v): def VectorObject34DType_neg3D_impl(v): return v.scale3D(-1) return VectorObject34DType_neg3D_impl @numba.extending.overload_attribute(VectorObject4DType, "neg4D") def VectorObject4DType_neg4D(v): def VectorObject4DType_neg4D_impl(v): return v.scale4D(-1) return VectorObject4DType_neg4D_impl # momentum aliases ############################################################ @numba.extending.overload_attribute(MomentumObject2DType, "px") @numba.extending.overload_attribute(MomentumObject3DType, "px") @numba.extending.overload_attribute(MomentumObject4DType, "px") def MomentumObject234DType_px(v): def MomentumObject234DType_px_impl(v): return v.x return MomentumObject234DType_px_impl @numba.extending.overload_attribute(MomentumObject2DType, "py") @numba.extending.overload_attribute(MomentumObject3DType, "py") @numba.extending.overload_attribute(MomentumObject4DType, "py") def MomentumObject234DType_py(v): def MomentumObject234DType_py_impl(v): return v.y return MomentumObject234DType_py_impl @numba.extending.overload_attribute(MomentumObject2DType, "pt") @numba.extending.overload_attribute(MomentumObject3DType, "pt") @numba.extending.overload_attribute(MomentumObject4DType, "pt") def MomentumObject234DType_pt(v): def MomentumObject234DType_pt_impl(v): return v.rho return MomentumObject234DType_pt_impl @numba.extending.overload_attribute(MomentumObject2DType, "pt2") @numba.extending.overload_attribute(MomentumObject3DType, "pt2") @numba.extending.overload_attribute(MomentumObject4DType, "pt2") def MomentumObject234DType_pt2(v): def MomentumObject234DType_pt2_impl(v): return v.rho2 return MomentumObject234DType_pt2_impl @numba.extending.overload_attribute(MomentumObject3DType, "pz") @numba.extending.overload_attribute(MomentumObject4DType, "pz") def MomentumObject34DType_pz(v): def MomentumObject34DType_pz_impl(v): return v.z return MomentumObject34DType_pz_impl @numba.extending.overload_attribute(MomentumObject3DType, "pseudorapidity") @numba.extending.overload_attribute(MomentumObject4DType, "pseudorapidity") def MomentumObject34DType_pseudorapidity(v): def MomentumObject34DType_pseudorapidity_impl(v): return v.eta return MomentumObject34DType_pseudorapidity_impl @numba.extending.overload_attribute(MomentumObject3DType, "p") @numba.extending.overload_attribute(MomentumObject4DType, "p") def MomentumObject34DType_p(v): def MomentumObject34DType_p_impl(v): return v.mag return MomentumObject34DType_p_impl @numba.extending.overload_attribute(MomentumObject3DType, "p2") @numba.extending.overload_attribute(MomentumObject4DType, "p2") def MomentumObject34DType_p2(v): def MomentumObject34DType_p2_impl(v): return v.mag2 return MomentumObject34DType_p2_impl @numba.extending.overload_attribute(MomentumObject4DType, "E") def MomentumObject4DType_E(v): def MomentumObject4DType_E_impl(v): return v.t return MomentumObject4DType_E_impl @numba.extending.overload_attribute(MomentumObject4DType, "energy") def MomentumObject4DType_energy(v): def MomentumObject4DType_energy_impl(v): return v.t return MomentumObject4DType_energy_impl @numba.extending.overload_attribute(MomentumObject4DType, "E2") def MomentumObject4DType_E2(v): def MomentumObject4DType_E2_impl(v): return v.t2 return MomentumObject4DType_E2_impl @numba.extending.overload_attribute(MomentumObject4DType, "energy2") def MomentumObject4DType_energy2(v): def MomentumObject4DType_energy2_impl(v): return v.t2 return MomentumObject4DType_energy2_impl @numba.extending.overload_attribute(MomentumObject4DType, "M") def MomentumObject4DType_M(v): def MomentumObject4DType_M_impl(v): return v.tau return MomentumObject4DType_M_impl @numba.extending.overload_attribute(MomentumObject4DType, "mass") def MomentumObject4DType_mass(v): def MomentumObject4DType_mass_impl(v): return v.tau return MomentumObject4DType_mass_impl @numba.extending.overload_attribute(MomentumObject4DType, "M2") def MomentumObject4DType_M2(v): def MomentumObject4DType_M2_impl(v): return v.tau2 return MomentumObject4DType_M2_impl @numba.extending.overload_attribute(MomentumObject4DType, "mass2") def MomentumObject4DType_mass2(v): def MomentumObject4DType_mass2_impl(v): return v.tau2 return MomentumObject4DType_mass2_impl @numba.extending.overload_attribute(MomentumObject4DType, "transverse_energy") def MomentumObject4DType_transverse_energy(v): def MomentumObject4DType_transverse_energy_impl(v): return v.Et return MomentumObject4DType_transverse_energy_impl @numba.extending.overload_attribute(MomentumObject4DType, "transverse_energy2") def MomentumObject4DType_transverse_energy2(v): def MomentumObject4DType_transverse_energy2_impl(v): return v.Et2 return MomentumObject4DType_transverse_energy2_impl @numba.extending.overload_attribute(MomentumObject4DType, "transverse_mass") def MomentumObject4DType_transverse_mass(v): def MomentumObject4DType_transverse_mass_impl(v): return v.Mt return MomentumObject4DType_transverse_mass_impl @numba.extending.overload_attribute(MomentumObject4DType, "transverse_mass2") def MomentumObject4DType_transverse_mass2(v): def MomentumObject4DType_transverse_mass2_impl(v): return v.Mt2 return MomentumObject4DType_transverse_mass2_impl # NumPy functions in Numba #################################################### @numba.extending.overload(numpy.absolute) def numpy_absolute(v): if isinstance(v, VectorObject2DType): def numpy_absolute_impl(v): return v.rho return numpy_absolute_impl elif isinstance(v, VectorObject3DType): def numpy_absolute_impl(v): return v.mag return numpy_absolute_impl elif isinstance(v, VectorObject4DType): def numpy_absolute_impl(v): return v.tau return numpy_absolute_impl @numba.extending.overload(numpy.square) def numpy_square(v): if isinstance(v, VectorObject2DType): def numpy_square_impl(v): return v.rho2 return numpy_square_impl elif isinstance(v, VectorObject3DType): def numpy_square_impl(v): return v.mag2 return numpy_square_impl elif isinstance(v, VectorObject4DType): def numpy_square_impl(v): return v.tau2 return numpy_square_impl @numba.extending.overload(numpy.sqrt) def numpy_sqrt(v): if isinstance(v, VectorObject2DType): def numpy_sqrt_impl(v): return v.rho2**0.25 return numpy_sqrt_impl elif isinstance(v, VectorObject3DType): def numpy_sqrt_impl(v): return v.mag2**0.25 return numpy_sqrt_impl elif isinstance(v, VectorObject4DType): def numpy_sqrt_impl(v): return v.tau2**0.25 return numpy_sqrt_impl @numba.extending.overload(numpy.cbrt) def numpy_cbrt(v): if isinstance(v, VectorObject2DType): def numpy_cbrt_impl(v): return v.rho2**0.16666666666666666 return numpy_cbrt_impl elif isinstance(v, VectorObject3DType): def numpy_cbrt_impl(v): return v.mag2**0.16666666666666666 return numpy_cbrt_impl elif isinstance(v, VectorObject4DType): def numpy_cbrt_impl(v): return v.tau2**0.16666666666666666 return numpy_cbrt_impl @numba.extending.overload(numpy.power) def numpy_power(v, expo): if isinstance(v, VectorObject2DType) and isinstance( expo, (numba.types.Integer, numba.types.Float) ): def numpy_power_impl(v, expo): if expo == 2: return v.rho2 else: return v.rho**expo return numpy_power_impl elif isinstance(v, VectorObject3DType) and isinstance( expo, (numba.types.Integer, numba.types.Float) ): def numpy_power_impl(v, expo): if expo == 2: return v.mag2 else: return v.mag**expo return numpy_power_impl elif isinstance(v, VectorObject4DType) and isinstance( expo, (numba.types.Integer, numba.types.Float) ): def numpy_power_impl(v, expo): if expo == 2: return v.tau2 else: return v.tau**expo return numpy_power_impl @numba.extending.overload(numpy.add) def numpy_add(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def numpy_add_impl(v1, v2): return v1.add(v2) return numpy_add_impl @numba.extending.overload(numpy.subtract) def numpy_subtract(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def numpy_subtract_impl(v1, v2): return v1.subtract(v2) return numpy_subtract_impl @numba.extending.overload(numpy.multiply) def numpy_multiply(a, b): if isinstance( a, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(b, (numba.types.Integer, numba.types.Float)): def numpy_multiply_impl(a, b): return a.scale(b) return numpy_multiply_impl elif isinstance(a, (numba.types.Integer, numba.types.Float)) and isinstance( b, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ): def numpy_multiply_impl(a, b): return b.scale(a) return numpy_multiply_impl @numba.extending.overload(numpy.true_divide) def numpy_true_divide(a, b): if isinstance( a, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(b, (numba.types.Integer, numba.types.Float)): def numpy_true_divide_impl(a, b): return a.scale(1 / b) return numpy_true_divide_impl @numba.extending.overload(numpy.negative) def numpy_negative(v): if isinstance(v, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def numpy_negative_impl(v): return v.scale(-1) return numpy_negative_impl @numba.extending.overload(numpy.positive) def numpy_positive(v): if isinstance(v, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def numpy_positive_impl(v): return v return numpy_positive_impl @numba.extending.overload(numpy.matmul) def numpy_matmul(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def numpy_matmul_impl(v1, v2): return v1.dot(v2) return numpy_matmul_impl # unary operator overloading ################################################## @numba.extending.overload(abs) def operator_abs(v): if isinstance(v, VectorObject2DType): def operator_abs_impl(v): return v.rho return operator_abs_impl elif isinstance(v, VectorObject3DType): def operator_abs_impl(v): return v.mag return operator_abs_impl elif isinstance(v, VectorObject4DType): def operator_abs_impl(v): return v.tau return operator_abs_impl @numba.extending.overload(operator.neg) def operator_neg(v): if isinstance(v, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_neg_impl(v): return v.scale(-1) return operator_neg_impl @numba.extending.overload(operator.pos) def operator_pos(v): if isinstance(v, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_pos_impl(v): return v return operator_pos_impl @numba.extending.overload(bool) def operator_truth(v): if isinstance(v, VectorObject2DType): def operator_truth_impl(v): return v.rho2 != 0 return operator_truth_impl elif isinstance(v, VectorObject3DType): def operator_truth_impl(v): return v.mag2 != 0 return operator_truth_impl elif isinstance(v, VectorObject4DType) and issubclass( v.temporaltype.instance_class, TemporalT ): def operator_truth_impl(v): return v.mag2 != 0 or v.t != 0 return operator_truth_impl elif isinstance(v, VectorObject4DType) and issubclass( v.temporaltype.instance_class, TemporalTau ): def operator_truth_impl(v): return v.mag2 != 0 or v.tau != 0 return operator_truth_impl # binary operator overloading ################################################# @numba.extending.overload(operator.eq) def operator_eq(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_eq_impl(v1, v2): return v1.equal(v2) return operator_eq_impl @numba.extending.overload(operator.ne) def operator_ne(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_ne_impl(v1, v2): return v1.not_equal(v2) return operator_ne_impl @numba.extending.overload(operator.add) def operator_add(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_add_impl(v1, v2): return v1.add(v2) return operator_add_impl @numba.extending.overload(operator.sub) def operator_sub(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_sub_impl(v1, v2): return v1.subtract(v2) return operator_sub_impl @numba.extending.overload(operator.mul) def operator_mul(a, b): if isinstance( a, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and not isinstance( b, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ): def operator_mul_impl(a, b): return a.scale(b) return operator_mul_impl elif not isinstance( a, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance(b, (VectorObject2DType, VectorObject3DType, VectorObject4DType)): def operator_mul_impl(a, b): return b.scale(a) return operator_mul_impl @numba.extending.overload(operator.truediv) def operator_truediv(a, b): if isinstance( a, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and not isinstance( b, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ): def operator_truediv_impl(a, b): return a.scale(1.0 / b) return operator_truediv_impl @numba.extending.overload(operator.pow) def operator_pow(a, b): if ( isinstance(a, VectorObject2DType) and isinstance(b, numba.types.Literal) and b.literal_value == 2 ): def operator_pow_impl(a, b): return a.rho2 return operator_pow_impl elif ( isinstance(a, VectorObject3DType) and isinstance(b, numba.types.Literal) and b.literal_value == 2 ): def operator_pow_impl(a, b): return a.mag2 return operator_pow_impl elif ( isinstance(a, VectorObject4DType) and isinstance(b, numba.types.Literal) and b.literal_value == 2 ): def operator_pow_impl(a, b): return a.tau2 return operator_pow_impl elif isinstance( a, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and not isinstance( b, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ): def operator_pow_impl(a, b): return abs(a) ** b return operator_pow_impl if hasattr(operator, "matmul"): @numba.extending.overload(operator.matmul) def operator_matmul(v1, v2): if isinstance( v1, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ) and isinstance( v2, (VectorObject2DType, VectorObject3DType, VectorObject4DType) ): def operator_matmul_impl(v1, v2): return v1.dot(v2) return operator_matmul_impl # helper functions for Awkward Numba backend ################################## @numba.jit(nopython=True) def _awkward_numba_xy(record): return AzimuthalObjectXY(record["x"], record["y"]) @numba.jit(nopython=True) def _awkward_numba_pxy(record): return AzimuthalObjectXY(record["px"], record["y"]) @numba.jit(nopython=True) def _awkward_numba_xpy(record): return AzimuthalObjectXY(record["x"], record["py"]) @numba.jit(nopython=True) def _awkward_numba_pxpy(record): return AzimuthalObjectXY(record["px"], record["py"]) @numba.jit(nopython=True) def _awkward_numba_rhophi(record): return AzimuthalObjectRhoPhi(record["rho"], record["phi"]) @numba.jit(nopython=True) def _awkward_numba_ptphi(record): return AzimuthalObjectRhoPhi(record["pt"], record["phi"]) @numba.jit(nopython=True) def _awkward_numba_z(record): return LongitudinalObjectZ(record["z"]) @numba.jit(nopython=True) def _awkward_numba_pz(record): return LongitudinalObjectZ(record["pz"]) @numba.jit(nopython=True) def _awkward_numba_theta(record): return LongitudinalObjectTheta(record["theta"]) @numba.jit(nopython=True) def _awkward_numba_eta(record): return LongitudinalObjectEta(record["eta"]) @numba.jit(nopython=True) def _awkward_numba_t(record): return TemporalObjectT(record["t"]) @numba.jit(nopython=True) def _awkward_numba_E(record): return TemporalObjectT(record["E"]) @numba.jit(nopython=True) def _awkward_numba_e(record): return TemporalObjectT(record["e"]) @numba.jit(nopython=True) def _awkward_numba_energy(record): return TemporalObjectT(record["energy"]) @numba.jit(nopython=True) def _awkward_numba_tau(record): return TemporalObjectTau(record["tau"]) @numba.jit(nopython=True) def _awkward_numba_M(record): return TemporalObjectTau(record["M"]) @numba.jit(nopython=True) def _awkward_numba_m(record): return TemporalObjectTau(record["m"]) @numba.jit(nopython=True) def _awkward_numba_mass(record): return TemporalObjectTau(record["mass"]) vector-1.6.3/src/vector/backends/awkward.py000066400000000000000000002251251503546127100207130ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Defines behaviors for Awkward Array. New arrays created with the .. code-block:: python vector.Array(...) vector.awk(...) vector.zip(...) function will have these behaviors built in (and will pass them to any derived arrays). Alternatively, you can .. code-block:: python vector.register_awkward() to install the behaviors globally, so that any record named ``Vector2D``, ``Vector3D``, ``Vector4D``, ``Momentum2D``, ``Momentum3D``, or ``Momentum4D`` will have these properties and methods. The Awkward-Vectors-in-Numba extension is also implemented here, since it requires two non-strict dependencies of Vector: Awkward and Numba. Awkward's ``ak.behavior`` manages this non-strictness well. """ from __future__ import annotations import functools import numbers import operator import types import typing import awkward as ak import numpy import vector from vector._methods import ( Azimuthal, AzimuthalRhoPhi, AzimuthalXY, Longitudinal, LongitudinalEta, LongitudinalTheta, LongitudinalZ, Lorentz, LorentzMomentum, Momentum, Planar, PlanarMomentum, Spatial, SpatialMomentum, Temporal, TemporalT, TemporalTau, Vector2D, Vector3D, Vector4D, VectorProtocol, ) from vector._typeutils import BoolCollection, Protocol, ScalarCollection from vector.backends.numpy import VectorNumpy2D, VectorNumpy3D, VectorNumpy4D from vector.backends.object import ( AzimuthalObjectRhoPhi, AzimuthalObjectXY, LongitudinalObjectEta, LongitudinalObjectTheta, LongitudinalObjectZ, TemporalObjectT, TemporalObjectTau, VectorObject2D, VectorObject3D, VectorObject4D, ) # Throws an error if awkward is too old vector._import_awkward() ArrayOrRecord = typing.TypeVar("ArrayOrRecord", bound=typing.Union[ak.Array, ak.Record]) Array = typing.TypeVar("Array") behavior: typing.Any = {} def _touch(array: Array) -> Array: # make sure that touching is only done on Awkward arrays if isinstance(array, (ak.Array, ak.Record)) and ak.backend(array) == "typetracer": return ak.typetracer.touch_data(array) return array # coordinates classes are a formality for Awkward ############################# class CoordinatesAwkward: lib: types.ModuleType = numpy class AzimuthalAwkward(CoordinatesAwkward, Azimuthal): """ Azimuthal class for the Awkward backend. See - - :meth:`AzimuthalAwkward.from_fields` - :meth:`AzimuthalAwkward.from_momentum_fields` to construct azimuthal type objects. """ def __repr__(self) -> str: return f"{type(self).__name__}{self.elements}" @classmethod def from_fields(cls, array: ak.Array) -> AzimuthalAwkward: """ Create a :class:`vector.backends.awkward.AzimuthalAwkwardXY` or a :class:`vector.backends.awkward.AzimuthalAwkwardRhoPhi`, depending on the fields in ``array``. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"x": [1, 2]}, {"y": [1]}]) >>> az = vector.backends.awkward.AzimuthalAwkward.from_fields(a) >>> az AzimuthalAwkwardXY(, ) >>> az.elements (, ) """ fields = ak.fields(array) if "x" in fields and "y" in fields: return AzimuthalAwkwardXY(_touch(array["x"]), _touch(array["y"])) elif "rho" in fields and "phi" in fields: return AzimuthalAwkwardRhoPhi(_touch(array["rho"]), _touch(array["phi"])) else: raise ValueError( "array does not have azimuthal coordinates (x, y or rho, phi): " f"{', '.join(fields)}" ) @classmethod def from_momentum_fields(cls, array: ak.Array) -> AzimuthalAwkward: """ Create a :class:`vector.backends.awkward.AzimuthalAwkwardXY` or a :class:`vector.backends.awkward.AzimuthalAwkwardRhoPhi`, depending on the fields in ``array``, allowing momentum synonyms. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"px": [1, 2]}, {"py": [1]}]) >>> az = vector.backends.awkward.AzimuthalAwkward.from_momentum_fields(a) >>> az AzimuthalAwkwardXY(, ) >>> az.elements (, ) """ fields = ak.fields(array) if "x" in fields and "y" in fields: return AzimuthalAwkwardXY(_touch(array["x"]), _touch(array["y"])) elif "x" in fields and "py" in fields: return AzimuthalAwkwardXY(_touch(array["x"]), _touch(array["py"])) elif "px" in fields and "y" in fields: return AzimuthalAwkwardXY(_touch(array["px"]), _touch(array["y"])) elif "px" in fields and "py" in fields: return AzimuthalAwkwardXY(_touch(array["px"]), _touch(array["py"])) elif "rho" in fields and "phi" in fields: return AzimuthalAwkwardRhoPhi(_touch(array["rho"]), _touch(array["phi"])) elif "pt" in fields and "phi" in fields: return AzimuthalAwkwardRhoPhi(_touch(array["pt"]), _touch(array["phi"])) else: raise ValueError( "array does not have azimuthal coordinates (x/px, y/py or rho/pt, phi): " f"{', '.join(fields)}" ) class LongitudinalAwkward(CoordinatesAwkward, Longitudinal): """ Longitudinal class for the Awkward backend. See - - :meth:`LongitudinalAwkward.from_fields` - :meth:`LongitudinalAwkward.from_momentum_fields` to construct longitudinal type objects. """ def __repr__(self) -> str: return f"{type(self).__name__}{self.elements}" @classmethod def from_fields(cls, array: ak.Array) -> LongitudinalAwkward: """ Create a :class:`vector.backends.awkward.LongitudinalAwkwardZ`, a :class:`vector.backends.awkward.LongitudinalAwkwardTheta`, or a :class:`vector.backends.awkward.LongitudinalAwkwardEta`, depending on the fields in ``array``. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"theta": [1, 0]}]) >>> l = vector.backends.awkward.LongitudinalAwkward.from_fields(a) >>> l LongitudinalAwkwardTheta(,) >>> l.elements (,) """ fields = ak.fields(array) if "z" in fields: return LongitudinalAwkwardZ(_touch(array["z"])) elif "theta" in fields: return LongitudinalAwkwardTheta(_touch(array["theta"])) elif "eta" in fields: return LongitudinalAwkwardEta(_touch(array["eta"])) else: raise ValueError( "array does not have longitudinal coordinates (z or theta or eta): " f"{', '.join(fields)}" ) @classmethod def from_momentum_fields(cls, array: ak.Array) -> LongitudinalAwkward: """ Create a :class:`vector.backends.awkward.LongitudinalAwkwardZ`, a :class:`vector.backends.awkward.LongitudinalAwkwardTheta`, or a :class:`vector.backends.awkward.LongitudinalAwkwardEta`, depending on the fields in ``array``, allowing momentum synonyms. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"theta": [1, 0]}]) >>> l = vector.backends.awkward.LongitudinalAwkward.from_momentum_fields(a) >>> l LongitudinalAwkwardTheta(,) >>> l.elements (,) """ fields = ak.fields(array) if "z" in fields: return LongitudinalAwkwardZ(_touch(array["z"])) elif "pz" in fields: return LongitudinalAwkwardZ(_touch(array["pz"])) elif "theta" in fields: return LongitudinalAwkwardTheta(_touch(array["theta"])) elif "eta" in fields: return LongitudinalAwkwardEta(_touch(array["eta"])) else: raise ValueError( "array does not have longitudinal coordinates (z/pz or theta or eta): " f"{', '.join(fields)}" ) class TemporalAwkward(CoordinatesAwkward, Temporal): """ Temporal class for the Awkward backend. See - - :meth:`TemporalAwkward.from_fields` - :meth:`TemporalAwkward.from_momentum_fields` to construct longitudinal type objects. """ def __repr__(self) -> str: return f"{type(self).__name__}{self.elements}" @classmethod def from_fields(cls, array: ak.Array) -> TemporalAwkward: """ Create a :class:`vector.backends.awkward.TemporalT` or a :class:`vector.backends.awkward.TemporalTau`, depending on the fields in ``array``. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"tau": [1, 0]}]) >>> t = vector.backends.awkward.TemporalAwkward.from_fields(a) >>> t TemporalAwkwardTau(,) >>> t.elements (,) """ fields = ak.fields(array) if "t" in fields: return TemporalAwkwardT(_touch(array["t"])) elif "tau" in fields: return TemporalAwkwardTau(_touch(array["tau"])) else: raise ValueError( "array does not have temporal coordinates (t or tau): " f"{', '.join(fields)}" ) @classmethod def from_momentum_fields(cls, array: ak.Array) -> TemporalAwkward: """ Create a :class:`vector.backends.awkward.TemporalT` or a :class:`vector.backends.awkward.TemporalTau`, depending on the fields in ``array``, allowing momentum synonyms. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"mass": [1, 0]}]) >>> t = vector.backends.awkward.TemporalAwkward.from_momentum_fields(a) >>> t TemporalAwkwardTau(,) >>> t.elements (,) """ fields = ak.fields(array) if "t" in fields: return TemporalAwkwardT(_touch(array["t"])) elif "E" in fields: return TemporalAwkwardT(_touch(array["E"])) elif "e" in fields: return TemporalAwkwardT(_touch(array["e"])) elif "energy" in fields: return TemporalAwkwardT(_touch(array["energy"])) elif "tau" in fields: return TemporalAwkwardTau(_touch(array["tau"])) elif "M" in fields: return TemporalAwkwardTau(_touch(array["M"])) elif "m" in fields: return TemporalAwkwardTau(_touch(array["m"])) elif "mass" in fields: return TemporalAwkwardTau(_touch(array["mass"])) else: raise ValueError( "array does not have temporal coordinates (t/E/e/energy or tau/M/m/mass): " f"{', '.join(fields)}" ) class AzimuthalAwkwardXY(AzimuthalAwkward, AzimuthalXY): """ Class for the ``x`` and ``y`` (azimuthal) coordinates of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"x": [1, 2]}, {"y": [1]}]) >>> az = vector.backends.awkward.AzimuthalAwkwardXY(a["x"], a["y"]) >>> az AzimuthalAwkwardXY(, ) >>> az.elements (, ) """ __slots__ = ("x", "y") def __init__(self, x: typing.Any, y: typing.Any) -> None: self.x = x self.y = y @property def elements(self) -> tuple[ArrayOrRecord, ArrayOrRecord]: """ Azimuthal coordinates (``x`` and ``y``) as a tuple. Examples: >>> import vector >>> az = vector.backends.awkward.AzimuthalAwkwardXY([1, 2, 3], [1, 2]) >>> az.elements ([1, 2, 3], [1, 2]) """ return (self.x, self.y) class AzimuthalAwkwardRhoPhi(AzimuthalAwkward, AzimuthalRhoPhi): """ Class for the ``rho`` and ``phi`` (azimuthal) coordinates of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"rho": [1, 2]}, {"phi": [1]}]) >>> az = vector.backends.awkward.AzimuthalAwkwardRhoPhi(a["rho"], a["phi"]) >>> az AzimuthalAwkwardRhoPhi(, ) >>> az.elements (, ) """ __slots__ = ("phi", "rho") def __init__(self, rho: typing.Any, phi: typing.Any) -> None: self.rho = rho self.phi = phi @property def elements(self) -> tuple[ArrayOrRecord, ArrayOrRecord]: """ Azimuthal coordinates (``rho`` and ``phi``) as a tuple. Examples: >>> import vector >>> az = vector.backends.awkward.AzimuthalAwkwardRhoPhi([1, 2, 3], [1, 2]) >>> az.elements ([1, 2, 3], [1, 2]) """ return (self.rho, self.phi) class LongitudinalAwkwardZ(LongitudinalAwkward, LongitudinalZ): """ Class for the ``z`` (longitudinal) coordinate of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"z": [1, 2]}]) >>> l = vector.backends.awkward.LongitudinalAwkwardZ(a["z"]) >>> l LongitudinalAwkwardZ(,) >>> l.elements (,) """ __slots__ = ("z",) def __init__(self, z: typing.Any) -> None: self.z = z @property def elements(self) -> tuple[ArrayOrRecord]: """ Longitudinal coordinates (``z``) as a tuple. Examples: >>> import vector >>> l = vector.backends.awkward.LongitudinalAwkwardZ(5) >>> l.elements (5,) """ return (self.z,) class LongitudinalAwkwardTheta(LongitudinalAwkward, LongitudinalTheta): """ Class for the ``theta`` (longitudinal) coordinate of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"theta": [1, 2]}]) >>> l = vector.backends.awkward.LongitudinalAwkwardTheta(a["theta"]) >>> l LongitudinalAwkwardTheta(,) >>> l.elements (,) """ __slots__ = ("theta",) def __init__(self, theta: typing.Any) -> None: self.theta = theta @property def elements(self) -> tuple[ArrayOrRecord]: """ Longitudinal coordinates (``theta``) as a tuple. Examples: >>> import vector >>> l = vector.backends.awkward.LongitudinalAwkwardTheta(5) >>> l.elements (5,) """ return (self.theta,) class LongitudinalAwkwardEta(LongitudinalAwkward, LongitudinalEta): """ Class for the ``eta`` (longitudinal) coordinate of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"eta": [1, 2]}]) >>> l = vector.backends.awkward.LongitudinalAwkwardEta(a["eta"]) >>> l LongitudinalAwkwardEta(,) >>> l.elements (,) """ __slots__ = ("eta",) def __init__(self, eta: typing.Any) -> None: self.eta = eta @property def elements(self) -> tuple[ArrayOrRecord]: """ Longitudinal coordinates (``eta``) as a tuple. Examples: >>> import vector >>> l = vector.backends.awkward.LongitudinalAwkwardEta(5) >>> l.elements (5,) """ return (self.eta,) class TemporalAwkwardT(TemporalAwkward, TemporalT): """ Class for the ``t`` (temporal) coordinate of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"t": [1, 2]}]) >>> t = vector.backends.awkward.TemporalAwkwardT(a["t"]) >>> t TemporalAwkwardT(,) >>> t.elements (,) """ __slots__ = ("t",) def __init__(self, t: typing.Any) -> None: self.t = t @property def elements(self) -> tuple[ArrayOrRecord]: """ Temporal coordinates (``t``) as a tuple. Examples: >>> import vector >>> t = vector.backends.awkward.TemporalAwkwardT(5) >>> t.elements (5,) """ return (self.t,) class TemporalAwkwardTau(TemporalAwkward, TemporalTau): """ Class for the ``tau`` (temporal) coordinate of Awkward backend. Examples: >>> import vector >>> import awkward as ak >>> a = ak.Array([{"tau": [1, 2]}]) >>> t = vector.backends.awkward.TemporalAwkwardTau(a["tau"]) >>> t TemporalAwkwardTau(,) >>> t.elements (,) """ __slots__ = ("tau",) def __init__(self, tau: typing.Any) -> None: self.tau = tau @property def elements(self) -> tuple[ArrayOrRecord]: """ Temporal coordinates (``tau``) as a tuple. Examples: >>> import vector >>> t = vector.backends.awkward.TemporalAwkwardTau(5) >>> t.elements (5,) """ return (self.tau,) def _class_to_name(cls: type[VectorProtocol]) -> str: # respect the type of classes inheriting VectorAwkward classes is_vector = "vector.backends" in cls.__module__ if issubclass(cls, Momentum): if issubclass(cls, Vector2D): return "Momentum2D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector3D): return "Momentum3D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector4D): return "Momentum4D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector2D): return "Vector2D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector3D): return "Vector3D" if is_vector else cls.__name__[:-5] if issubclass(cls, Vector4D): return "Vector4D" if is_vector else cls.__name__[:-5] raise AssertionError(repr(cls)) # the vector class ############################################################ def _yes_record( x: ak.Array, ) -> float | ak.Record | None: return x[0] def _no_record(x: ak.Array) -> ak.Array | None: return x # Type for mixing in Awkward later class AwkwardProtocol(Protocol): def __getitem__(self, where: typing.Any) -> float | ak.Array | ak.Record | None: ... class _lib(typing.NamedTuple): # noqa: PLW1641 """a wrapper that respects the numpy-like interface of awkward-array and the module interface of numpy""" module: types.ModuleType nplike: ak._nplikes.numpy_like.NumpyLike def __eq__(self, other: typing.Any) -> bool: if isinstance(other, _lib): return self.module is other.module and self.nplike is other.nplike else: return self.module is other def __ne__(self, other: typing.Any) -> bool: return not self.__eq__(other) def __getattr__(self, name: str) -> typing.Any: if fun := getattr(self.nplike, name, False): return fun else: return getattr(self.module, name) class VectorAwkward: """Mixin class for Awkward vectors.""" @property def lib(self): # type:ignore[no-untyped-def] if ( nplike := self.layout.backend.nplike # type:ignore[attr-defined] ) is ak._nplikes.typetracer.TypeTracer.instance(): return _lib(module=numpy, nplike=nplike) return numpy def _wrap_result( self: AwkwardProtocol, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as an array of scalars or an array of vectors. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result if all(not isinstance(x, ak.Array) for x in result): maybe_record = _yes_record result = [ ak.Array(x.layout.array[x.layout.at : x.layout.at + 1]) if isinstance(x, ak.Record) else ak.Array([x]) for x in result ] else: maybe_record = _no_record if ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): first = next(x for x in result if isinstance(x, ak.Array)) result = [ x if isinstance(x, ak.Array) else ak.broadcast_arrays(first, x)[1] for x in result ] names = [] arrays = [] if returns[0] is AzimuthalXY: names.extend(["x", "y"]) arrays.extend([result[0], result[1]]) elif returns[0] is AzimuthalRhoPhi: names.extend(["rho", "phi"]) arrays.extend([result[0], result[1]]) fields = ak.fields(self) if num_vecargs == 1: for name in fields: if name not in ( "x", "y", "rho", "pt", "phi", ): names.append(name) arrays.append(self[name]) if "t" in fields or "tau" in fields: cls = cls.ProjectionClass4D elif "z" in fields or "theta" in fields or "eta" in fields: cls = cls.ProjectionClass3D else: cls = cls.ProjectionClass2D return maybe_record( ak.zip( dict(zip(names, arrays)), depth_limit=first.layout.purelist_depth, with_name=_class_to_name(cls), behavior=None if vector._awkward_registered else first.behavior, ) ) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): first = next(x for x in result if isinstance(x, ak.Array)) result = [ x if isinstance(x, ak.Array) else ak.broadcast_arrays(first, x)[1] for x in result ] names = [] arrays = [] if returns[0] is AzimuthalXY: names.extend(["x", "y"]) arrays.extend([result[0], result[1]]) elif returns[0] is AzimuthalRhoPhi: names.extend(["rho", "phi"]) arrays.extend([result[0], result[1]]) if num_vecargs == 1: for name in ak.fields(self): if name not in ( "x", "y", "rho", "pt", "phi", "z", "pz", "theta", "eta", "t", "tau", "m", "M", "mass", "e", "E", "energy", ): names.append(name) arrays.append(self[name]) return maybe_record( ak.zip( dict(zip(names, arrays)), depth_limit=first.layout.purelist_depth, with_name=_class_to_name(cls.ProjectionClass2D), behavior=None if vector._awkward_registered else first.behavior, ) ) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): first = next(x for x in result if isinstance(x, ak.Array)) result = [ x if isinstance(x, ak.Array) else ak.broadcast_arrays(first, x)[1] for x in result ] names = [] arrays = [] if returns[0] is AzimuthalXY: names.extend(["x", "y"]) arrays.extend([result[0], result[1]]) elif returns[0] is AzimuthalRhoPhi: names.extend(["rho", "phi"]) arrays.extend([result[0], result[1]]) if returns[1] is LongitudinalZ: names.append("z") arrays.append(result[2]) elif returns[1] is LongitudinalTheta: names.append("theta") arrays.append(result[2]) elif returns[1] is LongitudinalEta: names.append("eta") arrays.append(result[2]) fields = ak.fields(self) if num_vecargs == 1: for name in fields: if name not in ( "x", "y", "rho", "pt", "phi", "z", "pz", "theta", "eta", ): names.append(name) arrays.append(self[name]) if "t" in fields or "tau" in fields: cls = cls.ProjectionClass4D else: cls = cls.ProjectionClass3D return maybe_record( ak.zip( dict(zip(names, arrays)), depth_limit=first.layout.purelist_depth, with_name=_class_to_name(cls), behavior=None if vector._awkward_registered else first.behavior, ) ) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and returns[2] is None ): first = next(x for x in result if isinstance(x, ak.Array)) result = [ x if isinstance(x, ak.Array) else ak.broadcast_arrays(first, x)[1] for x in result ] names = [] arrays = [] if returns[0] is AzimuthalXY: names.extend(["x", "y"]) arrays.extend([result[0], result[1]]) elif returns[0] is AzimuthalRhoPhi: names.extend(["rho", "phi"]) arrays.extend([result[0], result[1]]) if returns[1] is LongitudinalZ: names.append("z") arrays.append(result[2]) elif returns[1] is LongitudinalTheta: names.append("theta") arrays.append(result[2]) elif returns[1] is LongitudinalEta: names.append("eta") arrays.append(result[2]) if num_vecargs == 1: for name in ak.fields(self): if name not in ( "x", "y", "rho", "pt", "phi", "z", "pz", "theta", "eta", "t", "tau", "m", "M", "mass", "e", "E", "energy", ): names.append(name) arrays.append(self[name]) return maybe_record( ak.zip( dict(zip(names, arrays)), depth_limit=first.layout.purelist_depth, with_name=_class_to_name(cls.ProjectionClass3D), behavior=None if vector._awkward_registered else first.behavior, ) ) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): first = next(x for x in result if isinstance(x, ak.Array)) result = [ x if isinstance(x, ak.Array) else ak.broadcast_arrays(first, x)[1] for x in result ] names = [] arrays = [] if returns[0] is AzimuthalXY: names.extend(["x", "y"]) arrays.extend([result[0], result[1]]) elif returns[0] is AzimuthalRhoPhi: names.extend(["rho", "phi"]) arrays.extend([result[0], result[1]]) if returns[1] is LongitudinalZ: names.append("z") arrays.append(result[2]) elif returns[1] is LongitudinalTheta: names.append("theta") arrays.append(result[2]) elif returns[1] is LongitudinalEta: names.append("eta") arrays.append(result[2]) if returns[2] is TemporalT: names.append("t") arrays.append(result[3]) elif returns[2] is TemporalTau: names.append("tau") arrays.append(result[3]) if num_vecargs == 1: for name in ak.fields(self): if name not in ( "x", "y", "rho", "pt", "phi", "z", "pz", "theta", "eta", "t", "tau", "m", "M", "mass", "e", "E", "energy", ): names.append(name) arrays.append(self[name]) return maybe_record( ak.zip( dict(zip(names, arrays)), depth_limit=first.layout.purelist_depth, with_name=_class_to_name(cls.ProjectionClass4D), behavior=None if vector._awkward_registered else first.behavior, ) ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function( self: AwkwardProtocol, func: typing.Callable, # type: ignore[type-arg] ) -> typing.Callable: # type: ignore[type-arg] return awkward_transform(func) _placeholder = object() class bind(functools.partial): # type: ignore[type-arg] def __call__(self, *args: typing.Any, **keywords: typing.Any) -> typing.Any: keywords = {**self.keywords, **keywords} iargs = iter(args) args = tuple(next(iargs) if arg is _placeholder else arg for arg in self.args) return self.func(*args, *iargs, **keywords) class awkward_transform: """Apply a batch of transformations in a single layout traversal""" def __init__(self, func: typing.Callable) -> None: # type: ignore[type-arg] self.func = func def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Callable: # type: ignore[type-arg] # '__awkward_transform_allowed__' is a flag that can be set to False to disable this wrapping, e.g. for no-ops if not getattr(self.func, "__awkward_transform_allowed__", True): return self.func(*args, **kwargs) # only pos args assert not kwargs # prepare the function and its args; we're currently assuming non-nested input args args2bind, awkward_arrays = [], [] n_orig_akarrays = 0 for arg in args: if isinstance(arg, (ak.highlevel.Array, ak.highlevel.Record)): n_orig_akarrays += 1 awkward_arrays.append(arg) args2bind.append(_placeholder) elif ak._regularize.is_array_like(arg): awkward_arrays.append(ak.Array(ak.to_layout(arg))) args2bind.append(_placeholder) else: args2bind.append(arg) # this means we're working with awkward-arrays and we should group operations with ak.transform if n_orig_akarrays > 0: def transformer( layouts: ak.contents.Content | tuple[ak.contents.Content, ...], **kwargs: typing.Any, ) -> ak.contents.Content: if not isinstance(layouts, tuple): layouts = (layouts,) if all(layout.is_numpy for layout in layouts): # setup parameter propagation, taken from awkward's broadcasting machinery options = kwargs["options"] rule = options.pop("broadcast_parameters_rule", None) try: parameters_factory = ( ak._broadcasting.BROADCAST_RULE_TO_FACTORY_IMPL[rule] ) except KeyError: raise ValueError( f"`broadcast_parameters_rule` should be one of {[str(x) for x in ak._broadcasting.BroadcastParameterRule]}, " f"but this routine received `{rule}`" ) from None # apply the function to the numpy arrays, first we need to 'partial out' all non-awkward array arguments out_numpys = bind(self.func, *args2bind)( *(map(operator.attrgetter("data"), layouts)) ) # if the function returns a single array, wrap it in a tuple if not isinstance(out_numpys, tuple): out_numpys = (out_numpys,) # propagate parameters out_params = parameters_factory( tuple(map(operator.attrgetter("parameters"), layouts)), len(out_numpys), ) # wrap the numpy arrays in awkward arrays out_arrays = tuple( ak.contents.NumpyArray(data=data, parameters=parameters) for data, parameters in zip(out_numpys, out_params) ) # if the function returns a single array, return it directly if len(out_arrays) == 1: return out_arrays[0] # return the awkward array(s) return out_arrays # apply to all arrays return ak.transform(transformer, *awkward_arrays) # if we arrived here this is because we're not working with awkward arrays (or have awkward-array installed). # that also means that if we call nested functions that are all decorated with this decorator, we'll end up here, # which is great, because that still means only one awkward broadcasting traversal! return self.func(*args, **kwargs) class VectorAwkward2D(VectorAwkward, Planar, Vector2D): """ Two dimensional vector class for the Awkward backend. Two dimensional awkward vectors for the users are defined using the :class:`VectorArray2D` class. See :class:`MomentumAwkward2D` for momentum vectors. """ @property def azimuthal(self) -> AzimuthalAwkward: """ Returns an Azimuthal type object. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], ... ) >>> a.azimuthal AzimuthalAwkwardXY(, ) >>> a.azimuthal.elements (, ) """ return AzimuthalAwkward.from_fields(self) class MomentumAwkward2D(PlanarMomentum, VectorAwkward2D): """ Two dimensional momentum vector class for the Awkward backend. Two dimensional momentum vectors for the users are defined using the :class:`MomentumArray2D` class. See :class:`VectorAwkward2D` for momentum vectors. """ @property def azimuthal(self) -> AzimuthalAwkward: """ Returns an Azimuthal type object containing the azimuthal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"px": 1, "py": 1.1}, {"px": 2, "py": 2.2}], ... ) >>> a.azimuthal AzimuthalAwkwardXY(, ) >>> a.azimuthal.elements (, ) """ return AzimuthalAwkward.from_momentum_fields(self) class VectorAwkward3D(VectorAwkward, Spatial, Vector3D): """ Three dimensional vector class for the Awkward backend. Three dimensional awkward vectors for the users are defined using the :class:`VectorArray3D` class. See :class:`MomentumAwkward3D` for momentum vectors. """ @property def azimuthal(self) -> AzimuthalAwkward: """ Returns an Azimuthal type object containing the azimuthal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}], ... ) >>> a.azimuthal AzimuthalAwkwardXY(, ) >>> a.azimuthal.elements (, ) """ return AzimuthalAwkward.from_fields(self) @property def longitudinal(self) -> LongitudinalAwkward: """ Returns a Longitudinal type object containing the longitudinal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}], ... ) >>> a.longitudinal LongitudinalAwkwardZ(,) >>> a.longitudinal.elements (,) """ return LongitudinalAwkward.from_fields(self) class MomentumAwkward3D(SpatialMomentum, VectorAwkward3D): """ Three dimensional momentum vector class for the Awkward backend. Three dimensional momentum vectors for the users are defined using the :class:`MomentumArray3D` class. See :class:`VectorAwkward3D` for momentum vectors. """ @property def azimuthal(self) -> AzimuthalAwkward: """ Returns an Azimuthal type object containing the azimuthal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"px": 1, "py": 1.1, "z": 0.1}, {"px": 2, "py": 2.2, "z": 0.2}], ... ) >>> a.azimuthal AzimuthalAwkwardXY(, ) >>> a.azimuthal.elements (, ) """ return AzimuthalAwkward.from_momentum_fields(self) @property def longitudinal(self) -> LongitudinalAwkward: """ Returns a Longitudinal type object containing the longitudinal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"px": 1, "py": 1.1, "z": 0.1}, {"px": 2, "py": 2.2, "z": 0.2}], ... ) >>> a.longitudinal LongitudinalAwkwardZ(,) >>> a.longitudinal.elements (,) """ return LongitudinalAwkward.from_momentum_fields(self) class VectorAwkward4D(VectorAwkward, Lorentz, Vector4D): """ Four dimensional momentum vector class for the Awkward backend. Four dimensional momentum vectors for the users are defined using the :class:`MomentumArray4D` class. See :class:`VectorAwkward4D` for momentum vectors. """ @property def azimuthal(self) -> AzimuthalAwkward: """ Returns an Azimuthal type object containing the azimuthal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"x": 1, "y": 1.1, "z": 0.1, "tau": 1}, {"x": 2, "y": 2.2, "z": 0.2, "tau": 3}], ... ) >>> a.azimuthal AzimuthalAwkwardXY(, ) >>> a.azimuthal.elements (, ) """ return AzimuthalAwkward.from_fields(self) @property def longitudinal(self) -> LongitudinalAwkward: """ Returns a Longitudinal type object containing the longitudinal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"x": 1, "y": 1.1, "z": 0.1, "tau": 1}, {"x": 2, "y": 2.2, "z": 0.2, "tau": 3}], ... ) >>> a.longitudinal LongitudinalAwkwardZ(,) >>> a.longitudinal.elements (,) """ return LongitudinalAwkward.from_fields(self) @property def temporal(self) -> TemporalAwkward: """ Returns a Temporal type object containing the temporal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"x": 1, "y": 1.1, "z": 0.1, "tau": 1}, {"x": 2, "y": 2.2, "z": 0.2, "tau": 3}], ... ) >>> a.temporal TemporalAwkwardTau(,) >>> a.temporal.elements (,) """ return TemporalAwkward.from_fields(self) class MomentumAwkward4D(LorentzMomentum, VectorAwkward4D): """ Four dimensional momentum vector class for the Awkward backend. Four dimensional momentum vectors for the users are defined using the :class:`MomentumArray4D` class. See :class:`VectorAwkward4D` for momentum vectors. """ @property def azimuthal(self) -> AzimuthalAwkward: """ Returns an Azimuthal type object containing the azimuthal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"px": 1, "py": 1.1, "z": 0.1, "m": 1}, {"px": 2, "py": 2.2, "z": 0.2, "m": 3}], ... ) >>> a.azimuthal AzimuthalAwkwardXY(, ) >>> a.azimuthal.elements (, ) """ return AzimuthalAwkward.from_momentum_fields(self) @property def longitudinal(self) -> LongitudinalAwkward: """ Returns a Longitudinal type object containing the longitudinal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"px": 1, "py": 1.1, "z": 0.1, "m": 1}, {"px": 2, "py": 2.2, "z": 0.2, "m": 3}], ... ) >>> a.longitudinal LongitudinalAwkwardZ(,) >>> a.longitudinal.elements (,) """ return LongitudinalAwkward.from_momentum_fields(self) @property def temporal(self) -> TemporalAwkward: """ Returns a Temporal type object containing the temporal coordinates. Use the ``elements`` property of the created object to access the coordinates. Examples: >>> import vector >>> a = vector.Array( ... [{"px": 1, "py": 1.1, "z": 0.1, "m": 1}, {"px": 2, "py": 2.2, "z": 0.2, "m": 3}], ... ) >>> a.temporal TemporalAwkwardTau(,) >>> a.temporal.elements (,) """ return TemporalAwkward.from_momentum_fields(self) # ak.Array and ak.Record subclasses ########################################### class VectorArray2D(VectorAwkward2D, ak.Array): # type: ignore[misc] """ Defines ``awkward`` behavior for a 2 dimensional vector. See :class:`VectorRecord2D` and :class:`MomentumArray2D` for the corresponding ``Record`` and ``Momentum`` classes. """ def allclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: """Like ``np.ndarray.allclose``, but for VectorArray2D.""" return ak.all(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan)) behavior["*", "Vector2D"] = VectorArray2D class VectorRecord2D(VectorAwkward2D, ak.Record): # type: ignore[misc] """ Defines ``awkward`` behavior for a 2 dimensional vector record. See :class:`VectorArray2D` and :class:`MomentumArray2D` for the corresponding ``Vector`` and ``Momentum`` classes. """ behavior["Vector2D"] = VectorRecord2D class VectorArray3D(VectorAwkward3D, ak.Array): # type: ignore[misc] """ Defines ``awkward`` behavior for a 3 dimensional vector. See :class:`VectorRecord3D` and :class:`MomentumArray3D` for the corresponding ``Record`` and ``Momentum`` classes. """ def allclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: """Like ``np.ndarray.allclose``, but for VectorArray3D.""" return ak.all(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan)) behavior["*", "Vector3D"] = VectorArray3D class VectorRecord3D(VectorAwkward3D, ak.Record): # type: ignore[misc] """ Defines ``awkward`` behavior for a 3 dimensional vector record. See :class:`VectorArray3D` and :class:`MomentumArray3D` for the corresponding ``Vector`` and ``Momentum`` classes. """ behavior["Vector3D"] = VectorRecord3D class VectorArray4D(VectorAwkward4D, ak.Array): # type: ignore[misc] """ Defines ``awkward`` behavior for a 4 dimensional vector. See :class:`VectorRecord4D` and :class:`MomentumArray4D` for the corresponding ``Record`` and ``Momentum`` classes. """ def allclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: """Like ``np.ndarray.allclose``, but for VectorArray4D.""" return ak.all(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan)) behavior["*", "Vector4D"] = VectorArray4D class VectorRecord4D(VectorAwkward4D, ak.Record): # type: ignore[misc] """ Defines ``awkward`` behavior for a 4 dimensional vector record. See :class:`VectorArray4D` and :class:`MomentumArray4D` for the corresponding ``Vector`` and ``Momentum`` classes. """ behavior["Vector4D"] = VectorRecord4D class MomentumArray2D(MomentumAwkward2D, ak.Array): # type: ignore[misc] """ Defines ``awkward`` behavior for a 2 dimensional momentum vector. See :class:`MomentumRecord2D` and :class:`VectorArray2D` for the corresponding ``Record`` and ``Vector`` classes. """ def allclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: """Like ``np.ndarray.allclose``, but for MomentumArray4D.""" return ak.all(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan)) behavior["*", "Momentum2D"] = MomentumArray2D class MomentumRecord2D(MomentumAwkward2D, ak.Record): # type: ignore[misc] """ Defines ``awkward`` behavior for a 2 dimensional momentum record. See :class:`MomentumArray2D` and :class:`VectorArray2D` for the corresponding ``Momentum`` and ``Vector`` classes. """ behavior["Momentum2D"] = MomentumRecord2D class MomentumArray3D(MomentumAwkward3D, ak.Array): # type: ignore[misc] """ Defines ``awkward`` behavior for a 3 dimensional momentum vector. See :class:`MomentumRecord3D` and :class:`VectorArray3D` for the corresponding ``Record`` and ``Vector`` classes. """ def allclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: return ak.all(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan)) behavior["*", "Momentum3D"] = MomentumArray3D class MomentumRecord3D(MomentumAwkward3D, ak.Record): # type: ignore[misc] """ Defines ``awkward`` behavior for a 3 dimensional momentum record. See :class:`MomentumArray3D` and :class:`VectorArray3D` for the corresponding ``Momentum`` and ``Vector`` classes. """ behavior["Momentum3D"] = MomentumRecord3D class MomentumArray4D(MomentumAwkward4D, ak.Array): # type: ignore[misc] """ Defines ``awkward`` behavior for a 4 dimensional momentum vector. See :class:`MomentumRecord4D` and :class:`VectorArray4D` for the corresponding ``Record`` and ``Vector`` classes. """ def allclose( self, other: VectorProtocol, rtol: ScalarCollection = 1e-05, atol: ScalarCollection = 1e-08, equal_nan: BoolCollection = False, ) -> BoolCollection: return ak.all(self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan)) behavior["*", "Momentum4D"] = MomentumArray4D class MomentumRecord4D(MomentumAwkward4D, ak.Record): # type: ignore[misc] """ Defines ``awkward`` behavior for a 4 dimensional momentum record. See :class:`MomentumArray4D` and :class:`VectorArray4D` for the corresponding ``Momentum`` and ``Vector`` classes. """ behavior["Momentum4D"] = MomentumRecord4D # NumPy functions, which also affect operator overloading ##################### behavior[numpy.absolute, "Vector2D"] = lambda v: v.rho behavior[numpy.absolute, "Vector3D"] = lambda v: v.mag behavior[numpy.absolute, "Vector4D"] = lambda v: v.tau behavior[numpy.absolute, "Momentum2D"] = lambda v: v.rho behavior[numpy.absolute, "Momentum3D"] = lambda v: v.mag behavior[numpy.absolute, "Momentum4D"] = lambda v: v.tau behavior[numpy.square, "Vector2D"] = lambda v: v.rho2 behavior[numpy.square, "Vector3D"] = lambda v: v.mag2 behavior[numpy.square, "Vector4D"] = lambda v: v.tau2 behavior[numpy.square, "Momentum2D"] = lambda v: v.rho2 behavior[numpy.square, "Momentum3D"] = lambda v: v.mag2 behavior[numpy.square, "Momentum4D"] = lambda v: v.tau2 behavior[numpy.sqrt, "Vector2D"] = lambda v: v.rho2**0.25 behavior[numpy.sqrt, "Vector3D"] = lambda v: v.mag2**0.25 behavior[numpy.sqrt, "Vector4D"] = lambda v: v.tau2**0.25 behavior[numpy.sqrt, "Momentum2D"] = lambda v: v.rho2**0.25 behavior[numpy.sqrt, "Momentum3D"] = lambda v: v.mag2**0.25 behavior[numpy.sqrt, "Momentum4D"] = lambda v: v.tau2**0.25 behavior[numpy.cbrt, "Vector2D"] = lambda v: v.rho2**0.16666666666666666 behavior[numpy.cbrt, "Vector3D"] = lambda v: v.mag2**0.16666666666666666 behavior[numpy.cbrt, "Vector4D"] = lambda v: v.tau2**0.16666666666666666 behavior[numpy.cbrt, "Momentum2D"] = lambda v: v.rho2**0.16666666666666666 behavior[numpy.cbrt, "Momentum3D"] = lambda v: v.mag2**0.16666666666666666 behavior[numpy.cbrt, "Momentum4D"] = lambda v: v.tau2**0.16666666666666666 behavior[numpy.power, "Vector2D", numbers.Real] = ( lambda v, expo: v.rho2 if expo == 2 else v.rho**expo ) behavior[numpy.power, "Vector3D", numbers.Real] = ( lambda v, expo: v.mag2 if expo == 2 else v.mag**expo ) behavior[numpy.power, "Vector4D", numbers.Real] = ( lambda v, expo: v.tau2 if expo == 2 else v.tau**expo ) behavior[numpy.power, "Momentum2D", numbers.Real] = ( lambda v, expo: v.rho2 if expo == 2 else v.rho**expo ) behavior[numpy.power, "Momentum3D", numbers.Real] = ( lambda v, expo: v.mag2 if expo == 2 else v.mag**expo ) behavior[numpy.power, "Momentum4D", numbers.Real] = ( lambda v, expo: v.tau2 if expo == 2 else v.tau**expo ) behavior["__cast__", VectorNumpy2D] = lambda v: vector.Array(v) behavior["__cast__", VectorNumpy3D] = lambda v: vector.Array(v) behavior["__cast__", VectorNumpy4D] = lambda v: vector.Array(v) for left in ( "Vector2D", "Vector3D", "Vector4D", "Momentum2D", "Momentum3D", "Momentum4D", VectorObject2D, VectorObject3D, VectorObject4D, ): for right in ( "Vector2D", "Vector3D", "Vector4D", "Momentum2D", "Momentum3D", "Momentum4D", VectorObject2D, VectorObject3D, VectorObject4D, ): if not (isinstance(left, type) and isinstance(right, type)): behavior[numpy.add, left, right] = lambda v1, v2: v1.add(v2) behavior[numpy.subtract, left, right] = lambda v1, v2: v1.subtract(v2) behavior[numpy.matmul, left, right] = lambda v1, v2: v1.dot(v2) behavior[numpy.equal, left, right] = lambda v1, v2: v1.equal(v2) behavior[numpy.not_equal, left, right] = lambda v1, v2: v1.not_equal(v2) for name in ( "Vector2D", "Vector3D", "Vector4D", "Momentum2D", "Momentum3D", "Momentum4D", ): behavior[numpy.multiply, name, numbers.Real] = lambda v, factor: v.scale(factor) behavior[numpy.multiply, numbers.Real, name] = lambda factor, v: v.scale(factor) behavior[numpy.negative, name] = lambda v: v.scale(-1) behavior[numpy.positive, name] = lambda v: v behavior[numpy.true_divide, name, numbers.Real] = lambda v, denom: v.scale( 1 / denom ) # class object cross-references ############################################### VectorArray2D.ProjectionClass2D = VectorArray2D VectorArray2D.ProjectionClass3D = VectorArray3D VectorArray2D.ProjectionClass4D = VectorArray4D VectorArray2D.GenericClass = VectorArray2D VectorArray2D.MomentumClass = MomentumArray2D VectorRecord2D.ProjectionClass2D = VectorRecord2D VectorRecord2D.ProjectionClass3D = VectorRecord3D VectorRecord2D.ProjectionClass4D = VectorRecord4D VectorRecord2D.GenericClass = VectorRecord2D VectorRecord2D.MomentumClass = MomentumRecord2D MomentumArray2D.ProjectionClass2D = MomentumArray2D MomentumArray2D.ProjectionClass3D = MomentumArray3D MomentumArray2D.ProjectionClass4D = MomentumArray4D MomentumArray2D.GenericClass = VectorArray2D MomentumArray2D.MomentumClass = MomentumArray2D MomentumRecord2D.ProjectionClass2D = MomentumRecord2D MomentumRecord2D.ProjectionClass3D = MomentumRecord3D MomentumRecord2D.ProjectionClass4D = MomentumRecord4D MomentumRecord2D.GenericClass = VectorRecord2D MomentumRecord2D.MomentumClass = MomentumRecord2D VectorArray3D.ProjectionClass2D = VectorArray2D VectorArray3D.ProjectionClass3D = VectorArray3D VectorArray3D.ProjectionClass4D = VectorArray4D VectorArray3D.GenericClass = VectorArray3D VectorArray3D.MomentumClass = MomentumArray3D VectorRecord3D.ProjectionClass2D = VectorRecord2D VectorRecord3D.ProjectionClass3D = VectorRecord3D VectorRecord3D.ProjectionClass4D = VectorRecord4D VectorRecord3D.GenericClass = VectorRecord3D VectorRecord3D.MomentumClass = MomentumRecord3D MomentumArray3D.ProjectionClass2D = MomentumArray2D MomentumArray3D.ProjectionClass3D = MomentumArray3D MomentumArray3D.ProjectionClass4D = MomentumArray4D MomentumArray3D.GenericClass = VectorArray3D MomentumArray3D.MomentumClass = MomentumArray3D MomentumRecord3D.ProjectionClass2D = MomentumRecord2D MomentumRecord3D.ProjectionClass3D = MomentumRecord3D MomentumRecord3D.ProjectionClass4D = MomentumRecord4D MomentumRecord3D.GenericClass = VectorRecord3D MomentumRecord3D.MomentumClass = MomentumRecord3D VectorArray4D.ProjectionClass2D = VectorArray2D VectorArray4D.ProjectionClass3D = VectorArray3D VectorArray4D.ProjectionClass4D = VectorArray4D VectorArray4D.GenericClass = VectorArray4D VectorArray4D.MomentumClass = MomentumArray4D VectorRecord4D.ProjectionClass2D = VectorRecord2D VectorRecord4D.ProjectionClass3D = VectorRecord3D VectorRecord4D.ProjectionClass4D = VectorRecord4D VectorRecord4D.GenericClass = VectorRecord4D VectorRecord4D.MomentumClass = MomentumRecord4D MomentumArray4D.ProjectionClass2D = MomentumArray2D MomentumArray4D.ProjectionClass3D = MomentumArray3D MomentumArray4D.ProjectionClass4D = MomentumArray4D MomentumArray4D.GenericClass = VectorArray4D MomentumArray4D.MomentumClass = MomentumArray4D MomentumRecord4D.ProjectionClass2D = MomentumRecord2D MomentumRecord4D.ProjectionClass3D = MomentumRecord3D MomentumRecord4D.ProjectionClass4D = MomentumRecord4D MomentumRecord4D.GenericClass = VectorRecord4D MomentumRecord4D.MomentumClass = MomentumRecord4D # implementation of behaviors in Numba ######################################## def _lookup_field(record_type: typing.Any, name: str) -> int: return record_type.fields.index(name) def _arraytype_of(awkwardtype: typing.Any, component: str) -> typing.Any: import numba if isinstance(awkwardtype, ak._connect.numba.layout.NumpyArrayType): return awkwardtype.arraytype elif isinstance(awkwardtype, ak._connect.numba.layout.IndexedArrayType): return _arraytype_of(awkwardtype.contenttype, component) raise numba.TypingError( f"vector components like {component!r} must be NumpyArrayType, not {awkwardtype}" ) def _aztype_of(recordarraytype: typing.Any, is_momentum: bool) -> typing.Any: import numba cls: type[AzimuthalObjectXY] | type[AzimuthalObjectRhoPhi] x_index = None y_index = None rho_index = None phi_index = None if is_momentum: try: x_index = _lookup_field(recordarraytype, "px") except ValueError: x_index = None if x_index is None: try: x_index = _lookup_field(recordarraytype, "x") except ValueError: x_index = None if is_momentum: try: y_index = _lookup_field(recordarraytype, "py") except ValueError: y_index = None if y_index is None: try: y_index = _lookup_field(recordarraytype, "y") except ValueError: y_index = None if is_momentum: try: rho_index = _lookup_field(recordarraytype, "pt") except ValueError: rho_index = None if rho_index is None: try: rho_index = _lookup_field(recordarraytype, "rho") except ValueError: rho_index = None try: phi_index = _lookup_field(recordarraytype, "phi") except ValueError: phi_index = None if x_index is not None and y_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[x_index], "x").dtype coord2 = _arraytype_of(recordarraytype.contenttypes[y_index], "y").dtype cls = AzimuthalObjectXY elif rho_index is not None and phi_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[rho_index], "rho").dtype coord2 = _arraytype_of(recordarraytype.contenttypes[phi_index], "phi").dtype cls = AzimuthalObjectRhoPhi elif is_momentum: raise numba.TypingError( f"{recordarraytype} is missing azimuthal fields: px/py (x/y) or pt/phi (rho/phi)" ) else: raise numba.TypingError( f"{recordarraytype} is missing azimuthal fields: x/y or rho/phi" ) return numba.typeof(cls(coord1.cast_python_value(0), coord2.cast_python_value(0))) def _ltype_of(recordarraytype: typing.Any, is_momentum: bool) -> typing.Any: import numba cls: ( type[LongitudinalObjectZ] | type[LongitudinalObjectTheta] | type[LongitudinalObjectEta] ) z_index = None theta_index = None eta_index = None if is_momentum: try: z_index = _lookup_field(recordarraytype, "pz") except ValueError: z_index = None if z_index is None: try: z_index = _lookup_field(recordarraytype, "z") except ValueError: z_index = None try: theta_index = _lookup_field(recordarraytype, "theta") except ValueError: theta_index = None try: eta_index = _lookup_field(recordarraytype, "eta") except ValueError: eta_index = None if z_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[z_index], "z").dtype cls = LongitudinalObjectZ elif theta_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[theta_index], "theta").dtype cls = LongitudinalObjectTheta elif eta_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[eta_index], "eta").dtype cls = LongitudinalObjectEta elif is_momentum: raise numba.TypingError( f"{recordarraytype} is missing longitudinal fields: pz (z) or theta or eta" ) else: raise numba.TypingError( f"{recordarraytype} is missing longitudinal fields: z or theta or eta" ) return numba.typeof(cls(coord1.cast_python_value(0))) def _ttype_of(recordarraytype: typing.Any, is_momentum: bool) -> typing.Any: import numba cls: type[TemporalObjectT] | type[TemporalObjectTau] t_index = None tau_index = None if is_momentum: try: t_index = _lookup_field(recordarraytype, "E") except ValueError: t_index = None if is_momentum and t_index is None: try: t_index = _lookup_field(recordarraytype, "e") except ValueError: t_index = None if is_momentum and t_index is None: try: t_index = _lookup_field(recordarraytype, "energy") except ValueError: t_index = None if t_index is None: try: t_index = _lookup_field(recordarraytype, "t") except ValueError: t_index = None if is_momentum: try: tau_index = _lookup_field(recordarraytype, "M") except ValueError: tau_index = None if is_momentum and tau_index is None: try: tau_index = _lookup_field(recordarraytype, "m") except ValueError: tau_index = None if is_momentum and tau_index is None: try: tau_index = _lookup_field(recordarraytype, "mass") except ValueError: tau_index = None if tau_index is None: try: tau_index = _lookup_field(recordarraytype, "tau") except ValueError: tau_index = None if t_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[t_index], "t").dtype cls = TemporalObjectT elif tau_index is not None: coord1 = _arraytype_of(recordarraytype.contenttypes[tau_index], "tau").dtype cls = TemporalObjectTau elif is_momentum: raise numba.TypingError( f"{recordarraytype} is missing temporal fields: E/e/energy (t) or M/m/mass (tau)" ) else: raise numba.TypingError( f"{recordarraytype} is missing temporal fields: t or tau" ) return numba.typeof(cls(coord1.cast_python_value(0))) def _numba_typer_Vector2D(viewtype: typing.Any) -> typing.Any: import vector.backends._numba_object # These exist, but the file is type ignored, so can't be found return vector.backends._numba_object.VectorObject2DType( # type: ignore[attr-defined] _aztype_of(viewtype.arrayviewtype.type, False) ) def _numba_typer_Vector3D(viewtype: typing.Any) -> typing.Any: import vector.backends._numba_object return vector.backends._numba_object.VectorObject3DType( # type: ignore[attr-defined] _aztype_of(viewtype.arrayviewtype.type, False), _ltype_of(viewtype.arrayviewtype.type, False), ) def _numba_typer_Vector4D(viewtype: typing.Any) -> typing.Any: import vector.backends._numba_object return vector.backends._numba_object.VectorObject4DType( # type: ignore[attr-defined] _aztype_of(viewtype.arrayviewtype.type, False), _ltype_of(viewtype.arrayviewtype.type, False), _ttype_of(viewtype.arrayviewtype.type, False), ) def _numba_typer_Momentum2D(viewtype: typing.Any) -> typing.Any: import vector.backends._numba_object return vector.backends._numba_object.MomentumObject2DType( # type: ignore[attr-defined] _aztype_of(viewtype.arrayviewtype.type, True) ) def _numba_typer_Momentum3D(viewtype: typing.Any) -> typing.Any: import vector.backends._numba_object return vector.backends._numba_object.MomentumObject3DType( # type: ignore[attr-defined] _aztype_of(viewtype.arrayviewtype.type, True), _ltype_of(viewtype.arrayviewtype.type, True), ) def _numba_typer_Momentum4D(viewtype: typing.Any) -> typing.Any: import vector.backends._numba_object return vector.backends._numba_object.MomentumObject4DType( # type: ignore[attr-defined] _aztype_of(viewtype.arrayviewtype.type, True), _ltype_of(viewtype.arrayviewtype.type, True), _ttype_of(viewtype.arrayviewtype.type, True), ) def _numba_lower( context: typing.Any, builder: typing.Any, sig: typing.Any, args: typing.Any ) -> typing.Any: from vector.backends._numba_object import ( # type: ignore[attr-defined] _awkward_numba_E, _awkward_numba_e, _awkward_numba_energy, _awkward_numba_eta, _awkward_numba_M, _awkward_numba_m, _awkward_numba_mass, _awkward_numba_ptphi, _awkward_numba_pxpy, _awkward_numba_pxy, _awkward_numba_pz, _awkward_numba_rhophi, _awkward_numba_t, _awkward_numba_tau, _awkward_numba_theta, _awkward_numba_xpy, _awkward_numba_xy, _awkward_numba_z, ) vectorcls = sig.return_type.instance_class fields = sig.args[0].arrayviewtype.type.fields if issubclass(vectorcls, (VectorObject2D, VectorObject3D, VectorObject4D)): if issubclass(sig.return_type.azimuthaltype.instance_class, AzimuthalXY): if "x" in fields and "y" in fields: azimuthal = _awkward_numba_xy elif "x" in fields and "py" in fields: azimuthal = _awkward_numba_xpy elif "px" in fields and "y" in fields: azimuthal = _awkward_numba_pxy elif "px" in fields and "py" in fields: azimuthal = _awkward_numba_pxpy else: raise AssertionError elif issubclass(sig.return_type.azimuthaltype.instance_class, AzimuthalRhoPhi): if "rho" in fields and "phi" in fields: azimuthal = _awkward_numba_rhophi elif "pt" in fields and "phi" in fields: azimuthal = _awkward_numba_ptphi else: raise AssertionError if issubclass(vectorcls, (VectorObject3D, VectorObject4D)): if issubclass(sig.return_type.longitudinaltype.instance_class, LongitudinalZ): if "z" in fields: longitudinal = _awkward_numba_z elif "pz" in fields: longitudinal = _awkward_numba_pz else: raise AssertionError elif issubclass( sig.return_type.longitudinaltype.instance_class, LongitudinalTheta ): longitudinal = _awkward_numba_theta elif issubclass( sig.return_type.longitudinaltype.instance_class, LongitudinalEta ): longitudinal = _awkward_numba_eta if issubclass(vectorcls, VectorObject4D): if issubclass(sig.return_type.temporaltype.instance_class, TemporalT): if "t" in fields: temporal = _awkward_numba_t elif "E" in fields: temporal = _awkward_numba_E elif "e" in fields: temporal = _awkward_numba_e elif "energy" in fields: temporal = _awkward_numba_energy else: raise AssertionError elif issubclass(sig.return_type.temporaltype.instance_class, TemporalTau): if "tau" in fields: temporal = _awkward_numba_tau elif "M" in fields: temporal = _awkward_numba_M elif "m" in fields: temporal = _awkward_numba_m elif "mass" in fields: temporal = _awkward_numba_mass else: raise AssertionError if issubclass(vectorcls, VectorObject2D): def impl(record: typing.Any) -> typing.Any: return vectorcls(azimuthal(record)) elif issubclass(vectorcls, VectorObject3D): def impl(record: typing.Any) -> typing.Any: return vectorcls(azimuthal(record), longitudinal(record)) elif issubclass(vectorcls, VectorObject4D): def impl(record: typing.Any) -> typing.Any: return vectorcls(azimuthal(record), longitudinal(record), temporal(record)) return context.compile_internal(builder, impl, sig, args) behavior["__numba_typer__", "Vector2D"] = _numba_typer_Vector2D behavior["__numba_typer__", "Vector3D"] = _numba_typer_Vector3D behavior["__numba_typer__", "Vector4D"] = _numba_typer_Vector4D behavior["__numba_typer__", "Momentum2D"] = _numba_typer_Momentum2D behavior["__numba_typer__", "Momentum3D"] = _numba_typer_Momentum3D behavior["__numba_typer__", "Momentum4D"] = _numba_typer_Momentum4D behavior["__numba_lower__", "Vector2D"] = _numba_lower behavior["__numba_lower__", "Vector3D"] = _numba_lower behavior["__numba_lower__", "Vector4D"] = _numba_lower behavior["__numba_lower__", "Momentum2D"] = _numba_lower behavior["__numba_lower__", "Momentum3D"] = _numba_lower behavior["__numba_lower__", "Momentum4D"] = _numba_lower def _reduce_sum( array: VectorArray2D | VectorArray3D | VectorArray4D | MomentumArray2D | MomentumArray3D | MomentumArray4D, mask_identity: bool, ) -> VectorProtocol: fields = {} if isinstance(array, Lorentz): fields["t"] = numpy.sum(array.t, axis=1) if isinstance(array, Spatial): fields["z"] = numpy.sum(array.z, axis=1) assert isinstance(array, Planar) fields["x"] = numpy.sum(array.x, axis=1) fields["y"] = numpy.sum(array.y, axis=1) layout = ak.to_layout(array) return ak.zip( fields, behavior=array.behavior, with_name=layout.purelist_parameter("__record__"), ) def _reduce_count( array: VectorArray2D | VectorArray3D | VectorArray4D | MomentumArray2D | MomentumArray3D | MomentumArray4D, mask_identity: bool, ) -> VectorProtocol: first_field = array[array.fields[0]] return ak.count(first_field, axis=1) def _reduce_count_nonzero( array: VectorArray2D | VectorArray3D | VectorArray4D | MomentumArray2D | MomentumArray3D | MomentumArray4D, mask_identity: bool, ) -> ScalarCollection: is_nonzero = array.rho2 != 0 if isinstance(array, Spatial): is_nonzero = numpy.logical_or(is_nonzero, array.z != 0) if isinstance(array, Lorentz): is_nonzero = numpy.logical_or(is_nonzero, array.t2 != 0) return ak.count_nonzero(is_nonzero, axis=1) for reducer, impl in ( (ak.sum, _reduce_sum), (ak.count, _reduce_count), (ak.count_nonzero, _reduce_count_nonzero), ): for dim in range(2, 5): behavior[reducer, f"Vector{dim}D"] = impl behavior[reducer, f"Momentum{dim}D"] = impl vector-1.6.3/src/vector/backends/awkward_constructors.py000066400000000000000000000335651503546127100235500ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import typing import numpy def _recname(is_momentum: bool, dimension: int) -> str: name = "Momentum" if is_momentum else "Vector" return f"{name}{dimension}D" def _check_names( projectable: typing.Any, fieldnames: list[str] ) -> tuple[bool, int, list[str], typing.Any]: complaint1 = "duplicate coordinates (through momentum-aliases): " + ", ".join( repr(x) for x in fieldnames ) complaint2 = ( "unrecognized combination of coordinates, allowed combinations are:\n\n" " (2D) x= y=\n" " (2D) rho= phi=\n" " (3D) x= y= z=\n" " (3D) x= y= theta=\n" " (3D) x= y= eta=\n" " (3D) rho= phi= z=\n" " (3D) rho= phi= theta=\n" " (3D) rho= phi= eta=\n" " (4D) x= y= z= t=\n" " (4D) x= y= z= tau=\n" " (4D) x= y= theta= t=\n" " (4D) x= y= theta= tau=\n" " (4D) x= y= eta= t=\n" " (4D) x= y= eta= tau=\n" " (4D) rho= phi= z= t=\n" " (4D) rho= phi= z= tau=\n" " (4D) rho= phi= theta= t=\n" " (4D) rho= phi= theta= tau=\n" " (4D) rho= phi= eta= t=\n" " (4D) rho= phi= eta= tau=" ) is_momentum = False dimension = 0 names = [] columns = [] if "x" in fieldnames and "y" in fieldnames: if dimension != 0: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 2 names.extend(["x", "y"]) columns.extend([projectable["x"], projectable["y"]]) fieldnames.remove("x") fieldnames.remove("y") if "rho" in fieldnames and "phi" in fieldnames: if dimension != 0: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 2 names.extend(["rho", "phi"]) columns.extend([projectable["rho"], projectable["phi"]]) fieldnames.remove("rho") fieldnames.remove("phi") if "x" in fieldnames and "py" in fieldnames: is_momentum = True if dimension != 0: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 2 names.extend(["x", "y"]) columns.extend([projectable["x"], projectable["py"]]) fieldnames.remove("x") fieldnames.remove("py") if "px" in fieldnames and "y" in fieldnames: is_momentum = True if dimension != 0: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 2 names.extend(["x", "y"]) columns.extend([projectable["px"], projectable["y"]]) fieldnames.remove("px") fieldnames.remove("y") if "px" in fieldnames and "py" in fieldnames: is_momentum = True if dimension != 0: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 2 names.extend(["x", "y"]) columns.extend([projectable["px"], projectable["py"]]) fieldnames.remove("px") fieldnames.remove("py") if "pt" in fieldnames and "phi" in fieldnames: is_momentum = True if dimension != 0: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 2 names.extend(["rho", "phi"]) columns.extend([projectable["pt"], projectable["phi"]]) fieldnames.remove("pt") fieldnames.remove("phi") if "z" in fieldnames: if dimension != 2: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 3 names.append("z") columns.append(projectable["z"]) fieldnames.remove("z") if "theta" in fieldnames: if dimension != 2: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 3 names.append("theta") columns.append(projectable["theta"]) fieldnames.remove("theta") if "eta" in fieldnames: if dimension != 2: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 3 names.append("eta") columns.append(projectable["eta"]) fieldnames.remove("eta") if "pz" in fieldnames: is_momentum = True if dimension != 2: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 3 names.append("z") columns.append(projectable["pz"]) fieldnames.remove("pz") if "t" in fieldnames: if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("t") columns.append(projectable["t"]) fieldnames.remove("t") if "tau" in fieldnames: if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("tau") columns.append(projectable["tau"]) fieldnames.remove("tau") if "E" in fieldnames: is_momentum = True if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("t") columns.append(projectable["E"]) fieldnames.remove("E") if "e" in fieldnames: is_momentum = True if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("t") columns.append(projectable["e"]) fieldnames.remove("e") if "energy" in fieldnames: is_momentum = True if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("t") columns.append(projectable["energy"]) fieldnames.remove("energy") if "M" in fieldnames: is_momentum = True if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("tau") columns.append(projectable["M"]) fieldnames.remove("M") if "m" in fieldnames: is_momentum = True if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("tau") columns.append(projectable["m"]) fieldnames.remove("m") if "mass" in fieldnames: is_momentum = True if dimension != 3: raise TypeError(complaint1 if is_momentum else complaint2) dimension = 4 names.append("tau") columns.append(projectable["mass"]) fieldnames.remove("mass") if dimension == 0: raise TypeError(complaint1 if is_momentum else complaint2) for name in fieldnames: names.append(name) columns.append(projectable[name]) return is_momentum, dimension, names, columns def _is_type_safe(array_type: typing.Any) -> None: import awkward while isinstance( array_type, ( awkward.types.ArrayType, awkward.types.RegularType, awkward.types.ListType, awkward.types.OptionType, ), ): array_type = array_type.content msg = "a coordinate must be of the type int or float" if not isinstance(array_type, awkward.types.RecordType): raise TypeError(msg) contents = array_type.contents for field_type in contents: if isinstance(field_type, awkward.types.OptionType): field_type = ( # noqa: PLW2901 field_type.content ) if not isinstance(field_type, awkward.types.NumpyType): raise TypeError(msg) dt = field_type.primitive if ( not dt.startswith("int") and not dt.startswith("uint") and not dt.startswith("float") ): raise TypeError(msg) def Array(*args: typing.Any, **kwargs: typing.Any) -> typing.Any: """ Constructs an Awkward Array of vectors, whose type is determined by the fields of the record array (which may be nested within lists or other non-record structures). All allowed signatures for ``ak.Array`` can be used in this function. The array must contain records with the following combinations of field names: - (2D) ``x``, ``y`` - (2D) ``rho``, ``phi`` - (3D) ``x``, ``y``, ``z`` - (3D) ``x``, ``y``, ``theta`` - (3D) ``x``, ``y``, ``eta`` - (3D) ``rho``, ``phi``, ``z`` - (3D) ``rho``, ``phi``, ``theta`` - (3D) ``rho``, ``phi``, ``eta`` - (4D) ``x``, ``y``, ``z``, ``t`` - (4D) ``x``, ``y``, ``z``, ``tau``` - (4D) ``x``, ``y``, ``theta``, ``t``` - (4D) ``x``, ``y``, ``theta``, ``tau``` - (4D) ``x``, ``y``, ``eta``, ``t``` - (4D) ``x``, ``y``, ``eta``, ``tau``` - (4D) ``rho``, ``phi``, ``z``, ``t``` - (4D) ``rho``, ``phi``, ``z``, ``tau``` - (4D) ``rho``, ``phi``, ``theta``, ``t``` - (4D) ``rho``, ``phi``, ``theta``, ``tau``` - (4D) ``rho``, ``phi``, ``eta``, ``t``` - (4D) ``rho``, ``phi``, ``eta``, ``tau``` in which - ``px`` may be substituted for ``x`` - ``py`` may be substituted for ``y`` - ``pt`` may be substituted for ``rho`` - ``pz`` may be substituted for ``z`` - ``E`` may be substituted for ``t`` - ``e`` may be substituted for ``t`` - ``energy`` may be substituted for ``t`` - ``M`` may be substituted for ``tau`` - ``m`` may be substituted for ``tau`` - ``mass`` may be substituted for ``tau`` to make the vector a momentum vector. No constraints are placed on the types of the vector fields, though if they are not numbers, mathematical operations will fail. Usually, you want them to be integers or floating-point numbers. """ import awkward import vector import vector.backends.awkward # don't pass dask_awkward arrays in ak.Array akarray = ( awkward.Array(*args, **kwargs) if isinstance(args[0], (list, numpy.ndarray)) else args[0] ) array_type = akarray.type _is_type_safe(array_type) fields = awkward.fields(akarray) is_momentum, dimension, names, arrays = _check_names(akarray, fields.copy()) # don't execute for dask_awkward arrays if isinstance(args[0], (awkward.Array, list, numpy.ndarray)): needs_behavior = not vector._awkward_registered for x in arrays: if needs_behavior: if x.behavior is None: x.behavior = vector.backends.awkward.behavior else: x.behavior = dict(x.behavior) x.behavior.update(vector.backends.awkward.behavior) else: x.behavior = None needs_behavior = False assert 2 <= dimension <= 4, f"Dimension must be between 2-4, not {dimension}" return awkward.with_name( awkward.zip( dict(__builtins__["zip"](names, arrays)), # type: ignore[index] depth_limit=akarray.layout.purelist_depth, ), _recname(is_momentum, dimension), behavior=vector.backends.awkward.behavior, ) def zip(arrays: dict[str, typing.Any], depth_limit: int | None = None) -> typing.Any: """ Constructs an Awkward Array of vectors, whose type is determined by the fields of the record array (which may be nested within lists or other non-record structures). This function accepts a subset of ``ak.zip``'s arguments. Args: arrays (dict of str to array-like): Arrays, lists, etc. to zip together. Unlike ``ak.zip``, this must be a dict with string keys to determine the coordinate system of the arrays; it may not be a tuple. depth_limit (None or int): If None, attempt to fully broadcast the ``array`` to all levels. If an int, limit the number of dimensions that get broadcasted. The minimum value is ``1``, for no broadcasting. The array must contain records with the following combinations of field names: - (2D) ``x``, ``y`` - (2D) ``rho``, ``phi`` - (3D) ``x``, ``y``, ``z`` - (3D) ``x``, ``y``, ``theta`` - (3D) ``x``, ``y``, ``eta`` - (3D) ``rho``, ``phi``, ``z`` - (3D) ``rho``, ``phi``, ``theta`` - (3D) ``rho``, ``phi``, ``eta`` - (4D) ``x``, ``y``, ``z``, ``t`` - (4D) ``x``, ``y``, ``z``, ``tau``` - (4D) ``x``, ``y``, ``theta``, ``t``` - (4D) ``x``, ``y``, ``theta``, ``tau``` - (4D) ``x``, ``y``, ``eta``, ``t``` - (4D) ``x``, ``y``, ``eta``, ``tau``` - (4D) ``rho``, ``phi``, ``z``, ``t``` - (4D) ``rho``, ``phi``, ``z``, ``tau``` - (4D) ``rho``, ``phi``, ``theta``, ``t``` - (4D) ``rho``, ``phi``, ``theta``, ``tau``` - (4D) ``rho``, ``phi``, ``eta``, ``t``` - (4D) ``rho``, ``phi``, ``eta``, ``tau``` in which - ``px`` may be substituted for ``x`` - ``py`` may be substituted for ``y`` - ``pt`` may be substituted for ``rho`` - ``pz`` may be substituted for ``z`` - ``E`` may be substituted for ``t`` - ``e`` may be substituted for ``t`` - ``energy`` may be substituted for ``t`` - ``M`` may be substituted for ``tau`` - ``m`` may be substituted for ``tau`` - ``mass`` may be substituted for ``tau`` to make the vector a momentum vector. """ import awkward import vector import vector.backends.awkward if not isinstance(arrays, dict): raise TypeError("argument passed to vector.zip must be a dictionary") is_momentum, dimension, names, columns = _check_names(arrays, list(arrays.keys())) behavior = None if not vector._awkward_registered: behavior = dict(vector.backends.awkward.behavior) return awkward.zip( dict(__builtins__["zip"](names, columns)), # type: ignore[index] depth_limit=depth_limit, with_name=_recname(is_momentum, dimension), behavior=behavior, ) vector-1.6.3/src/vector/backends/numba_numpy.py000066400000000000000000000010001503546127100215650ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Implements VectorNumpys in Numba. Only __getitem__ should be overloaded to return lowered VectorObjects, and maybe the ``vector.array`` constructor should also be handled. As mentioned in scikit-hep/vector#43, this capability is blocked by numba/numba#6148. """ vector-1.6.3/src/vector/backends/numpy.py000066400000000000000000002343201503546127100204200ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Defines behaviors for NumPy Array. New arrays created with the .. code-block:: python vector.array(...) function will have these behaviors built in (and will pass them to any derived arrays). """ from __future__ import annotations import collections.abc import typing import numpy import vector.backends.object from vector._methods import ( Azimuthal, AzimuthalRhoPhi, AzimuthalXY, Longitudinal, LongitudinalEta, LongitudinalTheta, LongitudinalZ, Lorentz, LorentzMomentum, Momentum, Planar, PlanarMomentum, Spatial, SpatialMomentum, Temporal, TemporalT, TemporalTau, Vector, Vector2D, Vector3D, Vector4D, VectorProtocol, _aztype, _coordinate_class_to_names, _coordinate_order, _handler_of, _ltype, _repr_momentum_to_generic, _ttype, ) from vector._typeutils import BoolCollection, FloatArray, ScalarCollection ArrayLike = ScalarCollection T = typing.TypeVar("T", bound="VectorNumpy") V = typing.TypeVar("V") SameVectorNumpyType = typing.TypeVar("SameVectorNumpyType", bound="VectorNumpy") def _reduce_sum( a: T, axis: int | None = None, dtype: typing.Any = None, out: typing.Any = None, keepdims: bool = False, initial: typing.Any = None, where: typing.Any = None, ) -> T: if where is not None: raise ValueError("cannot invoke reducer with `where` argument") if initial is not None: raise ValueError("cannot invoke reducer with `initial` argument") if out is not None: raise ValueError("cannot invoke reducer with `out` argument") if dtype is not None: raise ValueError("cannot invoke reducer with `dtype` argument") fields: dict[str, typing.Any] = {} if isinstance(a, Lorentz): fields["E"] = numpy.sum(a.t, axis=axis, keepdims=keepdims) if isinstance(a, Spatial): fields["pz"] = numpy.sum(a.z, axis=axis, keepdims=keepdims) assert isinstance(a, Planar) fields["px"] = numpy.sum(a.x, axis=axis, keepdims=keepdims) fields["py"] = numpy.sum(a.y, axis=axis, keepdims=keepdims) # Convert between representations if not isinstance(a, Momentum): fields = {_repr_momentum_to_generic[n]: v for n, v in fields.items()} return array(fields) def _reduce_count_nonzero( a: T, axis: int | None = None, *, keepdims: bool = False ) -> ScalarCollection: if isinstance(a, Planar): is_nonzero = a.rho2 != 0 else: raise AssertionError if isinstance(a, Spatial): is_nonzero = numpy.logical_or(is_nonzero, a.z != 0) if isinstance(a, Lorentz): is_nonzero = numpy.logical_or(is_nonzero, a.t2 != 0) return numpy.count_nonzero(is_nonzero, axis=axis, keepdims=keepdims) def _array_from_columns(columns: dict[str, ArrayLike]) -> ArrayLike: """ Converts a dictionary (or columns) of coordinates to an array. Args: columns (dict): The dictionary of coordinates to be converted. Returns: ``np.ndarray``: A structured array of coordinates. Examples: >>> import vector >>> vector.backends.numpy._array_from_columns({"x": [1, 2, 3], "y": [1, 2, 4]}) array([(1., 1.), (2., 2.), (3., 4.)], dtype=[('x', ' None: if isinstance(where, str): if is_momentum: where = _repr_momentum_to_generic.get(where, where) array.view(numpy.ndarray)[where] = what else: if not hasattr(what, "dtype") or what.dtype.names is None: raise TypeError( "right-hand side of assignment must be a structured array with " "the same fields as " + type(array).__name__ ) tofill = array[where] for name in what.dtype.names: if is_momentum: generic = _repr_momentum_to_generic.get(name, name) tofill[generic] = what[name] def _getitem( array: VectorNumpy2D | VectorNumpy3D | VectorNumpy4D, where: typing.Any, is_momentum: bool, ) -> float | FloatArray: """ Implementation for the ``__getitem__`` method. See :class:`GetItem` for more details. """ if isinstance(where, str): if is_momentum: where = _repr_momentum_to_generic.get(where, where) return array.view(numpy.ndarray)[where] else: out = numpy.ndarray.__getitem__(array, where) if not isinstance(out, numpy.void): return out azimuthal, longitudinal, temporal = None, None, None if hasattr(array, "_azimuthal_type"): azimuthal = array._azimuthal_type.ObjectClass( *(out[x] for x in _coordinate_class_to_names[_aztype(array)]) ) if hasattr(array, "_longitudinal_type"): longitudinal = array._longitudinal_type.ObjectClass( *(out[x] for x in _coordinate_class_to_names[_ltype(array)]) ) if hasattr(array, "_temporal_type"): temporal = array._temporal_type.ObjectClass( *(out[x] for x in _coordinate_class_to_names[_ttype(array)]) ) if temporal is not None: return array.ObjectClass( azimuthal=azimuthal, longitudinal=longitudinal, # type: ignore[arg-type] temporal=temporal, # type: ignore[arg-type] ) elif longitudinal is not None: return array.ObjectClass(azimuthal=azimuthal, longitudinal=longitudinal) # type: ignore[arg-type] elif azimuthal is not None: return array.ObjectClass(azimuthal=azimuthal) elif issubclass(array.ObjectClass, vector.backends.object.AzimuthalObject): return array.ObjectClass(*tuple(out)[:2]) # type: ignore[arg-type] elif issubclass(array.ObjectClass, vector.backends.object.LongitudinalObject): coords = ( out.view(numpy.ndarray)[0] if len(out) == 1 # type: ignore[arg-type] else out.view(numpy.ndarray)[2] ) return array.ObjectClass(coords) else: coords = ( out.view(numpy.ndarray)[0] if len(out) == 1 # type: ignore[arg-type] else out.view(numpy.ndarray)[3] ) return array.ObjectClass(coords) def _array_repr( array: VectorNumpy2D | VectorNumpy3D | VectorNumpy4D, is_momentum: bool, ) -> str: """ Constructs the value for ``__repr__`` function of the provided VectorNumpy class. """ name = type(array).__name__ vanilla_array = array.view(numpy.ndarray) return name + repr(vanilla_array)[5:].replace("\n ", "\n" + " " * len(name)) def _has( array: VectorNumpy2D | VectorNumpy3D | VectorNumpy4D | MomentumNumpy2D | MomentumNumpy3D | MomentumNumpy4D | CoordinatesNumpy, names: tuple[str, ...], ) -> bool: """ Checks if a NumPy vector has the provided coordinate attributes. Args: array (NumPy vector): A NumPy Vector whose coordinate attributes have to be checked. names (tuple): Names of the attributes. Returns: bool: If the attribute exists or not. Examples: >>> from vector.backends.numpy import _has >>> from vector._methods import _coordinate_class_to_names >>> vec = vector.array([ ... (1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5) ... ], dtype=[("x", float), ("y", float)]) >>> _has(vec, ("x", "y")) True >>> _has(vec, _coordinate_class_to_names[vector._methods.AzimuthalXY]) True >>> _has(vec, _coordinate_class_to_names[vector._methods.AzimuthalRhoPhi]) False """ dtype_names = array.dtype.names if dtype_names is None: dtype_names = () return all(x in dtype_names for x in names) def _toarrays( result: tuple[ScalarCollection, ...] | ScalarCollection, ) -> tuple[FloatArray, ...]: """ Converts a tuple of values to a tuple of ``numpy.array``s. Args: result (tuple): A tuple of values to be converted. Returns: tuple: A tuple of ``numpy.array``. Examples >>> from vector.backends.numpy import _toarrays >>> _toarrays((1, 2, 3, 4)) (array([1.]), array([2.]), array([3.]), array([4.])) >>> _toarrays((1, 2, (1, 2, 3))) (array([1.]), array([2.]), array([[1., 2., 3.]])) >>> _toarrays((1, 2, (1, 2, False))) (array([1.]), array([2.]), array([[1., 2., 0.]])) """ istuple = True if not isinstance(result, tuple): istuple = False result = (result,) result = tuple( x if isinstance(x, numpy.ndarray) else numpy.array([x], numpy.float64) for x in result ) if istuple: return result else: return result[0] def _shape_of(result: tuple[FloatArray, ...] | ScalarCollection) -> tuple[int, ...]: """ Calculates the shape of a tuple of ``numpy.array``s. The shape returned is the highest (numerical) value of the shapes present in the tuple. Args: result (tuple): A tuple of ``numpy.array``s. Returns: tuple: The calculated shape. Examples: >>> from vector.backends.numpy import _shape_of >>> import numpy as np >>> _shape_of((np.array([1]), np.array([2]))) (1,) >>> _shape_of((np.array([1]), np.array([2, 8]), np.array([0]))) (2,) >>> _shape_of((np.array([1]), np.array([2, 8]))) (2,) """ if not isinstance(result, tuple): result = (result,) shape: list[int] | None = None for x in result: if hasattr(x, "shape"): thisshape = list(x.shape) elif isinstance(x, collections.abc.Sized): thisshape = [len(x)] if shape is None or thisshape[0] > shape[0]: shape = thisshape assert shape is not None return tuple(shape) def _is_type_safe( array: VectorNumpy2D | VectorNumpy3D | VectorNumpy4D | MomentumNumpy2D | MomentumNumpy3D | MomentumNumpy4D | CoordinatesNumpy, ) -> None: for i in range(len(array.dtype)): # type: ignore[arg-type] if not issubclass( array.dtype[i].type, (numpy.integer, numpy.floating) ) or issubclass(array.dtype[i].type, numpy.timedelta64): raise TypeError( "a coordinate must be of the type numpy.integer or numpy.floating" ) class GetItem: _IS_MOMENTUM: typing.ClassVar[bool] @typing.overload def __getitem__(self, where: str) -> FloatArray: ... @typing.overload def __getitem__(self, where: typing.Any) -> float | FloatArray: ... def __getitem__(self, where: typing.Any) -> float | FloatArray: return _getitem(self, where, self.__class__._IS_MOMENTUM) # type: ignore[arg-type] class CoordinatesNumpy: """Coordinates class for the Numpy backend.""" lib = numpy dtype: numpy.dtype[typing.Any] class AzimuthalNumpy(CoordinatesNumpy, Azimuthal): """Azimuthal class for the NumPy backend.""" ObjectClass: type[vector.backends.object.AzimuthalObject] class LongitudinalNumpy(CoordinatesNumpy, Longitudinal): """Longitudinal class for the NumPy backend.""" ObjectClass: type[vector.backends.object.LongitudinalObject] class TemporalNumpy(CoordinatesNumpy, Temporal): """Temporal class for the NumPy backend.""" ObjectClass: type[vector.backends.object.TemporalObject] class AzimuthalNumpyXY(AzimuthalNumpy, AzimuthalXY, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """ Class for the ``x`` and ``y`` (azimuthal) coordinates of NumPy backend. Creates a structured NumPy array and returns it as an AzimuthalNumpyXY object. Examples: >>> import vector >>> vector.backends.numpy.AzimuthalNumpyXY([(1, 1), (2.1, 3.1)], dtype=[("x", float), ("y", float)]) AzimuthalNumpyXY([(1. , 1. ), (2.1, 3.1)], dtype=[('x', ' AzimuthalNumpyXY: if "dtype" in kwargs: AzimuthalNumpyXY.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("x", "y")): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y")' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, AzimuthalNumpyXY): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, AzimuthalNumpyXY): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray, FloatArray]: """ Azimuthal coordinates (``x`` and ``y``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.AzimuthalNumpyXY([(1, 1), (2.1, 3.1)], dtype=[("x", float), ("y", float)]) >>> vec.elements (array([1. , 2.1]), array([1. , 3.1])) """ return (self["x"], self["y"]) @property def x(self) -> FloatArray: """The ``x`` coordinates.""" return self["x"] @property def y(self) -> FloatArray: """The ``y`` coordinates.""" return self["y"] class AzimuthalNumpyRhoPhi(AzimuthalNumpy, AzimuthalRhoPhi, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """ Class for the ``rho`` and ``phi`` (azimuthal) coordinates of NumPy backend. Creates a structured NumPy array and returns it as an AzimuthalNumpyXY object. Examples: >>> import vector >>> vector.backends.numpy.AzimuthalNumpyRhoPhi([(1, 1), (2.1, 3.1)], dtype=[("rho", float), ("phi", float)]) AzimuthalNumpyRhoPhi([(1. , 1. ), (2.1, 3.1)], dtype=[('rho', ' AzimuthalNumpyRhoPhi: if "dtype" in kwargs: AzimuthalNumpyRhoPhi.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("rho", "phi")): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("rho", "phi")' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, AzimuthalNumpyRhoPhi): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, AzimuthalNumpyRhoPhi): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray, FloatArray]: """ Azimuthal coordinates (``rho`` and ``phi``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.AzimuthalNumpyRhoPhi([(1, 1), (2.1, 3.1)], dtype=[("rho", float), ("phi", float)]) >>> vec.elements (array([1. , 2.1]), array([1. , 3.1])) """ return (self["rho"], self["phi"]) @property def rho(self) -> FloatArray: """The ``rho`` coordinates.""" return self["rho"] @property def phi(self) -> FloatArray: """The ``phi`` coordinates.""" return self["phi"] class LongitudinalNumpyZ(LongitudinalNumpy, LongitudinalZ, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """ Class for the ``z`` (longitudinal) coordinate of NumPy backend. Creates a structured NumPy array and returns it as a LongitudinalNumpyZ object. Examples: >>> import vector >>> vector.backends.numpy.LongitudinalNumpyZ([(1), (2.1)], dtype=[("z", float)]) LongitudinalNumpyZ([(1. ,), (2.1,)], dtype=[('z', ' LongitudinalNumpyZ: if "dtype" in kwargs: LongitudinalNumpyZ.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("z",)): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "z"' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, LongitudinalNumpyZ): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, LongitudinalNumpyZ): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray]: """ Longitudinal coordinates (``z``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.LongitudinalNumpyZ([(1), (2.1)], dtype=[("z", float)]) >>> vec.elements (array([1. , 2.1]),) """ return (self["z"],) @property def z(self) -> FloatArray: """The ``z`` coordinates.""" return self["z"] class LongitudinalNumpyTheta(LongitudinalNumpy, LongitudinalTheta, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """ Class for the ``theta`` (longitudinal) coordinate of NumPy backend. Creates a structured NumPy array and returns it as a LongitudinalNumpyTheta object. Examples: >>> import vector >>> vector.backends.numpy.LongitudinalNumpyTheta([(1), (2.1)], dtype=[("theta", float)]) LongitudinalNumpyTheta([(1. ,), (2.1,)], dtype=[('theta', ' LongitudinalNumpyTheta: if "dtype" in kwargs: LongitudinalNumpyTheta.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("theta",)): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "theta"' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, LongitudinalNumpyTheta): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, LongitudinalNumpyTheta): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray]: """ Longitudinal coordinates (``theta``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.LongitudinalNumpyTheta([(1), (2.1)], dtype=[("theta", float)]) >>> vec.elements (array([1. , 2.1]),) """ return (self["theta"],) @property def theta(self) -> FloatArray: """The ``theta`` coordinates.""" return self["theta"] class LongitudinalNumpyEta(LongitudinalNumpy, LongitudinalEta, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """ Class for the ``eta`` (longitudinal) coordinate of NumPy backend. Creates a structured NumPy array and returns it as a LongitudinalNumpyEta object. Examples: >>> import vector >>> vector.backends.numpy.LongitudinalNumpyEta([(1), (2.1)], dtype=[("eta", float)]) LongitudinalNumpyEta([(1. ,), (2.1,)], dtype=[('eta', ' LongitudinalNumpyEta: if "dtype" in kwargs: LongitudinalNumpyEta.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("eta",)): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "eta"' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, LongitudinalNumpyEta): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, LongitudinalNumpyEta): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray]: """ Longitudinal coordinates (``eta``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.LongitudinalNumpyTheta([(1), (2.1)], dtype=[("theta", float)]) >>> vec.elements (array([1. , 2.1]),) """ return (self["eta"],) @property def eta(self) -> FloatArray: """The ``eta`` coordinates.""" return self["eta"] class TemporalNumpyT(TemporalNumpy, TemporalT, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """ Class for the ``t`` (temporal) coordinate of NumPy backend. Creates a structured NumPy array and returns it as a TemporalNumpyT object. Examples: >>> import vector >>> vector.backends.numpy.TemporalNumpyT([(1), (2.1)], dtype=[("t", float)]) TemporalNumpyT([(1. ,), (2.1,)], dtype=[('t', ' TemporalNumpyT: if "dtype" in kwargs: TemporalNumpyT.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("t",)): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "t"' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, TemporalNumpyT): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, TemporalNumpyT): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray]: """ Temporal coordinates (``t``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.TemporalNumpyT([(1), (2.1)], dtype=[("t", float)]) >>> vec.elements (array([1. , 2.1]),) """ return (self["t"],) @property def t(self) -> FloatArray: """The ``t`` coordinates.""" return self["t"] class TemporalNumpyTau(TemporalNumpy, TemporalTau, GetItem, FloatArray): # type: ignore[misc] # noqa: PLW1641 """Class for the ``tau`` (temporal) coordinate of NumPy backend.""" ObjectClass = vector.backends.object.TemporalObjectTau _IS_MOMENTUM = False dtype: numpy.dtype[typing.Any] = numpy.dtype([("tau", numpy.float64)]) def __new__(cls, *args: typing.Any, **kwargs: typing.Any) -> TemporalNumpyTau: if "dtype" in kwargs: TemporalNumpyTau.dtype = numpy.dtype(kwargs["dtype"]) return numpy.array(*args, **kwargs).view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if not _has(self, ("tau",)): raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "tau"' ) def __eq__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, TemporalNumpyTau): return False return all(coord1 == coord2 for coord1, coord2 in zip(self, other)) def __ne__(self, other: typing.Any) -> bool: if self.dtype != other.dtype or not isinstance(other, TemporalNumpyTau): return True return any(coord1 != coord2 for coord1, coord2 in zip(self, other)) @property def elements(self) -> tuple[FloatArray]: """ Temporal coordinates (``tau``) as a tuple. Each coordinate is a NumPy array of values and not a vector. Examples: >>> import vector >>> vec = vector.backends.numpy.TemporalNumpyTau([(1), (2.1)], dtype=[("tau", float)]) >>> vec.elements (array([1. , 2.1]),) """ return (self["tau"],) @property def tau(self) -> FloatArray: """The ``tau`` coordinates.""" return self["tau"] class VectorNumpy(Vector, GetItem): # noqa: PLW1641 """Mixin class for NumPy vectors.""" lib = numpy dtype: numpy.dtype[typing.Any] def allclose( self, other: VectorProtocol, rtol: float | FloatArray = 1e-05, atol: float | FloatArray = 1e-08, equal_nan: bool | FloatArray = False, ) -> BoolCollection: """Like ``np.ndarray.allclose``, but for VectorNumpy.""" return self.isclose(other, rtol=rtol, atol=atol, equal_nan=equal_nan).all() def sum( self: SameVectorNumpyType, axis: int | None = None, dtype: numpy.dtype[typing.Any] | str | None = None, out: ArrayLike | None = None, keepdims: bool = False, initial: typing.Any = None, where: typing.Any = None, ) -> SameVectorNumpyType: return typing.cast( SameVectorNumpyType, numpy.sum( self, axis=axis, dtype=dtype, out=out, keepdims=keepdims, initial=initial, where=where, ), ) def __eq__(self, other: typing.Any) -> typing.Any: return numpy.equal(self, other) def __ne__(self, other: typing.Any) -> typing.Any: return numpy.not_equal(self, other) def __reduce__(self) -> str | tuple[typing.Any, ...]: pickled_state = super().__reduce__() new_state = (*pickled_state[2], self.__dict__) return pickled_state[0], pickled_state[1], new_state def __setstate__(self, state: typing.Any) -> None: self.__dict__.update(state[-1]) super().__setstate__(state[0:-1]) # type: ignore[misc] def __array_ufunc__( self, ufunc: typing.Any, method: typing.Any, *inputs: typing.Any, **kwargs: typing.Any, ) -> typing.Any: """ Implements NumPy's ``ufunc``s for ``VectorNumpy``. The current implementation includes ``numpy.absolute``, ``numpy.add``, ``numpy.subtract``, ``numpy.multiply``, ``numpy.positive``, ``numpy.negative``, ``numpy.true_divide``, ``numpy.power``, ``numpy.square``, ``numpy.sqrt``, ``numpy.cbrt``, ``numpy.matmul``, ``numpy.equal``, and ``numpy.not_equal``. """ if not isinstance(_handler_of(*inputs), VectorNumpy): # Let a higher-precedence backend handle it. return NotImplemented outputs: tuple[VectorNumpy, ...] = kwargs.get("out", ()) if any(not isinstance(x, VectorNumpy) for x in outputs): raise TypeError( "ufunc operating on VectorNumpys can only use the 'out' keyword " "with another VectorNumpy" ) if ( ufunc is numpy.absolute and len(inputs) == 1 and isinstance(inputs[0], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.absolute' is scalar, cannot fill a VectorNumpy with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho elif isinstance(inputs[0], Vector3D): return inputs[0].mag elif isinstance(inputs[0], Vector4D): return inputs[0].tau elif ( ufunc is numpy.add and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[0].add(inputs[1]) for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.subtract and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[0].subtract(inputs[1]) for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.multiply and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = inputs[0].scale(inputs[1]) for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.multiply and len(inputs) == 2 and not isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[1].scale(inputs[0]) for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.negative and len(inputs) == 1 and isinstance(inputs[0], Vector) ): result = inputs[0].scale(-1) for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.positive and len(inputs) == 1 and isinstance(inputs[0], Vector) ): return inputs[0] elif ( ufunc is numpy.true_divide and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = inputs[0].scale(1 / inputs[1]) for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.power and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = numpy.absolute(inputs[0]) ** inputs[1] for output in outputs: assert output.dtype.names is not None for name in output.dtype.names: output[name] = result[name] # type: ignore[index] return result elif ( ufunc is numpy.square and len(inputs) == 1 and isinstance(inputs[0], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.square' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 elif ufunc is numpy.sqrt and len(inputs) == 1 and isinstance(inputs[0], Vector): if len(outputs) != 0: raise TypeError( "output of 'numpy.sqrt' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 ** 0.25 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 ** 0.25 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 ** 0.25 elif ufunc is numpy.cbrt and len(inputs) == 1 and isinstance(inputs[0], Vector): if len(outputs) != 0: raise TypeError( "output of 'numpy.cbrt' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 ** 0.16666666666666666 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 ** 0.16666666666666666 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 ** 0.16666666666666666 elif ( ufunc is numpy.matmul and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.matmul' is scalar, cannot fill a VectorObject with 'out'" ) return inputs[0].dot(inputs[1]) elif ( ufunc is numpy.equal and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.equal' is scalar, cannot fill a VectorObject with 'out'" ) return inputs[0].equal(inputs[1]) elif ( ufunc is numpy.not_equal and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.equal' is scalar, cannot fill a VectorObject with 'out'" ) return inputs[0].not_equal(inputs[1]) else: return NotImplemented def __array_function__( self, func: typing.Any, types: typing.Any, args: typing.Any, kwargs: typing.Any ) -> typing.Any: """ Implements NumPy's function for ``VectorNumpy`` and its subclasses. The current implementation includes ``numpy.isclose`` and ``numpy.allclose``. """ if func is numpy.isclose: return type(self).isclose(*args, **kwargs) elif func is numpy.allclose: return type(self).allclose(*args, **kwargs) elif func is numpy.sum: return _reduce_sum(*args, **kwargs) elif func is numpy.count_nonzero: return _reduce_count_nonzero(*args, **kwargs) else: return NotImplemented class VectorNumpy2D(VectorNumpy, Planar, Vector2D, FloatArray): # type: ignore[misc] """ Two dimensional vector class for the NumPy backend. This class can be directly used to construct two dimensional NumPy vectors. For two dimensional Momentum NumPy vectors see :class:`vector.backends.numpy.MomentumNumpy2D`. Examples: >>> import vector >>> vec = vector.VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], ... dtype=[('x', float), ('y', float)]) >>> vec VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', '>> import numpy as np >>> import vector >>> arr = np.array([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], ... dtype=[('x', float), ('y', float)]) >>> arr.view(vector.VectorNumpy2D) VectorNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', ' VectorNumpy2D: """Returns the object of ``VectorNumpy2D``. Behaves as ``__init__`` in this case.""" if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict): array = _array_from_columns(args[0]) else: array = numpy.array(*args, **kwargs) return array.view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): self._azimuthal_type = AzimuthalNumpyRhoPhi else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y") or ("rho", "phi")' ) _is_type_safe(self) def __str__(self) -> str: return str(self.view(numpy.ndarray)) def __repr__(self) -> str: return _array_repr(self, False) @property def azimuthal(self) -> AzimuthalNumpy: """ Returns the azimuthal type class for the given ``VectorNumpy2D`` object. Examples: >>> import vector >>> vec = vector.array([ ... (1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5) ... ], dtype=[("x", float), ("y", float)]) >>> vec.azimuthal AzimuthalNumpyXY([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', ' typing.Any: """ Wraps the raw result of a compute function as an array of scalars or an array of vectors. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif len(returns) == 1 or ( (len(returns) == 2 and returns[1] is None) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): result = _toarrays(result) dtype = [ (name, result[i].dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]) ] out = numpy.empty(_shape_of(result), dtype=dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]): out[name] = result[i] return out.view(cls.ProjectionClass2D) elif ( (len(returns) == 2 or (len(returns) == 3 and returns[2] is None)) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 return out.view(cls.ProjectionClass3D) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[2]]: dtype.append((name, result[i].dtype)) i += 1 out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[2]]: out[name] = result[i] i += 1 return out.view(cls.ProjectionClass4D) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, False) class MomentumNumpy2D(PlanarMomentum, VectorNumpy2D): """ Two dimensional momentum vector class for the NumPy backend. This class can be directly used to construct two dimensional NumPy momentum vectors. For two dimensional NumPy vectors see :class:`vector.backends.numpy.VectorNumpy2D`. Examples: >>> import vector >>> vec = vector.MomentumNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], ... dtype=[('px', float), ('py', float)]) >>> vec MomentumNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', '>> import numpy as np >>> import vector >>> arr = np.array([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], ... dtype=[('px', float), ('py', float)]) >>> arr.view(vector.MomentumNumpy2D) MomentumNumpy2D([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', ' None: if obj is None: return self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): self._azimuthal_type = AzimuthalNumpyRhoPhi else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y") or ("rho", "phi") or ("px", "py") or ("pt", "phi")' ) _is_type_safe(self) def __repr__(self) -> str: return _array_repr(self, True) def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, True) class VectorNumpy3D(VectorNumpy, Spatial, Vector3D, FloatArray): # type: ignore[misc] """ Three dimensional vector class for the NumPy backend. This class can be directly used to construct three dimensional NumPy vectors. For three dimensional Momentum NumPy vectors see :class:`vector.backends.numpy.MomentumNumpy3D`. Examples: >>> import vector >>> vec = vector.VectorNumpy3D([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], ... dtype=[('x', float), ('y', float), ('z', float)]) >>> vec VectorNumpy3D([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], dtype=[('x', '>> import numpy as np >>> import vector >>> arr = np.array([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], ... dtype=[('x', float), ('y', float), ('z', float)]) >>> arr.view(vector.VectorNumpy3D) VectorNumpy3D([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], dtype=[('x', ' VectorNumpy3D: """Returns the object of ``VectorNumpy3D``. Behaves as ``__init__`` in this case.""" if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict): array = _array_from_columns(args[0]) else: array = numpy.array(*args, **kwargs) return array.view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): self._azimuthal_type = AzimuthalNumpyRhoPhi else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y") or ("rho", "phi")' ) if _has(self, ("z",)): self._longitudinal_type = LongitudinalNumpyZ elif _has(self, ("theta",)): self._longitudinal_type = LongitudinalNumpyTheta elif _has(self, ("eta",)): self._longitudinal_type = LongitudinalNumpyEta else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "z" or "theta" or "eta"' ) _is_type_safe(self) def __str__(self) -> str: return str(self.view(numpy.ndarray)) def __repr__(self) -> str: return _array_repr(self, False) @property def azimuthal(self) -> AzimuthalNumpy: """ Returns the azimuthal type class for the given ``VectorNumpy3D`` object. Example: >>> import vector >>> vec = vector.array( ... [ ... (1.1, 2.1, 3.1), ... (1.2, 2.2, 3.2), ... (1.3, 2.3, 3.3), ... (1.4, 2.4, 4.4), ... (1.5, 2.5, 5.5) ... ], dtype=[("x", float), ("y", float), ("z", float)] ... ) >>> vec.azimuthal AzimuthalNumpyXY([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', ' LongitudinalNumpy: """ Returns the longitudinal type class for the given ``VectorNumpy3D`` object. Example: >>> import vector >>> vec = vector.array( ... [ ... (1.1, 2.1, 3.1), ... (1.2, 2.2, 3.2), ... (1.3, 2.3, 3.3), ... (1.4, 2.4, 4.4), ... (1.5, 2.5, 5.5) ... ], dtype=[("x", float), ("y", float), ("z", float)] ... ) >>> vec.longitudinal LongitudinalNumpyZ([(3.1,), (3.2,), (3.3,), (4.4,), (5.5,)], dtype=[('z', ' typing.Any: """ Wraps the raw result of a compute function as an array of scalars or an array of vectors. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): result = _toarrays(result) dtype = [ (name, result[i].dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]) ] for name in _coordinate_class_to_names[_ltype(self)]: dtype.append((name, self.dtype[name])) out = numpy.empty(_shape_of(result), dtype=dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]): out[name] = result[i] for name in _coordinate_class_to_names[_ltype(self)]: out[name] = self[name] return out.view(cls.ProjectionClass3D) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): result = _toarrays(result) dtype = [ (name, result[i].dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]) ] out = numpy.empty(_shape_of(result), dtype=dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]): out[name] = result[i] return out.view(cls.ProjectionClass2D) elif len(returns) == 2 or ( (len(returns) == 3 and returns[2] is None) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 return out.view(cls.ProjectionClass3D) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[2]]: dtype.append((name, result[i].dtype)) i += 1 out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[2]]: out[name] = result[i] i += 1 return out.view(cls.ProjectionClass4D) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, False) class MomentumNumpy3D(SpatialMomentum, VectorNumpy3D): """ Three dimensional momentum vector class for the NumPy backend. This class can be directly used to construct three dimensional NumPy momentum vectors. For three dimensional NumPy vectors see :class:`vector.backends.numpy.VectorNumpy3D`. Examples: >>> import vector >>> vec = vector.MomentumNumpy3D([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], ... dtype=[('px', float), ('py', float), ('pz', float)]) >>> vec MomentumNumpy3D([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], dtype=[('x', '>> import numpy as np >>> import vector >>> arr = np.array([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], ... dtype=[('px', float), ('py', float), ('pz', float)]) >>> arr.view(vector.MomentumNumpy3D) MomentumNumpy3D([(1.1, 2.1, 3.1), (1.2, 2.2, 3.2), (1.3, 2.3, 3.3), (1.4, 2.4, 3.4), (1.5, 2.5, 3.5)], dtype=[('x', ' None: if obj is None: return self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): self._azimuthal_type = AzimuthalNumpyRhoPhi else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y") or ("rho", "phi") or ("px", "py") or ("pt", "phi")' ) if _has(self, ("z",)): self._longitudinal_type = LongitudinalNumpyZ elif _has(self, ("theta",)): self._longitudinal_type = LongitudinalNumpyTheta elif _has(self, ("eta",)): self._longitudinal_type = LongitudinalNumpyEta else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "z" or "theta" or "eta" or "pz"' ) _is_type_safe(self) def __repr__(self) -> str: return _array_repr(self, True) def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, True) class VectorNumpy4D(VectorNumpy, Lorentz, Vector4D, FloatArray): # type: ignore[misc] """ Four dimensional vector class for the NumPy backend. This class can be directly used to construct four dimensional NumPy vectors. For four dimensional Momentum NumPy vectors see :class:`vector.backends.numpy.MomentumNumpy4D`. Examples: >>> import vector >>> vec = vector.VectorNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], ... dtype=[('x', float), ('y', float), ('z', float), ('t', float)]) >>> vec VectorNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], dtype=[('x', '>> import numpy as np >>> import vector >>> arr = np.array([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], ... dtype=[('x', float), ('y', float), ('z', float), ('t', float)]) >>> arr.view(vector.VectorNumpy4D) VectorNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], dtype=[('x', ' VectorNumpy4D: """Returns the object of ``VectorNumpy4D``. Behaves as ``__init__`` in this case.""" if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict): array = _array_from_columns(args[0]) else: array = numpy.array(*args, **kwargs) return array.view(cls) def __array_finalize__(self, obj: typing.Any) -> None: if obj is None: return if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): self._azimuthal_type = AzimuthalNumpyRhoPhi else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y") or ("rho", "phi")' ) if _has(self, ("z",)): self._longitudinal_type = LongitudinalNumpyZ elif _has(self, ("theta",)): self._longitudinal_type = LongitudinalNumpyTheta elif _has(self, ("eta",)): self._longitudinal_type = LongitudinalNumpyEta else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "z" or "theta" or "eta"' ) if _has(self, ("t",)): self._temporal_type = TemporalNumpyT elif _has(self, ("tau",)): self._temporal_type = TemporalNumpyTau else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "t" or "tau"' ) _is_type_safe(self) def __str__(self) -> str: return str(self.view(numpy.ndarray)) def __repr__(self) -> str: return _array_repr(self, False) @property def azimuthal(self) -> AzimuthalNumpy: """ Returns the azimuthal type class for the given ``VectorNumpy4D`` object. Example: >>> import vector >>> vec = vector.array( ... [ ... (1.1, 2.1, 3.1, 4.1), ... (1.2, 2.2, 3.2, 4.2), ... (1.3, 2.3, 3.3, 4.3), ... (1.4, 2.4, 3.4, 4.4), ... (1.5, 2.5, 3.5, 4.5) ... ], dtype=[("x", float), ("y", float), ("z", float), ("t", float)] ... ) >>> vec.azimuthal AzimuthalNumpyXY([(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[('x', ' LongitudinalNumpy: """ Returns the longitudinal type class for the given ``VectorNumpy4D`` object. Example: >>> import vector >>> vec = vector.array( ... [ ... (1.1, 2.1, 3.1, 4.1), ... (1.2, 2.2, 3.2, 4.2), ... (1.3, 2.3, 3.3, 4.3), ... (1.4, 2.4, 3.4, 4.4), ... (1.5, 2.5, 3.5, 4.5) ... ], dtype=[("x", float), ("y", float), ("z", float), ("t", float)] ... ) >>> vec.longitudinal LongitudinalNumpyZ([(3.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[('z', ' TemporalNumpy: """ Returns the temporal type class for the given ``VectorNumpy4D`` object. Example: >>> import vector >>> vec = vector.array( ... [ ... (1.1, 2.1, 3.1, 4.1), ... (1.2, 2.2, 3.2, 4.2), ... (1.3, 2.3, 3.3, 4.3), ... (1.4, 2.4, 3.4, 4.4) ... ], dtype=[("x", float), ("y", float), ("z", float), ("t", float)] ... ) >>> vec.temporal TemporalNumpyT([(4.1,), (4.2,), (4.3,), (4.4,)], dtype=[('t', ' typing.Any: """ Wraps the raw result of a compute function as an array of scalars or an array of vectors. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): result = _toarrays(result) dtype = [ (name, result[i].dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]) ] for name in _coordinate_class_to_names[_ltype(self)]: dtype.append((name, self.dtype[name])) for name in _coordinate_class_to_names[_ttype(self)]: dtype.append((name, self.dtype[name])) out = numpy.empty(_shape_of(result), dtype=dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]): out[name] = result[i] for name in _coordinate_class_to_names[_ltype(self)]: out[name] = self[name] for name in _coordinate_class_to_names[_ttype(self)]: out[name] = self[name] return out.view(cls.ProjectionClass4D) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): result = _toarrays(result) dtype = [ (name, result[i].dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]) ] out = numpy.empty(_shape_of(result), dtype=dtype) for i, name in enumerate(_coordinate_class_to_names[returns[0]]): out[name] = result[i] return out.view(cls.ProjectionClass2D) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[_ttype(self)]: dtype.append((name, self.dtype[name])) out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[_ttype(self)]: out[name] = self[name] return out.view(cls.ProjectionClass4D) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and returns[2] is None ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 return out.view(cls.ProjectionClass3D) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): result = _toarrays(result) dtype = [] i = 0 for name in _coordinate_class_to_names[returns[0]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[1]]: dtype.append((name, result[i].dtype)) i += 1 for name in _coordinate_class_to_names[returns[2]]: dtype.append((name, result[i].dtype)) i += 1 out = numpy.empty(_shape_of(result), dtype=dtype) i = 0 for name in _coordinate_class_to_names[returns[0]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[1]]: out[name] = result[i] i += 1 for name in _coordinate_class_to_names[returns[2]]: out[name] = result[i] i += 1 return out.view(cls.ProjectionClass4D) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, False) class MomentumNumpy4D(LorentzMomentum, VectorNumpy4D): """ Four dimensional momentum vector class for the NumPy backend. This class can be directly used to construct four dimensional NumPy momentum vectors. For three dimensional NumPy vectors see :class:`vector.backends.numpy.VectorNumpy4D`. Examples: >>> import vector >>> vec = vector.MomentumNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], ... dtype=[('px', float), ('py', float), ('pz', float), ('t', float)]) >>> vec MomentumNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], dtype=[('x', '>> import numpy as np >>> import vector >>> arr = np.array([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], ... dtype=[('px', float), ('py', float), ('pz', float), ('t', float)]) >>> arr.view(vector.MomentumNumpy4D) MomentumNumpy4D([(1.1, 2.1, 3.1, 4.1), (1.2, 2.2, 3.2, 4.2), (1.3, 2.3, 3.3, 4.3), (1.4, 2.4, 3.4, 4.4), (1.5, 2.5, 3.5, 4.5)], dtype=[('x', ' None: if obj is None: return self.dtype.names = tuple( _repr_momentum_to_generic.get(x, x) for x in (self.dtype.names or ()) ) if _has(self, ("x", "y")): self._azimuthal_type = AzimuthalNumpyXY elif _has(self, ("rho", "phi")): self._azimuthal_type = AzimuthalNumpyRhoPhi else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'fields ("x", "y") or ("rho", "phi") or ("px", "py") or ("pt", "phi")' ) if _has(self, ("z",)): self._longitudinal_type = LongitudinalNumpyZ elif _has(self, ("theta",)): self._longitudinal_type = LongitudinalNumpyTheta elif _has(self, ("eta",)): self._longitudinal_type = LongitudinalNumpyEta else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "z" or "theta" or "eta" or "pz"' ) if _has(self, ("t",)): self._temporal_type = TemporalNumpyT elif _has(self, ("tau",)): self._temporal_type = TemporalNumpyTau else: raise TypeError( f"{type(self).__name__} must have a structured dtype containing " 'field "t" or "tau" or "E" or "e" or "energy" or "M" or "m" or "mass"' ) _is_type_safe(self) def __repr__(self) -> str: return _array_repr(self, True) def __setitem__(self, where: typing.Any, what: typing.Any) -> None: return _setitem(self, where, what, True) def array(*args: typing.Any, **kwargs: typing.Any) -> VectorNumpy: """ Constructs a NumPy array of vectors, whose type is determined by the dtype of the structured array or Pandas-style "columns" argument. All allowed signatures for ``np.array`` can be used in this function, plus one more: .. code-block:: python vector.array({"x": x_column, "y": y_column}) to make an array with ``dtype=[("x", x_column.dtype), ("y", y_column.dtype)]``. The array must have structured ``dtype`` (i.e. ``dtype.names is not None``) and the following combinations of names are allowed: - (2D) ``x``, ``y`` - (2D) ``rho``, ``phi`` - (3D) ``x``, ``y``, ``z`` - (3D) ``x``, ``y``, ``theta`` - (3D) ``x``, ``y``, ``eta`` - (3D) ``rho``, ``phi``, ``z`` - (3D) ``rho``, ``phi``, ``theta`` - (3D) ``rho``, ``phi``, ``eta`` - (4D) ``x``, ``y``, ``z``, ``t`` - (4D) ``x``, ``y``, ``z``, ``tau``` - (4D) ``x``, ``y``, ``theta``, ``t``` - (4D) ``x``, ``y``, ``theta``, ``tau``` - (4D) ``x``, ``y``, ``eta``, ``t``` - (4D) ``x``, ``y``, ``eta``, ``tau``` - (4D) ``rho``, ``phi``, ``z``, ``t``` - (4D) ``rho``, ``phi``, ``z``, ``tau``` - (4D) ``rho``, ``phi``, ``theta``, ``t``` - (4D) ``rho``, ``phi``, ``theta``, ``tau``` - (4D) ``rho``, ``phi``, ``eta``, ``t``` - (4D) ``rho``, ``phi``, ``eta``, ``tau``` in which - ``px`` may be substituted for ``x`` - ``py`` may be substituted for ``y`` - ``pt`` may be substituted for ``rho`` - ``pz`` may be substituted for ``z`` - ``E`` may be substituted for ``t`` - ``e`` may be substituted for ``t`` - ``energy`` may be substituted for ``t`` - ``M`` may be substituted for ``tau`` - ``m`` may be substituted for ``tau`` - ``mass`` may be substituted for ``tau`` to make the vector a momentum vector. """ names: tuple[str, ...] if len(args) == 1 and len(kwargs) == 0 and isinstance(args[0], dict): names = tuple(args[0].keys()) elif "dtype" in kwargs: names = numpy.dtype(kwargs["dtype"]).names or () elif len(args) >= 2: names = numpy.dtype(args[1]).names or () else: names = () cls: type[VectorNumpy] is_momentum = any(x in _repr_momentum_to_generic for x in names) if any(x in ("t", "E", "e", "energy", "tau", "M", "m", "mass") for x in names): cls = MomentumNumpy4D if is_momentum else VectorNumpy4D elif any(x in ("z", "pz", "theta", "eta") for x in names): cls = MomentumNumpy3D if is_momentum else VectorNumpy3D else: cls = MomentumNumpy2D if is_momentum else VectorNumpy2D return cls(*args, **kwargs) VectorNumpy2D.ProjectionClass2D = VectorNumpy2D VectorNumpy2D.ProjectionClass3D = VectorNumpy3D VectorNumpy2D.ProjectionClass4D = VectorNumpy4D VectorNumpy2D.GenericClass = VectorNumpy2D VectorNumpy2D.MomentumClass = MomentumNumpy2D MomentumNumpy2D.ProjectionClass2D = MomentumNumpy2D MomentumNumpy2D.ProjectionClass3D = MomentumNumpy3D MomentumNumpy2D.ProjectionClass4D = MomentumNumpy4D MomentumNumpy2D.GenericClass = VectorNumpy2D MomentumNumpy2D.MomentumClass = MomentumNumpy2D VectorNumpy3D.ProjectionClass2D = VectorNumpy2D VectorNumpy3D.ProjectionClass3D = VectorNumpy3D VectorNumpy3D.ProjectionClass4D = VectorNumpy4D VectorNumpy3D.GenericClass = VectorNumpy3D VectorNumpy3D.MomentumClass = MomentumNumpy3D MomentumNumpy3D.ProjectionClass2D = MomentumNumpy2D MomentumNumpy3D.ProjectionClass3D = MomentumNumpy3D MomentumNumpy3D.ProjectionClass4D = MomentumNumpy4D MomentumNumpy3D.GenericClass = VectorNumpy3D MomentumNumpy3D.MomentumClass = MomentumNumpy3D VectorNumpy4D.ProjectionClass2D = VectorNumpy2D VectorNumpy4D.ProjectionClass3D = VectorNumpy3D VectorNumpy4D.ProjectionClass4D = VectorNumpy4D VectorNumpy4D.GenericClass = VectorNumpy4D VectorNumpy4D.MomentumClass = MomentumNumpy4D MomentumNumpy4D.ProjectionClass2D = MomentumNumpy2D MomentumNumpy4D.ProjectionClass3D = MomentumNumpy3D MomentumNumpy4D.ProjectionClass4D = MomentumNumpy4D MomentumNumpy4D.GenericClass = VectorNumpy4D MomentumNumpy4D.MomentumClass = MomentumNumpy4D vector-1.6.3/src/vector/backends/object.py000066400000000000000000003052001503546127100205120ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Defines behaviors for Object vectors. New vectors created with the respective classes .. code-block:: python vector.VectorObject2D(...) vector.VectorObject3D(...) will have these behaviors built in (and will pass them to any derived objects). The class methods can also be used to construct object type vectors - .. code-block:: python vec = vector.VectorObject2D.from_xy(1, 2) Additionally, object type vectors can also be constructed using - .. code-block:: python vector.obj(...) function. """ from __future__ import annotations import numbers import typing import numpy from vector._methods import ( Azimuthal, AzimuthalRhoPhi, AzimuthalXY, Longitudinal, LongitudinalEta, LongitudinalTheta, LongitudinalZ, Lorentz, LorentzMomentum, Planar, PlanarMomentum, SameVectorType, Spatial, SpatialMomentum, Temporal, TemporalT, TemporalTau, Vector, Vector2D, Vector3D, Vector4D, VectorProtocol, _aztype, _coordinate_class_to_names, _handler_of, _ltype, _repr_generic_to_momentum, _repr_momentum_to_generic, _ttype, ) from vector._typeutils import FloatArray class CoordinatesObject: """Coordinates class for the Object backend.""" class AzimuthalObject(CoordinatesObject, Azimuthal): """Azimuthal class for the Object backend.""" class LongitudinalObject(CoordinatesObject, Longitudinal): """Longitudinal class for the Object backend.""" class TemporalObject(CoordinatesObject, Temporal): """Temporal class for the Object backend.""" class TupleXY(typing.NamedTuple): """``x`` and ``y`` coordinates as a ``NamedTuple``.""" x: float y: float class AzimuthalObjectXY(AzimuthalObject, AzimuthalXY, TupleXY): """ Class for the ``x`` and ``y`` (azimuthal) coordinates of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float, float]: """ Azimuthal coordinates (``x`` and ``y``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.VectorObject2D(x=1, y=2) >>> az = v.azimuthal >>> az.elements (1, 2) """ return (self.x, self.y) class TupleRhoPhi(typing.NamedTuple): """``rho`` and ``phi`` coordinates as a ``NamedTuple``.""" rho: float phi: float class AzimuthalObjectRhoPhi(AzimuthalObject, AzimuthalRhoPhi, TupleRhoPhi): """ Class for the ``rho`` and ``phi`` (azimuthal) coordinates of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float, float]: """ Azimuthal coordinates (``rho`` and ``phi``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.VectorObject2D(rho=1, phi=2) >>> az = v.azimuthal >>> az.elements (1, 2) """ return (self.rho, self.phi) class TupleZ(typing.NamedTuple): """``z`` coordinate as a ``NamedTuple``.""" z: float class LongitudinalObjectZ(LongitudinalObject, LongitudinalZ, TupleZ): """ Class for the ``z`` (longitudinal) coordinate of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float]: """ Longitudinal coordinates (``z``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.obj(rho=1, phi=2, z=3) >>> lc = v.longitudinal >>> lc.elements (3,) """ return (self.z,) class TupleTheta(typing.NamedTuple): """``theta`` coordinates as a ``NamedTuple``.""" theta: float class LongitudinalObjectTheta(LongitudinalObject, LongitudinalTheta, TupleTheta): """ Class for the ``theta`` (longitudinal) coordinate of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float]: """ Longitudinal coordinates (``theta``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.obj(rho=1, phi=2, theta=3) >>> lc = v.longitudinal >>> lc.elements (3,) """ return (self.theta,) class TupleEta(typing.NamedTuple): """``eta`` coordinate as a ``NamedTuple``.""" eta: float class LongitudinalObjectEta(LongitudinalObject, LongitudinalEta, TupleEta): """ Class for the ``eta`` (longitudinal) coordinate of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float]: """ Longitudinal coordinates (``theta``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.obj(rho=1, phi=2, eta=3) >>> lc = v.longitudinal >>> lc.elements (3,) """ return (self.eta,) class TupleT(typing.NamedTuple): """``t`` coordinate as a ``NamedTuple``.""" t: float class TemporalObjectT(TemporalObject, TemporalT, TupleT): """ Class for the ``t`` (temporal) coordinate of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float]: """ Temporal coordinates (``t``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.obj(rho=1, phi=2, theta=3, t=4) >>> tc = v.temporal >>> tc.elements (4,) """ return (self.t,) class TupleTau(typing.NamedTuple): """``tau`` coordinate as a ``NamedTuple``.""" tau: float class TemporalObjectTau(TemporalObject, TemporalTau, TupleTau): """ Class for the ``tau`` (temporal) coordinate of Object backend. Use the ``elements`` property to retrieve the coordinates. """ @property def elements(self) -> tuple[float]: """ Temporal coordinates (``tau``) as a tuple. Each coordinate is a scalar and not a vector. Examples: >>> import vector >>> v = vector.obj(rho=1, phi=2, theta=3, tau=4) >>> tc = v.temporal >>> tc.elements (4,) """ return (self.tau,) _coord_object_type = { AzimuthalXY: AzimuthalObjectXY, AzimuthalRhoPhi: AzimuthalObjectRhoPhi, LongitudinalZ: LongitudinalObjectZ, LongitudinalTheta: LongitudinalObjectTheta, LongitudinalEta: LongitudinalObjectEta, TemporalT: TemporalObjectT, TemporalTau: TemporalObjectTau, } def _replace_data(obj: typing.Any, result: typing.Any) -> typing.Any: if not isinstance(result, VectorObject): raise TypeError(f"can only assign a single vector to {type(obj).__name__}") if isinstance(result, (VectorObject2D, VectorObject3D, VectorObject4D)): if isinstance(obj.azimuthal, AzimuthalObjectXY): obj.azimuthal = AzimuthalObjectXY(result.x, result.y) elif isinstance(obj.azimuthal, AzimuthalObjectRhoPhi): obj.azimuthal = AzimuthalObjectRhoPhi(result.rho, result.phi) else: raise AssertionError(type(obj)) if isinstance(result, (VectorObject3D, VectorObject4D)): if isinstance(obj.longitudinal, LongitudinalObjectZ): obj.longitudinal = LongitudinalObjectZ(result.z) elif isinstance(obj.longitudinal, LongitudinalObjectTheta): obj.longitudinal = LongitudinalObjectTheta(result.theta) elif isinstance(obj.longitudinal, LongitudinalObjectEta): obj.longitudinal = LongitudinalObjectEta(result.eta) else: raise AssertionError(type(obj)) if isinstance(result, VectorObject4D): if isinstance(obj.temporal, TemporalObjectT): obj.temporal = TemporalObjectT(result.t) elif isinstance(obj.temporal, TemporalObjectTau): obj.temporal = TemporalObjectTau(result.tau) else: raise AssertionError(type(obj)) return obj class VectorObject(Vector): # noqa: PLW1641 """Mixin class for Object vectors.""" lib = numpy def __eq__(self, other: typing.Any) -> typing.Any: return numpy.equal(self, other) def __ne__(self, other: typing.Any) -> typing.Any: return numpy.not_equal(self, other) def __abs__(self) -> float: return numpy.absolute(self) def __add__(self, other: VectorProtocol) -> VectorProtocol: return numpy.add(self, other) def __radd__(self, other: VectorProtocol) -> VectorProtocol: return numpy.add(other, self) def __iadd__(self: SameVectorType, other: VectorProtocol) -> SameVectorType: return _replace_data(self, numpy.add(self, other)) def __sub__(self, other: VectorProtocol) -> VectorProtocol: return numpy.subtract(self, other) def __rsub__(self, other: VectorProtocol) -> VectorProtocol: return numpy.subtract(other, self) def __isub__(self: SameVectorType, other: VectorProtocol) -> SameVectorType: return _replace_data(self, numpy.subtract(self, other)) def __mul__(self, other: float) -> VectorProtocol: return numpy.multiply(self, other) def __rmul__(self, other: float) -> VectorProtocol: return numpy.multiply(other, self) def __imul__(self: SameVectorType, other: float) -> SameVectorType: return _replace_data(self, numpy.multiply(self, other)) def __neg__(self: SameVectorType) -> SameVectorType: return numpy.negative(self) def __pos__(self: SameVectorType) -> SameVectorType: return numpy.positive(self) def __truediv__(self, other: float) -> VectorProtocol: return numpy.true_divide(self, other) def __rtruediv__(self, other: float) -> VectorProtocol: return numpy.true_divide(other, self) def __itruediv__(self: SameVectorType, other: float) -> VectorProtocol: return _replace_data(self, numpy.true_divide(self, other)) def __pow__(self, other: float) -> float: return numpy.square(self) if other == 2 else numpy.power(self, other) def __matmul__(self, other: VectorProtocol) -> float: return numpy.matmul(self, other) def __array_ufunc__( self, ufunc: typing.Any, method: typing.Any, *inputs: typing.Any, **kwargs: typing.Any, ) -> typing.Any: """ Implements NumPy's ``ufunc``s for ``VectorObject`` and its subclasses. The current implementation includes ``numpy.absolute``, ``numpy.add``, ``numpy.subtract``, ``numpy.multiply``, ``numpy.positive``, ``numpy.negative``, ``numpy.true_divide``, ``numpy.power``, ``numpy.square``, ``numpy.sqrt``, ``numpy.cbrt``, ``numpy.matmul``, ``numpy.equal``, and ``numpy.not_equal``. """ if not isinstance(_handler_of(*inputs), VectorObject): # Let a higher-precedence backend handle it. return NotImplemented outputs = kwargs.get("out", ()) if any(not isinstance(x, VectorObject) for x in outputs): raise TypeError( "ufunc operating on VectorObjects can only use the 'out' keyword " "with another VectorObject" ) if ( ufunc is numpy.absolute and len(inputs) == 1 and isinstance(inputs[0], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.absolute' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho elif isinstance(inputs[0], Vector3D): return inputs[0].mag elif isinstance(inputs[0], Vector4D): return inputs[0].tau elif ( ufunc is numpy.add and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[0].add(inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.subtract and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[0].subtract(inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.multiply and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = inputs[0].scale(inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.multiply and len(inputs) == 2 and not isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[1].scale(inputs[0]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.negative and len(inputs) == 1 and isinstance(inputs[0], Vector) ): result = inputs[0].scale(-1) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.positive and len(inputs) == 1 and isinstance(inputs[0], Vector) ): return inputs[0] elif ( ufunc is numpy.true_divide and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = inputs[0].scale(1 / inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.power and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = numpy.absolute(inputs[0]) ** inputs[1] for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.square and len(inputs) == 1 and isinstance(inputs[0], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.square' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 elif ufunc is numpy.sqrt and len(inputs) == 1 and isinstance(inputs[0], Vector): if len(outputs) != 0: raise TypeError( "output of 'numpy.sqrt' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 ** 0.25 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 ** 0.25 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 ** 0.25 elif ufunc is numpy.cbrt and len(inputs) == 1 and isinstance(inputs[0], Vector): if len(outputs) != 0: raise TypeError( "output of 'numpy.cbrt' is scalar, cannot fill a VectorObject with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 ** 0.16666666666666666 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 ** 0.16666666666666666 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 ** 0.16666666666666666 elif ( ufunc is numpy.matmul and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.matmul' is scalar, cannot fill a VectorObject with 'out'" ) return inputs[0].dot(inputs[1]) elif ( ufunc is numpy.equal and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.equal' is scalar, cannot fill a VectorObject with 'out'" ) return inputs[0].equal(inputs[1]) elif ( ufunc is numpy.not_equal and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.equal' is scalar, cannot fill a VectorObject with 'out'" ) return inputs[0].not_equal(inputs[1]) else: return NotImplemented class VectorObject2D(VectorObject, Planar, Vector2D): """ Two dimensional vector class for the object backend. Examples: >>> import vector >>> vec = vector.VectorObject2D(x=1, y=2) >>> vec.x, vec.y (1, 2) >>> vec = vector.VectorObject2D(rho=1, phi=2) >>> vec.rho, vec.phi (1, 2) >>> vec = vector.VectorObject2D(azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2)) >>> vec.x, vec.y (1, 2) The following class methods can also be used to construct 2D object type vectors - - :meth:`VectorObject2D.from_xy` - :meth:`VectorObject2D.from_rhophi` Additionally, the :func:`vector.obj` function can also be used to construct 2D object type vectors. For two dimensional momentum vector objects, see :class:`vector.backends.object.MomentumObject2D`. """ __slots__ = ("azimuthal",) azimuthal: AzimuthalObject @classmethod def from_xy(cls, x: float, y: float) -> VectorObject2D: """ Constructs a ``VectorObject2D`` from Cartesian coordinates. Use :class:`vector.backends.object.MomentumObject2D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject2D.from_xy(1, 2) >>> vec VectorObject2D(x=1, y=2) """ _is_type_safe(locals()) return cls(azimuthal=AzimuthalObjectXY(x, y)) @classmethod def from_rhophi(cls, rho: float, phi: float) -> VectorObject2D: """ Constructs a ``VectorObject2D`` from polar coordinates. Use :class:`vector.backends.object.MomentumObject2D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject2D.from_rhophi(1, 2) >>> vec VectorObject2D(rho=1, phi=2) """ _is_type_safe(locals()) return cls(azimuthal=AzimuthalObjectRhoPhi(rho, phi)) def __init__( self, azimuthal: AzimuthalObject | None = None, **kwargs: float ) -> None: _is_type_safe(kwargs) for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v if not kwargs and azimuthal is not None: self.azimuthal = azimuthal elif kwargs and azimuthal is None: if set(kwargs) == {"x", "y"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) elif set(kwargs) == {"rho", "phi"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) else: complaint = """unrecognized combination of coordinates, allowed combinations are:\n x= y= rho= phi=""".replace(" ", " ") if type(self) is VectorObject2D: raise TypeError(complaint) else: raise TypeError(f"{complaint}\n\nor their momentum equivalents") else: raise TypeError("must give Azimuthal if not giving keyword arguments") def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] out = [f"{x}={getattr(self.azimuthal, x)}" for x in aznames] return "VectorObject2D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import VectorNumpy2D return VectorNumpy2D( self.azimuthal.elements, dtype=[ (x, numpy.float64) for x in _coordinate_class_to_names[_aztype(self)] ], ) def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as a scalar or a vector. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( (len(returns) == 1 or (len(returns) == 2 and returns[1] is None)) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) return cls.ProjectionClass2D(azimuthal=azcoords) elif len(returns) == 2 or ( (len(returns) == 3 and returns[2] is None) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) return cls.ProjectionClass3D(azimuthal=azcoords, longitudinal=lcoords) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) tcoords = _coord_object_type[returns[2]](result[3]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=tcoords ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func @property def x(self) -> float: return super().x @x.setter def x(self, x: float) -> None: self.azimuthal = AzimuthalObjectXY(x, self.y) @property def y(self) -> float: return super().y @y.setter def y(self, y: float) -> None: self.azimuthal = AzimuthalObjectXY(self.x, y) @property def rho(self) -> float: return super().rho @rho.setter def rho(self, rho: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(rho, self.phi) @property def phi(self) -> float: return super().phi @phi.setter def phi(self, phi: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(self.rho, phi) class MomentumObject2D(PlanarMomentum, VectorObject2D): """ Two dimensional momentum vector class for the object backend. Examples: >>> import vector >>> vec = vector.MomentumObject2D(px=1, py=2) >>> vec.px, vec.py (1, 2) >>> vec = vector.MomentumObject2D(pt=1, phi=2) >>> vec.pt, vec.phi (1, 2) >>> vec = vector.MomentumObject2D(azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2)) >>> vec.px, vec.py (1, 2) The :func:`vector.obj` function can also be used to construct 2D momentum object type vectors. For two dimensional vector objects, see :class:`vector.backends.object.VectorObject2D`. """ def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] out = [] for x in aznames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.azimuthal, x)}") return "MomentumObject2D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import MomentumNumpy2D return MomentumNumpy2D( self.azimuthal.elements, dtype=[ (x, numpy.float64) for x in _coordinate_class_to_names[_aztype(self)] ], ) @property def px(self) -> float: return super().px @px.setter def px(self, px: float) -> None: self.azimuthal = AzimuthalObjectXY(px, self.py) @property def py(self) -> float: return super().py @py.setter def py(self, py: float) -> None: self.azimuthal = AzimuthalObjectXY(self.px, py) @property def pt(self) -> float: return super().pt @pt.setter def pt(self, pt: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(pt, self.phi) class VectorObject3D(VectorObject, Spatial, Vector3D): """ Three dimensional vector class for the object backend. Use the class methods - Examples: >>> import vector >>> vec = vector.VectorObject3D(x=1, y=2, z=3) >>> vec.x, vec.y, vec.z (1, 2, 3) >>> vec = vector.VectorObject3D(rho=1, phi=2, eta=3) >>> vec.rho, vec.phi, vec.eta (1, 2, 3) >>> vec = vector.VectorObject3D( ... azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), ... longitudinal=vector.backends.object.LongitudinalObjectTheta(3) ... ) >>> vec.x, vec.y, vec.theta (1, 2, 3) The following class methods can also be used to construct 3D object type vectors - - :meth:`VectorObject3D.from_xyz` - :meth:`VectorObject3D.from_xytheta` - :meth:`VectorObject3D.from_xyeta` - :meth:`VectorObject3D.from_rhophiz` - :meth:`VectorObject3D.from_rhophitheta` - :meth:`VectorObject3D.from_rhophieta` Additionally, the :func:`vector.obj` function can also be used to construct 3D object type vectors. For three dimensional momentum vector objects, see :class:`vector.backends.object.MomentumObject3D`. """ __slots__ = ("azimuthal", "longitudinal") azimuthal: AzimuthalObject longitudinal: LongitudinalObject @classmethod def from_xyz(cls, x: float, y: float, z: float) -> VectorObject3D: """ Constructs a ``VectorObject3D`` from Cartesian coordinates. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject3D.from_xyz(1, 1, 1) >>> vec VectorObject3D(x=1, y=1, z=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectZ(z) ) @classmethod def from_xytheta(cls, x: float, y: float, theta: float) -> VectorObject3D: r""" Constructs a ``VectorObject3D`` from Cartesian azimuthal coordinates and a polar angle $\theta$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject3D.from_xytheta(1, 1, 1) >>> vec VectorObject3D(x=1, y=1, theta=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectTheta(theta), ) @classmethod def from_xyeta(cls, x: float, y: float, eta: float) -> VectorObject3D: r""" Constructs a ``VectorObject3D`` from Cartesian coordinates and a pseudorapidity $\eta$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject3D.from_xyeta(1, 1, 1) >>> vec VectorObject3D(x=1, y=1, eta=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectEta(eta) ) @classmethod def from_rhophiz(cls, rho: float, phi: float, z: float) -> VectorObject3D: """ Constructs a ``VectorObject3D`` from polar azimuthal coordinates and a Cartesian longitudinal coordinate $z$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject3D.from_rhophiz(1, 1, 1) >>> vec VectorObject3D(rho=1, phi=1, z=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectZ(z), ) @classmethod def from_rhophitheta(cls, rho: float, phi: float, theta: float) -> VectorObject3D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates and a polar angle $\theta$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject3D.from_rhophitheta(1, 1, 1) >>> vec VectorObject3D(rho=1, phi=1, theta=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectTheta(theta), ) @classmethod def from_rhophieta(cls, rho: float, phi: float, eta: float) -> VectorObject3D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates and a pseudorapidity $\eta$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject3D.from_rhophieta(1, 1, 1) >>> vec VectorObject3D(rho=1, phi=1, eta=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectEta(eta), ) def __init__( self, azimuthal: AzimuthalObject | None = None, longitudinal: LongitudinalObject | None = None, **kwargs: float, ) -> None: _is_type_safe(kwargs) for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v if not kwargs and azimuthal is not None and longitudinal is not None: self.azimuthal = azimuthal self.longitudinal = longitudinal elif kwargs and azimuthal is None and longitudinal is None: if set(kwargs) == {"x", "y", "z"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectZ(kwargs["z"]) elif set(kwargs) == {"x", "y", "eta"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) elif set(kwargs) == {"x", "y", "theta"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) elif set(kwargs) == {"rho", "phi", "z"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectZ(kwargs["z"]) elif set(kwargs) == {"rho", "phi", "eta"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) elif set(kwargs) == {"rho", "phi", "theta"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) else: complaint = """unrecognized combination of coordinates, allowed combinations are:\n x= y= z= x= y= theta= x= y= eta= rho= phi= z= rho= phi= theta= rho= phi= eta=""".replace(" ", " ") if type(self) is VectorObject3D: raise TypeError(complaint) else: raise TypeError(f"{complaint}\n\nor their momentum equivalents") else: raise TypeError( "must give Azimuthal and Longitudinal if not giving keyword arguments" ) def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] out = [f"{x}={getattr(self.azimuthal, x)}" for x in aznames] for x in lnames: out.append(f"{x}={getattr(self.longitudinal, x)}") return "VectorObject3D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import VectorNumpy3D return VectorNumpy3D( self.azimuthal.elements + self.longitudinal.elements, dtype=[ (x, numpy.float64) for x in _coordinate_class_to_names[_aztype(self)] + _coordinate_class_to_names[_ltype(self)] ], ) def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as a scalar or a vector. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) return cls.ProjectionClass3D( azimuthal=azcoords, longitudinal=self.longitudinal ) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) return cls.ProjectionClass2D(azimuthal=azcoords) elif len(returns) == 2 or ( (len(returns) == 3 and returns[2] is None) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) return cls.ProjectionClass3D(azimuthal=azcoords, longitudinal=lcoords) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) tcoords = _coord_object_type[returns[2]](result[3]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=tcoords ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func @property def x(self) -> float: return super().x @x.setter def x(self, x: float) -> None: self.azimuthal = AzimuthalObjectXY(x, self.y) @property def y(self) -> float: return super().y @y.setter def y(self, y: float) -> None: self.azimuthal = AzimuthalObjectXY(self.x, y) @property def rho(self) -> float: return super().rho @rho.setter def rho(self, rho: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(rho, self.phi) @property def phi(self) -> float: return super().phi @phi.setter def phi(self, phi: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(self.rho, phi) @property def z(self) -> float: return super().z @z.setter def z(self, z: float) -> None: self.longitudinal = LongitudinalObjectZ(z) @property def theta(self) -> float: return super().theta @theta.setter def theta(self, theta: float) -> None: self.longitudinal = LongitudinalObjectTheta(theta) @property def eta(self) -> float: return super().eta @eta.setter def eta(self, eta: float) -> None: self.longitudinal = LongitudinalObjectEta(eta) class MomentumObject3D(SpatialMomentum, VectorObject3D): """ Three dimensional momentum vector class for the object backend. Examples: >>> import vector >>> vec = vector.MomentumObject3D(px=1, py=2, pz=3) >>> vec.px, vec.py, vec.pz (1, 2, 3) >>> vec = vector.MomentumObject3D(pt=1, phi=2, pz=3) >>> vec.pt, vec.phi, vec.pz (1, 2, 3) >>> vec = vector.MomentumObject3D( ... azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), ... longitudinal=vector.backends.object.LongitudinalObjectTheta(3) ... ) >>> vec.x, vec.y, vec.theta (1, 2, 3) The :func:`vector.obj` function can also be used to construct 3D momentum object type vectors. For three dimensional vector objects, see :class:`vector.backends.object.VectorObject3D`. """ def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] out = [] for x in aznames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.azimuthal, x)}") for x in lnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.longitudinal, x)}") return "MomentumObject3D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import MomentumNumpy3D return MomentumNumpy3D( self.azimuthal.elements + self.longitudinal.elements, dtype=[ (x, numpy.float64) for x in _coordinate_class_to_names[_aztype(self)] + _coordinate_class_to_names[_ltype(self)] ], ) @property def px(self) -> float: return super().px @px.setter def px(self, px: float) -> None: self.azimuthal = AzimuthalObjectXY(px, self.py) @property def py(self) -> float: return super().py @py.setter def py(self, py: float) -> None: self.azimuthal = AzimuthalObjectXY(self.px, py) @property def pt(self) -> float: return super().pt @pt.setter def pt(self, pt: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(pt, self.phi) @property def pz(self) -> float: return super().pz @pz.setter def pz(self, pz: float) -> None: self.longitudinal = LongitudinalObjectZ(pz) class VectorObject4D(VectorObject, Lorentz, Vector4D): """ Four dimensional vector class for the object backend. Use the class methods - Examples: >>> import vector >>> vec = vector.VectorObject4D(x=1, y=2, z=3, t=4) >>> vec.x, vec.y, vec.z, vec.t (1, 2, 3, 4) >>> vec = vector.VectorObject4D(rho=1, phi=2, eta=3, tau=4) >>> vec.rho, vec.phi, vec.eta, vec.tau (1, 2, 3, 4) >>> vec = vector.VectorObject4D( ... azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), ... longitudinal=vector.backends.object.LongitudinalObjectTheta(3), ... temporal=vector.backends.object.TemporalObjectTau(4) ... ) >>> vec.x, vec.y, vec.theta, vec.tau (1, 2, 3, 4) The following class methods can also be used to construct 4D object type vectors - - :meth:`VectorObject4D.from_xyzt` - :meth:`VectorObject4D.from_xythetat` - :meth:`VectorObject4D.from_xyetat` - :meth:`VectorObject4D.from_rhophizt` - :meth:`VectorObject4D.from_rhophithetat` - :meth:`VectorObject4D.from_rhophietat` - :meth:`VectorObject4D.from_xyztau` - :meth:`VectorObject4D.from_xythetatau` - :meth:`VectorObject4D.from_xyetatau` - :meth:`VectorObject4D.from_rhophiztau` - :meth:`VectorObject4D.from_rhophithetatau` - :meth:`VectorObject4D.from_rhophietatau` Additionally, the :func:`vector.obj` function can also be used to construct 4D object type vectors. For four dimensional momentum vector objects, see :class:`vector.backends.object.MomentumObject4D`. """ __slots__ = ("azimuthal", "longitudinal", "temporal") azimuthal: AzimuthalObject longitudinal: LongitudinalObject temporal: TemporalObject @classmethod def from_xyzt( cls, x: float, y: float, z: float, t: float, ) -> VectorObject4D: """ Constructs a ``VectorObject3D`` from Cartesian coordinates and a time coordinate $t$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_xyzt(1, 1, 1, 1) >>> vec VectorObject4D(x=1, y=1, z=1, t=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectZ(z), temporal=TemporalObjectT(t), ) @classmethod def from_xyztau( cls, x: float, y: float, z: float, tau: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from Cartesian coordinates and a proper time coordinate $\tau$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_xyztau(1, 1, 1, 1) >>> vec VectorObject4D(x=1, y=1, z=1, tau=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectZ(z), temporal=TemporalObjectTau(tau), ) @classmethod def from_xythetat( cls, x: float, y: float, theta: float, t: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from Cartesian azimuthal coordinates, a polar angle $\theta$, and a time coordinate $t$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_xythetat(1, 1, 1, 1) >>> vec VectorObject4D(x=1, y=1, theta=1, t=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectTheta(theta), temporal=TemporalObjectT(t), ) @classmethod def from_xythetatau( cls, x: float, y: float, theta: float, tau: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from Cartesian azimuthal coordinates, a polar angle $\theta$, and a proper time coordinate $\tau$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_xythetatau(1, 1, 1, 1) >>> vec VectorObject4D(x=1, y=1, theta=1, tau=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectTheta(theta), temporal=TemporalObjectTau(tau), ) @classmethod def from_xyetat( cls, x: float, y: float, eta: float, t: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from Cartesian coordinates, a pseudorapidity $\eta$, and a time coordinate $t$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_xyetat(1, 1, 1, 1) >>> vec VectorObject4D(x=1, y=1, eta=1, t=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectEta(eta), temporal=TemporalObjectT(t), ) @classmethod def from_xyetatau( cls, x: float, y: float, eta: float, tau: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from Cartesian coordinates, a pseudorapidity $\eta$, and a proper time coordinate $\tau$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_xyetatau(1, 1, 1, 1) >>> vec VectorObject4D(x=1, y=1, eta=1, tau=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectXY(x, y), longitudinal=LongitudinalObjectEta(eta), temporal=TemporalObjectTau(tau), ) @classmethod def from_rhophizt( cls, rho: float, phi: float, z: float, t: float, ) -> VectorObject4D: """ Constructs a ``VectorObject3D`` from polar azimuthal coordinates, a Cartesian longitudinal coordinate $z$, and a time coordinate $t$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_rhophizt(1, 1, 1, 1) >>> vec VectorObject4D(rho=1, phi=1, z=1, t=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectZ(z), temporal=TemporalObjectT(t), ) @classmethod def from_rhophiztau( cls, rho: float, phi: float, z: float, tau: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates, a Cartesian longitudinal coordinate $z$, and a proper time coordinate $\tau$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_rhophiztau(1, 1, 1, 1) >>> vec VectorObject4D(rho=1, phi=1, z=1, tau=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectZ(z), temporal=TemporalObjectTau(tau), ) @classmethod def from_rhophithetat( cls, rho: float, phi: float, theta: float, t: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates, a polar angle $\theta$, and a time coordinate $t$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_rhophithetat(1, 1, 1, 1) >>> vec VectorObject4D(rho=1, phi=1, theta=1, t=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectTheta(theta), temporal=TemporalObjectT(t), ) @classmethod def from_rhophithetatau( cls, rho: float, phi: float, theta: float, tau: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates, a polar angle $\theta$, and a proper time coordinate $\tau$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_rhophithetatau(1, 1, 1, 1) >>> vec VectorObject4D(rho=1, phi=1, theta=1, tau=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectTheta(theta), temporal=TemporalObjectTau(tau), ) @classmethod def from_rhophietat( cls, rho: float, phi: float, eta: float, t: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates, a pseudorapidity $\eta$, and a time coordinate $t$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_rhophietat(1, 1, 1, 1) >>> vec VectorObject4D(rho=1, phi=1, eta=1, t=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectEta(eta), temporal=TemporalObjectT(t), ) @classmethod def from_rhophietatau( cls, rho: float, phi: float, eta: float, tau: float, ) -> VectorObject4D: r""" Constructs a ``VectorObject3D`` from polar azimuthal coordinates, a pseudorapidity $\eta$, and a proper time coordinate $\tau$. Use :class:`vector.backends.object.MomentumObject3D` to construct a vector with momentum properties and methods. Examples: >>> import vector >>> vec = vector.VectorObject4D.from_rhophietatau(1, 1, 1, 1) >>> vec VectorObject4D(rho=1, phi=1, eta=1, tau=1) """ _is_type_safe(locals()) return cls( azimuthal=AzimuthalObjectRhoPhi(rho, phi), longitudinal=LongitudinalObjectEta(eta), temporal=TemporalObjectTau(tau), ) def __init__( self, azimuthal: AzimuthalObject | None = None, longitudinal: LongitudinalObject | None = None, temporal: TemporalObject | None = None, **kwargs: float, ) -> None: _is_type_safe(kwargs) for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v if ( not kwargs and azimuthal is not None and longitudinal is not None and temporal is not None ): self.azimuthal = azimuthal self.longitudinal = longitudinal self.temporal = temporal elif kwargs and azimuthal is None and longitudinal is None and temporal is None: if set(kwargs) == {"x", "y", "z", "t"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectZ(kwargs["z"]) self.temporal = TemporalObjectT(kwargs["t"]) elif set(kwargs) == {"x", "y", "eta", "t"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) self.temporal = TemporalObjectT(kwargs["t"]) elif set(kwargs) == {"x", "y", "theta", "t"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) self.temporal = TemporalObjectT(kwargs["t"]) elif set(kwargs) == {"rho", "phi", "z", "t"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectZ(kwargs["z"]) self.temporal = TemporalObjectT(kwargs["t"]) elif set(kwargs) == {"rho", "phi", "eta", "t"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) self.temporal = TemporalObjectT(kwargs["t"]) elif set(kwargs) == {"rho", "phi", "theta", "t"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) self.temporal = TemporalObjectT(kwargs["t"]) elif set(kwargs) == {"x", "y", "z", "tau"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectZ(kwargs["z"]) self.temporal = TemporalObjectTau(kwargs["tau"]) elif set(kwargs) == {"x", "y", "eta", "tau"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) self.temporal = TemporalObjectTau(kwargs["tau"]) elif set(kwargs) == {"x", "y", "theta", "tau"}: self.azimuthal = AzimuthalObjectXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) self.temporal = TemporalObjectTau(kwargs["tau"]) elif set(kwargs) == {"rho", "phi", "z", "tau"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectZ(kwargs["z"]) self.temporal = TemporalObjectTau(kwargs["tau"]) elif set(kwargs) == {"rho", "phi", "eta", "tau"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectEta(kwargs["eta"]) self.temporal = TemporalObjectTau(kwargs["tau"]) elif set(kwargs) == {"rho", "phi", "theta", "tau"}: self.azimuthal = AzimuthalObjectRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalObjectTheta(kwargs["theta"]) self.temporal = TemporalObjectTau(kwargs["tau"]) else: complaint = """unrecognized combination of coordinates, allowed combinations are:\n x= y= z= tau= x= y= theta= t= x= y= theta= tau= x= y= eta= t= x= y= z= t= x= y= eta= tau= rho= phi= z= t= rho= phi= z= tau= rho= phi= theta= t= rho= phi= theta= tau= rho= phi= eta= t= rho= phi= eta= tau=""".replace(" ", " ") if type(self) is VectorObject4D: raise TypeError(complaint) else: raise TypeError(f"{complaint}\n\nor their momentum equivalents") else: raise TypeError( "must give Azimuthal, Longitudinal, and Temporal if not giving keyword arguments" ) def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] tnames = _coordinate_class_to_names[_ttype(self)] out = [f"{x}={getattr(self.azimuthal, x)}" for x in aznames] for x in lnames: out.append(f"{x}={getattr(self.longitudinal, x)}") for x in tnames: out.append(f"{x}={getattr(self.temporal, x)}") return "VectorObject4D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import VectorNumpy4D return VectorNumpy4D( self.azimuthal.elements + self.longitudinal.elements + self.temporal.elements, dtype=[ (x, numpy.float64) for x in _coordinate_class_to_names[_aztype(self)] + _coordinate_class_to_names[_ltype(self)] + _coordinate_class_to_names[_ttype(self)] ], ) def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as a scalar or a vector. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=self.longitudinal, temporal=self.temporal, ) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) return cls.ProjectionClass2D(azimuthal=azcoords) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=self.temporal ) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and returns[2] is None ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) return cls.ProjectionClass3D(azimuthal=azcoords, longitudinal=lcoords) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): azcoords = _coord_object_type[returns[0]](result[0], result[1]) lcoords = _coord_object_type[returns[1]](result[2]) tcoords = _coord_object_type[returns[2]](result[3]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=tcoords ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func @property def x(self) -> float: return super().x @x.setter def x(self, x: float) -> None: self.azimuthal = AzimuthalObjectXY(x, self.y) @property def y(self) -> float: return super().y @y.setter def y(self, y: float) -> None: self.azimuthal = AzimuthalObjectXY(self.x, y) @property def rho(self) -> float: return super().rho @rho.setter def rho(self, rho: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(rho, self.phi) @property def phi(self) -> float: return super().phi @phi.setter def phi(self, phi: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(self.rho, phi) @property def z(self) -> float: return super().z @z.setter def z(self, z: float) -> None: self.longitudinal = LongitudinalObjectZ(z) @property def theta(self) -> float: return super().theta @theta.setter def theta(self, theta: float) -> None: self.longitudinal = LongitudinalObjectTheta(theta) @property def eta(self) -> float: return super().eta @eta.setter def eta(self, eta: float) -> None: self.longitudinal = LongitudinalObjectEta(eta) @property def t(self) -> float: return super().t @t.setter def t(self, t: float) -> None: self.temporal = TemporalObjectT(t) @property def tau(self) -> float: return super().tau @tau.setter def tau(self, tau: float) -> None: self.temporal = TemporalObjectTau(tau) class MomentumObject4D(LorentzMomentum, VectorObject4D): """ Four dimensional momentum vector class for the object backend. Examples: >>> import vector >>> vec = vector.MomentumObject4D(px=1, py=2, pz=3, t=4) >>> vec.px, vec.py, vec.pz, vec.t (1, 2, 3, 4) >>> vec = vector.MomentumObject4D(pt=1, phi=2, pz=3, M=4) >>> vec.pt, vec.phi, vec.pz, vec.M (1, 2, 3, 4) >>> vec = vector.MomentumObject4D( ... azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), ... longitudinal=vector.backends.object.LongitudinalObjectTheta(3), ... temporal=vector.backends.object.TemporalObjectTau(4) ... ) >>> vec.x, vec.y, vec.theta, vec.tau (1, 2, 3, 4) The :func:`vector.obj` function can also be used to construct 4D momentum object type vectors. For four dimensional vector objects, see :class:`vector.backends.object.VectorObject4D`. """ def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] tnames = _coordinate_class_to_names[_ttype(self)] out = [] for x in aznames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.azimuthal, x)}") for x in lnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.longitudinal, x)}") for x in tnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.temporal, x)}") return "MomentumObject4D(" + ", ".join(out) + ")" def __array__(self) -> FloatArray: from vector.backends.numpy import MomentumNumpy4D return MomentumNumpy4D( self.azimuthal.elements + self.longitudinal.elements + self.temporal.elements, dtype=[ (x, numpy.float64) for x in _coordinate_class_to_names[_aztype(self)] + _coordinate_class_to_names[_ltype(self)] + _coordinate_class_to_names[_ttype(self)] ], ) @property def px(self) -> float: return super().px @px.setter def px(self, px: float) -> None: self.azimuthal = AzimuthalObjectXY(px, self.py) @property def py(self) -> float: return super().py @py.setter def py(self, py: float) -> None: self.azimuthal = AzimuthalObjectXY(self.px, py) @property def pt(self) -> float: return super().pt @pt.setter def pt(self, pt: float) -> None: self.azimuthal = AzimuthalObjectRhoPhi(pt, self.phi) @property def pz(self) -> float: return super().pz @pz.setter def pz(self, pz: float) -> None: self.longitudinal = LongitudinalObjectZ(pz) @property def E(self) -> float: return super().E @E.setter def E(self, E: float) -> None: self.temporal = TemporalObjectT(E) @property def e(self) -> float: return super().e @e.setter def e(self, e: float) -> None: self.temporal = TemporalObjectT(e) @property def energy(self) -> float: return super().energy @energy.setter def energy(self, energy: float) -> None: self.temporal = TemporalObjectT(energy) @property def M(self) -> float: return super().M @M.setter def M(self, M: float) -> None: self.temporal = TemporalObjectTau(M) @property def m(self) -> float: return super().m @m.setter def m(self, m: float) -> None: self.temporal = TemporalObjectTau(m) @property def mass(self) -> float: return super().mass @mass.setter def mass(self, mass: float) -> None: self.temporal = TemporalObjectTau(mass) def _is_type_safe(coordinates: dict[str, typing.Any]) -> None: coords = coordinates.copy() if "cls" in coords: del coords["cls"] for _, value in coords.items(): if not issubclass(type(value), numbers.Real) or isinstance(value, bool): raise TypeError("a coordinate must be of the type int or float") def _gather_coordinates( planar_class: type[VectorObject2D], spatial_class: type[VectorObject3D], lorentz_class: type[VectorObject4D], coordinates: dict[str, typing.Any], ) -> typing.Any: """ Helper function for :func:`vector.backends.object.obj`. Constructs and returns a 2D, 3D, or 4D ``VectorObject`` or ``MomentumObject`` with the provided coordinates (dictionary), planar (``VectorObject2D`` or ``MomentumObject2D``), spatial (``VectorObject3D`` or ``MomentumObject3D``), and lorentz (``VectorObject4D`` or ``MomentumObject4D``) classes. """ azimuthal: None | (AzimuthalObjectXY | AzimuthalObjectRhoPhi) = None if "x" in coordinates and "y" in coordinates: if "rho" in coordinates or "phi" in coordinates: raise TypeError("specify x= and y= or rho= and phi=, but not both") azimuthal = AzimuthalObjectXY(coordinates.pop("x"), coordinates.pop("y")) elif "rho" in coordinates and "phi" in coordinates: if "x" in coordinates or "y" in coordinates: raise TypeError("specify x= and y= or rho= and phi=, but not both") azimuthal = AzimuthalObjectRhoPhi( coordinates.pop("rho"), coordinates.pop("phi") ) longitudinal: None | ( LongitudinalObjectZ | LongitudinalObjectTheta | LongitudinalObjectEta ) = None if "z" in coordinates: if "theta" in coordinates or "eta" in coordinates: raise TypeError("specify z= or theta= or eta=, but not more than one") longitudinal = LongitudinalObjectZ(coordinates.pop("z")) elif "theta" in coordinates: if "eta" in coordinates: raise TypeError("specify z= or theta= or eta=, but not more than one") longitudinal = LongitudinalObjectTheta(coordinates.pop("theta")) elif "eta" in coordinates: longitudinal = LongitudinalObjectEta(coordinates.pop("eta")) temporal: TemporalObjectT | TemporalObjectTau | None = None if "t" in coordinates: if "tau" in coordinates: raise TypeError("specify t= or tau=, but not more than one") temporal = TemporalObjectT(coordinates.pop("t")) elif "tau" in coordinates: temporal = TemporalObjectTau(coordinates.pop("tau")) if not coordinates: if azimuthal is not None and longitudinal is None and temporal is None: return planar_class(azimuthal=azimuthal) if azimuthal is not None and longitudinal is not None and temporal is None: return spatial_class(azimuthal=azimuthal, longitudinal=longitudinal) if azimuthal is not None and longitudinal is not None and temporal is not None: return lorentz_class( azimuthal=azimuthal, longitudinal=longitudinal, temporal=temporal ) raise TypeError( "unrecognized combination of coordinates, allowed combinations are:\n\n" " (2D) x= y=\n" " (2D) rho= phi=\n" " (3D) x= y= z=\n" " (3D) x= y= theta=\n" " (3D) x= y= eta=\n" " (3D) rho= phi= z=\n" " (3D) rho= phi= theta=\n" " (3D) rho= phi= eta=\n" " (4D) x= y= z= t=\n" " (4D) x= y= z= tau=\n" " (4D) x= y= theta= t=\n" " (4D) x= y= theta= tau=\n" " (4D) x= y= eta= t=\n" " (4D) x= y= eta= tau=\n" " (4D) rho= phi= z= t=\n" " (4D) rho= phi= z= tau=\n" " (4D) rho= phi= theta= t=\n" " (4D) rho= phi= theta= tau=\n" " (4D) rho= phi= eta= t=\n" " (4D) rho= phi= eta= tau=" ) @typing.overload def obj(*, x: float, y: float) -> VectorObject2D: ... @typing.overload def obj(*, x: float, py: float) -> MomentumObject2D: ... @typing.overload def obj(*, px: float, y: float) -> MomentumObject2D: ... @typing.overload def obj(*, px: float, py: float) -> MomentumObject2D: ... @typing.overload def obj(*, rho: float, phi: float) -> VectorObject2D: ... @typing.overload def obj(*, pt: float, phi: float) -> MomentumObject2D: ... @typing.overload def obj(*, x: float, y: float, z: float) -> VectorObject3D: ... @typing.overload def obj(*, x: float, y: float, pz: float) -> MomentumObject3D: ... @typing.overload def obj(*, x: float, py: float, z: float) -> MomentumObject3D: ... @typing.overload def obj(*, x: float, py: float, pz: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, y: float, z: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, y: float, pz: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, py: float, z: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, py: float, pz: float) -> MomentumObject3D: ... @typing.overload def obj(*, rho: float, phi: float, z: float) -> VectorObject3D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float) -> MomentumObject3D: ... @typing.overload def obj(*, pt: float, phi: float, z: float) -> MomentumObject3D: ... @typing.overload def obj(*, pt: float, phi: float, pz: float) -> MomentumObject3D: ... @typing.overload def obj(*, x: float, y: float, theta: float) -> VectorObject3D: ... @typing.overload def obj(*, x: float, py: float, theta: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, y: float, theta: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, py: float, theta: float) -> MomentumObject3D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float) -> VectorObject3D: ... @typing.overload def obj(*, pt: float, phi: float, theta: float) -> MomentumObject3D: ... @typing.overload def obj(*, x: float, y: float, eta: float) -> VectorObject3D: ... @typing.overload def obj(*, x: float, py: float, eta: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, y: float, eta: float) -> MomentumObject3D: ... @typing.overload def obj(*, px: float, py: float, eta: float) -> MomentumObject3D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float) -> VectorObject3D: ... @typing.overload def obj(*, pt: float, phi: float, eta: float) -> MomentumObject3D: ... @typing.overload def obj(*, x: float, y: float, z: float, t: float) -> VectorObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, t: float) -> VectorObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, z: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, pz: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, t: float) -> VectorObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, t: float) -> VectorObject4D: ... @typing.overload def obj(*, pt: float, phi: float, theta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, t: float) -> VectorObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, t: float) -> VectorObject4D: ... @typing.overload def obj(*, pt: float, phi: float, eta: float, t: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, tau: float) -> VectorObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, tau: float) -> VectorObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, ptau: float, phi: float, z: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, ptau: float, phi: float, pz: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, tau: float) -> VectorObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, tau: float) -> VectorObject4D: ... @typing.overload def obj(*, ptau: float, phi: float, theta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, tau: float) -> VectorObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, tau: float) -> VectorObject4D: ... @typing.overload def obj(*, ptau: float, phi: float, eta: float, tau: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, pE: float, phi: float, z: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, pE: float, phi: float, pz: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, pE: float, phi: float, theta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, pE: float, phi: float, eta: float, E: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, pe: float, phi: float, z: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, pe: float, phi: float, pz: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, pe: float, phi: float, theta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, pe: float, phi: float, eta: float, e: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, z: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, pz: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, theta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, eta: float, energy: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, pM: float, phi: float, z: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, pM: float, phi: float, pz: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, pM: float, phi: float, theta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, pM: float, phi: float, eta: float, M: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, pm: float, phi: float, z: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, pm: float, phi: float, pz: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, pm: float, phi: float, theta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, pm: float, phi: float, eta: float, m: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, z: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, pz: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, z: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, pz: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, z: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, pz: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, z: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, pz: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, z: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, pz: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, z: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, pz: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, theta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, theta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, theta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, theta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, theta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, theta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, y: float, eta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, x: float, py: float, eta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, y: float, eta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, px: float, py: float, eta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, rho: float, phi: float, eta: float, mass: float) -> MomentumObject4D: ... @typing.overload def obj(*, pt: float, phi: float, eta: float, mass: float) -> MomentumObject4D: ... def obj(**coordinates: float) -> VectorObject: """ Constructs a single ``Object`` type vector, whose type is determined by the keyword-only arguments to this function. Allowed combinations are: - (2D) ``x``, ``y`` - (2D) ``rho``, ``phi`` - (3D) ``x``, ``y``, ``z`` - (3D) ``x``, ``y``, ``theta`` - (3D) ``x``, ``y``, ``eta`` - (3D) ``rho``, ``phi``, ``z`` - (3D) ``rho``, ``phi``, ``theta`` - (3D) ``rho``, ``phi``, ``eta`` - (4D) ``x``, ``y``, ``z``, ``t`` - (4D) ``x``, ``y``, ``z``, ``tau``` - (4D) ``x``, ``y``, ``theta``, ``t``` - (4D) ``x``, ``y``, ``theta``, ``tau``` - (4D) ``x``, ``y``, ``eta``, ``t``` - (4D) ``x``, ``y``, ``eta``, ``tau``` - (4D) ``rho``, ``phi``, ``z``, ``t``` - (4D) ``rho``, ``phi``, ``z``, ``tau``` - (4D) ``rho``, ``phi``, ``theta``, ``t``` - (4D) ``rho``, ``phi``, ``theta``, ``tau``` - (4D) ``rho``, ``phi``, ``eta``, ``t``` - (4D) ``rho``, ``phi``, ``eta``, ``tau``` in which - ``px`` may be substituted for ``x`` - ``py`` may be substituted for ``y`` - ``pt`` may be substituted for ``rho`` - ``pz`` may be substituted for ``z`` - ``E`` may be substituted for ``t`` - ``e`` may be substituted for ``t`` - ``energy`` may be substituted for ``t`` - ``M`` may be substituted for ``tau`` - ``m`` may be substituted for ``tau`` - ``mass`` may be substituted for ``tau`` to make the vector a momentum vector. Alternatively, the :class:`vector.VectorObject2D`, :class:`vector.VectorObject3D`, and :class:`vector.VectorObject4D` classes (with momentum subclasses) have explicit constructors: - :meth:`vector.VectorObject2D.from_xy` - :meth:`vector.VectorObject2D.from_rhophi` - :meth:`vector.VectorObject3D.from_xyz` - :meth:`vector.VectorObject3D.from_xytheta` - :meth:`vector.VectorObject3D.from_xyeta` - :meth:`vector.VectorObject3D.from_rhophiz` - :meth:`vector.VectorObject3D.from_rhophitheta` - :meth:`vector.VectorObject3D.from_rhophieta` - :meth:`vector.VectorObject4D.from_xyzt` - :meth:`vector.VectorObject4D.from_xyztau` - :meth:`vector.VectorObject4D.from_xythetat` - :meth:`vector.VectorObject4D.from_xythetatau` - :meth:`vector.VectorObject4D.from_xyetat` - :meth:`vector.VectorObject4D.from_xyetatau` - :meth:`vector.VectorObject4D.from_rhophizt` - :meth:`vector.VectorObject4D.from_rhophiztau` - :meth:`vector.VectorObject4D.from_rhophithetat` - :meth:`vector.VectorObject4D.from_rhophithetatau` - :meth:`vector.VectorObject4D.from_rhophietat` - :meth:`vector.VectorObject4D.from_rhophietatau` """ is_momentum = False generic_coordinates = {} _is_type_safe(coordinates) if "px" in coordinates: is_momentum = True generic_coordinates["x"] = coordinates.pop("px") if "py" in coordinates: is_momentum = True generic_coordinates["y"] = coordinates.pop("py") if "pt" in coordinates: is_momentum = True generic_coordinates["rho"] = coordinates.pop("pt") if "pz" in coordinates: is_momentum = True generic_coordinates["z"] = coordinates.pop("pz") if "E" in coordinates: is_momentum = True generic_coordinates["t"] = coordinates.pop("E") if "e" in coordinates: is_momentum = True generic_coordinates["t"] = coordinates.pop("e") if "energy" in coordinates and "t" not in generic_coordinates: is_momentum = True generic_coordinates["t"] = coordinates.pop("energy") if "M" in coordinates: is_momentum = True generic_coordinates["tau"] = coordinates.pop("M") if "m" in coordinates: is_momentum = True generic_coordinates["tau"] = coordinates.pop("m") if "mass" in coordinates and "tau" not in generic_coordinates: is_momentum = True generic_coordinates["tau"] = coordinates.pop("mass") for x in list(coordinates): if x not in generic_coordinates: generic_coordinates[x] = coordinates.pop(x) if len(coordinates) != 0: raise TypeError( "duplicate coordinates (through momentum-aliases): " + ", ".join(repr(x) for x in coordinates) ) if is_momentum: return _gather_coordinates( MomentumObject2D, MomentumObject3D, MomentumObject4D, generic_coordinates ) else: return _gather_coordinates( VectorObject2D, VectorObject3D, VectorObject4D, generic_coordinates ) VectorObject2D.ProjectionClass2D = VectorObject2D VectorObject2D.ProjectionClass3D = VectorObject3D VectorObject2D.ProjectionClass4D = VectorObject4D VectorObject2D.GenericClass = VectorObject2D VectorObject2D.MomentumClass = MomentumObject2D MomentumObject2D.ProjectionClass2D = MomentumObject2D MomentumObject2D.ProjectionClass3D = MomentumObject3D MomentumObject2D.ProjectionClass4D = MomentumObject4D MomentumObject2D.GenericClass = VectorObject2D MomentumObject2D.MomentumClass = MomentumObject2D VectorObject3D.ProjectionClass2D = VectorObject2D VectorObject3D.ProjectionClass3D = VectorObject3D VectorObject3D.ProjectionClass4D = VectorObject4D VectorObject3D.GenericClass = VectorObject3D VectorObject3D.MomentumClass = MomentumObject3D MomentumObject3D.ProjectionClass2D = MomentumObject2D MomentumObject3D.ProjectionClass3D = MomentumObject3D MomentumObject3D.ProjectionClass4D = MomentumObject4D MomentumObject3D.GenericClass = VectorObject3D MomentumObject3D.MomentumClass = MomentumObject3D VectorObject4D.ProjectionClass2D = VectorObject2D VectorObject4D.ProjectionClass3D = VectorObject3D VectorObject4D.ProjectionClass4D = VectorObject4D VectorObject4D.GenericClass = VectorObject4D VectorObject4D.MomentumClass = MomentumObject4D MomentumObject4D.ProjectionClass2D = MomentumObject2D MomentumObject4D.ProjectionClass3D = MomentumObject3D MomentumObject4D.ProjectionClass4D = MomentumObject4D MomentumObject4D.GenericClass = VectorObject4D MomentumObject4D.MomentumClass = MomentumObject4D vector-1.6.3/src/vector/backends/sympy.py000066400000000000000000001522641503546127100204370ustar00rootroot00000000000000from __future__ import annotations import typing import numpy import sympy # type: ignore[import-untyped] import vector from vector._methods import ( Azimuthal, AzimuthalRhoPhi, AzimuthalXY, Longitudinal, LongitudinalEta, LongitudinalTheta, LongitudinalZ, Lorentz, LorentzMomentum, Planar, PlanarMomentum, SameVectorType, Spatial, SpatialMomentum, Temporal, TemporalT, TemporalTau, Vector, Vector2D, Vector3D, Vector4D, VectorProtocol, _aztype, _coordinate_class_to_names, _handler_of, _ltype, _repr_generic_to_momentum, _repr_momentum_to_generic, _ttype, ) class _lib: """a wrapper that maps numpy functions to sympy functions (or custom implementations)""" # functions modified specifically for sympy # TODO: https://github.com/scikit-hep/vector/issues/615 # should NOT be used as a replacement for np.where as it works specifically for the # case of vector/_compute/spatial/eta.py. `where` returns the second argument and # ignores the third (first argument is a boolean condition) because the purpose of # this function is to handle exceptional values — we know that the "normal" values # are in the second argument and the "exceptional" ones are in the third argument. # remove once https://github.com/cupy/cupy/issues/9143 is fixed. def where(self, val1: sympy.Expr, val2: sympy.Expr, val3: sympy.Expr) -> sympy.Expr: return val2 def nan_to_num(self, val: sympy.Expr, **kwargs: typing.Any) -> sympy.Expr: return val def maximum(self, val1: sympy.Expr | int, val2: sympy.Expr | int) -> sympy.Expr: return val1 if isinstance(val1, sympy.Expr) else val2 def minimum(self, val1: sympy.Expr | int, val2: sympy.Expr | int) -> sympy.Expr: return val1 if isinstance(val1, sympy.Expr) else val2 def arcsin(self, val: sympy.Expr) -> sympy.Expr: return sympy.asin(val) def arccos(self, val: sympy.Expr) -> sympy.Expr: return sympy.acos(val) def arctan(self, val: sympy.Expr) -> sympy.Expr: return sympy.atan(val) def arctan2(self, val1: sympy.Expr, val2: sympy.Expr) -> sympy.Expr: return sympy.atan2(val1, val2) def arcsinh(self, val: sympy.Expr) -> sympy.Expr: return sympy.asinh(val) def arccosh(self, val: sympy.Expr) -> sympy.Expr: return sympy.acosh(val) def arctanh(self, val: sympy.Expr) -> sympy.Expr: return sympy.atanh(val) def absolute(self, val: sympy.Expr) -> sympy.Expr: return sympy.Abs(val) def isclose( self, val1: sympy.Expr, val2: sympy.Expr, *args: float, ) -> sympy.Equality: return sympy.Eq(val1, val2) def copysign(self, val1: sympy.Expr, val2: sympy.Expr) -> sympy.Expr: return val1 @property def inf(self) -> sympy.Expr: return sympy.oo # same named functions def sign(self, val: int | float) -> sympy.Expr: return numpy.sign(val) def sqrt(self, val: sympy.Expr) -> sympy.Expr: return sympy.sqrt(val) def exp(self, val: sympy.Expr) -> sympy.Expr: return sympy.exp(val) def log(self, val: sympy.Expr) -> sympy.Expr: return sympy.log(val) def sin(self, val: sympy.Expr) -> sympy.Expr: return sympy.sin(val) def cos(self, val: sympy.Expr) -> sympy.Expr: return sympy.cos(val) def tan(self, val: sympy.Expr) -> sympy.Expr: return sympy.tan(val) def sinh(self, val: sympy.Expr) -> sympy.Expr: return sympy.sinh(val) def cosh(self, val: sympy.Expr) -> sympy.Expr: return sympy.cosh(val) def tanh(self, val: sympy.Expr) -> sympy.Expr: return sympy.tanh(val) @property def pi(self) -> sympy.Expr: return sympy.pi class CoordinatesSympy: """Coordinates class for the SymPy backend.""" lib = _lib() class AzimuthalSympy(CoordinatesSympy, Azimuthal): """Azimuthal class for the SymPy backend.""" ObjectClass: type[vector.backends.object.AzimuthalObject] class LongitudinalSympy(CoordinatesSympy, Longitudinal): """Longitudinal class for the SymPy backend.""" ObjectClass: type[vector.backends.object.LongitudinalObject] class TemporalSympy(CoordinatesSympy, Temporal): """Temporal class for the SymPy backend.""" ObjectClass: type[vector.backends.object.TemporalObject] class AzimuthalSympyXY(AzimuthalSympy, AzimuthalXY): """ Class for the ``x`` and ``y`` (azimuthal) coordinates of SymPy backend. Examples: >>> import vector; import sympy >>> vector.backends.sympy.AzimuthalSympyXY(sympy.Symbol("x"), sympy.Symbol("y")) AzimuthalSympyXY(x=x, y=y) """ def __init__(self, x: sympy.Symbol, y: sympy.Symbol): self.x = x self.y = y def __repr__(self) -> str: return f"AzimuthalSympyXY(x={self.x!r}, y={self.y!r})" @property def elements(self) -> tuple[sympy.Symbol, sympy.Symbol]: """ Azimuthal coordinates (``x`` and ``y``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector; import sympy >>> vec = vector.backends.sympy.AzimuthalSympyXY(sympy.Symbol("x"), sympy.Symbol("y")) >>> vec.elements (x, y) """ return (self.x, self.y) class AzimuthalSympyRhoPhi(AzimuthalSympy, AzimuthalRhoPhi): """ Class for the ``rho`` and ``phi`` (azimuthal) coordinates of SymPy backend. Examples: >>> import vector; import sympy >>> vector.backends.sympy.AzimuthalSympyRhoPhi(sympy.Symbol("rho"), sympy.Symbol("phi")) AzimuthalSympyRhoPhi(rho=rho, phi=phi) """ def __init__(self, rho: sympy.Symbol, phi: sympy.Symbol): self.rho = rho self.phi = phi def __repr__(self) -> str: return f"AzimuthalSympyRhoPhi(rho={self.rho!r}, phi={self.phi!r})" @property def elements(self) -> tuple[sympy.Symbol, sympy.Symbol]: """ Azimuthal coordinates (``rho`` and ``phi``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector >>> vec = vector.backends.sympy.AzimuthalSympyRhoPhi(sympy.Symbol("rho"), sympy.Symbol("phi")) >>> vec.elements (rho, phi) """ return (self.rho, self.phi) class LongitudinalSympyZ(LongitudinalSympy, LongitudinalZ): """ Class for the ``z`` (longitudinal) coordinate of SymPy backend. Examples: >>> import vector; import sympy >>> vector.backends.sympy.LongitudinalSympyZ(sympy.Symbol("z")) LongitudinalSympyZ(z=z) """ def __init__(self, z: sympy.Symbol): self.z = z def __repr__(self) -> str: return f"LongitudinalSympyZ(z={self.z!r})" @property def elements(self) -> tuple[sympy.Symbol]: """ Longitudinal coordinates (``z``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector; import sympy >>> vec = vector.backends.sympy.LongitudinalSympyZ(sympy.Symbol("z")) >>> vec.elements (z,) """ return (self.z,) class LongitudinalSympyTheta(LongitudinalSympy, LongitudinalTheta): """ Class for the ``theta`` (longitudinal) coordinate of SymPy backend. Examples: >>> import vector; import sympy >>> vector.backends.sympy.LongitudinalSympyTheta(sympy.Symbol("theta")) LongitudinalSympyTheta(theta=theta) """ def __init__(self, theta: sympy.Symbol): self.theta = theta def __repr__(self) -> str: return f"LongitudinalSympyTheta(theta={self.theta!r})" @property def elements(self) -> tuple[sympy.Symbol]: """ Longitudinal coordinates (``theta``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector; import sympy >>> vec = vector.backends.sympy.LongitudinalSympyTheta(sympy.Symbol("theta")) >>> vec.elements (theta,) """ return (self.theta,) class LongitudinalSympyEta(LongitudinalSympy, LongitudinalEta): """ Class for the ``eta`` (longitudinal) coordinate of SymPy backend. Examples: >>> import vector; import sympy >>> vector.backends.sympy.LongitudinalSympyEta(sympy.Symbol("eta")) LongitudinalSympyEta(eta=eta) """ def __init__(self, eta: sympy.Symbol): self.eta = eta def __repr__(self) -> str: return f"LongitudinalSympyEta(eta={self.eta!r})" @property def elements(self) -> tuple[sympy.Symbol]: """ Longitudinal coordinates (``eta``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector >>> vec = vector.backends.sympy.LongitudinalSympyEta(sympy.Symbol("eta")) >>> vec.elements (eta,) """ return (self.eta,) class TemporalSympyT(TemporalSympy, TemporalT): """ Class for the ``t`` (temporal) coordinate of SymPy backend. Examples: >>> import vector >>> vector.backends.sympy.TemporalSympyT(sympy.Symbol("t")) TemporalSympyT(t=t) """ def __init__(self, t: sympy.Symbol): self.t = t def __repr__(self) -> str: return f"TemporalSympyT(t={self.t!r})" @property def elements(self) -> tuple[sympy.Symbol]: """ Temporal coordinates (``t``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector >>> vec = vector.backends.sympy.TemporalSympyT(sympy.Symbol("t")) >>> vec.elements (t,) """ return (self.t,) class TemporalSympyTau(TemporalSympy, TemporalTau): """ Class for the ``tau`` (temporal) coordinate of SymPy backend. Examples: >>> import vector >>> vector.backends.sympy.TemporalSympyTau(sympy.Symbol("tau")) TemporalSympyTau(tau=tau) """ def __init__(self, tau: sympy.Symbol): self.tau = tau def __repr__(self) -> str: return f"TemporalSympyTau(tau={self.tau!r})" @property def elements(self) -> tuple[sympy.Symbol]: """ Temporal coordinates (``tau``) as a tuple. Each coordinate is a SymPy expression and not a vector. Examples: >>> import vector >>> vec = vector.backends.sympy.TemporalSympyTau(sympy.Symbol("tau")) >>> vec.elements (tau,) """ return (self.tau,) _coord_sympy_type = { AzimuthalXY: AzimuthalSympyXY, AzimuthalRhoPhi: AzimuthalSympyRhoPhi, LongitudinalZ: LongitudinalSympyZ, LongitudinalTheta: LongitudinalSympyTheta, LongitudinalEta: LongitudinalSympyEta, TemporalT: TemporalSympyT, TemporalTau: TemporalSympyTau, } def _is_type_safe(coordinates: dict[str, typing.Any]) -> None: if not all(isinstance(coord, sympy.Expr) for coord in coordinates.values()): raise TypeError("coordinates must be a sympy expression") def _replace_data(obj: typing.Any, result: typing.Any) -> typing.Any: if not isinstance(result, VectorSympy): raise TypeError(f"can only assign a single vector to {type(obj).__name__}") if isinstance(result, (VectorSympy2D, VectorSympy3D, VectorSympy4D)): if isinstance(obj.azimuthal, AzimuthalSympyXY): obj.azimuthal = AzimuthalSympyXY(result.x, result.y) elif isinstance(obj.azimuthal, AzimuthalSympyRhoPhi): obj.azimuthal = AzimuthalSympyRhoPhi(result.rho, result.phi) else: raise AssertionError(type(obj)) if isinstance(result, (VectorSympy3D, VectorSympy4D)): if isinstance(obj.longitudinal, LongitudinalSympyZ): obj.longitudinal = LongitudinalSympyZ(result.z) elif isinstance(obj.longitudinal, LongitudinalSympyTheta): obj.longitudinal = LongitudinalSympyTheta(result.theta) elif isinstance(obj.longitudinal, LongitudinalSympyEta): obj.longitudinal = LongitudinalSympyEta(result.eta) else: raise AssertionError(type(obj)) if isinstance(result, VectorSympy4D): if isinstance(obj.temporal, TemporalSympyT): obj.temporal = TemporalSympyT(result.t) elif isinstance(obj.temporal, TemporalSympyTau): obj.temporal = TemporalSympyTau(result.tau) else: raise AssertionError(type(obj)) return obj class VectorSympy(Vector): # noqa: PLW1641 """Mixin class for Sympy vectors.""" lib = _lib() def __eq__(self, other: typing.Any) -> typing.Any: return numpy.equal(self, other) def __ne__(self, other: typing.Any) -> typing.Any: return numpy.not_equal(self, other) def __abs__(self) -> float: return numpy.absolute(self) def __add__(self, other: VectorProtocol) -> VectorProtocol: return numpy.add(self, other) def __radd__(self, other: VectorProtocol) -> VectorProtocol: return numpy.add(other, self) def __iadd__(self: SameVectorType, other: VectorProtocol) -> SameVectorType: return _replace_data(self, numpy.add(self, other)) def __sub__(self, other: VectorProtocol) -> VectorProtocol: return numpy.subtract(self, other) def __rsub__(self, other: VectorProtocol) -> VectorProtocol: return numpy.subtract(other, self) def __isub__(self: SameVectorType, other: VectorProtocol) -> SameVectorType: return _replace_data(self, numpy.subtract(self, other)) def __mul__(self, other: float) -> VectorProtocol: return numpy.multiply(self, other) def __rmul__(self, other: float) -> VectorProtocol: return numpy.multiply(other, self) def __imul__(self: SameVectorType, other: float) -> SameVectorType: return _replace_data(self, numpy.multiply(self, other)) def __neg__(self: SameVectorType) -> SameVectorType: return numpy.negative(self) def __pos__(self: SameVectorType) -> SameVectorType: return numpy.positive(self) def __truediv__(self, other: float) -> VectorProtocol: return numpy.true_divide(self, other) def __rtruediv__(self, other: float) -> VectorProtocol: return numpy.true_divide(other, self) def __itruediv__(self: SameVectorType, other: float) -> VectorProtocol: return _replace_data(self, numpy.true_divide(self, other)) def __pow__(self, other: float) -> float: return numpy.power(self, other) def __matmul__(self, other: VectorProtocol) -> float: return numpy.matmul(self, other) def __array_ufunc__( self, ufunc: typing.Any, method: typing.Any, *inputs: typing.Any, **kwargs: typing.Any, ) -> typing.Any: """ Implements NumPy's ``ufunc``s for ``VectorSympy`` and its subclasses. The current implementation includes ``numpy.absolute``, ``numpy.add``, ``numpy.subtract``, ``numpy.multiply``, ``numpy.positive``, ``numpy.negative``, ``numpy.true_divide``, ``numpy.power``, ``numpy.square``, ``numpy.sqrt``, ``numpy.cbrt``, ``numpy.matmul``, ``numpy.equal``, and ``numpy.not_equal``. """ if not isinstance(_handler_of(*inputs), VectorSympy): # Let a higher-precedence backend handle it. return NotImplemented outputs = kwargs.get("out", ()) if any(not isinstance(x, VectorSympy) for x in outputs): raise TypeError( "ufunc operating on VectorSympy can only use the 'out' keyword " "with another VectorSympy" ) if ( ufunc is numpy.absolute and len(inputs) == 1 and isinstance(inputs[0], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.absolute' is scalar, cannot fill a VectorSympy with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho elif isinstance(inputs[0], Vector3D): return inputs[0].mag elif isinstance(inputs[0], Vector4D): return inputs[0].tau elif ( ufunc is numpy.add and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[0].add(inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.subtract and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[0].subtract(inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.multiply and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = inputs[0].scale(inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.multiply and len(inputs) == 2 and not isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): result = inputs[1].scale(inputs[0]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.negative and len(inputs) == 1 and isinstance(inputs[0], Vector) ): result = inputs[0].scale(-1) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.positive and len(inputs) == 1 and isinstance(inputs[0], Vector) ): return inputs[0] elif ( ufunc is numpy.true_divide and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = inputs[0].scale(1 / inputs[1]) for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.power and len(inputs) == 2 and isinstance(inputs[0], Vector) and not isinstance(inputs[1], Vector) ): result = numpy.absolute(inputs[0]) ** inputs[1] for output in outputs: _replace_data(output, result) return result elif ( ufunc is numpy.square and len(inputs) == 1 and isinstance(inputs[0], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.square' is scalar, cannot fill a VectorSympy with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 elif ufunc is numpy.sqrt and len(inputs) == 1 and isinstance(inputs[0], Vector): if len(outputs) != 0: raise TypeError( "output of 'numpy.sqrt' is scalar, cannot fill a VectorSympy with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 ** 0.25 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 ** 0.25 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 ** 0.25 elif ufunc is numpy.cbrt and len(inputs) == 1 and isinstance(inputs[0], Vector): if len(outputs) != 0: raise TypeError( "output of 'numpy.cbrt' is scalar, cannot fill a VectorSympy with 'out'" ) if isinstance(inputs[0], Vector2D): return inputs[0].rho2 ** 0.16666666666666666 elif isinstance(inputs[0], Vector3D): return inputs[0].mag2 ** 0.16666666666666666 elif isinstance(inputs[0], Vector4D): return inputs[0].tau2 ** 0.16666666666666666 elif ( ufunc is numpy.matmul and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.matmul' is scalar, cannot fill a VectorSympy with 'out'" ) return inputs[0].dot(inputs[1]) elif ( ufunc is numpy.equal and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.equal' is scalar, cannot fill a VectorSympy with 'out'" ) return inputs[0].equal(inputs[1]) elif ( ufunc is numpy.not_equal and len(inputs) == 2 and isinstance(inputs[0], Vector) and isinstance(inputs[1], Vector) ): if len(outputs) != 0: raise TypeError( "output of 'numpy.equal' is scalar, cannot fill a VectorSympy with 'out'" ) return inputs[0].not_equal(inputs[1]) else: return NotImplemented class VectorSympy2D(VectorSympy, Planar, Vector2D): """ Two dimensional vector class for the SymPy backend. Examples: >>> import vector; import sympy >>> vec = vector.VectorSympy2D(x=sympy.Symbol("x"), y=sympy.Symbol("y")) >>> vec.x, vec.y (x, y) >>> vec = vector.VectorSympy2D(rho=sympy.Symbol("rho"), phi=sympy.Symbol("phi")) >>> vec.rho, vec.phi (rho, phi) For two dimensional momentum SymPy vectors, see :class:`vector.backends.sympy.MomentumSympy2D`. """ __slots__ = ("azimuthal",) azimuthal: AzimuthalSympy def __init__(self, azimuthal: AzimuthalSympy | None = None, **kwargs: sympy.Symbol): for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v if not kwargs and azimuthal is not None: self.azimuthal = azimuthal elif kwargs and azimuthal is None: _is_type_safe(kwargs) if set(kwargs) == {"x", "y"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) elif set(kwargs) == {"rho", "phi"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) else: complaint = """unrecognized combination of coordinates, allowed combinations are:\n x= y= rho= phi=""".replace(" ", " ") if type(self) is VectorSympy2D: raise TypeError(complaint) else: raise TypeError(f"{complaint}\n\nor their momentum equivalents") else: raise TypeError("must give Azimuthal if not giving keyword arguments") def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] out = [f"{x}={getattr(self.azimuthal, x)}" for x in aznames] return "VectorSympy2D(" + ", ".join(out) + ")" @property def x(self) -> sympy.Symbol: return super().x @x.setter def x(self, x: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(x, self.y) @property def y(self) -> sympy.Symbol: return super().y @y.setter def y(self, y: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(self.x, y) @property def rho(self) -> sympy.Symbol: return super().rho @rho.setter def rho(self, rho: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(rho, self.phi) @property def phi(self) -> sympy.Symbol: return super().phi @phi.setter def phi(self, phi: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(self.rho, phi) def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as a scalar or a vector. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( (len(returns) == 1 or (len(returns) == 2 and returns[1] is None)) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) return cls.ProjectionClass2D(azimuthal=azcoords) elif len(returns) == 2 or ( (len(returns) == 3 and returns[2] is None) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) return cls.ProjectionClass3D(azimuthal=azcoords, longitudinal=lcoords) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) tcoords = _coord_sympy_type[returns[2]](result[3]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=tcoords ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func class MomentumSympy2D(PlanarMomentum, VectorSympy2D): """ Two dimensional momentum vector class for the SymPy backend. Examples: >>> import vector; import sympy >>> vec = vector.MomentumSympy2D(px=sympy.Symbol("px"), py=sympy.Symbol("py")) >>> vec.px, vec.py (px, py) >>> vec = vector.MomentumSympy2D(pt=sympy.Symbol("pt"), phi=sympy.Symbol("phi")) >>> vec.pt, vec.phi (pt, phi) For two dimensional SymPy vectors, see :class:`vector.backends.object.VectorSympy2D`. """ def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] out = [] for x in aznames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.azimuthal, x)}") return "MomentumSympy2D(" + ", ".join(out) + ")" @property def px(self) -> sympy.Symbol: return super().px @px.setter def px(self, px: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(px, self.py) @property def py(self) -> sympy.Symbol: return super().py @py.setter def py(self, py: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(self.px, py) @property def pt(self) -> sympy.Symbol: return super().pt @pt.setter def pt(self, pt: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(pt, self.phi) class VectorSympy3D(VectorSympy, Spatial, Vector3D): """ Three dimensional vector class for the SymPy backend. Examples: >>> import vector; import sympy >>> vec = vector.VectorSympy3D(x=sympy.Symbol("x"), y=sympy.Symbol("y"), z=sympy.Symbol("z")) >>> vec.x, vec.y, vec.z (x, y, z) >>> vec = vector.VectorSympy3D(rho=sympy.Symbol("rho"), phi=sympy.Symbol("phi"), eta=sympy.Symbol("eta")) >>> vec.rho, vec.phi, vec.eta (rho, phi, eta) For three dimensional momentum SymPy vectors, see :class:`vector.backends.object.MomentumSympy3D`. """ __slots__ = ("azimuthal", "longitudinal") azimuthal: AzimuthalSympy longitudinal: LongitudinalSympy def __init__( self, azimuthal: AzimuthalSympy | None = None, longitudinal: LongitudinalSympy | None = None, **kwargs: sympy.Symbol, ): for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v if not kwargs and azimuthal is not None and longitudinal is not None: self.azimuthal = azimuthal self.longitudinal = longitudinal elif kwargs and azimuthal is None and longitudinal is None: _is_type_safe(kwargs) if set(kwargs) == {"x", "y", "z"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyZ(kwargs["z"]) elif set(kwargs) == {"x", "y", "eta"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) elif set(kwargs) == {"x", "y", "theta"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) elif set(kwargs) == {"rho", "phi", "z"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyZ(kwargs["z"]) elif set(kwargs) == {"rho", "phi", "eta"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) elif set(kwargs) == {"rho", "phi", "theta"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) else: complaint = """unrecognized combination of coordinates, allowed combinations are:\n x= y= z= x= y= theta= x= y= eta= rho= phi= z= rho= phi= theta= rho= phi= eta=""".replace(" ", " ") if type(self) is VectorSympy3D: raise TypeError(complaint) else: raise TypeError(f"{complaint}\n\nor their momentum equivalents") else: raise TypeError( "must give Azimuthal and Longitudinal if not giving keyword arguments" ) def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] out = [f"{x}={getattr(self.azimuthal, x)}" for x in aznames] for x in lnames: out.append(f"{x}={getattr(self.longitudinal, x)}") return "VectorSympy3D(" + ", ".join(out) + ")" def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as a scalar or a vector. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) return cls.ProjectionClass3D( azimuthal=azcoords, longitudinal=self.longitudinal ) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) return cls.ProjectionClass2D(azimuthal=azcoords) elif len(returns) == 2 or ( (len(returns) == 3 and returns[2] is None) and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) return cls.ProjectionClass3D(azimuthal=azcoords, longitudinal=lcoords) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) tcoords = _coord_sympy_type[returns[2]](result[3]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=tcoords ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func @property def x(self) -> sympy.Symbol: return super().x @x.setter def x(self, x: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(x, self.y) @property def y(self) -> sympy.Symbol: return super().y @y.setter def y(self, y: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(self.x, y) @property def rho(self) -> sympy.Symbol: return super().rho @rho.setter def rho(self, rho: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(rho, self.phi) @property def phi(self) -> sympy.Symbol: return super().phi @phi.setter def phi(self, phi: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(self.rho, phi) @property def z(self) -> sympy.Symbol: return super().z @z.setter def z(self, z: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyZ(z) @property def theta(self) -> sympy.Symbol: return super().theta @theta.setter def theta(self, theta: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyTheta(theta) @property def eta(self) -> sympy.Symbol: return super().eta @eta.setter def eta(self, eta: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyEta(eta) class MomentumSympy3D(SpatialMomentum, VectorSympy3D): """ Three dimensional momentum vector class for the SymPy backend. Examples: >>> import vector; import sympy >>> vec = vector.MomentumSympy3D(px=sympy.Symbol("px"), py=sympy.Symbol("py"), pz=sympy.Symbol("pz")) >>> vec.px, vec.py, vec.pz (px, py, pz) >>> vec = vector.MomentumSympy3D(pt=sympy.Symbol("pt"), phi=sympy.Symbol("phi"), pz=sympy.Symbol("pz")) >>> vec.pt, vec.phi, vec.pz (pt, phi, pz) For three dimensional SymPy vectors, see :class:`vector.backends.sympy.VectorSympy3D`. """ def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] out = [] for x in aznames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.azimuthal, x)}") for x in lnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.longitudinal, x)}") return "MomentumSympy3D(" + ", ".join(out) + ")" @property def px(self) -> sympy.Symbol: return super().px @px.setter def px(self, px: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(px, self.py) @property def py(self) -> sympy.Symbol: return super().py @py.setter def py(self, py: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(self.px, py) @property def pt(self) -> sympy.Symbol: return super().pt @pt.setter def pt(self, pt: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(pt, self.phi) @property def pz(self) -> sympy.Symbol: return super().pz @pz.setter def pz(self, pz: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyZ(pz) class VectorSympy4D(VectorSympy, Lorentz, Vector4D): """ Four dimensional vector class for the SymPy backend. Examples: >>> import vector; import sympy >>> vec = vector.VectorSympy4D(x=sympy.Symbol("x"), y=sympy.Symbol("y"), z=sympy.Symbol("z"), t=sympy.Symbol("t")) >>> vec.x, vec.y, vec.z, vec.t (x, y, z, t) >>> vec = vector.VectorSympy4D(rho=sympy.Symbol("rho"), phi=sympy.Symbol("phi"), eta=sympy.Symbol("eta"), tau=sympy.Symbol("tau")) >>> vec.rho, vec.phi, vec.eta, vec.tau (rho, phi, eta, tau) For four dimensional momentum SymPy vectors, see :class:`vector.backends.sympy.MomentumSympy4D`. """ __slots__ = ("azimuthal", "longitudinal", "temporal") azimuthal: AzimuthalSympy longitudinal: LongitudinalSympy temporal: TemporalSympy def __init__( self, azimuthal: AzimuthalSympy | None = None, longitudinal: LongitudinalSympy | None = None, temporal: TemporalSympy | None = None, **kwargs: sympy.Symbol, ): for k, v in kwargs.copy().items(): kwargs.pop(k) kwargs[_repr_momentum_to_generic.get(k, k)] = v if ( not kwargs and azimuthal is not None and longitudinal is not None and temporal is not None ): self.azimuthal = azimuthal self.longitudinal = longitudinal self.temporal = temporal elif kwargs and azimuthal is None and longitudinal is None and temporal is None: _is_type_safe(kwargs) if set(kwargs) == {"x", "y", "z", "t"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyZ(kwargs["z"]) self.temporal = TemporalSympyT(kwargs["t"]) elif set(kwargs) == {"x", "y", "eta", "t"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) self.temporal = TemporalSympyT(kwargs["t"]) elif set(kwargs) == {"x", "y", "theta", "t"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) self.temporal = TemporalSympyT(kwargs["t"]) elif set(kwargs) == {"rho", "phi", "z", "t"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyZ(kwargs["z"]) self.temporal = TemporalSympyT(kwargs["t"]) elif set(kwargs) == {"rho", "phi", "eta", "t"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) self.temporal = TemporalSympyT(kwargs["t"]) elif set(kwargs) == {"rho", "phi", "theta", "t"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) self.temporal = TemporalSympyT(kwargs["t"]) elif set(kwargs) == {"x", "y", "z", "tau"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyZ(kwargs["z"]) self.temporal = TemporalSympyTau(kwargs["tau"]) elif set(kwargs) == {"x", "y", "eta", "tau"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) self.temporal = TemporalSympyTau(kwargs["tau"]) elif set(kwargs) == {"x", "y", "theta", "tau"}: self.azimuthal = AzimuthalSympyXY(kwargs["x"], kwargs["y"]) self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) self.temporal = TemporalSympyTau(kwargs["tau"]) elif set(kwargs) == {"rho", "phi", "z", "tau"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyZ(kwargs["z"]) self.temporal = TemporalSympyTau(kwargs["tau"]) elif set(kwargs) == {"rho", "phi", "eta", "tau"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyEta(kwargs["eta"]) self.temporal = TemporalSympyTau(kwargs["tau"]) elif set(kwargs) == {"rho", "phi", "theta", "tau"}: self.azimuthal = AzimuthalSympyRhoPhi(kwargs["rho"], kwargs["phi"]) self.longitudinal = LongitudinalSympyTheta(kwargs["theta"]) self.temporal = TemporalSympyTau(kwargs["tau"]) else: complaint = """unrecognized combination of coordinates, allowed combinations are:\n x= y= z= tau= x= y= theta= t= x= y= theta= tau= x= y= eta= t= x= y= z= t= x= y= eta= tau= rho= phi= z= t= rho= phi= z= tau= rho= phi= theta= t= rho= phi= theta= tau= rho= phi= eta= t= rho= phi= eta= tau=""".replace(" ", " ") if type(self) is VectorSympy4D: raise TypeError(complaint) else: raise TypeError(f"{complaint}\n\nor their momentum equivalents") else: raise TypeError( "must give Azimuthal, Longitudinal, and Temporal if not giving keyword arguments" ) def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] tnames = _coordinate_class_to_names[_ttype(self)] out = [f"{x}={getattr(self.azimuthal, x)}" for x in aznames] for x in lnames: out.append(f"{x}={getattr(self.longitudinal, x)}") for x in tnames: out.append(f"{x}={getattr(self.temporal, x)}") return "VectorSympy4D(" + ", ".join(out) + ")" def _wrap_result( self, cls: typing.Any, result: typing.Any, returns: typing.Any, num_vecargs: typing.Any, ) -> typing.Any: """ Wraps the raw result of a compute function as a scalar or a vector. Args: result: Value or tuple of values from a compute function. returns: Signature from a ``dispatch_map``. num_vecargs (int): Number of vector arguments in the function that would be treated on an equal footing (i.e. ``add`` has two, but ``rotate_axis`` has only one: the ``axis`` is secondary). """ if returns in ([float], [bool]): return result elif ( len(returns) == 1 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=self.longitudinal, temporal=self.temporal, ) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and returns[1] is None ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) return cls.ProjectionClass2D(azimuthal=azcoords) elif ( len(returns) == 2 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=self.temporal ) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and returns[2] is None ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) return cls.ProjectionClass3D(azimuthal=azcoords, longitudinal=lcoords) elif ( len(returns) == 3 and isinstance(returns[0], type) and issubclass(returns[0], Azimuthal) and isinstance(returns[1], type) and issubclass(returns[1], Longitudinal) and isinstance(returns[2], type) and issubclass(returns[2], Temporal) ): azcoords = _coord_sympy_type[returns[0]](result[0], result[1]) lcoords = _coord_sympy_type[returns[1]](result[2]) tcoords = _coord_sympy_type[returns[2]](result[3]) return cls.ProjectionClass4D( azimuthal=azcoords, longitudinal=lcoords, temporal=tcoords ) else: raise AssertionError(repr(returns)) def _wrap_dispatched_function(self, func: typing.Callable) -> typing.Callable: # type: ignore[type-arg] return func @property def x(self) -> sympy.Symbol: return super().x @x.setter def x(self, x: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(x, self.y) @property def y(self) -> sympy.Symbol: return super().y @y.setter def y(self, y: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(self.x, y) @property def rho(self) -> sympy.Symbol: return super().rho @rho.setter def rho(self, rho: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(rho, self.phi) @property def phi(self) -> sympy.Symbol: return super().phi @phi.setter def phi(self, phi: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(self.rho, phi) @property def z(self) -> sympy.Symbol: return super().z @z.setter def z(self, z: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyZ(z) @property def theta(self) -> sympy.Symbol: return super().theta @theta.setter def theta(self, theta: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyTheta(theta) @property def eta(self) -> sympy.Symbol: return super().eta @eta.setter def eta(self, eta: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyEta(eta) @property def t(self) -> sympy.Symbol: return super().t @t.setter def t(self, t: sympy.Symbol) -> None: self.temporal = TemporalSympyT(t) @property def tau(self) -> sympy.Symbol: return super().tau @tau.setter def tau(self, tau: sympy.Symbol) -> None: self.temporal = TemporalSympyTau(tau) class MomentumSympy4D(LorentzMomentum, VectorSympy4D): """ Four dimensional momentum vector class for the SymPy backend. Examples: >>> import vector; import sympy >>> vec = vector.MomentumSympy4D(px=sympy.Symbol("px"), py=sympy.Symbol("py"), pz=sympy.Symbol("pz"), t=sympy.Symbol("t")) >>> vec.px, vec.py, vec.pz, vec.t (px, py, pz, t) >>> vec = vector.MomentumSympy4D(pt=sympy.Symbol("pt"), phi=sympy.Symbol("phi"), pz=sympy.Symbol("pz"), M=sympy.Symbol("M")) >>> vec.pt, vec.phi, vec.pz, vec.M (pt, phi, pz, M) For four dimensional SymPy vectors, see :class:`vector.backends.sympy.VectorSympy4D`. """ def __repr__(self) -> str: aznames = _coordinate_class_to_names[_aztype(self)] lnames = _coordinate_class_to_names[_ltype(self)] tnames = _coordinate_class_to_names[_ttype(self)] out = [] for x in aznames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.azimuthal, x)}") for x in lnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.longitudinal, x)}") for x in tnames: y = _repr_generic_to_momentum.get(x, x) out.append(f"{y}={getattr(self.temporal, x)}") return "MomentumSympy4D(" + ", ".join(out) + ")" @property def px(self) -> sympy.Symbol: return super().px @px.setter def px(self, px: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(px, self.py) @property def py(self) -> sympy.Symbol: return super().py @py.setter def py(self, py: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyXY(self.px, py) @property def pt(self) -> sympy.Symbol: return super().pt @pt.setter def pt(self, pt: sympy.Symbol) -> None: self.azimuthal = AzimuthalSympyRhoPhi(pt, self.phi) @property def pz(self) -> sympy.Symbol: return super().pz @pz.setter def pz(self, pz: sympy.Symbol) -> None: self.longitudinal = LongitudinalSympyZ(pz) @property def E(self) -> sympy.Symbol: return super().E @E.setter def E(self, E: sympy.Symbol) -> None: self.temporal = TemporalSympyT(E) @property def e(self) -> sympy.Symbol: return super().e @e.setter def e(self, e: sympy.Symbol) -> None: self.temporal = TemporalSympyT(e) @property def energy(self) -> sympy.Symbol: return super().energy @energy.setter def energy(self, energy: sympy.Symbol) -> None: self.temporal = TemporalSympyT(energy) @property def M(self) -> sympy.Symbol: return super().M @M.setter def M(self, M: sympy.Symbol) -> None: self.temporal = TemporalSympyTau(M) @property def m(self) -> sympy.Symbol: return super().m @m.setter def m(self, m: sympy.Symbol) -> None: self.temporal = TemporalSympyTau(m) @property def mass(self) -> sympy.Symbol: return super().mass @mass.setter def mass(self, mass: sympy.Symbol) -> None: self.temporal = TemporalSympyTau(mass) VectorSympy2D.ProjectionClass2D = VectorSympy2D VectorSympy2D.ProjectionClass3D = VectorSympy3D VectorSympy2D.ProjectionClass4D = VectorSympy4D VectorSympy2D.GenericClass = VectorSympy2D VectorSympy2D.MomentumClass = MomentumSympy2D MomentumSympy2D.ProjectionClass2D = MomentumSympy2D MomentumSympy2D.ProjectionClass3D = MomentumSympy3D MomentumSympy2D.ProjectionClass4D = MomentumSympy4D MomentumSympy2D.GenericClass = VectorSympy2D MomentumSympy2D.MomentumClass = MomentumSympy2D VectorSympy3D.ProjectionClass2D = VectorSympy2D VectorSympy3D.ProjectionClass3D = VectorSympy3D VectorSympy3D.ProjectionClass4D = VectorSympy4D VectorSympy3D.GenericClass = VectorSympy3D VectorSympy3D.MomentumClass = MomentumSympy3D MomentumSympy3D.ProjectionClass2D = MomentumSympy2D MomentumSympy3D.ProjectionClass3D = MomentumSympy3D MomentumSympy3D.ProjectionClass4D = MomentumSympy4D MomentumSympy3D.GenericClass = VectorSympy3D MomentumSympy3D.MomentumClass = MomentumSympy3D VectorSympy4D.ProjectionClass2D = VectorSympy2D VectorSympy4D.ProjectionClass3D = VectorSympy3D VectorSympy4D.ProjectionClass4D = VectorSympy4D VectorSympy4D.GenericClass = VectorSympy4D VectorSympy4D.MomentumClass = MomentumSympy4D MomentumSympy4D.ProjectionClass2D = MomentumSympy2D MomentumSympy4D.ProjectionClass3D = MomentumSympy3D MomentumSympy4D.ProjectionClass4D = MomentumSympy4D MomentumSympy4D.GenericClass = VectorSympy4D MomentumSympy4D.MomentumClass = MomentumSympy4D vector-1.6.3/src/vector/py.typed000066400000000000000000000000001503546127100166050ustar00rootroot00000000000000vector-1.6.3/tests/000077500000000000000000000000001503546127100141715ustar00rootroot00000000000000vector-1.6.3/tests/__init__.py000066400000000000000000000003641503546127100163050ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/backends/000077500000000000000000000000001503546127100157435ustar00rootroot00000000000000vector-1.6.3/tests/backends/__init__.py000066400000000000000000000003641503546127100200570ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/backends/test_awkward.py000066400000000000000000001101051503546127100210120ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import functools import importlib.metadata import numbers import numpy as np import packaging.version import pytest import vector from vector import VectorObject2D ak = pytest.importorskip("awkward") pytestmark = pytest.mark.awkward # Record reducers were added before awkward==2.2.3, but had some bugs. awkward_without_record_reducers = packaging.version.Version( importlib.metadata.version("awkward") ) < packaging.version.Version("2.2.3") def _assert(bool1, bool2, backend): # this is a helper function that differentiates between backends: # - 'typetracer' backend: # this method does nothing (no-op). typetracers can't be compared as they do not have numerical values. # this is useful, because the expression that goes into this function (bool1 & bool2) will still be computed, # i.e. this makes sure that the computation does (at least) not raise an error. # - other backends: # convert awkward arrays to lists and compare their values with python's builtin list comparison if backend != "typetracer": if isinstance(bool1, ak.Array): bool1 = bool1.tolist() if isinstance(bool2, ak.Array): bool2 = bool2.tolist() else: assert bool1 == bool2 @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_dimension_conversion(backend): assert_backend = functools.partial(_assert, backend=backend) # 2D -> 3D vec = vector.Array( [ [{"x": 1, "y": 1.1}, {"x": 2, "y": 2.1}], [], ], backend=backend, ) # test alias assert_backend(ak.all(vec.to_3D(z=1).z == 1), True) assert_backend(ak.all(vec.to_3D(eta=1).eta == 1), True) assert_backend(ak.all(vec.to_3D(theta=1).theta == 1), True) assert_backend(vec.to_Vector3D(z=1).x, vec.x) assert_backend(vec.to_Vector3D(z=1).y, vec.y) # 2D -> 4D assert_backend(ak.all(vec.to_Vector4D(z=1, t=1).t == 1), True) assert_backend(ak.all(vec.to_Vector4D(z=1, t=1).z == 1), True) assert_backend(ak.all(vec.to_Vector4D(eta=1, t=1).eta == 1), True) assert_backend(ak.all(vec.to_Vector4D(eta=1, t=1).t == 1), True) assert_backend(ak.all(vec.to_Vector4D(theta=1, t=1).theta == 1), True) assert_backend(ak.all(vec.to_Vector4D(theta=1, t=1).t == 1), True) assert_backend(ak.all(vec.to_Vector4D(z=1, tau=1).z == 1), True) assert_backend(ak.all(vec.to_Vector4D(z=1, tau=1).tau == 1), True) assert_backend(ak.all(vec.to_Vector4D(eta=1, tau=1).eta == 1), True) assert_backend(ak.all(vec.to_Vector4D(eta=1, tau=1).tau == 1), True) assert_backend(ak.all(vec.to_Vector4D(theta=1, tau=1).theta == 1), True) assert_backend(ak.all(vec.to_Vector4D(theta=1, tau=1).tau == 1), True) # test alias assert_backend(vec.to_4D(z=1, t=1).x, vec.x) assert_backend(vec.to_4D(z=1, t=1).y, vec.y) # 3D -> 4D vec = vector.Array( [ [{"x": 1, "y": 1.1, "z": 1.2}, {"x": 2, "y": 2.1, "z": 2.2}], [], ], backend=backend, ) assert_backend(ak.all(vec.to_Vector4D(t=1).t == 1), True) assert_backend(ak.all(vec.to_Vector4D(tau=1).tau == 1), True) assert_backend(vec.to_Vector4D(t=1).x, vec.x) assert_backend(vec.to_Vector4D(t=1).y, vec.y) assert_backend(vec.to_Vector4D(t=1).z, vec.z) # check if momentum coords work vec = vector.Array( [ [{"px": 1, "py": 1.1}, {"px": 2, "py": 2.1}], [], ], backend=backend, ) assert_backend(ak.all(vec.to_Vector3D(pz=1).pz == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, m=1).pz == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, m=1).m == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, mass=1).mass == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, M=1).M == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, e=1).e == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, energy=1).energy == 1), True) assert_backend(ak.all(vec.to_Vector4D(pz=1, E=1).E == 1), True) vec = vector.Array( [ [{"px": 1, "py": 1.1, "pz": 1.2}, {"px": 2, "py": 2.1, "pz": 2.2}], [], ], backend=backend, ) assert_backend(ak.all(vec.to_Vector4D(m=1).m == 1), True) assert_backend(ak.all(vec.to_Vector4D(mass=1).mass == 1), True) assert_backend(ak.all(vec.to_Vector4D(M=1).M == 1), True) assert_backend(ak.all(vec.to_Vector4D(e=1).e == 1), True) assert_backend(ak.all(vec.to_Vector4D(energy=1).energy == 1), True) assert_backend(ak.all(vec.to_Vector4D(E=1).E == 1), True) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_type_checks(backend): with pytest.raises(TypeError): vector.Array( [ [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": [0.2]}], [], ], backend=backend, ) with pytest.raises(TypeError): vector.Array( [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": complex(1, 2)}], backend=backend, ) with pytest.raises(TypeError): vector.zip( [ [{"x": 1, "y": 1.1, "z": 0.1}, {"x": 2, "y": 2.2, "z": 0.2}], [], ] ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_basic(backend): assert_backend = functools.partial(_assert, backend=backend) array = vector.Array( [[{"x": 1, "y": 2}], [], [{"x": 3, "y": 4}]], backend=backend, ) assert isinstance(array, vector.backends.awkward.VectorArray2D) assert_backend(array.x, [[1], [], [3]]) assert_backend(array.y, [[2], [], [4]]) assert_backend(array.rho, [[2.23606797749979], [], [5]]) assert isinstance(array[2, 0], vector.backends.awkward.VectorRecord2D) if backend == "cpu": (a,), (), (c,) = array.phi.tolist() assert a == pytest.approx(1.1071487177940904) assert c == pytest.approx(0.9272952180016122) assert array[2, 0].rho == 5 assert array.deltaphi(array).tolist() == [[0], [], [0]] else: # at least make sure they run _ = array.phi _ = array[2, 0].rho _ = array.deltaphi(array) array = vector.Array( [[{"pt": 1, "phi": 2}], [], [{"pt": 3, "phi": 4}]], backend=backend, ) assert isinstance(array, vector.backends.awkward.MomentumArray2D) assert_backend(array.pt, [[1], [], [3]]) array = vector.Array( [ [{"x": 1, "y": 2, "z": 3, "wow": 99}], [], [{"x": 4, "y": 5, "z": 6, "wow": 123}], ], backend=backend, ) assert isinstance(array, vector.backends.awkward.VectorArray3D) assert_backend(array.wow, [[99], [], [123]]) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_rotateZ(backend): assert_backend = functools.partial(_assert, backend=backend) array = vector.Array( [[{"pt": 1, "phi": 0}], [], [{"pt": 2, "phi": 1}]], backend=backend, ) out = array.rotateZ(1) assert isinstance(out, vector.backends.awkward.MomentumArray2D) assert_backend(out, [[{"rho": 1, "phi": 1}], [], [{"rho": 2, "phi": 2}]]) array = vector.Array( [[{"x": 1, "y": 0, "wow": 99}], [], [{"x": 2, "y": 1, "wow": 123}]], backend=backend, ) out = array.rotateZ(0.1) assert isinstance(out, vector.backends.awkward.VectorArray2D) assert_backend(out.wow, [[99], [], [123]]) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_projection(backend): assert_backend = functools.partial(_assert, backend=backend) array = vector.Array( [ [{"x": 1, "y": 2, "z": 3, "wow": 99}], [], [{"x": 4, "y": 5, "z": 6, "wow": 123}], ], backend=backend, ) out = array.to_Vector2D() assert isinstance(out, vector.backends.awkward.VectorArray2D) assert_backend( out, [ [{"x": 1, "y": 2, "wow": 99}], [], [{"x": 4, "y": 5, "wow": 123}], ], ) out = array.to_Vector4D() assert isinstance(out, vector.backends.awkward.VectorArray4D) assert_backend( out, [ [{"x": 1, "y": 2, "z": 3, "t": 0, "wow": 99}], [], [{"x": 4, "y": 5, "z": 6, "t": 0, "wow": 123}], ], ) out = array.to_rhophietatau() assert isinstance(out, vector.backends.awkward.VectorArray4D) if backend == "cpu": (a,), (), (c,) = out.tolist() assert a == pytest.approx( { "rho": 2.23606797749979, "phi": 1.1071487177940904, "eta": 1.1035868415601453, "tau": 0, "wow": 99, } ) assert c == pytest.approx( { "rho": 6.4031242374328485, "phi": 0.8960553845713439, "eta": 0.8361481196083127, "tau": 0, "wow": 123, } ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_add(backend): assert_backend = functools.partial(_assert, backend=backend) one = vector.Array( [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], [], [{"x": 3, "y": 3.3}]], backend=backend, ) two = vector.Array( [{"x": 10, "y": 20}, {"x": 100, "y": 200}, {"x": 1000, "y": 2000}], backend=backend, ) assert isinstance(one.add(two), vector.backends.awkward.VectorArray2D) assert isinstance(two.add(one), vector.backends.awkward.VectorArray2D) assert_backend( one.add(two), [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 1003, "y": 2003.3}], ], ) assert_backend( two.add(one), [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 1003, "y": 2003.3}], ], ) two = vector.array({"x": [10, 100, 1000], "y": [20, 200, 2000]}) assert isinstance(one.add(two), vector.backends.awkward.VectorArray2D) assert isinstance(two.add(one), vector.backends.awkward.VectorArray2D) assert_backend( one.add(two), [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 1003, "y": 2003.3}], ], ) assert_backend( two.add(one), [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 1003, "y": 2003.3}], ], ) two = vector.obj(x=10, y=20) assert isinstance(one.add(two), vector.backends.awkward.VectorArray2D) assert isinstance(two.add(one), vector.backends.awkward.VectorArray2D) assert_backend( one.add(two), [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 13, "y": 23.3}], ], ) assert_backend( two.add(one), [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 13, "y": 23.3}], ], ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_ufuncs(backend): assert_backend = functools.partial(_assert, backend=backend) array = vector.Array( [[{"x": 3, "y": 4}, {"x": 5, "y": 12}], [], [{"x": 8, "y": 15}]], backend=backend, ) assert_backend(abs(array), [[5, 13], [], [17]]) assert_backend((array**2), [[5**2, 13**2], [], [17**2]]) one = vector.Array( [[{"x": 1, "y": 1.1}, {"x": 2, "y": 2.2}], [], [{"x": 3, "y": 3.3}]], backend=backend, ) two = vector.Array( [{"x": 10, "y": 20}, {"x": 100, "y": 200}, {"x": 1000, "y": 2000}], backend=backend, ) assert_backend( one + two, [ [{"x": 11, "y": 21.1}, {"x": 12, "y": 22.2}], [], [{"x": 1003, "y": 2003.3}], ], ) assert_backend( one * 10, [ [{"x": 10, "y": 11}, {"x": 20, "y": 22}], [], [{"x": 30, "y": 33}], ], ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_zip(backend): assert_backend = functools.partial(_assert, backend=backend) v = ak.to_backend(vector.zip({"x": [[], [1]], "y": [[], [1]]}), backend=backend) assert isinstance(v, vector.backends.awkward.VectorArray2D) assert isinstance(v[1], vector.backends.awkward.VectorArray2D) assert isinstance(v[1, 0], vector.backends.awkward.VectorRecord2D) assert_backend(v, [[], [{"x": 1, "y": 1}]]) assert_backend(v.x, [[], [1]]) assert_backend(v.y, [[], [1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_sum_2d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"rho": 1.0, "phi": 0.1}, {"rho": 4.0, "phi": 0.2}, ], [ {"rho": 1.0, "phi": 0.3}, {"rho": 4.0, "phi": 0.4}, {"rho": 1.0, "phi": 0.1}, ], ], backend=backend, ) # typetracer backend does not implement ak.almost_equal if backend != "typetracer": assert_backend( ak.almost_equal( ak.sum(v, axis=0, keepdims=True), vector.Array( [ [ {"x": 1.950340654403632, "y": 0.3953536233081677}, {"x": 7.604510287376507, "y": 2.3523506924148467}, {"x": 0.9950041652780258, "y": 0.09983341664682815}, ] ] ), ), True, ) assert_backend( ak.almost_equal( ak.sum(v, axis=0, keepdims=False), vector.Array( [ {"x": 1.950340654403632, "y": 0.3953536233081677}, {"x": 7.604510287376507, "y": 2.3523506924148467}, {"x": 0.9950041652780258, "y": 0.09983341664682815}, ] ), ), True, ) assert_backend( ak.almost_equal( ak.sum(v, axis=1, keepdims=True), ak.to_regular( vector.Array( [ [{"x": 4.915270, "y": 0.89451074}], [{"x": 5.63458463, "y": 1.95302699}], ] ) ), ), True, ) assert_backend( ak.almost_equal( ak.sum(v, axis=1, keepdims=False), vector.Array( [ {"x": 4.915270, "y": 0.89451074}, {"x": 5.63458463, "y": 1.95302699}, ] ), ), True, ) assert_backend( ak.almost_equal( ak.sum(v.mask[[False, True]], axis=1), vector.Array( [ {"x": 5.63458463, "y": 1.95302699}, ] )[[None, 0]], ), True, ) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_sum_3d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}, ], [ {"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}, {"x": 1, "y": 1, "z": 1}, ], ], backend=backend, ) assert_backend( ak.sum(v, axis=0, keepdims=True), [ [ {"x": 2, "y": 4, "z": 6}, {"x": 8, "y": 10, "z": 12}, {"x": 1, "y": 1, "z": 1}, ] ], ) assert_backend( ak.sum(v, axis=0, keepdims=False), [ {"x": 2, "y": 4, "z": 6}, {"x": 8, "y": 10, "z": 12}, {"x": 1, "y": 1, "z": 1}, ], ) assert_backend( ak.sum(v, axis=1, keepdims=True), [ [{"x": 5, "y": 7, "z": 9}], [{"x": 6, "y": 8, "z": 10}], ], ) assert_backend( ak.sum(v, axis=1, keepdims=False), [ {"x": 5, "y": 7, "z": 9}, {"x": 6, "y": 8, "z": 10}, ], ) assert_backend( ak.sum(v.mask[[False, True]], axis=1), [ None, {"x": 6, "y": 8, "z": 10}, ], ) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_sum_4d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1, "y": 2, "z": 3, "t": 4}, {"x": 4, "y": 5, "z": 6, "t": 2}, {"x": 0, "y": 0, "z": 0, "t": 3}, ], [ {"x": 1, "y": 2, "z": 3, "t": 8}, {"x": 4, "y": 5, "z": 6, "t": 0}, {"x": 1, "y": 1, "z": 1, "t": 0}, ], ], backend=backend, ) assert_backend( ak.sum(v, axis=0, keepdims=True), [ [ {"t": 12, "z": 6, "x": 2, "y": 4}, {"t": 2, "z": 12, "x": 8, "y": 10}, {"t": 3, "z": 1, "x": 1, "y": 1}, ] ], ) assert_backend( ak.sum(v, axis=0, keepdims=False), [ {"t": 12, "z": 6, "x": 2, "y": 4}, {"t": 2, "z": 12, "x": 8, "y": 10}, {"t": 3, "z": 1, "x": 1, "y": 1}, ], ) assert_backend( ak.sum(v, axis=1, keepdims=True), [ [{"t": 9, "z": 9, "x": 5, "y": 7}], [{"t": 8, "z": 10, "x": 6, "y": 8}], ], ) assert_backend( ak.sum(v, axis=1, keepdims=False), [ {"t": 9, "z": 9, "x": 5, "y": 7}, {"t": 8, "z": 10, "x": 6, "y": 8}, ], ) assert_backend( ak.sum(v.mask[[False, True]], axis=1), [ None, {"t": 8, "z": 10, "x": 6, "y": 8}, ], ) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_count_nonzero_2d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"rho": 1.0, "phi": 0.1}, {"rho": 4.0, "phi": 0.2}, {"rho": 0.0, "phi": 0.0}, ], [ {"rho": 1.0, "phi": 0.3}, {"rho": 4.0, "phi": 0.4}, {"rho": 1.0, "phi": 0.1}, ], ], backend=backend, ) assert_backend(ak.count_nonzero(v, axis=1), [2, 3]) assert_backend(ak.count_nonzero(v, axis=1, keepdims=True), [[2], [3]]) assert_backend(ak.count_nonzero(v, axis=0), [2, 2, 1]) assert_backend(ak.count_nonzero(v, axis=0, keepdims=True), [[2, 2, 1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_count_nonzero_3d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1.0, "y": 2.0, "theta": 0.1}, {"x": 4.0, "y": 5.0, "theta": 0.2}, {"x": 0.0, "y": 0.0, "theta": 0.0}, ], [ {"x": 1.0, "y": 2.0, "theta": 0.6}, {"x": 4.0, "y": 5.0, "theta": 1.3}, {"x": 1.0, "y": 1.0, "theta": 1.9}, ], ], backend=backend, ) assert_backend(ak.count_nonzero(v, axis=1), [2, 3]) assert_backend(ak.count_nonzero(v, axis=1, keepdims=True), [[2], [3]]) assert_backend(ak.count_nonzero(v, axis=0), [2, 2, 1]) assert_backend(ak.count_nonzero(v, axis=0, keepdims=True), [[2, 2, 1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_count_nonzero_4d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1.0, "y": 2.0, "z": 3.0, "t": 4.0}, {"x": 4.0, "y": 5.0, "z": 6.0, "t": 2.0}, {"x": 0.0, "y": 0.0, "z": 0.0, "t": 3.0}, ], [ {"x": 1.0, "y": 2.0, "z": 3.0, "t": 8.0}, {"x": 4.0, "y": 5.0, "z": 6.0, "t": 0.0}, {"x": 1.0, "y": 1.0, "z": 1.0, "t": 0.0}, ], ], backend=backend, ) assert_backend(ak.count_nonzero(v, axis=1), [3, 3]) assert_backend(ak.count_nonzero(v, axis=1, keepdims=True), [[3], [3]]) assert_backend(ak.count_nonzero(v, axis=0), [2, 2, 2]) assert_backend(ak.count_nonzero(v, axis=0, keepdims=True), [[2, 2, 2]]) v2 = vector.Array( [ [ {"x": 1, "y": 2, "z": 3, "t": 1}, {"x": 4, "y": 5, "z": 6, "t": 2}, {"x": 0, "y": 0, "z": 0, "t": 2}, ], [ {"x": 1, "y": 2, "z": 3, "t": 0}, {"x": 4, "y": 5, "z": 6, "t": 1}, {"x": 0, "y": 0, "z": 0, "t": 0}, ], ], backend=backend, ) assert_backend(ak.count_nonzero(v2, axis=1), [3, 2]) assert_backend(ak.count_nonzero(v2, axis=1, keepdims=True), [[3], [2]]) assert_backend(ak.count_nonzero(v2, axis=0), [2, 2, 1]) assert_backend(ak.count_nonzero(v2, axis=0, keepdims=True), [[2, 2, 1]]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_count_2d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1, "y": 2}, {"x": 4, "y": 5}, {"x": 0, "y": 0}, ], [ {"x": 1, "y": 2}, {"x": 4, "y": 5}, {"x": 1, "y": 1}, ], ], backend=backend, ) assert_backend(ak.count(v, axis=1), [3, 3]) assert_backend(ak.count(v, axis=1, keepdims=True), [[3], [3]]) assert_backend(ak.count(v, axis=0), [2, 2, 2]) assert_backend(ak.count(v, axis=0, keepdims=True), [[2, 2, 2]]) assert_backend(ak.count(v.mask[[False, True]], axis=1), [None, 3]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_count_3d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}, {"x": 0, "y": 0, "z": 0}, ], [ {"x": 1, "y": 2, "z": 3}, {"x": 4, "y": 5, "z": 6}, {"x": 1, "y": 1, "z": 1}, ], ] ) assert_backend(ak.count(v, axis=1), [3, 3]) assert_backend(ak.count(v, axis=1, keepdims=True), [[3], [3]]) assert_backend(ak.count(v, axis=0), [2, 2, 2]) assert_backend(ak.count(v, axis=0, keepdims=True), [[2, 2, 2]]) assert_backend(ak.count(v.mask[[False, True]], axis=1), [None, 3]) @pytest.mark.skipif( awkward_without_record_reducers, reason="record reducers were added before awkward==2.2.3, but had some bugs", ) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_count_4d(backend): assert_backend = functools.partial(_assert, backend=backend) v = vector.Array( [ [ {"x": 1, "y": 2, "z": 3, "t": 9}, {"x": 4, "y": 5, "z": 6, "t": 9}, {"x": 0, "y": 0, "z": 0, "t": 9}, ], [ {"x": 1, "y": 2, "z": 3, "t": 9}, {"x": 4, "y": 5, "z": 6, "t": 9}, {"x": 1, "y": 1, "z": 1, "t": 9}, ], ] ) assert_backend(ak.count(v, axis=1), [3, 3]) assert_backend(ak.count(v, axis=1, keepdims=True), [[3], [3]]) assert_backend(ak.count(v, axis=0), [2, 2, 2]) assert_backend(ak.count(v, axis=0, keepdims=True), [[2, 2, 2]]) assert_backend(ak.count(v.mask[[False, True]], axis=1), [None, 3]) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_like(backend): assert_backend = functools.partial(_assert, backend=backend) v1 = ak.to_backend( vector.zip( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], }, ), backend=backend, ) v2 = ak.to_backend( vector.zip( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], "z": [5.0, 1.0, 1.0], }, ), backend=backend, ) v3 = ak.to_backend( vector.zip( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], "z": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ), backend=backend, ) v1_v2 = ak.to_backend( vector.zip( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], }, ), backend=backend, ) v2_v1 = ak.to_backend( vector.zip( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [5.0, 1.0, 1.0], }, ), backend=backend, ) v2_v3 = ak.to_backend( vector.zip( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [10.0, 2.0, 2.0], }, ), backend=backend, ) v3_v2 = ak.to_backend( vector.zip( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [10.0, 2.0, 2.0], "t": [16.0, 31.0, 46.0], }, ), backend=backend, ) v1_v3 = ak.to_backend( vector.zip( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ), backend=backend, ) with pytest.raises(TypeError): v1 + v2 with pytest.raises(TypeError): v2 + v3 with pytest.raises(TypeError): v1 + v3 with pytest.raises(TypeError): v1 - v2 with pytest.raises(TypeError): v2 - v3 with pytest.raises(TypeError): v1 - v3 with pytest.raises(TypeError): v1.equal(v2) with pytest.raises(TypeError): v2.equal(v3) with pytest.raises(TypeError): v1.equal(v3) with pytest.raises(TypeError): v1.not_equal(v2) with pytest.raises(TypeError): v2.not_equal(v3) with pytest.raises(TypeError): v1.not_equal(v3) with pytest.raises(TypeError): v1.dot(v2) with pytest.raises(TypeError): v2.dot(v3) with pytest.raises(TypeError): v1.dot(v3) # 2D + 3D.like(2D) = 2D assert_backend(v1 + v2.like(v1), v1_v2) assert_backend(v2.like(v1) + v1, v1_v2) # 2D + 4D.like(2D) = 2D assert_backend(v1 + v3.like(v1), v1_v2) assert_backend(v3.like(v1) + v1, v1_v2) # 3D + 2D.like(3D) = 3D assert_backend(v2 + v1.like(v2), v2_v1) assert_backend(v1.like(v2) + v2, v2_v1) # 3D + 4D.like(3D) = 3D assert_backend(v2 + v3.like(v2), v2_v3) assert_backend(v3.like(v2) + v2, v2_v3) # 4D + 2D.like(4D) = 4D assert_backend(v3 + v1.like(v3), v1_v3) assert_backend(v1.like(v3) + v3, v1_v3) # 4D + 3D.like(4D) = 4D assert_backend(v3 + v2.like(v3), v3_v2) assert_backend(v2.like(v3) + v3, v3_v2) v1 = ak.to_backend( vector.zip( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], }, ), backend=backend, ) v2 = ak.to_backend( vector.zip( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], "pz": [5.0, 1.0, 1.0], }, ), backend=backend, ) v3 = ak.to_backend( vector.zip( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], "pz": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ), backend=backend, ) # 2D + 3D.like(2D) = 2D assert_backend(v1 + v2.like(v1), v1_v2) assert_backend(v2.like(v1) + v1, v1_v2) # 2D + 4D.like(2D) = 2D assert_backend(v1 + v3.like(v1), v1_v2) assert_backend(v3.like(v1) + v1, v1_v2) # 3D + 2D.like(3D) = 3D assert_backend(v2 + v1.like(v2), v2_v1) assert_backend(v1.like(v2) + v2, v2_v1) # 3D + 4D.like(3D) = 3D assert_backend(v2 + v3.like(v2), v2_v3) assert_backend(v3.like(v2) + v2, v2_v3) # 4D + 2D.like(4D) = 4D assert_backend(v3 + v1.like(v3), v1_v3) assert_backend(v1.like(v3) + v3, v1_v3) # 4D + 3D.like(4D) = 4D assert_backend(v3 + v2.like(v3), v3_v2) assert_backend(v2.like(v3) + v3, v3_v2) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_handler_of(backend): assert_backend = functools.partial(_assert, backend=backend) numpy_vec = ak.to_backend( vector.array( { "x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5], "z": [3.1, 3.2, 3.3, 3.4, 3.5], } ), backend=backend, ) awkward_vec2 = ak.to_backend( vector.zip( { "x": [1.0, 2.0, 3.0], "y": [-1.0, 2.0, 3.0], "z": [5.0, 10.0, 15.0], "t": [16.0, 31.0, 46.0], }, ), backend=backend, ) object_vec = VectorObject2D.from_xy(1.0, 1.0) protocol = vector._methods._handler_of(object_vec, numpy_vec, awkward_vec2) # chooses awkward backend assert_backend(protocol, awkward_vec2) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_momentum_coordinate_transforms(backend): awkward_vec = ak.to_backend( vector.zip( { "px": [1.0, 2.0, 3.0], "py": [-1.0, 2.0, 3.0], }, ), backend=backend, ) for t1 in "pxpy", "ptphi": for t2 in "pz", "eta", "theta": for t3 in "mass", "energy": transformed_object = getattr(awkward_vec, "to_" + t1)() assert isinstance( transformed_object, vector.backends.awkward.MomentumAwkward2D ) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) transformed_object = getattr(awkward_vec, "to_" + t1 + t2)() assert isinstance( transformed_object, vector.backends.awkward.MomentumAwkward3D ) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) assert hasattr(transformed_object, t2) transformed_object = getattr(awkward_vec, "to_" + t1 + t2 + t3)() assert isinstance( transformed_object, vector.backends.awkward.MomentumAwkward4D ) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) assert hasattr(transformed_object, t2) assert hasattr(transformed_object, t3) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_momentum_preservation(backend): v1 = ak.to_backend( vector.zip( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], }, ), backend=backend, ) v2 = ak.to_backend( vector.zip( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], "z": [5.0, 1.0, 1.0], }, ), backend=backend, ) v3 = ak.to_backend( vector.zip( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], "pz": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ), backend=backend, ) # momentum + generic = momentum # 2D + 3D.like(2D) = 2D assert isinstance(v1 + v2.like(v1), vector.backends.awkward.MomentumAwkward2D) assert isinstance(v2.like(v1) + v1, vector.backends.awkward.MomentumAwkward2D) # 2D + 4D.like(2D) = 2D assert isinstance(v1 + v3.like(v1), vector.backends.awkward.MomentumAwkward2D) assert isinstance(v3.like(v1) + v1, vector.backends.awkward.MomentumAwkward2D) # 3D + 2D.like(3D) = 3D assert isinstance(v2 + v1.like(v2), vector.backends.awkward.MomentumAwkward3D) assert isinstance(v1.like(v2) + v2, vector.backends.awkward.MomentumAwkward3D) # 3D + 4D.like(3D) = 3D assert isinstance(v2 + v3.like(v2), vector.backends.awkward.MomentumAwkward3D) assert isinstance(v3.like(v2) + v2, vector.backends.awkward.MomentumAwkward3D) # 4D + 2D.like(4D) = 4D assert isinstance(v3 + v1.like(v3), vector.backends.awkward.MomentumAwkward4D) assert isinstance(v1.like(v3) + v3, vector.backends.awkward.MomentumAwkward4D) # 4D + 3D.like(4D) = 4D assert isinstance(v3 + v2.like(v3), vector.backends.awkward.MomentumAwkward4D) assert isinstance(v2.like(v3) + v3, vector.backends.awkward.MomentumAwkward4D) @pytest.mark.parametrize("backend", ["cpu", "typetracer"]) def test_subclass_fields(backend): @ak.mixin_class(vector.backends.awkward.behavior) class TwoVector(vector.backends.awkward.MomentumAwkward2D): pass @ak.mixin_class(vector.backends.awkward.behavior) class ThreeVector(vector.backends.awkward.MomentumAwkward3D): pass @ak.mixin_class(vector.backends.awkward.behavior) class LorentzVector(vector.backends.awkward.MomentumAwkward4D): @ak.mixin_class_method(np.divide, {numbers.Number}) def divide(self, factor): return self.scale(1 / factor) LorentzVectorArray.ProjectionClass2D = TwoVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass3D = ThreeVectorArray # noqa: F821 LorentzVectorArray.ProjectionClass4D = LorentzVectorArray # noqa: F821 LorentzVectorArray.MomentumClass = LorentzVectorArray # noqa: F821 vec = ak.to_backend( ak.zip( { "pt": [[1, 2], [], [3], [4]], "eta": [[1.2, 1.4], [], [1.6], [3.4]], "phi": [[0.3, 0.4], [], [0.5], [0.6]], "energy": [[50, 51], [], [52], [60]], }, with_name="LorentzVector", behavior=vector.backends.awkward.behavior, ), backend=backend, ) assert vec.like(vector.obj(x=1, y=2)).fields == ["rho", "phi"] assert vec.like(vector.obj(x=1, y=2, z=3)).fields == ["rho", "phi", "eta"] assert (vec / 2).fields == ["rho", "phi", "eta", "t"] vector-1.6.3/tests/backends/test_awkward_numba.py000066400000000000000000000022361503546127100222010ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector ak = pytest.importorskip("awkward") numba = pytest.importorskip("numba") pytest.importorskip("vector.backends._numba_object") pytestmark = [pytest.mark.numba, pytest.mark.awkward] def test(): @numba.njit def extract(x): return x[2][0] array = vector.Array([[{"x": 1, "y": 2}], [], [{"x": 3, "y": 4}, {"x": 5, "y": 6}]]) out = extract(array) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(3) assert out.y == pytest.approx(4) array = vector.Array( [[{"x": 1, "y": 2, "z": 3, "E": 4}], [], [{"x": 5, "y": 6, "z": 7, "E": 15}]] ) out = extract(array) assert isinstance(out, vector.backends.object.MomentumObject4D) assert out.x == pytest.approx(5) assert out.y == pytest.approx(6) assert out.z == pytest.approx(7) assert out.t == pytest.approx(15) vector-1.6.3/tests/backends/test_dask_awkward.py000066400000000000000000000030211503546127100220120ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector dak = pytest.importorskip("dask_awkward") ak = pytest.importorskip("awkward") def test_constructor(): x = dak.from_awkward( ak.Array([{"x": 1, "y": 2}, {"x": 1.1, "y": 2.2}]), npartitions=1 ) vec = vector.Array(x) assert isinstance(vec, dak.Array) assert isinstance(vec.compute(), vector.backends.awkward.VectorAwkward2D) assert ak.all(vec.x.compute() == ak.Array([1, 1.1])) assert ak.all(vec.y.compute() == ak.Array([2, 2.2])) def test_necessary_columns(): vec = vector.Array([[{"pt": 1, "phi": 2}], [], [{"pt": 3, "phi": 4}]]) dak_vec = dak.from_awkward(vec, npartitions=1) cols = next(iter(dak.report_necessary_columns(dak_vec).values())) # this may seem weird at first: why would one need "phi" and "rho", if one asked for "pt"? # the reason is that vector will build internally a class with "phi" and "rho", # see: https://github.com/scikit-hep/vector/blob/608da2d55a74eed25635fd408d1075b568773c99/src/vector/backends/awkward.py#L166-L167 # So, even if one asks for "pt", "phi" and "rho" are as well in order to build the vector class in the first place. # (the same argument holds true for all other vector classes) assert cols == frozenset({"phi", "rho"}) vector-1.6.3/tests/backends/test_numba_numpy.py000066400000000000000000000013651503546127100217130ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector import vector.backends.object numba = pytest.importorskip("numba") import vector.backends.numba_numpy # noqa: E402 pytestmark = pytest.mark.numba def test_pass_through(): @numba.njit def pass_through(obj): return obj array = vector.array({"px": [1, 2, 3], "py": [10, 20, 30]}) assert isinstance(array, vector.backends.numpy.VectorNumpy) assert isinstance(pass_through(array), vector.backends.numpy.VectorNumpy) vector-1.6.3/tests/backends/test_numba_object.py000066400000000000000000001523371503546127100220170ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import sys import numpy import pytest import vector import vector.backends.object numba = pytest.importorskip("numba") import vector.backends._numba_object # noqa: E402 pytestmark = pytest.mark.numba def test_namedtuples(): @numba.njit def get_x(obj): return obj.x assert get_x(vector.backends.object.AzimuthalObjectXY(1, 2.2)) == 1 assert get_x(vector.backends.object.AzimuthalObjectXY(1.1, 2)) == 1.1 def test_VectorObjectType(): # These tests verify that the reference counts for Python objects touched in # the lowered Numba code do not increase or decrease with the number of times # the function is run. @numba.njit def zero(obj): return None @numba.njit def one(obj): return obj @numba.njit def two(obj): return obj, obj obj = vector.obj(x=1, y=2) assert (sys.getrefcount(obj), sys.getrefcount(obj.azimuthal)) == (2, 2) class_refs = None for _ in range(10): zero(obj) assert (sys.getrefcount(obj), sys.getrefcount(obj.azimuthal)) == (2, 2) if class_refs is None: class_refs = sys.getrefcount(vector.backends.object.VectorObject2D) assert class_refs + 1 == sys.getrefcount(vector.backends.object.VectorObject2D) class_refs = None for _ in range(10): a = one(obj) assert (sys.getrefcount(obj), sys.getrefcount(obj.azimuthal)) == (2, 2) assert (sys.getrefcount(a), sys.getrefcount(a.azimuthal)) == (2, 2) if class_refs is None: class_refs = sys.getrefcount(vector.backends.object.VectorObject2D) assert class_refs + 1 == sys.getrefcount(vector.backends.object.VectorObject2D) class_refs = None for _ in range(10): a, b = two(obj) assert (sys.getrefcount(obj), sys.getrefcount(obj.azimuthal)) == (2, 2) assert ( sys.getrefcount(a), sys.getrefcount(a.azimuthal), sys.getrefcount(b), sys.getrefcount(b.azimuthal), ) == (2, 2, 2, 2) if class_refs is None: class_refs = sys.getrefcount(vector.backends.object.VectorObject2D) assert class_refs + 1 == sys.getrefcount(vector.backends.object.VectorObject2D) # These tests just check that the rest of the implementations are sane. obj = vector.obj(x=1, y=2) out = one(obj) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) obj = vector.obj(px=1, py=2) out = one(obj) assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) obj = vector.obj(x=1, y=2, z=3) out = one(obj) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(3) obj = vector.obj(px=1, py=2, pz=3) out = one(obj) assert isinstance(out, vector.backends.object.MomentumObject3D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(3) obj = vector.obj(x=1, y=2, z=3, t=4) out = one(obj) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(3) assert out.t == pytest.approx(4) obj = vector.obj(px=1, py=2, pz=3, t=4) out = one(obj) assert isinstance(out, vector.backends.object.MomentumObject4D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(3) assert out.t == pytest.approx(4) def test_VectorObjectconstructor(): @numba.njit def vector_xy(): return vector.backends.object.VectorObject2D( vector.backends.object.AzimuthalObjectXY(1, 2.2) ) @numba.njit def vector_rhophi(): return vector.backends.object.VectorObject2D( vector.backends.object.AzimuthalObjectRhoPhi(1, 2.2) ) @numba.njit def momentum_xy(): return vector.backends.object.MomentumObject2D( vector.backends.object.AzimuthalObjectXY(1, 2.2) ) @numba.njit def momentum_rhophi(): return vector.backends.object.MomentumObject2D( vector.backends.object.AzimuthalObjectRhoPhi(1, 2.2) ) @numba.njit def vector_xyz(): return vector.backends.object.VectorObject3D( vector.backends.object.AzimuthalObjectXY(1, 2.2), vector.backends.object.LongitudinalObjectZ(3), ) @numba.njit def momentum_xyz(): return vector.backends.object.MomentumObject3D( vector.backends.object.AzimuthalObjectXY(1, 2.2), vector.backends.object.LongitudinalObjectZ(3), ) @numba.njit def vector_rhophitheta(): return vector.backends.object.VectorObject3D( vector.backends.object.AzimuthalObjectRhoPhi(1, 2.2), vector.backends.object.LongitudinalObjectTheta(3), ) @numba.njit def momentum_rhophitheta(): return vector.backends.object.MomentumObject3D( vector.backends.object.AzimuthalObjectRhoPhi(1, 2.2), vector.backends.object.LongitudinalObjectTheta(3), ) @numba.njit def vector_xyzt(): return vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 2.2), vector.backends.object.LongitudinalObjectZ(3), vector.backends.object.TemporalObjectT(4), ) @numba.njit def momentum_xyzt(): return vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(1, 2.2), vector.backends.object.LongitudinalObjectZ(3), vector.backends.object.TemporalObjectT(4), ) @numba.njit def vector_rhophietatau(): return vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 2.2), vector.backends.object.LongitudinalObjectEta(3), vector.backends.object.TemporalObjectTau(4), ) @numba.njit def momentum_rhophietatau(): return vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 2.2), vector.backends.object.LongitudinalObjectEta(3), vector.backends.object.TemporalObjectTau(4), ) out = vector_xy() assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2.2) out = vector_rhophi() assert isinstance(out, vector.backends.object.VectorObject2D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(2.2) out = momentum_xy() assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2.2) out = momentum_rhophi() assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(2.2) out = vector_xyz() assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2.2) assert out.z == pytest.approx(3) out = momentum_xyz() assert isinstance(out, vector.backends.object.MomentumObject3D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2.2) assert out.z == pytest.approx(3) out = vector_rhophitheta() assert isinstance(out, vector.backends.object.VectorObject3D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(2.2) assert out.theta == pytest.approx(3) out = momentum_rhophitheta() assert isinstance(out, vector.backends.object.MomentumObject3D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(2.2) assert out.theta == pytest.approx(3) out = vector_xyzt() assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2.2) assert out.z == pytest.approx(3) assert out.t == pytest.approx(4) out = momentum_xyzt() assert isinstance(out, vector.backends.object.MomentumObject4D) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2.2) assert out.z == pytest.approx(3) assert out.t == pytest.approx(4) out = vector_rhophietatau() assert isinstance(out, vector.backends.object.VectorObject4D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(2.2) assert out.eta == pytest.approx(3) assert out.tau == pytest.approx(4) out = momentum_rhophietatau() assert isinstance(out, vector.backends.object.MomentumObject4D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(2.2) assert out.eta == pytest.approx(3) assert out.tau == pytest.approx(4) def test_projections(): @numba.njit def to_Vector2D(x): return x.to_Vector2D() @numba.njit def to_Vector3D(x): return x.to_Vector3D() @numba.njit def to_Vector4D(x): return x.to_Vector4D() assert isinstance( to_Vector2D(vector.obj(x=1.1, y=2.2)), vector.backends.object.VectorObject2D ) assert isinstance( to_Vector2D(vector.obj(x=1.1, y=2.2, z=3.3)), vector.backends.object.VectorObject2D, ) assert isinstance( to_Vector2D(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4)), vector.backends.object.VectorObject2D, ) assert isinstance( to_Vector2D(vector.obj(px=1.1, py=2.2)), vector.backends.object.MomentumObject2D, ) assert isinstance( to_Vector2D(vector.obj(px=1.1, py=2.2, pz=3.3)), vector.backends.object.MomentumObject2D, ) assert isinstance( to_Vector2D(vector.obj(px=1.1, py=2.2, pz=3.3, E=4.4)), vector.backends.object.MomentumObject2D, ) assert isinstance( to_Vector3D(vector.obj(x=1.1, y=2.2)), vector.backends.object.VectorObject3D ) assert isinstance( to_Vector3D(vector.obj(x=1.1, y=2.2, z=3.3)), vector.backends.object.VectorObject3D, ) assert isinstance( to_Vector3D(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4)), vector.backends.object.VectorObject3D, ) assert isinstance( to_Vector3D(vector.obj(px=1.1, py=2.2)), vector.backends.object.MomentumObject3D, ) assert isinstance( to_Vector3D(vector.obj(px=1.1, py=2.2, pz=3.3)), vector.backends.object.MomentumObject3D, ) assert isinstance( to_Vector3D(vector.obj(px=1.1, py=2.2, pz=3.3, E=4.4)), vector.backends.object.MomentumObject3D, ) assert isinstance( to_Vector4D(vector.obj(x=1.1, y=2.2)), vector.backends.object.VectorObject4D ) assert isinstance( to_Vector4D(vector.obj(x=1.1, y=2.2, z=3.3)), vector.backends.object.VectorObject4D, ) assert isinstance( to_Vector4D(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4)), vector.backends.object.VectorObject4D, ) assert isinstance( to_Vector4D(vector.obj(px=1.1, py=2.2)), vector.backends.object.MomentumObject4D, ) assert isinstance( to_Vector4D(vector.obj(px=1.1, py=2.2, pz=3.3)), vector.backends.object.MomentumObject4D, ) assert isinstance( to_Vector4D(vector.obj(px=1.1, py=2.2, pz=3.3, E=4.4)), vector.backends.object.MomentumObject4D, ) def test_conversions(): @numba.njit def to_xy(x): return x.to_xy() @numba.njit def to_rhophi(x): return x.to_rhophi() @numba.njit def to_xyz(x): return x.to_xyz() @numba.njit def to_rhophiz(x): return x.to_rhophiz() @numba.njit def to_xytheta(x): return x.to_xytheta() @numba.njit def to_rhophitheta(x): return x.to_rhophitheta() @numba.njit def to_xyeta(x): return x.to_xyeta() @numba.njit def to_rhophieta(x): return x.to_rhophieta() @numba.njit def to_xyzt(x): return x.to_xyzt() @numba.njit def to_rhophizt(x): return x.to_rhophizt() @numba.njit def to_xythetat(x): return x.to_xythetat() @numba.njit def to_rhophithetat(x): return x.to_rhophithetat() @numba.njit def to_xyetat(x): return x.to_xyetat() @numba.njit def to_rhophietat(x): return x.to_rhophietat() @numba.njit def to_xyztau(x): return x.to_xyztau() @numba.njit def to_rhophiztau(x): return x.to_rhophiztau() @numba.njit def to_xythetatau(x): return x.to_xythetatau() @numba.njit def to_rhophithetatau(x): return x.to_rhophithetatau() @numba.njit def to_xyetatau(x): return x.to_xyetatau() @numba.njit def to_rhophietatau(x): return x.to_rhophietatau() for v in ( vector.obj(x=1.1, y=2.2), vector.obj(px=1.1, py=2.2), vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(px=1.1, py=2.2, pz=3.3), vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.obj(px=1.1, py=2.2, pz=3.3, E=4.4), ): print(v) out = to_xy(v) assert isinstance(out, vector.backends.object.VectorObject2D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject2D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) out = to_rhophi(v) assert isinstance(out, vector.backends.object.VectorObject2D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject2D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectRhoPhi) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) out = to_xyz(v) assert isinstance(out, vector.backends.object.VectorObject3D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) if isinstance(v, vector.backends.object.VectorObject2D): assert out.z == pytest.approx(0) else: assert out.z == pytest.approx(3.3) out = to_rhophiz(v) assert isinstance(out, vector.backends.object.VectorObject3D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectRhoPhi) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) if isinstance(v, vector.backends.object.VectorObject2D): assert out.z == pytest.approx(0) else: assert out.z == pytest.approx(3.3) out = to_xytheta(v) assert isinstance(out, vector.backends.object.VectorObject3D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectTheta ) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) if isinstance(v, vector.backends.object.VectorObject2D): assert out.theta == pytest.approx(0) else: assert out.z == pytest.approx(3.3) out = to_rhophitheta(v) assert isinstance(out, vector.backends.object.VectorObject3D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectRhoPhi) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectTheta ) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) if isinstance(v, vector.backends.object.VectorObject2D): assert out.theta == pytest.approx(0) else: assert out.z == pytest.approx(3.3) out = to_xyeta(v) assert isinstance(out, vector.backends.object.VectorObject3D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectEta ) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) if isinstance(v, vector.backends.object.VectorObject2D): assert out.eta == pytest.approx(0) else: assert out.z == pytest.approx(3.3) out = to_rhophietatau(v) assert isinstance(out, vector.backends.object.VectorObject4D) if isinstance(v, vector._methods.Momentum): assert isinstance(out, vector.backends.object.MomentumObject4D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectRhoPhi) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectEta ) assert isinstance(out.temporal, vector.backends.object.TemporalObjectTau) assert out.x == pytest.approx(1.1) assert out.y == pytest.approx(2.2) if isinstance(v, vector.backends.object.VectorObject2D): assert out.eta == pytest.approx(0) else: assert out.z == pytest.approx(3.3) if isinstance( v, ( vector.backends.object.VectorObject2D, vector.backends.object.VectorObject3D, ), ): assert out.tau == pytest.approx(0) else: assert out.t == pytest.approx(4.4) def test_factory(): @numba.njit def vector_xy(): return vector.obj(x=2, y=3.3) @numba.njit def momentum_xy(): return vector.obj(px=2, py=3.3) @numba.njit def vector_rhophi(): return vector.obj(rho=2, phi=3.3) @numba.njit def momentum_rhophi(): return vector.obj(pt=2, phi=3.3) @numba.njit def vector_xyz(): return vector.obj(x=2, y=3.3, z=5) @numba.njit def momentum_xyz(): return vector.obj(x=2, y=3.3, pz=5) @numba.njit def vector_rhophieta(): return vector.obj(rho=2, phi=3.3, eta=5) @numba.njit def momentum_rhophieta(): return vector.obj(pt=2, phi=3.3, eta=5) @numba.njit def vector_xyztau(): return vector.obj(x=2, y=3.3, z=5, tau=10) @numba.njit def momentum_xyztau(): return vector.obj(x=2, y=3.3, z=5, m=10) @numba.njit def vector_rhophizt(): return vector.obj(rho=2, phi=3.3, z=5, t=10) @numba.njit def momentum_rhophizt(): return vector.obj(rho=2, phi=3.3, z=5, energy=10) out = vector_xy() assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(3.3) out = momentum_xy() assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(3.3) out = vector_rhophi() assert isinstance(out, vector.backends.object.VectorObject2D) assert out.rho == pytest.approx(2) assert out.phi == pytest.approx(3.3) out = momentum_rhophi() assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.pt == pytest.approx(2) assert out.phi == pytest.approx(3.3) out = vector_xyz() assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(3.3) assert out.z == pytest.approx(5) out = momentum_xyz() assert isinstance(out, vector.backends.object.MomentumObject3D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(3.3) assert out.z == pytest.approx(5) out = vector_rhophieta() assert isinstance(out, vector.backends.object.VectorObject3D) assert out.rho == pytest.approx(2) assert out.phi == pytest.approx(3.3) assert out.eta == pytest.approx(5) out = momentum_rhophieta() assert isinstance(out, vector.backends.object.MomentumObject3D) assert out.pt == pytest.approx(2) assert out.phi == pytest.approx(3.3) assert out.eta == pytest.approx(5) out = vector_xyztau() assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(3.3) assert out.z == pytest.approx(5) assert out.tau == pytest.approx(10) out = momentum_xyztau() assert isinstance(out, vector.backends.object.MomentumObject4D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(3.3) assert out.z == pytest.approx(5) assert out.tau == pytest.approx(10) out = vector_rhophizt() assert isinstance(out, vector.backends.object.VectorObject4D) assert out.rho == pytest.approx(2) assert out.phi == pytest.approx(3.3) assert out.z == pytest.approx(5) assert out.t == pytest.approx(10) out = momentum_rhophizt() assert isinstance(out, vector.backends.object.MomentumObject4D) assert out.pt == pytest.approx(2) assert out.phi == pytest.approx(3.3) assert out.z == pytest.approx(5) assert out.E == pytest.approx(10) # noqa: SIM300 def test_property_float(): @numba.njit def get_x(v): return v.x @numba.njit def get_z(v): return v.z @numba.njit def get_t(v): return v.t @numba.njit def get_Et(v): return v.Et assert get_x(vector.obj(x=1.1, y=2)) == pytest.approx(1.1) assert get_x(vector.obj(px=1.1, py=2)) == pytest.approx(1.1) assert get_x(vector.obj(x=1.1, y=2, z=3)) == pytest.approx(1.1) assert get_x(vector.obj(px=1.1, py=2, pz=3)) == pytest.approx(1.1) assert get_x(vector.obj(x=1.1, y=2, z=3, t=4)) == pytest.approx(1.1) assert get_x(vector.obj(px=1.1, py=2, pz=3, E=4)) == pytest.approx(1.1) assert get_x(vector.obj(rho=1, phi=0)) == pytest.approx(1) assert get_x(vector.obj(rho=1, phi=numpy.pi / 4)) == pytest.approx( 1 / numpy.sqrt(2) ) assert get_x(vector.obj(rho=1, phi=numpy.pi / 2)) == pytest.approx(0) with pytest.raises(numba.TypingError): get_z(vector.obj(x=1, y=2)) assert get_z(vector.obj(x=1, y=2, z=3)) == pytest.approx(3) assert get_z(vector.obj(px=1, py=2, pz=3)) == pytest.approx(3) with pytest.raises(numba.TypingError): get_t(vector.obj(x=1, y=2)) with pytest.raises(numba.TypingError): get_t(vector.obj(x=1, y=2, z=3)) assert get_t(vector.obj(x=1, y=2, z=3, t=4)) == pytest.approx(4) assert get_t(vector.obj(px=1, py=2, pz=3, E=4)) == pytest.approx(4) with pytest.raises(numba.TypingError): get_Et(vector.obj(x=1, y=2)) with pytest.raises(numba.TypingError): get_Et(vector.obj(x=1, y=2, z=3)) with pytest.raises(numba.TypingError): get_Et(vector.obj(x=1, y=2, z=3, t=4)) assert get_Et(vector.obj(px=1, py=2, pz=3, E=4)) == pytest.approx( numpy.sqrt(4**2 * (1**2 + 2**2) / (1**2 + 2**2 + 3**2)) ) def test_planar_method_float(): @numba.njit def get_deltaphi(v1, v2): return v1.deltaphi(v2) assert get_deltaphi(vector.obj(x=1, y=0), vector.obj(x=0, y=1)) == pytest.approx( -numpy.pi / 2 ) assert get_deltaphi(vector.obj(x=1, y=0), vector.obj(px=0, py=1)) == pytest.approx( -numpy.pi / 2 ) assert get_deltaphi( vector.obj(px=1, py=0), vector.obj(px=0, py=1) ) == pytest.approx(-numpy.pi / 2) def test_spatial_method_float(): @numba.njit def get_deltaeta(v1, v2): return v1.deltaeta(v2) assert get_deltaeta( vector.obj(x=1, y=0, eta=2.5), vector.obj(x=0, y=1, eta=1) ) == pytest.approx(1.5) def test_method_deltaphi(): @numba.njit def get_deltaphi(v1, v2): return v1.deltaphi(v2) assert get_deltaphi( vector.obj(rho=1.1, phi=2.2), vector.obj(rho=3, phi=4) ) == pytest.approx(2.2 - 4) assert get_deltaphi( vector.obj(rho=1.1, phi=2.2), vector.obj(rho=3, phi=4, z=5) ) == pytest.approx(2.2 - 4) assert get_deltaphi( vector.obj(rho=1.1, phi=2.2, z=3.3), vector.obj(rho=3, phi=4) ) == pytest.approx(2.2 - 4) assert get_deltaphi( vector.obj(rho=1.1, phi=2.2, z=3.3), vector.obj(rho=3, phi=4, z=5) ) == pytest.approx(2.2 - 4) def test_method_add(): @numba.njit def get_add(v1, v2): return v1.add(v2) out = get_add(vector.obj(x=1.1, y=2.2), vector.obj(x=3, y=4)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(x=1.1, y=2.2), vector.obj(x=3, y=4, z=5)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(x=3, y=4)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(x=3, y=4, z=5)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) assert out.z == pytest.approx(8.3) out = get_add(vector.obj(px=1.1, py=2.2), vector.obj(px=3, py=4)) assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(px=1.1, py=2.2), vector.obj(px=3, py=4, pz=5)) assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(px=1.1, py=2.2, pz=3.3), vector.obj(px=3, py=4)) assert isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(px=1.1, py=2.2, pz=3.3), vector.obj(px=3, py=4, pz=5)) assert isinstance(out, vector.backends.object.MomentumObject3D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) assert out.z == pytest.approx(8.3) out = get_add(vector.obj(x=1.1, y=2.2), vector.obj(px=3, py=4)) assert not isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(x=1.1, y=2.2), vector.obj(px=3, py=4, pz=5)) assert not isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(px=3, py=4)) assert not isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(px=3, py=4, pz=5)) assert not isinstance(out, vector.backends.object.MomentumObject3D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) assert out.z == pytest.approx(8.3) out = get_add(vector.obj(px=1.1, py=2.2), vector.obj(x=3, y=4)) assert not isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(px=1.1, py=2.2), vector.obj(x=3, y=4, z=5)) assert not isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(px=1.1, py=2.2, pz=3.3), vector.obj(x=3, y=4)) assert not isinstance(out, vector.backends.object.MomentumObject2D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) out = get_add(vector.obj(px=1.1, py=2.2, pz=3.3), vector.obj(x=3, y=4, z=5)) assert not isinstance(out, vector.backends.object.MomentumObject3D) assert out.x == pytest.approx(4.1) assert out.y == pytest.approx(6.2) assert out.z == pytest.approx(8.3) def test_method_isparallel(): @numba.njit def get_isparallel(v1, v2): return v1.is_parallel(v2) assert get_isparallel(vector.obj(x=1.1, y=2.2), vector.obj(x=2.2, y=4.4)) assert get_isparallel(vector.obj(x=1.1, y=2.2), vector.obj(px=2.2, py=4.4)) assert get_isparallel(vector.obj(px=1.1, py=2.2), vector.obj(px=2.2, py=4.4)) assert get_isparallel(vector.obj(x=1.1, y=2.2), vector.obj(x=2.2, y=4.4, z=0.0)) assert get_isparallel(vector.obj(x=1.1, y=2.2, z=0.0), vector.obj(x=2.2, y=4.4)) assert get_isparallel( vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(x=2.2, y=4.4, z=6.6) ) def test_method_isclose(): @numba.njit def get_isclose(v1, v2): return v1.isclose(v2) assert get_isclose(vector.obj(x=1.1, y=2.2), vector.obj(x=1.1, y=2.2)) assert get_isclose(vector.obj(x=1.1, y=2.2, z=3.3), vector.obj(x=1.1, y=2.2, z=3.3)) assert get_isclose( vector.obj(x=1.1, y=2.2, z=3.3, t=4.4), vector.obj(x=1.1, y=2.2, z=3.3, t=4.4) ) def test_method_rotateZ(): @numba.njit def get_rotateZ(v, angle): return v.rotateZ(angle) out = get_rotateZ(vector.obj(x=1, y=0), 0.1) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(0.9950041652780258) assert out.y == pytest.approx(0.09983341664682815) out = get_rotateZ(vector.obj(rho=1, phi=0), 0.1) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(0.1) out = get_rotateZ(vector.obj(x=1, y=0, z=2.2), 0.1) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.9950041652780258) assert out.y == pytest.approx(0.09983341664682815) assert out.z == pytest.approx(2.2) def test_method_transform2D(): @numba.njit def get_transform2D(v, obj): return v.transform2D(obj) obj = numba.typed.Dict() obj["xx"] = numpy.cos(0.1) obj["xy"] = -numpy.sin(0.1) obj["yx"] = numpy.sin(0.1) obj["yy"] = numpy.cos(0.1) out = get_transform2D(vector.obj(x=1, y=0), obj) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(0.9950041652780258) assert out.y == pytest.approx(0.09983341664682815) out = get_transform2D(vector.obj(rho=1, phi=0), obj) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.rho == pytest.approx(1) assert out.phi == pytest.approx(0.1) out = get_transform2D(vector.obj(x=1, y=0, z=2.2), obj) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.9950041652780258) assert out.y == pytest.approx(0.09983341664682815) assert out.z == pytest.approx(2.2) def test_method_unit(): @numba.njit def get_unit(v): return v.unit() out = get_unit(vector.obj(x=1, y=1)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(1 / numpy.sqrt(2)) assert out.y == pytest.approx(1 / numpy.sqrt(2)) out = get_unit(vector.obj(x=1, y=1, z=1)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(1 / numpy.sqrt(3)) assert out.y == pytest.approx(1 / numpy.sqrt(3)) assert out.z == pytest.approx(1 / numpy.sqrt(3)) out = get_unit(vector.obj(x=1, y=1, z=1, t=1)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(1 / numpy.sqrt(2)) assert out.y == pytest.approx(1 / numpy.sqrt(2)) assert out.z == pytest.approx(1 / numpy.sqrt(2)) assert out.t == pytest.approx(1 / numpy.sqrt(2)) def test_method_scale(): @numba.njit def get_scale(v): return v.scale(2) out = get_scale(vector.obj(x=1.1, y=2.2)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(2.2) assert out.y == pytest.approx(4.4) out = get_scale(vector.obj(x=1.1, y=2.2, z=3.3)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(2.2) assert out.y == pytest.approx(4.4) assert out.z == pytest.approx(6.6) out = get_scale(vector.obj(x=1.1, y=2.2, z=3.3, t=4.4)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(2.2) assert out.y == pytest.approx(4.4) assert out.z == pytest.approx(6.6) assert out.t == pytest.approx(8.8) def test_method_cross(): @numba.njit def get_cross(v1, v2): return v1.cross(v2) out = get_cross(vector.obj(x=0.1, y=0.2, z=0.3), vector.obj(x=0.4, y=0.5, z=0.6)) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert (out.x, out.y, out.z) == pytest.approx((-0.03, 0.06, -0.03)) out = get_cross( vector.obj(x=0.1, y=0.2, z=0.3, t=999), vector.obj(x=0.4, y=0.5, z=0.6) ) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert (out.x, out.y, out.z) == pytest.approx((-0.03, 0.06, -0.03)) out = get_cross( vector.obj(x=0.1, y=0.2, z=0.3), vector.obj(x=0.4, y=0.5, z=0.6, t=999) ) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert (out.x, out.y, out.z) == pytest.approx((-0.03, 0.06, -0.03)) out = get_cross( vector.obj(x=0.1, y=0.2, z=0.3, t=999), vector.obj(x=0.4, y=0.5, z=0.6, t=999) ) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert (out.x, out.y, out.z) == pytest.approx((-0.03, 0.06, -0.03)) def test_method_rotateX(): @numba.njit def get_rotateX(v): return v.rotateX(0.25) out = get_rotateX(vector.obj(x=0.1, y=0.2, z=0.3)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.1) assert out.y == pytest.approx(0.1195612965657721) assert out.z == pytest.approx(0.340154518364098) out = get_rotateX(vector.obj(x=0.1, y=0.2, z=0.3, t=0.4)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(0.1) assert out.y == pytest.approx(0.1195612965657721) assert out.z == pytest.approx(0.340154518364098) assert out.t == pytest.approx(0.4) def test_method_rotateY(): @numba.njit def get_rotateY(v): return v.rotateY(0.25) out = get_rotateY(vector.obj(x=0.1, y=0.2, z=0.3)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.17111242994742137) assert out.y == pytest.approx(0.2) assert out.z == pytest.approx(0.2659333305877411) out = get_rotateY(vector.obj(x=0.1, y=0.2, z=0.3, t=0.4)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(0.17111242994742137) assert out.y == pytest.approx(0.2) assert out.z == pytest.approx(0.2659333305877411) assert out.t == pytest.approx(0.4) def test_method_rotate_axis(): @numba.njit def get_rotate_axis(vec, axis): return vec.rotate_axis(axis, 0.25) axis = vector.obj(x=0.1, y=0.2, z=0.3) vec = vector.obj(x=0.4, y=0.5, z=0.6) out = get_rotate_axis(vec, axis) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) axis = vector.obj(x=0.1, y=0.2, z=0.3) vec = vector.obj(x=0.4, y=0.5, z=0.6, t=999) out = get_rotate_axis(vec, axis) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) assert out.t == pytest.approx(999) axis = vector.obj(x=0.1, y=0.2, z=0.3, t=999) vec = vector.obj(x=0.4, y=0.5, z=0.6) out = get_rotate_axis(vec, axis) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) axis = vector.obj(x=0.1, y=0.2, z=0.3, t=999) vec = vector.obj(x=0.4, y=0.5, z=0.6, t=999) out = get_rotate_axis(vec, axis) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) assert out.t == pytest.approx(999) def test_method_rotate_euler(): @numba.njit def get_rotate_euler(vec, phi, theta, psi): return vec.rotate_euler(phi, theta, psi, order="zxz") vec = vector.obj(x=0.4, y=0.5, z=0.6) out = get_rotate_euler(vec, 0.1, 0.2, 0.3) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.5956646364506655) assert out.y == pytest.approx(0.409927258162962) assert out.z == pytest.approx(0.4971350761081869) def test_method_rotate_quaternion(): @numba.njit def get_rotate_quaternion(vec, u, i, j, k): return vec.rotate_quaternion(u, i, j, k) vec = vector.obj(x=0.5, y=0.6, z=0.7) out = get_rotate_quaternion(vec, 0.1, 0.2, 0.3, 0.4) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.078) assert out.y == pytest.approx(0.18) assert out.z == pytest.approx(0.246) def test_method_transform3D(): @numba.njit def get_transform3D(v, obj): return v.transform3D(obj) obj = numba.typed.Dict() obj["xx"] = numpy.cos(0.1) obj["xy"] = -numpy.sin(0.1) obj["xz"] = 0 obj["yx"] = numpy.sin(0.1) obj["yy"] = numpy.cos(0.1) obj["yz"] = 0 obj["zx"] = 0 obj["zy"] = 0 obj["zz"] = 1 out = get_transform3D(vector.obj(x=1, y=0, z=99), obj) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(0.9950041652780258) assert out.y == pytest.approx(0.09983341664682815) assert out.z == pytest.approx(99) def test_method_boost(): @numba.njit def get_boost_p4(vec, p4): return vec.boost_p4(p4) @numba.njit def get_boost_beta3(vec, beta3): return vec.boost_beta3(beta3) @numba.njit def get_boost(vec, booster): return vec.boost(booster) out = get_boost_p4(vector.obj(x=1, y=2, z=3, t=4), vector.obj(x=5, y=6, z=7, t=15)) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) out = get_boost_beta3( vector.obj(x=1, y=2, z=3, t=4), vector.obj(x=5 / 15, y=6 / 15, z=7 / 15) ) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) out = get_boost(vector.obj(x=1, y=2, z=3, t=4), vector.obj(x=5, y=6, z=7, t=15)) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) out = get_boost( vector.obj(x=1, y=2, z=3, t=4), vector.obj(x=5 / 15, y=6 / 15, z=7 / 15) ) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) def test_method_boostX(): @numba.njit def get_boostX_beta(vec, beta): return vec.boostX(beta=beta) @numba.njit def get_boostX_gamma(vec, gamma): return vec.boostX(gamma=gamma) out = get_boostX_beta(vector.obj(x=3, y=2, z=1, t=4), -0.9428090415820634) assert out.x == pytest.approx(-2.313708498984761) assert out.y == pytest.approx(2) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) out = get_boostX_gamma(vector.obj(x=3, y=2, z=1, t=4), -3) assert out.x == pytest.approx(-2.313708498984761) assert out.y == pytest.approx(2) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) def test_method_to_beta3(): @numba.njit def get_to_beta3(vec): return vec.to_beta3() out = get_to_beta3(vector.obj(x=3, y=4, z=10, t=20)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_method_transform4D(): @numba.njit def get_transform4D(v, obj): return v.transform4D(obj) obj = numba.typed.Dict() obj["xx"] = numpy.cos(0.1) obj["xy"] = -numpy.sin(0.1) obj["xz"] = 0 obj["xt"] = 0 obj["yx"] = numpy.sin(0.1) obj["yy"] = numpy.cos(0.1) obj["yz"] = 0 obj["yt"] = 0 obj["zx"] = 0 obj["zy"] = 0 obj["zz"] = 1 obj["zt"] = 0 obj["tx"] = 0 obj["ty"] = 0 obj["tz"] = 0 obj["tt"] = 1 out = get_transform4D(vector.obj(x=1, y=0, z=99, t=123), obj) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(0.9950041652780258) assert out.y == pytest.approx(0.09983341664682815) assert out.z == pytest.approx(99) assert out.t == pytest.approx(123) def test_method_timespacelight_like(): @numba.njit def get_is_timelike(v): return v.is_timelike() @numba.njit def get_is_spacelike(v): return v.is_spacelike() @numba.njit def get_is_lightlike(v): return v.is_lightlike() assert get_is_timelike(vector.obj(x=3, y=4, z=0, t=10)) assert get_is_lightlike(vector.obj(x=3, y=4, z=0, t=5)) assert get_is_spacelike(vector.obj(x=3, y=4, z=0, t=2)) def test_momentum_alias(): @numba.njit def get_px(v): return v.px assert get_px(vector.obj(px=3, py=4)) == pytest.approx(3) with pytest.raises(numba.TypingError): get_px(vector.obj(x=3, y=4)) def test_operator_abs(): @numba.njit def get_abs(v): return abs(v) assert get_abs(vector.obj(x=3, y=4)) == pytest.approx(5) assert get_abs(vector.obj(x=3, y=4, z=5)) == pytest.approx(numpy.sqrt(50)) assert get_abs(vector.obj(x=3, y=4, z=5, tau=100)) == pytest.approx(100) def test_operator_neg(): @numba.njit def get_neg(v): return -v @numba.njit def get_pos(v): return +v out = get_neg(vector.obj(x=3, y=4)) assert out.x == -3 assert out.y == -4 out = get_pos(vector.obj(x=3, y=4)) assert out.x == 3 assert out.y == 4 def test_operator_bool(): @numba.njit def get_true(v): return bool(v) assert not get_true(vector.obj(x=0, y=0)) assert get_true(vector.obj(x=0, y=0.1)) assert not get_true(vector.obj(x=0, y=0, z=0)) assert get_true(vector.obj(x=0, y=0, z=0.1)) assert not get_true(vector.obj(x=0, y=0, z=0, t=0)) assert get_true(vector.obj(x=0, y=0, z=0.1, t=0)) assert get_true(vector.obj(x=0, y=0, z=0, t=0.1)) assert get_true(vector.obj(x=0, y=0, z=10, t=10)) def test_operator_truth(): @numba.njit def get_true(v): return True if v else False # noqa: SIM210 @numba.njit def get_false(v): return False if v else True # noqa: SIM211 assert not get_true(vector.obj(x=0, y=0)) assert get_true(vector.obj(x=0, y=0.1)) assert not get_true(vector.obj(x=0, y=0, z=0)) assert get_true(vector.obj(x=0, y=0, z=0.1)) assert not get_true(vector.obj(x=0, y=0, z=0, t=0)) assert get_true(vector.obj(x=0, y=0, z=0.1, t=0)) assert get_true(vector.obj(x=0, y=0, z=0, t=0.1)) assert get_true(vector.obj(x=0, y=0, z=10, t=10)) assert get_false(vector.obj(x=0, y=0)) assert not get_false(vector.obj(x=0, y=0.1)) assert get_false(vector.obj(x=0, y=0, z=0)) assert not get_false(vector.obj(x=0, y=0, z=0.1)) assert get_false(vector.obj(x=0, y=0, z=0, t=0)) assert not get_false(vector.obj(x=0, y=0, z=0.1, t=0)) assert not get_false(vector.obj(x=0, y=0, z=0, t=0.1)) assert not get_false(vector.obj(x=0, y=0, z=10, t=10)) def test_operator_eq(): @numba.njit def get_eq(v1, v2): return v1 == v2 @numba.njit def get_ne(v1, v2): return v1 != v2 assert get_eq(vector.obj(x=3, y=4), vector.obj(px=3, py=4)) assert not get_eq(vector.obj(x=3, y=4), vector.obj(x=4, y=3)) with pytest.raises(numba.TypingError): get_eq(vector.obj(x=3, y=4), vector.obj(x=3, y=4, z=99)) assert not get_ne(vector.obj(x=3, y=4), vector.obj(px=3, py=4)) assert get_ne(vector.obj(x=3, y=4), vector.obj(x=4, y=3)) with pytest.raises(numba.TypingError): get_ne(vector.obj(x=3, y=4), vector.obj(x=3, y=4, z=99)) def test_operator_add(): @numba.njit def get_add(v1, v2): return v1 + v2 @numba.njit def get_subtract(v1, v2): return v1 - v2 out = get_add(vector.obj(x=1, y=2), vector.obj(px=3, py=4)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(4) assert out.y == pytest.approx(6) out = get_add(vector.obj(x=1, y=2), vector.obj(x=3, y=4, z=5)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(4) assert out.y == pytest.approx(6) out = get_add(vector.obj(x=1, y=2, z=0), vector.obj(x=3, y=4, z=5)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(4) assert out.y == pytest.approx(6) assert out.z == pytest.approx(5) out = get_subtract(vector.obj(x=1, y=2), vector.obj(px=3, py=4)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-2) def test_operator_mul(): @numba.njit def get_mul(a, b): return a * b @numba.njit def get_div(a, b): return a / b out = get_mul(vector.obj(x=1, y=2), 2) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(4) out = get_mul(2, vector.obj(x=1, y=2)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(2) assert out.y == pytest.approx(4) out = get_div(vector.obj(x=1, y=2), 2) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(0.5) assert out.y == pytest.approx(1) with pytest.raises(numba.TypingError): get_div(2, vector.obj(x=1, y=2)) def test_operator_pow(): @numba.njit def get_pow(a, b): return a**b @numba.njit def get_square(a): return a**2 assert get_pow(vector.obj(x=1, y=2), 4) == pytest.approx(25) assert get_square(vector.obj(x=1, y=2)) == pytest.approx(5) def test_operator_matmul(): @numba.njit def get_matmul(v1, v2): return v1 @ v2 assert get_matmul(vector.obj(x=1, y=2), vector.obj(x=3, y=4)) == pytest.approx(11) def test_numpy_functions(): @numba.njit def get_absolute(v): return numpy.absolute(v) assert get_absolute(vector.obj(x=3, y=4)) == pytest.approx(5) @numba.njit def get_add(v1, v2): return numpy.add(v1, v2) out = get_add(vector.obj(x=3, y=4), vector.obj(x=10, y=20)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(13) assert out.y == pytest.approx(24) def test_scale2D_3D_4D(): # @numba.njit def get_scale2D(v, factor): return v.scale2D(factor) # @numba.njit def get_neg2D(v): return v.neg2D out = get_scale2D(vector.obj(x=1, y=2), 10) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(10) assert out.y == pytest.approx(20) out = get_scale2D(vector.obj(x=1, y=2, z=3), 10) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(10) assert out.y == pytest.approx(20) assert out.z == pytest.approx(3) out = get_scale2D(vector.obj(x=1, y=2, z=3, t=4), 10) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(10) assert out.y == pytest.approx(20) assert out.z == pytest.approx(3) assert out.t == pytest.approx(4) out = get_neg2D(vector.obj(x=1, y=2)) assert isinstance(out, vector.backends.object.VectorObject2D) assert out.x == pytest.approx(-1) assert out.y == pytest.approx(-2) out = get_neg2D(vector.obj(x=1, y=2, z=3)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(-1) assert out.y == pytest.approx(-2) assert out.z == pytest.approx(3) out = get_neg2D(vector.obj(x=1, y=2, z=3, t=4)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(-1) assert out.y == pytest.approx(-2) assert out.z == pytest.approx(3) assert out.t == pytest.approx(4) # @numba.njit def get_scale3D(v, factor): return v.scale3D(factor) # @numba.njit def get_neg3D(v): return v.neg3D out = get_scale3D(vector.obj(x=1, y=2, z=3), 10) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(10) assert out.y == pytest.approx(20) assert out.z == pytest.approx(30) out = get_scale3D(vector.obj(x=1, y=2, z=3, t=4), 10) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(10) assert out.y == pytest.approx(20) assert out.z == pytest.approx(30) assert out.t == pytest.approx(4) out = get_neg3D(vector.obj(x=1, y=2, z=3)) assert isinstance(out, vector.backends.object.VectorObject3D) assert out.x == pytest.approx(-1) assert out.y == pytest.approx(-2) assert out.z == pytest.approx(-3) out = get_neg3D(vector.obj(x=1, y=2, z=3, t=4)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(-1) assert out.y == pytest.approx(-2) assert out.z == pytest.approx(-3) assert out.t == pytest.approx(4) # @numba.njit def get_scale4D(v, factor): return v.scale4D(factor) # @numba.njit def get_neg4D(v): return v.neg4D out = get_scale4D(vector.obj(x=1, y=2, z=3, t=4), 10) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(10) assert out.y == pytest.approx(20) assert out.z == pytest.approx(30) assert out.t == pytest.approx(40) out = get_neg4D(vector.obj(x=1, y=2, z=3, t=4)) assert isinstance(out, vector.backends.object.VectorObject4D) assert out.x == pytest.approx(-1) assert out.y == pytest.approx(-2) assert out.z == pytest.approx(-3) assert out.t == pytest.approx(-4) def test_method_boostCM_of(): # @numba.njit def get_boostCM_of_p4(vec, p4): return vec.boostCM_of_p4(p4) # @numba.njit def get_boostCM_of_beta3(vec, beta3): return vec.boostCM_of_beta3(beta3) # @numba.njit def get_boostCM_of(vec, booster): return vec.boostCM_of(booster) vec = vector.obj(x=1, y=2, z=3, t=4) out = get_boostCM_of_p4(vec, vec) assert out.x == pytest.approx(0) assert out.y == pytest.approx(0) assert out.z == pytest.approx(0) assert out.t == pytest.approx(vec.tau) out = get_boostCM_of_beta3(vec, vec.to_beta3()) assert out.x == pytest.approx(0) assert out.y == pytest.approx(0) assert out.z == pytest.approx(0) assert out.t == pytest.approx(vec.tau) out = get_boostCM_of(vec, vec) assert out.x == pytest.approx(0) assert out.y == pytest.approx(0) assert out.z == pytest.approx(0) assert out.t == pytest.approx(vec.tau) out = get_boostCM_of(vec, vec.to_beta3()) assert out.x == pytest.approx(0) assert out.y == pytest.approx(0) assert out.z == pytest.approx(0) assert out.t == pytest.approx(vec.tau) vector-1.6.3/tests/backends/test_numpy.py000066400000000000000000000347451503546127100205410ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pickle import numpy import pytest import vector.backends.numpy def test_type_checks(): with pytest.raises(TypeError): vector.backends.numpy.VectorNumpy2D( (0, 0), dtype=[("x", numpy.float64), ("y", numpy.timedelta64)] ) with pytest.raises(TypeError): vector.backends.numpy.VectorNumpy2D( (0, 0), dtype=[("x", numpy.complex64), ("y", numpy.float64)] ) with pytest.raises(TypeError): vector.backends.numpy.VectorNumpy2D( [(0, 0), (0, 1), (3, 4)], dtype=[("x", numpy.complex64), ("y", numpy.float64)], ) with pytest.raises(TypeError): vector.backends.numpy.VectorNumpy3D( [([0], 0, 0), ([0], 1, 2), ([3], 4, 5)], dtype=[("x", list), ("y", numpy.float64), ("z", numpy.float64)], ) with pytest.raises(TypeError): vector.backends.numpy.VectorNumpy4D( [(0, 0, 0, 0), (0, 1, 2, 3), (3, 4, 5, 6)], [ ("x", complex), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) with pytest.raises(TypeError): vector.backends.numpy.MomentumNumpy2D( (0, 0), dtype=[("x", numpy.float64), ("y", numpy.timedelta64)] ) with pytest.raises(TypeError): vector.backends.numpy.MomentumNumpy2D( (0, 0), dtype=[("x", numpy.complex64), ("y", numpy.float64)] ) with pytest.raises(TypeError): vector.backends.numpy.MomentumNumpy3D( [(0, 0), (0, 1), (3, 4)], dtype=[("x", numpy.complex64), ("y", numpy.float64)], ) with pytest.raises(TypeError): vector.backends.numpy.MomentumNumpy4D( [([0], 0, 0), ([0], 1, 2), ([3], 4, 5)], dtype=[("x", list), ("y", numpy.float64), ("z", numpy.float64)], ) with pytest.raises(TypeError): vector.backends.numpy.MomentumNumpy4D( [(0, 0, 0, 0), (0, 1, 2, 3), (3, 4, 5, 6)], [ ("x", complex), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) def test_xy(): array = vector.backends.numpy.VectorNumpy2D( [(0, 0), (0, 1), (3, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64)] ) assert numpy.allclose(array.x, [0, 0, 3]) assert numpy.allclose(array.y, [0, 1, 4]) assert numpy.allclose(array.rho, [0, 1, 5]) assert numpy.allclose(array.phi, [0, math.atan2(1, 0), math.atan2(4, 3)]) def test_rhophi(): array = vector.backends.numpy.VectorNumpy2D( [(0, 10), (1, math.atan2(1, 0)), (5, math.atan2(4, 3))], dtype=[("rho", numpy.float64), ("phi", numpy.float64)], ) assert numpy.allclose(array.x, [0, 0, 3]) assert numpy.allclose(array.y, [0, 1, 4]) assert numpy.allclose(array.rho, [0, 1, 5]) assert numpy.allclose(array.phi, [10, math.atan2(1, 0), math.atan2(4, 3)]) def test_pickle_vector_numpy_2d(): array = vector.backends.numpy.VectorNumpy2D( [(0, 0), (0, 1), (3, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64)] ) array_pickled = pickle.dumps(array) array_new = pickle.loads(array_pickled) assert numpy.allclose(array_new.x, array.x) assert numpy.allclose(array_new.y, array.y) def test_pickle_momentum_numpy_2d(): array = vector.backends.numpy.MomentumNumpy2D( [(0, 0), (0, 1), (3, 4)], dtype=[("rho", numpy.float64), ("phi", numpy.float64)] ) array_pickled = pickle.dumps(array) array_new = pickle.loads(array_pickled) assert numpy.allclose(array_new.rho, array.rho) assert numpy.allclose(array_new.phi, array.phi) def test_pickle_vector_numpy_3d(): array = vector.backends.numpy.VectorNumpy3D( [(0, 0, 0), (0, 1, 1), (3, 4, 5)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) array_pickled = pickle.dumps(array) array_new = pickle.loads(array_pickled) assert numpy.allclose(array_new.x, array.x) assert numpy.allclose(array_new.y, array.y) assert numpy.allclose(array_new.z, array.z) def test_pickle_momentum_numpy_3d(): array = vector.backends.numpy.MomentumNumpy3D( [(0, 0, 0), (0, 1, 1), (3, 4, 5)], dtype=[ ("rho", numpy.float64), ("phi", numpy.float64), ("theta", numpy.float64), ], ) array_pickled = pickle.dumps(array) array_new = pickle.loads(array_pickled) assert numpy.allclose(array_new.rho, array.rho) assert numpy.allclose(array_new.phi, array.phi) assert numpy.allclose(array_new.theta, array.theta) def test_pickle_vector_numpy_4d(): array = vector.backends.numpy.VectorNumpy4D( [(0, 0, 0, 0), (0, 1, 1, 1), (3, 4, 5, 6)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) array_pickled = pickle.dumps(array) array_new = pickle.loads(array_pickled) assert numpy.allclose(array_new.x, array.x) assert numpy.allclose(array_new.y, array.y) assert numpy.allclose(array_new.z, array.z) assert numpy.allclose(array_new.t, array.t) def test_pickle_momentum_numpy_4d(): array = vector.backends.numpy.MomentumNumpy4D( [(0, 0, 0, 0), (0, 1, 1, 1), (3, 4, 5, 6)], dtype=[ ("rho", numpy.float64), ("phi", numpy.float64), ("theta", numpy.float64), ("tau", numpy.float64), ], ) array_pickled = pickle.dumps(array) array_new = pickle.loads(array_pickled) assert numpy.allclose(array_new.rho, array.rho) assert numpy.allclose(array_new.phi, array.phi) assert numpy.allclose(array_new.theta, array.theta) assert numpy.allclose(array_new.tau, array.tau) def test_sum_2d(): v = vector.VectorNumpy2D( [[(1, 0.1), (4, 0.2), (0, 0)], [(1, 0.3), (4, 0.4), (1, 0.1)]], dtype=[("rho", numpy.float64), ("phi", numpy.float64)], ) assert numpy.sum(v, axis=0, keepdims=False).allclose( vector.VectorNumpy2D( [ [ (1.950340654403632, 0.3953536233081677), (7.604510287376507, 2.3523506924148467), (0.9950041652780258, 0.09983341664682815), ] ], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) ) assert numpy.sum(v, axis=0, keepdims=True).allclose( vector.VectorNumpy2D( [ (1.950340654403632, 0.3953536233081677), (7.604510287376507, 2.3523506924148467), (0.9950041652780258, 0.09983341664682815), ], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) ) assert numpy.sum(v, axis=0).allclose( vector.VectorNumpy2D( [ (1.950340654403632, 0.3953536233081677), (7.604510287376507, 2.3523506924148467), (0.9950041652780258, 0.09983341664682815), ], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) ) assert numpy.sum(v, axis=0, keepdims=False).shape == (3,) assert numpy.sum(v, axis=0, keepdims=True).shape == (1, 3) assert numpy.sum(v, axis=0).shape == (3,) assert numpy.sum(v, axis=1, keepdims=True).allclose( vector.VectorNumpy2D( [[(4.91527048, 0.89451074)], [(5.63458463, 1.95302699)]], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) ) assert numpy.sum(v, axis=1, keepdims=False).allclose( vector.VectorNumpy2D( [(4.91527048, 0.89451074), (5.63458463, 1.95302699)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) ) assert numpy.sum(v, axis=1).allclose( vector.VectorNumpy2D( [(4.91527048, 0.89451074), (5.63458463, 1.95302699)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) ) assert numpy.sum(v, axis=1, keepdims=False).shape == (2,) assert numpy.sum(v, axis=1, keepdims=True).shape == (2, 1) assert numpy.sum(v, axis=1).shape == (2,) def test_sum_3d(): v = vector.VectorNumpy3D( [ [(1, 2, 0.1), (4, 5, 0.2), (0, 0, 0.04)], [(1, 2, 0.6), (4, 5, 1.3), (1, 1, 1.9)], ], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("theta", numpy.float64), ], ) assert numpy.sum(v, axis=0, keepdims=True).allclose( vector.VectorNumpy3D( [ [ (2.0, 4.0, 25.55454594), (8.0, 10.0, 33.36521103), (1.0, 1.0, -0.48314535), ] ], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) ) assert numpy.sum(v, axis=0, keepdims=False).allclose( vector.VectorNumpy3D( [ (2.0, 4.0, 25.55454594), (8.0, 10.0, 33.36521103), (1.0, 1.0, -0.48314535), ], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) ) assert numpy.sum(v, axis=0).allclose( vector.VectorNumpy3D( [ (2.0, 4.0, 25.55454594), (8.0, 10.0, 33.36521103), (1.0, 1.0, -0.48314535), ], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) ) assert numpy.sum(v, axis=0, keepdims=False).shape == (3,) assert numpy.sum(v, axis=0, keepdims=True).shape == (1, 3) assert numpy.sum(v, axis=0).shape == (3,) assert numpy.sum(v, axis=1, keepdims=True).allclose( vector.VectorNumpy3D( [[(5.0, 7.0, 53.87369799)], [(6.0, 8.0, 4.56291362)]], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) ) assert numpy.sum(v, axis=1, keepdims=False).allclose( vector.VectorNumpy3D( [(5.0, 7.0, 53.87369799), (6.0, 8.0, 4.56291362)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) ) assert numpy.sum(v, axis=1).allclose( vector.VectorNumpy3D( [(5.0, 7.0, 53.87369799), (6.0, 8.0, 4.56291362)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) ) assert numpy.sum(v, axis=1, keepdims=False).shape == (2,) assert numpy.sum(v, axis=1, keepdims=True).shape == (2, 1) assert numpy.sum(v, axis=1).shape == (2,) def test_sum_4d(): v = vector.VectorNumpy4D( [ [(1, 2, 3, 4), (4, 5, 6, 2), (0, 0, 0, 3)], [(1, 2, 3, 8), (4, 5, 6, 0), (1, 1, 1, 0)], ], dtype=[ ("x", numpy.int64), ("y", numpy.int64), ("z", numpy.int64), ("t", numpy.int64), ], ) assert numpy.sum(v, axis=0, keepdims=True).tolist() == [ [(2, 4, 6, 12), (8, 10, 12, 2), (1, 1, 1, 3)] ] assert numpy.sum(v, axis=0, keepdims=False).tolist() == [ (2, 4, 6, 12), (8, 10, 12, 2), (1, 1, 1, 3), ] assert numpy.sum(v, axis=0).tolist() == [ (2, 4, 6, 12), (8, 10, 12, 2), (1, 1, 1, 3), ] assert numpy.sum(v, axis=0, keepdims=False).shape == (3,) assert numpy.sum(v, axis=0, keepdims=True).shape == (1, 3) assert numpy.sum(v, axis=0).shape == (3,) assert numpy.sum(v, axis=1, keepdims=True).tolist() == [ [(5, 7, 9, 9)], [(6, 8, 10, 8)], ] assert numpy.sum(v, axis=1, keepdims=False).tolist() == [ (5, 7, 9, 9), (6, 8, 10, 8), ] assert numpy.sum(v, axis=1).tolist() == [ (5, 7, 9, 9), (6, 8, 10, 8), ] assert numpy.sum(v, axis=1, keepdims=False).shape == (2,) assert numpy.sum(v, axis=1, keepdims=True).shape == (2, 1) assert numpy.sum(v, axis=1).shape == (2,) def test_count_nonzero_2d(): v = vector.VectorNumpy2D( [ [(1, 0.1), (4, 0.2), (0, 0)], [(1, 0.3), (4, 0.4), (1, 0.1)], ], dtype=[("rho", numpy.float64), ("phi", numpy.float64)], ) assert numpy.count_nonzero(v, axis=1).tolist() == [2, 3] assert numpy.count_nonzero(v, axis=1, keepdims=True).tolist() == [[2], [3]] assert numpy.count_nonzero(v, axis=0).tolist() == [2, 2, 1] assert numpy.count_nonzero(v, axis=0, keepdims=True).tolist() == [[2, 2, 1]] def test_count_nonzero_3d(): v = vector.VectorNumpy3D( [ [(1, 2, 0.1), (4, 5, 0.2), (0, 0, 0)], [(1, 2, 0.6), (4, 5, 1.3), (1, 1, 1.9)], ], dtype=[("x", numpy.float64), ("y", numpy.float64), ("theta", numpy.float64)], ) assert numpy.count_nonzero(v, axis=1).tolist() == [2, 3] assert numpy.count_nonzero(v, axis=1, keepdims=True).tolist() == [[2], [3]] assert numpy.count_nonzero(v, axis=0).tolist() == [2, 2, 1] assert numpy.count_nonzero(v, axis=0, keepdims=True).tolist() == [[2, 2, 1]] def test_count_nonzero_4d(): v = vector.VectorNumpy4D( [ [(1, 2, 3, 4), (4, 5, 6, 2), (0, 0, 0, 3)], [(1, 2, 3, 8), (4, 5, 6, 0), (1, 1, 1, 0)], ], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert numpy.count_nonzero(v, axis=1).tolist() == [3, 3] assert numpy.count_nonzero(v, axis=1, keepdims=True).tolist() == [[3], [3]] assert numpy.count_nonzero(v, axis=0).tolist() == [2, 2, 2] assert numpy.count_nonzero(v, axis=0, keepdims=True).tolist() == [[2, 2, 2]] v2 = vector.VectorNumpy4D( [ [(1, 2, 3, 1), (4, 5, 6, 2), (0, 0, 0, 2)], [(1, 2, 3, 0), (4, 5, 6, 1), (0, 0, 0, 0)], ], dtype=[ ("x", numpy.int64), ("y", numpy.int64), ("z", numpy.int64), ("t", numpy.int64), ], ) assert numpy.count_nonzero(v2, axis=1).tolist() == [3, 2] assert numpy.count_nonzero(v2, axis=1, keepdims=True).tolist() == [[3], [2]] assert numpy.count_nonzero(v2, axis=0).tolist() == [2, 2, 1] assert numpy.count_nonzero(v2, axis=0, keepdims=True).tolist() == [[2, 2, 1]] vector-1.6.3/tests/backends/test_object.py000066400000000000000000000276111503546127100206310ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector def test_constructors_2D(): vec = vector.VectorObject2D(x=1, y=2) assert vec.x == 1 assert vec.y == 2 vec = vector.VectorObject2D(rho=1, phi=2) assert vec.rho == 1 assert vec.phi == 2 vec = vector.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2) ) assert vec.x == 1 assert vec.y == 2 with pytest.raises(TypeError): vector.VectorObject2D(rho=1, wow=2) with pytest.raises(TypeError): vector.VectorObject2D(rho=complex(1, 2), wow=2) with pytest.raises(TypeError): vector.VectorObject2D() vec = vector.MomentumObject2D(px=1, py=2) assert vec.px == 1 assert vec.py == 2 assert vec.x == 1 assert vec.y == 2 vec = vector.MomentumObject2D(x=1, py=2) assert vec.px == 1 assert vec.py == 2 assert vec.x == 1 assert vec.y == 2 vec = vector.MomentumObject2D(pt=1, phi=2) assert vec.pt == 1 assert vec.phi == 2 assert vec.rho == 1 assert vec.phi == 2 vec = vector.MomentumObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2) ) assert vec.px == 1 assert vec.py == 2 assert vec.x == 1 assert vec.y == 2 with pytest.raises(TypeError): vector.MomentumObject2D(rho=1, wow=2) with pytest.raises(TypeError): vector.MomentumObject2D(rho=False, wow=2) with pytest.raises(TypeError): vector.MomentumObject2D() for coord in ( "xy", "rhophi", ): with pytest.raises(TypeError): getattr(vector.VectorObject2D, "from_" + coord)(complex(1, 2), 2) with pytest.raises(TypeError): getattr(vector.MomentumObject2D, "from_" + coord)(complex(1, 2), 2) def test_constructors_3D(): vec = vector.VectorObject3D(x=1, y=2, z=3) assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 vec = vector.VectorObject3D(x=1, y=2, eta=3) assert vec.x == 1 assert vec.y == 2 assert vec.eta == 3 vec = vector.VectorObject3D(x=1, y=2, theta=3) assert vec.x == 1 assert vec.y == 2 assert vec.theta == 3 vec = vector.VectorObject3D(rho=1, phi=2, z=3) assert vec.rho == 1 assert vec.phi == 2 assert vec.z == 3 vec = vector.VectorObject3D(rho=1, phi=2, eta=3) assert vec.rho == 1 assert vec.phi == 2 assert vec.eta == 3 vec = vector.VectorObject3D(rho=1, phi=2, theta=3) assert vec.rho == 1 assert vec.phi == 2 assert vec.theta == 3 vec = vector.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), ) assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 with pytest.raises(TypeError): vector.VectorObject3D(rho=1, wow=2, z=3) with pytest.raises(TypeError): vector.VectorObject3D() with pytest.raises(TypeError): vector.VectorObject3D(x=complex(1, 2), y=2, z=3) vec = vector.MomentumObject3D(px=1, py=2, pz=3) assert vec.px == 1 assert vec.py == 2 assert vec.pz == 3 assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 vec = vector.MomentumObject3D(px=1, py=2, eta=3) assert vec.px == 1 assert vec.py == 2 assert vec.eta == 3 assert vec.x == 1 assert vec.y == 2 vec = vector.MomentumObject3D(px=1, py=2, theta=3) assert vec.px == 1 assert vec.py == 2 assert vec.theta == 3 assert vec.x == 1 assert vec.y == 2 vec = vector.MomentumObject3D(pt=1, phi=2, pz=3) assert vec.pt == 1 assert vec.phi == 2 assert vec.pz == 3 assert vec.rho == 1 assert vec.z == 3 vec = vector.MomentumObject3D(pt=1, phi=2, eta=3) assert vec.pt == 1 assert vec.phi == 2 assert vec.eta == 3 assert vec.rho == 1 vec = vector.MomentumObject3D(pt=1, phi=2, theta=3) assert vec.pt == 1 assert vec.phi == 2 assert vec.theta == 3 assert vec.rho == 1 vec = vector.MomentumObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), ) assert vec.px == 1 assert vec.py == 2 assert vec.pz == 3 assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 with pytest.raises(TypeError): vector.MomentumObject3D(rho=1, wow=2, pz=3) with pytest.raises(TypeError): vector.MomentumObject3D() with pytest.raises(TypeError): vector.MomentumObject3D(x=complex(1, 2), y=2, z=3) for coord in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): with pytest.raises(TypeError): getattr(vector.VectorObject3D, "from_" + coord)(complex(1, 2), 2, 3) with pytest.raises(TypeError): getattr(vector.MomentumObject3D, "from_" + coord)(complex(1, 2), 2, 3) def test_constructors_4D(): vec = vector.VectorObject4D(x=1, y=2, z=3, t=4) assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 assert vec.t == 4 vec = vector.VectorObject4D(x=1, y=2, eta=3, t=4) assert vec.x == 1 assert vec.y == 2 assert vec.eta == 3 assert vec.t == 4 vec = vector.VectorObject4D(x=1, y=2, theta=3, t=4) assert vec.x == 1 assert vec.y == 2 assert vec.theta == 3 assert vec.t == 4 vec = vector.VectorObject4D(rho=1, phi=2, z=3, t=4) assert vec.rho == 1 assert vec.phi == 2 assert vec.z == 3 assert vec.t == 4 vec = vector.VectorObject4D(rho=1, phi=2, eta=3, t=4) assert vec.rho == 1 assert vec.phi == 2 assert vec.eta == 3 assert vec.t == 4 vec = vector.VectorObject4D(rho=1, phi=2, theta=3, t=4) assert vec.rho == 1 assert vec.phi == 2 assert vec.theta == 3 assert vec.t == 4 vec = vector.VectorObject4D(x=1, y=2, z=3, tau=4) assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 assert vec.tau == 4 vec = vector.VectorObject4D(x=1, y=2, eta=3, tau=4) assert vec.x == 1 assert vec.y == 2 assert vec.eta == 3 assert vec.tau == 4 vec = vector.VectorObject4D(x=1, y=2, theta=3, tau=4) assert vec.x == 1 assert vec.y == 2 assert vec.theta == 3 assert vec.tau == 4 vec = vector.VectorObject4D(rho=1, phi=2, z=3, tau=4) assert vec.rho == 1 assert vec.phi == 2 assert vec.z == 3 assert vec.tau == 4 vec = vector.VectorObject4D(rho=1, phi=2, eta=3, tau=4) assert vec.rho == 1 assert vec.phi == 2 assert vec.eta == 3 assert vec.tau == 4 vec = vector.VectorObject4D(rho=1, phi=2, theta=3, tau=4) assert vec.rho == 1 assert vec.phi == 2 assert vec.theta == 3 assert vec.tau == 4 vec = vector.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 assert vec.t == 4 with pytest.raises(TypeError): vector.VectorObject4D(rho=1, wow=2, z=3, t=4) with pytest.raises(TypeError): vector.VectorObject4D() with pytest.raises(TypeError): vector.VectorObject4D(x=complex(1, 2), y=2, z=3, t=4) vec = vector.MomentumObject4D(px=1, py=2, pz=3, E=4) assert vec.px == 1 assert vec.py == 2 assert vec.pz == 3 assert vec.E == 4 vec = vector.MomentumObject4D(px=1, py=2, eta=3, E=4) assert vec.px == 1 assert vec.py == 2 assert vec.eta == 3 assert vec.E == 4 vec = vector.MomentumObject4D(px=1, py=2, theta=3, E=4) assert vec.px == 1 assert vec.py == 2 assert vec.theta == 3 assert vec.E == 4 vec = vector.MomentumObject4D(pt=1, phi=2, pz=3, E=4) assert vec.pt == 1 assert vec.phi == 2 assert vec.pz == 3 assert vec.E == 4 vec = vector.MomentumObject4D(pt=1, phi=2, eta=3, E=4) assert vec.pt == 1 assert vec.phi == 2 assert vec.eta == 3 assert vec.E == 4 vec = vector.MomentumObject4D(pt=1, phi=2, theta=3, E=4) assert vec.pt == 1 assert vec.phi == 2 assert vec.theta == 3 assert vec.E == 4 vec = vector.MomentumObject4D(px=1, py=2, pz=3, M=4) assert vec.px == 1 assert vec.py == 2 assert vec.pz == 3 assert vec.M == 4 vec = vector.MomentumObject4D(px=1, py=2, eta=3, tau=4) assert vec.px == 1 assert vec.py == 2 assert vec.eta == 3 assert vec.M == 4 vec = vector.MomentumObject4D(px=1, py=2, theta=3, tau=4) assert vec.px == 1 assert vec.py == 2 assert vec.theta == 3 assert vec.M == 4 vec = vector.MomentumObject4D(pt=1, phi=2, pz=3, tau=4) assert vec.pt == 1 assert vec.phi == 2 assert vec.pz == 3 assert vec.M == 4 vec = vector.MomentumObject4D(pt=1, phi=2, eta=3, tau=4) assert vec.pt == 1 assert vec.phi == 2 assert vec.eta == 3 assert vec.M == 4 vec = vector.MomentumObject4D(pt=1, phi=2, theta=3, tau=4) assert vec.pt == 1 assert vec.phi == 2 assert vec.theta == 3 assert vec.M == 4 vec = vector.MomentumObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) assert vec.px == 1 assert vec.py == 2 assert vec.pz == 3 assert vec.x == 1 assert vec.y == 2 assert vec.z == 3 assert vec.t == 4 with pytest.raises(TypeError): vector.MomentumObject4D(rho=1, wow=2, pz=3, t=4) with pytest.raises(TypeError): vector.MomentumObject4D() with pytest.raises(TypeError): vector.MomentumObject4D(x=complex(1, 2), y=2, z=3, t=4) for coord in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): with pytest.raises(TypeError): getattr(vector.VectorObject4D, "from_" + coord)(complex(1, 2), 2, 3, 4) with pytest.raises(TypeError): getattr(vector.MomentumObject4D, "from_" + coord)(complex(1, 2), 2, 3, 4) def test_array_casting(): obj = vector.obj(x=1, y=1) assert isinstance(obj, vector.VectorObject2D) assert isinstance(numpy.asanyarray(obj), vector.VectorNumpy2D) assert numpy.asanyarray(obj).shape == () obj = vector.obj(px=1, py=1) assert isinstance(obj, vector.MomentumObject2D) assert isinstance(numpy.asanyarray(obj), vector.MomentumNumpy2D) assert numpy.asanyarray(obj).shape == () obj = vector.obj(x=1, y=1, z=1) assert isinstance(obj, vector.VectorObject3D) assert isinstance(numpy.asanyarray(obj), vector.VectorNumpy3D) assert numpy.asanyarray(obj).shape == () obj = vector.obj(px=1, py=1, pz=1) assert isinstance(obj, vector.MomentumObject3D) assert isinstance(numpy.asanyarray(obj), vector.MomentumNumpy3D) assert numpy.asanyarray(obj).shape == () obj = vector.obj(x=1, y=1, z=1, t=1) assert isinstance(obj, vector.VectorObject4D) assert isinstance(numpy.asanyarray(obj), vector.VectorNumpy4D) assert numpy.asanyarray(obj).shape == () obj = vector.obj(px=1, py=1, pz=1, E=1) assert isinstance(obj, vector.MomentumObject4D) assert isinstance(numpy.asanyarray(obj), vector.MomentumNumpy4D) assert numpy.asanyarray(obj).shape == () with pytest.raises(TypeError): vector.obj(x=1, y=[1, 2]) with pytest.raises(TypeError): vector.obj(x=1, y=complex(1, 2)) with pytest.raises(TypeError): vector.obj(x=1, y=False) vector-1.6.3/tests/backends/test_operators.py000066400000000000000000000133521503546127100213760ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector v1 = vector.obj(x=1, y=5) a1 = vector.array({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]}) v2 = vector.obj(x=10, y=20) a2 = vector.array({"x": [10, 100, 1000, 10000], "y": [20, 200, 2000, 20000]}) def test_eq(): assert v1 == v1 assert not v1 == v2 assert (a1 == a1).all() assert not (a1 == a2).any() assert (v1 == a1).any() assert not (v1 == a1).all() assert (a1 == v1).any() assert not (a1 == v1).all() with pytest.raises(TypeError): v1.equal(v2.to_Vector3D()) with pytest.raises(TypeError): a1.equal(a2.to_Vector3D()) def test_ne(): assert not v1 != v1 assert v1 != v2 assert not (a1 != a1).any() assert (a1 != a2).all() assert (v1 != a1).any() assert not (v1 != a1).all() assert (a1 != v1).any() assert not (a1 != v1).all() with pytest.raises(TypeError): v1.not_equal(v2.to_Vector3D()) with pytest.raises(TypeError): a1.not_equal(a2.to_Vector3D()) def test_abs(): assert abs(v1) == pytest.approx(numpy.sqrt(1**2 + 5**2)) assert numpy.allclose( abs(a1), numpy.sqrt(numpy.array([1, 2, 3, 4]) ** 2 + numpy.array([5, 6, 7, 8]) ** 2), ) def test_add(): assert v1 + v2 == vector.obj(x=11, y=25) assert v1 + v2.to_Vector3D().like(v1) == vector.obj(x=11, y=25) assert numpy.allclose( a1 + a2, vector.array({"x": [11, 102, 1003, 10004], "y": [25, 206, 2007, 20008]}), ) assert numpy.allclose( a1 + a2.to_Vector3D().like(a1), vector.array({"x": [11, 102, 1003, 10004], "y": [25, 206, 2007, 20008]}), ) assert numpy.allclose( v1 + a2, vector.array({"x": [11, 101, 1001, 10001], "y": [25, 205, 2005, 20005]}), ) assert numpy.allclose( a2 + v1, vector.array({"x": [11, 101, 1001, 10001], "y": [25, 205, 2005, 20005]}), ) with pytest.raises(TypeError): v1 + 5 with pytest.raises(TypeError): 5 + v1 with pytest.raises(TypeError): v1 + v2.to_Vector3D() with pytest.raises(TypeError): a1 + a2.to_Vector3D() def test_sub(): assert v1 - v2 == vector.obj(x=-9, y=-15) assert numpy.allclose( a1 - a2, vector.array({"x": [-9, -98, -997, -9996], "y": [-15, -194, -1993, -19992]}), ) assert numpy.allclose( v1 - a2, vector.array({"x": [-9, -99, -999, -9999], "y": [-15, -195, -1995, -19995]}), ) assert numpy.allclose( a2 - v1, vector.array({"x": [9, 99, 999, 9999], "y": [15, 195, 1995, 19995]}), ) with pytest.raises(TypeError): v1 - 5 with pytest.raises(TypeError): 5 - v1 with pytest.raises(TypeError): v1 - v2.to_Vector3D() with pytest.raises(TypeError): a1 - a2.to_Vector3D() def test_mul(): assert v1 * 10 == vector.obj(x=10, y=50) assert 10 * v1 == vector.obj(x=10, y=50) assert numpy.allclose( a1 * 10, vector.array({"x": [10, 20, 30, 40], "y": [50, 60, 70, 80]}) ) assert numpy.allclose( 10 * a1, vector.array({"x": [10, 20, 30, 40], "y": [50, 60, 70, 80]}) ) with pytest.raises(TypeError): v1 * v2 with pytest.raises(TypeError): a1 * a2 def test_neg(): assert -v1 == vector.obj(x=-1, y=-5) assert numpy.allclose( -a1, vector.array({"x": [-1, -2, -3, -4], "y": [-5, -6, -7, -8]}) ) def test_pos(): assert +v1 == vector.obj(x=1, y=5) assert numpy.allclose(+a1, vector.array({"x": [1, 2, 3, 4], "y": [5, 6, 7, 8]})) def test_truediv(): assert v1 / 10 == vector.obj(x=0.1, y=0.5) with pytest.raises(TypeError): 10 / v1 assert numpy.allclose( a1 / 10, vector.array({"x": [0.1, 0.2, 0.3, 0.4], "y": [0.5, 0.6, 0.7, 0.8]}) ) with pytest.raises(TypeError): 10 / a1 with pytest.raises(TypeError): v1 / v2 with pytest.raises(TypeError): a1 / a2 def test_pow(): assert v1**2 == pytest.approx(1**2 + 5**2) with pytest.raises(TypeError): 2**v1 assert numpy.allclose( a1**2, numpy.array([1**2 + 5**2, 2**2 + 6**2, 3**2 + 7**2, 4**2 + 8**2]), ) with pytest.raises(TypeError): 2**a1 with pytest.raises(TypeError): v1**v2 with pytest.raises(TypeError): a1**a2 def test_matmul(): assert v1 @ v2 == pytest.approx(1 * 10 + 5 * 20) assert v2 @ v1 == pytest.approx(1 * 10 + 5 * 20) assert numpy.allclose( a1 @ a2, numpy.array( [ 1 * 10 + 5 * 20, 2 * 100 + 6 * 200, 3 * 1000 + 7 * 2000, 4 * 10000 + 8 * 20000, ] ), ) assert numpy.allclose( a2 @ a1, numpy.array( [ 1 * 10 + 5 * 20, 2 * 100 + 6 * 200, 3 * 1000 + 7 * 2000, 4 * 10000 + 8 * 20000, ] ), ) assert numpy.allclose( v1 @ a2, numpy.array( [ 1 * 10 + 5 * 20, 1 * 100 + 5 * 200, 1 * 1000 + 5 * 2000, 1 * 10000 + 5 * 20000, ] ), ) assert numpy.allclose( a2 @ v1, numpy.array( [ 1 * 10 + 5 * 20, 1 * 100 + 5 * 200, 1 * 1000 + 5 * 2000, 1 * 10000 + 5 * 20000, ] ), ) with pytest.raises(TypeError): v1 @ 5 with pytest.raises(TypeError): a1 @ 5 vector-1.6.3/tests/backends/test_sympy.py000066400000000000000000000350351503546127100205430ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") rho, phi = sympy.symbols("rho phi") z, eta, theta = sympy.symbols("z eta theta") t, tau = sympy.symbols("t tau") px, py = sympy.symbols("px py") pt = sympy.symbols("pt") pz = sympy.symbols("pz") M, m, mass, E, e, energy = sympy.symbols("M m mass E e energy") def test_construction(): # generic # 2D coords for vec_cls in (vector.VectorSympy2D, vector.VectorSympy3D, vector.VectorSympy4D): coords = {"x": x, "y": y} if vec_cls in (vector.VectorSympy3D, vector.VectorSympy4D): coords["z"] = z if vec_cls == vector.VectorSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.x == x assert vec.y == y assert vec.phi == sympy.atan2(y, x) assert vec.rho == sympy.sqrt(x**2 + y**2) assert isinstance(vec.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert vec.azimuthal.elements == (x, y) coords = {"rho": rho, "phi": phi} if vec_cls in (vector.VectorSympy3D, vector.VectorSympy4D): coords["z"] = z if vec_cls == vector.VectorSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.x == rho * sympy.cos(phi) assert vec.y == rho * sympy.sin(phi) assert vec.phi == phi assert vec.rho == rho assert isinstance(vec.azimuthal, vector.backends.sympy.AzimuthalSympyRhoPhi) assert vec.azimuthal.elements == (rho, phi) # 3D coords for vec_cls in (vector.VectorSympy3D, vector.VectorSympy4D): coords = {"x": x, "y": y, "z": z} if vec_cls == vector.VectorSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.z == z assert vec.eta == sympy.asinh(z / sympy.sqrt(x**2 + y**2)) assert vec.theta == sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) assert isinstance(vec.longitudinal, vector.backends.sympy.LongitudinalSympyZ) assert vec.longitudinal.elements == (z,) coords = {"x": x, "y": y, "eta": eta} if vec_cls == vector.VectorSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.z == sympy.sqrt(x**2 + y**2) * sympy.sinh(eta) assert vec.eta == eta assert vec.theta == 2.0 * sympy.atan(sympy.exp(-eta)) assert isinstance(vec.longitudinal, vector.backends.sympy.LongitudinalSympyEta) assert vec.longitudinal.elements == (eta,) coords = {"x": x, "y": y, "theta": theta} if vec_cls == vector.VectorSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.z == sympy.sqrt(x**2 + y**2) / sympy.tan(theta) assert vec.eta == -sympy.log(sympy.tan(0.5 * theta)) assert vec.theta == theta assert isinstance( vec.longitudinal, vector.backends.sympy.LongitudinalSympyTheta ) assert vec.longitudinal.elements == (theta,) # 4D coords for vec_cls in (vector.VectorSympy4D,): coords = {"x": x, "y": y, "z": z, "t": t} vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.t == t assert vec.tau == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert isinstance(vec.temporal, vector.backends.sympy.TemporalSympyT) assert vec.temporal.elements == (t,) coords = {"x": x, "y": y, "z": z, "tau": tau} vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.t == sympy.sqrt(tau**2 + x**2 + y**2 + z**2) assert vec.tau == tau assert isinstance(vec.temporal, vector.backends.sympy.TemporalSympyTau) assert vec.temporal.elements == (tau,) # momentum # 2D coords for vec_cls in ( vector.MomentumSympy2D, vector.MomentumSympy3D, vector.MomentumSympy4D, ): coords = {"px": px, "py": py} if vec_cls in (vector.MomentumSympy3D, vector.MomentumSympy4D): coords["z"] = z if vec_cls == vector.MomentumSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.px == px assert vec.py == py assert vec.phi == sympy.atan2(py, px) assert vec.rho == sympy.sqrt(px**2 + py**2) assert vec.pt == sympy.sqrt(px**2 + py**2) assert isinstance(vec.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert vec.azimuthal.elements == (px, py) coords = {"pt": pt, "phi": phi} if vec_cls in (vector.MomentumSympy3D, vector.MomentumSympy4D): coords["z"] = z if vec_cls == vector.MomentumSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.px == pt * sympy.cos(phi) assert vec.py == pt * sympy.sin(phi) assert vec.phi == phi assert vec.rho == pt assert vec.pt == pt assert isinstance(vec.azimuthal, vector.backends.sympy.AzimuthalSympyRhoPhi) assert vec.azimuthal.elements == (pt, phi) # 3D coords for vec_cls in (vector.MomentumSympy3D, vector.MomentumSympy4D): coords = {"px": px, "py": py, "pz": pz} if vec_cls == vector.MomentumSympy4D: coords["t"] = t vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.pz == pz assert vec.eta == sympy.asinh(pz / sympy.sqrt(px**2 + py**2)) assert vec.theta == sympy.acos(pz / sympy.sqrt(px**2 + py**2 + pz**2)) assert isinstance(vec.longitudinal, vector.backends.sympy.LongitudinalSympyZ) assert vec.longitudinal.elements == (pz,) # 4D coords for vec_cls in (vector.MomentumSympy4D,): coords = {"px": px, "py": py, "pz": pz, "E": E} vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.m == sympy.sqrt(sympy.Abs(px**2 + py**2 + pz**2 - E**2)) assert sympy.sqrt(sympy.Abs(px**2 + py**2 + pz**2 - E**2)) == vec.M assert vec.mass == sympy.sqrt(sympy.Abs(px**2 + py**2 + pz**2 - E**2)) assert vec.e == E assert vec.E == E assert vec.energy == E assert isinstance(vec.temporal, vector.backends.sympy.TemporalSympyT) assert vec.temporal.elements == (E,) coords = {"px": px, "py": py, "pz": pz, "M": M} vec = vec_cls(**coords) with pytest.raises(TypeError): vec = vec_cls(bad=1, **coords) assert vec.m == M assert vec.M == M assert vec.mass == M assert vec.e == sympy.sqrt(M**2 + px**2 + py**2 + pz**2) assert sympy.sqrt(M**2 + px**2 + py**2 + pz**2) == vec.E assert vec.energy == sympy.sqrt(M**2 + px**2 + py**2 + pz**2) assert isinstance(vec.temporal, vector.backends.sympy.TemporalSympyTau) assert vec.temporal.elements == (M,) def test_type_checks(): with pytest.raises(TypeError): vector.VectorSympy2D(x=1, y=2) with pytest.raises(TypeError): vector.VectorSympy2D() with pytest.raises(TypeError): vector.VectorSympy3D(x=1, y=2, z=3) with pytest.raises(TypeError): vector.VectorSympy3D() with pytest.raises(TypeError): vector.VectorSympy4D(x=1, y=2, z=3, t=4) with pytest.raises(TypeError): vector.VectorSympy4D() nx, ny, nz = sympy.symbols("nx ny nz") def test_eq(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1 == v1 # noqa: PLR0124 assert not v1 == v2 # noqa: SIM201 with pytest.raises(TypeError): v1.equal(v2.to_Vector3D()) def test_ne(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert not v1 != v1 # noqa: PLR0124,SIM202 assert v1 != v2 with pytest.raises(TypeError): v1.not_equal(v2.to_Vector3D()) def test_abs(): v1 = vector.VectorSympy2D(x=x, y=y) assert abs(v1) == sympy.sqrt(x**2 + y**2) v2 = vector.VectorSympy3D(x=x, y=y, z=z) assert abs(v2) == sympy.sqrt(x**2 + y**2 + z**2) v3 = vector.VectorSympy4D(x=x, y=y, z=z, t=t) assert abs(v3) == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) def test_add(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1 + v2 == vector.VectorSympy2D(x=x + nx, y=y + ny) assert v1 + v2.to_Vector3D().like(v1) == vector.VectorSympy2D(x=x + nx, y=y + ny) v1 += v2 assert v1 == vector.VectorSympy2D(x=x + nx, y=y + ny) with pytest.raises(TypeError): v1 + 5 with pytest.raises(TypeError): 5 + v1 with pytest.raises(TypeError): v1 + v2.to_Vector3D() def test_sub(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1 - v2 == vector.VectorSympy2D(x=x - nx, y=y - ny) v1 -= v2 assert v1 == vector.VectorSympy2D(x=x - nx, y=y - ny) with pytest.raises(TypeError): v1 - 5 with pytest.raises(TypeError): 5 - v1 with pytest.raises(TypeError): v1 - v2.to_Vector3D() def test_mul(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1 * 10 == vector.VectorSympy2D(x=x * 10, y=y * 10) assert 10 * v1 == vector.VectorSympy2D(x=10 * x, y=10 * y) v1 *= 10 assert v1 == vector.VectorSympy2D(x=10 * x, y=10 * y) with pytest.raises(TypeError): v1 * v2 def test_neg(): v1 = vector.VectorSympy2D(x=x, y=y) assert -v1 == vector.VectorSympy2D(x=-x, y=-y) def test_pos(): v1 = vector.VectorSympy2D(x=x, y=y) assert +v1 == vector.VectorSympy2D(x=x, y=y) def test_truediv(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1 / 10 == vector.VectorSympy2D(x=0.1 * x, y=0.1 * y) v1 /= 10 assert v1 == vector.VectorSympy2D(x=0.1 * x, y=0.1 * y) with pytest.raises(TypeError): 10 / v1 with pytest.raises(TypeError): v1 / v2 def test_pow(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1**2 == x**2 + y**2 with pytest.raises(TypeError): 2**v1 with pytest.raises(TypeError): v1**v2 def test_matmul(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy2D(x=nx, y=ny) assert v1 @ v2 == x * nx + y * ny assert v2 @ v1 == x * nx + y * ny with pytest.raises(TypeError): v1 @ 5 def test_reprs(): assert ( vector.backends.sympy.AzimuthalSympyXY(x, y).__repr__() == "AzimuthalSympyXY(x=x, y=y)" ) assert ( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi).__repr__() == "AzimuthalSympyRhoPhi(rho=rho, phi=phi)" ) assert ( vector.backends.sympy.LongitudinalSympyZ(z).__repr__() == "LongitudinalSympyZ(z=z)" ) assert ( vector.backends.sympy.LongitudinalSympyEta(eta).__repr__() == "LongitudinalSympyEta(eta=eta)" ) assert ( vector.backends.sympy.LongitudinalSympyTheta(theta).__repr__() == "LongitudinalSympyTheta(theta=theta)" ) assert vector.backends.sympy.TemporalSympyT(t).__repr__() == "TemporalSympyT(t=t)" assert ( vector.backends.sympy.TemporalSympyTau(tau).__repr__() == "TemporalSympyTau(tau=tau)" ) assert vector.VectorSympy2D(x=x, y=y).__repr__() == "VectorSympy2D(x=x, y=y)" assert ( vector.VectorSympy3D(x=x, y=y, z=z).__repr__() == "VectorSympy3D(x=x, y=y, z=z)" ) assert ( vector.VectorSympy4D(x=x, y=y, z=z, t=t).__repr__() == "VectorSympy4D(x=x, y=y, z=z, t=t)" ) assert ( vector.MomentumSympy2D(px=px, py=py).__repr__() == "MomentumSympy2D(px=px, py=py)" ) assert ( vector.MomentumSympy3D(px=px, py=py, pz=pz).__repr__() == "MomentumSympy3D(px=px, py=py, pz=pz)" ) assert ( vector.MomentumSympy4D(px=px, py=py, pz=pz, M=M).__repr__() == "MomentumSympy4D(px=px, py=py, pz=pz, mass=M)" ) def test_setters(): v = vector.VectorSympy2D(x=x, y=y) v.x = 2 * x assert v.x == 2 * x v.y = 2 * y assert v.y == 2 * y v.rho = 2 * rho assert v.rho == 2 * rho v.phi = 2 * phi assert v.phi == 2 * phi v = vector.MomentumSympy2D(px=px, py=py) v.px = 2 * px assert v.px == 2 * px v.py = 2 * py assert v.py == 2 * py v.pt = 2 * pt assert v.pt == 2 * pt v = vector.VectorSympy3D(x=x, y=y, z=z) v.x = 2 * x assert v.x == 2 * x v.y = 2 * y assert v.y == 2 * y v.z = 2 * z assert v.z == 2 * z v.eta = 2 * eta assert v.eta == 2 * eta v.theta = 2 * theta assert v.theta == 2 * theta v = vector.MomentumSympy3D(px=px, py=py, pz=pz) v.px = 2 * px assert v.px == 2 * px v.py = 2 * py assert v.py == 2 * py v.pt = 2 * pt assert v.pt == 2 * pt v.pz = 2 * pz assert v.pz == 2 * pz v = vector.VectorSympy4D(x=x, y=y, z=z, t=t) v.x = 2 * x assert v.x == 2 * x v.y = 2 * y assert v.y == 2 * y v.z = 2 * z assert v.z == 2 * z v.eta = 2 * eta assert v.eta == 2 * eta v.theta = 2 * theta assert v.theta == 2 * theta v.t = 2 * t assert v.t == 2 * t v.tau = 2 * tau assert v.tau == 2 * tau v = vector.MomentumSympy4D(px=px, py=py, pz=pz, m=M) v.px = 2 * px assert v.px == 2 * px v.py = 2 * py assert v.py == 2 * py v.pt = 2 * pt assert v.pt == 2 * pt v.pz = 2 * pz assert v.pz == 2 * pz v.m = 2 * m assert v.m == 2 * m v.mass = 2 * mass assert v.mass == 2 * mass v.M = 2 * M assert v.M == 2 * M v.e = 2 * e assert v.e == 2 * e v.energy = 2 * energy assert v.energy == 2 * energy v.E = 2 * E assert v.E == 2 * E vector-1.6.3/tests/compute/000077500000000000000000000000001503546127100156455ustar00rootroot00000000000000vector-1.6.3/tests/compute/__init__.py000066400000000000000000000003641503546127100177610ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/lorentz/000077500000000000000000000000001503546127100173425ustar00rootroot00000000000000vector-1.6.3/tests/compute/lorentz/__init__.py000066400000000000000000000003641503546127100214560ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/lorentz/test_Et.py000066400000000000000000000101111503546127100213150ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_xy_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_xy_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_xy_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_xy_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_xy_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_rhophi_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_rhophi_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_rhophi_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_rhophi_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_rhophi_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Et == pytest.approx(math.sqrt(80)) def test_rhophi_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et == pytest.approx(math.sqrt(80)) vector-1.6.3/tests/compute/lorentz/test_Et2.py000066400000000000000000000077041503546127100214150ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Et2 == pytest.approx(80) def test_xy_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et2 == pytest.approx(80) def test_xy_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Et2 == pytest.approx(80) def test_xy_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et2 == pytest.approx(80) def test_xy_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Et2 == pytest.approx(80) def test_xy_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et2 == pytest.approx(80) def test_rhophi_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Et2 == pytest.approx(80) def test_rhophi_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et2 == pytest.approx(80) def test_rhophi_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Et2 == pytest.approx(80) def test_rhophi_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et2 == pytest.approx(80) def test_rhophi_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Et2 == pytest.approx(80) def test_rhophi_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Et2 == pytest.approx(80) vector-1.6.3/tests/compute/lorentz/test_Mt.py000066400000000000000000000101251503546127100213320ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_xy_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_xy_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_xy_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_xy_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_xy_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_rhophi_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_rhophi_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_rhophi_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_rhophi_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_rhophi_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt == pytest.approx(math.sqrt(300)) def test_rhophi_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt == pytest.approx(math.sqrt(300)) vector-1.6.3/tests/compute/lorentz/test_Mt2.py000066400000000000000000000077201503546127100214230ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt2 == pytest.approx(300) def test_xy_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt2 == pytest.approx(300) def test_xy_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt2 == pytest.approx(300) def test_xy_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt2 == pytest.approx(300) def test_xy_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt2 == pytest.approx(300) def test_xy_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt2 == pytest.approx(300) def test_rhophi_z_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt2 == pytest.approx(300) def test_rhophi_z_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt2 == pytest.approx(300) def test_rhophi_theta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt2 == pytest.approx(300) def test_rhophi_theta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt2 == pytest.approx(300) def test_rhophi_eta_t(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.Mt2 == pytest.approx(300) def test_rhophi_eta_tau(): vec = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.Mt2 == pytest.approx(300) vector-1.6.3/tests/compute/lorentz/test_beta.py000066400000000000000000000107701503546127100216730ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.beta == pytest.approx(0.5590169943749475) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.beta == pytest.approx(0.5590169943749475) vector-1.6.3/tests/compute/lorentz/test_boostX_beta.py000066400000000000000000000026071503546127100232310ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(1), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.boostX(beta=-0.9428090415820634) assert out.x == pytest.approx(-2.313708498984761) assert out.y == pytest.approx(2) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostX(beta=-0.9428090415820634) assert out.x == pytest.approx(-2.313708498984761) assert out.y == pytest.approx(2) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/lorentz/test_boostX_gamma.py000066400000000000000000000025471503546127100234030ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(1), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.boostX(gamma=-3) assert out.x == pytest.approx(-2.313708498984761) assert out.y == pytest.approx(2) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostX(gamma=-3) assert out.x == pytest.approx(-2.313708498984761) assert out.y == pytest.approx(2) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/lorentz/test_boostY_beta.py000066400000000000000000000026071503546127100232320ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(2, 3), longitudinal=vector.backends.object.LongitudinalObjectZ(1), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.boostY(beta=-0.9428090415820634) assert out.x == pytest.approx(2) assert out.y == pytest.approx(-2.313708498984761) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostY(beta=-0.9428090415820634) assert out.x == pytest.approx(2) assert out.y == pytest.approx(-2.313708498984761) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/lorentz/test_boostY_gamma.py000066400000000000000000000025471503546127100234040ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(2, 3), longitudinal=vector.backends.object.LongitudinalObjectZ(1), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.boostY(gamma=-3) assert out.x == pytest.approx(2) assert out.y == pytest.approx(-2.313708498984761) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostY(gamma=-3) assert out.x == pytest.approx(2) assert out.y == pytest.approx(-2.313708498984761) assert out.z == pytest.approx(1) assert out.t == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/lorentz/test_boostZ_beta.py000066400000000000000000000026071503546127100232330ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.boostZ(beta=-0.9428090415820634) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(-2.313708498984761) assert out.t == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostZ(beta=-0.9428090415820634) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(-2.313708498984761) assert out.t == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/lorentz/test_boostZ_gamma.py000066400000000000000000000025471503546127100234050ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.boostZ(gamma=-3) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(-2.313708498984761) assert out.t == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostZ(gamma=-3) assert out.x == pytest.approx(1) assert out.y == pytest.approx(2) assert out.z == pytest.approx(-2.313708498984761) assert out.t == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/lorentz/test_boost_beta3.py000066400000000000000000000035361503546127100231660ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) beta = vector.backends.object.MomentumObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(5 / 15, 6 / 15), longitudinal=vector.backends.object.LongitudinalObjectZ(7 / 15), ) out = vec.boost_beta3(beta) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): tvec, tbeta = getattr(vec, "to_" + t1)(), getattr(beta, "to_" + t2)() out = tvec.boost_beta3(tbeta) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) vector-1.6.3/tests/compute/lorentz/test_boost_p4.py000066400000000000000000000040361503546127100225070ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) p4 = vector.backends.object.MomentumObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 6), longitudinal=vector.backends.object.LongitudinalObjectZ(7), temporal=vector.backends.object.TemporalObjectT(15), ) out = vec.boost_p4(p4) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec, tp4 = getattr(vec, "to_" + t1)(), getattr(p4, "to_" + t2)() out = tvec.boost_p4(tp4) assert out.x == pytest.approx(3.5537720741941676) assert out.y == pytest.approx(5.0645264890330015) assert out.z == pytest.approx(6.575280903871835) assert out.t == pytest.approx(9.138547120755076) vector-1.6.3/tests/compute/lorentz/test_deltaRapidityPhi.py000066400000000000000000000064461503546127100242250ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_lorentz_object(): v1 = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(1.0, 1.0), vector.backends.object.LongitudinalObjectZ(1.0), vector.backends.object.TemporalObjectTau(1.0), ) v2 = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(-1.0, -1.0), vector.backends.object.LongitudinalObjectZ(-1.0), vector.backends.object.TemporalObjectTau(1.0), ) expected_result = numpy.sqrt( # phi numpy.pi**2 # rapidity + ((0.5 * numpy.log(3 / 1) - 0.5 * numpy.log(1 / 3)) ** 2) ) assert v1.deltaRapidityPhi(v2) == pytest.approx(expected_result) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaRapidityPhi(tr2) == pytest.approx(expected_result) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(1.0, 1.0, 1.0, 1.0)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("tau", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(-1.0, -1.0, -1.0, 1.0)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("tau", numpy.float64), ], ) expected_result = numpy.sqrt( # phi numpy.pi**2 # rapidity + ((0.5 * numpy.log(3 / 1) - 0.5 * numpy.log(1 / 3)) ** 2) ) assert v1.deltaRapidityPhi(v2) == pytest.approx(expected_result) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert numpy.allclose(tr1.deltaRapidityPhi(tr2), expected_result) vector-1.6.3/tests/compute/lorentz/test_deltaRapidityPhi2.py000066400000000000000000000064261503546127100243050ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_lorentz_object(): v1 = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(1.0, 1.0), vector.backends.object.LongitudinalObjectZ(1.0), vector.backends.object.TemporalObjectTau(1.0), ) v2 = vector.backends.object.MomentumObject4D( vector.backends.object.AzimuthalObjectXY(-1.0, -1.0), vector.backends.object.LongitudinalObjectZ(-1.0), vector.backends.object.TemporalObjectTau(1.0), ) expected_result = ( # phi numpy.pi**2 # rapidity + ((0.5 * numpy.log(3 / 1) - 0.5 * numpy.log(1 / 3)) ** 2) ) assert v1.deltaRapidityPhi2(v2) == pytest.approx(expected_result) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaRapidityPhi2(tr2) == pytest.approx(expected_result) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(1.0, 1.0, 1.0, 1.0)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("tau", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(-1.0, -1.0, -1.0, 1.0)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("tau", numpy.float64), ], ) expected_result = ( # phi numpy.pi**2 # rapidity + ((0.5 * numpy.log(3 / 1) - 0.5 * numpy.log(1 / 3)) ** 2) ) assert v1.deltaRapidityPhi2(v2) == pytest.approx(expected_result) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert numpy.allclose(tr1.deltaRapidityPhi2(tr2), expected_result) vector-1.6.3/tests/compute/lorentz/test_gamma.py000066400000000000000000000102041503546127100220320ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(3, 4), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectT(20), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectZ(10), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectT(20), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectT(20), ) assert vec.gamma == pytest.approx(1.2060453783110545) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(5, 0), vector.backends.object.LongitudinalObjectEta(1.4436354751788103), vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.gamma == pytest.approx(1.2060453783110545) vector-1.6.3/tests/compute/lorentz/test_is_lightlike.py000066400000000000000000000154661503546127100234360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(1), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_lightlike() def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(0), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_lightlike() def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(1), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_lightlike() def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(0), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_lightlike() def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(1), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_lightlike() def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(0), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_lightlike() def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(1), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_lightlike() def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(0), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(-4.58257569495584), ) assert not vec.is_lightlike() def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(1), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_lightlike() def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(0), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(-4.58257569495584), ) assert not vec.is_lightlike() def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(1), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_lightlike() def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(0), ) assert vec.is_lightlike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(-4.58257569495584), ) assert not vec.is_lightlike() vector-1.6.3/tests/compute/lorentz/test_is_spacelike.py000066400000000000000000000235751503546127100234220ustar00rootroot00000000000000# Copyright (c) 019-024, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(0), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_spacelike() def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(-1), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_spacelike() def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(0), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_spacelike() def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(-1), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_spacelike() def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(0), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_spacelike() def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(-1), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_spacelike() def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(0), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_spacelike() def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(-1), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_spacelike() def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(0), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_spacelike() def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(-1), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_spacelike() def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(0), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(2), ) assert not vec.is_spacelike() def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(-1), ) assert vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_spacelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert not vec.is_spacelike() vector-1.6.3/tests/compute/lorentz/test_is_timelike.py000066400000000000000000000235311503546127100232550ustar00rootroot00000000000000# Copyright (c) 019-024, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(2), ) assert vec.is_timelike() def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(-1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert vec.is_timelike() def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(2), ) assert vec.is_timelike() def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(-1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert vec.is_timelike() def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(2), ) assert vec.is_timelike() def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(-1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectXY(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert vec.is_timelike() def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectT(2), ) assert vec.is_timelike() def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(-1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectZ(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert vec.is_timelike() def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectT(2), ) assert vec.is_timelike() def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(-1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectTheta(1.5707963267948966), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert vec.is_timelike() def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectT(2), ) assert vec.is_timelike() def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(-1), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(0), ) assert not vec.is_timelike() vec = vector.backends.object.VectorObject4D( vector.backends.object.AzimuthalObjectRhoPhi(1, 0), vector.backends.object.LongitudinalObjectEta(0), vector.backends.object.TemporalObjectTau(1.7320508075688772), ) assert vec.is_timelike() vector-1.6.3/tests/compute/lorentz/test_rapidity.py000066400000000000000000000110501503546127100225750ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.rapidity == pytest.approx(0.5493061443340549) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.rapidity == pytest.approx(0.5493061443340549) vector-1.6.3/tests/compute/lorentz/test_t.py000066400000000000000000000104241503546127100212170ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t == pytest.approx(20) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t == pytest.approx(20) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t == pytest.approx(20) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t == pytest.approx(20) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t == pytest.approx(20) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t == pytest.approx(20) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t == pytest.approx(20) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t == pytest.approx(20) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t == pytest.approx(20) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t == pytest.approx(20) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t == pytest.approx(20) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t == pytest.approx(20) vector-1.6.3/tests/compute/lorentz/test_t2.py000066400000000000000000000104541503546127100213040ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t2 == pytest.approx(400) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t2 == pytest.approx(400) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t2 == pytest.approx(400) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t2 == pytest.approx(400) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t2 == pytest.approx(400) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t2 == pytest.approx(400) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t2 == pytest.approx(400) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t2 == pytest.approx(400) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t2 == pytest.approx(400) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t2 == pytest.approx(400) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.t2 == pytest.approx(400) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.t2 == pytest.approx(400) vector-1.6.3/tests/compute/lorentz/test_tau.py000066400000000000000000000107101503546127100215430ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau == pytest.approx(16.583123951777) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau == pytest.approx(16.583123951777) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau == pytest.approx(16.583123951777) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau == pytest.approx(16.583123951777) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau == pytest.approx(16.583123951777) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau == pytest.approx(16.583123951777) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau == pytest.approx(16.583123951777) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau == pytest.approx(16.583123951777) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau == pytest.approx(16.583123951777) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau == pytest.approx(16.583123951777) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau == pytest.approx(16.583123951777) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau == pytest.approx(16.583123951777) vector-1.6.3/tests/compute/lorentz/test_tau2.py000066400000000000000000000105041503546127100216260ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau2 == pytest.approx(275) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau2 == pytest.approx(275) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau2 == pytest.approx(275) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau2 == pytest.approx(275) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau2 == pytest.approx(275) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau2 == pytest.approx(275) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau2 == pytest.approx(275) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau2 == pytest.approx(275) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau2 == pytest.approx(275) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau2 == pytest.approx(275) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) assert vec.tau2 == pytest.approx(275) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) assert vec.tau2 == pytest.approx(275) vector-1.6.3/tests/compute/lorentz/test_to_beta3.py000066400000000000000000000200541503546127100224540ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_xy_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_xy_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_xy_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_xy_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_xy_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(3 / 20) assert out.y == pytest.approx(4 / 20) assert out.z == pytest.approx(10 / 20) def test_rhophi_z_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectT(20), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(5 / 20) assert out.y == pytest.approx(0 / 20) assert out.z == pytest.approx(10 / 20) def test_rhophi_z_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(5 / 20) assert out.y == pytest.approx(0 / 20) assert out.z == pytest.approx(10 / 20) def test_rhophi_theta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectT(20), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(5 / 20) assert out.y == pytest.approx(0 / 20) assert out.z == pytest.approx(10 / 20) def test_rhophi_theta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(5 / 20) assert out.y == pytest.approx(0 / 20) assert out.z == pytest.approx(10 / 20) def test_rhophi_eta_t(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectT(20), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(5 / 20) assert out.y == pytest.approx(0 / 20) assert out.z == pytest.approx(10 / 20) def test_rhophi_eta_tau(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), temporal=vector.backends.object.TemporalObjectTau(16.583123951777), ) out = vec.to_beta3() assert isinstance(out, vector.backends.object.VectorObject3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == pytest.approx(5 / 20) assert out.y == pytest.approx(0 / 20) assert out.z == pytest.approx(10 / 20) vector-1.6.3/tests/compute/planar/000077500000000000000000000000001503546127100171225ustar00rootroot00000000000000vector-1.6.3/tests/compute/planar/__init__.py000066400000000000000000000003641503546127100212360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/planar/test_deltaphi.py000066400000000000000000000032031503546127100223230ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy_xy(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert v1.deltaphi(v2) == pytest.approx(math.atan2(2, 1) - math.atan2(4, 3)) def test_xy_rhophi(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(3, 4) ) assert v1.deltaphi(v2) == pytest.approx(math.atan2(2, 1) - 4) def test_rhophi_xy(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(1, 2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert v1.deltaphi(v2) == pytest.approx(2 - math.atan2(4, 3)) def test_rhophi_rhophi(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(1, 2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(3, 4) ) assert v1.deltaphi(v2) == pytest.approx(2 - 4) vector-1.6.3/tests/compute/planar/test_phi.py000066400000000000000000000013501503546127100213120ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert vec.phi == pytest.approx(math.atan2(4, 3)) def test_rhophi(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, math.atan2(4, 3)) ) assert vec.phi == pytest.approx(math.atan2(4, 3)) vector-1.6.3/tests/compute/planar/test_rho.py000066400000000000000000000013121503546127100213200ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert vec.rho == pytest.approx(5) def test_rhophi(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, math.atan2(4, 3)) ) assert vec.rho == pytest.approx(5) vector-1.6.3/tests/compute/planar/test_rho2.py000066400000000000000000000013161503546127100214060ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert vec.rho2 == pytest.approx(25) def test_rhophi(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, math.atan2(4, 3)) ) assert vec.rho2 == pytest.approx(25) vector-1.6.3/tests/compute/planar/test_rotateZ.py000066400000000000000000000033101503546127100221600ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_xy(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 0) ) assert vec.rotateZ(0.1).x == pytest.approx(0.9950041652780258) assert vec.rotateZ(0.1).y == pytest.approx(0.09983341664682815) array = vector.backends.numpy.VectorNumpy2D( [(0, 0), (1, 0), (0, 1)], dtype=[("x", numpy.float64), ("y", numpy.float64)] ) assert isinstance(array.rotateZ(0.1), vector.backends.numpy.VectorNumpy2D) out = array.rotateZ(0.1) assert out.dtype.names == ("x", "y") assert numpy.allclose(out.x, [0, 0.9950041652780258, -0.09983341664682815]) assert numpy.allclose(out.y, [0, 0.09983341664682815, 0.9950041652780258]) def test_rhophi(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(1, 0) ) assert vec.rotateZ(0.1).rho == pytest.approx(1) assert vec.rotateZ(0.1).phi == pytest.approx(0.1) array = vector.backends.numpy.VectorNumpy2D( [(0, 0), (1, 0), (0, 1)], dtype=[("rho", numpy.float64), ("phi", numpy.float64)] ) assert isinstance(array.rotateZ(0.1), vector.backends.numpy.VectorNumpy2D) out = array.rotateZ(0.1) assert out.dtype.names == ("rho", "phi") assert numpy.allclose(out.rho, [0, 1, 0]) assert numpy.allclose(out.phi, [0.1, 0.1, 1.1]) vector-1.6.3/tests/compute/planar/test_x.py000066400000000000000000000013061503546127100210020ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert vec.x == pytest.approx(3) def test_rhophi(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, math.atan2(4, 3)) ) assert vec.x == pytest.approx(3) vector-1.6.3/tests/compute/planar/test_y.py000066400000000000000000000013061503546127100210030ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) assert vec.y == pytest.approx(4) def test_rhophi(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, math.atan2(4, 3)) ) assert vec.y == pytest.approx(4) vector-1.6.3/tests/compute/spatial/000077500000000000000000000000001503546127100173025ustar00rootroot00000000000000vector-1.6.3/tests/compute/spatial/__init__.py000066400000000000000000000003641503546127100214160ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/spatial/test_costheta.py000066400000000000000000000041231503546127100225250ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.costheta == pytest.approx(math.cos(0.4636476090008061)) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.costheta == pytest.approx(math.cos(0.4636476090008061)) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.costheta == pytest.approx(math.cos(0.4636476090008061)) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.costheta == pytest.approx(math.cos(0.4636476090008061)) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.costheta == pytest.approx(math.cos(0.4636476090008061)) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.costheta == pytest.approx(math.cos(0.4636476090008061)) vector-1.6.3/tests/compute/spatial/test_cottheta.py000066400000000000000000000041531503546127100225310ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.cottheta == pytest.approx(1 / math.tan(0.4636476090008061)) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.cottheta == pytest.approx(1 / math.tan(0.4636476090008061)) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.cottheta == pytest.approx(1 / math.tan(0.4636476090008061)) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.cottheta == pytest.approx(1 / math.tan(0.4636476090008061)) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.cottheta == pytest.approx(1 / math.tan(0.4636476090008061)) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.cottheta == pytest.approx(1 / math.tan(0.4636476090008061)) vector-1.6.3/tests/compute/spatial/test_cross.py000066400000000000000000000060611503546127100220470ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) with pytest.raises(TypeError): out = v1.to_Vector4D().cross(v2) out = v1.cross(v2) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert (out.x, out.y, out.z) == pytest.approx((-0.03, 0.06, -0.03)) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.cross(transformed2) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectZ ) assert (out.x, out.y, out.z) == pytest.approx((-0.03, 0.06, -0.03)) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) with pytest.raises(TypeError): out = v1.to_Vector4D().cross(v2) out = v1.cross(v2) assert isinstance(out, vector.backends.numpy.VectorNumpy3D) assert out.dtype.names == ("x", "y", "z") assert (out[0].x, out[0].y, out[0].z) == pytest.approx((-0.03, 0.06, -0.03)) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.cross(transformed2) assert isinstance(out, vector.backends.numpy.VectorNumpy3D) assert out.dtype.names == ("x", "y", "z") assert (out[0].x, out[0].y, out[0].z) == pytest.approx((-0.03, 0.06, -0.03)) vector-1.6.3/tests/compute/spatial/test_deltaR.py000066400000000000000000000113511503546127100221270ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import numpy import pytest import vector.backends.numpy import vector.backends.object def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) assert v1.deltaR(v2) == pytest.approx(math.sqrt(0.116083865330319)) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaR(transformed2) == pytest.approx( math.sqrt(0.116083865330319) ) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.deltaR(v2)[0] == pytest.approx(math.sqrt(0.116083865330319)) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.deltaR(tr2)[0] == pytest.approx(math.sqrt(0.116083865330319)) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), temporal=vector.backends.object.TemporalObjectT(99), ) assert v1.deltaR(v2) == pytest.approx(math.sqrt(0.116083865330319)) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaR(tr2) == pytest.approx(math.sqrt(0.116083865330319)) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.4, 0.5, 0.6, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.deltaR(v2) == pytest.approx(math.sqrt(0.116083865330319)) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaR(tr2) == pytest.approx(math.sqrt(0.116083865330319)) vector-1.6.3/tests/compute/spatial/test_deltaR2.py000066400000000000000000000112141503546127100222070ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) assert v1.deltaR2(v2) == pytest.approx(0.116083865330319) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaR2(transformed2) == pytest.approx( 0.116083865330319 ) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.deltaR2(v2)[0] == pytest.approx(0.116083865330319) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.deltaR2(tr2)[0] == pytest.approx(0.116083865330319) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), temporal=vector.backends.object.TemporalObjectT(99), ) assert v1.deltaR2(v2) == pytest.approx(0.116083865330319) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaR2(tr2) == pytest.approx(0.116083865330319) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.4, 0.5, 0.6, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.deltaR2(v2) == pytest.approx(0.116083865330319) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaR2(tr2) == pytest.approx(0.116083865330319) vector-1.6.3/tests/compute/spatial/test_deltaangle.py000066400000000000000000000110761503546127100230200ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) assert v1.deltaangle(v2) == pytest.approx(0.225726) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaangle(transformed2) == pytest.approx(0.225726) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.deltaangle(v2)[0] == pytest.approx(0.225726) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.deltaangle(tr2)[0] == pytest.approx(0.225726) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), temporal=vector.backends.object.TemporalObjectT(99), ) assert v1.deltaangle(v2) == pytest.approx(0.225726) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaangle(tr2) == pytest.approx(0.225726) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.4, 0.5, 0.6, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.deltaangle(v2) == pytest.approx(0.225726) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaangle(tr2) == pytest.approx(0.225726) vector-1.6.3/tests/compute/spatial/test_deltaeta.py000066400000000000000000000112341503546127100224770ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) assert v1.deltaeta(v2) == pytest.approx(0.2674387219518324) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaeta(transformed2) == pytest.approx( 0.2674387219518324 ) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.deltaeta(v2)[0] == pytest.approx(0.2674387219518324) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.deltaeta(tr2)[0] == pytest.approx(0.2674387219518324) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), temporal=vector.backends.object.TemporalObjectT(99), ) assert v1.deltaeta(v2) == pytest.approx(0.2674387219518324) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaeta(tr2) == pytest.approx(0.2674387219518324) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.4, 0.5, 0.6, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.deltaeta(v2) == pytest.approx(0.2674387219518324) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaeta(tr2) == pytest.approx(0.2674387219518324) vector-1.6.3/tests/compute/spatial/test_eta.py000066400000000000000000000037541503546127100214750ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.eta == pytest.approx(1.4436354751788103) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.eta == pytest.approx(1.4436354751788103) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.eta == pytest.approx(1.4436354751788103) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.eta == pytest.approx(1.4436354751788103) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.eta == pytest.approx(1.4436354751788103) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.eta == pytest.approx(1.4436354751788103) vector-1.6.3/tests/compute/spatial/test_mag.py000066400000000000000000000037411503546127100214640ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.mag == pytest.approx(math.sqrt(125)) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.mag == pytest.approx(math.sqrt(125)) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.mag == pytest.approx(math.sqrt(125)) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.mag == pytest.approx(math.sqrt(125)) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.mag == pytest.approx(math.sqrt(125)) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.mag == pytest.approx(math.sqrt(125)) vector-1.6.3/tests/compute/spatial/test_mag2.py000066400000000000000000000036301503546127100215430ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.mag2 == pytest.approx(125) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.mag2 == pytest.approx(125) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.mag2 == pytest.approx(125) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.mag2 == pytest.approx(125) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.mag2 == pytest.approx(125) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.mag2 == pytest.approx(125) vector-1.6.3/tests/compute/spatial/test_rotateX.py000066400000000000000000000113731503546127100223460ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector._methods import vector.backends.numpy import vector.backends.object def test_spatial_object(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) out = vec.rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x == pytest.approx(0.1) assert out.y == pytest.approx(0.1195612965657721) assert out.z == pytest.approx(0.340154518364098) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": out = getattr(vec, "to_" + t)().rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x == pytest.approx(0.1) assert out.y == pytest.approx(0.1195612965657721) assert out.z == pytest.approx(0.340154518364098) def test_spatial_numpy(): vec = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) out = vec.rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.1) assert out[0].y == pytest.approx(0.1195612965657721) assert out[0].z == pytest.approx(0.340154518364098) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": out = getattr(vec, "to_" + t)().rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.1) assert out[0].y == pytest.approx(0.1195612965657721) assert out[0].z == pytest.approx(0.340154518364098) def test_lorentz_object(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) out = vec.rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x == pytest.approx(0.1) assert out.y == pytest.approx(0.1195612965657721) assert out.z == pytest.approx(0.340154518364098) for t in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): out = getattr(vec, "to_" + t)().rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x == pytest.approx(0.1) assert out.y == pytest.approx(0.1195612965657721) assert out.z == pytest.approx(0.340154518364098) def test_lorentz_numpy(): vec = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) out = vec.rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.1) assert out[0].y == pytest.approx(0.1195612965657721) assert out[0].z == pytest.approx(0.340154518364098) for t in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): out = getattr(vec, "to_" + t)().rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.1) assert out[0].y == pytest.approx(0.1195612965657721) assert out[0].z == pytest.approx(0.340154518364098) vector-1.6.3/tests/compute/spatial/test_rotateY.py000066400000000000000000000114131503546127100223420ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector._methods import vector.backends.numpy import vector.backends.object def test_spatial_object(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) out = vec.rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x == pytest.approx(0.17111242994742137) assert out.y == pytest.approx(0.2) assert out.z == pytest.approx(0.2659333305877411) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": out = getattr(vec, "to_" + t)().rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x == pytest.approx(0.17111242994742137) assert out.y == pytest.approx(0.2) assert out.z == pytest.approx(0.2659333305877411) def test_spatial_numpy(): vec = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) out = vec.rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.17111242994742137) assert out[0].y == pytest.approx(0.2) assert out[0].z == pytest.approx(0.2659333305877411) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": out = getattr(vec, "to_" + t)().rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.17111242994742137) assert out[0].y == pytest.approx(0.2) assert out[0].z == pytest.approx(0.2659333305877411) def test_lorentz_object(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) out = vec.rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x == pytest.approx(0.17111242994742137) assert out.y == pytest.approx(0.2) assert out.z == pytest.approx(0.2659333305877411) for t in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): out = getattr(vec, "to_" + t)().rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x == pytest.approx(0.17111242994742137) assert out.y == pytest.approx(0.2) assert out.z == pytest.approx(0.2659333305877411) def test_lorentz_numpy(): vec = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) out = vec.rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.17111242994742137) assert out[0].y == pytest.approx(0.2) assert out[0].z == pytest.approx(0.2659333305877411) for t in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): out = getattr(vec, "to_" + t)().rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out[0].x == pytest.approx(0.17111242994742137) assert out[0].y == pytest.approx(0.2) assert out[0].z == pytest.approx(0.2659333305877411) vector-1.6.3/tests/compute/spatial/test_rotate_axis.py000066400000000000000000000167031503546127100232440ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_spatial_object(): axis = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) out = vec.rotate_axis(axis, 0.25) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": taxis, tvec = ( getattr(axis, "to_" + t1)(), getattr(vec, "to_" + t2)(), ) out = tvec.rotate_axis(taxis, 0.25) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectZ ) assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) def test_spatial_numpy(): axis = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) vec = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) out = vec.rotate_axis(axis, 0.25) assert isinstance(out, vector.backends.numpy.VectorNumpy3D) assert out.dtype.names == ("x", "y", "z") assert out[0].x == pytest.approx(0.37483425404335763) assert out[0].y == pytest.approx(0.5383405688588193) assert out[0].z == pytest.approx(0.5828282027463345) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": taxis, tvec = ( getattr(axis, "to_" + t1)(), getattr(vec, "to_" + t2)(), ) out = tvec.rotate_axis(taxis, 0.25) assert isinstance(out, vector.backends.numpy.VectorNumpy3D) assert out.dtype.names == ("x", "y", "z") assert out[0].x == pytest.approx(0.37483425404335763) assert out[0].y == pytest.approx(0.5383405688588193) assert out[0].z == pytest.approx(0.5828282027463345) def test_lorentz_object(): axis = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(99), ) vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), temporal=vector.backends.object.TemporalObjectT(99), ) with pytest.raises(TypeError): out = vec.rotate_axis(axis, 0.25) out = vec.rotate_axis(axis.to_Vector3D(), 0.25) assert isinstance(out, vector.backends.object.VectorObject4D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert hasattr(out, "temporal") assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): taxis, tvec = ( getattr(axis, "to_" + t1)(), getattr(vec, "to_" + t2)(), ) out = tvec.rotate_axis(taxis, 0.25) assert isinstance(out, vector.backends.object.VectorObject4D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance( out.longitudinal, vector.backends.object.LongitudinalObjectZ ) assert hasattr(out, "temporal") assert out.x == pytest.approx(0.37483425404335763) assert out.y == pytest.approx(0.5383405688588193) assert out.z == pytest.approx(0.5828282027463345) def test_lorentz_numpy(): axis = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) vec = vector.backends.numpy.VectorNumpy4D( [(0.4, 0.5, 0.6, 99)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) with pytest.raises(TypeError): out = vec.rotate_axis(axis, 0.25) out = vec.rotate_axis(axis.to_Vector3D(), 0.25) assert isinstance(out, vector.backends.numpy.VectorNumpy4D) assert out.dtype.names == ("x", "y", "z", "t") assert out[0].x == pytest.approx(0.37483425404335763) assert out[0].y == pytest.approx(0.5383405688588193) assert out[0].z == pytest.approx(0.5828282027463345) for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): taxis, tvec = ( getattr(axis, "to_" + t1)(), getattr(vec, "to_" + t2)(), ) out = tvec.rotate_axis(taxis, 0.25) assert isinstance(out, vector.backends.numpy.VectorNumpy4D) assert out.dtype.names in { ("x", "y", "z", "t"), ("x", "y", "z", "tau"), } assert out[0].x == pytest.approx(0.37483425404335763) assert out[0].y == pytest.approx(0.5383405688588193) assert out[0].z == pytest.approx(0.5828282027463345) vector-1.6.3/tests/compute/spatial/test_rotate_euler.py000066400000000000000000000030341503546127100234050ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_spatial_object(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) out = vec.rotate_euler(0.1, 0.2, 0.3) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(0.5956646364506655) assert out.y == pytest.approx(0.409927258162962) assert out.z == pytest.approx(0.4971350761081869) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tvec = getattr(vec, "to_" + t)() out = tvec.rotate_euler(0.1, 0.2, 0.3) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(0.5956646364506655) assert out.y == pytest.approx(0.409927258162962) assert out.z == pytest.approx(0.4971350761081869) vector-1.6.3/tests/compute/spatial/test_rotate_quaternion.py000066400000000000000000000027421503546127100244630ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_spatial_object(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.5, 0.6), longitudinal=vector.backends.object.LongitudinalObjectZ(0.7), ) out = vec.rotate_quaternion(0.1, 0.2, 0.3, 0.4) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(0.078) assert out.y == pytest.approx(0.18) assert out.z == pytest.approx(0.246) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tvec = getattr(vec, "to_" + t)() out = tvec.rotate_quaternion(0.1, 0.2, 0.3, 0.4) assert isinstance(out, vector.backends.object.VectorObject3D) assert isinstance(out.azimuthal, vector.backends.object.AzimuthalObjectXY) assert isinstance(out.longitudinal, vector.backends.object.LongitudinalObjectZ) assert out.x == pytest.approx(0.078) assert out.y == pytest.approx(0.18) assert out.z == pytest.approx(0.246) vector-1.6.3/tests/compute/spatial/test_theta.py000066400000000000000000000037701503546127100220270ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.theta == pytest.approx(0.4636476090008061) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.theta == pytest.approx(0.4636476090008061) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.theta == pytest.approx(0.4636476090008061) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.theta == pytest.approx(0.4636476090008061) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.theta == pytest.approx(0.4636476090008061) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.theta == pytest.approx(0.4636476090008061) vector-1.6.3/tests/compute/spatial/test_z.py000066400000000000000000000036001503546127100211630ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.object def test_xy_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.z == pytest.approx(10) def test_xy_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.z == pytest.approx(10) def test_xy_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.z == pytest.approx(10) def test_rhophi_z(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(10), ) assert vec.z == pytest.approx(10) def test_rhophi_theta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectTheta(0.4636476090008061), ) assert vec.z == pytest.approx(10) def test_rhophi_eta(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectRhoPhi(5, 0), longitudinal=vector.backends.object.LongitudinalObjectEta(1.4436354751788103), ) assert vec.z == pytest.approx(10) vector-1.6.3/tests/compute/sympy/000077500000000000000000000000001503546127100170265ustar00rootroot00000000000000vector-1.6.3/tests/compute/sympy/lorentz/000077500000000000000000000000001503546127100205235ustar00rootroot00000000000000vector-1.6.3/tests/compute/sympy/lorentz/__init__.py000066400000000000000000000003641503546127100226370ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/sympy/lorentz/test_Et.py000066400000000000000000000152721503546127100225130ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et == t * sympy.sqrt(x**2 + y**2) / sympy.sqrt(x**2 + y**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_xy_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) # TODO: the expression blows up on simplifying? assert vec.Et == sympy.sqrt(x**2 + y**2) * sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / sympy.sqrt(x**2 + y**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_xy_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et.simplify() == t * sympy.sqrt(x**2 + y**2) / sympy.sqrt( x**2 + y**2 + z**2 ) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_xy_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Et.simplify() == sympy.sqrt(x**2 + y**2) * sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / sympy.sqrt(x**2 + y**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_xy_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et.simplify() == t / sympy.sqrt(z**2 / (x**2 + y**2) + 1) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_xy_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert ( sympy.simplify( vec.Et.simplify() - 2 * sympy.sqrt( 0.25 * (x**2 + y**2) * (sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1) ** 2 + sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / (sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1) ) == 0 ) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_rhophi_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et.simplify() == rho * t / sympy.sqrt(rho**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_rhophi_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Et.simplify() == rho * sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) / sympy.sqrt(rho**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_rhophi_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et.simplify() == rho * t / sympy.sqrt(rho**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_rhophi_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Et.simplify() == rho * sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) / sympy.sqrt(rho**2 + z**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_rhophi_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et.simplify() == t / sympy.sqrt(1 + z**2 / rho**2) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) def test_rhophi_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Et.simplify() == 2 * sympy.sqrt( 0.25 * rho**2 * (sympy.exp(2 * sympy.asinh(z / rho)) + 1) ** 2 + sympy.exp(2 * sympy.asinh(z / rho)) * sympy.Abs(rho**2 - t**2 + z**2) ) / (sympy.exp(2 * sympy.asinh(z / rho)) + 1) assert vec.Et.subs(values).evalf() == pytest.approx(math.sqrt(80)) vector-1.6.3/tests/compute/sympy/lorentz/test_Et2.py000066400000000000000000000135721503546127100225760ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et2 == t**2 * (x**2 + y**2) / (x**2 + y**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_xy_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Et2.simplify() == (x**2 + y**2) * ( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / (x**2 + y**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_xy_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et2.simplify() == t**2 * (x**2 + y**2) / (x**2 + y**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_xy_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Et2.simplify() == (x**2 + y**2) * ( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / (x**2 + y**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_xy_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et2.simplify() == t**2 / (z**2 / (x**2 + y**2) + 1) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_xy_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Et2.simplify() == 1.0 * x**2 + 1.0 * y**2 + sympy.Abs( -(t**2) + x**2 + y**2 + z**2 ) / (z**2 / (x**2 + y**2) + 1) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_rhophi_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et2.simplify() == rho**2 * t**2 / (rho**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_rhophi_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Et2.simplify() == rho**2 * ( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) / (rho**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_rhophi_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et2.simplify() == rho**2 * t**2 / (rho**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_rhophi_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Et2.simplify() == rho**2 * ( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) / (rho**2 + z**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_rhophi_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Et2.simplify() == t**2 / (1 + z**2 / rho**2) assert vec.Et2.subs(values).evalf() == pytest.approx(80) def test_rhophi_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Et2.simplify() == 1.0 * rho**2 + sympy.Abs(rho**2 - t**2 + z**2) / ( 1 + z**2 / rho**2 ) assert vec.Et2.subs(values).evalf() == pytest.approx(80) vector-1.6.3/tests/compute/sympy/lorentz/test_Mt.py000066400000000000000000000133721503546127100225220ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt == sympy.sqrt(t**2 - z**2) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_xy_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Mt.simplify() == sympy.sqrt( x**2 + y**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_xy_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_xy_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Mt.simplify() == sympy.sqrt( x**2 + y**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_xy_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_xy_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Mt.simplify() == sympy.sqrt( x**2 + y**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_rhophi_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_rhophi_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Mt.simplify() == sympy.sqrt(rho**2 + sympy.Abs(rho**2 - t**2 + z**2)) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_rhophi_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_rhophi_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Mt.simplify() == sympy.sqrt(rho**2 + sympy.Abs(rho**2 - t**2 + z**2)) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_rhophi_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt.simplify() == sympy.sqrt(t**2 - z**2) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) def test_rhophi_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Mt.simplify() == sympy.sqrt(rho**2 + sympy.Abs(rho**2 - t**2 + z**2)) assert vec.Mt.subs(values).evalf() == pytest.approx(math.sqrt(300)) vector-1.6.3/tests/compute/sympy/lorentz/test_Mt2.py000066400000000000000000000126701503546127100226040ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt2 == t**2 - z**2 assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_xy_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Mt2.simplify() == x**2 + y**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_xy_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt2.simplify() == t**2 - z**2 assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_xy_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Mt2.simplify() == x**2 + y**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_xy_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt2.simplify() == t**2 - z**2 assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_xy_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.Mt2.simplify() == x**2 + y**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_rhophi_z_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt2.simplify() == t**2 - z**2 assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_rhophi_z_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Mt2.simplify() == rho**2 + sympy.Abs(rho**2 - t**2 + z**2) assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_rhophi_theta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt2.simplify() == t**2 - z**2 assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_rhophi_theta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Mt2.simplify() == rho**2 + sympy.Abs(rho**2 - t**2 + z**2) assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_rhophi_eta_t(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.Mt2.simplify() == t**2 - z**2 assert vec.Mt2.subs(values).evalf() == pytest.approx(300) def test_rhophi_eta_tau(): vec = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.Mt2.simplify() == rho**2 + sympy.Abs(rho**2 - t**2 + z**2) assert vec.Mt2.subs(values).evalf() == pytest.approx(300) vector-1.6.3/tests/compute/sympy/lorentz/test_beta.py000066400000000000000000000162101503546127100230470ustar00rootroot00000000000000# Copyright (c) t19-t24, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.beta == sympy.sqrt(x**2 + y**2 + z**2) / t assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_xy_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.beta.simplify() == sympy.sqrt(x**2 + y**2 + z**2) / sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_xy_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.beta.simplify() == sympy.sqrt(x**2 + y**2 + z**2) / t assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_xy_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.beta.simplify() == sympy.sqrt( x**4 + 2 * x**2 * y**2 + x**2 * z**2 + y**4 + y**2 * z**2 ) / ( sympy.sqrt(x**2 + y**2) * sympy.sqrt(x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_xy_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert ( vec.beta.simplify() == 1.0 * sympy.sqrt(x**2 + y**2) * sympy.sqrt(z**2 / (x**2 + y**2) + 1) / t ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_xy_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert ( sympy.simplify( vec.beta.simplify() - 0.5 * sympy.sqrt(x**2 + y**2) * (sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1) / sympy.sqrt( 0.25 * (x**2 + y**2) * (sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1) ** 2 + sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) ) == 0 ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_rhophi_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.beta == sympy.sqrt(rho**2 + z**2) / t assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.beta.simplify() == sympy.sqrt(rho**2 + z**2) / sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.beta.simplify() == sympy.sqrt(rho**2 + z**2) / t assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.beta.simplify() == sympy.sqrt(rho**2 + z**2) / sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.beta.simplify() == 1.0 * rho * sympy.sqrt(1 + z**2 / rho**2) / t assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.beta.simplify() == 0.5 * rho * ( sympy.exp(2 * sympy.asinh(z / rho)) + 1 ) / sympy.sqrt( 0.25 * rho**2 * (sympy.exp(2 * sympy.asinh(z / rho)) + 1) ** 2 + sympy.exp(2 * sympy.asinh(z / rho)) * sympy.Abs(rho**2 - t**2 + z**2) ) assert vec.beta.subs(values).evalf() == pytest.approx(0.5590169943749475) vector-1.6.3/tests/compute/sympy/lorentz/test_boostX_beta.py000066400000000000000000000033231503546127100244060ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 3, y: 2, z: 1, t: 4} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.boostX(beta=-0.9428090415820634) # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values).evalf() == pytest.approx(-2.313708498984761) assert out.y == y assert out.z == z # cannot equate the sympy expressions because of floating point errors assert out.t.subs(values).evalf() == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostX(beta=-0.9428090415820634) assert out.x.subs(values).evalf() == pytest.approx(-2.313708498984761) assert out.y.subs(values).evalf() == pytest.approx(2) assert out.z.subs(values).evalf() == pytest.approx(1) assert out.t.subs(values).evalf() == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/sympy/lorentz/test_boostX_gamma.py000066400000000000000000000027461503546127100245650ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 3, y: 2, z: 1, t: 10} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.boostX(gamma=3) assert out.x == 2 * sympy.sqrt(2) * t + 3 * x assert out.y == y assert out.z == z assert out.t == 3 * t + 2 * sympy.sqrt(2) * x for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostX(gamma=3) assert out.x.subs(values).evalf() == pytest.approx(37.2842712474619) assert out.y.subs(values).evalf() == pytest.approx(2) assert out.z.subs(values).evalf() == pytest.approx(1) assert out.t.subs(values).evalf() == pytest.approx(38.4852813742386) vector-1.6.3/tests/compute/sympy/lorentz/test_boostY_beta.py000066400000000000000000000043151503546127100244110ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 2, y: 3, z: 1, t: 4} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.boostY(beta=-0.9428090415820634) assert out.x == x # cannot equate the sympy expressions because of floating point errors assert out.y.subs(values).evalf() == pytest.approx(-2.313708498984761) assert out.z == z # cannot equate the sympy expressions because of floating point errors assert out.t.subs(values).evalf() == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostY(beta=-0.9428090415820634) assert ( out.x.subs(values).evalf() == pytest.approx(2) if not isinstance(out.x, (float, int)) else out.x == pytest.approx(2) ) assert ( out.y.subs(values).evalf() == pytest.approx(-2.313708498984761) if not isinstance(out.y, (float, int)) else out.y == pytest.approx(-2.313708498984761) ) assert ( out.z.subs(values).evalf() == pytest.approx(1) if not isinstance(out.z, (float, int)) else out.z == pytest.approx(1) ) assert ( out.t.subs(values).evalf() == pytest.approx(3.5147186257614287) if not isinstance(out.t, (float, int)) else out.t == pytest.approx(3.5147186257614287) ) vector-1.6.3/tests/compute/sympy/lorentz/test_boostY_gamma.py000066400000000000000000000027451503546127100245650ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 2, y: 3, z: 1, t: 4} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.boostY(gamma=3) assert out.x == x assert out.y == 2 * sympy.sqrt(2) * t + 3 * y assert out.z == z assert out.t == 3 * t + 2 * sympy.sqrt(2) * y for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostY(gamma=3) assert out.x.subs(values).evalf() == pytest.approx(2) assert out.y.subs(values).evalf() == pytest.approx(20.3137084989848) assert out.z.subs(values).evalf() == pytest.approx(1) assert out.t.subs(values).evalf() == pytest.approx(20.4852813742386) vector-1.6.3/tests/compute/sympy/lorentz/test_boostZ_beta.py000066400000000000000000000031701503546127100244100ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 1, y: 2, z: 3, t: 4} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.boostZ(beta=-0.9428090415820634) assert out.x == x assert out.y == y # cannot equate the sympy expressions because of floating point errors assert out.z.subs(values) == pytest.approx(-2.313708498984761) assert out.t.subs(values) == pytest.approx(3.5147186257614287) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostZ(beta=-0.9428090415820634) assert out.x.subs(values).evalf() == pytest.approx(1) assert out.y.subs(values).evalf() == pytest.approx(2) assert out.z.subs(values).evalf() == pytest.approx(-2.313708498984761) assert out.t.subs(values).evalf() == pytest.approx(3.5147186257614287) vector-1.6.3/tests/compute/sympy/lorentz/test_boostZ_gamma.py000066400000000000000000000030231503546127100245540ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 1, y: 2, z: 3, t: 4} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.boostZ(gamma=3) assert out.x == x assert out.y == y assert out.z.subs(values).evalf() == pytest.approx(20.3137084989848) assert out.t.subs(values).evalf() == pytest.approx(20.4852813742386) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.boostZ(gamma=3) assert out.x.subs(values).evalf() == pytest.approx(1) assert out.y.subs(values).evalf() == pytest.approx(2) assert out.z.subs(values).evalf() == pytest.approx(20.3137084989848) assert out.t.subs(values).evalf() == pytest.approx(20.4852813742386) vector-1.6.3/tests/compute/sympy/lorentz/test_boost_beta3.py000066400000000000000000000056701503546127100243500ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, px, py, pz = sympy.symbols("x y z t px py pz", real=True) values = {x: 1, y: 2, z: 3, t: 4, px: 5 / 15, py: 6 / 15, pz: 7 / 15} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) beta = vector.MomentumSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(px, py), longitudinal=vector.backends.sympy.LongitudinalSympyZ(pz), ) out = vec.boost_beta3(beta) assert out.x.simplify() == ( -px * py * y - px * pz * z - px * t * sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - px * t + py**2 * x + pz**2 * x - x * sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - x ) / (px**2 + py**2 + pz**2 - sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - 1) assert out.y.simplify() == ( px**2 * y - px * py * x - py * pz * z - py * t * sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - py * t + pz**2 * y - y * sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - y ) / (px**2 + py**2 + pz**2 - sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - 1) assert out.z.simplify() == ( px**2 * z - px * pz * x + py**2 * z - py * pz * y - pz * t * sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - pz * t - z * sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - z ) / (px**2 + py**2 + pz**2 - sympy.sqrt(-(px**2) - py**2 - pz**2 + 1) - 1) assert out.t.simplify() == (px * x + py * y + pz * z + t) / sympy.sqrt( -(px**2) - py**2 - pz**2 + 1 ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): tvec, tbeta = getattr(vec, "to_" + t1)(), getattr(beta, "to_" + t2)() out = tvec.boost_beta3(tbeta) assert out.x.subs(values).evalf() == pytest.approx(3.5537720741941676) assert out.y.subs(values).evalf() == pytest.approx(5.0645264890330015) assert out.z.subs(values).evalf() == pytest.approx(6.575280903871835) assert out.t.subs(values).evalf() == pytest.approx(9.138547120755076) vector-1.6.3/tests/compute/sympy/lorentz/test_boost_p4.py000066400000000000000000000063061503546127100236720ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, px, py, pz, M = sympy.symbols("x y z t px py pz M", real=True) values = {x: 1, y: 2, z: 3, t: 4, px: 5, py: 6, pz: 7, M: 15} def test(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) p4 = vector.MomentumSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(px, py), longitudinal=vector.backends.sympy.LongitudinalSympyZ(pz), temporal=vector.backends.sympy.TemporalSympyT(M), ) out = vec.boost_p4(p4) assert out.x.simplify() == ( M**2 * x + M * px * t + M * x * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) + px * py * y + px * pz * z + px * t * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - py**2 * x - pz**2 * x ) / (M**2 + M * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - px**2 - py**2 - pz**2) assert out.y.simplify() == ( M**2 * y + M * py * t + M * y * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - px**2 * y + px * py * x + py * pz * z + py * t * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - pz**2 * y ) / (M**2 + M * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - px**2 - py**2 - pz**2) assert out.z.simplify() == ( M**2 * z + M * pz * t + M * z * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - px**2 * z + px * pz * x - py**2 * z + py * pz * y + pz * t * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) ) / (M**2 + M * sympy.sqrt(M**2 - px**2 - py**2 - pz**2) - px**2 - py**2 - pz**2) assert out.t.simplify() == (M * t + px * x + py * y + pz * z) / sympy.sqrt( M**2 - px**2 - py**2 - pz**2 ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec, tp4 = getattr(vec, "to_" + t1)(), getattr(p4, "to_" + t2)() out = tvec.boost_p4(tp4) assert out.x.subs(values).evalf() == pytest.approx(3.5537720741941676) assert out.y.subs(values).evalf() == pytest.approx(5.0645264890330015) assert out.z.subs(values).evalf() == pytest.approx(6.575280903871835) assert out.t.subs(values).evalf() == pytest.approx(9.138547120755076) vector-1.6.3/tests/compute/sympy/lorentz/test_deltaRapidityPhi.py000066400000000000000000000047551503546127100254070ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, nx, ny, nz, nt = sympy.symbols("x y z t nx ny nz nt", real=True) values = {x: 1.0, y: 1.0, z: 1.0, t: 1.0, nx: -1.0, ny: -1.0, nz: -1.0, nt: 1.0} def test_lorentz_object(): v1 = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau(t), ) v2 = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(nx, ny), vector.backends.sympy.LongitudinalSympyZ(nz), vector.backends.sympy.TemporalSympyTau(nt), ) assert v1.deltaRapidityPhi(v2).simplify() == sympy.sqrt( 0.25 * ( sympy.log( (-nz - sympy.sqrt(nt**2 + nx**2 + ny**2 + nz**2)) / (nz - sympy.sqrt(nt**2 + nx**2 + ny**2 + nz**2)) ) - sympy.log( (-z - sympy.sqrt(t**2 + x**2 + y**2 + z**2)) / (z - sympy.sqrt(t**2 + x**2 + y**2 + z**2)) ) ) ** 2 + ( sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) ** 2 ) expected_result = sympy.sqrt( # phi sympy.pi**2 # rapidity + ((0.5 * sympy.log(3 / 1) - 0.5 * sympy.log(1 / 3)) ** 2) ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaRapidityPhi(tr2).subs(values).evalf() == pytest.approx( expected_result.evalf() ) vector-1.6.3/tests/compute/sympy/lorentz/test_deltaRapidityPhi2.py000066400000000000000000000047421503546127100254650ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, nx, ny, nz, nt = sympy.symbols("x y z t nx ny nz nt", real=True) values = {x: 1.0, y: 1.0, z: 1.0, t: 1.0, nx: -1.0, ny: -1.0, nz: -1.0, nt: 1.0} def test_lorentz_sympy(): v1 = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau(t), ) v2 = vector.MomentumSympy4D( vector.backends.sympy.AzimuthalSympyXY(nx, ny), vector.backends.sympy.LongitudinalSympyZ(nz), vector.backends.sympy.TemporalSympyTau(nt), ) assert ( v1.deltaRapidityPhi2(v2).simplify() == 0.25 * ( sympy.log( (-nz - sympy.sqrt(nt**2 + nx**2 + ny**2 + nz**2)) / (nz - sympy.sqrt(nt**2 + nx**2 + ny**2 + nz**2)) ) - sympy.log( (-z - sympy.sqrt(t**2 + x**2 + y**2 + z**2)) / (z - sympy.sqrt(t**2 + x**2 + y**2 + z**2)) ) ) ** 2 + ( sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) ** 2 ) expected_result = ( # phi sympy.pi**2 # rapidity + ((0.5 * sympy.log(3 / 1) - 0.5 * sympy.log(1 / 3)) ** 2) ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaRapidityPhi2(tr2).subs(values).evalf() == pytest.approx( expected_result.evalf() ) vector-1.6.3/tests/compute/sympy/lorentz/test_gamma.py000066400000000000000000000164511503546127100232250ustar00rootroot00000000000000# Copyright (c) t19-t24, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.gamma == t / sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_xy_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.gamma.simplify() == sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / sympy.Abs(sympy.sqrt(-(t**2) + x**2 + y**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_xy_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.gamma.simplify() == t / sympy.Abs(sympy.sqrt(t**2 - x**2 - y**2 - z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_xy_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.gamma.simplify() == sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / sympy.Abs(sympy.sqrt(-(t**2) + x**2 + y**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_xy_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.gamma == t / sympy.sqrt( sympy.Abs( t**2 - (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) ) ) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_xy_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.gamma == sympy.sqrt( (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) / sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_rhophi_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.gamma.simplify() == t / sympy.Abs(sympy.sqrt(rho**2 - t**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.gamma.simplify() == sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) / sympy.Abs(sympy.sqrt(rho**2 - t**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.gamma.simplify() == t / sympy.Abs(sympy.sqrt(-(rho**2) + t**2 - z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.gamma.simplify() == sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) / sympy.Abs(sympy.sqrt(rho**2 - t**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.gamma == t / sympy.sqrt( sympy.Abs( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / rho))) ** 2 * sympy.exp(2 * sympy.asinh(z / rho)) - t**2 ) ) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.gamma == sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / rho))) ** 2 * sympy.exp(2 * sympy.asinh(z / rho)) + sympy.Abs(rho**2 - t**2 + z**2) ) / sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) assert vec.gamma.subs(values).evalf() == pytest.approx(1.2060453783110545) vector-1.6.3/tests/compute/sympy/lorentz/test_is_lightlike.py000066400000000000000000000115241503546127100246060ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 1, y: 0, rho: 1, phi: 0, z: 0, t: 1} def test_xy_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.is_lightlike().subs(values) def test_xy_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.is_lightlike().subs(values) def test_xy_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.is_lightlike().subs(values) def test_xy_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.is_lightlike().subs(values) def test_xy_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.is_lightlike().subs(values) def test_xy_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.is_lightlike().subs(values) def test_rhophi_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.is_lightlike().subs(values) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.is_lightlike().subs(values) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.is_lightlike().subs(values) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.is_lightlike().subs(values) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert vec.is_lightlike().subs(values) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.is_lightlike().subs(values) vector-1.6.3/tests/compute/sympy/lorentz/test_is_spacelike.py000066400000000000000000000326111503546127100245720ustar00rootroot00000000000000# Copyright (c) 019-024, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, space_t, light_t, time_t = sympy.symbols( "x y rho phi z space_t light_t time_t", real=True ) values = {x: 1, y: 0, rho: 1, phi: 0, z: 0, space_t: 0, light_t: 1, time_t: 2} def test_xy_z_t(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyZ(z), # vector.backends.sympy.TemporalSympyT(space_t), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(time_t), ) assert not vec.is_spacelike().subs(values) def test_xy_z_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyZ(z), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(-(space_t**2) + x**2 + y**2 + z**2))), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(light_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(time_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) def test_xy_theta_t(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyTheta(sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2))), # vector.backends.sympy.TemporalSympyT(space_t), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(time_t), ) assert not vec.is_spacelike().subs(values) def test_xy_theta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyTheta(sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2))), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(-(space_t**2) + x**2 + y**2 + z**2))), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(light_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(time_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) def test_xy_eta_t(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / sympy.sqrt(x**2 + y**2))), # vector.backends.sympy.TemporalSympyT(space_t), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(time_t), ) assert not vec.is_spacelike().subs(values) def test_xy_eta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / sympy.sqrt(x**2 + y**2))), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(-(space_t**2) + x**2 + y**2 + z**2))), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(light_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(time_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) def test_rhophi_z_t(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyZ(z), # vector.backends.sympy.TemporalSympyT(space_t), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(time_t), ) assert not vec.is_spacelike().subs(values) def test_rhophi_z_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyZ(z), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(rho**2 - space_t**2 + z**2))), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - light_t**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - time_t**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) def test_rhophi_theta_t(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyTheta(sympy.acos(z / sympy.sqrt(rho**2 + z**2))), # vector.backends.sympy.TemporalSympyT(space_t), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(time_t), ) assert not vec.is_spacelike().subs(values) def test_rhophi_theta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyTheta(sympy.acos(z / sympy.sqrt(rho**2 + z**2))), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(rho**2 - space_t**2 + z**2))), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - light_t**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - time_t**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) def test_rhophi_eta_t(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), # vector.backends.sympy.TemporalSympyT(space_t), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(time_t), ) assert not vec.is_spacelike().subs(values) def test_rhophi_eta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(rho**2 - space_t**2 + z**2))), # ) # assert vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - light_t**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - time_t**2 + z**2)) ), ) assert not vec.is_spacelike().subs(values) vector-1.6.3/tests/compute/sympy/lorentz/test_is_timelike.py000066400000000000000000000312631503546127100244370ustar00rootroot00000000000000# Copyright (c) 019-024, Jonas Eschle, Jim Pivarski, Eduardo Rodrigues, and Henry Schreiner. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, space_t, light_t, time_t = sympy.symbols( "x y rho phi z space_t light_t time_t", real=True ) values = {x: 1, y: 0, rho: 1, phi: 0, z: 0, space_t: 0, light_t: 1, time_t: 2} def test_xy_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(space_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(time_t), ) assert vec.is_timelike().subs(values) def test_xy_z_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyZ(z), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(-(space_t**2) + x**2 + y**2 + z**2))), # ) # assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(light_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(time_t**2) + x**2 + y**2 + z**2)) ), ) assert vec.is_timelike().subs(values) def test_xy_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(space_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(time_t), ) assert vec.is_timelike().subs(values) def test_xy_theta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyTheta(sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2))), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(-(space_t**2) + x**2 + y**2 + z**2))), # ) # assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(light_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(time_t**2) + x**2 + y**2 + z**2)) ), ) assert vec.is_timelike().subs(values) def test_xy_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(space_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(time_t), ) assert vec.is_timelike().subs(values) def test_xy_eta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyXY(x, y), # vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / sympy.sqrt(x**2 + y**2))), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(-(space_t**2) + x**2 + y**2 + z**2))), # ) # assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(light_t**2) + x**2 + y**2 + z**2)) ), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(time_t**2) + x**2 + y**2 + z**2)) ), ) assert vec.is_timelike().subs(values) def test_rhophi_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(space_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(time_t), ) assert vec.is_timelike().subs(values) def test_rhophi_z_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyZ(z), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(rho**2 - space_t**2 + z**2))), # ) # assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - light_t**2 + z**2)) ), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - time_t**2 + z**2)) ), ) assert vec.is_timelike().subs(values) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(space_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(time_t), ) assert vec.is_timelike().subs(values) def test_rhophi_theta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyTheta(sympy.acos(z / sympy.sqrt(rho**2 + z**2))), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(rho**2 - space_t**2 + z**2))), # ) # assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - light_t**2 + z**2)) ), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - time_t**2 + z**2)) ), ) assert vec.is_timelike().subs(values) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(space_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(light_t), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(time_t), ) assert vec.is_timelike().subs(values) def test_rhophi_eta_tau(): # the following test fails, but it represent the t**2 < mag**2 case # so it should be okay for it to fail # vec = vector.VectorSympy4D( # vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), # vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), # vector.backends.sympy.TemporalSympyTau(sympy.sqrt(sympy.Abs(rho**2 - space_t**2 + z**2))), # ) # assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - light_t**2 + z**2)) ), ) assert not vec.is_timelike().subs(values) vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - time_t**2 + z**2)) ), ) assert vec.is_timelike().subs(values) vector-1.6.3/tests/compute/sympy/lorentz/test_rapidity.py000066400000000000000000000206661503546127100237730ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.rapidity == 0.5 * sympy.log((t + z) / (t - z)) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_xy_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.rapidity.simplify() == 0.5 * sympy.log( (-z - sympy.sqrt(x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2))) / (z - sympy.sqrt(x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2))) ) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_xy_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.rapidity.simplify() == 0.5 * sympy.log((t + z) / (t - z)) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_xy_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.rapidity.simplify() == 0.5 * sympy.log( (-z - sympy.sqrt(x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2))) / (z - sympy.sqrt(x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2))) ) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_xy_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.rapidity.simplify() == 0.5 * sympy.log((t + z) / (t - z)) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_xy_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) # TODO: why won't sympy equate the expressions without double # simplifying? assert ( sympy.simplify( vec.rapidity.simplify() - 0.5 * sympy.log( ( -z - sympy.sqrt( ( 0.25 * (x**2 + y**2) * ( sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1 ) ** 2 + sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) ) ) / ( z - sympy.sqrt( ( 0.25 * (x**2 + y**2) * ( sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1 ) ** 2 + sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) ) ) ) ) == 0 ) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_rhophi_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.rapidity.simplify() == 0.5 * sympy.log((t + z) / (t - z)) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.rapidity.simplify() == 0.5 * sympy.log( (-z - sympy.sqrt(rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2))) / (z - sympy.sqrt(rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2))) ) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.rapidity.simplify() == 0.5 * sympy.log((-t - z) / (-t + z)) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.rapidity.simplify() == 0.5 * sympy.log( (-z - sympy.sqrt(rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2))) / (z - sympy.sqrt(rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2))) ) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.rapidity.simplify() == 0.5 * sympy.log((t + z) / (t - z)) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.rapidity == 0.5 * sympy.log( ( z + sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / rho))) ** 2 * sympy.exp(2 * sympy.asinh(z / rho)) + sympy.Abs(rho**2 - t**2 + z**2) ) ) / ( -z + sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / rho))) ** 2 * sympy.exp(2 * sympy.asinh(z / rho)) + sympy.Abs(rho**2 - t**2 + z**2) ) ) ) assert vec.rapidity.subs(values).evalf() == pytest.approx(0.5493061443340549) vector-1.6.3/tests/compute/sympy/lorentz/test_t.py000066400000000000000000000133441503546127100224040ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t == t assert vec.t.subs(values).evalf() == pytest.approx(20) def test_xy_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.t.simplify() == sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.t.subs(values).evalf() == pytest.approx(20) def test_xy_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t.simplify() == t assert vec.t.subs(values).evalf() == pytest.approx(20) def test_xy_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.t.simplify() == sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.t.subs(values).evalf() == pytest.approx(20) def test_xy_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t.simplify() == t assert vec.t.subs(values).evalf() == pytest.approx(20) def test_xy_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.t == sympy.sqrt( (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert vec.t.subs(values).evalf() == pytest.approx(20) def test_rhophi_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t.simplify() == t assert vec.t.subs(values).evalf() == pytest.approx(20) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.t.simplify() == sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) assert vec.t.subs(values).evalf() == pytest.approx(20) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t.simplify() == t assert vec.t.subs(values).evalf() == pytest.approx(20) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.t.simplify() == sympy.sqrt( rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) ) assert vec.t.subs(values).evalf() == pytest.approx(20) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t.simplify() == t assert vec.t.subs(values).evalf() == pytest.approx(20) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.t == sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / rho))) ** 2 * sympy.exp(2 * sympy.asinh(z / rho)) + sympy.Abs(rho**2 - t**2 + z**2) ) assert vec.t.subs(values).evalf() == pytest.approx(20) vector-1.6.3/tests/compute/sympy/lorentz/test_t2.py000066400000000000000000000131401503546127100224600ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t2 == t**2 assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_xy_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.t2.simplify() == x**2 + y**2 + z**2 + sympy.Abs( -(t**2) + x**2 + y**2 + z**2 ) assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_xy_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t2.simplify() == t**2 assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_xy_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.t2.simplify() == x**2 + y**2 + z**2 + sympy.Abs( -(t**2) + x**2 + y**2 + z**2 ) assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_xy_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t2.simplify() == t**2 assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_xy_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.t2 == ( 0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) ) ** 2 * (x**2 + y**2) * sympy.exp( 2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_rhophi_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t2.simplify() == t**2 assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.t2.simplify() == rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t2.simplify() == t**2 assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.t2.simplify() == rho**2 + z**2 + sympy.Abs(rho**2 - t**2 + z**2) assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.t2.simplify() == t**2 assert vec.t2.subs(values).evalf() == pytest.approx(400) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.t2.simplify() == 0.25 * rho**2 * (4 + 4 * z**2 / rho**2) + sympy.Abs( rho**2 - t**2 + z**2 ) assert vec.t2.subs(values).evalf() == pytest.approx(400) vector-1.6.3/tests/compute/sympy/lorentz/test_tau.py000066400000000000000000000137401503546127100227320ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_xy_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_xy_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_xy_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_xy_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau == sympy.sqrt( sympy.Abs( t**2 - (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) ) ) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_xy_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_rhophi_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(1.0 * rho**2 - t**2 + 1.0 * z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.tau.simplify() == sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) assert vec.tau.subs(values).evalf() == pytest.approx(16.583123951777) vector-1.6.3/tests/compute/sympy/lorentz/test_tau2.py000066400000000000000000000131371503546127100230140ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau2 == t**2 - x**2 - y**2 - z**2 assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_xy_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.tau2.simplify() == sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_xy_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau2.simplify() == t**2 - x**2 - y**2 - z**2 assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_xy_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.tau2.simplify() == sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_xy_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau2 == t**2 - ( 0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) ) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_xy_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyXY(x, y), vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) assert vec.tau2.simplify() == sympy.Abs(-(t**2) + x**2 + y**2 + z**2) assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_rhophi_z_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau2.simplify() == -(rho**2) + t**2 - z**2 assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyZ(z), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.tau2.simplify() == sympy.Abs(rho**2 - t**2 + z**2) assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau2.simplify() == -(rho**2) + t**2 - z**2 assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.tau2.simplify() == sympy.Abs(rho**2 - t**2 + z**2) assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyT(t), ) assert vec.tau2.simplify() == -0.25 * rho**2 * (4 + 4 * z**2 / rho**2) + t**2 assert vec.tau2.subs(values).evalf() == pytest.approx(275) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(rho**2 - t**2 + z**2)) ), ) assert vec.tau2.simplify() == sympy.Abs(rho**2 - t**2 + z**2) assert vec.tau2.subs(values).evalf() == pytest.approx(275) vector-1.6.3/tests/compute/sympy/lorentz/test_to_beta3.py000066400000000000000000000334261503546127100236440ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z, t = sympy.symbols("x y rho phi z t", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10, t: 20} def test_xy_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == x / t assert out.y == y / t assert out.z == z / t assert out.x.subs(values).evalf() == pytest.approx(3 / 20) assert out.y.subs(values).evalf() == pytest.approx(4 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_xy_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == x / sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.y == y / sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.z == z / sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.x.subs(values).evalf() == pytest.approx(3 / 20) assert out.y.subs(values).evalf() == pytest.approx(4 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_xy_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == x / t assert out.y == y / t assert out.z.simplify() == z / sympy.Abs(t) assert out.x.subs(values).evalf() == pytest.approx(3 / 20) assert out.y.subs(values).evalf() == pytest.approx(4 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_xy_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x.simplify() == x / sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.y.simplify() == y / sympy.sqrt( x**2 + y**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.z.simplify() == z * sympy.sqrt(x**2 + y**2) / sympy.sqrt( x**2 * (x**2 + y**2 + z**2) + x**2 * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) + y**2 * (x**2 + y**2 + z**2) + y**2 * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.x.subs(values).evalf() == pytest.approx(3 / 20) assert out.y.subs(values).evalf() == pytest.approx(4 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_xy_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == x / t assert out.y == y / t assert out.z == z * sympy.sqrt(x**2 / t**2 + y**2 / t**2) / sympy.sqrt(x**2 + y**2) assert out.x.subs(values).evalf() == pytest.approx(3 / 20) assert out.y.subs(values).evalf() == pytest.approx(4 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_xy_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == x / sympy.sqrt( (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.y == y / sympy.sqrt( (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * (x**2 + y**2) * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) # TODO: why won't sympy equate the expressions without double # simplifying? assert ( sympy.simplify( out.z - z * sympy.sqrt( sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) / ( 0.25 * (x**2 + y**2) * (sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + 1) ** 2 + sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) ) ) == 0 ) assert out.x.subs(values).evalf() == pytest.approx(3 / 20) assert out.y.subs(values).evalf() == pytest.approx(4 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_rhophi_z_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == rho * sympy.cos(phi) / t assert out.y == rho * sympy.sin(phi) / t assert out.z == z / t assert out.x.subs(values).evalf() == pytest.approx(5 / 20) assert out.y.subs(values).evalf() == pytest.approx(0 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_rhophi_z_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == rho * sympy.cos(phi) / sympy.sqrt( rho**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.y == rho * sympy.sin(phi) / sympy.sqrt( rho**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.z == z / sympy.sqrt( rho**2 + z**2 + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.x.subs(values).evalf() == pytest.approx(5 / 20) assert out.y.subs(values).evalf() == pytest.approx(0 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_rhophi_theta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == rho * sympy.cos(phi) / t assert out.y == rho * sympy.sin(phi) / t assert out.z.simplify() == rho * z / (t * sympy.sqrt(x**2 + y**2)) assert out.x.subs(values).evalf() == pytest.approx(5 / 20) assert out.y.subs(values).evalf() == pytest.approx(0 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_rhophi_theta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == rho * sympy.cos(phi) / sympy.sqrt( rho**2 / (-(z**2) / (x**2 + y**2 + z**2) + 1) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.y == rho * sympy.sin(phi) / sympy.sqrt( rho**2 / (-(z**2) / (x**2 + y**2 + z**2) + 1) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.z.simplify() == rho * z / sympy.sqrt( rho**2 * (x**2 + y**2 + z**2) + (x**2 + y**2) * sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.x.subs(values).evalf() == pytest.approx(5 / 20) assert out.y.subs(values).evalf() == pytest.approx(0 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_rhophi_eta_t(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == rho * sympy.cos(phi) / t assert out.y == rho * sympy.sin(phi) / t assert out.z == rho * z / (t * sympy.sqrt(x**2 + y**2)) assert out.x.subs(values).evalf() == pytest.approx(5 / 20) assert out.y.subs(values).evalf() == pytest.approx(0 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) def test_rhophi_eta_tau(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), temporal=vector.backends.sympy.TemporalSympyTau( sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) ), ) out = vec.to_beta3() assert isinstance(out, vector.backends.sympy.VectorSympy3D) assert type(vec.azimuthal) == type(out.azimuthal) # noqa: E721 assert type(vec.longitudinal) == type(out.longitudinal) # noqa: E721 assert out.x == rho * sympy.cos(phi) / sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.y == rho * sympy.sin(phi) / sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) assert out.z == rho * z / ( sympy.sqrt(x**2 + y**2) * sympy.sqrt( rho**2 * (0.5 + 0.5 * sympy.exp(-2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2)))) ** 2 * sympy.exp(2 * sympy.asinh(z / sympy.sqrt(x**2 + y**2))) + sympy.Abs(-(t**2) + x**2 + y**2 + z**2) ) ) assert out.x.subs(values).evalf() == pytest.approx(5 / 20) assert out.y.subs(values).evalf() == pytest.approx(0 / 20) assert out.z.subs(values).evalf() == pytest.approx(10 / 20) vector-1.6.3/tests/compute/sympy/planar/000077500000000000000000000000001503546127100203035ustar00rootroot00000000000000vector-1.6.3/tests/compute/sympy/planar/__init__.py000066400000000000000000000003641503546127100224170ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/sympy/planar/test_deltaphi.py000066400000000000000000000042071503546127100235110ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, nx, ny = sympy.symbols("x y nx ny") values = {x: 1, y: 2, nx: 3, ny: 4} def test_xy_xy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny)) assert ( v1.deltaphi(v2) == sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) assert v1.deltaphi(v2).subs(values).evalf() == pytest.approx( math.atan2(2, 1) - math.atan2(4, 3) ) def test_xy_rhophi(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(nx, ny) ) assert ( v1.deltaphi(v2) == sympy.Mod(-ny + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) assert v1.deltaphi(v2).subs(values).evalf() == pytest.approx(math.atan2(2, 1) - 4) def test_rhophi_xy(): v1 = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(x, y) ) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny)) assert ( v1.deltaphi(v2) == sympy.Mod(y - sympy.atan2(ny, nx) + sympy.pi, 2 * sympy.pi) - sympy.pi ) assert v1.deltaphi(v2).subs(values).evalf() == pytest.approx(2 - math.atan2(4, 3)) def test_rhophi_rhophi(): v1 = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(x, y) ) v2 = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(nx, ny) ) assert v1.deltaphi(v2) == sympy.Mod(-ny + y + sympy.pi, 2 * sympy.pi) - sympy.pi assert v1.deltaphi(v2).subs(values).evalf() == pytest.approx(2 - 4) vector-1.6.3/tests/compute/sympy/planar/test_phi.py000066400000000000000000000017421503546127100225000ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") values = {x: 3, y: 4} def test_xy(): vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert vec.phi == pytest.approx(sympy.atan2(y, x)) assert vec.phi.subs(values).evalf() == pytest.approx(math.atan2(4, 3)) def test_rhophi(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi( sympy.sqrt(x**2 + y**2), sympy.atan2(y, x) ) ) assert vec.phi == pytest.approx(sympy.atan2(y, x)) assert vec.phi.subs(values).evalf() == pytest.approx(math.atan2(4, 3)) vector-1.6.3/tests/compute/sympy/planar/test_rho.py000066400000000000000000000016251503546127100225100ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") values = {x: 3, y: 4} def test_xy(): vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert vec.rho == sympy.sqrt(x**2 + y**2) assert vec.rho.subs(values) == pytest.approx(5) def test_rhophi(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi( sympy.sqrt(x**2 + y**2), sympy.atan2(y, x) ) ) assert vec.rho == sympy.sqrt(x**2 + y**2) assert vec.rho.subs(values) == pytest.approx(5) vector-1.6.3/tests/compute/sympy/planar/test_rho2.py000066400000000000000000000016031503546127100225660ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") values = {x: 3, y: 4} def test_xy(): vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert vec.rho2 == x**2 + y**2 assert vec.rho2.subs(values) == pytest.approx(25) def test_rhophi(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi( sympy.sqrt(x**2 + y**2), sympy.atan2(y, x) ) ) assert vec.rho2 == x**2 + y**2 assert vec.rho2.subs(values) == pytest.approx(25) vector-1.6.3/tests/compute/sympy/planar/test_rotateZ.py000066400000000000000000000023431503546127100233460ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi = sympy.symbols("x y rho phi") values = {x: 1, y: 0, rho: 1, phi: 0} def test_xy(): vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert vec.rotateZ(1).x == x * sympy.cos(1) - y * sympy.sin(1) assert vec.rotateZ(1).y == x * sympy.sin(1) + y * sympy.cos(1) assert vec.rotateZ(0.1).x.subs(values) == pytest.approx(0.9950041652780258) assert vec.rotateZ(0.1).y.subs(values) == pytest.approx(0.09983341664682815) def test_rhophi(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi) ) assert vec.rotateZ(1).rho == rho assert vec.rotateZ(1).phi == sympy.Mod(phi + 1 + sympy.pi, 2 * sympy.pi) - sympy.pi assert vec.rotateZ(0.1).rho.subs(values) == pytest.approx(1) assert vec.rotateZ(0.1).phi.subs(values) == pytest.approx(0.1) vector-1.6.3/tests/compute/sympy/planar/test_x.py000066400000000000000000000015411503546127100221640ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") values = {x: 3, y: 4} def test_xy(): vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert vec.x == x assert vec.x.subs(values) == pytest.approx(3) def test_rhophi(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi( sympy.sqrt(x**2 + y**2), sympy.atan2(y, x) ) ) assert vec.x == x assert vec.x.subs(values) == pytest.approx(3) vector-1.6.3/tests/compute/sympy/planar/test_y.py000066400000000000000000000015411503546127100221650ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") values = {x: 3, y: 4} def test_xy(): vec = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert vec.y == y assert vec.y.subs(values) == pytest.approx(4) def test_rhophi(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi( sympy.sqrt(x**2 + y**2), sympy.atan2(y, x) ) ) assert vec.y == y assert vec.y.subs(values) == pytest.approx(4) vector-1.6.3/tests/compute/sympy/spatial/000077500000000000000000000000001503546127100204635ustar00rootroot00000000000000vector-1.6.3/tests/compute/sympy/spatial/__init__.py000066400000000000000000000003641503546127100225770ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. vector-1.6.3/tests/compute/sympy/spatial/test_costheta.py000066400000000000000000000060301503546127100237050ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, theta = sympy.symbols("x y rho phi theta", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, theta: 0.4636476090008061} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ( sympy.sqrt(x**2 + y**2) / sympy.tan(theta) ), ) assert vec.costheta.simplify() == sympy.Abs( sympy.cos(theta) * sympy.tan(theta) ) / sympy.tan(theta) assert vec.costheta.subs(values).evalf() == pytest.approx( math.cos(0.4636476090008061) ) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta(theta), ) assert vec.costheta == sympy.cos(theta) assert vec.costheta.subs(values).evalf() == pytest.approx( math.cos(0.4636476090008061) ) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( -sympy.log(sympy.tan(0.5 * theta)) ), ) assert vec.costheta.simplify() == sympy.cos( 2.0 * sympy.atan(sympy.tan(0.5 * theta)) ) assert vec.costheta.subs(values).evalf() == pytest.approx( math.cos(0.4636476090008061) ) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(rho / sympy.tan(theta)), ) assert vec.costheta.simplify() == sympy.Abs( sympy.cos(theta) * sympy.tan(theta) ) / sympy.tan(theta) assert vec.costheta.subs(values).evalf() == pytest.approx( math.cos(0.4636476090008061) ) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta(theta), ) assert vec.costheta == sympy.cos(theta) assert vec.costheta.subs(values).evalf() == pytest.approx( math.cos(0.4636476090008061) ) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta( -sympy.log(sympy.tan(0.5 * theta)) ), ) assert vec.costheta.simplify() == sympy.cos( 2.0 * sympy.atan(sympy.tan(0.5 * theta)) ) assert vec.costheta.subs(values).evalf() == pytest.approx( math.cos(0.4636476090008061) ) vector-1.6.3/tests/compute/sympy/spatial/test_cottheta.py000066400000000000000000000055631503546127100237200ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, theta = sympy.symbols("x y rho phi theta", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, theta: 0.4636476090008061} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ( sympy.sqrt(x**2 + y**2) / sympy.tan(theta) ), ) assert vec.cottheta == 1 / sympy.tan(theta) assert vec.cottheta.subs(values).evalf() == pytest.approx( 1 / math.tan(0.4636476090008061) ) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta(theta), ) assert vec.cottheta == 1 / sympy.tan(theta) assert vec.cottheta.subs(values).evalf() == pytest.approx( 1 / math.tan(0.4636476090008061) ) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( -sympy.log(sympy.tan(0.5 * theta)) ), ) assert vec.cottheta == 1 / sympy.tan(2.0 * sympy.atan(sympy.tan(0.5 * theta))) assert vec.cottheta.subs(values).evalf() == pytest.approx( 1 / math.tan(0.4636476090008061) ) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(rho / sympy.tan(theta)), ) assert vec.cottheta == 1 / sympy.tan(theta) assert vec.cottheta.subs(values).evalf() == pytest.approx( 1 / math.tan(0.4636476090008061) ) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta(theta), ) assert vec.cottheta == 1 / sympy.tan(theta) assert vec.cottheta.subs(values).evalf() == pytest.approx( 1 / math.tan(0.4636476090008061) ) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta( -sympy.log(sympy.tan(0.5 * theta)) ), ) assert vec.cottheta == 1 / sympy.tan(2.0 * sympy.atan(sympy.tan(0.5 * theta))) assert vec.cottheta.subs(values).evalf() == pytest.approx( 1 / math.tan(0.4636476090008061) ) vector-1.6.3/tests/compute/sympy/spatial/test_cross.py000066400000000000000000000040671503546127100232340ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, nx, ny, nz = sympy.symbols("x y z nx ny nz", real=True) values = {x: 0.1, y: 0.2, z: 0.3, nx: 0.4, ny: 0.5, nz: 0.6} def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) with pytest.raises(TypeError): out = v1.to_Vector4D().cross(v2) out = v1.cross(v2) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) assert (out.x, out.y, out.z) == ( -ny * z + nz * y, nx * z - nz * x, -nx * y + ny * x, ) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.cross(transformed2) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance( out.longitudinal, vector.backends.sympy.LongitudinalSympyZ ) assert ( out.x.subs(values), out.y.subs(values), out.z.subs(values), ) == pytest.approx((-0.03, 0.06, -0.03)) vector-1.6.3/tests/compute/sympy/spatial/test_deltaR.py000066400000000000000000000063521503546127100233150ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, nx, ny, nz, t = sympy.symbols("x y z nx ny nz t", real=True) values = {x: 0.1, y: 0.2, z: 0.3, nx: 0.4, ny: 0.5, nz: 0.6} def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) assert v1.deltaR(v2) == sympy.sqrt( ( sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) ** 2 + ( -sympy.asinh(nz / sympy.sqrt(nx**2 + ny**2)) + sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ) ** 2 ) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaR(transformed2).subs(values) == pytest.approx( math.sqrt(0.116083865330319) ) def test_lorentz_sympy(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert v1.deltaR(v2) == sympy.sqrt( ( sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) ** 2 + ( -sympy.asinh(nz / sympy.sqrt(nx**2 + ny**2)) + sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ) ** 2 ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaR(tr2).subs(values) == pytest.approx( math.sqrt(0.116083865330319) ) vector-1.6.3/tests/compute/sympy/spatial/test_deltaR2.py000066400000000000000000000062511503546127100233750ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, nx, ny, nz, t = sympy.symbols("x y z nx ny nz t", real=True) values = {x: 0.1, y: 0.2, z: 0.3, nx: 0.4, ny: 0.5, nz: 0.6} def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) assert ( v1.deltaR2(v2) == ( sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) ** 2 + ( -sympy.asinh(nz / sympy.sqrt(nx**2 + ny**2)) + sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ) ** 2 ) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaR2(transformed2).subs(values) == pytest.approx( 0.116083865330319 ) def test_lorentz_sympy(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert ( v1.deltaR2(v2) == ( sympy.Mod(-sympy.atan2(ny, nx) + sympy.atan2(y, x) + sympy.pi, 2 * sympy.pi) - sympy.pi ) ** 2 + ( -sympy.asinh(nz / sympy.sqrt(nx**2 + ny**2)) + sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ) ** 2 ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaR2(tr2).subs(values) == pytest.approx(0.116083865330319) vector-1.6.3/tests/compute/sympy/spatial/test_deltaangle.py000066400000000000000000000055221503546127100242000ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, nx, ny, nz, t = sympy.symbols("x y z nx ny nz t", real=True, positive=True) values = {x: 0.1, y: 0.2, z: 0.3, nx: 0.4, ny: 0.5, nz: 0.6} def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) assert v1.deltaangle(v2) == sympy.acos( (nx * x + ny * y + nz * z) / (sympy.sqrt(nx**2 + ny**2 + nz**2) * sympy.sqrt(x**2 + y**2 + z**2)) ) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaangle(transformed2).subs(values) == pytest.approx( 0.225726 ) def test_lorentz_sympy(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert v1.deltaangle(v2) == sympy.acos( (nx * x + ny * y + nz * z) / (sympy.sqrt(nx**2 + ny**2 + nz**2) * sympy.sqrt(x**2 + y**2 + z**2)) ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaangle(tr2).subs(values) == pytest.approx(0.225726) vector-1.6.3/tests/compute/sympy/spatial/test_deltaeta.py000066400000000000000000000054231503546127100236630ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, nx, ny, nz, t = sympy.symbols("x y z nx ny nz t", real=True) values = {x: 0.1, y: 0.2, z: 0.3, nx: 0.4, ny: 0.5, nz: 0.6} def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) assert v1.deltaeta(v2) == -sympy.asinh( nz / sympy.sqrt(nx**2 + ny**2) ) + sympy.asinh(z / sympy.sqrt(x**2 + y**2)) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.deltaeta(transformed2).subs(values) == pytest.approx( 0.2674387219518324 ) def test_lorentz_sympy(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert v1.deltaeta(v2) == -sympy.asinh( nz / sympy.sqrt(nx**2 + ny**2) ) + sympy.asinh(z / sympy.sqrt(x**2 + y**2)) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.deltaeta(tr2).subs(values) == pytest.approx(0.2674387219518324) vector-1.6.3/tests/compute/sympy/spatial/test_eta.py000066400000000000000000000052541503546127100226530ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, eta = sympy.symbols("x y rho phi eta", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, eta: 1.4436354751788103} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ( sympy.sqrt(x**2 + y**2) * sympy.sinh(eta) ), ) assert vec.eta.simplify() == sympy.asinh(sympy.sinh(eta)) assert vec.eta.subs(values).evalf() == pytest.approx(1.4436354751788103) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( 2.0 * sympy.atan(sympy.exp(-eta)) ), ) assert vec.eta.simplify() == -sympy.log( sympy.tan(1.0 * sympy.atan(sympy.exp(-eta))) ) assert vec.eta.subs(values).evalf() == pytest.approx(1.4436354751788103) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta(eta), ) assert vec.eta == eta assert vec.eta.subs(values).evalf() == pytest.approx(1.4436354751788103) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(rho * sympy.sinh(eta)), ) assert vec.eta.simplify() == sympy.asinh(sympy.sinh(eta)) assert vec.eta.subs(values).evalf() == pytest.approx(1.4436354751788103) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( 2.0 * sympy.atan(sympy.exp(-eta)) ), ) assert vec.eta.simplify() == -sympy.log( sympy.tan(1.0 * sympy.atan(sympy.exp(-eta))) ) assert vec.eta.subs(values).evalf() == pytest.approx(1.4436354751788103) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(eta), ) assert vec.eta == eta assert vec.eta.subs(values).evalf() == pytest.approx(1.4436354751788103) vector-1.6.3/tests/compute/sympy/spatial/test_mag.py000066400000000000000000000053101503546127100226370ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import math import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z = sympy.symbols("x y rho phi z", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert vec.mag == sympy.sqrt(x**2 + y**2 + z**2) assert vec.mag.subs(values).evalf() == pytest.approx(math.sqrt(125)) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), ) assert vec.mag.simplify() == sympy.sqrt(x**2 + y**2 + z**2) assert vec.mag.subs(values).evalf() == pytest.approx(math.sqrt(125)) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), ) assert vec.mag.simplify() == 1.0 * sympy.sqrt(x**2 + y**2) * sympy.sqrt( z**2 / (x**2 + y**2) + 1 ) assert vec.mag.subs(values).evalf() == pytest.approx(math.sqrt(125)) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert vec.mag == sympy.sqrt(rho**2 + z**2) assert vec.mag.subs(values).evalf() == pytest.approx(math.sqrt(125)) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), ) assert vec.mag.simplify() == sympy.sqrt(rho**2 + z**2) assert vec.mag.subs(values).evalf() == pytest.approx(math.sqrt(125)) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), ) assert vec.mag.simplify() == 1.0 * rho * sympy.sqrt(1 + z**2 / rho**2) assert vec.mag.subs(values).evalf() == pytest.approx(math.sqrt(125)) vector-1.6.3/tests/compute/sympy/spatial/test_mag2.py000066400000000000000000000052761503546127100227340ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z = sympy.symbols("x y rho phi z", real=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert vec.mag2 == x**2 + y**2 + z**2 assert vec.mag2.subs(values).evalf() == pytest.approx(125) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), ) assert vec.mag2.simplify() == x**2 + y**2 + z**2 assert vec.mag2.subs(values).evalf() == pytest.approx(125) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), ) # TODO: why won't sympy equate the expressions without double # simplifying? assert ( sympy.simplify( vec.mag2.simplify() - 0.25 * (x**2 + y**2) * (4 * z**2 / (x**2 + y**2) + 4) ) == 0 ) assert vec.mag2.subs(values).evalf() == pytest.approx(125) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert vec.mag2 == rho**2 + z**2 assert vec.mag2.subs(values).evalf() == pytest.approx(125) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), ) assert vec.mag2.simplify() == rho**2 + z**2 assert vec.mag2.subs(values).evalf() == pytest.approx(125) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), ) assert vec.mag2.simplify() == 0.25 * rho**2 * (4 + 4 * z**2 / rho**2) assert vec.mag2.subs(values).evalf() == pytest.approx(125) vector-1.6.3/tests/compute/sympy/spatial/test_rotateX.py000066400000000000000000000056171503546127100235330ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector import vector._methods sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, time = sympy.symbols("x y z time", real=True) values = {x: 0.1, y: 0.2, z: 0.3} def test_spatial_sympy(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) out = vec.rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x == x # cannot equate the sympy expressions because of floating point errors assert out.y.subs(values) == pytest.approx(0.1195612965657721) assert out.z.subs(values) == pytest.approx(0.340154518364098) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": out = getattr(vec, "to_" + t)().rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x.subs(values) == pytest.approx(0.1) assert out.y.subs(values) == pytest.approx(0.1195612965657721) assert out.z.subs(values) == pytest.approx(0.340154518364098) def test_lorentz_sympy(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(time), ) out = vec.rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x == x # cannot equate the sympy expressions because of floating point errors assert out.y.subs(values) == pytest.approx(0.1195612965657721) assert out.z.subs(values) == pytest.approx(0.340154518364098) for t in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): out = getattr(vec, "to_" + t)().rotateX(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x.subs(values) == pytest.approx(0.1) assert out.y.subs(values) == pytest.approx(0.1195612965657721) assert out.z.subs(values) == pytest.approx(0.340154518364098) vector-1.6.3/tests/compute/sympy/spatial/test_rotateY.py000066400000000000000000000056271503546127100235350ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector import vector._methods sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, time = sympy.symbols("x y z time", real=True) values = {x: 0.1, y: 0.2, z: 0.3} def test_spatial_sympy(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) out = vec.rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values) == pytest.approx(0.17111242994742137) assert out.y == y assert out.z.subs(values) == pytest.approx(0.2659333305877411) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": out = getattr(vec, "to_" + t)().rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert out.x.subs(values) == pytest.approx(0.17111242994742137) assert out.y.subs(values) == pytest.approx(0.2) assert out.z.subs(values) == pytest.approx(0.2659333305877411) def test_lorentz_sympy(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(time), ) out = vec.rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values) == pytest.approx(0.17111242994742137) assert out.y == y assert out.z.subs(values) == pytest.approx(0.2659333305877411) for t in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): out = getattr(vec, "to_" + t)().rotateY(0.25) assert isinstance(out.azimuthal, vector._methods.AzimuthalXY) assert isinstance(out.longitudinal, vector._methods.LongitudinalZ) assert hasattr(out, "temporal") assert out.x.subs(values) == pytest.approx(0.17111242994742137) assert out.y.subs(values) == pytest.approx(0.2) assert out.z.subs(values) == pytest.approx(0.2659333305877411) vector-1.6.3/tests/compute/sympy/spatial/test_rotate_axis.py000066400000000000000000000104471503546127100244240ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, nx, ny, nz, t = sympy.symbols("x y z nx ny nz t", real=True) values = {x: 0.1, y: 0.2, z: 0.3, nx: 0.4, ny: 0.5, nz: 0.6} def test_spatial_sympy(): axis = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) out = vec.rotate_axis(axis, 0.25) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) assert out.x.subs(values) == pytest.approx(0.374834254043358) assert out.y.subs(values) == pytest.approx(0.5383405688588193) assert out.z.subs(values) == pytest.approx(0.5828282027463345) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": taxis, tvec = ( getattr(axis, "to_" + t1)(), getattr(vec, "to_" + t2)(), ) out = tvec.rotate_axis(taxis, 0.25) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance( out.longitudinal, vector.backends.sympy.LongitudinalSympyZ ) assert out.x.subs(values) == pytest.approx(0.374834254043358) assert out.y.subs(values) == pytest.approx(0.5383405688588193) assert out.z.subs(values) == pytest.approx(0.5828282027463345) def test_lorentz_sympy(): axis = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(t), ) with pytest.raises(TypeError): out = vec.rotate_axis(axis, 0.25) out = vec.rotate_axis(axis.to_Vector3D(), 0.25) assert isinstance(out, vector.VectorSympy4D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) assert hasattr(out, "temporal") assert out.x.subs(values) == pytest.approx(0.374834254043358) assert out.y.subs(values) == pytest.approx(0.5383405688588193) assert out.z.subs(values) == pytest.approx(0.5828282027463345) for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): taxis, tvec = ( getattr(axis, "to_" + t1)(), getattr(vec, "to_" + t2)(), ) out = tvec.rotate_axis(taxis, 0.25) assert isinstance(out, vector.VectorSympy4D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance( out.longitudinal, vector.backends.sympy.LongitudinalSympyZ ) assert hasattr(out, "temporal") assert out.x.subs(values) == pytest.approx(0.374834254043358) assert out.y.subs(values) == pytest.approx(0.5383405688588193) assert out.z.subs(values) == pytest.approx(0.5828282027463345) vector-1.6.3/tests/compute/sympy/spatial/test_rotate_euler.py000066400000000000000000000035031503546127100245670ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z = sympy.symbols("x y z", real=True) values = {x: 0.4, y: 0.5, z: 0.6} def test_spatial_sympy(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) out = vec.rotate_euler(0.1, 0.2, 0.3) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values) == pytest.approx(0.5956646364506655) assert out.y.subs(values) == pytest.approx(0.409927258162962) assert out.z.subs(values) == pytest.approx(0.4971350761081869) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tvec = getattr(vec, "to_" + t)() out = tvec.rotate_euler(0.1, 0.2, 0.3) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values) == pytest.approx(0.5956646364506655) assert out.y.subs(values) == pytest.approx(0.409927258162962) assert out.z.subs(values) == pytest.approx(0.4971350761081869) vector-1.6.3/tests/compute/sympy/spatial/test_rotate_quaternion.py000066400000000000000000000034111503546127100256360ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z = sympy.symbols("x y z", real=True) values = {x: 0.5, y: 0.6, z: 0.7} def test_spatial_sympy(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) out = vec.rotate_quaternion(0.1, 0.2, 0.3, 0.4) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values) == pytest.approx(0.078) assert out.y.subs(values) == pytest.approx(0.18) assert out.z.subs(values) == pytest.approx(0.246) for t in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tvec = getattr(vec, "to_" + t)() out = tvec.rotate_quaternion(0.1, 0.2, 0.3, 0.4) assert isinstance(out, vector.VectorSympy3D) assert isinstance(out.azimuthal, vector.backends.sympy.AzimuthalSympyXY) assert isinstance(out.longitudinal, vector.backends.sympy.LongitudinalSympyZ) # cannot equate the sympy expressions because of floating point errors assert out.x.subs(values) == pytest.approx(0.078) assert out.y.subs(values) == pytest.approx(0.18) assert out.z.subs(values) == pytest.approx(0.246) vector-1.6.3/tests/compute/sympy/spatial/test_theta.py000066400000000000000000000054471503546127100232130ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, theta = sympy.symbols("x y rho phi theta", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, theta: 0.4636476090008061} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ( sympy.sqrt(x**2 + y**2) / sympy.tan(theta) ), ) assert vec.theta.simplify() == sympy.acos( sympy.Abs(sympy.cos(theta) * sympy.tan(theta)) / sympy.tan(theta) ) assert vec.theta.subs(values).evalf() == pytest.approx(0.4636476090008061) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta(theta), ) assert vec.theta == theta assert vec.theta.subs(values).evalf() == pytest.approx(0.4636476090008061) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( -sympy.log(sympy.tan(0.5 * theta)) ), ) assert vec.theta.simplify() == 2.0 * sympy.atan(sympy.tan(0.5 * theta)) assert vec.theta.subs(values).evalf() == pytest.approx(0.4636476090008061) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(rho / sympy.tan(theta)), ) assert vec.theta.simplify() == sympy.acos( sympy.Abs(sympy.cos(theta) * sympy.tan(theta)) / sympy.tan(theta) ) assert vec.theta.subs(values).evalf() == pytest.approx(0.4636476090008061) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta(theta), ) assert vec.theta == theta assert vec.theta.subs(values).evalf() == pytest.approx(0.4636476090008061) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta( -sympy.log(sympy.tan(0.5 * theta)) ), ) assert vec.theta.simplify() == 2.0 * sympy.atan(sympy.tan(0.5 * theta)) assert vec.theta.subs(values).evalf() == pytest.approx(0.4636476090008061) vector-1.6.3/tests/compute/sympy/spatial/test_z.py000066400000000000000000000045701503546127100223530ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, rho, phi, z = sympy.symbols("x y rho phi z", real=True, positive=True) values = {x: 3, y: 4, rho: 5, phi: 0, z: 10} def test_xy_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert vec.z == z assert vec.z.subs(values).evalf() == pytest.approx(10) def test_xy_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) ), ) assert vec.z.simplify() == z assert vec.z.subs(values).evalf() == pytest.approx(10) def test_xy_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyEta( sympy.asinh(z / sympy.sqrt(x**2 + y**2)) ), ) assert vec.z.simplify() == z assert vec.z.subs(values).evalf() == pytest.approx(10) def test_rhophi_z(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert vec.z == z assert vec.z.subs(values).evalf() == pytest.approx(10) def test_rhophi_theta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyTheta( sympy.acos(z / sympy.sqrt(rho**2 + z**2)) ), ) assert vec.z.simplify() == z assert vec.z.subs(values).evalf() == pytest.approx(10) def test_rhophi_eta(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyRhoPhi(rho, phi), longitudinal=vector.backends.sympy.LongitudinalSympyEta(sympy.asinh(z / rho)), ) assert vec.z.simplify() == z assert vec.z.subs(values).evalf() == pytest.approx(10) vector-1.6.3/tests/compute/sympy/test_add.py000066400000000000000000000076511503546127100212000ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, nx, ny, nz, nt = sympy.symbols("x y z t nx ny nz nt", real=True) values = {x: 3, y: 4, z: 2, t: 10, nx: 5, ny: 12, nz: 4, nt: 15} def test_planar_object(): v1 = vector.backends.sympy.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y) ) v2 = vector.backends.sympy.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny) ) out = v1.add(v2) assert out.x == x + nx assert out.y == y + ny for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x.subs(values).evalf() == pytest.approx(8) assert out.y.subs(values).evalf() == pytest.approx(16) def test_spatial_object(): v1 = vector.backends.sympy.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.backends.sympy.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) out = v1.add(v2) assert out.x == x + nx assert out.y == y + ny assert out.z == z + nz for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x.subs(values).evalf() == pytest.approx(8) assert out.y.subs(values).evalf() == pytest.approx(16) assert out.z.subs(values).evalf() == pytest.approx(6) def test_lorentz_object(): v1 = vector.backends.sympy.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.backends.sympy.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(nt), ) out = v1.add(v2) assert out.x == x + nx assert out.y == y + ny assert out.z == z + nz assert out.t == t + nt for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x.subs(values).evalf() == pytest.approx(8) assert out.y.subs(values).evalf() == pytest.approx(16) assert out.z.subs(values).evalf() == pytest.approx(6) assert out.t.subs(values).evalf() == pytest.approx(25) vector-1.6.3/tests/compute/sympy/test_conversions.py000066400000000000000000000507751503546127100230250ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y = sympy.symbols("x y") rho, phi = sympy.symbols("rho phi") z, eta, theta = sympy.symbols("z eta theta") t, tau = sympy.symbols("t tau") px, py = sympy.symbols("px py") pt = sympy.symbols("pt") pz = sympy.symbols("pz") M, E = sympy.symbols("M E") def test_conversion_2D(): v = vector.VectorSympy2D(x=x, y=y) tv = v.to_Vector2D() assert isinstance(tv, vector.VectorSympy2D) assert tv.x == x assert tv.y == y tv = v.to_Vector3D() assert isinstance(tv, vector.VectorSympy3D) assert tv.x == x assert tv.y == y assert tv.z == 0.0 tv = v.to_Vector4D() assert isinstance(tv, vector.VectorSympy4D) assert tv.x == x assert tv.y == y assert tv.z == 0.0 assert tv.t == 0.0 for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.VectorSympy2D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.VectorSympy3D) assert tv.x == x assert tv.y == y assert getattr(tv, longitudinal) == 0.0 for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.VectorSympy4D) assert tv.x == x assert tv.y == y assert getattr(tv, longitudinal) == 0.0 assert getattr(tv, temporal) == 0.0 def test_momentum_conversion_2D(): v = vector.MomentumSympy2D(px=px, py=py) tv = v.to_Vector2D() assert isinstance(tv, vector.MomentumSympy2D) assert tv.x == px assert tv.y == py tv = v.to_Vector3D() assert isinstance(tv, vector.MomentumSympy3D) assert tv.x == px assert tv.y == py assert tv.z == 0.0 tv = v.to_Vector4D() assert isinstance(tv, vector.MomentumSympy4D) assert tv.x == px assert tv.y == py assert tv.z == 0.0 assert tv.t == 0.0 for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.MomentumSympy2D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.MomentumSympy3D) assert tv.x == px assert tv.y == py assert getattr(tv, longitudinal) == 0.0 for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.MomentumSympy4D) assert tv.x == px assert tv.y == py assert getattr(tv, longitudinal) == 0.0 assert getattr(tv, temporal) == 0.0 def test_conversion_3D(): v = vector.VectorSympy3D(x=x, y=y, z=z) tv = v.to_Vector2D() assert isinstance(tv, vector.VectorSympy2D) assert tv.x == x assert tv.y == y tv = v.to_Vector3D() assert isinstance(tv, vector.VectorSympy3D) assert tv.x == x assert tv.y == y assert tv.z == z tv = v.to_Vector4D() assert isinstance(tv, vector.VectorSympy4D) assert tv.x == x assert tv.y == y assert tv.z == z assert tv.t == 0.0 for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.VectorSympy2D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.VectorSympy3D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) if longitudinal == "z": assert tv.z == z elif longitudinal == "eta": assert tv.eta == sympy.asinh(z / sympy.sqrt(x**2 + y**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.VectorSympy4D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) if longitudinal == "z": assert tv.z == z elif longitudinal == "eta": assert tv.eta == sympy.asinh(z / sympy.sqrt(x**2 + y**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) assert getattr(tv, temporal) == pytest.approx(0) def test_momentum_conversion_3D(): v = vector.MomentumSympy3D(px=px, py=py, pz=pz) tv = v.to_Vector2D() assert isinstance(tv, vector.MomentumSympy2D) assert tv.x == px assert tv.y == py tv = v.to_Vector3D() assert isinstance(tv, vector.MomentumSympy3D) assert tv.x == px assert tv.y == py assert tv.z == pz tv = v.to_Vector4D() assert isinstance(tv, vector.MomentumSympy4D) assert tv.x == px assert tv.y == py assert tv.z == pz assert tv.t == 0.0 for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.MomentumSympy2D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.MomentumSympy3D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) if longitudinal == "z": assert tv.z == pz elif longitudinal == "eta": assert tv.eta == sympy.asinh(pz / sympy.sqrt(px**2 + py**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos(pz / sympy.sqrt(px**2 + py**2 + pz**2)) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.MomentumSympy4D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) if longitudinal == "z": assert tv.z == pz elif longitudinal == "eta": assert tv.eta == sympy.asinh(pz / sympy.sqrt(px**2 + py**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos( pz / sympy.sqrt(px**2 + py**2 + pz**2) ) assert getattr(tv, temporal) == pytest.approx(0) def test_conversion_4D(): v = vector.VectorSympy4D(x=x, y=y, z=z, t=t) tv = v.to_Vector2D() assert isinstance(tv, vector.VectorSympy2D) assert tv.x == x assert tv.y == y tv = v.to_Vector3D() assert isinstance(tv, vector.VectorSympy3D) assert tv.x == x assert tv.y == y assert tv.z == z tv = v.to_Vector4D() assert isinstance(tv, vector.VectorSympy4D) assert tv.x == x assert tv.y == y assert tv.z == z assert tv.t == t for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.VectorSympy2D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.VectorSympy3D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) if longitudinal == "z": assert tv.z == z elif longitudinal == "eta": assert tv.eta == sympy.asinh(z / sympy.sqrt(x**2 + y**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.VectorSympy4D) if azimuthal == "xy": assert tv.x == x assert tv.y == y elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(x**2 + y**2) assert tv.phi == sympy.atan2(y, x) if longitudinal == "z": assert tv.z == z elif longitudinal == "eta": assert tv.eta == sympy.asinh(z / sympy.sqrt(x**2 + y**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos(z / sympy.sqrt(x**2 + y**2 + z**2)) if temporal == "t": assert tv.t == t elif temporal == "tau": assert tv.tau == sympy.sqrt(sympy.Abs(-(t**2) + x**2 + y**2 + z**2)) def test_momentum_conversion_4D(): v = vector.MomentumSympy4D(px=px, py=py, pz=pz, E=E) tv = v.to_Vector2D() assert isinstance(tv, vector.MomentumSympy2D) assert tv.x == px assert tv.y == py tv = v.to_Vector3D() assert isinstance(tv, vector.MomentumSympy3D) assert tv.x == px assert tv.y == py assert tv.z == pz tv = v.to_Vector4D() assert isinstance(tv, vector.MomentumSympy4D) assert tv.x == px assert tv.y == py assert tv.z == pz assert tv.t == E for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.MomentumSympy2D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.MomentumSympy3D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) if longitudinal == "z": assert tv.z == pz elif longitudinal == "eta": assert tv.eta == sympy.asinh(pz / sympy.sqrt(px**2 + py**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos(pz / sympy.sqrt(px**2 + py**2 + pz**2)) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.MomentumSympy4D) if azimuthal == "xy": assert tv.x == px assert tv.y == py elif azimuthal == "rhophi": assert tv.rho == sympy.sqrt(px**2 + py**2) assert tv.phi == sympy.atan2(py, px) if longitudinal == "z": assert tv.z == pz elif longitudinal == "eta": assert tv.eta == sympy.asinh(pz / sympy.sqrt(px**2 + py**2)) elif longitudinal == "theta": assert tv.theta == sympy.acos( pz / sympy.sqrt(px**2 + py**2 + pz**2) ) if temporal == "t": assert tv.t == E elif temporal == "tau": assert tv.tau == sympy.sqrt( sympy.Abs(-(E**2) + px**2 + py**2 + pz**2) ) def test_conversion_with_coords(): # 2D -> 3D vec = vector.VectorSympy2D(x=x, y=y) assert vec.to_Vector3D(z=z).z == z assert vec.to_Vector3D(eta=eta).eta == eta assert vec.to_Vector3D(theta=theta).theta == theta # test alias assert vec.to_3D(z=z).x == vec.x assert vec.to_3D(z=z).y == vec.y # 2D -> 4D assert vec.to_Vector4D(z=z, t=t).z == z assert vec.to_Vector4D(z=z, t=t).t == t assert vec.to_Vector4D(eta=eta, t=t).eta == eta assert vec.to_Vector4D(eta=eta, t=t).t == t assert vec.to_Vector4D(theta=theta, t=t).theta == theta assert vec.to_Vector4D(theta=theta, t=t).t == t assert vec.to_Vector4D(z=z, tau=tau).z == z assert vec.to_Vector4D(z=z, tau=tau).tau == tau assert vec.to_Vector4D(eta=eta, tau=tau).eta == eta assert vec.to_Vector4D(eta=eta, tau=tau).tau == tau assert vec.to_Vector4D(theta=theta, tau=tau).theta == theta assert vec.to_Vector4D(theta=theta, tau=tau).tau == tau # test alias assert vec.to_4D(z=z, t=t).x == vec.x assert vec.to_4D(z=z, t=t).y == vec.y # 3D -> 4D vec = vector.VectorSympy3D(x=px, y=py, z=pz) # test alias assert vec.to_4D(t=t).t == t assert vec.to_4D(tau=tau).tau == tau assert vec.to_Vector4D(t=t).x == vec.x assert vec.to_Vector4D(t=t).y == vec.y assert vec.to_Vector4D(t=t).z == vec.z # check if momentum coords work vec = vector.MomentumSympy2D(px=px, py=py) assert vec.to_Vector3D(pz=pz).pz == pz # test both alias and original methods assert vec.to_4D(pz=pz, m=M).pz == pz assert vec.to_4D(pz=pz, m=M).m == M assert vec.to_4D(pz=pz, mass=M).mass == M assert vec.to_4D(pz=pz, M=M).M == M assert vec.to_Vector4D(pz=pz, e=E).e == E assert vec.to_Vector4D(pz=pz, energy=E).energy == E assert vec.to_Vector4D(pz=pz, E=E).E == E vec = vector.MomentumSympy3D(px=px, py=py, pz=pz) # test both alias and original methods assert vec.to_4D(m=M).m == M assert vec.to_4D(mass=M).mass == M assert vec.to_4D(M=M).M == M assert vec.to_Vector4D(e=E).e == E assert vec.to_Vector4D(energy=E).energy == E assert vec.to_Vector4D(E=E).E == E def test_like(): v1 = vector.VectorSympy2D(x=x, y=y) v2 = vector.VectorSympy3D(x=x, y=y, z=z) v3 = vector.VectorSympy4D(x=x, y=y, z=z, t=t) # 2D + 3D.like(2D) = 2D assert v1 + v2.like(v1) == vector.VectorSympy2D(x=2 * x, y=2 * y) assert v2.like(v1) + v1 == vector.VectorSympy2D(x=2 * x, y=2 * y) # 2D + 4D.like(2D) = 2D assert v1 + v3.like(v1) == vector.VectorSympy2D(x=2 * x, y=2 * y) assert v3.like(v1) + v1 == vector.VectorSympy2D(x=2 * x, y=2 * y) # 3D + 2D.like(3D) = 3D assert v2 + v1.like(v2) == vector.VectorSympy3D(x=2 * x, y=2 * y, z=z) assert v1.like(v2) + v2 == vector.VectorSympy3D(x=2 * x, y=2 * y, z=z) # 3D + 4D.like(3D) = 3D assert v2 + v3.like(v2) == vector.VectorSympy3D(x=2 * x, y=2 * y, z=2 * z) assert v3.like(v2) + v2 == vector.VectorSympy3D(x=2 * x, y=2 * y, z=2 * z) # 4D + 2D.like(4D) = 4D assert v3 + v1.like(v3) == vector.VectorSympy4D(x=2 * x, y=2 * y, z=z, t=t) assert v1.like(v3) + v3 == vector.VectorSympy4D(x=2 * x, y=2 * y, z=z, t=t) # 4D + 3D.like(4D) = 4D assert v3 + v2.like(v3) == vector.VectorSympy4D(x=2 * x, y=2 * y, z=2 * z, t=t) assert v2.like(v3) + v3 == vector.VectorSympy4D(x=2 * x, y=2 * y, z=2 * z, t=t) v1 = vector.MomentumSympy2D(px=px, py=py) v2 = vector.MomentumSympy3D(px=px, py=py, pz=pz) v3 = vector.MomentumSympy4D(px=px, py=py, pz=pz, E=E) # order should not matter # 2D + 3D.like(2D) = 2D assert v1 + v2.like(v1) == vector.MomentumSympy2D(px=2 * px, py=2 * py) assert v2.like(v1) + v1 == vector.MomentumSympy2D(px=2 * px, py=2 * py) # 2D + 4D.like(2D) = 2D assert v1 + v3.like(v1) == vector.MomentumSympy2D(px=2 * px, py=2 * py) assert v3.like(v1) + v1 == vector.MomentumSympy2D(px=2 * px, py=2 * py) # 3D + 2D.like(3D) = 3D assert v2 + v1.like(v2) == vector.MomentumSympy3D(px=2 * px, py=2 * py, pz=pz) assert v1.like(v2) + v2 == vector.MomentumSympy3D(px=2 * px, py=2 * py, pz=pz) # 3D + 4D.like(3D) = 3D assert v2 + v3.like(v2) == vector.MomentumSympy3D(px=2 * px, py=2 * py, pz=2 * pz) assert v3.like(v2) + v2 == vector.MomentumSympy3D(px=2 * px, py=2 * py, pz=2 * pz) # 4D + 2D.like(4D) = 4D assert v3 + v1.like(v3) == vector.MomentumSympy4D(px=2 * px, py=2 * py, pz=pz, E=E) assert v1.like(v3) + v3 == vector.MomentumSympy4D(px=2 * px, py=2 * py, pz=pz, E=E) # 4D + 3D.like(4D) = 4D assert v3 + v2.like(v3) == vector.MomentumSympy4D( px=2 * px, py=2 * py, pz=2 * pz, E=E ) assert v2.like(v3) + v3 == vector.MomentumSympy4D( px=2 * px, py=2 * py, pz=2 * pz, E=E ) def test_momentum_preservation(): v1 = vector.MomentumSympy2D(px=px, py=py) v2 = vector.VectorSympy3D(x=x, y=y, z=z) v3 = vector.MomentumSympy4D(px=px, py=py, pz=pz, t=t) # momentum + generic = momentum # 2D + 3D.like(2D) = 2D assert isinstance(v1 + v2.like(v1), vector.MomentumSympy2D) assert isinstance(v2.like(v1) + v1, vector.MomentumSympy2D) # 2D + 4D.like(2D) = 2D assert isinstance(v1 + v3.like(v1), vector.MomentumSympy2D) assert isinstance(v3.like(v1) + v1, vector.MomentumSympy2D) # 3D + 2D.like(3D) = 3D assert isinstance(v2 + v1.like(v2), vector.MomentumSympy3D) assert isinstance(v1.like(v2) + v2, vector.MomentumSympy3D) # 3D + 4D.like(3D) = 3D assert isinstance(v2 + v3.like(v2), vector.MomentumSympy3D) assert isinstance(v3.like(v2) + v2, vector.MomentumSympy3D) # 4D + 2D.like(4D) = 4D assert isinstance(v3 + v1.like(v3), vector.MomentumSympy4D) assert isinstance(v1.like(v3) + v3, vector.MomentumSympy4D) # 4D + 3D.like(4D) = 4D assert isinstance(v3 + v2.like(v3), vector.MomentumSympy4D) assert isinstance(v2.like(v3) + v3, vector.MomentumSympy4D) def test_momentum_coordinate_transforms(): vec = vector.MomentumSympy2D(px=px, py=py) for t1 in "pxpy", "ptphi": for t2 in "pz", "eta", "theta": for t3 in "mass", "energy": transformed_object = getattr(vec, "to_" + t1)() assert isinstance(transformed_object, vector.MomentumSympy2D) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) transformed_object = getattr(vec, "to_" + t1 + t2)() assert isinstance(transformed_object, vector.MomentumSympy3D) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) assert hasattr(transformed_object, t2) transformed_object = getattr(vec, "to_" + t1 + t2 + t3)() assert isinstance(transformed_object, vector.MomentumSympy4D) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) assert hasattr(transformed_object, t2) assert hasattr(transformed_object, t3) vector-1.6.3/tests/compute/sympy/test_dot.py000066400000000000000000000062511503546127100212310ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, nx, ny, nz, nt = sympy.symbols("x y z t nx ny nz nt", real=True) values = {x: 1, y: 2, z: 3, t: 10, nx: 5, ny: 6, nz: 7, nt: 20} def test_planar_object(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny)) assert v1.dot(v2) == nx * x + ny * y for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.dot(transformed2).subs(values).evalf() == pytest.approx( 17 ) def test_spatial_object(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) assert v1.dot(v2) == nx * x + ny * y + nz * z for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.dot(transformed2).subs(values).evalf() == pytest.approx( 38 ) def test_lorentz_object(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(nt), ) assert v1.dot(v2) == nt * t - nx * x - ny * y - nz * z for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.dot(tr2).subs(values).evalf() == pytest.approx(162) vector-1.6.3/tests/compute/sympy/test_equal.py000066400000000000000000000054361503546127100215560ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t") def test_planar_sympy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert v1.equal(v2) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) transformed1.equal(transformed2) def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert v1.equal(v2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) transformed1.equal(transformed2) def test_lorentz_sympy(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert v1.equal(v2) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() tr1.equal(tr2) vector-1.6.3/tests/compute/sympy/test_is_antiparallel.py000066400000000000000000000050411503546127100236020ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, nx, ny, mx, my, z, nz, mz = sympy.symbols("x y nx ny mx my z nz mz", real=True) values = { x: 0.1, y: 0.2, nx: -0.3, ny: -0.6, mx: 0.3, my: 0.6, z: 0.3, nz: -0.9, mz: 0.9, } def test_planar_sympy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny)) v3 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(mx, my)) assert v1.is_antiparallel(-v1).subs(values) assert v2.is_antiparallel(-v2).subs(values) assert v3.is_antiparallel(-v3).subs(values) assert v1.is_antiparallel(v2).subs(values) assert v2.is_antiparallel(v3).subs(values) assert not v1.is_antiparallel(v3).subs(values) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_antiparallel(tr2).subs(values) def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) v3 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(mx, my), longitudinal=vector.backends.sympy.LongitudinalSympyZ(mz), ) assert v1.is_antiparallel(-v1).subs(values) assert v2.is_antiparallel(-v2).subs(values) assert v3.is_antiparallel(-v3).subs(values) assert v1.is_antiparallel(v2).subs(values) assert v2.is_antiparallel(v3).subs(values) assert not v1.is_antiparallel(v3).subs(values) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_antiparallel(tr2).subs(values) vector-1.6.3/tests/compute/sympy/test_is_parallel.py000066400000000000000000000047011503546127100227300ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, nx, ny, mx, my, z, nz, mz = sympy.symbols("x y nx ny mx my z nz mz", real=True) values = {x: 0.1, y: 0.2, nx: 0.3, ny: 0.6, mx: 0.1, my: 0.5, z: 0.3, nz: 0.9, mz: 0.9} def test_planar_sympy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny)) v3 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(mx, my)) assert v1.is_parallel(v1).subs(values) assert v2.is_parallel(v2).subs(values) assert v3.is_parallel(v3).subs(values) assert v1.is_parallel(v2).subs(values) assert not v1.is_parallel(v3).subs(values) assert not v2.is_parallel(v3).subs(values) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_parallel(tr2).subs(values) def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) v3 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(mx, my), longitudinal=vector.backends.sympy.LongitudinalSympyZ(mz), ) assert v1.is_parallel(v1).subs(values) assert v2.is_parallel(v2).subs(values) assert v3.is_parallel(v3).subs(values) assert v1.is_parallel(v2).subs(values) assert not v1.is_parallel(v3).subs(values) assert not v2.is_parallel(v3).subs(values) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_parallel(tr2).subs(values) vector-1.6.3/tests/compute/sympy/test_is_perpendicular.py000066400000000000000000000050411503546127100237670ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, nx, ny, mx, my, z, nz, mz = sympy.symbols("x y nx ny mx my z nz mz", real=True) values = {x: 0, y: 1, nx: 1, ny: 0, mx: 0.1, my: 0.5, z: 1, nz: 0, mz: 0.9} def test_planar_sympy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(0, 1)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(1, 0)) v3 = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(0.1, 0.5) ) assert not v1.is_perpendicular(v1).subs(values) assert not v2.is_perpendicular(v2).subs(values) assert not v3.is_perpendicular(v3).subs(values) assert v1.is_perpendicular(v2).subs(values) assert not v1.is_perpendicular(v3).subs(values) assert not v2.is_perpendicular(v3).subs(values) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_perpendicular(tr2).subs(values) def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(0, 1), longitudinal=vector.backends.sympy.LongitudinalSympyZ(1), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(1, 0), longitudinal=vector.backends.sympy.LongitudinalSympyZ(0), ) v3 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(0.1, 0.5), longitudinal=vector.backends.sympy.LongitudinalSympyZ(0.9), ) assert not v1.is_perpendicular(v1).subs(values) assert not v2.is_perpendicular(v2).subs(values) assert not v3.is_perpendicular(v3).subs(values) assert v1.is_perpendicular(v2).subs(values) assert not v1.is_perpendicular(v3).subs(values) assert not v2.is_perpendicular(v3).subs(values) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_perpendicular(tr2).subs(values) vector-1.6.3/tests/compute/sympy/test_not_equal.py000066400000000000000000000055021503546127100224300ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t") def test_planar_sympy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) assert not v1.not_equal(v2) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) transformed1.not_equal(transformed2) def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) assert not v1.not_equal(v2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) transformed1.not_equal(transformed2) def test_lorentz_sympy(): v1 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) assert not v1.not_equal(v2) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() tr1.not_equal(tr2) vector-1.6.3/tests/compute/sympy/test_scale.py000066400000000000000000000262321503546127100215330ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t = sympy.symbols("x y z t", real=True) values = {x: 1, y: 2, z: 3, t: 10} def test_planar_posfactor(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert out.x == x * 1.75 assert out.y == y * 1.75 for t1 in ("xy", "rhophi"): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * 1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * 1.75) def test_planar_negfactor(): vec = vector.VectorSympy2D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert out.x == x * -1.75 assert out.y == y * -1.75 for t1 in ("xy", "rhophi"): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * -1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * -1.75) def test_spatial_posfactor(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert out.x == x * 1.75 assert out.y == y * 1.75 assert out.z == z * 1.75 for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * 1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * 1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * 1.75) def test_spatial_negfactor(): vec = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert out.x == x * -1.75 assert out.y == y * -1.75 assert out.z == z * -1.75 for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * -1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * -1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * -1.75) def test_lorentz_postime_posfactor(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == x * 1.75 assert out.y == y * 1.75 assert out.z == z * 1.75 assert out.t == t * 1.75 for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * 1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * 1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * 1.75) assert out.t.subs(values).evalf() == pytest.approx(10 * 1.75) def test_lorentz_postime_negfactor(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == x * -1.75 assert out.y == y * -1.75 assert out.z == z * -1.75 assert out.t == t * -1.75 for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * -1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * -1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * -1.75) assert out.t.subs(values).evalf() == pytest.approx(10 * -1.75) for t1 in ( "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * -1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * -1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * -1.75) assert out.t.subs(values).evalf() == pytest.approx(10 * 1.75) def test_lorentz_negtime_posfactor(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(-t), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == x * 1.75 assert out.y == y * 1.75 assert out.z == z * 1.75 assert out.t == -t * 1.75 for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * 1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * 1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * 1.75) assert out.t.subs(values).evalf() == pytest.approx(10 * -1.75) for t1 in ( "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * 1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * 1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * 1.75) assert out.t.subs(values).evalf() == pytest.approx(10 * 1.75) def test_lorentz_negtime_negfactor(): vec = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(-t), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == x * -1.75 assert out.y == y * -1.75 assert out.z == z * -1.75 assert out.t == -t * -1.75 for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * -1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * -1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * -1.75) assert out.t.subs(values).evalf() == pytest.approx(-10 * -1.75) for t1 in ( "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x.subs(values).evalf() == pytest.approx(1 * -1.75) assert out.y.subs(values).evalf() == pytest.approx(2 * -1.75) assert out.z.subs(values).evalf() == pytest.approx(3 * -1.75) assert out.t.subs(values).evalf() == pytest.approx(10 * 1.75) vector-1.6.3/tests/compute/sympy/test_subtract.py000066400000000000000000000075611503546127100222770ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, t, nx, ny, nz, nt = sympy.symbols("x y z t nx ny nz nt", real=True) values = {x: 3, y: 4, z: 2, t: 40, nx: 5, ny: 12, nz: 4, nt: 20} def test_planar_sympy(): v1 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) v2 = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny)) out = v1.subtract(v2) assert out.x == x - nx assert out.y == y - ny for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x.subs(values).evalf() == pytest.approx(-2) assert out.y.subs(values).evalf() == pytest.approx(-8) def test_spatial_sympy(): v1 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) v2 = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), ) out = v1.subtract(v2) assert out.x == x - nx assert out.y == y - ny assert out.z == z - nz for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x.subs(values).evalf() == pytest.approx(-2) assert out.y.subs(values).evalf() == pytest.approx(-8) assert out.z.subs(values).evalf() == pytest.approx(-2) def test_lorentz_sympy(): v1 = vector.backends.sympy.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(t), ) v2 = vector.backends.sympy.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(nx, ny), longitudinal=vector.backends.sympy.LongitudinalSympyZ(nz), temporal=vector.backends.sympy.TemporalSympyT(nt), ) out = v1.subtract(v2) assert out.x == x - nx assert out.y == y - ny assert out.z == z - nz assert out.t == t - nt for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x.subs(values).evalf() == pytest.approx(-2) assert out.y.subs(values).evalf() == pytest.approx(-8) assert out.z.subs(values).evalf() == pytest.approx(-2) assert out.t.subs(values).evalf() == pytest.approx(20) vector-1.6.3/tests/compute/sympy/test_unit.py000066400000000000000000000056671503546127100214340ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector sympy = pytest.importorskip("sympy") pytestmark = pytest.mark.sympy x, y, z, time = sympy.symbols("x y z time") values = {x: 1, y: 2, z: 3, time: 10} def test_planar_sympy(): v = vector.VectorSympy2D(azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y)) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert u.rho.simplify() == 1 for t1 in "xy", "rhophi": t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert ( u.rho == 1 if isinstance(u.rho, int) else u.rho.subs(values).evalf() == 1.0 ) def test_spatial_sympy(): v = vector.VectorSympy3D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert type(u.longitudinal) is type(v.longitudinal) assert u.mag.simplify() == 1 for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert type(u.longitudinal) is type(t.longitudinal) assert u.mag.subs(values).evalf() == 1.0 def test_lorentz_sympy(): v = vector.VectorSympy4D( azimuthal=vector.backends.sympy.AzimuthalSympyXY(x, y), longitudinal=vector.backends.sympy.LongitudinalSympyZ(z), temporal=vector.backends.sympy.TemporalSympyT(time), ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert type(u.longitudinal) is type(v.longitudinal) assert type(u.temporal) is type(v.temporal) assert u.tau.simplify() == sympy.sqrt( sympy.Abs( (-(time**2) + x**2 + y**2 + z**2) / sympy.Abs(-(time**2) + x**2 + y**2 + z**2) ) ) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert type(u.longitudinal) is type(t.longitudinal) assert type(u.temporal) is type(t.temporal) assert ( u.tau == 1 if isinstance(u.tau, int) else u.tau.subs(values).evalf() == 1.0 ) vector-1.6.3/tests/compute/test_add.py000066400000000000000000000160431503546127100200120ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 12) ) out = v1.add(v2) assert out.x == pytest.approx(8) assert out.y == pytest.approx(16) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x == pytest.approx(8) assert out.y == pytest.approx(16) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(3, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(5, 12)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) out = v1.add(v2) assert out.x[0] == pytest.approx(8) assert out.y[0] == pytest.approx(16) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x[0] == pytest.approx(8) assert out.y[0] == pytest.approx(16) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(2), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 12), longitudinal=vector.backends.object.LongitudinalObjectZ(4), ) out = v1.add(v2) assert out.x == pytest.approx(8) assert out.y == pytest.approx(16) assert out.z == pytest.approx(6) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x == pytest.approx(8) assert out.y == pytest.approx(16) assert out.z == pytest.approx(6) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(3, 4, 2)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(5, 12, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) out = v1.add(v2) assert out.x[0] == pytest.approx(8) assert out.y[0] == pytest.approx(16) assert out.z[0] == pytest.approx(6) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x[0] == pytest.approx(8) assert out.y[0] == pytest.approx(16) assert out.z[0] == pytest.approx(6) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(2), temporal=vector.backends.object.TemporalObjectT(10), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 12), longitudinal=vector.backends.object.LongitudinalObjectZ(4), temporal=vector.backends.object.TemporalObjectT(15), ) out = v1.add(v2) assert out.x == pytest.approx(8) assert out.y == pytest.approx(16) assert out.z == pytest.approx(6) assert out.t == pytest.approx(25) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x == pytest.approx(8) assert out.y == pytest.approx(16) assert out.z == pytest.approx(6) assert out.t == pytest.approx(25) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(3, 4, 2, 10)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(5, 12, 4, 15)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) out = v1.add(v2) assert out.x[0] == pytest.approx(8) assert out.y[0] == pytest.approx(16) assert out.z[0] == pytest.approx(6) assert out.t[0] == pytest.approx(25) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.add(transformed2) assert out.x[0] == pytest.approx(8) assert out.y[0] == pytest.approx(16) assert out.z[0] == pytest.approx(6) assert out.t[0] == pytest.approx(25) vector-1.6.3/tests/compute/test_conversions.py000066400000000000000000001025341503546127100216330ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test_VectorObject2D(): v = vector.obj(x=1, y=2) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.object.VectorObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.object.VectorObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(0) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.object.VectorObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(0) assert tv.t == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.object.VectorObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.object.VectorObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert getattr(tv, longitudinal) == pytest.approx(0) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.object.VectorObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert getattr(tv, longitudinal) == pytest.approx(0) assert getattr(tv, temporal) == pytest.approx(0) def test_MomentumObject2D(): v = vector.obj(px=1, py=2) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.object.MomentumObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.object.MomentumObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(0) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.object.MomentumObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(0) assert tv.t == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.object.MomentumObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.object.MomentumObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert getattr(tv, longitudinal) == pytest.approx(0) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.object.MomentumObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert getattr(tv, longitudinal) == pytest.approx(0) assert getattr(tv, temporal) == pytest.approx(0) def test_VectorNumpy2D(): v = vector.array({"x": [1, 1, 1], "y": [2, 2, 2]}) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.numpy.VectorNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.numpy.VectorNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(0) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.numpy.VectorNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(0) assert tv.t[0] == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert getattr(tv, longitudinal)[0] == pytest.approx(0) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert getattr(tv, longitudinal)[0] == pytest.approx(0) assert getattr(tv, temporal)[0] == pytest.approx(0) def test_MomentumNumpy2D(): v = vector.array({"px": [1, 1, 1], "py": [2, 2, 2]}) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(0) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(0) assert tv.t[0] == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert getattr(tv, longitudinal)[0] == pytest.approx(0) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert getattr(tv, longitudinal)[0] == pytest.approx(0) assert getattr(tv, temporal)[0] == pytest.approx(0) def test_VectorObject3D(): v = vector.obj(x=1, y=2, z=3) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.object.VectorObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.object.VectorObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.object.VectorObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert tv.t == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.object.VectorObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.object.VectorObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.object.VectorObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert getattr(tv, temporal) == pytest.approx(0) def test_MomentumObject3D(): v = vector.obj(px=1, py=2, pz=3) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.object.MomentumObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.object.MomentumObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.object.MomentumObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert tv.t == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.object.MomentumObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.object.MomentumObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.object.MomentumObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert getattr(tv, temporal) == pytest.approx(0) def test_VectorNumpy3D(): v = vector.array({"x": [1, 1, 1], "y": [2, 2, 2], "z": [3, 3, 3]}) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.numpy.VectorNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.numpy.VectorNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.numpy.VectorNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert tv.t[0] == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert getattr(tv, temporal)[0] == pytest.approx(0) def test_MomentumNumpy3D(): v = vector.array({"px": [1, 1, 1], "py": [2, 2, 2], "pz": [3, 3, 3]}) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert tv.t[0] == pytest.approx(0) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert getattr(tv, temporal)[0] == pytest.approx(0) def test_VectorObject4D(): v = vector.obj(x=1, y=2, z=3, t=4) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.object.VectorObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.object.VectorObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.object.VectorObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert tv.t == pytest.approx(4) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.object.VectorObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.object.VectorObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.object.VectorObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert tv.t == pytest.approx(4) def test_MomentumObject4D(): v = vector.obj(px=1, py=2, pz=3, E=4) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.object.MomentumObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.object.MomentumObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.object.MomentumObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert tv.t == pytest.approx(4) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.object.MomentumObject2D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.object.MomentumObject3D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.object.MomentumObject4D) assert tv.x == pytest.approx(1) assert tv.y == pytest.approx(2) assert tv.z == pytest.approx(3) assert tv.t == pytest.approx(4) def test_VectorNumpy4D(): v = vector.array({"x": [1, 1, 1], "y": [2, 2, 2], "z": [3, 3, 3], "t": [4, 4, 4]}) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.numpy.VectorNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.numpy.VectorNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.numpy.VectorNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert tv.t[0] == pytest.approx(4) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.numpy.VectorNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert tv.t[0] == pytest.approx(4) def test_MomentumNumpy4D(): v = vector.array( {"px": [1, 1, 1], "py": [2, 2, 2], "pz": [3, 3, 3], "E": [4, 4, 4]} ) tv = v.to_Vector2D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) tv = v.to_Vector3D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) tv = v.to_Vector4D() assert isinstance(tv, vector.backends.numpy.MomentumNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert tv.t[0] == pytest.approx(4) for azimuthal in "xy", "rhophi": tv = getattr(v, "to_" + azimuthal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy2D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) for longitudinal in "z", "theta", "eta": tv = getattr(v, "to_" + azimuthal + longitudinal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy3D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) for temporal in "t", "tau": tv = getattr(v, "to_" + azimuthal + longitudinal + temporal)() assert isinstance(tv, vector.backends.numpy.MomentumNumpy4D) assert tv.x[0] == pytest.approx(1) assert tv.y[0] == pytest.approx(2) assert tv.z[0] == pytest.approx(3) assert tv.t[0] == pytest.approx(4) def test_conversion_with_coords_object(): # 2D -> 3D vec = vector.VectorObject2D(x=1, y=2) assert vec.to_Vector3D(z=1).z == 1 assert vec.to_Vector3D(eta=1).eta == 1 assert vec.to_Vector3D(theta=1).theta == 1 # test alias assert vec.to_3D(z=1).x == vec.x assert vec.to_3D(z=1).y == vec.y # 2D -> 4D assert vec.to_Vector4D(z=1, t=1).z == 1 assert vec.to_Vector4D(z=1, t=1).t == 1 assert vec.to_Vector4D(eta=1, t=1).eta == 1 assert vec.to_Vector4D(eta=1, t=1).t == 1 assert vec.to_Vector4D(theta=1, t=1).theta == 1 assert vec.to_Vector4D(theta=1, t=1).t == 1 assert vec.to_Vector4D(z=1, tau=1).z == 1 assert vec.to_Vector4D(z=1, tau=1).tau == 1 assert vec.to_Vector4D(eta=1, tau=1).eta == 1 assert vec.to_Vector4D(eta=1, tau=1).tau == 1 assert vec.to_Vector4D(theta=1, tau=1).theta == 1 assert vec.to_Vector4D(theta=1, tau=1).tau == 1 # test alias assert vec.to_4D(z=1, t=1).x == vec.x assert vec.to_4D(z=1, t=1).y == vec.y # 3D -> 4D vec = vector.VectorObject3D(x=1, y=2, z=3) # test alias assert vec.to_4D(t=1).t == 1 assert vec.to_4D(tau=1).tau == 1 assert vec.to_Vector4D(t=1).x == vec.x assert vec.to_Vector4D(t=1).y == vec.y assert vec.to_Vector4D(t=1).z == vec.z # check if momentum coords work vec = vector.MomentumObject2D(px=1, py=2) assert vec.to_Vector3D(pz=1).pz == 1 # test both alias and original methods assert vec.to_4D(pz=1, m=1).pz == 1 assert vec.to_4D(pz=1, m=1).m == 1 assert vec.to_4D(pz=1, mass=1).mass == 1 assert vec.to_4D(pz=1, M=1).M == 1 assert vec.to_Vector4D(pz=1, e=1).e == 1 assert vec.to_Vector4D(pz=1, energy=1).energy == 1 assert vec.to_Vector4D(pz=1, E=1).E == 1 vec = vector.MomentumObject3D(px=1, py=2, pz=3) # test both alias and original methods assert vec.to_4D(m=1).m == 1 assert vec.to_4D(mass=1).mass == 1 assert vec.to_4D(M=1).M == 1 assert vec.to_Vector4D(e=1).e == 1 assert vec.to_Vector4D(energy=1).energy == 1 assert vec.to_Vector4D(E=1).E == 1 def test_conversion_with_coords_numpy(): # 2D -> 3D vec = vector.VectorNumpy2D( [(1.0, 1.0), (2.0, 2.0)], dtype=[("x", float), ("y", float)], ) assert all(vec.to_Vector3D(z=1).z == 1) assert all(vec.to_Vector3D(eta=1).eta == 1) assert all(vec.to_Vector3D(theta=1).theta == 1) # test alias assert all(vec.to_3D(z=1).x == vec.x) assert all(vec.to_3D(z=1).y == vec.y) # 2D -> 4D assert all(vec.to_Vector4D(z=1, t=1).t == 1) assert all(vec.to_Vector4D(z=1, t=1).z == 1) assert all(vec.to_Vector4D(eta=1, t=1).eta == 1) assert all(vec.to_Vector4D(eta=1, t=1).t == 1) assert all(vec.to_Vector4D(theta=1, t=1).theta == 1) assert all(vec.to_Vector4D(theta=1, t=1).t == 1) assert all(vec.to_Vector4D(z=1, tau=1).z == 1) assert all(vec.to_Vector4D(z=1, tau=1).tau == 1) assert all(vec.to_Vector4D(eta=1, tau=1).eta == 1) assert all(vec.to_Vector4D(eta=1, tau=1).tau == 1) assert all(vec.to_Vector4D(theta=1, tau=1).theta == 1) assert all(vec.to_Vector4D(theta=1, tau=1).tau == 1) # test alias assert all(vec.to_4D(z=1, t=1).x == vec.x) assert all(vec.to_4D(z=1, t=1).y == vec.y) # 3D -> 4D vec = vector.VectorNumpy3D( [(1.0, 1.0, 1.0), (2.0, 2.0, 2.0)], dtype=[("x", float), ("y", float), ("z", float)], ) assert all(vec.to_Vector4D(t=1).t == 1) assert all(vec.to_Vector4D(tau=1).tau == 1) # test alias assert all(vec.to_4D(t=1).x == vec.x) assert all(vec.to_4D(t=1).y == vec.y) assert all(vec.to_4D(t=1).z == vec.z) # check if momentum coords work vec = vector.MomentumNumpy2D( [(1.0, 1.0), (2.0, 2.0)], dtype=[("px", float), ("py", float)], ) assert all(vec.to_Vector3D(pz=1).pz == 1) # test both alias and original methods assert all(vec.to_4D(pz=1, m=1).pz == 1) assert all(vec.to_4D(pz=1, m=1).m == 1) assert all(vec.to_4D(pz=1, mass=1).mass == 1) assert all(vec.to_4D(pz=1, M=1).M == 1) assert all(vec.to_Vector4D(pz=1, e=1).e == 1) assert all(vec.to_Vector4D(pz=1, energy=1).energy == 1) assert all(vec.to_Vector4D(pz=1, E=1).E == 1) vec = vector.MomentumNumpy3D( [(1.0, 1.0, 1.0), (2.0, 2.0, 2.0)], dtype=[("px", float), ("py", float), ("pz", float)], ) # test both alias and original methods assert all(vec.to_4D(m=1).m == 1) assert all(vec.to_4D(mass=1).mass == 1) assert all(vec.to_4D(M=1).M == 1) assert all(vec.to_Vector4D(e=1).e == 1) assert all(vec.to_Vector4D(energy=1).energy == 1) assert all(vec.to_Vector4D(E=1).E == 1) def test_like_object(): v1 = vector.obj(x=0.1, y=0.2) v2 = vector.obj(x=1, y=2, z=3) v3 = vector.obj(x=10, y=20, z=30, t=40) # 2D + 3D.like(2D) = 2D assert v1 + v2.like(v1) == vector.obj(x=1.1, y=2.2) assert v2.like(v1) + v1 == vector.obj(x=1.1, y=2.2) # 2D + 4D.like(2D) = 2D assert v1 + v3.like(v1) == vector.obj(x=10.1, y=20.2) assert v3.like(v1) + v1 == vector.obj(x=10.1, y=20.2) # 3D + 2D.like(3D) = 3D assert v2 + v1.like(v2) == vector.obj(x=1.1, y=2.2, z=3) assert v1.like(v2) + v2 == vector.obj(x=1.1, y=2.2, z=3) # 3D + 4D.like(3D) = 3D assert v2 + v3.like(v2) == vector.obj(x=11, y=22, z=33) assert v3.like(v2) + v2 == vector.obj(x=11, y=22, z=33) # 4D + 2D.like(4D) = 4D assert v3 + v1.like(v3) == vector.obj(x=10.1, y=20.2, z=30.0, t=40.0) assert v1.like(v3) + v3 == vector.obj(x=10.1, y=20.2, z=30.0, t=40.0) # 4D + 3D.like(4D) = 4D assert v3 + v2.like(v3) == vector.obj(x=11, y=22, z=33, t=40) assert v2.like(v3) + v3 == vector.obj(x=11, y=22, z=33, t=40) v1 = vector.obj(px=0.1, py=0.2) v2 = vector.obj(px=1, py=2, pz=3) v3 = vector.obj(px=10, py=20, pz=30, t=40) # order should not matter # 2D + 3D.like(2D) = 2D assert v1 + v2.like(v1) == vector.obj(px=1.1, py=2.2) assert v2.like(v1) + v1 == vector.obj(px=1.1, py=2.2) # 2D + 4D.like(2D) = 2D assert v1 + v3.like(v1) == vector.obj(px=10.1, py=20.2) assert v3.like(v1) + v1 == vector.obj(px=10.1, py=20.2) # 3D + 2D.like(3D) = 3D assert v2 + v1.like(v2) == vector.obj(px=1.1, py=2.2, pz=3) assert v1.like(v2) + v2 == vector.obj(px=1.1, py=2.2, pz=3) # 3D + 4D.like(3D) = 3D assert v2 + v3.like(v2) == vector.obj(px=11, py=22, pz=33) assert v3.like(v2) + v2 == vector.obj(px=11, py=22, pz=33) # 4D + 2D.like(4D) = 4D assert v3 + v1.like(v3) == vector.obj(px=10.1, py=20.2, pz=30.0, E=40.0) assert v1.like(v3) + v3 == vector.obj(px=10.1, py=20.2, pz=30.0, E=40.0) # 4D + 3D.like(4D) = 4D assert v3 + v2.like(v3) == vector.obj(px=11, py=22, pz=33, E=40) assert v2.like(v3) + v3 == vector.obj(px=11, py=22, pz=33, E=40) v1 = vector.obj(px=0.1, py=0.2) v2 = vector.obj(x=1, y=2, z=3) v3 = vector.obj(px=10, py=20, pz=30, t=40) def test_like_numpy(): v1 = vector.array( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], }, ) v2 = vector.array( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], "z": [5.0, 1.0, 1.0], }, ) v3 = vector.array( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], "z": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ) v1_v2 = vector.array( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], }, ) v2_v1 = vector.array( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [5.0, 1.0, 1.0], }, ) v2_v3 = vector.array( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [10.0, 2.0, 2.0], }, ) v3_v2 = vector.array( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [10.0, 2.0, 2.0], "t": [16.0, 31.0, 46.0], }, ) v1_v3 = vector.array( { "x": [20.0, 40.0, 60.0], "y": [-20.0, 40.0, 60.0], "z": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ) # 2D + 3D.like(2D) = 2D assert all(v1 + v2.like(v1) == v1_v2) assert all(v2.like(v1) + v1 == v1_v2) # 2D + 4D.like(2D) = 2D assert all(v1 + v3.like(v1) == v1_v2) assert all(v3.like(v1) + v1 == v1_v2) # 3D + 2D.like(3D) = 3D assert all(v2 + v1.like(v2) == v2_v1) assert all(v1.like(v2) + v2 == v2_v1) # 3D + 4D.like(3D) = 3D assert all(v2 + v3.like(v2) == v2_v3) assert all(v3.like(v2) + v2 == v2_v3) # 4D + 2D.like(4D) = 4D assert all(v3 + v1.like(v3) == v1_v3) assert all(v1.like(v3) + v3 == v1_v3) # 4D + 3D.like(4D) = 4D assert all(v3 + v2.like(v3) == v3_v2) assert all(v2.like(v3) + v3 == v3_v2) v1 = vector.array( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], }, ) v2 = vector.array( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], "pz": [5.0, 1.0, 1.0], }, ) v3 = vector.array( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], "pz": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ) pv1_v2 = vector.array( { "px": [20.0, 40.0, 60.0], "py": [-20.0, 40.0, 60.0], }, ) pv2_v1 = vector.array( { "px": [20.0, 40.0, 60.0], "py": [-20.0, 40.0, 60.0], "pz": [5.0, 1.0, 1.0], }, ) pv2_v3 = vector.array( { "px": [20.0, 40.0, 60.0], "py": [-20.0, 40.0, 60.0], "pz": [10.0, 2.0, 2.0], }, ) pv3_v2 = vector.array( { "px": [20.0, 40.0, 60.0], "py": [-20.0, 40.0, 60.0], "pz": [10.0, 2.0, 2.0], "t": [16.0, 31.0, 46.0], }, ) pv1_v3 = vector.array( { "px": [20.0, 40.0, 60.0], "py": [-20.0, 40.0, 60.0], "pz": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ) # 2D + 3D.like(2D) = 2D assert all(v1 + v2.like(v1) == pv1_v2) assert all(v2.like(v1) + v1 == pv1_v2) # 2D + 4D.like(2D) = 2D assert all(v1 + v3.like(v1) == pv1_v2) assert all(v3.like(v1) + v1 == pv1_v2) # 3D + 2D.like(3D) = 3D assert all(v2 + v1.like(v2) == pv2_v1) assert all(v1.like(v2) + v2 == pv2_v1) # 3D + 4D.like(3D) = 3D assert all(v2 + v3.like(v2) == pv2_v3) assert all(v3.like(v2) + v2 == pv2_v3) # 4D + 2D.like(4D) = 4D assert all(v3 + v1.like(v3) == pv1_v3) assert all(v1.like(v3) + v3 == pv1_v3) # 4D + 3D.like(4D) = 4D assert all(v3 + v2.like(v3) == pv3_v2) assert all(v2.like(v3) + v3 == pv3_v2) def test_momentum_preservation_object(): v1 = vector.obj(px=0.1, py=0.2) v2 = vector.obj(x=1, y=2, z=3) v3 = vector.obj(px=10, py=20, pz=30, t=40) # momentum + generic = momentum # 2D + 3D.like(2D) = 2D assert isinstance(v1 + v2.like(v1), vector.MomentumObject2D) assert isinstance(v2.like(v1) + v1, vector.MomentumObject2D) # 2D + 4D.like(2D) = 2D assert isinstance(v1 + v3.like(v1), vector.MomentumObject2D) assert isinstance(v3.like(v1) + v1, vector.MomentumObject2D) # 3D + 2D.like(3D) = 3D assert isinstance(v2 + v1.like(v2), vector.MomentumObject3D) assert isinstance(v1.like(v2) + v2, vector.MomentumObject3D) # 3D + 4D.like(3D) = 3D assert isinstance(v2 + v3.like(v2), vector.MomentumObject3D) assert isinstance(v3.like(v2) + v2, vector.MomentumObject3D) # 4D + 2D.like(4D) = 4D assert isinstance(v3 + v1.like(v3), vector.MomentumObject4D) assert isinstance(v1.like(v3) + v3, vector.MomentumObject4D) # 4D + 3D.like(4D) = 4D assert isinstance(v3 + v2.like(v3), vector.MomentumObject4D) assert isinstance(v2.like(v3) + v3, vector.MomentumObject4D) def test_momentum_preservation_numpy(): v1 = vector.array( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], }, ) v2 = vector.array( { "x": [10.0, 20.0, 30.0], "y": [-10.0, 20.0, 30.0], "z": [5.0, 1.0, 1.0], }, ) v3 = vector.array( { "px": [10.0, 20.0, 30.0], "py": [-10.0, 20.0, 30.0], "pz": [5.0, 1.0, 1.0], "t": [16.0, 31.0, 46.0], }, ) # momentum + generic = momentum # 2D + 3D.like(2D) = 2D assert isinstance(v1 + v2.like(v1), vector.MomentumNumpy2D) assert isinstance(v2.like(v1) + v1, vector.MomentumNumpy2D) # 2D + 4D.like(2D) = 2D assert isinstance(v1 + v3.like(v1), vector.MomentumNumpy2D) assert isinstance(v3.like(v1) + v1, vector.MomentumNumpy2D) # 3D + 2D.like(3D) = 3D assert isinstance(v2 + v1.like(v2), vector.MomentumNumpy3D) assert isinstance(v1.like(v2) + v2, vector.MomentumNumpy3D) # 3D + 4D.like(3D) = 3D assert isinstance(v2 + v3.like(v2), vector.MomentumNumpy3D) assert isinstance(v3.like(v2) + v2, vector.MomentumNumpy3D) # 4D + 2D.like(4D) = 4D assert isinstance(v3 + v1.like(v3), vector.MomentumNumpy4D) assert isinstance(v1.like(v3) + v3, vector.MomentumNumpy4D) # 4D + 3D.like(4D) = 4D assert isinstance(v3 + v2.like(v3), vector.MomentumNumpy4D) assert isinstance(v2.like(v3) + v3, vector.MomentumNumpy4D) vector-1.6.3/tests/compute/test_dot.py000066400000000000000000000132251503546127100200470ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.3, 0.4) ) assert v1.dot(v2) == pytest.approx(0.11) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.dot(transformed2) == pytest.approx(0.11) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(0.3, 0.4)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert v1.dot(v2)[0] == pytest.approx(0.11) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.dot(tr2)[0] == pytest.approx(0.11) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.4, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.6), ) assert v1.dot(v2) == pytest.approx(0.32) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.dot(transformed2) == pytest.approx(0.32) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.4, 0.5, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.dot(v2)[0] == pytest.approx(0.32) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.dot(tr2)[0] == pytest.approx(0.32) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.5, 0.6), longitudinal=vector.backends.object.LongitudinalObjectZ(0.7), temporal=vector.backends.object.TemporalObjectT(0.8), ) assert v1.dot(v2) == pytest.approx(-0.06) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.dot(tr2) == pytest.approx(-0.06) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.5, 0.6, 0.7, 0.8)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.dot(v2) == pytest.approx(-0.06) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.dot(tr2) == pytest.approx(-0.06) vector-1.6.3/tests/compute/test_equal.py000066400000000000000000000134751503546127100203770ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) assert v1.equal(v2) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose transformed1.equal(transformed2) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert v1.equal(v2).all() for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose tr1.equal(tr2) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) assert v1.equal(v2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose transformed1.equal(transformed2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.equal(v2).all() for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose tr1.equal(tr2) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) assert v1.equal(v2) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() # can't test for exact equality due to round-off: see test_isclose tr1.equal(tr2) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.equal(v2).all() for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() # can't test for exact equality due to round-off: see test_isclose tr1.equal(tr2) vector-1.6.3/tests/compute/test_is_antiparallel.py000066400000000000000000000103041503546127100224170ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(-0.3, -0.6) ) v3 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.3, 0.6) ) assert v1.is_antiparallel(-v1) assert v2.is_antiparallel(-v2) assert v3.is_antiparallel(-v3) assert v1.is_antiparallel(v2) assert v2.is_antiparallel(v3) assert not v1.is_antiparallel(v3) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_antiparallel(tr2) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(-0.3, -0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v3 = vector.backends.numpy.VectorNumpy2D( [(0.3, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert v1.is_antiparallel(-v1) assert v2.is_antiparallel(-v2) assert v3.is_antiparallel(-v3) assert v1.is_antiparallel(v2) assert v2.is_antiparallel(v3) assert not v1.is_antiparallel(v3) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_antiparallel(tr2) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(-0.3, -0.6), longitudinal=vector.backends.object.LongitudinalObjectZ(-0.9), ) v3 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.3, 0.6), longitudinal=vector.backends.object.LongitudinalObjectZ(0.9), ) assert v1.is_antiparallel(-v1) assert v2.is_antiparallel(-v2) assert v3.is_antiparallel(-v3) assert v1.is_antiparallel(v2) assert v2.is_antiparallel(v3) assert not v1.is_antiparallel(v3) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_antiparallel(tr2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(-0.3, -0.6, -0.9)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v3 = vector.backends.numpy.VectorNumpy3D( [(0.3, 0.6, 0.9)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.is_antiparallel(-v1) assert v2.is_antiparallel(-v2) assert v3.is_antiparallel(-v3) assert v1.is_antiparallel(v2) assert v2.is_antiparallel(v3) assert not v1.is_antiparallel(v3) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_antiparallel(tr2) vector-1.6.3/tests/compute/test_is_parallel.py000066400000000000000000000101161503546127100215440ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.3, 0.6) ) v3 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.5) ) assert v1.is_parallel(v1) assert v2.is_parallel(v2) assert v3.is_parallel(v3) assert v1.is_parallel(v2) assert not v1.is_parallel(v3) assert not v2.is_parallel(v3) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_parallel(tr2) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(0.3, 0.6)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v3 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.5)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert v1.is_parallel(v1) assert v2.is_parallel(v2) assert v3.is_parallel(v3) assert v1.is_parallel(v2) assert not v1.is_parallel(v3) assert not v2.is_parallel(v3) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_parallel(tr2) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.3, 0.6), longitudinal=vector.backends.object.LongitudinalObjectZ(0.9), ) v3 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.9), ) assert v1.is_parallel(v1) assert v2.is_parallel(v2) assert v3.is_parallel(v3) assert v1.is_parallel(v2) assert not v1.is_parallel(v3) assert not v2.is_parallel(v3) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_parallel(tr2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.3, 0.6, 0.9)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v3 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.5, 0.9)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.is_parallel(v1) assert v2.is_parallel(v2) assert v3.is_parallel(v3) assert v1.is_parallel(v2) assert not v1.is_parallel(v3) assert not v2.is_parallel(v3) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_parallel(tr2) vector-1.6.3/tests/compute/test_is_perpendicular.py000066400000000000000000000103421503546127100226060ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0, 1) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 0) ) v3 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.5) ) assert not v1.is_perpendicular(v1) assert not v2.is_perpendicular(v2) assert not v3.is_perpendicular(v3) assert v1.is_perpendicular(v2) assert not v1.is_perpendicular(v3) assert not v2.is_perpendicular(v3) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_perpendicular(tr2) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0, 1)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(1, 0)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v3 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.5)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert not v1.is_perpendicular(v1) assert not v2.is_perpendicular(v2) assert not v3.is_perpendicular(v3) assert v1.is_perpendicular(v2) assert not v1.is_perpendicular(v3) assert not v2.is_perpendicular(v3) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_perpendicular(tr2) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0, 1), longitudinal=vector.backends.object.LongitudinalObjectZ(1), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 0), longitudinal=vector.backends.object.LongitudinalObjectZ(0), ) v3 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.5), longitudinal=vector.backends.object.LongitudinalObjectZ(0.9), ) assert not v1.is_perpendicular(v1) assert not v2.is_perpendicular(v2) assert not v3.is_perpendicular(v3) assert v1.is_perpendicular(v2) assert not v1.is_perpendicular(v3) assert not v2.is_perpendicular(v3) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_perpendicular(tr2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0, 1, 1)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(1, 0, 0)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v3 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.5, 0.9)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert not v1.is_perpendicular(v1) assert not v2.is_perpendicular(v2) assert not v3.is_perpendicular(v3) assert v1.is_perpendicular(v2) assert not v1.is_perpendicular(v3) assert not v2.is_perpendicular(v3) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.is_perpendicular(tr2) vector-1.6.3/tests/compute/test_isclose.py000066400000000000000000000126611503546127100207250ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) assert v1.isclose(v2) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.isclose(transformed2) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert v1.isclose(v2).all() for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.isclose(tr2).all() def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) assert v1.isclose(v2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert transformed1.isclose(transformed2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert v1.isclose(v2).all() for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) assert tr1.isclose(tr2).all() def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) assert v1.isclose(v2) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.isclose(tr2) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert v1.isclose(v2) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() assert tr1.isclose(tr2).all() vector-1.6.3/tests/compute/test_not_equal.py000066400000000000000000000136051503546127100212520ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) assert not v1.not_equal(v2) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose transformed1.not_equal(transformed2) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) assert not v1.not_equal(v2).all() for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose tr1.not_equal(tr2) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) assert not v1.not_equal(v2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose transformed1.not_equal(transformed2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) assert not v1.not_equal(v2).all() for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": tr1, tr2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) # can't test for exact equality due to round-off: see test_isclose tr1.not_equal(tr2) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) assert not v1.not_equal(v2) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() # can't test for exact equality due to round-off: see test_isclose tr1.not_equal(tr2) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) assert not v1.not_equal(v2).all() for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tr1, tr2 = getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)() # can't test for exact equality due to round-off: see test_isclose tr1.not_equal(tr2) vector-1.6.3/tests/compute/test_scale.py000066400000000000000000000255301503546127100203520ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import pytest import vector.backends.numpy import vector.backends.object def test_planar_posfactor(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) for t1 in ("xy", "rhophi"): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) def test_planar_negfactor(): vec = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) for t1 in ("xy", "rhophi"): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) def test_spatial_posfactor(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) def test_spatial_negfactor(): vec = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) for t1 in ( "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) def test_lorentz_postime_posfactor(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) assert out.t == pytest.approx(4 * 1.75) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) assert out.t == pytest.approx(4 * 1.75) def test_lorentz_postime_negfactor(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(4), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) assert out.t == pytest.approx(4 * -1.75) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) assert out.t == pytest.approx(4 * -1.75) for t1 in ( "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) assert out.t == pytest.approx(6.06217782649107) def test_lorentz_negtime_posfactor(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(-1.5), ) out = vec.scale(1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) assert out.t == pytest.approx(-1.5 * 1.75) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) assert out.t == pytest.approx(-1.5 * 1.75) for t1 in ( "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * 1.75) assert out.y == pytest.approx(2 * 1.75) assert out.z == pytest.approx(3 * 1.75) assert out.t == pytest.approx(1.5 * 1.75) def test_lorentz_negtime_negfactor(): vec = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(1, 2), longitudinal=vector.backends.object.LongitudinalObjectZ(3), temporal=vector.backends.object.TemporalObjectT(-1.5), ) out = vec.scale(-1.75) assert type(out.azimuthal) == type(vec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(vec.longitudinal) # noqa: E721 assert type(out.temporal) == type(vec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) assert out.t == pytest.approx(-1.5 * -1.75) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) assert out.t == pytest.approx(-1.5 * -1.75) for t1 in ( "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): tvec = getattr(vec, "to_" + t1)() out = tvec.scale(-1.75) assert type(out.azimuthal) == type(tvec.azimuthal) # noqa: E721 assert type(out.longitudinal) == type(tvec.longitudinal) # noqa: E721 assert type(out.temporal) == type(tvec.temporal) # noqa: E721 assert out.x == pytest.approx(1 * -1.75) assert out.y == pytest.approx(2 * -1.75) assert out.z == pytest.approx(3 * -1.75) assert out.t == pytest.approx(8.880280119455692) vector-1.6.3/tests/compute/test_subtract.py000066400000000000000000000161571503546127100211170ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_planar_object(): v1 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4) ) v2 = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 12) ) out = v1.subtract(v2) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-8) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-8) def test_planar_numpy(): v1 = vector.backends.numpy.VectorNumpy2D( [(3, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy2D( [(5, 12)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) out = v1.subtract(v2) assert out.x[0] == pytest.approx(-2) assert out.y[0] == pytest.approx(-8) for t1 in "xy", "rhophi": for t2 in "xy", "rhophi": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x[0] == pytest.approx(-2) assert out.y[0] == pytest.approx(-8) def test_spatial_object(): v1 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(2), ) v2 = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 12), longitudinal=vector.backends.object.LongitudinalObjectZ(4), ) out = v1.subtract(v2) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-8) assert out.z == pytest.approx(-2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-8) assert out.z == pytest.approx(-2) def test_spatial_numpy(): v1 = vector.backends.numpy.VectorNumpy3D( [(3, 4, 2)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) v2 = vector.backends.numpy.VectorNumpy3D( [(5, 12, 4)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) out = v1.subtract(v2) assert out.x[0] == pytest.approx(-2) assert out.y[0] == pytest.approx(-8) assert out.z[0] == pytest.approx(-2) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": for t2 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x[0] == pytest.approx(-2) assert out.y[0] == pytest.approx(-8) assert out.z[0] == pytest.approx(-2) def test_lorentz_object(): v1 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(3, 4), longitudinal=vector.backends.object.LongitudinalObjectZ(2), temporal=vector.backends.object.TemporalObjectT(20), ) v2 = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(5, 12), longitudinal=vector.backends.object.LongitudinalObjectZ(4), temporal=vector.backends.object.TemporalObjectT(15), ) out = v1.subtract(v2) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-8) assert out.z == pytest.approx(-2) assert out.t == pytest.approx(5) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x == pytest.approx(-2) assert out.y == pytest.approx(-8) assert out.z == pytest.approx(-2) assert out.t == pytest.approx(5) def test_lorentz_numpy(): v1 = vector.backends.numpy.VectorNumpy4D( [(3, 4, 2, 20)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) v2 = vector.backends.numpy.VectorNumpy4D( [(5, 12, 4, 15)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) out = v1.subtract(v2) assert out.x[0] == pytest.approx(-2) assert out.y[0] == pytest.approx(-8) assert out.z[0] == pytest.approx(-2) assert out.t[0] == pytest.approx(5) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): for t2 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): transformed1, transformed2 = ( getattr(v1, "to_" + t1)(), getattr(v2, "to_" + t2)(), ) out = transformed1.subtract(transformed2) assert out.x[0] == pytest.approx(-2) assert out.y[0] == pytest.approx(-8) assert out.z[0] == pytest.approx(-2) assert out.t[0] == pytest.approx(5) vector-1.6.3/tests/compute/test_unit.py000066400000000000000000000116461503546127100202450ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import numpy import pytest import vector.backends.numpy import vector.backends.object def test_planar_object(): v = vector.backends.object.VectorObject2D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2) ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert u.rho == pytest.approx(1) for t1 in "xy", "rhophi": t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert u.rho == pytest.approx(1) def test_planar_numpy(): v = vector.backends.numpy.VectorNumpy2D( [(0.1, 0.2)], dtype=[("x", numpy.float64), ("y", numpy.float64)], ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert u.rho[0] == pytest.approx(1) for t1 in "xy", "rhophi": t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert u.rho[0] == pytest.approx(1) def test_spatial_object(): v = vector.backends.object.VectorObject3D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert type(u.longitudinal) is type(v.longitudinal) assert u.mag == pytest.approx(1) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert type(u.longitudinal) is type(t.longitudinal) assert u.mag == pytest.approx(1) def test_spatial_numpy(): v = vector.backends.numpy.VectorNumpy3D( [(0.1, 0.2, 0.3)], dtype=[("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64)], ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert type(u.longitudinal) is type(v.longitudinal) assert u.mag[0] == pytest.approx(1) for t1 in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert type(u.longitudinal) is type(t.longitudinal) assert u.mag[0] == pytest.approx(1) def test_lorentz_object(): v = vector.backends.object.VectorObject4D( azimuthal=vector.backends.object.AzimuthalObjectXY(0.1, 0.2), longitudinal=vector.backends.object.LongitudinalObjectZ(0.3), temporal=vector.backends.object.TemporalObjectT(0.4), ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert type(u.longitudinal) is type(v.longitudinal) assert type(u.temporal) is type(v.temporal) assert u.tau == pytest.approx(1) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert type(u.longitudinal) is type(t.longitudinal) assert type(u.temporal) is type(t.temporal) assert u.tau == pytest.approx(1) def test_lorentz_numpy(): v = vector.backends.numpy.VectorNumpy4D( [(0.1, 0.2, 0.3, 0.4)], dtype=[ ("x", numpy.float64), ("y", numpy.float64), ("z", numpy.float64), ("t", numpy.float64), ], ) u = v.unit() assert type(u) is type(v) assert type(u.azimuthal) is type(v.azimuthal) assert type(u.longitudinal) is type(v.longitudinal) assert type(u.temporal) is type(v.temporal) assert u.tau[0] == pytest.approx(1) for t1 in ( "xyzt", "xythetat", "xyetat", "rhophizt", "rhophithetat", "rhophietat", "xyztau", "xythetatau", "xyetatau", "rhophiztau", "rhophithetatau", "rhophietatau", ): t = getattr(v, "to_" + t1)() u = t.unit() assert type(u) is type(t) assert type(u.azimuthal) is type(t.azimuthal) assert type(u.longitudinal) is type(t.longitudinal) assert type(u.temporal) is type(t.temporal) assert u.tau[0] == pytest.approx(1) vector-1.6.3/tests/samples/000077500000000000000000000000001503546127100156355ustar00rootroot00000000000000vector-1.6.3/tests/samples/issue-161-v2.pkl000066400000000000000000007142751503546127100203470ustar00rootroot00000000000000X { "class": "RecordArray", "contents": { "px": { "class": "NumpyArray", "itemsize": 8, "format": "d", "primitive": "float64", "form_key": "node1" }, "py": { "class": "NumpyArray", "itemsize": 8, "format": "d", "primitive": "float64", "form_key": "node2" }, "pz": { "class": "NumpyArray", "itemsize": 8, "format": "d", "primitive": "float64", "form_key": "node3" }, "E": { "class": "NumpyArray", "itemsize": 8, "format": "d", "primitive": "float64", "form_key": "node4" }, "area": { "class": "NumpyArray", "itemsize": 8, "format": "d", "primitive": "float64", "form_key": "node5" }, "constituents": { "class": "ListArray64", "starts": "i64", "stops": "i64", "content": { "class": "RecordArray", "contents": { "rho": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 4, "format": "f", "primitive": "float32", "form_key": "node9" }, "form_key": "node8" }, "phi": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 4, "format": "f", "primitive": "float32", "form_key": "node11" }, "form_key": "node10" }, "eta": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 4, "format": "f", "primitive": "float32", "form_key": "node13" }, "form_key": "node12" }, "tau": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 4, "format": "f", "primitive": "float32", "form_key": "node15" }, "form_key": "node14" }, "charged": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 4, "format": "f", "primitive": "float32", "form_key": "node17" }, "form_key": "node16" }, "label": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 4, "format": "i", "primitive": "int32", "form_key": "node19" }, "form_key": "node18" }, "index": { "class": "IndexedArray64", "index": "i64", "content": { "class": "NumpyArray", "itemsize": 8, "format": "l", "primitive": "int64", "form_key": "node21" }, "form_key": "node20" }, "jf_index": { "class": "NumpyArray", "itemsize": 8, "format": "l", "primitive": "int64", "form_key": "node22" } }, "parameters": { "__record__": "Momentum4D" }, "form_key": "node7" }, "form_key": "node6" }, "matching": { "class": "NumpyArray", "itemsize": 8, "format": "l", "primitive": "int64", "form_key": "node23" }, "angularity": { "class": "NumpyArray", "itemsize": 8, "format": "d", "primitive": "float64", "form_key": "node24" } }, "parameters": { "__record__": "Momentum4D" }, "form_key": "node0" } K}( node20-indexnumpy.core.multiarray _reconstructnumpyndarrayKCbR(KMhdtypei8R(K?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb node21-datahhKh R(KM hi8R(KhNNNJJKtbBL                                                                                                           tb node22-datahhKh R(KMhBp@                                                                                                                                     tb node23-datahhKh R(KKhBtb node24-datahhKh R(KKhf8R(KhNNNJJKtbB:4~?zd?u[?nZ)vK?vz>?+Xt?fY?jO !? d9Ǵ?ܮ?bx??AJ?VGw?iV??mMlD?ͭ?C3?FNZ?L5ӹ?ިX?x/J:?h?qfPƞ?.T?ҕٳ?KOZ?Hma?pՙ'箤?OV4#?)KIm?)7x?4ZR"?9?|9)r?lcUq? i,p?B^ ?3Ȼ?,U8?(M?3z?`kz?y62T?lXq?["'+?eH?"jR?li?]?N"z?tW®?@?s孲?b ?f{ԧ?Ńj.N?LSnɨ?0buS?[z)?%$O?;,V#=?OPeQ?t?s*,ޤ? ?&&Ma?\f ?x@?$ ?> t?u??SqwCı?CZ5?Ϧ ?ղͥ?vC?6^0?[?V?y_bt? (l?EH!e?zX.˓?k)7?Ͱֽ?۠zǮ?5??;G?#a?@ ?zɉџ?+?y$!?ua|?QEL?Ɩ9? 공?^W>W̳?}0´?Rl?(3 Z?+?zId͋O?ɍ?[.?:XN?^U?7d?pgp?4Q??:?^%?g\G?j?5?/?V6m?MDwmٸ?SN?5?#A?-1ѫ?d?;)2?.zF?Xұ?gJ?Ѷ?0!?!. ?~}?CƢ?k#釂?`c0L?Y? ?]r vu?ő?q]?r3v?]駧֛?VDԹ ?-޸WR?Z?>ZPy+?yp0?)d?Uh?>?#^^ٴ?g,p?>??GTʇ?lX;?Ky??"?b+۴?SQ?RZ?$?.Z?U?"Qл??ZmS?l A>?Fo #??v?P!@٣?[68EY&??6?#Ĩ?JL?9N.?8eid(?!bKv?he?Dbu?Hn1?0˼?duIL?,"yo? ɷ?Г[?J׾?aK?53?ˠ?2R3?`?| m??,Z ?s?h?>ag?Jy L?׳x?S][?0ΡI?tb node1-datahhKh R(KKh7Bw<%@j-HG@ I5@G<&}B@A@+%dt7?@#@g?q" .:@-{C@`r@h<-q@ؖ&@n?@@p5t1 d22@@AhK@݃<f7@E@& Q;E,@ڱ6;:@`:@!1@"6I@M kCM"7MF@CE@e#7 @6@/;4@}FA@3 <@!J>8q @/¨&<Uy+0@|?Wߞ.1>'F@tC@ĠU@??g@1#@{?J7@̣UA@p@5@M9ՄAX"+ UT @` Rl<@9gJ8<=)WF7 il 1@v=@[B87D@@qE@p)@zx8Cx=@a7)@|A?/@k5(@t>@|H@ @'\>8+>@2Ei0@Z?@-B@A6@Pg0@ɫ@S>@^4=xH 4@~ֽ2O|3L":@B)MMŪ@v7k8@:4H5@(6@i5߆9r5@"OG@ A@?D(@rEg7|UL@;jH@";@V!@_#@Mq+m>@A%9@/G AB4@3@R~:@7 X06(;uW3@-@_(!B@SI4Z$@qx:@?I@^@@ A@O@:4h@һD@xK?N>@37@2KG D!@muD3@_E@@H@*6;@(j>@B>&v#@DAzC@ Д&@7@DKCw@X8@\U5>D ="N15IB=?F;!@K2.@qN@@/\?@ 2+B6I35@-NM,} h:@m:Z @a-#$@'V&81@J8-4@¶AJ@Ax:{r1#:Mk>P@Wk5@5(@06@&H@p/@N1@K8@/- ,95%@6@@t@@l-)3fB@@q U4@4(A4t*@@y2zd3.-@Ɛ< i;I=@R@dN@X!@;M H)cH@tb node2-datahhKh R(KKh7B$C@@3e6Tr8@a_@Qܯ@Л @-@@>+AZ#7Es?@@>V@R9`:2@".@@@@[)@2<9@k%L$H=GUP*F}c0 5@ :>@@mK4=ZNd?^52F11@w8C9<}-@Z2W;0]84%)'C$@̬Ao @iy6@ K<v?Ӌ'@'ȹ2p= pOI+Bi@29EC@O}@2eD@v @>@B@k1{)E=]:@93y7@k1@@0`^x9'@&@PtK@y@!^1C=N A@g4>h53~ %@]@@ȗ &@E@!M|3@u?C @DgFO*1h&w$@1@Sp-@B@ 5?sD@=8@O@b5%QEB@=.@p@lGD5?8@u|'@r]6@?!}݂1}J/:@>%\@)"7@M\@B@@A@92#ڽ;wo@RA@89@to9y9@@@9f3@=3@@,?3`֔&P@2@X9@V?@{%@A9@jKFVA,/ EE@L!-*@ņ9@>#D@43!>@SD@f:?@-@'Q+8@>h¿j5;On,@G2@(z6j D:W".m3!@/cN@qL@Ʀ@58%tb7 node3-datahhKh R(KKh7BA/.ϰ?&# $@@_3& c@,[&@ 6p#@TN>j!@\@.`?꨹@o]@MA!@,LR\(XEp @5b a 6@HV)p ? @] j̻" ?S@@,iM$@E8e^@MY@/pm~ho@4T4a @(72@]5@"Ǚ2e'@p% @9߸@2& @P%ul~@@`[@} jy@+'?f$ @ H1 ,@hs@) L @t@f"@@(&_@%"S@d @,(@ @.VRl1@@y&@q![ǣFB@̿f@hGh*B7#Q)(@@*W8T~QX"@d7?e^0h 2@!@3Ӫ'GQ"_uY-@C‡? 8լ_ү@FGK@ڬO!{*0ٛ@?ʆ̿`s|A @HOj?d@a%@y"A@HD#C u @`qw( *@wm@`*@ }@? ?y*@ony?*t+,L&@u?ov?%+@?-|1п,P"!@SU?tB 8A(@E@,? G.@)C+@@@H@@\P@;@m` @ @Ph@dEtaϟ@!h?8 @bZ$5@_&8?@7N @167&@@@(@K$u#է%\?쮅@@1/@@HC!0bc~Q @*r+#P?Yl!@plY@w@- W j@j"@0J$7! @c?('ӿ{J] NXX"@Uul@P˥?pn @jf @@t̏@&!?N!2?n{,u@fo@fm2 Eq@!=@hM8{@t !dB@_ͲO `.# @g@@ @RI*@ԡU?ޅ@`@'C#@'@٠(@k,*@]7Mtb node4-datahhKh R(KKh7B$}D@h=L@3J@k}6@]B@۽C@@jC@6B@>C@LF@;6G@@@,T;W@4>@+F@RUh/@ F@e%@:4@<@@A@bN@ #B@|`M@YYM@T҅@@@@@aZ@@@V@<{J@-@W<@@G <@Ti?@RE@.@J@=@6%@m?@VG@F@:@B@W7@;H<@TS?@{F@<@u@@lْ)@lX5@0@@ R@٠5@JG@:@lE@$<2@ H@ZkM@F@D?@,C@T1@ ЕE@HC@y F@`>@5-w:@5E@Rw2@SB@a1@R@d_?@@K@b?@E@qA@j7@k B@M>@TE@F>@:@~I@t*@Q@fL@V[H@ PB@ݾF@25 I@vY@@^C@@ lu;@@B@I@q?D@0@G@BI@RF@)N@R@uDF@}B@uj?@#2K@PJ@=@ A@3A@y45@4@t=@C@ZC@ڭGP@%|N@E!@7|9@5DC@7C@b-;@o6@m:@0^6@c~wI@?qQ@KJ@D%|L@C@I~E@oT@?N@kI@=@((@ƙ23@@0@@ P@I6<@L@N@(5@;@H@%>@wA@yU3@N.D@2-L@jC@[K@f8D@J D@`" P@@"G@ yD@B@Vc5@A@OE@K@O@@xW?@(M2@9G@'oG@E@kc:@G@93@ܛI@5A@+A@TQI@[C@ [B@d8@V~F@C@j)C@fZQD@G@ۙ8@dA@A*xD@B0D@-|:@_2@fM@H@=%^E@Y\@̸?@ODAC@H@VB@:V2@@Q@~A@C@@@=@@`N@u=@B@l7@)S?@h7@x9@ 2@?@e>5@ $L@:@A@4@P@KyU@BE@W=4@7@F@߈0@P7@gBA@GE@Iv1@@C@36E@"%B@܀@@"e7A@F3@l@@-"H@K5@o0EA@GB@ІA@-;@D>@pmF@d#@2A@v}?@8O@DO@~N@9I@cI@tb node5-datahhKh R(KKh7B"?,?"?F?F?F?"?"?x?x?T?F?x?"?F?F?`\?"?F?T?F?`?,?x?F?;?x?;? ?,?F?F?F?F?x?F?F?F?,?`?F?"?,?;?F?"?"?F?x?"??"?"?x?F?F?T?"?"?`?;?l?;?"?F?T?F?@?x?;?"?x?F?x?F?x?x?F?F?"?T?"?"?x?;?F?x?`?F?,?F?"?"?T?F?x?,?F?"?F?F?"?"?F?,?x?T?x?"?"?x?F?T?"?F?Qr?@?x?F?x?"?F?`\? ?,?F?,?T?F?F?"?F?F?"?,?"?T?F?;?F?x?,?"?x?`\?"?"?T?x?"?;?F?T?"?F?"?,?;?`\?`\?F?x?,?F?F?;?F?F?`?T?"?F?"?"?"?F?"?T?T?F?x?F?T?"?;?"?F?F?"?F?"?F?F?F?T?;?,?F?F?x?x?"?F?T?F?F?F?"?F?x?F?F?"?x?x?,?,?x?"?,?,?F?F?"?;?F?F?F?@?F?`?F?F?"?T?;?"?"?F?F?,?x?"?F?F?"?x?"?F?`?F?,?F?F?F?tb node6-startshhKh R(KKhB #)/8@EObqz (/25:EJSY`hrz #+19?EJPahp{ "%,4<FSYbhoz #+9BFOYhotx $,27@IOUY`irv *29LPW`fqz%*-4:@IOS^hrxtb node6-stopshhKh R(KKhB #)/8@EO]qz (/25:EJSY`hrz~  +19?EJPZhpr "%,4<FSYbhou #+1BFOYdotx $,27@IOUY`irv &29@PW`fqx%*-2:@IOS^hnx~tb node8-indexhhKh R(KMhBp@  !"#&'+*(),-./0123456789:=>?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb node9-datahhKh R(KM hf4R(KhNNNJJKtbBD&s>Ř;>y@#?{?w@d@A@!?a=?&@I@bWAbAQEANA)!@?¨AA>@?B?ߓ?JX@qA?d)? ?<@S@d?@@eZA[?,>s?9MA(@,> A@ @+>?H׻??@.^>Ag><@@@ҝ@F@?@"aA?@ ;>5?E>?e@qTB@@"KR@±$Az%?G?l?Ĵw?n ?5@@K.hAx;An@sh?b?4]@>@е@0>NU\?}>3?.:@%??^5?y@P@ >@4@?~2? TA?Ъ!??D?J>7@0AUo@t$Au@@Z@SV@Ґ @Ui@@LAi(A6@??$@x@p?^.@p@r?&Ap@x?CHP@g@>?@?S?Ż?Lk?Iu??c@?>Aa7?NJ?0N?@?y@+_-@L(>T@AIQA?k?K?+l?r&?l[A7U>G@?Æ?@~?y@(8x?%?֟?\j@@!?6Ax@t?~>P?Oy?+l??{@C*@k1>U@+@5n/@%@A?E>UP??>NdO?'?x->E @ļU>?VH?U)@~&@6Ǹ?[>,@-q8@\@A ?4?n} @<@탯AiA-T@@h?f?\{>@@?mA4A@ˑA\>٬"AP8@M8@q@Af>@3@>@êABy@g{>32?>Vd#? ?)f@ŝ? n?XR?p1??vA9E>?\|>Aa A麰?^A4@AAR@W@ @!@0/@i??C@n@hj@ j@e?w@@ &@sA<ۨAgcAL.?V?|YAj>; z?*?@@Ru?rxAM(?>0?L@2?O?@?Sl1A?#9@@@pWA6N>uV%?vA?,?*l>ڟ?}@?B?Hr>\@䆰@|?}>@@,@C[@?ty@^E>@@C@">'n?A?@w@1V@,%@b/1AԨ?A@;@B@ޔ@$CA>>ɜ?;i@Nv@?>ŗ(@u.@,C5A @×>{@~@t&AY?1>AD ?S?M@Gq5@ AA@8W>?D?Z?>?ώ?(@B?@1 @-x@Y?@ebAߘA0?`{@@]?A@e^>7@9@JX*?w ?9??"i@2<@@@'7Apw@k?O?@??j@A@i\?N?==??%@5*@@A|U@@=Z?ÔS@/?>!@@@7@)_N@TE@0Ant>>~ A&5?t ?h@[@@nt@@y@?8?"?&?P6@~u@,?z<A@?>^>?r@e 7@W@㦅?@?ABA>Gwd><Ab?ϋ?d?A-)@HAQE?R?|{??[>t@՚?gW?- @{A-g>X?Du?H@ @@--?@5@FA$AC@A}0?ӏ=@.??>K?iIM@Afc@w@5AUACH6??F@>3?e?m??L ?>A\A>?@#@kY?JAtxt?@LA|?D@p@ ?-?4Œ@hAW@͂?@E.An?@J?Q AB<@WAA~^?(???=0@-5@@h?h?@-,@-@>Hb?6l?A?~@RI@@e??<^A'j@F&>̷?8??>h?,{?[?@I?C?ֵ?@C>@lAYrT>D[?l??fAh@X@-G???, @?NJ.@?ʜ_@Ι@A4a@8ȑ@r>5@?y@w@,B?p.?q҉?)?P?Z?C?O$AA?VA48?C>_>i@>C?@GBU@s3sAO@RaA]?|?.OY@@Ul|A&xA!?c?`[?׬?Q@.ȿ?Z`A<@kW&@A>@L:AΑ,@?&A4U?y??]j@n_z@՛@ R/A@^@9jAv)|?@+?X?O,?_.@?rA[??Od?O'@x>(?@Au?`g@?j@#@l7^?g>?Nd??+A!"@B?>qJ?@AA A>+j?(f@!O@UA+@lx@=;?z?ќ@PA>ѵAѵ?ܱA$@k@EAq?XX>?%?1?_>@F?'3?J>A)A}?RA >!>jS>d@ ANg`A`Mv@„?ԽD@??!?>Rh@a@#I@`AG(9As?lT>K? @ AeAnK?^]>V@1?Eڤ?М>e? |@}>@ap@b@NV@p@ _@sK@>> ? ?1A5?ec?6Q@ʾ@@d@|A@4@Qj?K>Ԧ>?#@Il?T?o_A4>@8?/@G>A(A5@A ?!9?0^'@*C@$?Z??/Az@A鏋?@P@ADA2>>??x?A?@7??x?݀@ ?f? @??a-@P?_.@E! ANm@45@H@Y@\>b>?d?s@?Z?@,@@ޣ@ ??[??]^A+AS@g>>/?8?;o?;>? ?@W @Hv?d=@eAm:?l>S\>y?+w?p?@sA#@45@@G1?:4??Z@p#x@K??T>?Kb;@@m@#?I?f?>a9@D@@x@s~Ax?@@L?J@;@@.AGu>>)>@@ϴ?@@aͽ?X@gf@^w@=@ V@ܲ@xt>wN?2?}?LP@ͳ@18@Rf@ՉlAGAa> A?q?=D?U@f0s@ A@,@qA+>?8?>?@! @%,@z@|A?@F>n9??r@?@?)Y@8(AA@QA:u?`q?߰?@ >?ҧ?n7@A5@5N@!9@j4A~AA>U?l'@AV>@BAݘA7?)>.N?|?2?u?qU@=@'@oJR@@AA(,8?p3@ @1o@i@A4A)I1@,>8Y>P@z]k?KmlA%Ag?M?F?(??nr??"@>m@/AA@ö@A?>?77?@l@O>5V??Rz>v@!X@BA?AG7?\X?88?l#?e?Z @yD@f@eL@\A(>ɰ?A @,M@6~@@?@CAE>a>+Q?ki? Y??Ғl@@?@~e@A>*?-r??\?'@rB@@@S@@vc?"?Q??9@x"@A'? q?@F?1?I?H@Z@ @2@M?cqH@0@ &@2.@u?}J>5@h:@WAdA?~A?T?6?1ݗ?P@2@G@S9@@q@9@%AA?.BA!]?+t>.?=A?gC?H:?@3 AΪBu@b͋@@@M?@?P>_?3@P[@t<@bJ@S$?q@{?.? @J@ ?@LA!AA >k?Wr?T??;?œ>??w?j$@<>@CM@ A 6@ Y>35>Y5@w?ٞ?k(:@1??&e@TS?]A!?a?ͤ?@e@?#A@%?,@ܴ)A?;>\U?F??G_?A@U#?7@MA?k?2?8@SÎA@C|?hJ@9?kG@WA@{@ ?f>j~?@ ?@L\A@AMA/@-6?>Bç?(A9@#A!AeAV@)f?Co?@@\ARM@@~A*@?D@5??Q??b@o\?8@*G@B@AxAX@?{@?@u@@WĶ@IFA@pF ?,A q9@߻A3? 9?>\?@'@B?q@@1A@Ak1AA0j??c0??04@j&@z?{?=?A賥A,?A>? ?rr @X?"M@b@NA>@3@BJAl?F?@Ș0@ム?@[2@O@@u-@0X@̝??F @@׌?&Mn?@?k>m@ŸZA9R??>'>h@!9?@u>?@:b@JF@:3VA @9Ao@i>{?6!@*f@@zA$~@,?`4?;?A-ַ>tMAj@@PAA5@?D@7o?@A@^d@<?XAj;?>@<@?#@ @N@@?bݾ@t@? @@OA ?? ә@ӝ@˿AYAC?qP?>K#@b]/??'???Jh@(@+@@nA^>zֻ>+?3?œV?t@%A@r@P@'?df'>M?+@@(A7x+@>O@@Sm@ܑ?@z?@3@drA>f?xBA,BA?CA?>nbj@uN?ˬ@, >=?@>Vd@A#@?1?vr?.Ԛ@ۻS@@To#As@AO '?M@,@]?<@? @u@?o@0A@?wx>'@Ýu?B?c@|@?/uu?d2?@?L?cK?*E?#1>Ѡ>g?;@^5?A@ҏA>Q?d[@VD?5~?U> @E?>XA_{AՏA/VEA(?5?ً?>ҵ@1>?{>>?4A@KAiE?QZ?h> ?y@ho>>/?'@n8@F@\o@&A@@.&q@ V@A?'@?A@+@`'A[/A?a??i@ѳ%@[[?5)FA٠?>!E@??Ss,?@?AAe>L>?fH?v?I?|?+x?q@l}@q=A ?#g? @Ȃ.As??_??;>g?VAu/AdA?'4+@[A@!@;&A@#??4@*?au>?7z@49&@y+?D?/@-V+@-A>6?z9@f@@R5@bgA2@54?F @Y?[@#^?A @ 0>>oy>>?>"@@@kT?q?.8?*$>|@t@@?A>9(?R]OA!g@V??j? AM@M8@|@AW~@گL?@DvA@$ @??O?27??S@z?G@ج@t @9Aw?B^@#S@H@@\?Aa1@I?AK?iN>?g?k@,_?>?e@$@vgAߊ/??h@W?@@?A8mIAX>?"C?>]۪@ @MA @K=@Ě?h6?@@??8 A" @ B? 9@A7X\Al@5A_N@{>zm???Ry?hu?R@l!?̺{@u@B?"`@@I@q@Ҕ?J@"y@ AA@TA0??FW?޶?X(@@E??@e@E^7A@R@??D?Sܻ>k@ӫ>|O?5?$p>}3@,@ڹ@@ ?x|@3j@wA(ֵ>7?k?^?뤫>z|~@@#A*NĄ@ w?n.@? @ @AB?<>S@+=u?:@Ҷ?%??;k@A?G9R?@3AN?@A1A=@J1@.{A@ԗ@zAW?>'@f?; ?z@7?@]g@-@A A4.AA?YZ\@;k\AOK>AA4@;AAߦA?\d>??/!@*3:@i~@L(@?$2@yA*??es?5AK:@x[@AѺ?1?ؼ?1<@$t?9@@RB@}@YAIK?^ BAVӎA<*?AkA?C>.?;@???t@uA_-A?8:?hU>{5 @R@&?i? $?@2?@]??>Qá?#AfA A`??nk(>7@H>N>:?_&? @A ?*??@ާ@V@K6@`?@6M?Hd?˙>M%@U?{F@o@KuA(@Ȗ? @_@S5@ZҲ?>@-e_A>Y@,}?].@ ?A"@MS@a@q%?`/?s@/LBUA?C?V@z>?6?;?3@ >c.@ _>@p*@x@O?z@p?@MP@S@[?@'@@`I@[K?Sn@Zr@Q@ж@#o?Dz?(??<@ّ?ٜ?U@{@@Z??H?l'@Y@3A{;@4(A޾>P?rAZ?[? ?X??? #?G@D?E@iDAΆ?C@T?tS?Q@Ny@>7?9>#@>>A״A֝g??l @?M>i??@/@??h>U?=1?x@@@nX@y^@g@@Av?-T?D@b>R?J@ARx@3?4w?_#?>C?@P@Y?Aи@7???ֱ@xw3@"@7>}|@?.)H@?P"A%?>mR@ah?Cw?>e=>@ܾ? A̛AzA?O=@&?0?1@@5ANAA>E>?],?fZy? ?7w@N@&@z\@@:A$Aca^? @7?遞>?=?8?:=@PH@i@]@y@!@r?>?o/?>0? G@KA|vB?pS@7??l@ڊ?6BY?@?R@6?\۪>Ϡ@Sc?@?y?]C@@Q?A?{&?Y@* 6AAAjAF ?7>?n?0$?H?n@^@A?l>P?Νy?@?>#>@@>$@P?bwA @>[@(o?@l?-@+i@L@(?> @to@;k?M@'>0@*X@@?l?^EA@~>"\>??K?f @Π@MAtM?Y#@i@Aڢ>`s>?> @Ε?ln4AEA-5F@@.?_g?Yh?j9?!6@e?N>A@A@vC?0D@@A"@.Ё?@R@(@VAt??~>P @1?_As?E? ?@@3*@@EA?(>?%G?1t?FӴ>D?RK>@?R? ?i@ vN@@a@B? ?ʋAAZd>%#>UC>?\o@TAM@?Y?øA?'Ax@A˼>Kٱ>)k?Ut?/ ?s|?k?$?I @W@r|@ \C?@A>?A"@&]@{?o]??%?5Aǐ?C;@G|@w@?{> @AW@L'G@'@+۔@B^kAA@`\?n?½?A?s#@?g?Xso?4>3O@@n?BAy?F?,4@s???3@?>/?@r @ Ak?#?n#@@?@0jO@_%@P@tAZZ??ֲ?M@l5?x?f@x@ AA^^@@I1@RKg@'xA! @ALV?W ?<>.G?o;?hu?=?ԄS@Gf@@9h>@?\?s@"yAm@DAtb node10-indexhhKh R(KMhBp@  !"#&'+*(),-./0123456789:=>?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb node11-datahhKh R(KM hzBD&?*?)???0?X?P?X@;a@[[@mi@h d@b@ϥf@ d@>Z?>ݬ>?4p>>,<=->Q{@@e@F@@\ċ@N@ꈍ@v@B?@@p=Ii>{I>4Lb>W>JJ>x>Y>u>V>]>g@ю@@@։@@ʍ@G[@@Ϲ@@PM@|@@(X@G@8:@C@@ @x@@䡛@[@U@j('@k@I&@ȩ@%#@@c*@=!@z!@*@Ue@[t@dx@=M~@s@E@߀@@@S@ @:@@+@x@g@/@֓@Վ@ɹ@@Y?(?Jo?V?s??=د??$?۩?!?ϩ??;?s$?;? [?a@g@A @@@1@ӄ@ @/*@[W/?e/?.?b>Kz>> I>D ?:z?g>)F>6>!>>D۹?u??|i??0?%? @X@V@L@[a@?'@c@@,l@,@=@g@Ͳ@+K@I1@@OmH?/f?~ ?8q?'??G? Y?C?@@ȴ@e@\@I@@ U@W'@@@^@F@dw@h @G@O@9@_@@U@@2i@@@G@@@'@@@a@L@U@=W@{@U@Ԩ@!h@@@ @Lӆ@Е@Ղ@@@O@.V@@@ <@|^@W@o@j@]Յ@p@El@3e@e@)h@u?Y?%[?W?}H?,??@ۈ?% @@D@?ŧ?}?E?W?}P?2 _@@(s@0x@o@`v@ir@*t@zs@>u@|@Lc@U@U\@\Ha@&a@_@i@0j@ok@qk@kj@k@t@ rn@S}p@K o@S@i@>@"P[@@b@l@3k@k4r@~q@R@3GJ>pT@q@>Z@K @>x@/B>< =X@P@. @x@s@@t@U@@q@ }@(@@@N@m@1R@zp@Ѣ@d@yϢ@@< ????e@기@@ &@ Z@2%@@@J@"@h@eG@I@~@(@ₓ@|@NJ@2ϐ@z@ď@F@PD@@@@{@ {@i&{@ p@u@m'>@@@@@@_n@#@@o@ҳ@ض@L@O@¾@k@B@@@@k@T@@F@B%@*#+@0@!;:@w4@*@&@ @w@@,@<[@֞@@@?]@= =Ty>b>M#b>d@hi@7WA@FY@HU@qK@UJ@O@/5P@<٢>*S?g.?r?>"W>N%>>T=@z@ͻ<9Q@ @6V@P@ŞP@FL@y9@H/@]1@Lk1@X{@逎@n@25}@@8@d@y@٨@Qq@f@eo@ N@k@ ɋ@ъ@f@@XM@kt@c@9r@_*@I@@q@I@͆@ш@B@C@@ @@@T@Y@~@B @y@lSy@:}@S|@@8@n@!@@+@E@䘔@@@5׌@o@x@A@@i?[@@@ @3?r?#@O0?=?$@A@?@1@0@"%6@vh?-dJ?\S_?=w?1>?:q>(?H/>>>>>ǽ@@po@κ@@@ @>7>t=?n>>H #?r??7#?,???nU4?ףT?=?B ?Qe?h?f?:i?@@@N@@3@YX@9\@ٱ@ל@4@w,??M?Q"?ݾ?}A?0{@?>}>T>?3U?_?Ǩ???V?N?= ??Z?8,U@'I@ޛ@ @|@oњ@Ϫ@@f-@@ ؓ@.,@@O@!@ @@꼘@n@d@TP@|@@o.@@)@@g}@QN@T>q^??2??d#?T%?ˡ@N)@@iͲ@@p@@L@H@S@IV@U@nd?!f?'H?F@8@i@Tb)@z@=%@cz#@^"@ @z;"@@ @@ڋ@@v=@?X@I$@?'@DY@)@Vc @@ M?z??.!@;C?_?k?U@D@@@ǝ@㏝@3b@B@F@M@a%N@N@N@ K@)@#5@52@j0@5,@o.@i_@gA?@&U=@<@9@!+<@9:@rl@R\@rs@r[r@a@\@ c@A1U@P&Q@_@X@RQ@EX@}S@U:@_|@@#{@V @oi@,<@6%@@s@E@:z@y@y@6N@Wu@@pz@hd|@z@/~@U@k@{{@/z@{@)|@?od?t?w?B+@,@?1?g@. @)5@@@ɛ?d4?Y?r??0q?&@ϰ$@o]T@OC@C@XF@qF@9K@z@AZ@l@j@Dl@}`@fp@zj@f@Xh@Dv@Ul@ d@ř}@q@0-t@a@@@Cţ@1@Ö@ߠ@@ޝ@R@R@|=k=h;>?>l%?P>}!>g>E#?IV?޼??|? ?#?A?y?|@*@ ̋@Ǐ@@Ն@/p@ss@@׌@@&Ɖ@@o @ӊ@?s@\@{$@Qy!@@@3@@B@f?@z@ @<@@Ap4@?@p@ @Y{@S@q@x@sA@C@=@|@@ɕ@w@*@,@NǠ@#@ڟ@Mx@i@y=@ي*@?@R @?=A?B@>@@90@ʅ@@u@y@I|@z@$Tz@Ri@}@w@@D@m@?@_@q\@@d@E@@@@ރ@Z@q@RDs@+@zƀ@|@!@UE}@Q?X?Q?ĉ?ؤ?a? j??j?8?\?HX??:@@8@׺@@8@58@X>@=@k7>5>ƹ<"=?G??u?(?xø?Ti?>?M?n?D??1?ƅ?k2@G?(@J@i@/Q @;o@@@@H?ZJ?w?z?a??{{?n?f?$l?@_w@h@vG@.gc@WUW@2Y@Z@1[@Z@?8?j>?%u???Z?0?4???E?!? y? ?o?@??K?~?LO?:?Տ?t?C?w@m6@޶@ŷ@d@@@$[@=u@OE\@LRm@@SZ@\@]@a@W}a@ g@Cd@?U'?'?b?`?t?N@}@~D@)@^@e @@O@\0-@X!@M,@ (#@@D???h?[??o?ϐ??u??y ?#@'@ڋ@ӆ@lݘ@T@@) @@@U@d#@Ӷ@bζ@G@@U@8@@"@ϼ@2@ 4%@ @xJ4@/3@84-@>%@l?v=d>k^> >+">e> ==o >_>Q@}@w@Ӏ@$@}@(@S*|@@vSy@{@L@S@c@lT@__@Gc@Dh^@g͖?#??+?B?&c:?lI?r@?^A?D?wn?g?S?-[?s?t@?@2@a>@>@{B@@@fM@QZk@YB@mQ@!Q@Q@b@X_@\@m_@\@[@n\@bH]@L@W@T@cF@;@IF@[7@l7<@f>@<@9 ?=? E>>I@@6@B@pU@;@*@si@ @ E@m4@pa@"@9@ @b@/@\@@2@!@S@Э@0%@?@t@(@^@2ϧ@xu@+@@4ڠ@] @8?_???ۈ??ew?9~?e?nn??+?)?{ ?>Ad>T??v*?l?RF@"T@C@. S@\@#N@NN@a@7U@ 4R@@$IB@7@Tl0@3@2@ξP@@@9:@B@=@ =X=ѐ:?h>>,&>d>i>HG>i>Ni>?@?;?i?T???!?V?>P?jX?_I?k^?&\?j?oa?ԏ@@a@$@@QΉ@<@@4@0r@p@;|@Uw@tv@˸@p@@`j@Aa@t@'@@n@g,@\L?@*nL@|F@ F@?㒬?{b???c???)?g??~??y=L@+)@(@E@w@@@@Լ@@j>d@k˼@9S@Ŗ,=@7 @,@]@=^@U@@R@a@@Og@@T@/ͼ@t2@Ⱦ@U @>d>َW?, q?Q?3 ?(J?aK?$s?q?{S??o@!W@Yq@s@sm@]n@o@q@F@r.@ܗ@b@z@@9@m;@@@Õ@A@|@趻@@$?J@?>F? ?Њ?q@Xq@o@Em@Un@S@@@~4@)@@H@q@$@>@>[>2>&?U?#> >>d=C=[;<;@?Z ?:?$??jW?9>@6@O@$@W@'@!@@F@c@!$ @H@-@@@R@HL@Of@^@$Nj@w@,p@[p@Yfp@.V@sM@՟G@N8?@<@;@ϩ@@nҟ@ @ @e@)@y4@|@@@2?8?l?Խ@a$;d@5>k >m >df=nE=%>@h,>>4> >=Ll@+V@\@V@sIg@`c@W]@|`@]@}o@ќ@@L@/@@ @0@C@@SZ@7r@Z:?u ?@?%?T? %l?%uw?{d? ????D??Sɛ?f@h@/@qê@t@u@i@@?M=? >^'?@ )?;?j>?O?[>c?#?q??@*ɲ@oa@qd@$I@=@v@(@lЮ@[@ @2@`0@?`@~ @{{%@$@T@w@@@%@=N@=@@@0@@A@A@7@ <@.@e//@W2@3@8@I@_i@Gh@_L@MVb@YX@x?Y?L????6???@f@?F?Ojܲ/>H>IF>;>>j>{>mɣ>#^@@@&2@]@s@~ @k @ @c]@I@F@ۤ@@Zi@U@i@zn@@.@: @),&@@^@7@@M@ls/@'Y@[Y@l8@F^O@B@\F@9*@.@K@{dN@A@[D@dE@ C@F@&@v@@\@m/@'A@!ԟ@ޞ@J@9S@&g@#X@U@V@c@dY@†`@\]@]@·@of@L@M@@;ɲ@>@h'@o8M@(.@vK4@k$@ӿ&@@@=@8@ѭ4@Y6@yH@0@?2@>@B@@<@R`7@?}#??ՙ?4#@4>@2@:@9@J7@w2@Ys4@?6@]s2@C@E@ID@M@h@@ҍ@[@@>@)@@@@@@l@[@-@dé@8`@Jf@D>__?+4??M?D???>;?T2?'?.>@Ck;@7@G@8D@K@G@pE@LO@6Z@j@@C@0Q@W@0@YP@&P@)@@@Ρ@uk@/@g@JO@~@f?bt??x7@NK@r@@c)@sb8@7@1('@)'@ 1@/@5@V3@-@/@">l>>UY>^r/>@>q=w@_U@n;@u@竒@@o@ɬ@QR@j@@f@ @v@@25@ ?bz'@@.>@?@!:@S<@@{h'@+"@۬-@C`'@RV+@M$@k$@S@@6 @@i @:@ne?8&??Լ?W?v?y?"w?ҧ? @d@Iα@^@I@9@`@b@@@r@@zή@@@mo@n@oT@tD@Tۥ@Ҩ@ɚ@t@;C@̡@@@C@@9@ƭ@ۖ@/d@\@ז@}@w@95@ц@n@(@Ɇ@L!@&@9o@0v@h~@cMv@s@Vy@Jg@z @@Bz?@h?G?J?"?y?ig?0i?PP@@S@@,9@~3@ma#@Z0@T+@Lv,@V,@U@=L@pxN@*O@q@?@ @I?@@S?%D @@[@@d@ @@+m@@y@Ʋ@{@0@l?/?q?l?T????d?~?4??q%?U?,???Yb@o@ @/@1[@@C@kB@*!@T?`@ @"J-@>p@)6@[@ɢT@,T@L/W@?U@\I@M@\@5==K>$@=e{=S@0?e@4?XU@Mx@9q@==@ ;|W@e=@<"@D`@=?@S@)?$>I=TU3=Ҟ=$?t>+>>}>p>M>^>@?????????ח @F@Q@r?nl?Ⱦ?8#??A??y?Ҭ?8?@ke@ o@|f@r@9~o@>}!@@@p)@~.@j,!@ҙ+@--@ '7@!+@w0@.@\/@x@@xH@/ӯ@@@r@1@@@[@)o@@ڦ@+@ڝ@wv@,@<[@@3@Sƨ@E|@:@͟@i@E@٥@@ @@ @ū@@b2>ô@{>q@@' >xC>^)<ő0?+! ?_G?zp,?@,@O>@}!@@Ю@nP@Z@„@r=@l@@E]@@ b@p@N@e@Gl@gp@h@gk@#j@/ @qD@l@! @Nj @|?߃?%γ?б?{??pg??W????@߲@u:@W@@@&@@@@~?T`?9?l?r?Q?O?Ʊ?Ug<>`]}>3>a4@n5@b$@@ݭ@8>!>Ї>۴c>g`>r@ރ@M@Ah@m@l@Su@;)m@P"@?&@@@@QK@@@@@ܻ@@ܰ@6@↴@bܷ@I@@t@@/@@@?>T@IW@[>L>>='@S@@@L@@ @@H@-@A@P@P@S@T@0T@N@o2@8 ? > >>m?> @_@-@@:-@ @@du@F@_@PZ"@l@K@@)@&@q@Y@{@@7@~@V@K}@@{@&~@@@p@wՠ@@)@.@-@j̢@@s@>@@L@zo@@~@t}@b^@f@Av@j@Nm@p@>q@Fs@*%?b?Ԫ? ?T@`@l@U@Q@;@M[;>3=o>ө>>޻>ҁ>4@5@ @y_@Zc@@@߻@@o@G@B @h@k_@.ʰ?N?!?.v?@!K@F@`QQ@D>@]);@~I@LC@C@@@K@JG@g?@/?@ B@p C@>@=`@@@G@w@$@}@ɋ@"@ߋ@8a@@g@۩@ei@tb node12-indexhhKh R(KMhBp@  !"#&'+*(),-./0123456789:=>?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb?5 node13-datahhKh R(KM hzBD&g'& nB׽7,7*; FFg#6󠣾{/z>9н IB nrq̘~\S,mfX>+sݽ?c~>ql>m>δ>܎>پ-[Zlɖֲ kgXtnvgĔ`D#>r5CaսMN>Q9H=En<m=>@>8==r>h >A>~> >> =;a=9TU}2qV|?pXx>>[ر><?" ?~>HH?>_ʲ>>a>>>t>u@='=@=]>E===k> [=9!7|ԽX<꨼R˕ͽjQ=+2= ?+>͋>>g>w>/>$6>mO>Yɽ>@޽\>x:>g=1=x۩=,>uN>>i7>E>4TX >$:v.[;Hb:>oՀG4> \ >Z}+<l>N;Ӥܾ\.FB<=B׶>">B->zr>Qߚ>Þ>.?ݫ>Z9>v>K>¼zB==9>==T>ס=9ֽj@낽(]C%4r F3K>1) ս?+y>F>c6=>3>ػ> 0>R>R>> e>|s ߻֬W==% >=1 99>+&NgTOuLɚ;'輵k>Q;>1}>Iw>>V>D|VIжؾQ^<@fe$R8P*>i>>aS>I>[>h>2o>&=!H=>*B>!>2 y>(Q>G>l"> > >*>Ab=sex>c : O3 d{Ϭ"=w>8>ʸ=澳jL¾>=ĽX㍼ѽ½m<F#ཬyP݇>Dgih2Mkv>;$>,N=솽X< ¼\D1Ͼ?R<씽Hc'*Rx5t[rܾ>B9O^==\=^,>) >\e;L<+>h&>y==B?Uo>.{>?ھ#^>%QU 'V~`߽&Ƚae}=7 = A<.; 1ܾ*쟾YLjՎ*N>U>>e#>>>dO>m>>+P> ?QD>ZoбmxB>6>T$>m>>>s/>=B=: =P>&&>T!>%>ܭtZ{ cxcG}뭾7(5&,>"̇9q ؍u1<"S I ?I?{<B6b)$>6w>s>G>(ľF8;o<eyص@9IA$i5Ӿ!˽vjྜӾhоܜ>΋e=9>=k>T>g=U>>wB?:>>?*>> >P>L>;ؤR=F~=$u=Q=]H>:u\XT>QA>~#׃˾`=FN > 霾l`}% k?/==->>XŐ>yqi>#b>bm<3dFf#=CP!ټ xн<㢾d;//9/'e>彼>ff"u"=J۽SxdN!N==$-\H?+$>ؽ|?=J=-=C>>ó= >>f>t=u=k)=?==M?=Af=ݍj>1 >Y>;>%WC(<*a=F`<2tOʽͽ 3ܾ>W*?T?y >>d+>=<פ=m=<P=/.X6*&->rqF>; Δ։4V<~<_u>Y;> >%4>عS>ӺQ3>\T&6< ngH?MJ]-XӾ‹%gN>.===1?j;>x&oc>&>>m$>\Y>|P>F^>uM>EݽYqwmR#L?w>a"u>x>Q>A>@>g!B>gpbS1(;M4- y=HFsWbJcE?,?!=&9>)>j>>&f$ܺ>^=>&m>D=V>^ ??Q> >FS? ?V?> >3 >F>\>^pi^KA\xEXs#a贾+>˩>L> 4ov>J_ڽI5SU?> =?>l>*>Z>Y>D>W߾? "c]gWuD+l~W>2'T /?[?(>={Iy>_<3=;0:>]1>H=q= >v>תo>^s^>vIK>C>>sY>_9xg=zbDsZH $V!;(>`R>>O>:T>a>2/>=>O:>,[>ֈ>>8>->>䈙>=+@> =q5>I>C>ƾ#tNu=bIA=ɩ=m<_JcxU߽q?]>O=d>˚>OP>a>T>BKG>-q<3;$&>>7>>=(jshvZMr5.?=Q>X>E>f=>m=?>,>>>EѾd^;MpM=nD+sFо־u Ǿl9a\m=i=ýAϛvdbD׽L [Vhz%pNc>}>ڣ=3.>[>Y?>>c>w ?QIe)В=c,l=E?g>>HzK>!4>T1L>G1>?>y> >Zi Ц<(V|=T'=;r|=y #=n?jr=\k>TfP`j|⽫FYgvBQu৩ H̚w`oW5?sD>6<=Nڼ>>5|>$R>N>~\i>\>^>F>>Ծmʻ0>:J==V=4YY澙=⛾ L];CϽ?^Lz[<Ȕny1-?dȾU6s]= =3yc[>~ 4>D=~?^>> >x>z=g=>J= =>x>r>C۽X3>Qw0>9whkχp6c/K{)>ɽ>UݣB==ȨqU<9,=đCŤ=b= =X40zXf@R5纾 ɾ>*&'&7q>2clր>W>>=OX>>>>"U>>Q>ʥz==*>/f[=Bg)A=P"88P6M;==;L o>o *qˑfG ) kl>A[Ƣ܇'9磾c&|%&eƇ2`sAp?n>c%4B }5=:G==|p=X=TMn\?=>> $>2߼] >Hy@A5Υľ6ܾ؎s i9_Ϳe刼^iFyk*GR\_s$}$1 @.Xrs$M33jnvxy/R?nxs:>t>Gp >lx.>>[^>}>A>;T׿C34QIH<'͆]IѼd$eR==ϲ=Y>=>> '>jS>PW=nѲ>l>f>]p>6>Q?>i&O>k{" dv 8>j!7>=`R=x+===2:s]/^NԾevE0<+BWß< , ;E=">2W>-w >>"m>`&>)>%u>*Aj9l@rMIQK˽Jþ*)5>{^ Y>X! >==a >3>9>=v&u>>C^]>E{>n>$D=!=E<<2Y>si>>>,>s>->ߐ>gL>^>Y>;tV!> I>>;k>cF,>ix>5h>˚>>dJ=E>r;x!=}b > ~=-/C1>ۃ> S<=^=̴>S0>}*>4 >=5>~L>cBG>]a>KK>| >>/6>Y> *>=_C>XB>>n>(A>{L>6L>a0??@>ɻ˼4hM`>p>a>ߏ>l=C쾦|>+Be>bS>J<虖=޽߯fL=0޽N:r~v 5 j>94=%Ou><>|>.>|>Ђm>]n>"Y]>u>`>xg:|@=>f,k =^D;!=iud9=˦>9>v>~>Q<>-sF>s7=?>=a>UU>QX>D>z=mW >"x=ҽ=@g=GKھ6D>qCxll=U"_1;Ǿ>˾I|H/}>m>վfov>uB/A>qr>V >]>X+D8;zNVd`]Ͼ~$}]oÜjO7>pr> E>=H>{ >5>@>~&>Q͐-=)=Yu=9=+~>.)>Οɽf=b>䠾l@r|م?--:Zd?>=j>WH>8ͥ>C˽.?=m ?xC;1G=w 3=?=>Nb>d)G>,>Xq>)a>=َ>rՈ==< = z>>=7'QV6??>">h3> >‰>̌>>>Q>4=>/n{ W?GMx=ڧ>6m@&Fn>,=> >@pYG/ d 0jރuRD@V+?n^C= g< +=#= >"p>4ITZ>K>P>D o;en;'=IL.O=Xf8sOh^tɾ->Ge>->g:>t*+(˽l^)z ?yPrto >J_e4=Zz=>̤>s=ʡ>tڃ>qv,>=v>=v=4١uԡ[p.Ilլy[oV>G>z>k>{=> >/>vB>L>t>'= ==$eS >p.=>>RT1>>kD>BV=wH[ã&=>80x2J2딾a9>haf>u= 65)ꊺS+=JV<a>L<R5D?x>V>|%>i>=D>0#r>B'>t%?>l!?Y<"`7_>&Ք==K~=L\>!> G >^>=>H)MT=Z ,ƾPp NfNt~۾V> ƻ@(ⅽJ>hsu>ت=y>H>h>+>z?>@1 ><s~uh*nӽ}H'HD ^Ā/X>kT5+r=E-=7=v5=";;>Z=o=-=C=1Θ=~>' /> >>O > <=^LR=*>3 >=e>>1=кˀY$t"ɽO~D2[ְ"az[v>ɾ"2\ ?eo<~5KHȽ_ztL7ߕ>/E{¼Gm=ۻ,н*wч.=>=~9=y >>">^%>/i@q)u=DҾ-cҾĥ0s-ok՜σS>=.Jԧ^OLoXk>><>5bN>z::>.=/><CFE[?=נ=%fC_ƽQ<k?>/vqn 7nԛV>&?'>o>L>?>/>6>>mʾb>Id$ T!ݎA3)j_>fi]>q|h[5Q< ξþbI1~h瘽o2;n=8>U>[}?>d0%8_>S>6=Gҽ >eբ>"=<>=Y>QJ>'.>aо>v,eV.5+SI+>{׽Lｺ>,v@ý>w2>a>+>/>쿁>aR"L.A=W?=Q/m Wkr#KY־ƾ =ܮ=xd}Qv=%&<[>>Kd=ma:!!>߾Wd(>W>o=fɺ=>T~5>>=Nl">>p(?e B=y>˝G>6>K>z>Jν5Aj>#fyLͽՃ>=Ѿ{dF\nnm >>Q;fNe)I?8<3>@>,>p]>8> b>91>z>6 gCVqD;?>3=>E'>#>G?[><=ꪪ<=Sk=Q;B=lR=i'rr*M5;,olQdCf"|e=p>9Y{hsojᾆn=i̽Rfn;aܾU=LE>N,>>'>m<>[[OL 1`Pƾzq v">Kkz z>r=Ƕ=F%Wi=*>gUsf9 =h4a='X2wʽo3u㽂&ھ섥Wwȁ ur+fQo>l;fr fU7,'D \ &ϛ>>Ȕ*3 Uz9Fɓ >)A?}h>F=hU>5Y>D==.o=h6>иHɈ;G ?k>`O>94=>hC>=]=aa=bߖ>'t޽ > > Y>j> 6>>A> 7> ">P#>j%=T%&S>3>wm>3> >hD&=={=ƃv=ה=F`J>3r{L>qT>۔>,=>4>0+>CK=2>^;>)<==)<|7==];Tن]iH 82N/y 7>]G>8Q> x6Ѽ=f=pw>Sz> l>ƺN>r4#1f>|>>>)=߾0R>T>^^M=g+=P|=8ͽ "G,*M }?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb node15-datahhKh R(KM hzBD&V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>V>tb node16-indexhhKh R(KMhBp@  !"#&'+*(),-./0123456789:=>?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb node17-datahhKh R(KM hzBD&@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@tb node18-indexhhKh R(KMhBp@  !"#&'+*(),-./0123456789:=>?@ABCDEFIGHLMNORPQSTU[]^_`abcdefghiVYZWXlmnopqrstuvwxyz{|}~     #$%&'*+,-./02534;<=>?@ABECDHGIJMNOPQRSTUVWZ[\]_`acdefghjklmnuvwxyz{|}~      !"#$%&',-./0213456789:;*)+>?@ABCDGHIJMPNORSTUVWXYZ[^_`abfeghijklmstuxvwyz{|}~     !"&#$%-./01235467,'()*89:;<>=?@ABCDEFHIJKLMNOPQRSTUWXYZ_`abcdefghkijnopqruvstxzy{|}      !#$%&',+-0./23546789:;=<>?@ABEFGKHJMILOPQRSTUVWXYZ\[]^_`abjkndcelmfghiqrstyz{|}~      !"#$%&'()+,-./0123456789:;<=>?@EFCDGHIJKLMNOPQRSTUVWXYZ[\]^_`abedcfghijklmnopqrstuvy|}~z{ !"#$%*+,-.359:;<678=>?@ABCDEFHIGJKQRSUTV[\]^`abcdefhijklmnoprstuvwxyz~     !"#$%&'*+,-01324569:;<=>?@ACBDIJKLMNOPEHFGRSTUVWXY\]^`_abcdefighjklnmoqprstuvwxy|}~      !"#$%&'()*./012>A?@=:;89<CEFDGJKLMNOPQRSTUVWXYZ[\_`abcdefghimnopqrstu{|}~                      " ! # $ % * + , 0 / - . 1 2 5 6 7 : ; 8 9 < ? A @ = > D E F G J K L H I X Y Z [ \ ] S T U Q R V _ a b ` c d e f g h i j k l m n o p r q t s u v w y z { | } ~ tb& node19-datahhKh R(KM hi4R(KhNNNJJKtbBD&2 ,0 /;9?>@AHMKRQUabdgILJf\hcfgb>@^h]iJUT VPWQRpifj nmkg}~ouG]#qls grCAYK FJLGEM IDHCKBAJE> TXRQPO[WY*46 321WOSQVh3kcideB79 : ?@=< 67?; =DEy^c31\ac\b`a^=,.5YKOMOK NHIMV`^[\_RS^\L]UY_ ZDG'"}MqrXTVJIFMQHN=<?;@W>HKQPTO S+-%,*lmfgicbd6?<=74.23+*-]S^/7]\YX^ <B > 0/136248kl5+*6-/21X$'\! e(,)-2d+CEBJ?:<AV_R0Lrnq'652-0/ oM HFKJ ()_jablidkfgVZE >;<@B*)31,0b "JFIGCED?GEA>FB89k< 7l~]\^uv$#a,*.QSO\WYRZPU 5\XGMYRTNIV\F`fcGFJI!|qw B@) 1-,U.T}pvwxZjbehg`I8us6r685  b;iBACZCF EDK@BClJ=<E; !&(*+1,0lkdjn i 68_ r YVkMhlZim6xtGPFLRlbk; a:8<>145z ZXta]^Y/r.q#3-2z(+&)*y%214 03():&' * ,@B C;A9>:C =BD8?GHcd_\SPRVd fg`A5? 394:7\[]hbe_dag@C IjhFGBkl?63524bzhidj`usch imk ^) #$'" nDE LN@ZVXTP W $&'b%+-c,-)$+6Z?E9F> C<@]`\Z^DA2U34;.72ONFG5;>?D<a^mcejhl/6),$! #<^]Y[\e6<@=;>C omGEpnIHJQTULP8:51*,-.+;7%32FHKI89; G< FA @CD?=B\X WSUoFHWVLRQP)%*+(mn.:1-($#6%+*&H=B:E?@< @7; =U:9<qr6l=g78fdgYZ mr 6: s05  1D WTUR P+/2r3qBDEikhGjJIVLKPTQ;8 9 53 6gk_e^ij`%&103246 J <:A (E)oEpDFCoclkmhj"/NHGFO%.*')UQSW @>?8;=5<|~|OLUgMOjLi_h[]a)20:1}63!J%- /012KC>BJEFGLHSVPUQO?@2;: =769lpfs]dcmi7L8 53MG8 6 \Dd54eDiCGJHQPO-)$' VI H BCUa`^ 56897/qpe+WqSpLoPN[ZRY*O,('/&$#-0.+\5 ]IXGF`\caWUpqOSrA?=jwD{ExV[fb_WXZc).4/ > <p5o3?6nqFqpqpikegd6 pqaTX^Z][(\+'[`-ZOG J a][S_UY"&r*)pigojfkmW\XLNUZSVqPXrT=EAB qvumplo^; q>:C@ADGGURSJPLHloS tqrspmn{|~YXWV|zcd kijlm(*d)c +heab_XYTPSUI5:D3=;C9^`(hWbH.6PK; :*eij3,0154'$,) 184s31*05t)-.2W}9481NwL0S45`~\b^*,-/) rAF E gcXj0e_[`8 :><RUSPOWdZ]^-+{z $ ;sEMAPJDG\! &(qm9>A1 ;2/467'0) +b*.c,-8; 9=> <zihea^bg9=?.d`ebjfh;=8796:O KHCG"),$# a1^ /0\_]o nljh, 13/be RXWQT$D3E[QRTlYmVSXwz}y431*,0.DIBs=oSTX[@GcFbnCutJvEGImaqochrikep=C&')8 dFcY^KM}~P 344 0.1@3B6=:>79< JH;?KAD@>MCIGFADECAHEJ$J&+1-3(9 t235 7 ;@?-,+56273.01/4:*+ $&(*./0796jc\ U_d `TXY[tbu.vector-1.6.3/tests/test_compute_features.py000066400000000000000000000262111503546127100211560ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. """ Ensures that new or modified vector.compute.* functions don't break any existing or future backends by using unsupportable Python language features. Compute functions are highly restricted, a least common denominator for all the backends we *ever* want to support. The functions themselves are duck-typed: arguments could be numbers, NumPy arrays, Awkward Arrays, and potentially TensorFlow/Torch/JAX/etc. An ``if`` statement on individual numbers would have to be ``np.where` or a masked assignment in NumPy, so ``if`` is not allowed. JAX traces a function for JIT-compilation and autodifferentiation by passing a "tracer" object through it, and that object can only follow one code path, another reason to exclude ``if`` statements. Loops are even more problematic. This suite of tests statically analyzes all of the compute functions by decompiling their bytecode with uncompyle6 (on Python 3.8; will have to be modified slightly every few years). Some compute functions are dynamically generated, so they don't all have an AST to inspect. The sieve has been defined narrowly: compute functions can use more functions, binary operators, and possibly more language features than are allowed here. Expanding this set of rules is therefore allowed and encouraged. The test failure and requirement to expand the rules is intended to force you to think about new features, to ask yourself if they can be supported by all current and hoped-for backends, and whether a (formally) simpler implementation is possible. """ from __future__ import annotations import collections import contextlib import inspect import sys import pytest import vector._compute.lorentz import vector._compute.planar import vector._compute.spatial uncompyle6 = pytest.importorskip("uncompyle6") spark_parser = pytest.importorskip("spark_parser") pytestmark = pytest.mark.dis Context = collections.namedtuple("Context", ["name", "closure"]) functions = dict( [ ( f"{y.__name__}({', '.join(repr(v) if isinstance(v, str) else v.__name__ for v in w)})", z[0], ) for x, y in inspect.getmembers( vector._compute.planar, predicate=inspect.ismodule ) if hasattr(y, "dispatch_map") for w, z in y.dispatch_map.items() ] + [ ( f"{y.__name__}({', '.join(repr(v) if isinstance(v, str) else v.__name__ for v in w)})", z[0], ) for x, y in inspect.getmembers( vector._compute.spatial, predicate=inspect.ismodule ) if hasattr(y, "dispatch_map") for w, z in y.dispatch_map.items() ] + [ ( f"{y.__name__}({', '.join(repr(v) if isinstance(v, str) else v.__name__ for v in w)})", z[0], ) for x, y in inspect.getmembers( vector._compute.lorentz, predicate=inspect.ismodule ) if hasattr(y, "dispatch_map") for w, z in y.dispatch_map.items() ] ) python_version = f"{sys.version_info[0]}.{sys.version_info[1]}" is_pypy = "__pypy__" in sys.builtin_module_names try: uncompyle6.scanner.get_scanner(python_version, is_pypy=is_pypy) except RuntimeError as err: is_unsupported = True unsupported_message = str(err) else: is_unsupported = False unsupported_message = "" @pytest.mark.skipif(is_unsupported, reason=unsupported_message) @pytest.mark.slow @pytest.mark.parametrize("signature", functions.keys()) def test(signature): analyze_function(functions[signature]) def analyze_function(function): if function not in analyze_function.done: closure = dict(function.__globals__) if function.__closure__ is not None: for var, cell in zip(function.__code__.co_freevars, function.__closure__): # the cell has not been filled yet, so ignore it with contextlib.suppress(ValueError): closure[var] = cell.cell_contents analyze_code(function.__code__, Context(function.__name__, closure)) analyze_function.done.add(function) analyze_function.done = set() def analyze_code(code, context): # this block is all uncompyle6 parser = uncompyle6.parser.get_python_parser( python_version, debug_parser=dict(spark_parser.DEFAULT_DEBUG), compile_mode="exec", is_pypy=is_pypy, ) scanner = uncompyle6.scanner.get_scanner(python_version, is_pypy=is_pypy) tokens, customize = scanner.ingest(code, code_objects={}, show_asm=False) parsed = uncompyle6.parser.parse(parser, tokens, customize, code) # now the disassembled bytecodes have been parsed into a tree for us to walk analyze_body(parsed, context) def analyze_body(node, context): assert node.kind == "stmts" assert len(node) >= 1 for statement in node[:-1]: analyze_assignment(statement, context) analyze_return(node[-1], context) def analyze_assignment(node, context): assert node.kind == "sstmt" assert len(node) == 1 assert node[0].kind == "assign", ( "only assignments and a final 'return' are allowed (and not tuple-assignment)" ) assert len(node[0]) == 2 assert node[0][1].kind == "store" if node[0][1][0].kind == "STORE_FAST": analyze_expression(expr(node[0][0]), context) elif node[0][1][0].kind == "unpack": assert len(node[0][1][0]) >= 2 assert node[0][1][0][0].kind.startswith("UNPACK_SEQUENCE") for item in node[0][1][0][1:]: assert item.kind == "store" assert len(item) == 1 assert item[0].kind == "STORE_FAST" else: print(node[0][1][0]) raise AssertionError("what is this?") def expr(node): assert node.kind == "expr" assert len(node) == 1 return node[0] def is_pi(node): return ( node.kind == "attribute" and len(node) == 2 and expr(node[0]).kind == "LOAD_FAST" and expr(node[0]).attr == "lib" and node[1].kind == "LOAD_ATTR" and node[1].attr == "pi" ) def is_nan_to_num(node): if node.kind != "call_kw36" or len(node) < 3: return False function = expr(node[0]) return ( function.kind == "attribute" and expr(function[0]).attr == "lib" and function[1].attr == "nan_to_num" ) def analyze_return(node, context): assert node.kind == "sstmt" assert len(node) == 1 assert node[0].kind == "return", "compute function must end with a 'return'" assert len(node[0]) == 2 assert node[0][0].kind in ("ret_expr", "return_expr") assert len(node[0][0]) == 1 expr(node[0][0][0]) assert node[0][1].kind == "RETURN_VALUE" if node[0][0][0][0].kind == "tuple": assert len(node[0][0][0][0]) >= 2, "returning an empty tuple?" assert node[0][0][0][0][-1].kind.startswith("BUILD_TUPLE") for item in node[0][0][0][0][:-1]: analyze_expression(expr(item), context) else: analyze_expression(node[0][0][0][0], context) def analyze_expression(node, context): if node.kind == "LOAD_FAST": # Don't bother checking to see if this variable has been defined. # Unit checks test that if the coverage is complete. pass elif node.kind == "LOAD_CONST": assert isinstance(node.attr, (int, float)) elif is_pi(node): pass elif node.kind == "unary_op": assert len(node) == 2 analyze_expression(expr(node[0]), context) assert node[1].kind == "unary_operator" assert len(node[1]) == 1 analyze_unary_operator(node[1][0], context) elif node.kind == "bin_op": assert len(node) == 3 analyze_expression(expr(node[0]), context) analyze_expression(expr(node[1]), context) assert node[2].kind == "binary_operator" assert len(node[2]) == 1 analyze_binary_operator(node[2][0], context) elif node.kind == "compare": assert len(node) == 1 assert node[0].kind == "compare_single", "only do single comparisons" assert len(node[0]) == 3 analyze_expression(expr(node[0][0]), context) analyze_expression(expr(node[0][1]), context) assert node[0][2].kind == "COMPARE_OP" assert node[0][2].attr in allowed_comparisons, ( f"add {node[0][2].attr!r} to allowed_comparisons" ) elif node.kind == "call": assert len(node) >= 2 assert node[-1].kind.startswith("CALL_METHOD") or node[-1].kind.startswith( "CALL_FUNCTION" ) analyze_callable(expr(node[0]), context) for argument in node[1:-1]: expr_arg = argument[0] if argument.kind == "pos_arg" else argument assert expr_arg.kind == "expr", "only positional arguments" analyze_expression(expr(expr_arg), context) elif is_nan_to_num(node): analyze_expression(expr(node[1]), context) else: print(node) raise AssertionError("what is this?") def analyze_unary_operator(node, context): assert node.kind in allowed_unary_operators, ( f"add {node.kind!r} to allowed_unary_operators" ) def analyze_binary_operator(node, context): assert node.kind in allowed_binary_operators, ( f"add {node.kind!r} to allowed_binary_operators" ) def analyze_callable(node, context): if node.kind == "attribute37": assert len(node) == 2 module = expr(node[0]) assert module.kind in {"LOAD_FAST", "LOAD_GLOBAL"} assert node[1].kind == "LOAD_METHOD" if module.attr == "lib": assert node[1].attr in allowed_lib_functions, ( f"add {node[1].attr!r} to allowed_lib_functions" ) else: module_name = ".".join( context.closure.get(module.attr).__name__.split(".")[:-1] ) assert module_name in ( "vector._compute.planar", "vector._compute.spatial", "vector._compute.lorentz", ) elif node.kind in {"LOAD_GLOBAL", "LOAD_DEREF"}: function = context.closure.get(node.attr) assert function is not None, f"unrecognized function in scope: {node.attr!r}" analyze_function(function) else: print(node) raise AssertionError("what is this?") allowed_unary_operators = [ "UNARY_NEGATIVE", ] allowed_binary_operators = [ "BINARY_ADD", "BINARY_SUBTRACT", "BINARY_MULTIPLY", "BINARY_TRUE_DIVIDE", "BINARY_MODULO", "BINARY_POWER", "BINARY_AND", ] allowed_comparisons = [ "==", "!=", "<", ">", ] allowed_lib_functions = [ "absolute", "sign", "copysign", "maximum", "minimum", "sqrt", "exp", "log", "sin", "cos", "tan", "arcsin", "arccos", "arctan", "arctan2", "sinh", "cosh", "tanh", "arcsinh", "arccosh", "arctanh", "isclose", # TODO: https://github.com/scikit-hep/vector/issues/615 # remove once https://github.com/cupy/cupy/issues/9143 # is fixed. "where", ] vector-1.6.3/tests/test_issues.py000066400000000000000000000133041503546127100171160ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import os import pickle import numpy as np import pytest import vector def test_issue_99(): ak = pytest.importorskip("awkward") vector.register_awkward() vec = ak.Array([{"x": 1.0, "y": 2.0, "z": 3.0}], with_name="Vector3D") assert vec.to_xyz().tolist() == [{"x": 1.0, "y": 2.0, "z": 3.0}] assert vec[0].to_xyz().tolist() == {"x": 1.0, "y": 2.0, "z": 3.0} assert vec[0].to_rhophiz().tolist() == { "rho": 2.23606797749979, "phi": 1.1071487177940904, "z": 3.0, } def test_issue_161(): ak = pytest.importorskip("awkward") nb = pytest.importorskip("numba") vector.register_awkward() @nb.njit def repro(generator_like_jet_constituents): for sublist in generator_like_jet_constituents: s = 0 for generator_like_constituent in sublist: s += generator_like_constituent.pt file_path = os.path.join("tests", "samples", "issue-161-v2.pkl") with open(file_path, "rb") as f: a = ak.from_buffers(*pickle.load(f)) repro(generator_like_jet_constituents=a.constituents) def test_issue_443(): ak = pytest.importorskip("awkward") vector.register_awkward() assert vector.array({"E": [1], "px": [1], "py": [1], "pz": [1]}) ** 2 == np.array( [-2.0] ) assert ak.zip( {"E": [1], "px": [1], "py": [1], "pz": [1]}, with_name="Momentum4D" ) ** 2 == ak.Array([-2]) assert vector.obj(E=1, px=1, py=1, pz=1) ** 2 == -2 def test_issue_194(): vec2d = vector.VectorNumpy2D( { "x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5], } ) az1 = vector.backends.numpy.AzimuthalNumpyXY( [(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[("x", float), ("y", float)], ) az2 = vector.backends.numpy.AzimuthalNumpyXY( [(1.1, 3.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[("x", float), ("y", float)], ) azp1 = vector.backends.numpy.AzimuthalNumpyRhoPhi( [(1.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[("rho", float), ("phi", float)], ) azp2 = vector.backends.numpy.AzimuthalNumpyRhoPhi( [(2.1, 2.1), (1.2, 2.2), (1.3, 2.3), (1.4, 2.4), (1.5, 2.5)], dtype=[("rho", float), ("phi", float)], ) assert vec2d.azimuthal == az1 assert vec2d.azimuthal != az2 assert vec2d.azimuthal != azp1 assert az1 != az2 assert not az1 == azp1 # noqa: SIM201 assert not azp1 == az1 # noqa: SIM201 assert azp1 != az1 assert azp1 == azp1 # noqa: PLR0124 assert azp1 != azp2 vec3d = vector.VectorNumpy3D( { "x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5], "z": [3.1, 3.2, 3.3, 3.4, 3.5], } ) lg1 = vector.backends.numpy.LongitudinalNumpyZ( [(3.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[("z", float)] ) lg2 = vector.backends.numpy.LongitudinalNumpyZ( [(4.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[("z", float)] ) lgeta1 = vector.backends.numpy.LongitudinalNumpyEta( [(3.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[("eta", float)] ) lgeta2 = vector.backends.numpy.LongitudinalNumpyEta( [(4.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[("eta", float)] ) lgtheta1 = vector.backends.numpy.LongitudinalNumpyTheta( [(3.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[("theta", float)] ) lgtheta2 = vector.backends.numpy.LongitudinalNumpyTheta( [(4.1,), (3.2,), (3.3,), (3.4,), (3.5,)], dtype=[("theta", float)] ) assert vec3d.azimuthal == az1 assert vec3d.longitudinal == lg1 assert vec3d.longitudinal != lg2 assert vec3d.longitudinal != lgeta1 assert lg1 != lg2 assert not lg1 == lgeta1 # noqa: SIM201 assert not lgeta1 == lg1 # noqa: SIM201 assert lgeta1 != lg1 assert lgeta1 == lgeta1 # noqa: PLR0124 assert lgeta1 != lgeta2 assert lgtheta1 == lgtheta1 # noqa: PLR0124 assert lgtheta1 != lgtheta2 assert lgtheta1 != lgeta1 assert not lgtheta1 == lgeta1 # noqa: SIM201 vec4d = vector.VectorNumpy4D( { "x": [1.1, 1.2, 1.3, 1.4, 1.5], "y": [2.1, 2.2, 2.3, 2.4, 2.5], "z": [3.1, 3.2, 3.3, 3.4, 3.5], "t": [4.1, 4.2, 4.3, 4.4, 4.5], } ) tm1 = vector.backends.numpy.TemporalNumpyT( [(4.1,), (4.2,), (4.3,), (4.4,), (4.5,)], dtype=[("t", float)] ) tm2 = vector.backends.numpy.TemporalNumpyT( [(5.1,), (4.2,), (4.3,), (4.4,), (4.5,)], dtype=[("t", float)] ) tmtau1 = vector.backends.numpy.TemporalNumpyTau( [(4.1,), (4.2,), (4.3,), (4.4,), (4.5,)], dtype=[("tau", float)] ) tmtau2 = vector.backends.numpy.TemporalNumpyTau( [(5.1,), (4.2,), (4.3,), (4.4,), (4.5,)], dtype=[("tau", float)] ) assert vec4d.azimuthal == az1 assert vec4d.longitudinal == lg1 assert vec4d.temporal == tm1 assert vec4d.temporal != tm2 assert vec4d.temporal != tmtau1 assert tm1 != tm2 assert not tm1 == tmtau1 # noqa: SIM201 assert not tmtau1 == tm1 # noqa: SIM201 assert tmtau1 != tm1 assert tmtau1 == tmtau1 # noqa: PLR0124 assert tmtau1 != tmtau2 def test_issue_463(): v = vector.obj(x=1, y=1, z=1) for transform in "xyz", "xytheta", "xyeta", "rhophiz", "rhophitheta", "rhophieta": trv = getattr(v, "to_" + transform)() assert trv.deltaangle(trv) == 0.0 vector-1.6.3/tests/test_methods.py000066400000000000000000000056101503546127100172470ustar00rootroot00000000000000# Copyright (c) 2019-2025, Saransh Chopra, Henry Schreiner, Eduardo Rodrigues, Jonas Eschle, and Jim Pivarski. # # Distributed under the 3-clause BSD license, see accompanying file LICENSE # or https://github.com/scikit-hep/vector for details. from __future__ import annotations import vector from vector import ( MomentumNumpy2D, MomentumNumpy3D, MomentumNumpy4D, MomentumObject2D, MomentumObject3D, MomentumObject4D, VectorObject4D, ) def test_handler_of(): object_a = VectorObject4D.from_xyzt(0.0, 0.0, 0.0, 0.0) object_b = VectorObject4D.from_xyzt(1.0, 1.0, 1.0, 1.0) protocol = vector._methods._handler_of(object_a, object_b) assert protocol == object_a def test_momentum_coordinate_transforms(): numpy_vec = vector.array( { "px": [1.0, 2.0, 3.0], "py": [-1.0, 2.0, 3.0], }, ) object_vec = MomentumObject2D(px=0.0, py=0.0) for t1 in "pxpy", "ptphi": for t2 in "pz", "eta", "theta": for t3 in "mass", "energy": transformed_object = getattr(object_vec, "to_" + t1)() assert isinstance(transformed_object, MomentumObject2D) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) transformed_object = getattr(object_vec, "to_" + t1 + t2)() assert isinstance(transformed_object, MomentumObject3D) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) assert hasattr(transformed_object, t2) transformed_object = getattr(object_vec, "to_" + t1 + t2 + t3)() assert isinstance(transformed_object, MomentumObject4D) assert hasattr(transformed_object, t1[:2]) assert hasattr(transformed_object, t1[2:]) assert hasattr(transformed_object, t2) assert hasattr(transformed_object, t3) transformed_numpy = getattr(numpy_vec, "to_" + t1)() assert isinstance(transformed_numpy, MomentumNumpy2D) assert hasattr(transformed_numpy, t1[:2]) assert hasattr(transformed_numpy, t1[2:]) transformed_numpy = getattr(numpy_vec, "to_" + t1 + t2)() assert isinstance(transformed_numpy, MomentumNumpy3D) assert hasattr(transformed_numpy, t1[:2]) assert hasattr(transformed_numpy, t1[2:]) assert hasattr(transformed_numpy, t2) transformed_numpy = getattr(numpy_vec, "to_" + t1 + t2 + t3)() assert isinstance(transformed_numpy, MomentumNumpy4D) assert hasattr(transformed_numpy, t1[:2]) assert hasattr(transformed_numpy, t1[2:]) assert hasattr(transformed_numpy, t2) assert hasattr(transformed_numpy, t3) vector-1.6.3/tests/test_notebooks.py000066400000000000000000000021331503546127100176040ustar00rootroot00000000000000from __future__ import annotations import sys from pathlib import Path import papermill as pm import pytest @pytest.fixture def common_kwargs(tmpdir): outputnb = tmpdir.join("output.ipynb") return { "output_path": str(outputnb), "kernel_name": f"python{sys.version_info.major}", "progress_bar": False, } def test_object(common_kwargs): execution_dir = Path.cwd() / "docs" / "src" pm.execute_notebook(execution_dir / "object.ipynb", **common_kwargs) def test_numpy(common_kwargs): execution_dir = Path.cwd() / "docs" / "src" pm.execute_notebook(execution_dir / "numpy.ipynb", **common_kwargs) def test_awkward(common_kwargs): execution_dir = Path.cwd() / "docs" / "src" pm.execute_notebook(execution_dir / "awkward.ipynb", **common_kwargs) def test_numba(common_kwargs): execution_dir = Path.cwd() / "docs" / "src" pm.execute_notebook(execution_dir / "numba.ipynb", **common_kwargs) def test_sympy(common_kwargs): execution_dir = Path.cwd() / "docs" / "src" pm.execute_notebook(execution_dir / "sympy.ipynb", **common_kwargs)