pax_global_header00006660000000000000000000000064151603244610014514gustar00rootroot0000000000000052 comment=22b847b6f4d2b8bb674578a8b652a0db73ddc33a pint-xarray-0.6.1/000077500000000000000000000000001516032446100137765ustar00rootroot00000000000000pint-xarray-0.6.1/.codecov.yml000066400000000000000000000004511516032446100162210ustar00rootroot00000000000000codecov: ci: # by default, codecov doesn't recognize azure as a CI provider - dev.azure.com require_ci_to_pass: yes coverage: status: project: default: # Require 1% coverage, i.e., always succeed target: 1 patch: false changes: false comment: off pint-xarray-0.6.1/.gitattributes000066400000000000000000000001721516032446100166710ustar00rootroot00000000000000# SCM syntax highlighting & preventing 3-way merges pixi.lock merge=binary linguist-language=YAML linguist-generated=true pint-xarray-0.6.1/.github/000077500000000000000000000000001516032446100153365ustar00rootroot00000000000000pint-xarray-0.6.1/.github/PULL_REQUEST_TEMPLATE.md000066400000000000000000000004541516032446100211420ustar00rootroot00000000000000 - [ ] Closes #xxxx - [ ] Tests added - [ ] Passes `pre-commit run --all-files` - [ ] User visible changes (including notable bug fixes) are documented in `whats-new.rst` - [ ] New functions/methods are listed in `api.rst` pint-xarray-0.6.1/.github/dependabot.yml000066400000000000000000000001661516032446100201710ustar00rootroot00000000000000version: 2 updates: - package-ecosystem: "github-actions" directory: "/" schedule: interval: "weekly" pint-xarray-0.6.1/.github/release.yml000066400000000000000000000001261516032446100175000ustar00rootroot00000000000000changelog: exclude: authors: - dependabot[bot] - pre-commit-ci[bot] pint-xarray-0.6.1/.github/workflows/000077500000000000000000000000001516032446100173735ustar00rootroot00000000000000pint-xarray-0.6.1/.github/workflows/ci-additional.yml000066400000000000000000000016321516032446100226210ustar00rootroot00000000000000name: CI Additional on: push: branches: - main pull_request: branches: - main concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true jobs: doctests: name: Doctests runs-on: ubuntu-latest if: github.repository == 'xarray-contrib/pint-xarray' env: FORCE_COLOR: 3 steps: - name: checkout the repository uses: actions/checkout@v6 with: # need to fetch all tags to get a correct version fetch-depth: 0 # fetch all branches and tags - name: setup environment uses: prefix-dev/setup-pixi@a0af7a228712d6121d37aba47adf55c1332c9c2e # 0.9.4 with: environments: "ci-py313" - name: import pint-xarray run: | pixi run -e ci-py313 python -c 'import pint_xarray' - name: run doctests run: | pixi run -e ci-py313 doctests pint-xarray-0.6.1/.github/workflows/ci.yml000066400000000000000000000066471516032446100205260ustar00rootroot00000000000000# adapted from xarray's ci name: CI on: push: branches: [main] pull_request: branches: [main] workflow_dispatch: concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true jobs: detect-skip-ci-trigger: name: "Detect CI Trigger: [skip-ci]" if: github.event_name == 'push' || github.event_name == 'pull_request' runs-on: ubuntu-slim outputs: triggered: ${{ steps.detect-trigger.outputs.trigger-found }} steps: - uses: actions/checkout@v6 with: fetch-depth: 2 - uses: xarray-contrib/ci-trigger@v1 id: detect-trigger with: keyword: "[skip-ci]" cache-pixi-lock: name: "Cache pixi lock" needs: detect-skip-ci-trigger runs-on: ubuntu-slim if: | always() && github.repository == 'xarray-contrib/pint-xarray' && ( github.event_name == 'workflow_dispatch' || github.event_name == 'push' || ( github.event_name == 'pull_request' && ( needs.detect-skip-ci-trigger.outputs.triggered == 'false' && !contains(github.event.pull_request.labels.*.name, 'skip-ci') ) ) ) outputs: cache-key: ${{ steps.pixi-lock.outputs.cache-key }} pixi-version: ${{ steps.pixi-lock.outputs.pixi-version }} steps: - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 with: persist-credentials: false - uses: Parcels-code/pixi-lock/create-and-cache@5cbd7a155fa24aa0711ec0a9aad9ede0819e84a1 # v0.1.0 id: pixi-lock ci: name: ${{ matrix.os }} ${{ matrix.env }} runs-on: ${{ matrix.os }} needs: cache-pixi-lock defaults: run: shell: bash -l {0} env: FORCE_COLOR: 3 strategy: fail-fast: false matrix: env: ["ci-py311", "ci-py313", "ci-py314"] os: ["ubuntu-latest", "macos-latest", "windows-latest"] steps: - name: checkout the repository uses: actions/checkout@v6 with: # need to fetch all tags to get a correct version fetch-depth: 0 # fetch all branches and tags - uses: Parcels-code/pixi-lock/restore@5cbd7a155fa24aa0711ec0a9aad9ede0819e84a1 # v0.1.0 with: cache-key: ${{ needs.cache-pixi-lock.outputs.cache-key }} - name: setup environment uses: prefix-dev/setup-pixi@a0af7a228712d6121d37aba47adf55c1332c9c2e # 0.9.4 with: pixi-version: "${{ needs.cache-pixi-lock.outputs.pixi-version }}" frozen: true cache: true environments: "${{ matrix.env }}" - name: investigate env variables run: | echo PYTHON_VERSION=$(pixi run -e ${{ matrix.env }} python -c 'import sys; print(".".join(map(str, sys.version_info[:2])))') >> $GITHUB_ENV echo RUNNER_OS="${{ matrix.os }}" - name: import pint-xarray run: | pixi run -e ${{ matrix.env }} python -c 'import pint_xarray' - name: run tests if: success() id: status run: | pixi run -e ${{ matrix.env }} tests --cov-report=xml - name: Upload code coverage to Codecov uses: codecov/codecov-action@v5.5.3 with: token: "${{ secrets.CODECOV_TOKEN }}" files: ./coverage.xml flags: unittests env_vars: RUNNER_OS,PYTHON_VERSION name: codecov-umbrella fail_ci_if_error: false pint-xarray-0.6.1/.github/workflows/nightly.yml000066400000000000000000000046121516032446100215770ustar00rootroot00000000000000# adapted from xarray's nightly CI name: Nightly CI on: push: branches: [main] pull_request: branches: [main] schedule: - cron: "0 0 * * *" # Daily "At 00:00" UTC workflow_dispatch: concurrency: group: ${{ github.workflow }}-${{ github.ref }} cancel-in-progress: true jobs: detect-test-upstream-trigger: name: "Detect CI Trigger: [test-upstream]" if: | github.repository_owner == 'xarray-contrib' && (github.event_name == 'push' || github.event_name == 'pull_request') runs-on: ubuntu-slim outputs: triggered: ${{ steps.detect-trigger.outputs.trigger-found }} steps: - uses: actions/checkout@v6 with: fetch-depth: 2 - uses: xarray-contrib/ci-trigger@v1.2 id: detect-trigger with: keyword: "[test-upstream]" upstream-dev: name: upstream-dev runs-on: ubuntu-latest needs: detect-test-upstream-trigger if: | always() && github.repository_owner == 'xarray-contrib' && ( (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') || needs.detect-test-upstream-trigger.outputs.triggered == 'true' || ( github.event_name == 'pull_request' && contains(github.event.pull_request.labels.*.name, 'run-upstream') ) ) outputs: artifacts_availability: ${{ steps.status.outputs.ARTIFACTS_AVAILABLE }} env: FORCE_COLOR: 3 steps: - name: checkout the repository uses: actions/checkout@v6 with: # need to fetch all tags to get a correct version fetch-depth: 0 # fetch all branches and tags - name: setup environment uses: prefix-dev/setup-pixi@a0af7a228712d6121d37aba47adf55c1332c9c2e # 0.9.4 with: environments: "nightly" locked: false frozen: false cache: false - name: import pint-xarray run: | pixi run -e nightly python -c 'import pint_xarray' - name: run tests if: success() id: status run: | pixi run -e nightly tests -rf --report-log=pytest-log.jsonl - name: report failures if: | failure() && steps.status.outcome == 'failure' && github.event_name == 'schedule' uses: scientific-python/issue-from-pytest-log@v1 with: log-path: pytest-log.jsonl pint-xarray-0.6.1/.github/workflows/pypi.yaml000066400000000000000000000030471516032446100212440ustar00rootroot00000000000000name: Upload Package to PyPI on: release: types: - published jobs: build-artifacts: runs-on: ubuntu-latest if: github.repository == 'xarray-contrib/pint-xarray' steps: - uses: actions/checkout@v6 with: fetch-depth: 0 - uses: actions/setup-python@v6 name: Install Python with: python-version: "3.x" - name: Install dependencies run: | python -m pip install --upgrade pip python -m pip install build twine - name: Build tarball and wheels run: | git clean -xdf git restore -SW . python -m build --outdir dist/ . - name: Check built artifacts run: | python -m twine check --strict dist/* pwd if [ -f dist/pint-xarray-0.0.0.tar.gz ]; then echo "❌ INVALID VERSION NUMBER" exit 1 else echo "✅ Looks good" fi - uses: actions/upload-artifact@v7 with: name: releases path: dist upload-to-pypi: needs: build-artifacts if: github.event_name == 'release' runs-on: ubuntu-latest environment: name: pypi url: https://pypi.org/p/pint-xarray permissions: id-token: write steps: - uses: actions/download-artifact@v8 with: name: releases path: dist - name: Publish package to PyPI uses: pypa/gh-action-pypi-publish@ed0c53931b1dc9bd32cbe73a98c7f6766f8a527e with: verbose: true pint-xarray-0.6.1/.gitignore000066400000000000000000000035061516032446100157720ustar00rootroot00000000000000# 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/ docs/generated/ # 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/ # pixi environments .pixi/* !.pixi/config.toml pint-xarray-0.6.1/.pre-commit-config.yaml000066400000000000000000000024451516032446100202640ustar00rootroot00000000000000ci: autoupdate_schedule: weekly repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v6.0.0 hooks: - id: trailing-whitespace - id: end-of-file-fixer - id: check-docstring-first - repo: https://github.com/rbubley/mirrors-prettier rev: v3.8.1 hooks: - id: prettier args: ["--cache-location=.prettier_cache/cache"] - repo: https://github.com/ComPWA/taplo-pre-commit rev: v0.9.3 hooks: - id: taplo-format args: ["--option", "array_auto_collapse=false"] - id: taplo-lint args: ["--no-schema"] - repo: https://github.com/abravalheri/validate-pyproject rev: v0.25 hooks: - id: validate-pyproject - repo: https://github.com/psf/black-pre-commit-mirror rev: 26.3.1 hooks: - id: black-jupyter - repo: https://github.com/keewis/blackdoc rev: v0.4.6 hooks: - id: blackdoc additional_dependencies: ["black==26.3.1"] - id: blackdoc-autoupdate-black - repo: https://github.com/astral-sh/ruff-pre-commit rev: v0.15.7 hooks: - id: ruff args: ["--fix"] - repo: https://github.com/kynan/nbstripout rev: 0.9.1 hooks: - id: nbstripout args: - "--extra-keys='metadata.kernelspec metadata.language_info.version'" pint-xarray-0.6.1/.readthedocs.yaml000066400000000000000000000011331516032446100172230ustar00rootroot00000000000000version: 2 build: os: ubuntu-lts-latest tools: # just so RTD stops complaining python: "latest" jobs: post_checkout: - (git --no-pager log --pretty="tformat:%s" -1 | grep -vqF "[skip-rtd]") || exit 183 - git fetch --unshallow || true create_environment: - asdf plugin add pixi - asdf install pixi latest - asdf global pixi latest pre_install: - git update-index --assume-unchanged docs/conf.py install: - pixi install -e docs build: html: - pixi run -e docs build-docs-rtd sphinx: configuration: docs/conf.py pint-xarray-0.6.1/HOW_TO_RELEASE.rst000066400000000000000000000016301516032446100167270ustar00rootroot00000000000000Release process =============== 1. the release happens from `main` so make sure it is up-to-date: .. code:: sh git pull origin main 2. look at `whats-new.rst` and make sure it is complete and with references to issues and pull requests 3. open and merge a pull request with these changes 4. make sure the CI on main pass 5. check that the documentation build on readthedocs completed successfully 6. Fill in the release date and commit the release: .. code:: sh git commit -am "Release v0.X.Y" 7. Tag the release and push to main: .. code:: sh git tag -a v0.X.Y -m "v0.X.Y" git push origin --tags 8. Draft a release for the new tag on github. A CI will pick that up, build the project and push to PyPI. Be careful, this can't be undone. 9. Make sure readthedocs builds both `stable` and the new tag 10. Add a new section to `whats-new.rst` and push directly to main pint-xarray-0.6.1/LICENSE000066400000000000000000000261351516032446100150120ustar00rootroot00000000000000 Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: (a) You must give any other recipients of the Work or Derivative Works a copy of this License; and (b) You must cause any modified files to carry prominent notices stating that You changed the files; and (c) You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and (d) If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work. To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. pint-xarray-0.6.1/README.md000066400000000000000000000041751516032446100152640ustar00rootroot00000000000000[![CI](https://github.com/xarray-contrib/pint-xarray/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/xarray-contrib/pint-xarray/actions/workflows/ci.yml) [![code coverage](https://codecov.io/gh/xarray-contrib/pint-xarray/branch/main/graph/badge.svg)](https://codecov.io/gh/xarray-contrib/pint-xarray) [![docs](https://readthedocs.org/projects/pint-xarray/badge/?version=latest)](https://pint-xarray.readthedocs.io) [![PyPI version](https://img.shields.io/pypi/v/pint-xarray.svg)](https://pypi.org/project/pint-xarray) [![codestyle](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/python/black) [![conda-forge](https://img.shields.io/conda/vn/conda-forge/pint-xarray)](https://github.com/conda-forge/pint-xarray-feedstock) # pint-xarray A convenience wrapper for using [pint](https://pint.readthedocs.io) with [xarray](https://xarray.pydata.org). ## Usage To convert the variables of a `Dataset` to quantities: ```python In [1]: import pint_xarray ...: import xarray as xr In [2]: ds = xr.Dataset({"a": ("x", [0, 1, 2]), "b": ("y", [-3, 5, 1], {"units": "m"})}) ...: ds Out[2]: Dimensions: (x: 3, y: 3) Dimensions without coordinates: x, y Data variables: a (x) int64 0 1 2 b (y) int64 -3 5 1 In [3]: q = ds.pint.quantify(a="s") ...: q Out[3]: Dimensions: (x: 3, y: 3) Dimensions without coordinates: x, y Data variables: a (x) int64 [s] 0 1 2 b (y) int64 [m] -3 5 1 ``` to convert to different units: ```python In [4]: c = q.pint.to({"a": "ms", "b": "km"}) ...: c Out[4]: Dimensions: (x: 3, y: 3) Dimensions without coordinates: x, y Data variables: a (x) float64 [ms] 0.0 1e+03 2e+03 b (y) float64 [km] -0.003 0.005 0.001 ``` to convert back to non-quantities: ```python In [5]: d = c.pint.dequantify() ...: d Out[5]: Dimensions: (x: 3, y: 3) Dimensions without coordinates: x, y Data variables: a (x) float64 0.0 1e+03 2e+03 b (y) float64 -0.003 0.005 0.001 ``` For more, see the [documentation](https://pint-xarray.readthedocs.io) pint-xarray-0.6.1/conftest.py000066400000000000000000000011031516032446100161700ustar00rootroot00000000000000import pytest @pytest.fixture(autouse=True) def add_standard_imports(doctest_namespace, tmpdir): import numpy as np import pandas as pd import pint import xarray as xr import pint_xarray ureg = pint.UnitRegistry(force_ndarray_like=True) doctest_namespace["np"] = np doctest_namespace["pd"] = pd doctest_namespace["xr"] = xr doctest_namespace["pint"] = pint doctest_namespace["ureg"] = ureg doctest_namespace["pint_xarray"] = pint_xarray # always seed numpy.random to make the examples deterministic np.random.seed(0) pint-xarray-0.6.1/docs/000077500000000000000000000000001516032446100147265ustar00rootroot00000000000000pint-xarray-0.6.1/docs/api.rst000066400000000000000000000036771516032446100162460ustar00rootroot00000000000000API reference ============= This page contains a auto-generated summary of ``pint-xarray``'s API. .. autosummary:: :toctree: generated/ pint_xarray.unit_registry pint_xarray.setup_registry Dataset ------- .. autosummary:: :toctree: generated/ :template: autosummary/accessor_attribute.rst xarray.Dataset.pint.loc .. autosummary:: :toctree: generated/ :template: autosummary/accessor_method.rst xarray.Dataset.pint.quantify xarray.Dataset.pint.dequantify xarray.Dataset.pint.interp xarray.Dataset.pint.interp_like xarray.Dataset.pint.reindex xarray.Dataset.pint.reindex_like xarray.Dataset.pint.drop_sel xarray.Dataset.pint.sel xarray.Dataset.pint.to xarray.Dataset.pint.chunk xarray.Dataset.pint.ffill xarray.Dataset.pint.bfill xarray.Dataset.pint.interpolate_na DataArray --------- .. autosummary:: :toctree: generated/ :template: autosummary/accessor_attribute.rst xarray.DataArray.pint.loc xarray.DataArray.pint.magnitude xarray.DataArray.pint.units xarray.DataArray.pint.dimensionality xarray.DataArray.pint.registry .. autosummary:: :toctree: generated/ :template: autosummary/accessor_method.rst xarray.DataArray.pint.quantify xarray.DataArray.pint.dequantify xarray.DataArray.pint.interp xarray.DataArray.pint.interp_like xarray.DataArray.pint.reindex xarray.DataArray.pint.reindex_like xarray.DataArray.pint.drop_sel xarray.DataArray.pint.sel xarray.DataArray.pint.to xarray.DataArray.pint.chunk xarray.DataArray.pint.ffill xarray.DataArray.pint.bfill xarray.DataArray.pint.interpolate_na Wrapping quantity-unaware functions ----------------------------------- .. autosummary:: :toctree: generated/ pint_xarray.expects Exceptions ---------- .. autosummary:: :toctree: generated/ pint_xarray.errors.PintExceptionGroup Testing ------- .. autosummary:: :toctree: generated/ pint_xarray.testing.assert_units_equal pint-xarray-0.6.1/docs/conf.py000066400000000000000000000070341516032446100162310ustar00rootroot00000000000000# 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 # -- 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('.')) # -- Imports ----------------------------------------------------------------- import datetime as dt import sphinx_autosummary_accessors # need to import so accessors get registered import pint_xarray # noqa: F401 # -- Project information ----------------------------------------------------- year = dt.datetime.now().year project = "pint-xarray" author = f"{project} developers" copyright = f"{year}, {author}" github_url = "https://github.com/xarray-contrib/pint-xarray" # -- 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 = [ "sphinx.ext.intersphinx", "sphinx.ext.extlinks", "sphinx.ext.autosummary", "sphinx.ext.autodoc", "sphinx.ext.napoleon", "sphinx_autosummary_accessors", "IPython.sphinxext.ipython_directive", "IPython.sphinxext.ipython_console_highlighting", "nbsphinx", ] # Add any paths that contain templates here, relative to this directory. templates_path = ["_templates", sphinx_autosummary_accessors.templates_path] # 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"] # -- 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_rtd_theme" # 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"] # -- Extension configuration ------------------------------------------------- # extlinks extlinks = { "issue": (f"{github_url}/issues/%s", "GH%s"), "pull": (f"{github_url}/pull/%s", "PR%s"), } # autosummary autosummary_generate = True # autodoc autodoc_typehints = "none" # napoleon napoleon_use_param = False napoleon_use_rtype = True napoleon_preprocess_types = True napoleon_type_aliases = { "dict-like": ":term:`dict-like `", "mapping": ":term:`mapping`", "hashable": ":term:`hashable`", # xarray "Dataset": "~xarray.Dataset", "DataArray": "~xarray.DataArray", # pint / pint-xarray "unit-like": ":term:`unit-like`", } # nbsphinx nbsphinx_timeout = 600 nbsphinx_execute = "always" # -- Options for intersphinx extension --------------------------------------- intersphinx_mapping = { "python": ("https://docs.python.org/3/", None), "xarray": ("https://docs.xarray.dev/en/stable", None), "pint": ("https://pint.readthedocs.io/en/stable", None), "pandas": ("https://pandas.pydata.org/pandas-docs/stable", None), } pint-xarray-0.6.1/docs/contributing.rst000066400000000000000000000015461516032446100201750ustar00rootroot00000000000000Contributing ============ ``pint-xarray`` is developed on `github `_. Commit message tags ------------------- By default, the upstream dev CI is disabled on pull request and push events. You can override this behavior per commit by adding a [test-upstream] tag to the first line of the commit message. Linters / Autoformatters ------------------------ In order to keep code consistent, we use - `Black `_ for standardized code formatting - `blackdoc `_ for standardized code formatting in documentation - `Flake8 `_ for general code quality - `isort `_ for standardized order in imports. See also `flake8-isort `_. pint-xarray-0.6.1/docs/conversion.rst000066400000000000000000000013501516032446100176440ustar00rootroot00000000000000.. currentmodule:: xarray Converting units ================ .. ipython:: python :suppress: import xarray as xr When working with :py:class:`Dataset` or :py:class:`DataArray` objects with units, we frequently might want to convert the units. Suppose we have: .. ipython:: In [1]: ds = xr.Dataset( ...: {"a": ("x", [4, 8, 12, 16])}, coords={"u": ("x", [10, 20, 30, 40])} ...: ).pint.quantify({"a": "m", "u": "s"}) ...: ds In [2]: da = ds.a ...: da To convert the data to different units, we can use the :py:meth:`Dataset.pint.to` and :py:meth:`DataArray.pint.to` methods: .. ipython:: In [3]: ds.pint.to(a="feet", u="ks") In [4]: da.pint.to({da.name: "nautical_mile", "u": "ms"}) pint-xarray-0.6.1/docs/creation.rst000066400000000000000000000077331516032446100172760ustar00rootroot00000000000000.. currentmodule:: xarray Creating and saving objects with units ====================================== Attaching units --------------- .. ipython:: python :suppress: import pint import pint_xarray import xarray as xr Usually, when loading data from disk we get a :py:class:`Dataset` or :py:class:`DataArray` with units in attributes: .. ipython:: In [1]: ds = xr.Dataset( ...: { ...: "a": (("lon", "lat"), [[11.84, 3.12, 9.7], [7.8, 9.3, 14.72]]), ...: "b": (("lon", "lat"), [[13, 2, 7], [5, 4, 9]], {"units": "m"}), ...: }, ...: coords={"lat": [10, 20, 30], "lon": [74, 76]}, ...: ) ...: ds In [2]: da = ds.b ...: da In order to get :py:class:`pint.Quantity` instances, we can use the :py:meth:`Dataset.pint.quantify` or :py:meth:`DataArray.pint.quantify` methods: .. ipython:: In [3]: ds.pint.quantify() We can also override the units of a variable: .. ipython:: In [4]: ds.pint.quantify(b="km") In [5]: da.pint.quantify("degree") Overriding works even if there is no ``units`` attribute, so we could use this to attach units to a normal :py:class:`Dataset`: .. ipython:: In [6]: temporary_ds = xr.Dataset({"a": ("x", [0, 5, 10])}, coords={"x": [1, 2, 3]}) ...: temporary_ds.pint.quantify({"a": "m"}) Of course, we could use :py:class:`pint.Unit` instances instead of strings to specify units, too. .. note:: Unit objects tied to different registries cannot interact with each other. In order to avoid this, :py:meth:`DataArray.pint.quantify` and :py:meth:`Dataset.pint.quantify` will make sure only a single registry is used per ``xarray`` object. If we wanted to change the units of the data of a :py:class:`DataArray`, we could do so using the :py:attr:`DataArray.name` attribute: .. ipython:: In [7]: da.pint.quantify({da.name: "J", "lat": "degree", "lon": "degree"}) However, `xarray`_ currently doesn't support `units in indexes`_, so the new units were set as attributes. To really observe the changes the ``quantify`` methods make, we have to first swap the dimensions: .. ipython:: In [8]: ds_with_units = ds.swap_dims({"lon": "x", "lat": "y"}).pint.quantify( ...: {"lat": "degree", "lon": "degree"} ...: ) ...: ds_with_units In [9]: da_with_units = da.swap_dims({"lon": "x", "lat": "y"}).pint.quantify( ...: {"lat": "degree", "lon": "degree"} ...: ) ...: da_with_units By default, :py:meth:`Dataset.pint.quantify` and :py:meth:`DataArray.pint.quantify` will use the unit registry at :py:obj:`pint_xarray.unit_registry` (the :py:func:`application registry `). If we want a different registry, we can either pass it as the ``unit_registry`` parameter: .. ipython:: In [10]: ureg = pint.UnitRegistry(force_ndarray_like=True) ...: # set up the registry In [11]: da.pint.quantify("degree", unit_registry=ureg) or overwrite the default registry: .. ipython:: In [12]: pint_xarray.unit_registry = ureg In [13]: da.pint.quantify("degree") .. note:: To properly work with ``xarray``, the ``force_ndarray_like`` or ``force_ndarray`` options have to be enabled on the custom registry. Without it, python scalars wrapped by :py:class:`pint.Quantity` may raise errors or have their units stripped. Saving with units ----------------- In order to not lose the units when saving to disk, we first have to call the :py:meth:`Dataset.pint.dequantify` and :py:meth:`DataArray.pint.dequantify` methods: .. ipython:: In [10]: ds_with_units.pint.dequantify() In [11]: da_with_units.pint.dequantify() This will get the string representation of a :py:class:`pint.Unit` instance and attach it as a ``units`` attribute. The data of the variable will now be whatever `pint`_ wrapped. .. _pint: https://pint.readthedocs.io/en/stable/ .. _xarray: https://docs.xarray.dev/en/stable/ .. _units in indexes: https://github.com/pydata/xarray/issues/1603 pint-xarray-0.6.1/docs/examples.rst000066400000000000000000000001101516032446100172660ustar00rootroot00000000000000Examples ======== .. toctree:: :maxdepth: 1 examples/plotting pint-xarray-0.6.1/docs/examples/000077500000000000000000000000001516032446100165445ustar00rootroot00000000000000pint-xarray-0.6.1/docs/examples/plotting.ipynb000066400000000000000000000071201516032446100214470ustar00rootroot00000000000000{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# plotting quantified data" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "# to be able to read unit attributes following the CF conventions\n", "import cf_xarray.units # noqa: F401 # must be imported before pint_xarray\n", "import xarray as xr\n", "\n", "import pint_xarray # noqa: F401\n", "from pint_xarray import unit_registry as ureg\n", "\n", "xr.set_options(display_expand_data=False)" ] }, { "cell_type": "markdown", "id": "2", "metadata": {}, "source": [ "## load the data" ] }, { "cell_type": "code", "execution_count": null, "id": "3", "metadata": {}, "outputs": [], "source": [ "ds = xr.tutorial.open_dataset(\"air_temperature\")\n", "data = ds.air\n", "data" ] }, { "cell_type": "markdown", "id": "4", "metadata": {}, "source": [ "## quantify the data" ] }, { "cell_type": "markdown", "id": "5", "metadata": {}, "source": [ "
\n", "Note: this example uses the data provided by the xarray.tutorial functions. As such, the units attributes follow the CF conventions, which pint does not understand by default. To still be able to read them we are using the registry provided by cf-xarray.\n", "
" ] }, { "cell_type": "code", "execution_count": null, "id": "6", "metadata": {}, "outputs": [], "source": [ "quantified = data.pint.quantify()\n", "quantified" ] }, { "cell_type": "markdown", "id": "7", "metadata": {}, "source": [ "## work with the data" ] }, { "cell_type": "code", "execution_count": null, "id": "8", "metadata": {}, "outputs": [], "source": [ "monthly_means = quantified.pint.to(\"degC\").sel(time=\"2013\").groupby(\"time.month\").mean()\n", "monthly_means" ] }, { "cell_type": "markdown", "id": "9", "metadata": {}, "source": [ "Most operations will preserve the units but there are some which will drop them (see the [duck array integration status](https://xarray.pydata.org/en/stable/user-guide/duckarrays.html#missing-features) page). To work around that there are unit-aware versions on the `.pint` accessor. For example, to select data use `.pint.sel` instead of `.sel`:" ] }, { "cell_type": "code", "execution_count": null, "id": "10", "metadata": {}, "outputs": [], "source": [ "monthly_means.sel(\n", " lat=ureg.Quantity(4350, \"angular_minute\"),\n", " lon=ureg.Quantity(12000, \"angular_minute\"),\n", ")" ] }, { "cell_type": "markdown", "id": "11", "metadata": {}, "source": [ "## plot\n", "\n", "`xarray`'s plotting functions will cast the data to `numpy.ndarray`, so we need to \"dequantify\" first." ] }, { "cell_type": "code", "execution_count": null, "id": "12", "metadata": {}, "outputs": [], "source": [ "monthly_means.pint.dequantify(format=\"~P\").plot.imshow(col=\"month\", col_wrap=4)" ] } ], "metadata": { "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3" } }, "nbformat": 4, "nbformat_minor": 5 } pint-xarray-0.6.1/docs/index.rst000066400000000000000000000015401516032446100165670ustar00rootroot00000000000000pint-xarray =========== A convenience wrapper for using `pint`_ in `xarray`_ objects. .. _pint: https://pint.readthedocs.io/en/stable .. _xarray: https://xarray.pydata.org/en/stable .. warning:: This package is experimental, and new versions might introduce backwards incompatible changes. Documentation ------------- **Getting Started**: - :doc:`installation` - :doc:`examples` .. toctree:: :maxdepth: 1 :caption: Getting Started :hidden: installation examples **User Guide**: - :doc:`terminology` - :doc:`creation` - :doc:`conversion` .. toctree:: :maxdepth: 1 :caption: User Guide :hidden: terminology creation conversion **Help & Reference**: - :doc:`whats-new` - :doc:`api` - :doc:`contributing` .. toctree:: :maxdepth: 1 :caption: Help & Reference :hidden: whats-new api contributing pint-xarray-0.6.1/docs/installation.rst000066400000000000000000000007121516032446100201610ustar00rootroot00000000000000Installation ------------ Install from ``conda-forge``: .. code:: sh conda install -c conda-forge pint-xarray or from ``PyPI``: .. code:: sh python -m pip install pint-xarray or from source, either directly from github: .. code:: sh python -m pip install git+https://github.com/xarray-contrib/pint-xarray or from a local copy: .. code:: sh git clone https://github.com/xarray-contrib/pint-xarray python -m pip install ./pint-xarray pint-xarray-0.6.1/docs/terminology.rst000066400000000000000000000004161516032446100200310ustar00rootroot00000000000000Terminology =========== .. glossary:: unit-like A `pint`_ unit definition, as accepted by :py:class:`pint.Unit`. May be a :py:class:`str`, a :py:class:`pint.Unit` instance or :py:obj:`None`. .. _pint: https://pint.readthedocs.io/en/stable pint-xarray-0.6.1/docs/whats-new.rst000066400000000000000000000204461516032446100174030ustar00rootroot00000000000000.. currentmodule:: xarray What's new ========== 0.6.1 (23 Mar 2026) ------------------- - Properly rename the coordinate units (:pull:`351`). By `Justus Magin `_. - Don't rely on ``xarray``'s default behavior regarding ``attrs`` in :py:func:`pint_xarray.expects` (:issue:`360`, :pull:`367`) By `Justus Magin `_. - Fix quantifying a existing index and adjust the index conversion tests (:pull:`368`) By `Justus Magin `_. 0.6.0 (31 Aug 2025) ------------------- - Bump dependency versions (:pull:`313`): ============ ============== ============== dependency old minimum new minimum ============ ============== ============== python 3.10 3.11 xarray 2022.06.0 2023.07.0 numpy 1.23 1.26 pint 0.21 0.24 ============ ============== ============== By `Justus Magin `_. - Switch to using pixi for all dependency management (:pull:`314`). By `Justus Magin `_. - Added the :py:func:`pint_xarray.expects` decorator to allow wrapping quantity-unaware functions (:issue:`141`, :pull:`316`). By `Justus Magin `_ and `Tom Nicholas `_. - Follow the change in signature of :py:meth:`xarray.Index.equals` (:issue:`322`, :pull:`324`) By `Justus Magin `_. - Add units to the inline ``repr`` and define a custom ``repr`` (:issue:`308`, :pull:`325`) By `Justus Magin `_. - Collect multiple errors into a specific exception group (:pull:`329`) By `Justus Magin `_. 0.5.1 (10 Aug 2025) ------------------- - Pass ``sel`` options to the wrapped array (:pull:`304`, :issue:`303`) By `Bhavin Patel `_. - Support python 3.13 (:pull:`310`) By `Justus Magin `_. 0.5 (09 Jun 2025) ------------------ - drop support for python 3.9 (:pull:`266`) By `Justus Magin `_. - create a `PintIndex` to allow units on indexed coordinates (:pull:`163`, :issue:`162`) By `Justus Magin `_ and `Benoit Bovy `_. - fix :py:meth:`Dataset.pint.interp` and :py:meth:`DataArray.pint.interp` bug failing to pass through arguments (:pull:`270`, :issue:`267`) By `Martijn van der Marel `_ 0.4 (23 Jun 2024) ----------------- - adopt `SPEC0 `_ (:pull:`228`) This means that the supported versions change: ============ ============== ============== dependency old minimum new minimum ============ ============== ============== python 3.8 3.9 xarray 0.16.1 2022.06.0 numpy 1.17 1.23 pint 0.16 0.21 ============ ============== ============== By `Justus Magin `_. - add support for python 3.11 and 3.12 (:pull:`228`, :pull:`263`) By `Justus Magin `_. - ignore datetime units on attributes (:pull:`241`) By `Justus Magin `_. 0.3 (27 Jul 2022) ----------------- - drop support for python 3.7 (:pull:`153`) By `Justus Magin `_. - add support for python 3.10 (:pull:`155`) By `Justus Magin `_. - preserve :py:class:`pandas.MultiIndex` objects (:issue:`164`, :pull:`168`). By `Justus Magin `_. - fix "quantifying" dimension coordinates (:issue:`105`, :pull:`174`). By `Justus Magin `_. - allow using :py:meth:`DataArray.pint.quantify` and :py:meth:`Dataset.pint.quantify` as identity operators (:issue:`47`, :pull:`175`). By `Justus Magin `_. 0.2.1 (26 Jul 2021) ------------------- - allow special "no unit" values in :py:meth:`Dataset.pint.quantify` and :py:meth:`DataArray.pint.quantify` (:pull:`125`) By `Justus Magin `_. - convert the note about dimension coordinates saving their units in the attributes a warning (:issue:`124`, :pull:`126`) By `Justus Magin `_. - improve the documentation on the ``format`` parameter of :py:meth:`Dataset.pint.dequantify` and :py:meth:`DataArray.pint.dequantify` (:issue:`121`, :pull:`127`, :pull:`132`) By `Justus Magin `_. - use `cf-xarray `_'s unit registry in the plotting example (:issue:`107`, :pull:`128`). By `Justus Magin `_. 0.2 (May 10 2021) ----------------- - rewrite :py:meth:`Dataset.pint.quantify` and :py:meth:`DataArray.pint.quantify`, to use pint's ``UnitRegistry.parse_units`` instead of ``UnitRegistry.parse_expression`` (:issue:`40`) By `Tom Nicholas `_. - ensure the variables which causes the error is explicit if an error occurs in :py:meth:`Dataset.pint.quantify` and other methods (:pull:`43`, :issue:`91`) By `Tom Nicholas `_ and `Justus Magin `_. - refactor the internal conversion functions (:pull:`56`) By `Justus Magin `_. - allow converting indexes (except :py:class:`pandas.MultiIndex`) (:pull:`56`) By `Justus Magin `_. - document the reason for requiring the ``force_ndarray_like`` or ``force_ndarray`` options on unit registries (:pull:`59`) By `Justus Magin `_. - allow passing a format string to :py:meth:`Dataset.pint.dequantify` and :py:meth:`DataArray.pint.dequantify` (:pull:`49`) By `Justus Magin `_. - allow converting all data variables in a Dataset to the same units using :py:meth:`Dataset.pint.to` (:issue:`45`, :pull:`63`). By `Mika Pflüger `_. - update format of examples in docstrings (:pull:`64`). By `Mika Pflüger `_. - implement :py:meth:`Dataset.pint.sel` and :py:meth:`DataArray.pint.sel` (:pull:`60`). By `Justus Magin `_. - implement :py:attr:`Dataset.pint.loc` and :py:attr:`DataArray.pint.loc` (:pull:`79`). By `Justus Magin `_. - implement :py:meth:`Dataset.pint.drop_sel` and :py:meth:`DataArray.pint.drop_sel` (:pull:`73`). By `Justus Magin `_. - implement :py:meth:`Dataset.pint.chunk` and :py:meth:`DataArray.pint.chunk` (:pull:`83`). By `Justus Magin `_. - implement :py:meth:`Dataset.pint.reindex`, :py:meth:`Dataset.pint.reindex_like`, :py:meth:`DataArray.pint.reindex` and :py:meth:`DataArray.pint.reindex_like` (:pull:`69`). By `Justus Magin `_. - implement :py:meth:`Dataset.pint.interp`, :py:meth:`Dataset.pint.interp_like`, :py:meth:`DataArray.pint.interp` and :py:meth:`DataArray.pint.interp_like` (:pull:`72`, :pull:`76`, :pull:`97`). By `Justus Magin `_. - implement :py:meth:`Dataset.pint.ffill`, :py:meth:`Dataset.pint.bfill`, :py:meth:`DataArray.pint.ffill` and :py:meth:`DataArray.pint.bfill` (:pull:`78`). By `Justus Magin `_. - implement :py:meth:`Dataset.pint.interpolate_na` and :py:meth:`DataArray.pint.interpolate_na` (:pull:`82`). By `Justus Magin `_. - expose :py:func:`pint_xarray.setup_registry` as public API (:pull:`89`) By `Justus Magin `_. v0.1 (October 26 2020) ---------------------- - add initial draft of documentation (:pull:`13`, :pull:`20`) - implement :py:meth:`DataArray.pint.to` and :py:meth:`Dataset.pint.to` (:pull:`11`) - rewrite :py:meth:`DataArray.pint.quantify`, :py:meth:`Dataset.pint.quantify`, :py:meth:`DataArray.pint.dequantify` and :py:meth:`Dataset.pint.dequantify` (:pull:`17`) - expose :py:func:`pint_xarray.testing.assert_units_equal` as public API (:pull:`24`) - fix the :py:attr:`DataArray.pint.units`, :py:attr:`DataArray.pint.magnitude` and :py:attr:`DataArray.pint.dimensionality` properties and add docstrings for all three. (:pull:`31`) - use ``pint``'s application registry as a module-global registry (:pull:`32`) pint-xarray-0.6.1/licenses/000077500000000000000000000000001516032446100156035ustar00rootroot00000000000000pint-xarray-0.6.1/licenses/XARRAY_LICENSE000066400000000000000000000241121516032446100176760ustar00rootroot00000000000000Copyright 2014-2020, xarray developers Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 1. Definitions. "License" shall mean the terms and conditions for use, reproduction, and distribution as defined by Sections 1 through 9 of this document. "Licensor" shall mean the copyright owner or entity authorized by the copyright owner that is granting the License. "Legal Entity" shall mean the union of the acting entity and all other entities that control, are controlled by, or are under common control with that entity. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. "You" (or "Your") shall mean an individual or Legal Entity exercising permissions granted by this License. "Source" form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files. "Object" form shall mean any form resulting from mechanical transformation or translation of a Source form, including but not limited to compiled object code, generated documentation, and conversions to other media types. "Work" shall mean the work of authorship, whether in Source or Object form, made available under the License, as indicated by a copyright notice that is included in or attached to the work (an example is provided in the Appendix below). "Derivative Works" shall mean any work, whether in Source or Object form, that is based on (or derived from) the Work and for which the editorial revisions, annotations, elaborations, or other modifications represent, as a whole, an original work of authorship. For the purposes of this License, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, the Work and Derivative Works thereof. "Contribution" shall mean any work of authorship, including the original version of the Work and any modifications or additions to that Work or Derivative Works thereof, that is intentionally submitted to Licensor for inclusion in the Work by the copyright owner or by an individual or Legal Entity authorized to submit on behalf of the copyright owner. For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to the Licensor or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, the Licensor for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution." "Contributor" shall mean Licensor and any individual or Legal Entity on behalf of whom a Contribution has been received by Licensor and subsequently incorporated within the Work. 2. Grant of Copyright License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable copyright license to reproduce, prepare Derivative Works of, publicly display, publicly perform, sublicense, and distribute the Work and such Derivative Works in Source or Object form. 3. Grant of Patent License. Subject to the terms and conditions of this License, each Contributor hereby grants to You a perpetual, worldwide, non-exclusive, no-charge, royalty-free, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by such Contributor that are necessarily infringed by their Contribution(s) alone or by combination of their Contribution(s) with the Work to which such Contribution(s) was submitted. If You institute patent litigation against any entity (including a cross-claim or counterclaim in a lawsuit) alleging that the Work or a Contribution incorporated within the Work constitutes direct or contributory patent infringement, then any patent licenses granted to You under this License for that Work shall terminate as of the date such litigation is filed. 4. Redistribution. You may reproduce and distribute copies of the Work or Derivative Works thereof in any medium, with or without modifications, and in Source or Object form, provided that You meet the following conditions: You must give any other recipients of the Work or Derivative Works a copy of this License; and You must cause any modified files to carry prominent notices stating that You changed the files; and You must retain, in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work, excluding those notices that do not pertain to any part of the Derivative Works; and If the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation, if provided along with the Derivative Works; or, within a display generated by the Derivative Works, if and wherever such third-party notices normally appear. The contents of the NOTICE file are for informational purposes only and do not modify the License. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from the Work, provided that such additional attribution notices cannot be construed as modifying the License. You may add Your own copyright statement to Your modifications and may provide additional or different license terms and conditions for use, reproduction, or distribution of Your modifications, or for any such Derivative Works as a whole, provided Your use, reproduction, and distribution of the Work otherwise complies with the conditions stated in this License. 5. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in the Work by You to the Licensor shall be under the terms and conditions of this License, without any additional terms or conditions. Notwithstanding the above, nothing herein shall supersede or modify the terms of any separate license agreement you may have executed with Licensor regarding such Contributions. 6. Trademarks. This License does not grant permission to use the trade names, trademarks, service marks, or product names of the Licensor, except as required for reasonable and customary use in describing the origin of the Work and reproducing the content of the NOTICE file. 7. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Licensor provides the Work (and each Contributor provides its Contributions) on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. You are solely responsible for determining the appropriateness of using or redistributing the Work and assume any risks associated with Your exercise of permissions under this License. 8. Limitation of Liability. In no event and under no legal theory, whether in tort (including negligence), contract, or otherwise, unless required by applicable law (such as deliberate and grossly negligent acts) or agreed to in writing, shall any Contributor be liable to You for damages, including any direct, indirect, special, incidental, or consequential damages of any character arising as a result of this License or out of the use or inability to use the Work (including but not limited to damages for loss of goodwill, work stoppage, computer failure or malfunction, or any and all other commercial damages or losses), even if such Contributor has been advised of the possibility of such damages. 9. Accepting Warranty or Additional Liability. While redistributing the Work or Derivative Works thereof, You may choose to offer, and charge a fee for, acceptance of support, warranty, indemnity, or other liability obligations and/or rights consistent with this License. However, in accepting such obligations, You may act only on Your own behalf and on Your sole responsibility, not on behalf of any other Contributor, and only if You agree to indemnify, defend, and hold each Contributor harmless for any liability incurred by, or claims asserted against, such Contributor by reason of your accepting any such warranty or additional liability. END OF TERMS AND CONDITIONS APPENDIX: How to apply the Apache License to your work To apply the Apache License to your work, attach the following boilerplate notice, with the fields enclosed by brackets "[]" replaced with your own identifying information. (Don't include the brackets!) The text should be enclosed in the appropriate comment syntax for the file format. We also recommend that a file or class name and description of purpose be included on the same "printed page" as the copyright notice for easier identification within third-party archives. Copyright [yyyy] [name of copyright owner] Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at http://www.apache.org/licenses/LICENSE-2.0 Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. pint-xarray-0.6.1/pint_xarray/000077500000000000000000000000001516032446100163365ustar00rootroot00000000000000pint-xarray-0.6.1/pint_xarray/__init__.py000066400000000000000000000012771516032446100204560ustar00rootroot00000000000000from importlib.metadata import version import pint from pint_xarray import accessors, formatting, testing # noqa: F401 from pint_xarray._expects import expects from pint_xarray.accessors import default_registry as unit_registry from pint_xarray.accessors import setup_registry from pint_xarray.index import PintIndex try: __version__ = version("pint-xarray") except Exception: # pragma: no cover # Local copy or not installed with setuptools. # Disable minimum version checks on downstream libraries. __version__ = "999" pint.Quantity._repr_inline_ = formatting.inline_repr __all__ = [ "testing", "unit_registry", "setup_registry", "PintIndex", "expects", ] pint-xarray-0.6.1/pint_xarray/_expects.py000066400000000000000000000230161516032446100205240ustar00rootroot00000000000000import functools import inspect import itertools from inspect import Parameter import pint import pint.testing import xarray as xr from pint_xarray.accessors import get_registry from pint_xarray.conversion import extract_units, strip_units from pint_xarray.itertools import zip_mappings variable_parameters = (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD) def _number_of_results(result): if isinstance(result, tuple): return len(result) elif result is None: return 0 else: return 1 def expects(*args_units, return_value=None, **kwargs_units): """ Decorator which ensures the inputs and outputs of the decorated function are expressed in the expected units. Arguments to the decorated function are checked for the specified units, converting to those units if necessary, and then stripped of their units before being passed into the undecorated function. Therefore the undecorated function should expect unquantified DataArrays, Datasets, or numpy-like arrays, but with the values expressed in specific units. Parameters ---------- func : callable Function to decorate, which accepts zero or more xarray.DataArrays or numpy-like arrays as inputs, and may optionally return one or more xarray.DataArrays or numpy-like arrays. *args_units : unit-like or mapping of hashable to unit-like, optional Units to expect for each positional argument given to func. The decorator will first check that arguments passed to the decorated function possess these specific units (or will attempt to convert the argument to these units), then will strip the units before passing the magnitude to the wrapped function. A value of None indicates not to check that argument for units (suitable for flags and other non-data arguments). return_value : unit-like or list of unit-like or mapping of hashable to unit-like \ or list of mapping of hashable to unit-like, optional The expected units of the returned value(s), either as a single unit or as a list of units. The decorator will attach these units to the variables returned from the function. A value of None indicates not to attach any units to that return value (suitable for flags and other non-data results). **kwargs_units : mapping of hashable to unit-like, optional Unit to expect for each keyword argument given to func. The decorator will first check that arguments passed to the decorated function possess these specific units (or will attempt to convert the argument to these units), then will strip the units before passing the magnitude to the wrapped function. A value of None indicates not to check that argument for units (suitable for flags and other non-data arguments). Returns ------- return_values : Any Return values of the wrapped function, either a single value or a tuple of values. These will be given units according to ``return_value``. Raises ------ TypeError If any of the units are not a valid type. ValueError If the number of arguments or return values does not match the number of units specified. Also thrown if any parameter does not have a unit specified. See Also -------- pint.wraps Examples -------- Decorating a function which takes one quantified input, but returns a non-data value (in this case a boolean). >>> @expects("deg C") ... def above_freezing(temp): ... return temp > 0 ... Decorating a function which allows any dimensions for the array, but also accepts an optional `weights` keyword argument, which must be dimensionless. >>> @expects(None, weights="dimensionless") ... def mean(da, weights=None): ... if weights: ... return da.weighted(weights=weights).mean() ... else: ... return da.mean() ... """ def outer(func): signature = inspect.signature(func) params_units = signature.bind(*args_units, **kwargs_units) missing_params = [ name for name, p in signature.parameters.items() if p.kind not in variable_parameters and name not in params_units.arguments ] if missing_params: raise ValueError( "Missing units for the following parameters: " + ", ".join(map(repr, missing_params)) ) n_expected_results = _number_of_results(return_value) @functools.wraps(func) def wrapper(*args, **kwargs): nonlocal return_value params = signature.bind(*args, **kwargs) # don't apply defaults, as those can't be quantities and thus must # already be in the correct units spec_units = dict( enumerate( itertools.chain.from_iterable( spec.values() if isinstance(spec, dict) else (spec,) for spec in params_units.arguments.values() if spec is not None ) ) ) params_units_ = dict( enumerate( itertools.chain.from_iterable( ( extract_units(param) if isinstance(param, (xr.DataArray, xr.Dataset)) else (param.units,) ) for name, param in params.arguments.items() if isinstance(param, (xr.DataArray, xr.Dataset, pint.Quantity)) ) ) ) ureg = get_registry( None, dict(spec_units) if spec_units else {}, dict(params_units_) if params_units else {}, ) errors = [] for name, (value, units) in zip_mappings( params.arguments, params_units.arguments ): try: if units is None: if isinstance(value, pint.Quantity) or ( isinstance(value, (xr.DataArray, xr.Dataset)) and value.pint.units ): raise TypeError( "Passed in a quantity where none was expected" ) continue if isinstance(value, pint.Quantity): params.arguments[name] = value.m_as(units) elif isinstance(value, (xr.DataArray, xr.Dataset)): params.arguments[name] = strip_units(value.pint.to(units)) else: raise TypeError( f"Attempting to convert non-quantity {value} to {units}." ) except ( TypeError, pint.errors.UndefinedUnitError, pint.errors.DimensionalityError, ) as e: e.add_note( f"expects: raised while trying to convert parameter {name}" ) errors.append(e) if errors: raise ExceptionGroup("Errors while converting parameters", errors) result = func(*params.args, **params.kwargs) n_results = _number_of_results(result) if return_value is not None and ( (isinstance(result, tuple) ^ isinstance(return_value, tuple)) or (n_results != n_expected_results) ): message = "mismatched number of return values:" if n_results != n_expected_results: message += f" expected {n_expected_results} but got {n_results}." elif isinstance(result, tuple) and not isinstance(return_value, tuple): message += ( " expected a single return value but got a 1-sized tuple." ) else: message += ( " expected a 1-sized tuple but got a single return value." ) raise ValueError(message) if result is None: return if not isinstance(result, tuple): result = (result,) if not isinstance(return_value, tuple): return_value = (return_value,) final_result = [] errors = [] for index, (value, units) in enumerate(zip(result, return_value)): if units is not None: try: if isinstance(value, (xr.Dataset, xr.DataArray)): value = value.pint.quantify(units) else: value = ureg.Quantity(value, units) except Exception as e: e.add_note( f"expects: raised while trying to convert return value {index}" ) errors.append(e) final_result.append(value) if errors: raise ExceptionGroup("Errors while converting return values", errors) if n_results == 1: return final_result[0] return tuple(final_result) return wrapper return outer pint-xarray-0.6.1/pint_xarray/accessors.py000066400000000000000000001611041516032446100207000ustar00rootroot00000000000000# TODO is it possible to import pint-xarray from within xarray if pint is present? import itertools import pint from pint import Unit from xarray import register_dataarray_accessor, register_dataset_accessor from xarray.core.dtypes import NA from pint_xarray import conversion from pint_xarray.conversion import no_unit_values from pint_xarray.errors import create_exception_group _default = object() def setup_registry(registry): """set up the given registry for use with pint_xarray Namely, it enables ``force_ndarray_like`` to make sure results are always duck arrays. Parameters ---------- registry : pint.UnitRegistry The registry to modify """ if not registry.force_ndarray and not registry.force_ndarray_like: registry.force_ndarray_like = True return registry default_registry = setup_registry(pint.get_application_registry()) # TODO could/should we overwrite xr.open_dataset and xr.open_mfdataset to make # them apply units upon loading??? # TODO could even override the decode_cf kwarg? # TODO docstrings # TODO type hints def is_dict_like(obj): return hasattr(obj, "keys") and hasattr(obj, "__getitem__") def zip_mappings(*mappings, fill_value=None): """zip mappings by combining values for common keys into a tuple Works like itertools.zip_longest, so if a key is missing from a mapping, it is replaced by ``fill_value``. Parameters ---------- *mappings : dict-like The mappings to zip fill_value The value to use if a key is missing from a mapping. Returns ------- zipped : dict-like The zipped mapping """ keys = set(itertools.chain.from_iterable(mapping.keys() for mapping in mappings)) # TODO: could this be made more efficient using itertools.groupby? zipped = { key: tuple(mapping.get(key, fill_value) for mapping in mappings) for key in keys } return zipped def units_to_str_or_none(mapping, unit_format): formatter = str if not unit_format else lambda v: unit_format.format(v) return { key: formatter(value) if isinstance(value, Unit) else value for key, value in mapping.items() } # based on xarray.core.utils.either_dict_or_kwargs # https://github.com/pydata/xarray/blob/v0.15.1/xarray/core/utils.py#L249-L268 def either_dict_or_kwargs(positional, keywords, method_name): if positional not in (_default, None): if not is_dict_like(positional): raise ValueError( f"the first argument to .{method_name} must be a dictionary" ) if keywords: raise ValueError( "cannot specify both keyword and positional " f"arguments to .{method_name}" ) return positional else: return keywords def get_registry(unit_registry, new_units, existing_units): units = itertools.chain(new_units.values(), existing_units.values()) registries = {unit._REGISTRY for unit in units if isinstance(unit, Unit)} if unit_registry is None: if not registries: unit_registry = default_registry elif len(registries) == 1: (unit_registry,) = registries registries.add(unit_registry) if len(registries) > 1 or unit_registry not in registries: raise ValueError( "using multiple unit registries in the same object is not supported" ) if not unit_registry.force_ndarray_like and not unit_registry.force_ndarray: raise ValueError( "invalid registry. Please enable 'force_ndarray_like' or 'force_ndarray'." ) return unit_registry def _decide_units(units, registry, unit_attribute): if units is _default and unit_attribute in (None, _default): # or warn and return None? raise ValueError("no units given") elif units in no_unit_values or isinstance(units, Unit): # TODO what happens if they pass in a Unit from a different registry return units elif units is _default: if unit_attribute in no_unit_values: return unit_attribute if isinstance(unit_attribute, Unit): units = unit_attribute else: units = registry.parse_units(unit_attribute) else: units = registry.parse_units(units) return units class DatasetLocIndexer: __slots__ = ("ds",) def __init__(self, ds): self.ds = ds def __getitem__(self, indexers): if not is_dict_like(indexers): raise NotImplementedError("pandas-style indexing is not supported, yet") dims = self.ds.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units try: converted = conversion.convert_units(self.ds, indexer_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(indexers) stripped = conversion.strip_units(converted) converted_units = conversion.extract_units(converted) indexed = stripped.loc[stripped_indexers] return conversion.attach_units(indexed, converted_units) class DataArrayLocIndexer: __slots__ = ("da",) def __init__(self, da): self.da = da def __getitem__(self, indexers): if not is_dict_like(indexers): raise NotImplementedError("pandas-style indexing is not supported, yet") dims = self.da.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units try: converted = conversion.convert_units(self.da, indexer_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(indexers) stripped = conversion.strip_units(converted) converted_units = conversion.extract_units(converted) indexed = stripped.loc[stripped_indexers] return conversion.attach_units(indexed, converted_units) def __setitem__(self, indexers, values): if not is_dict_like(indexers): raise NotImplementedError("pandas-style indexing is not supported, yet") dims = self.da.dims unit_attrs = conversion.extract_unit_attributes(self.da) index_units = { name: units for name, units in unit_attrs.items() if name in dims } # convert the indexers to the index units try: converted = conversion.convert_indexer_units(indexers, index_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(converted) self.da.loc[stripped_indexers] = values @register_dataarray_accessor("pint") class PintDataArrayAccessor: """ Access methods for DataArrays with units using Pint. Methods and attributes can be accessed through the `.pint` attribute. """ def __init__(self, da): self.da = da def quantify(self, units=_default, unit_registry=None, **unit_kwargs): """ Attach units to the DataArray. Units can be specified as a pint.Unit or as a string, which will be parsed by the given unit registry. If no units are specified then the units will be parsed from the `'units'` entry of the DataArray's `.attrs`. Will raise a ValueError if the DataArray already contains a unit-aware array with a different unit. .. note:: Be aware that unless you're using ``dask`` this will load the data into memory. To avoid that, consider converting to ``dask`` first (e.g. using ``chunk``). .. note:: Also note that datetime units (i.e. ones that match ``{units} since {date}``) in unit attributes will be ignored, to avoid interfering with ``xarray``'s datetime encoding / decoding. Parameters ---------- units : unit-like or mapping of hashable to unit-like, optional Physical units to use for this DataArray. If a str or pint.Unit, will be used as the DataArray's units. If a dict-like, it should map a variable name to the desired unit (use the DataArray's name to refer to its data). If not provided, ``quantify`` will try to read them from ``DataArray.attrs['units']`` using pint's parser. The ``"units"`` attribute will be removed from all variables except from dimension coordinates. unit_registry : pint.UnitRegistry, optional Unit registry to be used for the units attached to this DataArray. If not given then a default registry will be created. **unit_kwargs Keyword argument form of units. Returns ------- quantified : DataArray DataArray whose wrapped array data will now be a Quantity array with the specified units. Notes ----- ``"none"`` and ``None`` can be used to mark variables that should not be quantified. Examples -------- >>> da = xr.DataArray( ... data=[0.4, 0.9, 1.7, 4.8, 3.2, 9.1], ... dims=["wavelength"], ... coords={"wavelength": [1e-4, 2e-4, 4e-4, 6e-4, 1e-3, 2e-3]}, ... ) >>> da.pint.quantify(units="Hz") Size: 48B Coordinates: * wavelength (wavelength) float64 48B 0.0001 0.0002 0.0004 0.0006 0.001 0.002 Don't quantify the data: >>> da = xr.DataArray( ... data=[0.4, 0.9], ... dims=["wavelength"], ... attrs={"units": "Hz"}, ... ) >>> da.pint.quantify(units=None) Size: 16B array([0.4, 0.9]) Dimensions without coordinates: wavelength Quantify with the same unit: >>> q = da.pint.quantify() >>> q Size: 16B Dimensions without coordinates: wavelength >>> q.pint.quantify("Hz") Size: 16B Dimensions without coordinates: wavelength """ if units is None or isinstance(units, (str, pint.Unit)): if self.da.name in unit_kwargs: raise ValueError( f"ambiguous values given for {repr(self.da.name)}:" f" {repr(units)} and {repr(unit_kwargs[self.da.name])}" ) unit_kwargs[self.da.name] = units units = None units = either_dict_or_kwargs(units, unit_kwargs, "quantify") registry = get_registry(unit_registry, units, conversion.extract_units(self.da)) unit_attrs = conversion.extract_unit_attributes(self.da) possible_new_units = zip_mappings(units, unit_attrs, fill_value=_default) new_units = {} invalid_units = {} for name, (unit, attr) in possible_new_units.items(): if unit not in (_default, None) or attr not in (_default, None): try: new_units[name] = _decide_units(unit, registry, attr) except (ValueError, pint.UndefinedUnitError) as e: if unit not in (_default, None): type = "parameter" reported_unit = unit else: type = "attribute" reported_unit = attr invalid_units[name] = (reported_unit, type, e) if invalid_units: raise create_exception_group(invalid_units, "parse") existing_units = { name: unit for name, unit in conversion.extract_units(self.da).items() if isinstance(unit, Unit) } overwritten_units = { name: (old, new) for name, (old, new) in zip_mappings( existing_units, new_units, fill_value=_default ).items() if old is not _default and new is not _default and old != new } if overwritten_units: errors = { name: ( new, ValueError( f"Cannot attach unit {repr(new)} to quantity: data " f"already has units {repr(old)}" ), ) for name, (old, new) in overwritten_units.items() } raise create_exception_group(errors, "attach") return self.da.pipe(conversion.strip_unit_attributes).pipe( conversion.attach_units, new_units ) def dequantify(self, format=None): r""" Convert the units of the DataArray to string attributes. Will replace ``.attrs['units']`` on each variable with a string representation of the ``pint.Unit`` instance. Parameters ---------- format : str, default: None The format specification (as accepted by pint) used for the string representations. If ``None``, the registry's default (:py:attr:`pint.UnitRegistry.default_format`) is used instead. Returns ------- dequantified : DataArray DataArray whose array data is unitless, and of the type that was previously wrapped by `pint.Quantity`. See Also -------- :doc:`pint:user/formatting` pint's string formatting guide Examples -------- >>> da = xr.DataArray([0, 1], dims="x") >>> q = da.pint.quantify("m / s") >>> q Size: 16B Dimensions without coordinates: x >>> q.pint.dequantify(format="P") Size: 16B array([0, 1]) Dimensions without coordinates: x Attributes: units: meter/second >>> q.pint.dequantify(format="~P") Size: 16B array([0, 1]) Dimensions without coordinates: x Attributes: units: m/s Use the registry's default format >>> pint_xarray.unit_registry.default_format = "~L" >>> q.pint.dequantify() Size: 16B array([0, 1]) Dimensions without coordinates: x Attributes: units: \frac{\mathrm{m}}{\mathrm{s}} """ units = conversion.extract_unit_attributes(self.da) units.update(conversion.extract_units(self.da)) unit_format = f"{{:{format}}}" if isinstance(format, str) else format units = units_to_str_or_none(units, unit_format) return ( self.da.pipe(conversion.strip_units) .pipe(conversion.strip_unit_attributes) .pipe(conversion.attach_unit_attributes, units) ) @property def magnitude(self): """the magnitude of the data or the data itself if not a quantity.""" data = self.da.data return getattr(data, "magnitude", data) @property def units(self): """the units of the data or :py:obj:`None` if not a quantity. Setting the units is possible, but only if the data is not already a quantity. """ return getattr(self.da.data, "units", None) @units.setter def units(self, units): self.da.data = conversion.array_attach_units(self.da.data, units) @property def dimensionality(self): """get the dimensionality of the data or :py:obj:`None` if not a quantity.""" return getattr(self.da.data, "dimensionality", None) @property def registry(self): # TODO is this a bad idea? (see GH issue #1071 in pint) return getattr(self.da.data, "_REGISTRY", None) @registry.setter def registry(self, _): raise AttributeError("Don't try to change the registry once created") def to(self, units=None, **unit_kwargs): """convert the quantities in a DataArray Parameters ---------- units : unit-like or mapping of hashable to unit-like, optional The units to convert to. If a unit name or ``pint.Unit`` object, convert the DataArray's data. If a dict-like, it has to map a variable name to a unit name or ``pint.Unit`` object. **unit_kwargs The kwargs form of ``units``. Can only be used for variable names that are strings and valid python identifiers. Returns ------- object : DataArray A new object with converted units. Examples -------- >>> da = xr.DataArray( ... data=np.linspace(0, 1, 5) * ureg.m, ... coords={"u": ("x", np.arange(5) * ureg.s)}, ... dims="x", ... name="arr", ... ) >>> da Size: 40B Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Convert the data >>> da.pint.to("mm") Size: 40B Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x >>> da.pint.to(ureg.mm) Size: 40B Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x >>> da.pint.to({da.name: "mm"}) Size: 40B Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Convert coordinates >>> da.pint.to({"u": ureg.ms}) Size: 40B Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x >>> da.pint.to(u="ms") Size: 40B Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x Convert both simultaneously >>> da.pint.to("mm", u="ms") Size: 40B Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x >>> da.pint.to({"arr": ureg.mm, "u": ureg.ms}) Size: 40B Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x >>> da.pint.to(arr="mm", u="ms") Size: 40B Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x """ if isinstance(units, (str, pint.Unit)): unit_kwargs[self.da.name] = units units = None elif units is not None and not is_dict_like(units): raise ValueError( "units must be either a string, a pint.Unit object or a dict-like," f" but got {units!r}" ) units = either_dict_or_kwargs(units, unit_kwargs, "to") return conversion.convert_units(self.da, units) def chunk(self, chunks, name_prefix="xarray-", token=None, lock=False): """unit-aware version of chunk Like :py:meth:`xarray.DataArray.chunk`, but chunking a quantity will change the wrapped type to ``dask``. .. note:: It is recommended to only use this when chunking in-memory arrays. To rechunk please use :py:meth:`xarray.DataArray.chunk`. See Also -------- xarray.DataArray.chunk xarray.Dataset.pint.chunk """ units = conversion.extract_units(self.da) stripped = conversion.strip_units(self.da) chunked = stripped.chunk( chunks, name_prefix=name_prefix, token=token, lock=lock ) return conversion.attach_units(chunked, units) def reindex( self, indexers=None, method=None, tolerance=None, copy=True, fill_value=NA, **indexers_kwargs, ): """unit-aware version of reindex Like :py:meth:`xarray.DataArray.reindex`, except the object's indexes are converted to the units of the indexers first. .. note:: ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to ``DataArray.reindex`` unmodified. See Also -------- xarray.Dataset.pint.reindex xarray.DataArray.pint.reindex_like xarray.DataArray.reindex """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") dims = self.da.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance # TODO: handle fill_value # convert the indexes to the indexer's units converted = conversion.convert_units(self.da, indexer_units) converted_units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) # index stripped_indexers = conversion.strip_indexer_units(indexers) indexed = stripped.reindex( stripped_indexers, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) return conversion.attach_units(indexed, converted_units) def reindex_like( self, other, method=None, tolerance=None, copy=True, fill_value=NA ): """unit-aware version of reindex_like Like :py:meth:`xarray.DataArray.reindex_like`, except the object's indexes are converted to the units of the indexers first. .. note:: ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to ``DataArray.reindex_like`` unmodified. See Also -------- xarray.Dataset.pint.reindex_like xarray.DataArray.pint.reindex xarray.DataArray.reindex_like """ indexer_units = conversion.extract_units(other) converted = conversion.convert_units(self.da, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) stripped_other = conversion.strip_units(other) # TODO: handle tolerance # TODO: handle fill_value reindexed = stripped.reindex_like( stripped_other, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) return conversion.attach_units(reindexed, units) def interp( self, coords=None, method="linear", assume_sorted=False, kwargs=None, **coords_kwargs, ): """unit-aware version of interp Like :py:meth:`xarray.DataArray.interp`, except the object's indexes are converted to the units of the indexers first. .. note:: ``method``, ``assume_sorted`` and ``kwargs`` are passed unmodified to ``DataArray.interp``. See Also -------- xarray.Dataset.pint.interp xarray.DataArray.pint.interp_like xarray.DataArray.interp """ indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp") dims = self.da.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units converted = conversion.convert_units(self.da, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) # index stripped_indexers = conversion.strip_indexer_units(indexers) interpolated = stripped.interp( stripped_indexers, method=method, assume_sorted=assume_sorted, kwargs=kwargs, ) return conversion.attach_units(interpolated, units) def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None): """unit-aware version of interp_like Like :py:meth:`xarray.DataArray.interp_like`, except the object's indexes are converted to the units of the indexers first. .. note:: ``method``, ``assume_sorted`` and ``kwargs`` are passed unmodified to ``DataArray.interp``. See Also -------- xarray.Dataset.pint.interp_like xarray.DataArray.pint.interp xarray.DataArray.interp_like """ indexer_units = conversion.extract_units(other) converted = conversion.convert_units(self.da, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) stripped_other = conversion.strip_units(other) interpolated = stripped.interp_like( stripped_other, method=method, assume_sorted=assume_sorted, kwargs=kwargs, ) return conversion.attach_units(interpolated, units) def sel( self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs ): """unit-aware version of sel Like :py:meth:`xarray.DataArray.sel`, except the object's indexes are converted to the units of the indexers first. .. note:: ``tolerance`` is not supported, yet. It will be passed through to ``DataArray.sel`` unmodified. See Also -------- xarray.Dataset.pint.sel xarray.DataArray.sel xarray.Dataset.sel """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel") dims = self.da.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance # convert the indexes to the indexer's units try: converted = conversion.convert_units(self.da, indexer_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(indexers) stripped = conversion.strip_units(converted) converted_units = conversion.extract_units(converted) indexed = stripped.sel( stripped_indexers, method=method, tolerance=tolerance, drop=drop, ) return conversion.attach_units(indexed, converted_units) @property def loc(self): """Unit-aware attribute for indexing .. note:: Position based indexing (e.g. ``ds.loc[1, 2:]``) is not supported, yet See Also -------- xarray.DataArray.loc """ return DataArrayLocIndexer(self.da) def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): """unit-aware version of drop_sel Just like :py:meth:`xarray.DataArray.drop_sel`, except the indexers are converted to the units of the object's indexes first. See Also -------- xarray.Dataset.pint.drop_sel xarray.DataArray.drop_sel xarray.Dataset.drop_sel """ indexers = either_dict_or_kwargs(labels, labels_kwargs, "drop_sel") dims = self.da.dims unit_attrs = conversion.extract_unit_attributes(self.da) index_units = { name: units for name, units in unit_attrs.items() if name in dims } # convert the indexers to the indexes units try: converted_indexers = conversion.convert_indexer_units(indexers, index_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(converted_indexers) indexed = self.da.drop_sel( stripped_indexers, errors=errors, ) return indexed def ffill(self, dim, limit=None): """unit-aware version of ffill Like :py:meth:`xarray.DataArray.ffill` but without stripping the data units. See Also -------- xarray.DataArray.ffill xarray.DataArray.pint.bfill """ units = conversion.extract_units(self.da) stripped = conversion.strip_units(self.da) filled = stripped.ffill(dim=dim, limit=limit) return conversion.attach_units(filled, units) def bfill(self, dim, limit=None): """unit-aware version of bfill Like :py:meth:`xarray.DataArray.bfill` but without stripping the data units. See Also -------- xarray.DataArray.bfill xarray.DataArray.pint.ffill """ units = conversion.extract_units(self.da) stripped = conversion.strip_units(self.da) filled = stripped.bfill(dim=dim, limit=limit) return conversion.attach_units(filled, units) def interpolate_na( self, dim=None, method="linear", limit=None, use_coordinate=True, max_gap=None, keep_attrs=None, **kwargs, ): """unit-aware version of interpolate_na Like :py:meth:`xarray.DataArray.interpolate_na` but without stripping the units on data or coordinates. .. note:: ``max_gap`` is not supported, yet, and will be passed through to ``DataArray.interpolate_na`` unmodified. See Also -------- xarray.Dataset.pint.interpolate_na xarray.DataArray.interpolate_na """ units = conversion.extract_units(self.da) stripped = conversion.strip_units(self.da) interpolated = stripped.interpolate_na( dim=dim, method=method, limit=limit, use_coordinate=use_coordinate, max_gap=max_gap, keep_attrs=keep_attrs, **kwargs, ) return conversion.attach_units(interpolated, units) @register_dataset_accessor("pint") class PintDatasetAccessor: """ Access methods for DataArrays with units using Pint. Methods and attributes can be accessed through the `.pint` attribute. """ def __init__(self, ds): self.ds = ds def quantify(self, units=_default, unit_registry=None, **unit_kwargs): """ Attach units to the variables of the Dataset. Units can be specified as a ``pint.Unit`` or as a string, which will be parsed by the given unit registry. If no units are specified then the units will be parsed from the ``"units"`` entry of the Dataset variable's ``.attrs``. Will raise a ValueError if any of the variables already contain a unit-aware array with a different unit. .. note:: Be aware that unless you're using ``dask`` this will load the data into memory. To avoid that, consider converting to ``dask`` first (e.g. using ``chunk``). .. note:: Also note that datetime units (i.e. ones that match ``{units} since {date}``) in unit attributes will be ignored, to avoid interfering with ``xarray``'s datetime encoding / decoding. Parameters ---------- units : mapping of hashable to unit-like, optional Physical units to use for particular DataArrays in this Dataset. It should map variable names to units (unit names or ``pint.Unit`` objects). If not provided, ``quantify`` will try to read them from ``Dataset[var].attrs['units']`` using pint's parser. The ``"units"`` attribute will be removed from all variables except from dimension coordinates. unit_registry : pint.UnitRegistry, optional Unit registry to be used for the units attached to each DataArray in this Dataset. If not given then a default registry will be created. **unit_kwargs Keyword argument form of ``units``. Returns ------- quantified : Dataset Dataset whose variables will now contain Quantity arrays with units. Notes ----- ``"none"`` and ``None`` can be used to mark variables that should not be quantified. Examples -------- >>> ds = xr.Dataset( ... {"a": ("x", [0, 3, 2], {"units": "m"}), "b": ("x", [5, -2, 1])}, ... coords={"x": [0, 1, 2], "u": ("x", [-1, 0, 1], {"units": "s"})}, ... ) >>> ds.pint.quantify() Size: 96B Dimensions: (x: 3) Coordinates: * x (x) int64 24B 0 1 2 u (x) int64 24B [s] -1 0 1 Data variables: a (x) int64 24B [m] 0 3 2 b (x) int64 24B 5 -2 1 >>> ds.pint.quantify({"b": "dm"}) Size: 96B Dimensions: (x: 3) Coordinates: * x (x) int64 24B 0 1 2 u (x) int64 24B [s] -1 0 1 Data variables: a (x) int64 24B [m] 0 3 2 b (x) int64 24B [dm] 5 -2 1 Don't quantify specific variables: >>> ds.pint.quantify({"a": None}) Size: 96B Dimensions: (x: 3) Coordinates: * x (x) int64 24B 0 1 2 u (x) int64 24B [s] -1 0 1 Data variables: a (x) int64 24B 0 3 2 b (x) int64 24B 5 -2 1 Quantify with the same unit: >>> q = ds.pint.quantify() >>> q Size: 96B Dimensions: (x: 3) Coordinates: * x (x) int64 24B 0 1 2 u (x) int64 24B [s] -1 0 1 Data variables: a (x) int64 24B [m] 0 3 2 b (x) int64 24B 5 -2 1 >>> q.pint.quantify({"a": "m"}) Size: 96B Dimensions: (x: 3) Coordinates: * x (x) int64 24B 0 1 2 u (x) int64 24B [s] -1 0 1 Data variables: a (x) int64 24B [m] 0 3 2 b (x) int64 24B 5 -2 1 """ units = either_dict_or_kwargs(units, unit_kwargs, "quantify") registry = get_registry(unit_registry, units, conversion.extract_units(self.ds)) unit_attrs = conversion.extract_unit_attributes(self.ds) possible_new_units = zip_mappings(units, unit_attrs, fill_value=_default) new_units = {} invalid_units = {} for name, (unit, attr) in possible_new_units.items(): if unit is not _default or attr not in (None, _default): try: new_units[name] = _decide_units(unit, registry, attr) except (ValueError, pint.UndefinedUnitError) as e: if unit is not _default: type = "parameter" reported_unit = unit else: type = "attribute" reported_unit = attr invalid_units[name] = (reported_unit, type, e) if invalid_units: raise create_exception_group(invalid_units, "parse") existing_units = { name: unit for name, unit in conversion.extract_units(self.ds).items() if isinstance(unit, Unit) } overwritten_units = { name: (old, new) for name, (old, new) in zip_mappings( existing_units, new_units, fill_value=_default ).items() if old is not _default and new is not _default and old != new } if overwritten_units: errors = { name: ( new, ValueError( f"Cannot attach unit {repr(new)} to quantity: data " f"already has units {repr(old)}" ), ) for name, (old, new) in overwritten_units.items() } raise create_exception_group(errors, "attach") return self.ds.pipe(conversion.strip_unit_attributes).pipe( conversion.attach_units, new_units ) def dequantify(self, format=None): r""" Convert units from the Dataset to string attributes. Will replace ``.attrs['units']`` on each variable with a string representation of the ``pint.Unit`` instance. Parameters ---------- format : str, default: None The format specification (as accepted by pint's unit formatter) used for the string representations. If ``None``, the registry's default (:py:attr:`pint.UnitRegistry.default_format`) is used instead. Returns ------- dequantified : Dataset Dataset whose data variables are unitless, and of the type that was previously wrapped by ``pint.Quantity``. See Also -------- :doc:`pint:user/formatting` pint's string formatting guide Examples -------- >>> ds = xr.Dataset({"a": ("x", [0, 1]), "b": ("y", [2, 3, 4])}) >>> q = ds.pint.quantify({"a": "m / s", "b": "s"}) >>> q Size: 40B Dimensions: (x: 2, y: 3) Dimensions without coordinates: x, y Data variables: a (x) int64 16B [m/s] 0 1 b (y) int64 24B [s] 2 3 4 >>> d = q.pint.dequantify(format="P") >>> d.a Size: 16B array([0, 1]) Dimensions without coordinates: x Attributes: units: meter/second >>> d.b Size: 24B array([2, 3, 4]) Dimensions without coordinates: y Attributes: units: second >>> d = q.pint.dequantify(format="~P") >>> d.a Size: 16B array([0, 1]) Dimensions without coordinates: x Attributes: units: m/s >>> d.b Size: 24B array([2, 3, 4]) Dimensions without coordinates: y Attributes: units: s Use the registry's default format >>> pint_xarray.unit_registry.default_format = "~L" >>> d = q.pint.dequantify() >>> d.a Size: 16B array([0, 1]) Dimensions without coordinates: x Attributes: units: \frac{\mathrm{m}}{\mathrm{s}} >>> d.b Size: 24B array([2, 3, 4]) Dimensions without coordinates: y Attributes: units: \mathrm{s} """ units = conversion.extract_unit_attributes(self.ds) units.update(conversion.extract_units(self.ds)) unit_format = f"{{:{format}}}" if isinstance(format, str) else format units = units_to_str_or_none(units, unit_format) return ( self.ds.pipe(conversion.strip_units) .pipe(conversion.strip_unit_attributes) .pipe(conversion.attach_unit_attributes, units) ) def to(self, units=None, **unit_kwargs): """convert the quantities in a Dataset Parameters ---------- units : unit-like or mapping of hashable to unit-like, optional The units to convert to. If a unit name or ``pint.Unit`` object, convert all the object's data variables. If a dict-like, it maps variable names to unit names or ``pint.Unit`` objects. **unit_kwargs The kwargs form of ``units``. Can only be used for variable names that are strings and valid python identifiers. Returns ------- object : Dataset A new object with converted units. Examples -------- >>> ds = xr.Dataset( ... data_vars={ ... "a": ("x", np.linspace(0, 1, 5) * ureg.m), ... "b": ("x", np.linspace(-1, 0, 5) * ureg.kg), ... }, ... coords={"u": ("x", np.arange(5) * ureg.s)}, ... ) >>> ds Size: 120B Dimensions: (x: 5) Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Data variables: a (x) float64 40B [m] 0.0 0.25 0.5 0.75 1.0 b (x) float64 40B [kg] -1.0 -0.75 -0.5 -0.25 0.0 Convert the data >>> ds.pint.to({"a": "mm", "b": ureg.g}) Size: 120B Dimensions: (x: 5) Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Data variables: a (x) float64 40B [mm] 0.0 250.0 500.0 750.0 1e+03 b (x) float64 40B [g] -1e+03 -750.0 -500.0 -250.0 0.0 >>> ds.pint.to(a=ureg.mm, b="g") Size: 120B Dimensions: (x: 5) Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Data variables: a (x) float64 40B [mm] 0.0 250.0 500.0 750.0 1e+03 b (x) float64 40B [g] -1e+03 -750.0 -500.0 -250.0 0.0 Convert coordinates >>> ds.pint.to({"u": ureg.ms}) Size: 120B Dimensions: (x: 5) Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x Data variables: a (x) float64 40B [m] 0.0 0.25 0.5 0.75 1.0 b (x) float64 40B [kg] -1.0 -0.75 -0.5 -0.25 0.0 >>> ds.pint.to(u="ms") Size: 120B Dimensions: (x: 5) Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x Data variables: a (x) float64 40B [m] 0.0 0.25 0.5 0.75 1.0 b (x) float64 40B [kg] -1.0 -0.75 -0.5 -0.25 0.0 Convert both simultaneously >>> ds.pint.to(a=ureg.mm, b=ureg.g, u="ms") Size: 120B Dimensions: (x: 5) Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x Data variables: a (x) float64 40B [mm] 0.0 250.0 500.0 750.0 1e+03 b (x) float64 40B [g] -1e+03 -750.0 -500.0 -250.0 0.0 >>> ds.pint.to({"a": "mm", "b": "g", "u": ureg.ms}) Size: 120B Dimensions: (x: 5) Coordinates: u (x) float64 40B [ms] 0.0 1e+03 2e+03 3e+03 4e+03 Dimensions without coordinates: x Data variables: a (x) float64 40B [mm] 0.0 250.0 500.0 750.0 1e+03 b (x) float64 40B [g] -1e+03 -750.0 -500.0 -250.0 0.0 Convert homogeneous data >>> ds = xr.Dataset( ... data_vars={ ... "a": ("x", np.linspace(0, 1, 5) * ureg.kg), ... "b": ("x", np.linspace(-1, 0, 5) * ureg.mg), ... }, ... coords={"u": ("x", np.arange(5) * ureg.s)}, ... ) >>> ds Size: 120B Dimensions: (x: 5) Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Data variables: a (x) float64 40B [kg] 0.0 0.25 0.5 0.75 1.0 b (x) float64 40B [mg] -1.0 -0.75 -0.5 -0.25 0.0 >>> ds.pint.to("g") Size: 120B Dimensions: (x: 5) Coordinates: u (x) int64 40B [s] 0 1 2 3 4 Dimensions without coordinates: x Data variables: a (x) float64 40B [g] 0.0 250.0 500.0 750.0 1e+03 b (x) float64 40B [g] -0.001 -0.00075 -0.0005 -0.00025 0.0 """ if isinstance(units, (str, pint.Unit)): unit_kwargs.update( {name: units for name in self.ds.keys() if name not in unit_kwargs} ) units = None elif units is not None and not is_dict_like(units): raise ValueError( "units must be either a string, a pint.Unit object or a dict-like," f" but got {units!r}" ) units = either_dict_or_kwargs(units, unit_kwargs, "to") return conversion.convert_units(self.ds, units) def chunk(self, chunks, name_prefix="xarray-", token=None, lock=False): """unit-aware version of chunk Like :py:meth:`xarray.Dataset.chunk`, but chunking a quantity will change the wrapped type to ``dask``. .. note:: It is recommended to only use this when chunking in-memory arrays. To rechunk please use :py:meth:`xarray.Dataset.chunk`. See Also -------- xarray.Dataset.chunk xarray.DataArray.pint.chunk """ units = conversion.extract_units(self.ds) stripped = conversion.strip_units(self.ds) chunked = stripped.chunk( chunks, name_prefix=name_prefix, token=token, lock=lock ) return conversion.attach_units(chunked, units) def reindex( self, indexers=None, method=None, tolerance=None, copy=True, fill_value=NA, **indexers_kwargs, ): """unit-aware version of reindex Like :py:meth:`xarray.Dataset.reindex`, except the object's indexes are converted to the units of the indexers first. .. note:: ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to ``Dataset.reindex`` unmodified. See Also -------- xarray.DataArray.pint.reindex xarray.Dataset.pint.reindex_like xarray.Dataset.reindex """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "reindex") dims = self.ds.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance # TODO: handle fill_value # convert the indexes to the indexer's units converted = conversion.convert_units(self.ds, indexer_units) converted_units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) # index stripped_indexers = conversion.strip_indexer_units(indexers) indexed = stripped.reindex( stripped_indexers, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) return conversion.attach_units(indexed, converted_units) def reindex_like( self, other, method=None, tolerance=None, copy=True, fill_value=NA ): """unit-aware version of reindex_like Like :py:meth:`xarray.Dataset.reindex_like`, except the object's indexes are converted to the units of the indexers first. .. note:: ``tolerance`` and ``fill_value`` are not supported, yet. They will be passed through to ``Dataset.reindex_like`` unmodified. See Also -------- xarray.DataArray.pint.reindex_like xarray.Dataset.pint.reindex xarray.Dataset.reindex_like """ indexer_units = conversion.extract_units(other) converted = conversion.convert_units(self.ds, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) stripped_other = conversion.strip_units(other) # TODO: handle tolerance # TODO: handle fill_value reindexed = stripped.reindex_like( stripped_other, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) return conversion.attach_units(reindexed, units) def interp( self, coords=None, method="linear", assume_sorted=False, kwargs=None, **coords_kwargs, ): """unit-aware version of interp Like :py:meth:`xarray.Dataset.interp`, except the object's indexes are converted to the units of the indexers first. .. note:: ``method``, ``assume_sorted`` and ``kwargs`` are passed unmodified to ``DataArray.interp``. See Also -------- xarray.DataArray.pint.interp xarray.Dataset.pint.interp_like xarray.Dataset.interp """ indexers = either_dict_or_kwargs(coords, coords_kwargs, "interp") dims = self.ds.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # convert the indexes to the indexer's units converted = conversion.convert_units(self.ds, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) # index stripped_indexers = conversion.strip_indexer_units(indexers) interpolated = stripped.interp( stripped_indexers, method=method, assume_sorted=assume_sorted, kwargs=kwargs, ) return conversion.attach_units(interpolated, units) def interp_like(self, other, method="linear", assume_sorted=False, kwargs=None): """unit-aware version of interp_like Like :py:meth:`xarray.Dataset.interp_like`, except the object's indexes are converted to the units of the indexers first. .. note:: ``method``, ``assume_sorted`` and ``kwargs`` are passed unmodified to ``DataArray.interp``. See Also -------- xarray.DataArray.pint.interp_like xarray.Dataset.pint.interp xarray.Dataset.interp_like """ indexer_units = conversion.extract_units(other) converted = conversion.convert_units(self.ds, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) stripped_other = conversion.strip_units(other) interpolated = stripped.interp_like( stripped_other, method=method, assume_sorted=assume_sorted, kwargs=kwargs, ) return conversion.attach_units(interpolated, units) def sel( self, indexers=None, method=None, tolerance=None, drop=False, **indexers_kwargs ): """unit-aware version of sel Like :py:meth:`xarray.Dataset.sel`, except the object's indexes are converted to the units of the indexers first. .. note:: ``tolerance`` is not supported, yet. It will be passed through to ``Dataset.sel`` unmodified. See Also -------- xarray.DataArray.pint.sel xarray.Dataset.sel xarray.DataArray.sel """ indexers = either_dict_or_kwargs(indexers, indexers_kwargs, "sel") dims = self.ds.dims indexer_units = conversion.extract_indexer_units(indexers) indexer_units = { name: indexer for name, indexer in indexer_units.items() if name in dims } # TODO: handle tolerance # convert the indexes to the indexer's units try: converted = conversion.convert_units(self.ds, indexer_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(indexers) stripped = conversion.strip_units(converted) converted_units = conversion.extract_units(converted) indexed = stripped.sel( stripped_indexers, method=method, tolerance=tolerance, drop=drop, ) return conversion.attach_units(indexed, converted_units) @property def loc(self): """Unit-aware attribute for indexing Only supports ``__getitem__``. .. note:: Position based indexing (e.g. ``ds.loc[1, 2:]``) is not supported, yet See Also -------- xarray.Dataset.loc """ return DatasetLocIndexer(self.ds) def drop_sel(self, labels=None, *, errors="raise", **labels_kwargs): """unit-aware version of drop_sel Just like :py:meth:`xarray.Dataset.drop_sel`, except the indexers are converted to the units of the object's indexes first. See Also -------- xarray.DataArray.pint.drop_sel xarray.Dataset.drop_sel xarray.DataArray.drop_sel """ indexers = either_dict_or_kwargs(labels, labels_kwargs, "drop_sel") dims = self.ds.dims unit_attrs = conversion.extract_unit_attributes(self.ds) index_units = { name: units for name, units in unit_attrs.items() if name in dims } # convert the indexers to the indexes units try: converted_indexers = conversion.convert_indexer_units(indexers, index_units) except ValueError as e: raise KeyError(*e.args) from e # index stripped_indexers = conversion.strip_indexer_units(converted_indexers) indexed = self.ds.drop_sel( stripped_indexers, errors=errors, ) return indexed def ffill(self, dim, limit=None): """unit-aware version of ffill Like :py:meth:`xarray.Dataset.ffill` but without stripping the data units. See Also -------- xarray.Dataset.ffill xarray.Dataset.pint.bfill """ units = conversion.extract_units(self.ds) stripped = conversion.strip_units(self.ds) filled = stripped.ffill(dim=dim, limit=limit) return conversion.attach_units(filled, units) def bfill(self, dim, limit=None): """unit-aware version of bfill Like :py:meth:`xarray.Dataset.bfill` but without stripping the data units. See Also -------- xarray.Dataset.bfill xarray.Dataset.pint.ffill """ units = conversion.extract_units(self.ds) stripped = conversion.strip_units(self.ds) filled = stripped.bfill(dim=dim, limit=limit) return conversion.attach_units(filled, units) def interpolate_na( self, dim=None, method="linear", limit=None, use_coordinate=True, max_gap=None, keep_attrs=None, **kwargs, ): """unit-aware version of interpolate_na Like :py:meth:`xarray.Dataset.interpolate_na` but without stripping the units on data or coordinates. .. note:: ``max_gap`` is not supported, yet, and will be passed through to ``Dataset.interpolate_na`` unmodified. See Also -------- xarray.DataArray.pint.interpolate_na xarray.Dataset.interpolate_na """ units = conversion.extract_units(self.ds) stripped = conversion.strip_units(self.ds) interpolated = stripped.interpolate_na( dim=dim, method=method, limit=limit, use_coordinate=use_coordinate, max_gap=max_gap, keep_attrs=keep_attrs, **kwargs, ) return conversion.attach_units(interpolated, units) pint-xarray-0.6.1/pint_xarray/compat.py000066400000000000000000000007211516032446100201730ustar00rootroot00000000000000import xarray as xr try: from xarray import call_on_dataset except ImportError: def call_on_dataset(func, obj, name, *args, **kwargs): if isinstance(obj, xr.DataArray): ds = obj.to_dataset(name=name) else: ds = obj result = func(ds, *args, **kwargs) if isinstance(obj, xr.DataArray) and isinstance(result, xr.Dataset): result = result.get(name).rename(obj.name) return result pint-xarray-0.6.1/pint_xarray/conversion.py000066400000000000000000000370311516032446100211010ustar00rootroot00000000000000import itertools import re import pint from xarray import Coordinates, DataArray, Dataset, IndexVariable, Variable from pint_xarray.compat import call_on_dataset from pint_xarray.errors import create_exception_group from pint_xarray.index import PintIndex no_unit_values = ("none", None) unit_attribute_name = "units" slice_attributes = ("start", "stop", "step") temporary_name = "" time_units_re = r"\w+" datetime_re = r"\d{4}-\d{2}-\d{2}(?:[ T]\d{2}:\d{2}:\d{2}(?:\.\d+)?)?" datetime_units_re = re.compile(rf"{time_units_re} since {datetime_re}") def is_datetime_unit(unit): return isinstance(unit, str) and datetime_units_re.match(unit) is not None def array_attach_units(data, unit): """attach a unit to the data Parameters ---------- data : array-like The data to attach units to. unit : pint.Unit The desired unit. Returns ------- quantity : pint.Quantity """ if unit in no_unit_values: return data if not isinstance(unit, pint.Unit): raise ValueError(f"cannot use {unit!r} as a unit") if isinstance(data, pint.Quantity): if data.units == unit: return data raise ValueError( f"Cannot attach unit {unit!r} to quantity: data " f"already has units {data.units}" ) registry = unit._REGISTRY return registry.Quantity(data, unit) def array_convert_units(data, unit): """convert the units of an array This is roughly the same as ``data.to(unit)``. Parameters ---------- data : quantity or array-like The data to convert. If it is not a quantity, it is assumed to be dimensionless. unit : str or pint.Unit The unit to convert to. If a string ``data`` has to be a quantity. Returns ------- result : pint.Quantity The converted data """ if unit is None: return data if not isinstance(unit, (str, pint.Unit)): raise ValueError(f"cannot use {unit!r} as a unit") elif isinstance(unit, str) and not isinstance(data, pint.Quantity): raise ValueError(f"cannot convert a non-quantity using {unit!r} as unit") registry = data._REGISTRY if isinstance(unit, str) else unit._REGISTRY if not isinstance(data, pint.Quantity): data = registry.Quantity(data, "dimensionless") return data.to(unit) def array_extract_units(data): """extract the units of an array If ``data`` is not a quantity, the units are ``None`` """ try: return data.units except AttributeError: return None def array_strip_units(data): """strip the units of a quantity""" try: return data.magnitude except AttributeError: return data def attach_units_variable(variable, units): if isinstance(variable, IndexVariable): new_obj = variable.copy() if units is not None: new_obj.attrs[unit_attribute_name] = units elif isinstance(variable, Variable): new_data = array_attach_units(variable.data, units) new_obj = variable.copy(data=new_data) else: raise ValueError(f"invalid type: {variable!r}") return new_obj def dataset_from_variables(variables, coordinate_names, indexes, attrs): data_vars = { name: var for name, var in variables.items() if name not in coordinate_names } coords = {name: var for name, var in variables.items() if name in coordinate_names} new_coords = Coordinates(coords, indexes=indexes) return Dataset(data_vars=data_vars, coords=new_coords, attrs=attrs) def attach_units_index(index, index_vars, units): if all(unit is None for unit in units.values()): # skip non-quantity indexed variables return index if isinstance(index, PintIndex): if index.units != units: raise ValueError( f"cannot attach units to quantified index: {index.units} != {units}" ) else: return index return PintIndex(index=index, units=units) def attach_units_dataset(obj, units): attached = {} rejected_vars = {} indexed_variables = obj.xindexes.variables for name, var in obj.variables.items(): if name in indexed_variables: continue unit = units.get(name) try: converted = attach_units_variable(var, unit) attached[name] = converted except ValueError as e: rejected_vars[name] = (unit, e) indexes, index_vars = obj.xindexes.copy_indexes() for idx, idx_vars in obj.xindexes.group_by_index(): idx_units = {name: units.get(name) for name in idx_vars.keys()} try: attached_idx = attach_units_index(idx, idx_vars, idx_units) indexes.update({k: attached_idx for k in idx_vars}) index_vars.update(attached_idx.create_variables(idx_vars)) except ValueError as e: rejected_vars[name] = (units, e) attached.update(index_vars) if rejected_vars: raise ValueError(rejected_vars) reordered = {name: attached[name] for name in obj.variables.keys()} return dataset_from_variables(reordered, obj._coord_names, indexes, obj.attrs) def attach_units(obj, units): if not isinstance(obj, (DataArray, Dataset)): raise ValueError(f"cannot attach units to {obj!r}: unknown type") if isinstance(obj, DataArray): units = units.copy() if obj.name in units: units[temporary_name] = units.get(obj.name) try: new_obj = call_on_dataset( attach_units_dataset, obj, name=temporary_name, units=units ) except ValueError as e: (rejected_vars,) = e.args if temporary_name in rejected_vars: rejected_vars[obj.name] = rejected_vars.pop(temporary_name) raise create_exception_group(rejected_vars, "attach") from None return new_obj def attach_unit_attributes(obj, units, attr="units"): new_obj = obj.copy() if isinstance(obj, DataArray): for name, var in itertools.chain([(obj.name, new_obj)], new_obj.coords.items()): unit = units.get(name) if unit is None: continue var.attrs[attr] = unit elif isinstance(obj, Dataset): for name, var in new_obj.variables.items(): unit = units.get(name) if unit is None: continue var.attrs[attr] = unit else: raise ValueError(f"cannot attach unit attributes to {obj!r}: unknown type") return new_obj def convert_units_variable(variable, units): if isinstance(variable, IndexVariable): if variable.level_names: # don't try to convert MultiIndexes return variable if units is not None: quantity = array_attach_units( variable.data, variable.attrs.get(unit_attribute_name) ) converted = array_convert_units(quantity, units) new_obj = variable.copy(data=array_strip_units(converted)) new_obj.attrs[unit_attribute_name] = array_extract_units(converted) else: new_obj = variable elif isinstance(variable, Variable): converted = array_convert_units(variable.data, units) new_obj = variable.copy(data=converted) else: raise ValueError(f"unknown type: {variable}") return new_obj def convert_units_index(index, index_vars, units): if not isinstance(index, PintIndex): raise ValueError("cannot convert non-quantified index") converted_vars = {} failed = {} for name, var in index_vars.items(): unit = units.get(name) try: converted = convert_units_variable(var, unit) converted_vars[name] = strip_units_variable(converted) except (ValueError, pint.errors.PintTypeError) as e: failed[name] = e if failed: # raise exception group raise ValueError("failed to convert index variables:", failed) # TODO: figure out how to pull out `options` converted_index = index.index.from_variables(converted_vars, options={}) return PintIndex(index=converted_index, units=units) def convert_units_dataset(obj, units): converted = {} failed = {} indexed_variables = obj.xindexes.variables for name, var in obj.variables.items(): if name in indexed_variables: continue unit = units.get(name) try: converted[name] = convert_units_variable(var, unit) except (ValueError, pint.errors.PintTypeError) as e: failed[name] = e indexes, index_vars = obj.xindexes.copy_indexes() for idx, idx_vars in obj.xindexes.group_by_index(): idx_units = {name: units.get(name) for name in idx_vars.keys()} if all(unit is None for unit in idx_units.values()): continue try: converted_index = convert_units_index(idx, idx_vars, idx_units) indexes.update({k: converted_index for k in idx_vars}) index_vars.update(converted_index.create_variables(idx_vars)) except (ValueError, pint.errors.PintTypeError) as e: names = tuple(idx_vars) failed[names] = e converted.update(index_vars) if failed: raise ValueError(failed) reordered = {name: converted[name] for name in obj.variables.keys()} return dataset_from_variables(reordered, obj._coord_names, indexes, obj.attrs) def convert_units(obj, units): if not isinstance(obj, (DataArray, Dataset)): raise ValueError(f"cannot convert object: {obj!r}: unknown type") if isinstance(obj, DataArray): units = units.copy() if obj.name in units: units[temporary_name] = units.pop(obj.name) try: new_obj = call_on_dataset( convert_units_dataset, obj, name=temporary_name, units=units ) except ValueError as e: (failed,) = e.args if temporary_name in failed: failed[obj.name] = failed.pop(temporary_name) raise create_exception_group(failed, "convert") from None return new_obj def extract_units_dataset(obj): return {name: array_extract_units(var.data) for name, var in obj.variables.items()} def extract_units(obj): if not isinstance(obj, (DataArray, Dataset)): raise ValueError(f"unknown type: {type(obj)}") unit_attributes = extract_unit_attributes(obj) units = call_on_dataset(extract_units_dataset, obj, name=temporary_name) if temporary_name in units: units[obj.name] = units.pop(temporary_name) units_ = unit_attributes.copy() units_.update({k: v for k, v in units.items() if v is not None}) return units_ def extract_unit_attributes_dataset(obj, attr="units"): all_units = {name: var.attrs.get(attr, None) for name, var in obj.variables.items()} return { name: unit for name, unit in all_units.items() if not is_datetime_unit(unit) } def extract_unit_attributes(obj, attr="units"): if not isinstance(obj, (DataArray, Dataset)): raise ValueError( f"cannot retrieve unit attributes from unknown type: {type(obj)}" ) units = call_on_dataset( extract_unit_attributes_dataset, obj, name=temporary_name, attr=attr ) if temporary_name in units: units[obj.name] = units.pop(temporary_name) return units def strip_units_variable(var): if not isinstance(var.data, pint.Quantity): return var data = array_strip_units(var.data) return var.copy(data=data) def strip_units_dataset(obj): variables = {name: strip_units_variable(var) for name, var in obj.variables.items()} indexes = { name: (index.index if isinstance(index, PintIndex) else index) for name, index in obj.xindexes.items() } return dataset_from_variables(variables, obj._coord_names, indexes, obj.attrs) def strip_units(obj): if not isinstance(obj, (DataArray, Dataset)): raise ValueError("cannot strip units from {obj!r}: unknown type") return call_on_dataset(strip_units_dataset, obj, name=temporary_name) def strip_unit_attributes_dataset(obj, attr="units"): new_obj = obj.copy() for var in new_obj.variables.values(): if is_datetime_unit(var.attrs.get(attr, "")): continue var.attrs.pop(attr, None) return new_obj def strip_unit_attributes(obj, attr="units"): if not isinstance(obj, (DataArray, Dataset)): raise ValueError(f"cannot strip unit attributes from unknown type: {type(obj)}") return call_on_dataset( strip_unit_attributes_dataset, obj, name=temporary_name, attr=attr ) def slice_extract_units(indexer): elements = {name: getattr(indexer, name) for name in slice_attributes} extracted_units = [ array_extract_units(value) for name, value in elements.items() if value is not None ] none_values = [_ is None for _ in extracted_units] if not extracted_units or all(none_values): # empty slice (slice(None)) or slice without units return None dimensionalities = { str(getattr(units, "dimensionality", "dimensionless")) for units in extracted_units } if len(dimensionalities) > 1: raise ValueError(f"incompatible units in {indexer}: {dimensionalities}") units = [_ for _ in extracted_units if _ is not None] if len(set(units)) == 1: return units[0] else: units_ = units[0] registry = units_._REGISTRY return registry.Quantity(1, units_).to_base_units().units def convert_units_slice(indexer, units): attrs = {name: getattr(indexer, name) for name in slice_attributes} converted = { name: array_convert_units(value, units) if value is not None else None for name, value in attrs.items() } args = [converted[name] for name in slice_attributes] return slice(*args) def convert_indexer_units(indexers, units): def convert(indexer, units): if isinstance(indexer, slice): return convert_units_slice(indexer, units) elif isinstance(indexer, DataArray): return convert_units(indexer, {None: units}) elif isinstance(indexer, Variable): return convert_units_variable(indexer, units) else: return array_convert_units(indexer, units) converted = {} invalid = {} for name, indexer in indexers.items(): indexer_units = units.get(name) try: converted[name] = convert(indexer, indexer_units) except (ValueError, pint.errors.PintTypeError) as e: invalid[name] = e if invalid: raise create_exception_group(invalid, "convert_indexers") return converted def extract_indexer_units(indexers): def extract(indexer): if isinstance(indexer, slice): return slice_extract_units(indexer) elif isinstance(indexer, (DataArray, Variable)): return array_extract_units(indexer.data) else: return array_extract_units(indexer) return {name: extract(indexer) for name, indexer in indexers.items()} def strip_indexer_units(indexers): def strip(indexer): if isinstance(indexer, slice): return slice( array_strip_units(indexer.start), array_strip_units(indexer.stop), array_strip_units(indexer.step), ) elif isinstance(indexer, DataArray): return strip_units(indexer) elif isinstance(indexer, Variable): return strip_units_variable(indexer) else: return array_strip_units(indexer) return {name: strip(indexer) for name, indexer in indexers.items()} pint-xarray-0.6.1/pint_xarray/errors.py000066400000000000000000000027751516032446100202370ustar00rootroot00000000000000from collections.abc import Hashable from typing import Any class PintExceptionGroup(ExceptionGroup, ValueError): """Exception group for errors related to unit operations Raised whenever there's the possibility of multiple errors. """ pass def _add_note(e: Exception, note: str) -> Exception: e.add_note(note) return e def create_exception_group(mapping: dict[Hashable, Any], op: str) -> ExceptionGroup: match op: case "attach": message = "Cannot attach units" errors = [ _add_note(e, f"cannot attach units to variable {key!r}: {unit}") for key, (unit, e) in mapping.items() ] case "parse": message = "Cannot parse units" errors = [ _add_note(e, f"invalid units for variable {key!r}: {unit} ({type})") for key, (unit, type, e) in mapping.items() ] case "convert": message = "Cannot convert variables" errors = [ _add_note(e, f"incompatible units for variable {key!r}") for key, e in mapping.items() ] case "convert_indexers": message = "Cannot convert indexers" errors = [ _add_note(e, f"incompatible units for indexer for {key!r}") for key, e in mapping.items() ] case _: # pragma: no cover raise ValueError("invalid op") return PintExceptionGroup(message, errors) pint-xarray-0.6.1/pint_xarray/formatting.py000066400000000000000000000162451516032446100210720ustar00rootroot00000000000000from itertools import zip_longest import numpy as np from xarray.core.options import OPTIONS # vendored from xarray.core.formatting def maybe_truncate(obj, maxlen=500): s = str(obj) if len(s) > maxlen: s = s[: (maxlen - 3)] + "..." return s # vendored from xarray.core.formatting def pretty_print(x, numchars: int): """Given an object `x`, call `str(x)` and format the returned string so that it is numchars long, padding with trailing spaces or truncating with ellipses as necessary """ s = maybe_truncate(x, numchars) return s + " " * max(numchars - len(s), 0) # vendored from xarray.core.formatting def _get_indexer_at_least_n_items(shape, n_desired, from_end): assert 0 < n_desired <= np.prod(shape) cum_items = np.cumprod(shape[::-1]) n_steps = np.argmax(cum_items >= n_desired) stop = int(np.ceil(float(n_desired) / np.r_[1, cum_items][n_steps])) indexer = ( ((-1 if from_end else 0),) * (len(shape) - 1 - n_steps) + ((slice(-stop, None) if from_end else slice(stop)),) + (slice(None),) * n_steps ) return indexer # vendored from xarray.core.formatting def first_n_items(array, n_desired): """Returns the first n_desired items of an array""" # Unfortunately, we can't just do array.flat[:n_desired] here because it # might not be a numpy.ndarray. Moreover, access to elements of the array # could be very expensive (e.g. if it's only available over DAP), so go out # of our way to get them in a single call to __getitem__ using only slices. if n_desired < 1: raise ValueError("must request at least one item") if array.size == 0: # work around for https://github.com/numpy/numpy/issues/5195 return [] if n_desired < array.size: indexer = _get_indexer_at_least_n_items(array.shape, n_desired, from_end=False) array = array[indexer] return np.asarray(array).flat[:n_desired] # vendored from xarray.core.formatting def last_n_items(array, n_desired): """Returns the last n_desired items of an array""" # Unfortunately, we can't just do array.flat[-n_desired:] here because it # might not be a numpy.ndarray. Moreover, access to elements of the array # could be very expensive (e.g. if it's only available over DAP), so go out # of our way to get them in a single call to __getitem__ using only slices. if (n_desired == 0) or (array.size == 0): return [] if n_desired < array.size: indexer = _get_indexer_at_least_n_items(array.shape, n_desired, from_end=True) array = array[indexer] return np.asarray(array).flat[-n_desired:] # based on xarray.core.formatting.format_item def format_item(x, quote_strings=True): """Returns a succinct summary of an object as a string""" if isinstance(x, (str, bytes)): return repr(x) if quote_strings else x elif isinstance(x, float): return f"{x:.4}" elif hasattr(x, "dtype") and np.issubdtype(x.dtype, np.floating): return f"{x.item():.4}" else: return str(x) # based on xarray.core.formatting.format_item def format_items(x): """Returns a succinct summaries of all items in a sequence as strings""" x = np.asarray(x) formatted = [format_item(xi) for xi in x] return formatted def summarize_attr(key, value, col_width=None): """Summary for __repr__ - use ``X.attrs[key]`` for full value.""" # Indent key and add ':', then right-pad if col_width is not None k_str = f" {key}:" if col_width is not None: k_str = pretty_print(k_str, col_width) # Replace tabs and newlines, so we print on one line in known width v_str = str(value).replace("\t", "\\t").replace("\n", "\\n") # Finally, truncate to the desired display width return maybe_truncate(f"{k_str} {v_str}", OPTIONS["display_width"]) # adapted from xarray.core.formatting def _diff_mapping_repr(a_mapping, b_mapping, title, summarizer, col_width=None): def extra_items_repr(extra_keys, mapping, ab_side): extra_repr = [summarizer(k, mapping[k], col_width) for k in extra_keys] if extra_repr: header = f"{title} only on the {ab_side} object:" return [header] + extra_repr else: return [] a_keys = set(a_mapping) b_keys = set(b_mapping) summary = [] diff_items = [] for k in a_keys & b_keys: compatible = a_mapping[k] == b_mapping[k] if not compatible: temp = [ summarizer(k, vars[k], col_width) for vars in (a_mapping, b_mapping) ] diff_items += [ab_side + s[1:] for ab_side, s in zip(("L", "R"), temp)] if diff_items: summary += [f"Differing {title.lower()}:"] + diff_items summary += extra_items_repr(a_keys - b_keys, a_mapping, "left") summary += extra_items_repr(b_keys - a_keys, b_mapping, "right") return "\n".join(summary) # vendored from xarray.core.formatting def format_array_flat(array, max_width: int): """Return a formatted string for as many items in the flattened version of array that will fit within max_width characters. """ # every item will take up at least two characters, but we always want to # print at least first and last items max_possibly_relevant = min( max(array.size, 1), max(int(np.ceil(max_width / 2.0)), 2) ) relevant_front_items = format_items( first_n_items(array, (max_possibly_relevant + 1) // 2) ) relevant_back_items = format_items(last_n_items(array, max_possibly_relevant // 2)) # interleave relevant front and back items: # [a, b, c] and [y, z] -> [a, z, b, y, c] relevant_items = sum( zip_longest(relevant_front_items, reversed(relevant_back_items)), () )[:max_possibly_relevant] cum_len = np.cumsum([len(s) + 1 for s in relevant_items]) - 1 if (array.size > 2) and ( (max_possibly_relevant < array.size) or (cum_len > max_width).any() ): padding = " ... " count = min( array.size, max(np.argmax(cum_len + len(padding) - 1 > max_width), 2) ) else: count = array.size padding = "" if (count <= 1) else " " num_front = (count + 1) // 2 num_back = count - num_front # note that num_back is 0 <--> array.size is 0 or 1 # <--> relevant_back_items is [] pprint_str = "".join( [ " ".join(relevant_front_items[:num_front]), padding, " ".join(relevant_back_items[-num_back:]), ] ) # As a final check, if it's still too long even with the limit in values, # replace the end with an ellipsis # NB: this will still returns a full 3-character ellipsis when max_width < 3 if len(pprint_str) > max_width: pprint_str = pprint_str[: max(max_width - 3, 0)] + "..." return pprint_str def inline_repr(quantity, max_width): magnitude = quantity.magnitude units = quantity.units units_repr = f"{units:~P}" if isinstance(magnitude, np.ndarray): data_repr = format_array_flat(magnitude, max_width - len(units_repr) - 3) else: data_repr = maybe_truncate(repr(magnitude), max_width - len(units_repr) - 3) return f"[{units_repr}] {data_repr}" pint-xarray-0.6.1/pint_xarray/index.py000066400000000000000000000103701516032446100200200ustar00rootroot00000000000000import inspect from xarray import Variable from xarray.core.indexes import Index, PandasIndex from pint_xarray import conversion class PintIndex(Index): def __init__(self, *, index, units): """create a unit-aware MetaIndex Parameters ---------- index : xarray.Index The wrapped index object. units : mapping of hashable to unit-like The units of the indexed coordinates """ if not isinstance(units, dict): raise TypeError( "Index units have to be a dict of coordinate names to units." ) self.index = index self.units = units def _replace(self, new_index): return self.__class__(index=new_index, units=self.units) def create_variables(self, variables=None): index_vars = self.index.create_variables(variables) index_vars_units = {} for name, var in index_vars.items(): data = conversion.array_attach_units(var.data, self.units[name]) var_units = Variable(var.dims, data, attrs=var.attrs, encoding=var.encoding) index_vars_units[name] = var_units return index_vars_units @classmethod def from_variables(cls, variables, options): if len(variables) != 1: raise ValueError("can only create a default index from single variables") units = options.pop("units", None) index = PandasIndex.from_variables(variables, options=options) return cls(index=index, units={index.index.name: units}) @classmethod def concat(cls, indexes, dim, positions): raise NotImplementedError() @classmethod def stack(cls, variables, dim): raise NotImplementedError() def unstack(self): raise NotImplementedError() def sel(self, labels, **options): converted_labels = conversion.convert_indexer_units(labels, self.units) stripped_labels = conversion.strip_indexer_units(converted_labels) return self.index.sel(stripped_labels, **options) def isel(self, indexers): subset = self.index.isel(indexers) if subset is None: return None return self._replace(subset) def join(self, other, how="inner"): raise NotImplementedError() def reindex_like(self, other): raise NotImplementedError() def equals(self, other, *, exclude=None): if not isinstance(other, PintIndex): return False # for now we require exactly matching units to avoid the potentially # expensive conversion if self.units != other.units: return False # TODO: # - remove try-except once we can drop xarray<2025.06.0 # - remove compat once we can require a version of xarray that completed # the deprecation cycle try: from xarray.core.indexes import _wrap_index_equals equals = _wrap_index_equals(self.index) kwargs = {"exclude": exclude} except ImportError: # pragma: no cover equals = self.index.equals signature = inspect.signature(self.index.equals) if "exclude" in signature.parameters: kwargs = {"exclude": exclude} else: kwargs = {} # Last to avoid the potentially expensive comparison return equals(other.index, **kwargs) def roll(self, shifts): return self._replace(self.index.roll(shifts)) def rename(self, name_dict, dims_dict): new_units = {new: self.units[old] for old, new in name_dict.items()} return self.__class__( index=self.index.rename(name_dict, dims_dict), units=new_units ) def __getitem__(self, indexer): return self._replace(self.index[indexer]) def _repr_inline_(self, max_width): name = self.__class__.__name__ wrapped_name = self.index.__class__.__name__ formatted_units = {n: f"{u:~P}" for n, u in self.units.items()} return f"{name}({wrapped_name}, units={formatted_units})" def __repr__(self): formatted_units = {n: f"{u:~P}" for n, u in self.units.items()} summary = f"<{self.__class__.__name__} (units={formatted_units})>" return "\n".join([summary, repr(self.index)]) pint-xarray-0.6.1/pint_xarray/itertools.py000066400000000000000000000014161516032446100207360ustar00rootroot00000000000000import itertools from functools import reduce def separate(predicate, iterable): evaluated = ((predicate(el), el) for el in iterable) key = lambda x: x[0] grouped = itertools.groupby(sorted(evaluated, key=key), key=key) groups = {label: [el for _, el in group] for label, group in grouped} return groups[False], groups[True] def unique(iterable): return list(dict.fromkeys(iterable)) def zip_mappings(*mappings): def common_keys(a, b): all_keys = unique(itertools.chain(a.keys(), b.keys())) intersection = set(a.keys()).intersection(b.keys()) return [key for key in all_keys if key in intersection] keys = list(reduce(common_keys, mappings)) for key in keys: yield key, tuple(m[key] for m in mappings) pint-xarray-0.6.1/pint_xarray/testing.py000066400000000000000000000016271516032446100203730ustar00rootroot00000000000000from pint_xarray import conversion, formatting def assert_units_equal(a, b): """assert that the units of two xarray objects are equal Raises an :py:exc:`AssertionError` if the units of both objects are not equal. ``units`` attributes and attached unit objects are compared separately. Parameters ---------- a, b : DataArray or Dataset The objects to compare """ __tracebackhide__ = True units_a = conversion.extract_units(a) units_b = conversion.extract_units(b) assert units_a == units_b, formatting._diff_mapping_repr( units_a, units_b, "Units", formatting.summarize_attr ) unit_attrs_a = conversion.extract_unit_attributes(a) unit_attrs_b = conversion.extract_unit_attributes(b) assert unit_attrs_a == unit_attrs_b, formatting._diff_mapping_repr( unit_attrs_a, unit_attrs_b, "Unit attrs", formatting.summarize_attr ) pint-xarray-0.6.1/pint_xarray/tests/000077500000000000000000000000001516032446100175005ustar00rootroot00000000000000pint-xarray-0.6.1/pint_xarray/tests/__init__.py000066400000000000000000000000001516032446100215770ustar00rootroot00000000000000pint-xarray-0.6.1/pint_xarray/tests/test_accessors.py000066400000000000000000002141741516032446100231070ustar00rootroot00000000000000import numpy as np import pandas as pd import pint import pytest import xarray as xr from numpy.testing import assert_array_equal from pint import Unit, UnitRegistry from pint_xarray import accessors, conversion from pint_xarray.errors import PintExceptionGroup from pint_xarray.index import PintIndex from pint_xarray.tests.utils import ( assert_equal, assert_identical, assert_units_equal, requires_bottleneck, requires_dask_array, requires_scipy, ) pytestmark = [ pytest.mark.filterwarnings("error::pint.UnitStrippedWarning"), ] # make sure scalars are converted to 0d arrays so quantities can # always be treated like ndarrays from pint_xarray import unit_registry Quantity = unit_registry.Quantity nan = np.nan def assert_all_str_or_none(mapping): __tracebackhide__ = True compared = { key: isinstance(value, str) or value is None for key, value in mapping.items() } not_passing = {key: value for key, value in mapping.items() if not compared[key]} check = all(compared.values()) assert check, f"Not all values are str or None: {not_passing}" @pytest.fixture def example_unitless_da(): array = np.linspace(0, 10, 20) x = np.arange(20) u = np.linspace(0, 1, 20) da = xr.DataArray( data=array, dims="x", coords={"x": ("x", x), "u": ("x", u, {"units": "hour"})}, attrs={"units": "m"}, ) return da @pytest.fixture() def example_quantity_da(): array = np.linspace(0, 10, 20) * unit_registry.m x = np.arange(20) u = np.linspace(0, 1, 20) * unit_registry.hour return xr.DataArray(data=array, dims="x", coords={"x": ("x", x), "u": ("x", u)}) class TestQuantifyDataArray: def test_attach_units_from_str(self, example_unitless_da): orig = example_unitless_da result = orig.pint.quantify("s") assert_array_equal(result.data.magnitude, orig.data) # TODO better comparisons for when you can't access the unit_registry? assert str(result.data.units) == "second" def test_attach_units_given_registry(self, example_unitless_da): orig = example_unitless_da ureg = UnitRegistry(force_ndarray=True) result = orig.pint.quantify("m", unit_registry=ureg) assert_array_equal(result.data.magnitude, orig.data) assert result.data.units == ureg.Unit("m") def test_attach_units_from_attrs(self, example_unitless_da): orig = example_unitless_da result = orig.pint.quantify() assert_array_equal(result.data.magnitude, orig.data) assert str(result.data.units) == "meter" remaining_attrs = conversion.extract_unit_attributes(result) assert {k: v for k, v in remaining_attrs.items() if v is not None} == {} def test_attach_units_from_str_attr_no_unit(self, example_unitless_da): orig = example_unitless_da orig.attrs["units"] = "none" result = orig.pint.quantify("m") assert_array_equal(result.data.magnitude, orig.data) assert str(result.data.units) == "meter" def test_attach_units_given_unit_objs(self, example_unitless_da): orig = example_unitless_da ureg = UnitRegistry(force_ndarray=True) result = orig.pint.quantify(ureg.Unit("m"), unit_registry=ureg) assert_array_equal(result.data.magnitude, orig.data) assert result.data.units == ureg.Unit("m") @pytest.mark.parametrize("no_unit_value", conversion.no_unit_values) def test_override_units(self, example_unitless_da, no_unit_value): orig = example_unitless_da result = orig.pint.quantify(no_unit_value, u=no_unit_value) with pytest.raises(AttributeError): result.data.units with pytest.raises(AttributeError): result["u"].data.units def test_error_when_changing_units(self, example_quantity_da): da = example_quantity_da with pytest.RaisesGroup( pytest.RaisesExc(ValueError, match="already has units"), match="Cannot attach units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): da.pint.quantify("s") def test_attach_no_units(self): arr = xr.DataArray([1, 2, 3], dims="x") quantified = arr.pint.quantify() assert_identical(quantified, arr) assert_units_equal(quantified, arr) def test_attach_no_new_units(self): da = xr.DataArray(unit_registry.Quantity([1, 2, 3], "m"), dims="x") quantified = da.pint.quantify() assert_identical(quantified, da) assert_units_equal(quantified, da) def test_attach_same_units(self): da = xr.DataArray(unit_registry.Quantity([1, 2, 3], "m"), dims="x") quantified = da.pint.quantify("m") assert_identical(quantified, da) assert_units_equal(quantified, da) def test_error_when_changing_units_dimension_coordinates(self): arr = xr.DataArray( [1, 2, 3], dims="x", coords={"x": ("x", [-1, 0, 1], {"units": unit_registry.Unit("m")})}, ) with pytest.RaisesGroup( pytest.RaisesExc(ValueError, match="already has units"), match="Cannot attach units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): arr.pint.quantify({"x": "s"}) def test_dimension_coordinate_array(self): ds = xr.Dataset(coords={"x": ("x", [10], {"units": "m"})}) arr = ds.x # does not actually quantify because `arr` wraps a IndexVariable # but we still get a `Unit` in the attrs q = arr.pint.quantify() assert isinstance(q.attrs["units"], Unit) def test_dimension_coordinate_array_already_quantified(self): ds = xr.Dataset(coords={"x": ("x", [10], {"units": unit_registry.Unit("m")})}) arr = ds.x with pytest.RaisesGroup( pytest.RaisesExc(ValueError, match="already has units"), match="Cannot attach units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): arr.pint.quantify({"x": "s"}) def test_dimension_coordinate_array_already_quantified_same_units(self): x = unit_registry.Quantity([10], "m") coords = xr.Coordinates( {"x": x}, indexes={ "x": PintIndex.from_variables( {"x": xr.Variable("x", x.magnitude)}, options={"units": x.units}, ), }, ) ds = xr.Dataset(coords=coords) arr = ds.x quantified = arr.pint.quantify({"x": "m"}) assert_identical(quantified, arr) assert_units_equal(quantified, arr) def test_error_on_nonsense_units(self, example_unitless_da): da = example_unitless_da with pytest.RaisesGroup( pytest.RaisesExc( pint.UndefinedUnitError, match=rf"{da.name}: .+ \(parameter\)" ), match="Cannot parse units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): da.pint.quantify(units="aecjhbav") def test_error_on_nonsense_units_attrs(self, example_unitless_da): da = example_unitless_da da.attrs["units"] = "aecjhbav" with pytest.RaisesGroup( pytest.RaisesExc( pint.UndefinedUnitError, match=rf"{da.name}: .+ \(attribute\)" ), match="Cannot parse units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): da.pint.quantify() def test_parse_integer_inverse(self): # Regression test for issue #40 da = xr.DataArray([10], attrs={"units": "m^-1"}) result = da.pint.quantify() assert result.pint.units == Unit("1 / meter") @pytest.mark.parametrize("formatter", ("", "P", "C")) @pytest.mark.parametrize("modifier", ("", "~")) def test_units_to_str_or_none(formatter, modifier): unit_format = f"{{:{modifier}{formatter}}}" unit_attrs = {None: "m", "a": "s", "b": "degC", "c": "degF", "d": "degK"} units = {key: unit_registry.Unit(value) for key, value in unit_attrs.items()} expected = {key: unit_format.format(value) for key, value in units.items()} actual = accessors.units_to_str_or_none(units, unit_format) assert expected == actual assert units == {key: unit_registry.Unit(value) for key, value in actual.items()} expected = {None: None} assert expected == accessors.units_to_str_or_none(expected, unit_format) class TestDequantifyDataArray: def test_strip_units(self, example_quantity_da): result = example_quantity_da.pint.dequantify() assert isinstance(result.data, np.ndarray) assert isinstance(result.coords["x"].data, np.ndarray) def test_attrs_reinstated(self, example_quantity_da): da = example_quantity_da result = da.pint.dequantify() units = conversion.extract_units(da) attrs = conversion.extract_unit_attributes(result) assert units == attrs assert_all_str_or_none(attrs) def test_roundtrip_data(self, example_unitless_da): orig = example_unitless_da quantified = orig.pint.quantify() result = quantified.pint.dequantify() assert_equal(result, orig) def test_multiindex(self): mindex = pd.MultiIndex.from_product([["a", "b"], [1, 2]], names=("lat", "lon")) da = xr.DataArray( np.arange(len(mindex)), dims="multi", coords={"multi": mindex} ) result = da.pint.dequantify() xr.testing.assert_identical(da, result) assert isinstance(result.indexes["multi"], pd.MultiIndex) class TestPropertiesDataArray: def test_magnitude_getattr(self, example_quantity_da): da = example_quantity_da actual = da.pint.magnitude assert not isinstance(actual, Quantity) def test_magnitude_getattr_unitless(self, example_unitless_da): da = example_unitless_da xr.testing.assert_duckarray_equal(da.pint.magnitude, da.data) def test_units_getattr(self, example_quantity_da): da = example_quantity_da actual = da.pint.units assert isinstance(actual, Unit) assert actual == unit_registry.m def test_units_setattr(self, example_quantity_da): da = example_quantity_da with pytest.raises(ValueError): da.pint.units = "s" def test_units_getattr_unitless(self, example_unitless_da): da = example_unitless_da assert da.pint.units is None def test_units_setattr_unitless(self, example_unitless_da): da = example_unitless_da da.pint.units = unit_registry.s assert da.pint.units == unit_registry.s @pytest.fixture() def example_unitless_ds(): users = np.linspace(0, 10, 20) funds = np.logspace(0, 10, 20) t = np.arange(20) ds = xr.Dataset( data_vars={"users": (["t"], users), "funds": (["t"], funds)}, coords={"t": t} ) ds["users"].attrs["units"] = "" ds["funds"].attrs["units"] = "pound" return ds @pytest.fixture() def example_quantity_ds(): users = np.linspace(0, 10, 20) * unit_registry.dimensionless funds = np.logspace(0, 10, 20) * unit_registry.pound t = np.arange(20) ds = xr.Dataset( data_vars={"users": (["t"], users), "funds": (["t"], funds)}, coords={"t": t} ) return ds class TestQuantifyDataSet: def test_attach_units_from_str(self, example_unitless_ds): orig = example_unitless_ds result = orig.pint.quantify() assert_array_equal(result["users"].data.magnitude, orig["users"].data) assert str(result["users"].data.units) == "dimensionless" def test_attach_units_given_registry(self, example_unitless_ds): orig = example_unitless_ds orig["users"].attrs.clear() result = orig.pint.quantify( {"users": "dimensionless"}, unit_registry=unit_registry ) assert_array_equal(result["users"].data.magnitude, orig["users"].data) assert str(result["users"].data.units) == "dimensionless" def test_attach_units_from_attrs(self, example_unitless_ds): orig = example_unitless_ds orig["users"].attrs.clear() result = orig.pint.quantify({"users": "dimensionless"}) assert_array_equal(result["users"].data.magnitude, orig["users"].data) assert str(result["users"].data.units) == "dimensionless" remaining_attrs = conversion.extract_unit_attributes(result) assert {k: v for k, v in remaining_attrs.items() if v is not None} == {} def test_attach_units_given_unit_objs(self, example_unitless_ds): orig = example_unitless_ds orig["users"].attrs.clear() dimensionless = unit_registry.Unit("dimensionless") result = orig.pint.quantify({"users": dimensionless}) assert_array_equal(result["users"].data.magnitude, orig["users"].data) assert str(result["users"].data.units) == "dimensionless" def test_attach_units_from_str_attr_no_unit(self, example_unitless_ds): orig = example_unitless_ds orig["users"].attrs["units"] = "none" result = orig.pint.quantify({"users": "m"}) assert_array_equal(result["users"].data.magnitude, orig["users"].data) assert str(result["users"].data.units) == "meter" @pytest.mark.parametrize("no_unit_value", conversion.no_unit_values) def test_override_units(self, example_unitless_ds, no_unit_value): orig = example_unitless_ds result = orig.pint.quantify({"users": no_unit_value}) assert ( getattr(result["users"].data, "units", "not a quantity") == "not a quantity" ) def test_error_when_already_units(self, example_quantity_ds): with pytest.RaisesGroup( pytest.RaisesExc(ValueError, match="already has units"), match="Cannot attach units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): example_quantity_ds.pint.quantify({"funds": "kg"}) def test_attach_no_units(self): ds = xr.Dataset({"a": ("x", [1, 2, 3])}) quantified = ds.pint.quantify() assert_identical(quantified, ds) assert_units_equal(quantified, ds) def test_attach_no_new_units(self): ds = xr.Dataset({"a": ("x", unit_registry.Quantity([1, 2, 3], "m"))}) quantified = ds.pint.quantify() assert_identical(quantified, ds) assert_units_equal(quantified, ds) def test_attach_same_units(self): ds = xr.Dataset({"a": ("x", unit_registry.Quantity([1, 2, 3], "m"))}) quantified = ds.pint.quantify({"a": "m"}) assert_identical(quantified, ds) assert_units_equal(quantified, ds) def test_error_when_changing_units_dimension_coordinates(self): ds = xr.Dataset( coords={"x": ("x", [-1, 0, 1], {"units": unit_registry.Unit("m")})}, ) with pytest.RaisesGroup( pytest.RaisesExc(ValueError, match="already has units"), match="Cannot attach units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): ds.pint.quantify({"x": "s"}) def test_error_on_nonsense_units(self, example_unitless_ds): ds = example_unitless_ds with pytest.RaisesGroup( pytest.RaisesExc( pint.UndefinedUnitError, match=r"'users': .+ \(parameter\)" ), match="Cannot parse units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): ds.pint.quantify(units={"users": "aecjhbav"}) def test_error_on_nonsense_units_attrs(self, example_unitless_ds): ds = example_unitless_ds ds.users.attrs["units"] = "aecjhbav" with pytest.RaisesGroup( pytest.RaisesExc( pint.UndefinedUnitError, match=r"'users': .+ \(attribute\)" ), match="Cannot parse units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): ds.pint.quantify() def test_error_indicates_problematic_variable(self, example_unitless_ds): ds = example_unitless_ds with pytest.RaisesGroup( pytest.RaisesExc( pint.UndefinedUnitError, match=r"'users': aecjhbav \(parameter\)" ), match="Cannot parse units", check=lambda eg: isinstance(eg, PintExceptionGroup), ): ds.pint.quantify(units={"users": "aecjhbav"}) def test_existing_units(self, example_quantity_ds): ds = example_quantity_ds.copy() ds.t.attrs["units"] = unit_registry.Unit("m") with pytest.raises(ValueError, match="Cannot attach"): ds.pint.quantify({"funds": "kg"}) def test_existing_units_dimension(self, example_quantity_ds): ds = example_quantity_ds.copy() ds.t.attrs["units"] = unit_registry.Unit("m") with pytest.raises(ValueError, match="Cannot attach"): ds.pint.quantify({"t": "s"}) class TestDequantifyDataSet: def test_strip_units(self, example_quantity_ds): result = example_quantity_ds.pint.dequantify() assert all( isinstance(var.data, np.ndarray) for var in result.variables.values() ) def test_attrs_reinstated(self, example_quantity_ds): ds = example_quantity_ds result = ds.pint.dequantify() units = conversion.extract_units(ds) # workaround for Unit("dimensionless") != str(Unit("dimensionless")) units = { key: str(value) if isinstance(value, Unit) else value for key, value in units.items() } attrs = conversion.extract_unit_attributes(result) assert units == attrs assert_all_str_or_none(attrs) def test_roundtrip_data(self, example_unitless_ds): orig = example_unitless_ds quantified = orig.pint.quantify() result = quantified.pint.dequantify() assert_equal(result, orig) result = quantified.pint.dequantify().pint.quantify() assert_equal(quantified, result) @pytest.mark.parametrize( ["obj", "units", "expected", "error"], ( pytest.param( xr.Dataset( {"a": ("x", Quantity([0, 1], "m")), "b": ("x", Quantity([2, 4], "s"))} ), {"a": "mm", "b": "ms"}, xr.Dataset( { "a": ("x", Quantity([0, 1000], "mm")), "b": ("x", Quantity([2000, 4000], "ms")), } ), None, id="Dataset-compatible units-data", ), pytest.param( xr.Dataset( {"a": ("x", Quantity([0, 1], "km")), "b": ("x", Quantity([2, 4], "cm"))} ), "m", xr.Dataset( { "a": ("x", Quantity([0, 1000], "m")), "b": ("x", Quantity([0.02, 0.04], "m")), } ), None, id="Dataset-compatible units-data-str", ), pytest.param( xr.Dataset( {"a": ("x", Quantity([0, 1], "m")), "b": ("x", Quantity([2, 4], "s"))} ), {"a": "ms", "b": "mm"}, None, ValueError, id="Dataset-incompatible units-data", ), pytest.param( xr.Dataset(coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={})), {"x": "ms"}, xr.Dataset( coords=xr.Coordinates({"x": Quantity([2000, 4000], "ms")}, indexes={}) ), None, id="Dataset-compatible units-dims-no index", ), pytest.param( xr.Dataset(coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={})), {"x": "mm"}, None, ValueError, id="Dataset-incompatible units-dims-no index", ), pytest.param( xr.DataArray(Quantity([0, 1], "m"), dims="x"), {None: "mm"}, xr.DataArray(Quantity([0, 1000], "mm"), dims="x"), None, id="DataArray-compatible units-data", ), pytest.param( xr.DataArray(Quantity([0, 1], "m"), dims="x"), "mm", xr.DataArray(Quantity([0, 1000], "mm"), dims="x"), None, id="DataArray-compatible units-data-str", ), pytest.param( xr.DataArray(Quantity([0, 1], "m"), dims="x", name="a"), {"a": "mm"}, xr.DataArray(Quantity([0, 1000], "mm"), dims="x", name="a"), None, id="DataArray-compatible units-data-by name", ), pytest.param( xr.DataArray(Quantity([0, 1], "m"), dims="x"), {None: "ms"}, None, ValueError, id="DataArray-incompatible units-data", ), pytest.param( xr.DataArray( [0, 1], dims="x", coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={}), ), {"x": "ms"}, xr.DataArray( [0, 1], dims="x", coords=xr.Coordinates({"x": Quantity([2000, 4000], "ms")}, indexes={}), ), None, id="DataArray-compatible units-dims-no index", ), pytest.param( xr.DataArray( [0, 1], dims="x", coords=xr.Coordinates({"x": Quantity([2, 4], "s")}, indexes={}), ), {"x": "mm"}, None, ValueError, id="DataArray-incompatible units-dims-no index", ), ), ) def test_to(obj, units, expected, error): if error is not None: with pytest.raises(error): obj.pint.to(units) else: actual = obj.pint.to(units) assert_units_equal(actual, expected) assert_identical(actual, expected) @pytest.mark.parametrize( ["obj", "indexers", "expected", "error"], ( pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.Dataset( { "x": ("x", [10, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60], {"units": unit_registry.Unit("s")}), } ), None, id="Dataset-identical units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, xr.Dataset( { "x": ("x", [1, 3], {"units": unit_registry.Unit("m")}), "y": ("y", [1], {"units": unit_registry.Unit("min")}), } ), None, id="Dataset-compatible units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, None, KeyError, id="Dataset-incompatible units", ), pytest.param( xr.Dataset( { "x": ( "x", [10, 20, 30], {"units": unit_registry.Unit("dm"), "long_name": "length"}, ), "y": ( "y", [60, 120], {"units": unit_registry.Unit("s"), "long_name": "time"}, ), } ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.Dataset( { "x": ( "x", [10, 30], {"units": unit_registry.Unit("dm"), "long_name": "length"}, ), "y": ( "y", [60], {"units": unit_registry.Unit("s"), "long_name": "time"}, ), } ), None, id="Dataset-coords with attrs", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.DataArray( [[0], [4]], dims=("x", "y"), coords={ "x": ("x", [10, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60], {"units": unit_registry.Unit("s")}), }, ), None, id="DataArray-identical units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, xr.DataArray( [[0], [4]], dims=("x", "y"), coords={ "x": ("x", [1, 3], {"units": unit_registry.Unit("m")}), "y": ("y", [1], {"units": unit_registry.Unit("min")}), }, ), None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, KeyError, id="DataArray-incompatible units", ), ), ) def test_sel(obj, indexers, expected, error): obj_ = obj.pint.quantify() if error is not None: with pytest.raises(error): obj_.pint.sel(indexers) else: expected_ = expected.pint.quantify() actual = obj_.pint.sel(indexers) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @pytest.mark.parametrize( ["obj", "indexers", "expected", "error"], ( pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.Dataset( { "x": ("x", [10, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60], {"units": unit_registry.Unit("s")}), } ), None, id="Dataset-identical units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, xr.Dataset( { "x": ("x", [1, 3], {"units": unit_registry.Unit("m")}), "y": ("y", [1], {"units": unit_registry.Unit("min")}), } ), None, id="Dataset-compatible units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, None, KeyError, id="Dataset-incompatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.DataArray( [[0], [4]], dims=("x", "y"), coords={ "x": ("x", [10, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60], {"units": unit_registry.Unit("s")}), }, ), None, id="DataArray-identical units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, xr.DataArray( [[0], [4]], dims=("x", "y"), coords={ "x": ("x", [1, 3], {"units": unit_registry.Unit("m")}), "y": ("y", [1], {"units": unit_registry.Unit("min")}), }, ), None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, KeyError, id="DataArray-incompatible units", ), ), ) def test_loc(obj, indexers, expected, error): obj_ = obj.pint.quantify() if error is not None: with pytest.raises(error): obj_.pint.loc[indexers] else: expected_ = expected.pint.quantify() actual = obj_.pint.loc[indexers] assert_units_equal(actual, expected_) assert_identical(actual, expected_) @pytest.mark.parametrize( ["obj", "indexers", "values", "expected", "error"], ( pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, [[-1], [-2]], xr.DataArray( [[-1, 1], [2, 3], [-2, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), None, id="coords-identical units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, [[-1], [-2]], xr.DataArray( [[-1, 1], [2, 3], [-2, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), None, id="coords-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, [[-1], [-2]], None, KeyError, id="coords-incompatible units", ), pytest.param( xr.DataArray( Quantity([[0, 1], [2, 3], [4, 5]], "m"), dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, Quantity([[-1], [-2]], "m"), xr.DataArray( Quantity([[-1, 1], [2, 3], [-2, 5]], "m"), dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), None, id="data-identical units", ), pytest.param( xr.DataArray( Quantity([[0, 1], [2, 3], [4, 5]], "m"), dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, Quantity([[-1], [-2]], "km"), xr.DataArray( Quantity([[-1000, 1], [2, 3], [-2000, 5]], "m"), dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), None, id="data-compatible units", ), pytest.param( xr.DataArray( Quantity([[0, 1], [2, 3], [4, 5]], "m"), dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, Quantity([[-1], [-2]], "s"), None, pint.DimensionalityError, id="data-incompatible units", ), ), ) def test_loc_setitem(obj, indexers, values, expected, error): if error is not None: with pytest.raises(error): obj.pint.loc[indexers] = values else: obj.pint.loc[indexers] = values assert_units_equal(obj, expected) assert_identical(obj, expected) @pytest.mark.parametrize( ["obj", "indexers", "expected", "error"], ( pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.Dataset( { "x": ("x", [20], {"units": unit_registry.Unit("dm")}), "y": ("y", [120], {"units": unit_registry.Unit("s")}), } ), None, id="Dataset-identical units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, xr.Dataset( { "x": ("x", [20], {"units": unit_registry.Unit("dm")}), "y": ("y", [120], {"units": unit_registry.Unit("s")}), } ), None, id="Dataset-compatible units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, None, KeyError, id="Dataset-incompatible units", ), pytest.param( xr.Dataset( { "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), } ), {"x": Quantity([10, 30], "m"), "y": Quantity([60], "min")}, None, KeyError, id="Dataset-compatible units-not found", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "dm"), "y": Quantity([60], "s")}, xr.DataArray( [[3]], dims=("x", "y"), coords={ "x": ("x", [20], {"units": unit_registry.Unit("dm")}), "y": ("y", [120], {"units": unit_registry.Unit("s")}), }, ), None, id="DataArray-identical units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([1, 3], "m"), "y": Quantity([1], "min")}, xr.DataArray( [[3]], dims=("x", "y"), coords={ "x": ("x", [20], {"units": unit_registry.Unit("dm")}), "y": ("y", [120], {"units": unit_registry.Unit("s")}), }, ), None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, KeyError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={ "x": ("x", [10, 20, 30], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), {"x": Quantity([10, 30], "m"), "y": Quantity([60], "min")}, None, KeyError, id="DataArray-compatible units-not found", ), ), ) def test_drop_sel(obj, indexers, expected, error): if error is not None: with pytest.raises(error): obj.pint.drop_sel(indexers) else: actual = obj.pint.drop_sel(indexers) assert_units_equal(actual, expected) assert_identical(actual, expected) @requires_dask_array @pytest.mark.parametrize( "obj", ( pytest.param( xr.Dataset( {"a": ("x", np.linspace(0, 1, 11))}, coords={"u": ("x", np.arange(11))}, ), id="Dataset-no units", ), pytest.param( xr.Dataset( { "a": ( "x", Quantity(np.linspace(0, 1, 11), "m"), ) }, coords={ "u": ( "x", Quantity(np.arange(11), "m"), ) }, ), id="Dataset-units", ), pytest.param( xr.DataArray( np.linspace(0, 1, 11), coords={ "u": ( "x", np.arange(11), ) }, dims="x", ), id="DataArray-no units", ), pytest.param( xr.DataArray( Quantity(np.linspace(0, 1, 11), "m"), coords={ "u": ( "x", Quantity(np.arange(11), "m"), ) }, dims="x", ), id="DataArray-units", ), ), ) def test_chunk(obj): actual = obj.pint.chunk({"x": 2}) expected = ( obj.pint.dequantify().chunk({"x": 2}).pint.quantify(unit_registry=unit_registry) ) assert_units_equal(actual, expected) assert_identical(actual, expected) @pytest.mark.parametrize( ["obj", "units", "indexers", "expected", "expected_units", "error"], ( pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 120, 240], "s")}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, {"x": Quantity([0, 1, 3, 5], "m"), "y": Quantity([0, 2, 4], "min")}, xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, None, {}, ValueError, id="Dataset-incompatible units", ), pytest.param( xr.Dataset( { "a": (("x", "y"), np.array([[0, 1], [2, 3], [4, 5]])), "x": [10, 20, 30], "y": [60, 120], } ), {"a": "kg"}, { "x": [15, 25], "y": [75, 105], }, xr.Dataset( { "a": (("x", "y"), np.array([[np.nan, np.nan], [np.nan, np.nan]])), "x": [15, 25], "y": [75, 105], } ), {"a": "kg"}, None, id="Dataset-data units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 240], "s")}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([1, 3, 5], "m"), "y": Quantity([0, 2], "min")}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, {}, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( np.array([[0, 1], [2, 3], [4, 5]]), dims=("x", "y"), coords={"x": [10, 20, 30], "y": [60, 120]}, ), {None: "kg"}, {"x": [15, 25], "y": [75, 105]}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={"x": [15, 25], "y": [75, 105]}, ), {None: "kg"}, None, id="DataArray-data units", ), ), ) def test_reindex(obj, units, indexers, expected, expected_units, error): obj_ = obj.pint.quantify(units) if error is not None: with pytest.raises(error): obj.pint.reindex(indexers) else: expected_ = expected.pint.quantify(expected_units) actual = obj_.pint.reindex(indexers) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @pytest.mark.parametrize( ["obj", "units", "other", "other_units", "expected", "expected_units", "error"], ( pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), {"x": "m", "y": "min"}, xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [1, 3]), "y": ("y", [1])}), {"x": "s", "y": "m"}, None, {}, ValueError, id="Dataset-incompatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}), {"x": "dm", "y": "s"}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), pytest.param( xr.Dataset( { "a": (("x", "y"), [[0, 1], [2, 3], [4, 5]]), "x": [10, 20, 30], "y": [60, 120], } ), {"a": "kg"}, xr.Dataset({"x": [15, 25], "y": [75, 105]}), {}, xr.Dataset( { "a": (("x", "y"), [[np.nan, np.nan], [np.nan, np.nan]]), "x": [15, 25], "y": [75, 105], } ), {"a": "kg"}, None, id="Dataset-data units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}), {"x": "m", "y": "min"}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30]), "y": ("y", [60])}), {"x": "s", "y": "m"}, None, {}, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": [10, 20, 30], "y": [60, 120]}, ), {"a": "kg"}, xr.Dataset({"x": [15, 25], "y": [75, 105]}), {}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={"x": [15, 25], "y": [75, 105]}, ), {"a": "kg"}, None, id="DataArray-data units", ), ), ) def test_reindex_like(obj, units, other, other_units, expected, expected_units, error): obj_ = obj.pint.quantify(units) other_ = other.pint.quantify(other_units) if error is not None: with pytest.raises(error): obj_.pint.reindex_like(other_) else: expected_ = expected.pint.quantify(expected_units) actual = obj_.pint.reindex_like(other_) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @requires_scipy @pytest.mark.parametrize( ["obj", "units", "indexers", "expected", "expected_units", "error", "kwargs"], ( pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 120, 240], "s")}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), {"x": "dm", "y": "s"}, None, None, id="Dataset-identical units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, {"x": Quantity([0, 1, 3, 5], "m"), "y": Quantity([0, 2, 4], "min")}, xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), {"x": "m", "y": "min"}, None, None, id="Dataset-compatible units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, {"x": Quantity([1, 3], "s"), "y": Quantity([1], "m")}, None, {}, ValueError, None, id="Dataset-incompatible units", ), pytest.param( xr.Dataset( { "a": (("x", "y"), np.array([[0, 1], [2, 3], [4, 5]])), "x": [10, 20, 30], "y": [60, 120], } ), {"a": "kg"}, { "x": [15, 25], "y": [75, 105], }, xr.Dataset( { "a": (("x", "y"), np.array([[1.25, 1.75], [3.25, 3.75]])), "x": [15, 25], "y": [75, 105], } ), {"a": "kg"}, None, None, id="Dataset-data units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([10, 30, 50], "dm"), "y": Quantity([0, 240], "s")}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), {"x": "dm", "y": "s"}, None, None, id="DataArray-identical units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([1, 3, 5], "m"), "y": Quantity([0, 2], "min")}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), {"x": "m", "y": "min"}, None, None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([10, 30], "s"), "y": Quantity([60], "m")}, None, {}, ValueError, None, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( np.array([[0, 1], [2, 3], [4, 5]]), dims=("x", "y"), coords={"x": [10, 20, 30], "y": [60, 120]}, ), {None: "kg"}, {"x": [15, 25], "y": [75, 105]}, xr.DataArray( [[1.25, 1.75], [3.25, 3.75]], dims=("x", "y"), coords={"x": [15, 25], "y": [75, 105]}, ), {None: "kg"}, None, None, id="DataArray-data units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, {"x": Quantity([1, 3, 5], "m"), "y": Quantity([0, 2], "min")}, xr.DataArray( [[0, 1], [0, 5], [0, 0]], dims=("x", "y"), coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), {"x": "m", "y": "min"}, None, {"bounds_error": False, "fill_value": 0}, id="DataArray-other parameters", ), ), ) def test_interp(obj, units, indexers, expected, expected_units, error, kwargs): obj_ = obj.pint.quantify(units) if error is not None: with pytest.raises(error): obj_.pint.interp(indexers, kwargs=kwargs) else: expected_ = expected.pint.quantify(expected_units) actual = obj_.pint.interp(indexers, kwargs=kwargs) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @requires_scipy @pytest.mark.parametrize( ["obj", "units", "other", "other_units", "expected", "expected_units", "error"], ( pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 120, 240])}), {"x": "dm", "y": "s"}, None, id="Dataset-identical units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), {"x": "m", "y": "min"}, xr.Dataset({"x": ("x", [0, 1, 3, 5]), "y": ("y", [0, 2, 4])}), {"x": "m", "y": "min"}, None, id="Dataset-compatible units", ), pytest.param( xr.Dataset({"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [1, 3]), "y": ("y", [1])}), {"x": "s", "y": "m"}, None, {}, ValueError, id="Dataset-incompatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}), {"x": "dm", "y": "s"}, xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [10, 30, 50]), "y": ("y", [0, 240])}, ), {"x": "dm", "y": "s"}, None, id="DataArray-identical units", ), pytest.param( xr.Dataset( { "a": (("x", "y"), [[0, 1], [2, 3], [4, 5]]), "x": [10, 20, 30], "y": [60, 120], } ), {"a": "kg"}, xr.Dataset({"x": [15, 25], "y": [75, 105]}), {}, xr.Dataset( { "a": (("x", "y"), [[1.25, 1.75], [3.25, 3.75]]), "x": [15, 25], "y": [75, 105], } ), {"a": "kg"}, None, id="Dataset-data units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}), {"x": "m", "y": "min"}, xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), coords={"x": ("x", [1, 3, 5]), "y": ("y", [0, 2])}, ), {"x": "m", "y": "min"}, None, id="DataArray-compatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": ("x", [10, 20, 30]), "y": ("y", [60, 120])}, ), {"x": "dm", "y": "s"}, xr.Dataset({"x": ("x", [10, 30]), "y": ("y", [60])}), {"x": "s", "y": "m"}, None, {}, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( [[0, 1], [2, 3], [4, 5]], dims=("x", "y"), coords={"x": [10, 20, 30], "y": [60, 120]}, ), {"a": "kg"}, xr.Dataset({"x": [15, 25], "y": [75, 105]}), {}, xr.DataArray( [[1.25, 1.75], [3.25, 3.75]], dims=("x", "y"), coords={"x": [15, 25], "y": [75, 105]}, ), {"a": "kg"}, None, id="DataArray-data units", ), ), ) def test_interp_like(obj, units, other, other_units, expected, expected_units, error): obj_ = obj.pint.quantify(units) other_ = other.pint.quantify(other_units) if error is not None: with pytest.raises(error): obj_.pint.interp_like(other_) else: expected_ = expected.pint.quantify(expected_units) actual = obj_.pint.interp_like(other_) assert_units_equal(actual, expected_) assert_identical(actual, expected_) @requires_bottleneck @pytest.mark.parametrize( ["obj", "expected"], ( pytest.param( xr.Dataset( {"a": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])}, coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])}, ), xr.Dataset( {"a": ("x", [nan, 0, 0, 1, 1, 1, 2, 2])}, coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])}, ), id="Dataset-no units", ), pytest.param( xr.Dataset( { "a": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, ), xr.Dataset( {"a": ("x", Quantity([nan, 0, 0, 1, 1, 1, 2, 2], "m"))}, coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, ), id="Dataset-units", ), pytest.param( xr.DataArray( [nan, 0, nan, 1, nan, nan, 2, nan], coords={ "u": ( "x", [nan, 0, nan, 1, nan, nan, 2, nan], ) }, dims="x", ), xr.DataArray( [nan, 0, 0, 1, 1, 1, 2, 2], coords={ "u": ( "x", [nan, 0, nan, 1, nan, nan, 2, nan], ) }, dims="x", ), id="DataArray-no units", ), pytest.param( xr.DataArray( Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, dims="x", ), xr.DataArray( Quantity([nan, 0, 0, 1, 1, 1, 2, 2], "m"), coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, dims="x", ), id="DataArray-units", ), ), ) def test_ffill(obj, expected): actual = obj.pint.ffill(dim="x") assert_identical(actual, expected) assert_units_equal(actual, expected) @requires_bottleneck @pytest.mark.parametrize( ["obj", "expected"], ( pytest.param( xr.Dataset( {"a": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])}, coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])}, ), xr.Dataset( {"a": ("x", [0, 0, 1, 1, 2, 2, 2, nan])}, coords={"u": ("x", [nan, 0, nan, 1, nan, nan, 2, nan])}, ), id="Dataset-no units", ), pytest.param( xr.Dataset( { "a": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, ), xr.Dataset( {"a": ("x", Quantity([0, 0, 1, 1, 2, 2, 2, nan], "m"))}, coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, ), id="Dataset-units", ), pytest.param( xr.DataArray( [nan, 0, nan, 1, nan, nan, 2, nan], coords={ "u": ( "x", [nan, 0, nan, 1, nan, nan, 2, nan], ) }, dims="x", ), xr.DataArray( [0, 0, 1, 1, 2, 2, 2, nan], coords={ "u": ( "x", [nan, 0, nan, 1, nan, nan, 2, nan], ) }, dims="x", ), id="DataArray-no units", ), pytest.param( xr.DataArray( Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, dims="x", ), xr.DataArray( Quantity([0, 0, 1, 1, 2, 2, 2, nan], "m"), coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, 2, nan], "m"), ) }, dims="x", ), id="DataArray-units", ), ), ) def test_bfill(obj, expected): actual = obj.pint.bfill(dim="x") assert_identical(actual, expected) assert_units_equal(actual, expected) @pytest.mark.parametrize( ["obj", "expected"], ( pytest.param( xr.Dataset( {"a": ("x", [nan, 0, nan, 1, nan, nan, nan, 2, nan])}, coords={"u": ("x", [nan, 0, nan, 1, nan, nan, nan, 2, nan])}, ), xr.Dataset( {"a": ("x", [nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan])}, coords={"u": ("x", [nan, 0, nan, 1, nan, nan, nan, 2, nan])}, ), id="Dataset-no units", ), pytest.param( xr.Dataset( { "a": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, ), xr.Dataset( {"a": ("x", Quantity([nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan], "m"))}, coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, ), id="Dataset-units", ), pytest.param( xr.DataArray( [nan, 0, nan, 1, nan, nan, nan, 2, nan], coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, dims="x", ), xr.DataArray( [nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan], coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, dims="x", ), id="DataArray-units", ), pytest.param( xr.DataArray( Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, dims="x", ), xr.DataArray( Quantity([nan, 0, 0.5, 1, 1.25, 1.5, 1.75, 2, nan], "m"), coords={ "u": ( "x", Quantity([nan, 0, nan, 1, nan, nan, nan, 2, nan], "m"), ) }, dims="x", ), id="DataArray-units", ), ), ) def test_interpolate_na(obj, expected): actual = obj.pint.interpolate_na(dim="x") assert_identical(actual, expected) assert_units_equal(actual, expected) pint-xarray-0.6.1/pint_xarray/tests/test_conversion.py000066400000000000000000000712311516032446100233020ustar00rootroot00000000000000import numpy as np import pandas as pd import pint import pytest from tlz.dicttoolz import dissoc from xarray import Coordinates, DataArray, Dataset, Variable from xarray.core.indexes import PandasIndex from pint_xarray import conversion from pint_xarray.errors import PintExceptionGroup from pint_xarray.index import PintIndex from pint_xarray.tests.utils import ( assert_array_equal, assert_array_units_equal, assert_identical, assert_indexer_units_equal, assert_indexers_equal, ) unit_registry = pint.UnitRegistry() Quantity = unit_registry.Quantity Unit = unit_registry.Unit pytestmark = pytest.mark.filterwarnings("error::pint.UnitStrippedWarning") def filter_none_values(mapping): return {k: v for k, v in mapping.items() if v is not None} def to_quantity(v, u): if u is None: return v return Quantity(v, u) def convert_quantity(q, u): if u is None: return q if not isinstance(q, Quantity): q = Quantity(q) return q.to(u) def strip_quantity(q): try: return q.magnitude except AttributeError: return q class TestArrayFunctions: @pytest.mark.parametrize( ["unit", "data", "expected", "match"], ( pytest.param( 1.2, np.array([0, 1]), None, "cannot use .+ as a unit", id="not a unit" ), pytest.param( 1, np.array([0, 1]), None, "cannot use .+ as a unit", id="no unit (1)" ), pytest.param( None, np.array([0, 1]), np.array([0, 1]), None, id="no unit (None)" ), pytest.param( "m", np.array([0, 1]), None, "cannot use .+ as a unit", id="string" ), pytest.param( Unit("m"), np.array([0, 1]), Quantity([0, 1], "m"), None, id="unit object", ), pytest.param( Unit("m"), Quantity(np.array([0, 1]), "s"), None, "already has units", id="unit object on quantity", ), pytest.param( Unit("m"), Quantity(np.array([0, 1]), "m"), Quantity(np.array([0, 1]), "m"), None, id="unit object on quantity with same unit", ), pytest.param( Unit("mm"), Quantity(np.array([0, 1]), "m"), None, "already has units", id="unit object on quantity with similar unit", ), ), ) def test_array_attach_units(self, data, unit, expected, match): if match is not None: with pytest.raises(ValueError, match=match): conversion.array_attach_units(data, unit) return actual = conversion.array_attach_units(data, unit) assert_array_units_equal(expected, actual) assert_array_equal(expected, actual) @pytest.mark.parametrize( ["unit", "data", "expected", "error", "match"], ( pytest.param( 1.2, np.array([0, 1, 2]), None, ValueError, "cannot use .+ as a unit", id="not a unit-ndarray", ), pytest.param( 1, np.array([0, 1, 2]), None, ValueError, "cannot use .+ as a unit", id="no unit (1)-ndarray", ), pytest.param( None, np.array([0, 1, 2]), np.array([0, 1, 2]), None, None, id="no unit (None)-ndarray", ), pytest.param( "mm", np.array([0, 1, 2]), None, ValueError, "cannot convert a non-quantity using .+ as unit", id="string-ndarray", ), pytest.param( Unit("deg"), np.array([0, np.pi / 2, np.pi]), Quantity([0, 90, 180], "deg"), None, None, id="dimensionless-ndarray", ), pytest.param( Unit("mm"), np.array([0, np.pi / 2, np.pi]), None, pint.DimensionalityError, None, id="unit-ndarray", ), pytest.param( "mm", Quantity([0, 1, 2], "m"), Quantity([0, 1000, 2000], "mm"), None, None, id="string-quantity", ), pytest.param( unit_registry.mm, Quantity([0, 1, 2], "m"), Quantity([0, 1000, 2000], "mm"), None, None, id="unit object", ), pytest.param( "s", Quantity([0, 1, 2], "m"), None, pint.DimensionalityError, None, id="quantity-incompatible unit", ), ), ) def test_array_convert_units(self, data, unit, expected, error, match): if error is not None: with pytest.raises(error, match=match): conversion.array_convert_units(data, unit) return actual = conversion.array_convert_units(data, unit) assert_array_equal(expected, actual) @pytest.mark.parametrize( ["data", "expected"], ( pytest.param(np.array([0, 1]), None, id="array_like"), pytest.param(Quantity([1, 2], "m"), Unit("m"), id="quantity"), ), ) def test_array_extract_units(self, data, expected): actual = conversion.array_extract_units(data) assert expected == actual @pytest.mark.parametrize( ["data", "expected"], ( pytest.param(np.array([1, 2]), np.array([1, 2]), id="array_like"), pytest.param(Quantity([1, 2], "m"), np.array([1, 2]), id="quantity"), ), ) def test_array_strip_units(self, data, expected): actual = conversion.array_strip_units(data) assert_array_equal(expected, actual) class TestXarrayFunctions: @pytest.mark.parametrize("type", ("Dataset", "DataArray")) @pytest.mark.parametrize( "units", ( pytest.param({}, id="empty units"), pytest.param({"a": None, "b": None, "u": None, "x": None}, id="no units"), pytest.param( {"a": unit_registry.m, "b": unit_registry.m, "u": None, "x": None}, id="data units", ), pytest.param( {"a": None, "b": None, "u": unit_registry.s, "x": None}, id="coord units", ), pytest.param( {"a": None, "b": None, "u": None, "x": unit_registry.m}, id="dim units" ), ), ) def test_attach_units(self, type, units): a = np.linspace(-1, 1, 5) b = np.linspace(0, 1, 5) x = np.linspace(0, 100, 5) u = np.arange(5) q_a = to_quantity(a, units.get("a")) q_b = to_quantity(b, units.get("b")) q_x = to_quantity(x, units.get("x")) q_u = to_quantity(u, units.get("u")) index = PandasIndex(x, dim="x") if units.get("x") is not None: index = PintIndex(index=index, units={"x": units.get("x")}) obj = Dataset({"a": ("x", a), "b": ("x", b)}, coords={"u": ("x", u), "x": x}) coords = Coordinates( coords={"u": Variable("x", q_u), "x": Variable("x", q_x)}, indexes={"x": index}, ) expected = Dataset( {"a": ("x", q_a), "b": ("x", q_b)}, coords=coords, ) if type == "DataArray": obj = obj["a"] expected = expected["a"] actual = conversion.attach_units(obj, units) assert_identical(actual, expected) if units.get("x") is None: assert not isinstance(actual.xindexes["x"], PintIndex) else: assert isinstance(actual.xindexes["x"], PintIndex) assert actual.xindexes["x"].units == {"x": units.get("x")} @pytest.mark.parametrize("type", ("DataArray", "Dataset")) def test_attach_unit_attributes(self, type): units = {"a": "K", "b": "hPa", "u": "m", "x": "s"} obj = Dataset( data_vars={"a": ("x", []), "b": ("x", [])}, coords={"x": [], "u": ("x", [])}, ) expected = Dataset( {"a": ("x", [], {"units": "K"}), "b": ("x", [], {"units": "hPa"})}, coords={"x": ("x", [], {"units": "s"}), "u": ("x", [], {"units": "m"})}, ) if type == "DataArray": obj = obj["a"] expected = expected["a"] actual = conversion.attach_unit_attributes(obj, units) assert_identical(actual, expected) @pytest.mark.parametrize("type", ("DataArray", "Dataset")) @pytest.mark.parametrize( ["variant", "units", "error", "suberrors"], ( pytest.param("none", {}, None, {}, id="none-no units"), pytest.param( "none", {"a": Unit("g"), "b": Unit("Pa"), "u": Unit("ms"), "x": Unit("mm")}, PintExceptionGroup, { "a": (pint.DimensionalityError, "'a'"), "b": (pint.DimensionalityError, "'b'"), "u": (pint.DimensionalityError, "'u'"), "x": (ValueError, "'x'"), }, id="none-with units", ), pytest.param("data", {}, None, {}, id="data-no units"), pytest.param( "data", {"a": Unit("g"), "b": Unit("Pa")}, None, {}, id="data-compatible units", ), pytest.param( "data", {"a": Unit("s"), "b": Unit("m")}, PintExceptionGroup, { "a": (pint.DimensionalityError, "'a'"), "b": (pint.DimensionalityError, "'b'"), }, id="data-incompatible units", ), pytest.param("dims", {}, None, {}, id="dims-no units"), pytest.param( "dims", {"x": Unit("mm")}, None, {}, id="dims-compatible units" ), pytest.param( "dims", {"x": Unit("ms")}, PintExceptionGroup, {"x": (ValueError, "'x'")}, id="dims-incompatible units", ), pytest.param("coords", {}, None, {}, id="coords-no units"), pytest.param( "coords", {"u": Unit("ms")}, None, {"u": (pint.DimensionalityError, "'u'")}, id="coords-compatible units", ), pytest.param( "coords", {"u": Unit("mm")}, PintExceptionGroup, { "u": ( pint.DimensionalityError, "incompatible units for variable 'u'", ) }, id="coords-incompatible units", ), ), ) def test_convert_units(self, type, variant, units, error, suberrors): variants = { "none": {"a": None, "b": None, "u": None, "x": None}, "data": {"a": Unit("kg"), "b": Unit("hPa"), "u": None, "x": None}, "coords": {"a": None, "b": None, "u": Unit("s"), "x": None}, "dims": {"a": None, "b": None, "u": None, "x": Unit("m")}, } a = np.linspace(-1, 1, 3) b = np.linspace(1, 2, 3) u = np.linspace(0, 100, 3) x = np.arange(3) original_units = variants.get(variant) q_a = to_quantity(a, original_units.get("a")) q_b = to_quantity(b, original_units.get("b")) q_u = to_quantity(u, original_units.get("u")) q_x = to_quantity(x, original_units.get("x")) x_index = PandasIndex(pd.Index(x), "x") if original_units.get("x") is not None: x_index = PintIndex(index=x_index, units={"x": original_units.get("x")}) obj = Dataset( { "a": ("x", q_a), "b": ("x", q_b), }, coords=Coordinates( {"u": ("x", q_u), "x": ("x", q_x)}, indexes={"x": x_index}, ), ) if type == "DataArray": obj = obj["a"] suberrors = dissoc(suberrors, "b") if error is not None: matchers = [ pytest.RaisesExc(err, match=match) for err, match in suberrors.values() ] with pytest.RaisesGroup( *matchers, match="Cannot convert variables", check=lambda eg: isinstance(eg, PintExceptionGroup), ): conversion.convert_units(obj, units) return expected_a = convert_quantity(q_a, units.get("a", original_units.get("a"))) expected_b = convert_quantity(q_b, units.get("b", original_units.get("b"))) expected_u = convert_quantity(q_u, units.get("u", original_units.get("u"))) expected_x = convert_quantity(q_x, units.get("x", original_units.get("x"))) expected_index = PandasIndex(pd.Index(strip_quantity(expected_x)), "x") expected_index_units = units.get("x", original_units.get("x")) if expected_index_units is not None: expected_index = PintIndex( index=expected_index, units={"x": expected_index_units} ) expected = Dataset( { "a": ("x", expected_a), "b": ("x", expected_b), }, coords=Coordinates( {"u": ("x", expected_u), "x": ("x", expected_x)}, indexes={"x": expected_index}, ), ) if type == "DataArray": expected = expected["a"] actual = conversion.convert_units(obj, units) assert conversion.extract_units(actual) == conversion.extract_units(expected) assert_identical(actual, expected) @pytest.mark.parametrize( "units", ( pytest.param({"a": None, "b": None, "u": None, "x": None}, id="none"), pytest.param( {"a": Unit("kg"), "b": Unit("hPa"), "u": None, "x": None}, id="data" ), pytest.param({"a": None, "b": None, "u": Unit("s"), "x": None}, id="coord"), pytest.param({"a": None, "b": None, "u": None, "x": Unit("m")}, id="dims"), ), ) @pytest.mark.parametrize("type", ("DataArray", "Dataset")) def test_extract_units(self, type, units): a = np.linspace(-1, 1, 2) b = np.linspace(0, 1, 2) u = np.linspace(0, 100, 2) x = np.arange(2) index = PandasIndex(x, "x") if units.get("x") is not None: index = PintIndex(index=index, units={"x": units.get("x")}) obj = Dataset( { "a": ("x", to_quantity(a, units.get("a"))), "b": ("x", to_quantity(b, units.get("b"))), }, coords=Coordinates( { "u": ("x", to_quantity(u, units.get("u"))), "x": ("x", to_quantity(x, units.get("x"))), }, indexes={"x": index}, ), ) if type == "DataArray": obj = obj["a"] units = units.copy() del units["b"] assert conversion.extract_units(obj) == units @pytest.mark.parametrize( ["obj", "expected"], ( pytest.param( DataArray( coords={ "x": ("x", [], {"units": "m"}), "u": ("x", [], {"units": "s"}), }, attrs={"units": "hPa"}, dims="x", ), {"x": "m", "u": "s", None: "hPa"}, id="DataArray", ), pytest.param( Dataset( data_vars={ "a": ("x", [], {"units": "K"}), "b": ("x", [], {"units": "hPa"}), }, coords={ "x": ("x", [], {"units": "m"}), "u": ("x", [], {"units": "s"}), }, ), {"a": "K", "b": "hPa", "x": "m", "u": "s"}, id="Dataset", ), pytest.param( Dataset(coords={"t": ("t", [], {"units": "seconds since 2000-01-01"})}), {}, id="datetime_unit", ), ), ) def test_extract_unit_attributes(self, obj, expected): actual = conversion.extract_unit_attributes(obj) assert expected == actual @pytest.mark.parametrize( ["obj", "expected"], ( pytest.param( DataArray( dims="x", data=Quantity([0, 4, 3], "kg"), coords=Coordinates( { "u": ("x", Quantity([2, 3, 4], "s")), "x": Quantity([0, 1, 2], "m"), }, indexes={}, ), ), {None: None, "u": None, "x": None}, id="DataArray", ), pytest.param( Dataset( data_vars={ "a": ("x", Quantity([3, 2, 5], "Pa")), "b": ("x", Quantity([0, 2, -1], "kg")), }, coords=Coordinates( { "u": ("x", Quantity([2, 3, 4], "s")), "x": Quantity([0, 1, 2], "m"), }, indexes={}, ), ), {"a": None, "b": None, "u": None, "x": None}, id="Dataset", ), ), ) def test_strip_units(self, obj, expected): actual = conversion.strip_units(obj) assert conversion.extract_units(actual) == expected @pytest.mark.parametrize( ["obj", "expected"], ( pytest.param( DataArray( coords={ "x": ("x", [], {"units": "m"}), "u": ("x", [], {"units": "s"}), }, attrs={"units": "hPa"}, dims="x", ), {"x": "m", "u": "s", None: "hPa"}, id="DataArray", ), pytest.param( Dataset( data_vars={ "a": ("x", [], {"units": "K"}), "b": ("x", [], {"units": "hPa"}), }, coords={ "x": ("x", [], {"units": "m"}), "u": ("x", [], {"units": "s"}), }, ), {"a": "K", "b": "hPa", "x": "m", "u": "s"}, id="Dataset", ), pytest.param( Dataset(coords={"t": ("t", [], {"units": "seconds since 2000-01-01"})}), {}, id="datetime_unit", ), ), ) def test_strip_unit_attributes(self, obj, expected): actual = conversion.strip_unit_attributes(obj) expected = {} assert ( filter_none_values(conversion.extract_unit_attributes(actual)) == expected ) class TestIndexerFunctions: @pytest.mark.parametrize( ["indexers", "units", "expected", "error", "suberrors"], ( pytest.param( {"x": 1}, {"x": None}, {"x": 1}, None, None, id="scalar-no units" ), pytest.param( {"x": 1}, {"x": "dimensionless"}, None, PintExceptionGroup, {"x": (ValueError, "'x'")}, id="scalar-dimensionless", ), pytest.param( {"x": Quantity(1, "m")}, {"x": Unit("dm")}, {"x": Quantity(10, "dm")}, None, None, id="scalar-units", ), pytest.param( {"x": np.array([1, 2])}, {"x": None}, {"x": np.array([1, 2])}, None, None, id="array-no units", ), pytest.param( {"x": Quantity([1, 2], "m")}, {"x": Unit("dm")}, {"x": Quantity([10, 20], "dm")}, None, None, id="array-units", ), pytest.param( {"x": Variable("x", [1, 2])}, {"x": None}, {"x": Variable("x", [1, 2])}, None, None, id="Variable-no units", ), pytest.param( {"x": Variable("x", Quantity([1, 2], "m"))}, {"x": Unit("dm")}, {"x": Variable("x", Quantity([10, 20], "dm"))}, None, None, id="Variable-units", ), pytest.param( {"x": DataArray([1, 2], dims="x")}, {"x": None}, {"x": DataArray([1, 2], dims="x")}, None, None, id="DataArray-no units", ), pytest.param( {"x": DataArray(Quantity([1, 2], "m"), dims="x")}, {"x": Unit("dm")}, {"x": DataArray(Quantity([10, 20], "dm"), dims="x")}, None, None, id="DataArray-units", ), pytest.param( {"x": slice(None)}, {"x": None}, {"x": slice(None)}, None, None, id="empty slice-no units", ), pytest.param( {"x": slice(1, None)}, {"x": None}, {"x": slice(1, None)}, None, None, id="slice-no units", ), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, {"x": Unit("m")}, {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, None, None, id="slice-identical units", ), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2000, "mm"))}, {"x": Unit("dm")}, {"x": slice(Quantity(10, "dm"), Quantity(20, "dm"))}, None, None, id="slice-compatible units", ), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, {"x": Unit("ms")}, None, PintExceptionGroup, {"x": (pint.DimensionalityError, "'x'")}, id="slice-incompatible units", ), pytest.param( {"x": slice(1000, Quantity(2000, "ms"))}, {"x": Unit("s")}, None, PintExceptionGroup, {"x": (pint.DimensionalityError, "'x'")}, id="slice-incompatible units-mixed", ), ), ) def test_convert_indexer_units(self, indexers, units, expected, error, suberrors): if error is not None: matchers = [ pytest.RaisesExc(err, match=match) for err, match in suberrors.values() ] with pytest.RaisesGroup( *matchers, match="Cannot convert indexers", check=lambda eg: isinstance(eg, PintExceptionGroup), ): conversion.convert_indexer_units(indexers, units) else: actual = conversion.convert_indexer_units(indexers, units) assert_indexers_equal(actual, expected) assert_indexer_units_equal(actual, expected) @pytest.mark.parametrize( ["indexers", "expected"], ( pytest.param({"x": 1}, {"x": None}, id="scalar-no units"), pytest.param({"x": Quantity(1, "m")}, {"x": Unit("m")}, id="scalar-units"), pytest.param({"x": np.array([1, 2])}, {"x": None}, id="array-no units"), pytest.param( {"x": Quantity([1, 2], "s")}, {"x": Unit("s")}, id="array-units" ), pytest.param( {"x": Variable("x", [1, 2])}, {"x": None}, id="Variable-no units" ), pytest.param( {"x": Variable("x", Quantity([1, 2], "m"))}, {"x": Unit("m")}, id="Variable-units", ), pytest.param( {"x": DataArray([1, 2], dims="x")}, {"x": None}, id="DataArray-no units" ), pytest.param( {"x": DataArray(Quantity([1, 2], "s"), dims="x")}, {"x": Unit("s")}, id="DataArray-units", ), pytest.param({"x": slice(None)}, {"x": None}, id="empty slice-no units"), pytest.param({"x": slice(1, None)}, {"x": None}, id="slice-no units"), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, {"x": Unit("m")}, id="slice-identical units", ), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2000, "mm"))}, {"x": Unit("m")}, id="slice-compatible units", ), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2, "ms"))}, ValueError, id="slice-incompatible units", ), pytest.param( {"x": slice(1, Quantity(2, "ms"))}, ValueError, id="slice-incompatible units-mixed", ), pytest.param( {"x": slice(1, Quantity(2, "rad"))}, {"x": Unit("rad")}, id="slice-incompatible units-mixed-dimensionless", ), ), ) def test_extract_indexer_units(self, indexers, expected): if isinstance(expected, type) and issubclass(expected, Exception): with pytest.raises(expected): conversion.extract_indexer_units(indexers) else: actual = conversion.extract_indexer_units(indexers) assert actual == expected @pytest.mark.parametrize( ["indexers", "expected"], ( pytest.param({"x": 1}, {"x": 1}, id="scalar-no units"), pytest.param({"x": Quantity(1, "m")}, {"x": 1}, id="scalar-units"), pytest.param( {"x": np.array([1, 2])}, {"x": np.array([1, 2])}, id="array-no units", ), pytest.param( {"x": Quantity([1, 2], "s")}, {"x": np.array([1, 2])}, id="array-units" ), pytest.param( {"x": Variable("x", [1, 2])}, {"x": Variable("x", [1, 2])}, id="Variable-no units", ), pytest.param( {"x": Variable("x", Quantity([1, 2], "m"))}, {"x": Variable("x", [1, 2])}, id="Variable-units", ), pytest.param( {"x": DataArray([1, 2], dims="x")}, {"x": DataArray([1, 2], dims="x")}, id="DataArray-no units", ), pytest.param( {"x": DataArray(Quantity([1, 2], "s"), dims="x")}, {"x": DataArray([1, 2], dims="x")}, id="DataArray-units", ), pytest.param( {"x": slice(None)}, {"x": slice(None)}, id="empty slice-no units" ), pytest.param( {"x": slice(1, None)}, {"x": slice(1, None)}, id="slice-no units" ), pytest.param( {"x": slice(Quantity(1, "m"), Quantity(2, "m"))}, {"x": slice(1, 2)}, id="slice-units", ), ), ) def test_strip_indexer_units(self, indexers, expected): actual = conversion.strip_indexer_units(indexers) assert_indexers_equal(actual, expected) pint-xarray-0.6.1/pint_xarray/tests/test_expects.py000066400000000000000000000221611516032446100225660ustar00rootroot00000000000000import re import pint import pytest import xarray as xr import pint_xarray from pint_xarray.testing import assert_units_equal ureg = pint_xarray.unit_registry class TestExpects: @pytest.mark.parametrize( ["values", "units", "expected"], ( ((ureg.Quantity(1, "m"), 2), ("mm", None, None), 500), ((ureg.Quantity(1, "m"), ureg.Quantity(0.5, "s")), ("mm", "ms", None), 2), ( (xr.DataArray(4).pint.quantify("km"), 2), ("m", None, None), xr.DataArray(2000), ), ( ( xr.DataArray([4, 2, 0]).pint.quantify("cm"), xr.DataArray([4, 2, 1]).pint.quantify("mg"), ), ("m", "g", None), xr.DataArray([10, 10, 0]), ), ( (ureg.Quantity(16, "m"), 2, ureg.Quantity(4, "s")), ("mm", None, "ms"), 2, ), ), ) def test_args(self, values, units, expected): @pint_xarray.expects(*units) def func(a, b, c=1): return a / (b * c) actual = func(*values) if isinstance(actual, xr.DataArray): xr.testing.assert_identical(actual, expected) elif isinstance(actual, pint.Quantity): pint.testing.assert_equal(actual, expected) else: assert actual == expected @pytest.mark.parametrize( ["value", "units", "error", "message", "multiple"], ( ( ureg.Quantity(1, "m"), (None, None), TypeError, "Passed in a quantity where none was expected", True, ), (1, ("m", None), TypeError, "Attempting to convert non-quantity", True), ( 1, (None,), ValueError, "Missing units for the following parameters: 'b'", False, ), ( ureg.Quantity(1, "m"), ("nonsense_unit", None), pint.errors.UndefinedUnitError, "'nonsense_unit' is not defined in the unit registry", True, ), ), ) def test_args_error(self, value, units, error, message, multiple): if multiple: root_error = ExceptionGroup root_message = "Errors while converting parameters" else: root_error = error root_message = message with pytest.raises(root_error, match=root_message) as excinfo: @pint_xarray.expects(*units) def func(a, b=1): return a * b func(value) if not multiple: return group = excinfo.value assert len(group.exceptions) == 1, f"Found {len(group.exceptions)} exceptions" exc = group.exceptions[0] assert isinstance( exc, error ), f"Unexpected exception type: {type(exc)}, expected {error}" if not re.search(message, str(exc)): raise AssertionError(f"exception {exc!r} did not match pattern {message!r}") @pytest.mark.parametrize( ["values", "units", "expected"], ( ( {"a": ureg.Quantity(1, "m"), "b": 2}, {"a": "mm", "b": None, "c": None}, 1000, ), ( {"a": 2, "b": ureg.Quantity(100, "cm")}, {"a": None, "b": "m", "c": None}, 4, ), ( {"a": ureg.Quantity(1, "m"), "b": ureg.Quantity(0.5, "s")}, {"a": "mm", "b": "ms", "c": None}, 4, ), ( {"a": xr.DataArray(4).pint.quantify("km"), "b": 2}, {"a": "m", "b": None, "c": None}, xr.DataArray(4000), ), ( { "a": xr.DataArray([4, 2, 0]).pint.quantify("cm"), "b": xr.DataArray([4, 2, 1]).pint.quantify("mg"), }, {"a": "m", "b": "g", "c": None}, xr.DataArray([20, 20, 0]), ), ), ) def test_kwargs(self, values, units, expected): @pint_xarray.expects(**units) def func(a, b, c=2): return a / b * c actual = func(**values) if isinstance(actual, xr.DataArray): xr.testing.assert_identical(actual, expected) elif isinstance(actual, pint.Quantity): pint.testing.assert_equal(actual, expected) else: assert actual == expected @pytest.mark.parametrize( ["values", "return_value_units", "expected"], ( ((1, 2), ("m", "s"), (ureg.Quantity(1, "m"), ureg.Quantity(2, "s"))), ((1, 2), "m / s", ureg.Quantity(0.5, "m / s")), ((1, 2), None, 0.5), ( (xr.DataArray(2), 2), ("m", "s"), (xr.DataArray(2).pint.quantify("m"), ureg.Quantity(2, "s")), ), ( (xr.DataArray(2), 2), "kg / m^2", xr.DataArray(1).pint.quantify("kg / m^2"), ), ), ) def test_return_value(self, values, return_value_units, expected): multiple = isinstance(return_value_units, tuple) @pint_xarray.expects(a=None, b=None, return_value=return_value_units) def func(a, b): if multiple: return a, b else: return a / b actual = func(*values) if isinstance(actual, xr.DataArray): xr.testing.assert_identical(actual, expected) elif isinstance(actual, pint.Quantity): pint.testing.assert_equal(actual, expected) else: assert actual == expected def test_return_value_none(self): @pint_xarray.expects(None) def func(a): return None actual = func(1) assert actual is None def test_return_value_none_error(self): @pint_xarray.expects(return_value="Hz") def func(): return None with pytest.raises( ValueError, match="mismatched number of return values: expected 1 but got 0.", ): func() @pytest.mark.parametrize( [ "return_value_units", "multiple_units", "error", "multiple_errors", "message", ], ( ( ("m", "s"), False, ValueError, False, "mismatched number of return values", ), ( "m", True, ValueError, False, "mismatched number of return values: expected 1 but got 2", ), ( ("m",), True, ValueError, False, "mismatched number of return values: expected 1 but got 2", ), (1, False, TypeError, True, "units must be of type"), (("m",), False, ValueError, False, ".*expected a 1-sized tuple.*"), ("m", 1, ValueError, False, ".*expected a single return value.*"), ), ) def test_return_value_error( self, return_value_units, multiple_units, error, multiple_errors, message ): if multiple_errors: root_error = ExceptionGroup root_message = "Errors while converting return values" else: root_error = error root_message = message with pytest.raises(root_error, match=root_message) as excinfo: @pint_xarray.expects(a=None, b=None, return_value=return_value_units) def func(a, b): if not isinstance(multiple_units, bool) and multiple_units == 1: print("return 1-tuple") return (a / b,) elif multiple_units: return a, b else: return a / b func(1, 2) if not multiple_errors: return group = excinfo.value assert len(group.exceptions) == 1, f"Found {len(group.exceptions)} exceptions" exc = group.exceptions[0] assert isinstance( exc, error ), f"Unexpected exception type: {type(exc)}, expected {error}" if not re.search(message, str(exc)): raise AssertionError(f"exception {exc!r} did not match pattern {message!r}") def test_datasets(self): @pint_xarray.expects({"m": "kg", "a": "m / s^2"}, return_value={"f": "newtons"}) def second_law(ds): f_da = ds["m"] * ds["a"] return f_da.to_dataset(name="f") ds = xr.Dataset({"m": 0.1, "a": 10}).pint.quantify( {"m": "tons", "a": "feet / second^2"} ) expected = xr.Dataset({"f": ds["m"] * ds["a"]}).pint.to("newtons") actual = second_law(ds) assert_units_equal(actual, expected) xr.testing.assert_allclose(actual, expected) pint-xarray-0.6.1/pint_xarray/tests/test_formatting.py000066400000000000000000000010511516032446100232600ustar00rootroot00000000000000import pint import pytest # only need to register _repr_inline_ import pint_xarray # noqa: F401 unit_registry = pint.UnitRegistry(force_ndarray_like=True) @pytest.mark.parametrize( ("length", "expected"), ( (40, "[N] 7.1 5.4 9.8 21.4 15.3"), (20, "[N] 7.1 5.4 ... 15.3"), (10, "[N] 7.1..."), (7, "[N] ..."), (3, "[N] ..."), ), ) def test_inline_repr(length, expected): quantity = unit_registry.Quantity([7.1, 5.4, 9.8, 21.4, 15.3], "N") assert quantity._repr_inline_(length) == expected pint-xarray-0.6.1/pint_xarray/tests/test_index.py000066400000000000000000000216251516032446100222260ustar00rootroot00000000000000import numpy as np import pandas as pd import pytest import xarray as xr from xarray.core.indexes import IndexSelResult, PandasIndex from pint_xarray import unit_registry as ureg from pint_xarray.index import PintIndex def indexer_equal(first, second): if type(first) is not type(second): return False if isinstance(first, np.ndarray): return np.all(first == second) else: return first == second @pytest.mark.parametrize( "base_index", [ PandasIndex(pd.Index([1, 2, 3]), dim="x"), PandasIndex(pd.Index([0.1, 0.2, 0.3]), dim="x"), PandasIndex(pd.Index([1j, 2j, 3j]), dim="y"), ], ) @pytest.mark.parametrize("units", [ureg.Unit("m"), ureg.Unit("s")]) def test_init(base_index, units): index = PintIndex(index=base_index, units={base_index.dim: units}) assert index.index.equals(base_index) assert index.units == {base_index.dim: units} def test_init_error(): base_index = PandasIndex(pd.Index([1, 2, 3]), dim="x") with pytest.raises(TypeError, match="dict of coordinate names to units"): PintIndex(index=base_index, units=ureg.Unit("s")) def test_replace(): old_index = PandasIndex([1, 2, 3], dim="x") new_index = PandasIndex([0.1, 0.2, 0.3], dim="x") old = PintIndex(index=old_index, units={"x": ureg.Unit("m")}) new = old._replace(new_index) assert new.index.equals(new_index) assert new.units == old.units # no mutation assert old.index.equals(old_index) @pytest.mark.parametrize( ["wrapped_index", "units", "expected"], ( pytest.param( PandasIndex(pd.Index([1, 2, 3]), dim="x"), {"x": ureg.Unit("m")}, {"x": xr.Variable("x", ureg.Quantity([1, 2, 3], "m"))}, ), pytest.param( PandasIndex(pd.Index([1j, 2j, 3j]), dim="y"), {"y": ureg.Unit("ms")}, {"y": xr.Variable("y", ureg.Quantity([1j, 2j, 3j], "ms"))}, ), ), ) def test_create_variables(wrapped_index, units, expected): index = PintIndex(index=wrapped_index, units=units) actual = index.create_variables() assert list(actual.keys()) == list(expected.keys()) assert all([actual[k].equals(expected[k]) for k in expected.keys()]) @pytest.mark.parametrize( ["labels", "expected"], ( ({"x": ureg.Quantity(1, "m")}, IndexSelResult(dim_indexers={"x": 0})), ({"x": ureg.Quantity(3000, "mm")}, IndexSelResult(dim_indexers={"x": 2})), ({"x": ureg.Quantity(0.002, "km")}, IndexSelResult(dim_indexers={"x": 1})), ( {"x": ureg.Quantity([0.002, 0.004], "km")}, IndexSelResult(dim_indexers={"x": np.array([1, 3])}), ), ( {"x": slice(ureg.Quantity(2, "m"), ureg.Quantity(3000, "mm"))}, IndexSelResult(dim_indexers={"x": slice(1, 3)}), ), ), ) def test_sel(labels, expected): index = PintIndex( index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("m")} ) actual = index.sel(labels) assert isinstance(actual, IndexSelResult) assert list(actual.dim_indexers.keys()) == list(expected.dim_indexers.keys()) assert all( [ indexer_equal(actual.dim_indexers[k], expected.dim_indexers[k]) for k in expected.dim_indexers.keys() ] ) @pytest.mark.parametrize( ["labels", "expected"], ( ( {"x": ureg.Quantity(1.1, "m")}, IndexSelResult(dim_indexers={"x": np.array(0)}), ), ( {"x": ureg.Quantity(3100, "mm")}, IndexSelResult(dim_indexers={"x": np.array(2)}), ), ( {"x": ureg.Quantity(0.0021, "km")}, IndexSelResult(dim_indexers={"x": np.array(1)}), ), ( {"x": ureg.Quantity([0.0021, 0.0041], "km")}, IndexSelResult(dim_indexers={"x": np.array([1, 3])}), ), ), ) def test_sel_nearest(labels, expected): index = PintIndex( index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("m")} ) actual = index.sel(labels, method="nearest") assert isinstance(actual, IndexSelResult) assert actual.dim_indexers.keys() == expected.dim_indexers.keys() assert all( indexer_equal(actual.dim_indexers[k], expected.dim_indexers[k]) for k in expected.dim_indexers.keys() ) @pytest.mark.parametrize( "indexers", ({"y": 0}, {"y": [1, 2]}, {"y": slice(0, None, 2)}, {"y": xr.Variable("y", [1])}), ) def test_isel(indexers): wrapped_index = PandasIndex(pd.Index([1, 2, 3, 4]), dim="y") index = PintIndex(index=wrapped_index, units={"y": ureg.Unit("s")}) actual = index.isel(indexers) wrapped_ = wrapped_index.isel(indexers) if wrapped_ is not None: expected = PintIndex( index=wrapped_index.isel(indexers), units={"y": ureg.Unit("s")} ) else: expected = None assert (actual is None and expected is None) or actual.equals(expected) @pytest.mark.parametrize( ["other", "expected"], ( ( PintIndex( index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("cm")}, ), True, ), (PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), False), ( PintIndex( index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("m")}, ), False, ), ( PintIndex( index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="y"), units={"y": ureg.Unit("cm")}, ), False, ), ( PintIndex( index=PandasIndex(pd.Index([1, 3, 3, 4]), dim="x"), units={"x": ureg.Unit("cm")}, ), False, ), ), ) def test_equals(other, expected): index = PintIndex( index=PandasIndex(pd.Index([1, 2, 3, 4]), dim="x"), units={"x": ureg.Unit("cm")} ) actual = index.equals(other) assert actual == expected @pytest.mark.filterwarnings("error") def test_align_equals_warning(): index1 = PintIndex( index=PandasIndex(pd.Index([1, 2]), dim="x"), units={"x": ureg.Unit("m")}, ) index2 = PintIndex( index=PandasIndex(pd.Index([0, 1, 2]), dim="y"), units={"y": ureg.Unit("m")}, ) ds = xr.Dataset( {"a": (["y", "x"], [[-1, 1], [0, 2], [1, 3]])}, coords=xr.Coordinates( {"x": [1, 2], "y": [0, 1, 2]}, indexes={"x": index1, "y": index2} ), ) # trigger comparison ds["a"] * ds["x"] * ds["y"] @pytest.mark.parametrize( ["shifts", "expected_index"], ( ({"x": 0}, PandasIndex(pd.Index([-2, -1, 0, 1, 2]), dim="x")), ({"x": 1}, PandasIndex(pd.Index([2, -2, -1, 0, 1]), dim="x")), ({"x": 2}, PandasIndex(pd.Index([1, 2, -2, -1, 0]), dim="x")), ({"x": -1}, PandasIndex(pd.Index([-1, 0, 1, 2, -2]), dim="x")), ({"x": -2}, PandasIndex(pd.Index([0, 1, 2, -2, -1]), dim="x")), ), ) def test_roll(shifts, expected_index): index = PintIndex( index=PandasIndex(pd.Index([-2, -1, 0, 1, 2]), dim="x"), units={"x": ureg.Unit("m")}, ) actual = index.roll(shifts) expected = index._replace(expected_index) assert actual.equals(expected) @pytest.mark.parametrize("dims_dict", ({"y": "x"}, {"y": "z"})) @pytest.mark.parametrize("name_dict", ({"y2": "y3"}, {"y2": "y1"})) def test_rename(name_dict, dims_dict): wrapped_index = PandasIndex(pd.Index([1, 2], name="y2"), dim="y") index = PintIndex(index=wrapped_index, units={"y2": ureg.Unit("m")}) actual = index.rename(name_dict, dims_dict) expected = PintIndex( index=wrapped_index.rename(name_dict, dims_dict), units={name_dict["y2"]: ureg.Unit("m")}, ) assert actual.equals(expected) @pytest.mark.parametrize("indexer", ([0], slice(0, 2))) def test_getitem(indexer): wrapped_index = PandasIndex(pd.Index([1, 2], name="y2"), dim="y") index = PintIndex(index=wrapped_index, units={"y": ureg.Unit("m")}) actual = index[indexer] expected = PintIndex(index=wrapped_index[indexer], units=index.units) assert actual.equals(expected) @pytest.mark.parametrize("wrapped_index", (PandasIndex(pd.Index([1, 2]), dim="x"),)) def test_repr_inline(wrapped_index): index = PintIndex(index=wrapped_index, units={"x": ureg.Unit("m")}) # TODO: parametrize actual = index._repr_inline_(90) assert "PintIndex" in actual assert wrapped_index.__class__.__name__ in actual assert "units" in actual @pytest.mark.parametrize("wrapped_index", (PandasIndex(pd.Index([1, 2]), dim="x"),)) def test_repr(wrapped_index): index = PintIndex(index=wrapped_index, units={"x": ureg.Unit("m")}) # TODO: parametrize actual = repr(index) assert "= 2023.07.0", "numpy >= 1.26", "pint >= 0.24", ] dynamic = ["version"] [project.urls] Home = "https://github.com/xarray-contrib/pint-xarray" Documentation = "https://pint-xarray.readthedocs.io/en/stable" [tool.setuptools.packages.find] include = [ "pint_xarray", "pint_xarray.tests", ] [build-system] requires = ["setuptools >= 77", "setuptools_scm >= 8"] build-backend = "setuptools.build_meta" [tool.setuptools_scm] fallback_version = "999" [tool.pytest.ini_options] junit_family = "xunit2" [tool.ruff] target-version = "py311" builtins = ["ellipsis"] exclude = [ ".git", ".eggs", "build", "dist", "__pycache__", ] line-length = 100 [tool.ruff.lint] # E402: module level import not at top of file # E501: line too long - let black worry about that # E731: do not assign a lambda expression, use a def ignore = [ "E402", "E501", "E731", "UP038", ] select = [ "F", # Pyflakes "E", # Pycodestyle "I", # isort "UP", # Pyupgrade "TID", # flake8-tidy-imports "W", ] extend-safe-fixes = [ "TID252", # absolute imports ] fixable = ["I", "TID252"] [tool.ruff.lint.isort] known-first-party = ["pint_xarray"] known-third-party = [ "xarray", ] [tool.ruff.lint.flake8-tidy-imports] # Disallow all relative imports. ban-relative-imports = "all" [tool.coverage.run] source = ["pint_xarray"] branch = true omit = ["pint_xarray/tests/*"] [tool.coverage.report] show_missing = true exclude_lines = ["pragma: no cover", "if TYPE_CHECKING"] pint-xarray-0.6.1/requirements.txt000066400000000000000000000000501516032446100172550ustar00rootroot00000000000000pint>=0.13 numpy>=1.17.1 xarray>=0.15.1