pax_global_header00006660000000000000000000000064146360266270014526gustar00rootroot0000000000000052 comment=14dd1f28f0b2f80956d64b880720bcd408a78645 pint-xarray-0.4/000077500000000000000000000000001463602662700136475ustar00rootroot00000000000000pint-xarray-0.4/.codecov.yml000066400000000000000000000004511463602662700160720ustar00rootroot00000000000000codecov: 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.4/.flake8000066400000000000000000000005711463602662700150250ustar00rootroot00000000000000[flake8] ignore = # E203: whitespace before ':' - doesn't work well with black # 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 # W503: line break before binary operator E203, E402, E501, E731, W503 exclude = .eggs doc builtins = ellipsis pint-xarray-0.4/.github/000077500000000000000000000000001463602662700152075ustar00rootroot00000000000000pint-xarray-0.4/.github/PULL_REQUEST_TEMPLATE.md000066400000000000000000000004541463602662700210130ustar00rootroot00000000000000 - [ ] 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.4/.github/dependabot.yml000066400000000000000000000001661463602662700200420ustar00rootroot00000000000000version: 2 updates: - package-ecosystem: "github-actions" directory: "/" schedule: interval: "weekly" pint-xarray-0.4/.github/workflows/000077500000000000000000000000001463602662700172445ustar00rootroot00000000000000pint-xarray-0.4/.github/workflows/ci-additional.yml000066400000000000000000000025151463602662700224730ustar00rootroot00000000000000name: 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' strategy: fail-fast: false matrix: python-version: ["3.12"] steps: - name: checkout uses: actions/checkout@v4 - name: setup python uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: initialize cache uses: actions/cache@v4 with: path: ~/.cache/pip key: ${{ runner.os }}-pip-py${{ matrix.python-version }}-${{ hashFiles('ci/requirements/**.txt') }} restore-keys: | ${{ runner.os }}-pip-py${{ matrix.python-version }}- - name: upgrade pip run: | python -m pip install --upgrade pip setuptools wheel - name: install dependencies run: | python -m pip install -r ci/requirements.txt - name: install pint-xarray run: | python -m pip install . - name: show versions run: | python -c 'import pint_xarray' - name: run doctests run: | python -m pytest --doctest-modules pint_xarray --ignore pint_xarray/tests pint-xarray-0.4/.github/workflows/ci.yml000066400000000000000000000052211463602662700203620ustar00rootroot00000000000000# 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-latest outputs: triggered: ${{ steps.detect-trigger.outputs.trigger-found }} steps: - uses: actions/checkout@v4 with: fetch-depth: 2 - uses: xarray-contrib/ci-trigger@v1 id: detect-trigger with: keyword: "[skip-ci]" ci: name: ${{ matrix.os }} py${{ matrix.python-version }} runs-on: ${{ matrix.os }} needs: detect-skip-ci-trigger defaults: run: shell: bash -l {0} if: | always() && github.repository == 'xarray-contrib/pint-xarray' && ( github.event_name == 'workflow_dispatch' || needs.detect-skip-ci-trigger.outputs.triggered == 'false' ) strategy: fail-fast: false matrix: python-version: ["3.9", "3.10", "3.11", "3.12"] os: ["ubuntu-latest", "macos-latest", "windows-latest"] steps: - name: checkout the repository uses: actions/checkout@v4 with: # need to fetch all tags to get a correct version fetch-depth: 0 # fetch all branches and tags - name: cache pip uses: actions/cache@v4 with: path: ~/.cache/pip key: pip-py${{ matrix.python-version }} restore-keys: | pip- - name: setup python uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: upgrade pip run: python -m pip install --upgrade pip setuptools wheel - name: install dependencies run: | python -m pip install -r ci/requirements.txt if [[ "${{matrix.python-version}}" == "3.9" ]]; then # remove after the release python -m pip install 'numpy<2.0' fi - name: install pint-xarray run: python -m pip install --no-deps . - name: show versions run: python -m pip list - name: import pint-xarray run: | python -c 'import pint_xarray' - name: run tests if: success() id: status run: | python -m pytest --cov=pint_xarray --cov-report=xml - name: Upload code coverage to Codecov uses: codecov/codecov-action@v4.5.0 with: file: ./coverage.xml flags: unittests env_vars: RUNNER_OS,PYTHON_VERSION name: codecov-umbrella fail_ci_if_error: false pint-xarray-0.4/.github/workflows/nightly.yml000066400000000000000000000051401463602662700214450ustar00rootroot00000000000000# 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-latest outputs: triggered: ${{ steps.detect-trigger.outputs.trigger-found }} steps: - uses: actions/checkout@v4 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') ) ) strategy: fail-fast: false matrix: python-version: ["3.12"] outputs: artifacts_availability: ${{ steps.status.outputs.ARTIFACTS_AVAILABLE }} steps: - name: checkout the repository uses: actions/checkout@v4 with: # need to fetch all tags to get a correct version fetch-depth: 0 # fetch all branches and tags - name: setup python uses: actions/setup-python@v5 with: python-version: ${{ matrix.python-version }} - name: upgrade pip run: python -m pip install --upgrade pip - name: install dependencies run: | python -m pip install -r ci/requirements.txt python -m pip install pytest-reportlog - name: install upstream-dev dependencies run: bash ci/install-upstream-dev.sh - name: install pint-xarray run: python -m pip install . - name: show versions run: python -m pip list - name: run tests if: success() id: status run: | python -m pytest -rf --report-log=pytest-log.jsonl - name: report failures if: | failure() && steps.tests.outcome == 'failure' && github.event_name == 'schedule' uses: xarray-contrib/issue-from-pytest-log@v1 with: log-path: pytest-log.jsonl pint-xarray-0.4/.github/workflows/parse_logs.py000066400000000000000000000046161463602662700217630ustar00rootroot00000000000000# type: ignore import argparse import functools import json import pathlib import textwrap from dataclasses import dataclass from pytest import CollectReport, TestReport @dataclass class SessionStart: pytest_version: str outcome: str = "status" @classmethod def _from_json(cls, json): json_ = json.copy() json_.pop("$report_type") return cls(**json_) @dataclass class SessionFinish: exitstatus: str outcome: str = "status" @classmethod def _from_json(cls, json): json_ = json.copy() json_.pop("$report_type") return cls(**json_) def parse_record(record): report_types = { "TestReport": TestReport, "CollectReport": CollectReport, "SessionStart": SessionStart, "SessionFinish": SessionFinish, } cls = report_types.get(record["$report_type"]) if cls is None: raise ValueError(f"unknown report type: {record['$report_type']}") return cls._from_json(record) @functools.singledispatch def format_summary(report): return f"{report.nodeid}: {report}" @format_summary.register def _(report: TestReport): message = report.longrepr.chain[0][1].message return f"{report.nodeid}: {message}" @format_summary.register def _(report: CollectReport): message = report.longrepr.split("\n")[-1].removeprefix("E").lstrip() return f"{report.nodeid}: {message}" def format_report(reports, py_version): newline = "\n" summaries = newline.join(format_summary(r) for r in reports) message = textwrap.dedent( """\
Python {py_version} Test Summary ``` {summaries} ```
""" ).format(summaries=summaries, py_version=py_version) return message if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("filepath", type=pathlib.Path) args = parser.parse_args() py_version = args.filepath.stem.split("-")[1] print("Parsing logs ...") lines = args.filepath.read_text().splitlines() reports = [parse_record(json.loads(line)) for line in lines] failed = [report for report in reports if report.outcome == "failed"] message = format_report(failed, py_version=py_version) output_file = pathlib.Path("pytest-logs.txt") print(f"Writing output file to: {output_file.absolute()}") output_file.write_text(message) pint-xarray-0.4/.github/workflows/pypi.yaml000066400000000000000000000030471463602662700211150ustar00rootroot00000000000000name: 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@v4 with: fetch-depth: 0 - uses: actions/setup-python@v5 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@v4 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@v4 with: name: releases path: dist - name: Publish package to PyPI uses: pypa/gh-action-pypi-publish@ec4db0b4ddc65acdf4bff5fa45ac92d78b56bdf0 with: verbose: true pint-xarray-0.4/.gitignore000066400000000000000000000034271463602662700156450ustar00rootroot00000000000000# 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/ pint-xarray-0.4/.pep8speaks.yml000066400000000000000000000000611463602662700165300ustar00rootroot00000000000000scanner: diff_only: False linter: flake8 pint-xarray-0.4/.pre-commit-config.yaml000066400000000000000000000020121463602662700201230ustar00rootroot00000000000000ci: autoupdate_schedule: weekly # https://pre-commit.com/ repos: - repo: https://github.com/pre-commit/pre-commit-hooks rev: v4.6.0 hooks: - id: trailing-whitespace - id: end-of-file-fixer - id: check-docstring-first - id: check-yaml # isort should run before black as black sometimes tweaks the isort output - repo: https://github.com/pycqa/isort rev: 5.13.2 hooks: - id: isort # https://github.com/python/black#version-control-integration - repo: https://github.com/psf/black rev: 24.4.2 hooks: - id: black-jupyter - repo: https://github.com/keewis/blackdoc rev: v0.3.9 hooks: - id: blackdoc additional_dependencies: ["black==24.4.2"] - id: blackdoc-autoupdate-black - repo: https://github.com/pycqa/flake8 rev: 7.1.0 hooks: - id: flake8 - repo: https://github.com/kynan/nbstripout rev: 0.7.1 hooks: - id: nbstripout args: [--extra-keys=metadata.kernelspec metadata.language_info.version] pint-xarray-0.4/.readthedocs.yaml000066400000000000000000000007051463602662700171000ustar00rootroot00000000000000version: 2 build: os: ubuntu-22.04 tools: python: "3.11" jobs: post_checkout: - (git --no-pager log --pretty="tformat:%s" -1 | grep -vqF "[skip-rtd]") || exit 183 - git fetch --unshallow || true pre_install: - git update-index --assume-unchanged docs/conf.py python: install: - requirements: docs/requirements.txt - method: pip path: . sphinx: configuration: docs/conf.py fail_on_warning: true pint-xarray-0.4/HOW_TO_RELEASE.rst000066400000000000000000000016301463602662700166000ustar00rootroot00000000000000Release 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.4/LICENSE000066400000000000000000000261351463602662700146630ustar00rootroot00000000000000 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.4/README.md000066400000000000000000000041571463602662700151350ustar00rootroot00000000000000[![CI](https://github.com/xarray-contrib/pint-xarray/workflows/CI/badge.svg?branch=main)](https://github.com/xarray-contrib/pint-xarray/actions?query=branch%3Amain) [![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.4/ci/000077500000000000000000000000001463602662700142425ustar00rootroot00000000000000pint-xarray-0.4/ci/install-upstream-dev.sh000077500000000000000000000005431463602662700206630ustar00rootroot00000000000000#!/usr/bin/env bash python -m pip install \ -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple \ --no-deps \ --pre \ --upgrade \ numpy \ scipy # until `scipy` has released a version compatible with `numpy>=2.0` python -m pip install --upgrade \ git+https://github.com/hgrecco/pint \ git+https://github.com/pydata/xarray pint-xarray-0.4/ci/requirements.txt000066400000000000000000000001361463602662700175260ustar00rootroot00000000000000pint!=0.24.0 numpy<2 scipy dask[array] bottleneck xarray isort black flake8 pytest pytest-cov pint-xarray-0.4/conftest.py000066400000000000000000000011031463602662700160410ustar00rootroot00000000000000import 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.4/docs/000077500000000000000000000000001463602662700145775ustar00rootroot00000000000000pint-xarray-0.4/docs/api.rst000066400000000000000000000033131463602662700161020ustar00rootroot00000000000000API 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 Testing ------- .. autosummary:: :toctree: generated/ pint_xarray.testing.assert_units_equal pint-xarray-0.4/docs/conf.py000066400000000000000000000070341463602662700161020ustar00rootroot00000000000000# 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.4/docs/contributing.rst000066400000000000000000000015461463602662700200460ustar00rootroot00000000000000Contributing ============ ``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.4/docs/conversion.rst000066400000000000000000000013501463602662700175150ustar00rootroot00000000000000.. 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.4/docs/creation.rst000066400000000000000000000077331463602662700171470ustar00rootroot00000000000000.. 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.4/docs/examples.rst000066400000000000000000000001101463602662700171370ustar00rootroot00000000000000Examples ======== .. toctree:: :maxdepth: 1 examples/plotting pint-xarray-0.4/docs/examples/000077500000000000000000000000001463602662700164155ustar00rootroot00000000000000pint-xarray-0.4/docs/examples/plotting.ipynb000066400000000000000000000070711463602662700213250ustar00rootroot00000000000000{ "cells": [ { "cell_type": "markdown", "id": "0", "metadata": {}, "source": [ "# plotting quantified data" ] }, { "cell_type": "code", "execution_count": null, "id": "1", "metadata": {}, "outputs": [], "source": [ "import xarray as xr\n", "\n", "# to be able to read unit attributes following the CF conventions\n", "import cf_xarray.units # must be imported before pint_xarray\n", "import pint_xarray\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.pint.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.4/docs/index.rst000066400000000000000000000015401463602662700164400ustar00rootroot00000000000000pint-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.4/docs/installation.rst000066400000000000000000000007121463602662700200320ustar00rootroot00000000000000Installation ------------ 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.4/docs/requirements.txt000066400000000000000000000002511463602662700200610ustar00rootroot00000000000000pint>=0.21 xarray>=2022.06.0 pooch netCDF4 cf-xarray>=0.6 sphinx sphinx_rtd_theme>=1.0 ipython ipykernel jupyter_client nbsphinx matplotlib sphinx-autosummary-accessors pint-xarray-0.4/docs/terminology.rst000066400000000000000000000003751463602662700177060ustar00rootroot00000000000000Terminology =========== .. glossary:: unit-like A `pint`_ unit definition, as accepted by :py:class:`pint.Unit`. May be either a :py:class:`str` or a :py:class:`pint.Unit` instance. .. _pint: https://pint.readthedocs.io/en/stable pint-xarray-0.4/docs/whats-new.rst000066400000000000000000000135431463602662700172540ustar00rootroot00000000000000.. currentmodule:: xarray What's new ========== 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.4/licenses/000077500000000000000000000000001463602662700154545ustar00rootroot00000000000000pint-xarray-0.4/licenses/XARRAY_LICENSE000066400000000000000000000241121463602662700175470ustar00rootroot00000000000000Copyright 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.4/pint_xarray/000077500000000000000000000000001463602662700162075ustar00rootroot00000000000000pint-xarray-0.4/pint_xarray/__init__.py000066400000000000000000000010321463602662700203140ustar00rootroot00000000000000from importlib.metadata import version import pint from . import accessors, formatting, testing # noqa: F401 from .accessors import default_registry as unit_registry from .accessors import setup_registry try: __version__ = version("pint-xarray") except Exception: # 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", ] pint-xarray-0.4/pint_xarray/accessors.py000066400000000000000000001577731463602662700205720ustar00rootroot00000000000000# 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 . import conversion from .conversion import no_unit_values from .errors import format_error_message _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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } return converted.loc[stripped_indexers] 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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } return converted.loc[stripped_indexers] 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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in converted.items() } 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``). .. warning:: As units in dimension coordinates are not supported until ``xarray`` changes the way it implements indexes, these units will be set as attributes. .. 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 ValueError(format_error_message(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 ValueError(format_error_message(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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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) # index stripped_indexers = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } indexed = converted.reindex( stripped_indexers, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) return indexed 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_unit_attributes(other) # TODO: handle tolerance # TODO: handle fill_value converted = conversion.convert_units(self.da, indexer_units) return converted.reindex_like( other, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) 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:: ``kwargs`` is 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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } interpolated = stripped.interp( stripped_indexers, method=method, assume_sorted=False, kwargs=None, ) 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:: ``kwargs`` is passed unmodified to ``DataArray.interp`` See Also -------- xarray.Dataset.pint.interp_like xarray.DataArray.pint.interp xarray.DataArray.interp_like """ indexer_units = conversion.extract_unit_attributes(other) converted = conversion.convert_units(self.da, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) interpolated = stripped.interp_like( 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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } indexed = converted.sel( stripped_indexers, method=method, tolerance=tolerance, drop=drop, ) return indexed @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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in converted_indexers.items() } 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``). .. warning:: As units in dimension coordinates are not supported until ``xarray`` changes the way it implements indexes, these units will be set as attributes. .. 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 ValueError(format_error_message(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 ValueError(format_error_message(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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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) # index stripped_indexers = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } indexed = converted.reindex( stripped_indexers, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) return indexed 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_unit_attributes(other) # TODO: handle tolerance # TODO: handle fill_value converted = conversion.convert_units(self.ds, indexer_units) return converted.reindex_like( other, method=method, tolerance=tolerance, copy=copy, fill_value=fill_value, ) 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:: ``kwargs`` is passed unmodified to ``Dataset.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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } interpolated = stripped.interp( stripped_indexers, method=method, assume_sorted=False, kwargs=None, ) 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:: ``kwargs`` is passed unmodified to ``Dataset.interp`` See Also -------- xarray.DataArray.pint.interp_like xarray.Dataset.pint.interp xarray.Dataset.interp_like """ indexer_units = conversion.extract_unit_attributes(other) converted = conversion.convert_units(self.ds, indexer_units) units = conversion.extract_units(converted) stripped = conversion.strip_units(converted) interpolated = stripped.interp_like( 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 = { name: conversion.extract_indexer_units(indexer) for name, indexer in indexers.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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in indexers.items() } indexed = converted.sel( stripped_indexers, method=method, tolerance=tolerance, drop=drop, ) return indexed @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 = { name: conversion.strip_indexer_units(indexer) for name, indexer in converted_indexers.items() } 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.4/pint_xarray/compat.py000066400000000000000000000007211463602662700200440ustar00rootroot00000000000000import 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.4/pint_xarray/conversion.py000066400000000000000000000303241463602662700207500ustar00rootroot00000000000000import itertools import re import pint from xarray import DataArray, Dataset, IndexVariable, Variable from .compat import call_on_dataset from .errors import format_error_message 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, coords, attrs): data_vars = {name: var for name, var in variables.items() if name not in coords} coords = {name: var for name, var in variables.items() if name in coords} return Dataset(data_vars=data_vars, coords=coords, attrs=attrs) def attach_units_dataset(obj, units): attached = {} rejected_vars = {} for name, var in obj.variables.items(): unit = units.get(name) try: converted = attach_units_variable(var, unit) attached[name] = converted except ValueError as e: rejected_vars[name] = (unit, e) if rejected_vars: raise ValueError(rejected_vars) return dataset_from_variables(attached, obj._coord_names, 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 ValueError(format_error_message(rejected_vars, "attach")) from e 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_dataset(obj, units): converted = {} failed = {} for name, var in obj.variables.items(): unit = units.get(name) try: converted[name] = convert_units_variable(var, unit) except (ValueError, pint.errors.PintTypeError) as e: failed[name] = e if failed: raise ValueError(failed) return dataset_from_variables(converted, obj._coord_names, 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 ValueError(format_error_message(failed, "convert")) from e 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()} return dataset_from_variables(variables, obj._coord_names, 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 ValueError(format_error_message(invalid, "convert_indexers")) return converted def extract_indexer_units(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) def strip_indexer_units(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) pint-xarray-0.4/pint_xarray/errors.py000066400000000000000000000024631463602662700201020ustar00rootroot00000000000000def format_error_message(mapping, op): sep = "\n " if len(mapping) == 1 else "\n -- " if op == "attach": message = "Cannot attach units:" message = sep.join( [message] + [ f"cannot attach units to variable {key!r}: {unit} (reason: {str(e)})" for key, (unit, e) in mapping.items() ] ) elif op == "parse": message = "Cannot parse units:" message = sep.join( [message] + [ f"invalid units for variable {key!r}: {unit} ({type}) (reason: {str(e)})" for key, (unit, type, e) in mapping.items() ] ) elif op == "convert": message = "Cannot convert variables:" message = sep.join( [message] + [ f"incompatible units for variable {key!r}: {error}" for key, error in mapping.items() ] ) elif op == "convert_indexers": message = "Cannot convert indexers:" message = sep.join( [message] + [ f"incompatible units for indexer for {key!r}: {error}" for key, error in mapping.items() ] ) else: raise ValueError("invalid op") return message pint-xarray-0.4/pint_xarray/formatting.py000066400000000000000000000162451463602662700207430ustar00rootroot00000000000000from 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.4/pint_xarray/testing.py000066400000000000000000000016151463602662700202410ustar00rootroot00000000000000from . 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.4/pint_xarray/tests/000077500000000000000000000000001463602662700173515ustar00rootroot00000000000000pint-xarray-0.4/pint_xarray/tests/__init__.py000066400000000000000000000000001463602662700214500ustar00rootroot00000000000000pint-xarray-0.4/pint_xarray/tests/test_accessors.py000066400000000000000000002112501463602662700227500ustar00rootroot00000000000000import 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 .. import accessors, conversion from .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 unit_registry = UnitRegistry(force_ndarray=True) 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.raises(ValueError, match="already has units"): 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.raises(ValueError, match="already has units"): 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.raises(ValueError): arr.pint.quantify({"x": "s"}) def test_dimension_coordinate_array_already_quantified_same_units(self): ds = xr.Dataset(coords={"x": ("x", [10], {"units": unit_registry.Unit("m")})}) 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.raises(ValueError, match=str(da.name)): 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.raises( ValueError, match=rf"{da.name}: {da.attrs['units']} \(attribute\)" ): 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.raises(ValueError, match="already has units"): 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.raises(ValueError, match="already has units"): ds.pint.quantify({"x": "s"}) def test_error_on_nonsense_units(self, example_unitless_ds): ds = example_unitless_ds with pytest.raises(ValueError): 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.raises( ValueError, match=rf"'users': {ds.users.attrs['units']} \(attribute\)" ): ds.pint.quantify() def test_error_indicates_problematic_variable(self, example_unitless_ds): ds = example_unitless_ds with pytest.raises(ValueError, match="'users'"): 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={"x": ("x", [2, 4], {"units": Unit("s")})}), {"x": "ms"}, xr.Dataset(coords={"x": ("x", [2000, 4000], {"units": Unit("ms")})}), None, id="Dataset-compatible units-dims", ), pytest.param( xr.Dataset(coords={"x": ("x", [2, 4], {"units": Unit("s")})}), {"x": "mm"}, None, ValueError, id="Dataset-incompatible units-dims", ), 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={"x": ("x", [2, 4], {"units": Unit("s")})} ), {"x": "ms"}, xr.DataArray( [0, 1], dims="x", coords={"x": ("x", [2000, 4000], {"units": Unit("ms")})}, ), None, id="DataArray-compatible units-dims", ), pytest.param( xr.DataArray( [0, 1], dims="x", coords={"x": ("x", [2, 4], {"units": Unit("s")})} ), {"x": "mm"}, None, ValueError, id="DataArray-incompatible units-dims", ), ), ) 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.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): if error is not None: with pytest.raises(error): obj.pint.sel(indexers) else: 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): if error is not None: with pytest.raises(error): obj.pint.loc[indexers] else: 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", "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, 50], "dm"), "y": Quantity([0, 120, 240], "s")}, xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 120, 240], {"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([0, 1, 3, 5], "m"), "y": Quantity([0, 2, 4], "min")}, xr.Dataset( { "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2, 4], {"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, ValueError, 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, 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], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 240], {"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, 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], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2], {"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, ValueError, id="DataArray-incompatible units", ), ), ) def test_reindex(obj, indexers, expected, error): if error is not None: with pytest.raises(error): obj.pint.reindex(indexers) else: actual = obj.pint.reindex(indexers) assert_units_equal(actual, expected) assert_identical(actual, expected) @pytest.mark.parametrize( ["obj", "other", "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")}), } ), xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), } ), xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 120, 240], {"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")}), } ), xr.Dataset( { "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), } ), xr.Dataset( { "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2, 4], {"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")}), } ), xr.Dataset( { "x": ("x", [1, 3], {"units": unit_registry.Unit("s")}), "y": ("y", [1], {"units": unit_registry.Unit("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], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), } ), xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={ "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 240], {"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")}), }, ), xr.Dataset( { "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), } ), xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), coords={ "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2], {"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")}), }, ), xr.Dataset( { "x": ("x", [10, 30], {"units": unit_registry.Unit("s")}), "y": ("y", [60], {"units": unit_registry.Unit("m")}), } ), None, ValueError, id="DataArray-incompatible units", ), ), ) def test_reindex_like(obj, other, expected, error): if error is not None: with pytest.raises(error): obj.pint.reindex_like(other) else: actual = obj.pint.reindex_like(other) assert_units_equal(actual, expected) assert_identical(actual, expected) @requires_scipy @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, 50], "dm"), "y": Quantity([0, 120, 240], "s")}, xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 120, 240], {"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([0, 1, 3, 5], "m"), "y": Quantity([0, 2, 4], "min")}, xr.Dataset( { "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2, 4], {"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, ValueError, id="Dataset-incompatible units", ), pytest.param( xr.Dataset( { "a": (("x", "y"), Quantity([[0, 1], [2, 3], [4, 5]], "kg")), "x": [10, 20, 30], "y": [60, 120], } ), { "x": [15, 25], "y": [75, 105], }, xr.Dataset( { "a": (("x", "y"), Quantity([[1.25, 1.75], [3.25, 3.75]], "kg")), "x": [15, 25], "y": [75, 105], } ), None, id="Dataset-data 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, 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], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 240], {"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, 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], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2], {"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, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( Quantity([[0, 1], [2, 3], [4, 5]], "kg"), dims=("x", "y"), coords={ "x": [10, 20, 30], "y": [60, 120], }, ), { "x": [15, 25], "y": [75, 105], }, xr.DataArray( Quantity([[1.25, 1.75], [3.25, 3.75]], "kg"), dims=("x", "y"), coords={ "x": [15, 25], "y": [75, 105], }, ), None, id="DataArray-data units", ), ), ) def test_interp(obj, indexers, expected, error): if error is not None: with pytest.raises(error): obj.pint.interp(indexers) else: actual = obj.pint.interp(indexers) assert_units_equal(actual, expected) assert_identical(actual, expected) @requires_scipy @pytest.mark.parametrize( ["obj", "other", "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")}), } ), xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 120, 240], {"units": unit_registry.Unit("s")}), } ), xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 120, 240], {"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")}), } ), xr.Dataset( { "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2, 4], {"units": unit_registry.Unit("min")}), } ), xr.Dataset( { "x": ("x", [0, 1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2, 4], {"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")}), } ), xr.Dataset( { "x": ("x", [1, 3], {"units": unit_registry.Unit("s")}), "y": ("y", [1], {"units": unit_registry.Unit("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], {"units": unit_registry.Unit("dm")}), "y": ("y", [60, 120], {"units": unit_registry.Unit("s")}), }, ), xr.Dataset( { "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), } ), xr.DataArray( [[np.nan, np.nan], [np.nan, np.nan], [np.nan, np.nan]], dims=("x", "y"), coords={ "x": ("x", [10, 30, 50], {"units": unit_registry.Unit("dm")}), "y": ("y", [0, 240], {"units": unit_registry.Unit("s")}), }, ), None, id="DataArray-identical units", ), pytest.param( xr.Dataset( { "a": (("x", "y"), Quantity([[0, 1], [2, 3], [4, 5]], "kg")), "x": [10, 20, 30], "y": [60, 120], } ), xr.Dataset( { "x": [15, 25], "y": [75, 105], } ), xr.Dataset( { "a": (("x", "y"), Quantity([[1.25, 1.75], [3.25, 3.75]], "kg")), "x": [15, 25], "y": [75, 105], } ), None, id="Dataset-data 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")}), }, ), xr.Dataset( { "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2], {"units": unit_registry.Unit("min")}), } ), xr.DataArray( [[np.nan, 1], [np.nan, 5], [np.nan, np.nan]], dims=("x", "y"), coords={ "x": ("x", [1, 3, 5], {"units": unit_registry.Unit("m")}), "y": ("y", [0, 2], {"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")}), }, ), xr.Dataset( { "x": ("x", [10, 30], {"units": unit_registry.Unit("s")}), "y": ("y", [60], {"units": unit_registry.Unit("m")}), } ), None, ValueError, id="DataArray-incompatible units", ), pytest.param( xr.DataArray( Quantity([[0, 1], [2, 3], [4, 5]], "kg"), dims=("x", "y"), coords={ "x": [10, 20, 30], "y": [60, 120], }, ), xr.Dataset( { "x": [15, 25], "y": [75, 105], } ), xr.DataArray( Quantity([[1.25, 1.75], [3.25, 3.75]], "kg"), dims=("x", "y"), coords={ "x": [15, 25], "y": [75, 105], }, ), None, id="DataArray-data units", ), ), ) def test_interp_like(obj, other, expected, error): if error is not None: with pytest.raises(error): obj.pint.interp_like(other) else: 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.4/pint_xarray/tests/test_conversion.py000066400000000000000000000630171463602662700231560ustar00rootroot00000000000000import numpy as np import pint import pytest from xarray import DataArray, Dataset, Variable from pint_xarray import conversion from .utils import ( assert_array_equal, assert_array_units_equal, assert_identical, assert_indexer_equal, assert_indexer_units_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_u = to_quantity(u, units.get("u")) units_x = units.get("x") obj = Dataset({"a": ("x", a), "b": ("x", b)}, coords={"u": ("x", u), "x": x}) expected = Dataset( {"a": ("x", q_a), "b": ("x", q_b)}, coords={"u": ("x", q_u), "x": x}, ) if units_x is not None: expected.x.attrs["units"] = units_x if type == "DataArray": obj = obj["a"] expected = expected["a"] actual = conversion.attach_units(obj, units) assert_identical(actual, expected) @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", "match"], ( pytest.param("none", {}, None, None, id="none-no units"), pytest.param( "none", {"a": Unit("g"), "b": Unit("Pa"), "u": Unit("ms"), "x": Unit("mm")}, ValueError, "(?s)Cannot convert variables:.+'u'", id="none-with units", ), pytest.param("data", {}, None, None, id="data-no units"), pytest.param( "data", {"a": Unit("g"), "b": Unit("Pa")}, None, None, id="data-compatible units", ), pytest.param( "data", {"a": Unit("s"), "b": Unit("m")}, ValueError, "(?s)Cannot convert variables:.+'a'", id="data-incompatible units", ), pytest.param( "dims", {}, None, None, id="dims-no units", ), pytest.param( "dims", {"x": Unit("mm")}, None, None, id="dims-compatible units", ), pytest.param( "dims", {"x": Unit("ms")}, ValueError, "(?s)Cannot convert variables:.+'x'", id="dims-incompatible units", ), pytest.param( "coords", {}, None, None, id="coords-no units", ), pytest.param( "coords", {"u": Unit("ms")}, None, None, id="coords-compatible units", ), pytest.param( "coords", {"u": Unit("mm")}, ValueError, "(?s)Cannot convert variables:.+'u'", id="coords-incompatible units", ), ), ) def test_convert_units(self, type, variant, units, error, match): 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")) obj = Dataset( { "a": ("x", q_a), "b": ("x", q_b), }, coords={ "u": ("x", q_u), "x": ("x", x, {"units": original_units.get("x")}), }, ) if type == "DataArray": obj = obj["a"] if error is not None: with pytest.raises(error, match=match): 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 = strip_quantity(convert_quantity(q_x, units.get("x"))) expected = Dataset( { "a": ("x", expected_a), "b": ("x", expected_b), }, coords={ "u": ("x", expected_u), "x": ( "x", expected_x, {"units": units.get("x", original_units.get("x"))}, ), }, ) if type == "DataArray": expected = expected["a"] actual = conversion.convert_units(obj, units) assert conversion.extract_units(actual) == conversion.extract_units(expected) assert_identical(expected, actual) @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) obj = Dataset( { "a": ("x", to_quantity(a, units.get("a"))), "b": ("x", to_quantity(b, units.get("b"))), }, coords={ "u": ("x", to_quantity(u, units.get("u"))), "x": ("x", x, {"units": units.get("x")}), }, ) 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=[0, 4, 3] * unit_registry.m, coords={"u": ("x", [2, 3, 4] * unit_registry.s)}, ), {None: None, "u": None}, id="DataArray", ), pytest.param( Dataset( data_vars={ "a": ("x", [3, 2, 5] * unit_registry.Pa), "b": ("x", [0, 2, -1] * unit_registry.kg), }, coords={"u": ("x", [2, 3, 4] * unit_registry.s)}, ), {"a": None, "b": None, "u": 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", "match"], ( pytest.param( {"x": 1}, {"x": None}, {"x": 1}, None, None, id="scalar-no units" ), pytest.param( {"x": 1}, {"x": "dimensionless"}, None, ValueError, "(?s)Cannot convert indexers:.+'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, ValueError, "(?s)Cannot convert indexers:.+'x'", id="slice-incompatible units", ), pytest.param( {"x": slice(1000, Quantity(2000, "ms"))}, {"x": Unit("s")}, None, ValueError, "(?s)Cannot convert indexers:.+'x'", id="slice-incompatible units-mixed", ), ), ) def test_convert_indexer_units(self, indexers, units, expected, error, match): if error is not None: with pytest.raises(error, match=match): conversion.convert_indexer_units(indexers, units) else: actual = conversion.convert_indexer_units(indexers, units) assert_indexer_equal(actual["x"], expected["x"]) assert_indexer_units_equal(actual["x"], expected["x"]) @pytest.mark.parametrize( ["indexer", "expected"], ( pytest.param(1, None, id="scalar-no units"), pytest.param(Quantity(1, "m"), Unit("m"), id="scalar-units"), pytest.param(np.array([1, 2]), None, id="array-no units"), pytest.param(Quantity([1, 2], "s"), Unit("s"), id="array-units"), pytest.param(Variable("x", [1, 2]), None, id="Variable-no units"), pytest.param( Variable("x", Quantity([1, 2], "m")), Unit("m"), id="Variable-units" ), pytest.param(DataArray([1, 2], dims="x"), None, id="DataArray-no units"), pytest.param( DataArray(Quantity([1, 2], "s"), dims="x"), Unit("s"), id="DataArray-units", ), pytest.param(slice(None), None, id="empty slice-no units"), pytest.param(slice(1, None), None, id="slice-no units"), pytest.param( slice(Quantity(1, "m"), Quantity(2, "m")), Unit("m"), id="slice-identical units", ), pytest.param( slice(Quantity(1, "m"), Quantity(2000, "mm")), Unit("m"), id="slice-compatible units", ), pytest.param( slice(Quantity(1, "m"), Quantity(2, "ms")), ValueError, id="slice-incompatible units", ), pytest.param( slice(1, Quantity(2, "ms")), ValueError, id="slice-incompatible units-mixed", ), pytest.param( slice(1, Quantity(2, "rad")), Unit("rad"), id="slice-incompatible units-mixed-dimensionless", ), ), ) def test_extract_indexer_units(self, indexer, expected): if expected is not None and not isinstance(expected, Unit): with pytest.raises(expected): conversion.extract_indexer_units(indexer) else: actual = conversion.extract_indexer_units(indexer) assert actual == expected @pytest.mark.parametrize( ["indexer", "expected"], ( pytest.param(1, 1, id="scalar-no units"), pytest.param(Quantity(1, "m"), 1, id="scalar-units"), pytest.param(np.array([1, 2]), np.array([1, 2]), id="array-no units"), pytest.param(Quantity([1, 2], "s"), np.array([1, 2]), id="array-units"), pytest.param( Variable("x", [1, 2]), Variable("x", [1, 2]), id="Variable-no units" ), pytest.param( Variable("x", Quantity([1, 2], "m")), Variable("x", [1, 2]), id="Variable-units", ), pytest.param( DataArray([1, 2], dims="x"), DataArray([1, 2], dims="x"), id="DataArray-no units", ), pytest.param( DataArray(Quantity([1, 2], "s"), dims="x"), DataArray([1, 2], dims="x"), id="DataArray-units", ), pytest.param(slice(None), slice(None), id="empty slice-no units"), pytest.param(slice(1, None), slice(1, None), id="slice-no units"), pytest.param( slice(Quantity(1, "m"), Quantity(2, "m")), slice(1, 2), id="slice-units", ), ), ) def test_strip_indexer_units(self, indexer, expected): actual = conversion.strip_indexer_units(indexer) if isinstance(indexer, DataArray): assert_identical(actual, expected) else: assert_array_equal(actual, expected) pint-xarray-0.4/pint_xarray/tests/test_formatting.py000066400000000000000000000010511463602662700231310ustar00rootroot00000000000000import 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.4/pint_xarray/tests/test_testing.py000066400000000000000000000031301463602662700224340ustar00rootroot00000000000000import pint import pytest import xarray as xr from pint_xarray import testing unit_registry = pint.UnitRegistry(force_ndarray_like=True) @pytest.mark.parametrize( ("a", "b", "error"), ( pytest.param( xr.DataArray(attrs={"units": "K"}), xr.DataArray(attrs={"units": "K"}), None, id="equal attrs", ), pytest.param( xr.DataArray(attrs={"units": "m"}), xr.DataArray(attrs={"units": "K"}), AssertionError, id="different attrs", ), pytest.param( xr.DataArray([10, 20] * unit_registry.K), xr.DataArray([50, 80] * unit_registry.K), None, id="equal units", ), pytest.param( xr.DataArray([10, 20] * unit_registry.K), xr.DataArray([50, 80] * unit_registry.dimensionless), AssertionError, id="different units", ), pytest.param( xr.Dataset({"a": ("x", [0, 10], {"units": "K"})}), xr.Dataset({"a": ("x", [20, 40], {"units": "K"})}), None, id="matching variables", ), pytest.param( xr.Dataset({"a": ("x", [0, 10], {"units": "K"})}), xr.Dataset({"b": ("x", [20, 40], {"units": "K"})}), AssertionError, id="mismatching variables", ), ), ) def test_assert_units_equal(a, b, error): if error is not None: with pytest.raises(error): testing.assert_units_equal(a, b) return testing.assert_units_equal(a, b) pint-xarray-0.4/pint_xarray/tests/utils.py000066400000000000000000000054331463602662700210700ustar00rootroot00000000000000import re from contextlib import contextmanager import numpy as np import pytest from pint import Quantity from xarray import DataArray, Variable from xarray.testing import assert_equal, assert_identical # noqa: F401 from ..conversion import ( array_strip_units, extract_indexer_units, strip_units, strip_units_variable, ) from ..testing import assert_units_equal # noqa: F401 def importorskip(name): try: __import__(name) has_name = True except ImportError: has_name = False return has_name, pytest.mark.skipif(not has_name, reason=f"{name} is not available") has_dask_array, requires_dask_array = importorskip("dask.array") has_scipy, requires_scipy = importorskip("scipy") has_bottleneck, requires_bottleneck = importorskip("bottleneck") @contextmanager def raises_regex(error, pattern): __tracebackhide__ = True with pytest.raises(error) as excinfo: yield message = str(excinfo.value) if not re.search(pattern, message): raise AssertionError( f"exception {excinfo.value!r} did not match pattern {pattern!r}" ) def assert_array_units_equal(a, b): __tracebackhide__ = True units_a = getattr(a, "units", None) units_b = getattr(b, "units", None) assert units_a == units_b def assert_array_equal(a, b): __tracebackhide__ = True a_ = getattr(a, "magnitude", a) b_ = getattr(b, "magnitude", b) np.testing.assert_array_equal(a_, b_) def assert_slice_equal(a, b): attrs = ("start", "stop", "step") values_a = tuple(getattr(a, name) for name in attrs) values_b = tuple(getattr(b, name) for name in attrs) stripped_a = tuple(array_strip_units(v) for v in values_a) stripped_b = tuple(array_strip_units(v) for v in values_b) assert ( stripped_a == stripped_b ), f"different values: {stripped_a!r} ←→ {stripped_b!r}" def assert_indexer_equal(a, b): __tracebackhide__ = True assert type(a) is type(b) if isinstance(a, slice): assert_slice_equal(a, b) elif isinstance(a, DataArray): stripped_a = strip_units(a) stripped_b = strip_units(b) assert_equal(stripped_a, stripped_b) elif isinstance(a, Variable): stripped_a = strip_units_variable(a) stripped_b = strip_units_variable(b) assert_equal(stripped_a, stripped_b) elif isinstance(a, (Quantity, np.ndarray)): assert_array_equal(a, b) else: a_ = array_strip_units(a) b_ = array_strip_units(b) assert a_ == b_, f"different values: {a_!r} ←→ {b_!r}" def assert_indexer_units_equal(a, b): __tracebackhide__ = True units_a = extract_indexer_units(a) units_b = extract_indexer_units(b) assert units_a == units_b, f"different units: {units_a!r} ←→ {units_b!r}" pint-xarray-0.4/pyproject.toml000066400000000000000000000026301463602662700165640ustar00rootroot00000000000000[project] name = "pint-xarray" authors = [ {name = "Tom Nicholas", email = "tomnicholas1@googlemail.com"} ] description = "Physical units interface to xarray using Pint" license = {text = "Apache-2"} readme = "README.md" classifiers = [ "Development Status :: 3 - Alpha", "Environment :: Console", "Intended Audience :: Science/Research", "License :: OSI Approved :: Apache Software License", "Operating System :: OS Independent", "Programming Language :: Python", "Programming Language :: Python :: 3.9", "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Topic :: Scientific/Engineering", ] requires-python = ">=3.9" dependencies = [ "numpy >= 1.23", "xarray >= 2022.06.0", "pint >= 0.21", ] 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 >= 64", "setuptools_scm >= 7.0"] build-backend = "setuptools.build_meta" [tool.setuptools_scm] fallback_version = "999" [tool.pytest.ini_options] junit_family = "xunit2" [tool.isort] profile = "black" skip_gitignore = "true" force_to_top = "true" default_section = "THIRDPARTY" known_first_party = "pint_xarray" pint-xarray-0.4/requirements.txt000066400000000000000000000000501463602662700171260ustar00rootroot00000000000000pint>=0.13 numpy>=1.17.1 xarray>=0.15.1