diff --git a/.bumpversion.cfg b/.bumpversion.cfg new file mode 100644 index 0000000..2c87cda --- /dev/null +++ b/.bumpversion.cfg @@ -0,0 +1,12 @@ +[bumpversion] +current_version = 2024.5.1 +commit = True +tag = True + +[bumpversion:file:pyproject.toml] +search = version = "{current_version}" +replace = version = "{new_version}" + +[bumpversion:file:earth2grid/__init__.py] +search = __version__ = '{current_version}' +replace = __version__ = '{new_version}' diff --git a/.editorconfig b/.editorconfig new file mode 100644 index 0000000..814d7b2 --- /dev/null +++ b/.editorconfig @@ -0,0 +1,24 @@ +# http://editorconfig.org + +root = true + +[*] +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true +charset = utf-8 +end_of_line = lf + +[*.bat] +indent_style = tab +end_of_line = crlf + +[LICENSE] +insert_final_newline = false + +[Makefile] +indent_style = tab + +[*.{yml, yaml}] +indent_size = 2 diff --git a/.github/ISSUE_TEMPLATE.md b/.github/ISSUE_TEMPLATE.md new file mode 100644 index 0000000..8b0fb3b --- /dev/null +++ b/.github/ISSUE_TEMPLATE.md @@ -0,0 +1,15 @@ +* Earth2 Grid Utilities version: +* Python version: +* Operating System: + +### Description + +Describe what you were trying to get done. +Tell us what happened, what went wrong, and what you expected to happen. + +### What I Did + +``` +Paste the command(s) you ran and the output. +If there was a crash, please include the traceback here. +``` diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml new file mode 100644 index 0000000..ad5345b --- /dev/null +++ b/.github/workflows/dev.yml @@ -0,0 +1,50 @@ +# This is a basic workflow to help you get started with Actions + +name: dev workflow + +# Controls when the action will run. +on: + # Triggers the workflow on push or pull request events but only for the master branch + push: + branches: [ master, main ] + pull_request: + branches: [ master, main ] + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +# A workflow run is made up of one or more jobs that can run sequentially or in parallel +jobs: + # This workflow contains a single job called "test" + test: + # The type of runner that the job will run on + strategy: + matrix: + python-versions: [3.6, 3.7, 3.8, 3.9] + os: [ubuntu-18.04, macos-latest, windows-latest] + runs-on: ${{ matrix.os }} + + # Steps represent a sequence of tasks that will be executed as part of the job + steps: + # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-versions }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install poetry tox tox-gh-actions + + - name: test with tox + run: + tox + + - name: list files + run: ls -l . + + - uses: codecov/codecov-action@v1 + with: + fail_ci_if_error: true + files: coverage.xml diff --git a/.github/workflows/preview.yml b/.github/workflows/preview.yml new file mode 100644 index 0000000..953a60e --- /dev/null +++ b/.github/workflows/preview.yml @@ -0,0 +1,50 @@ +# This is a basic workflow to help you get started with Actions + +name: stage & preview workflow + +# Controls when the action will run. +on: + # Triggers the workflow on push or pull request events but only for the master branch + push: + branches: [ master, main ] + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +# A workflow run is made up of one or more jobs that can run sequentially or in parallel +jobs: + publish_dev_build: + runs-on: ubuntu-latest + + strategy: + matrix: + python-versions: [ 3.8 ] + + steps: + - uses: actions/checkout@v2 + - uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-versions }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install poetry tox tox-gh-actions + + - name: test with tox + run: + tox + + - name: Build wheels and source tarball + run: | + poetry version $(poetry version --short)-dev.$GITHUB_RUN_NUMBER + poetry version --short + poetry build + + - name: publish to Test PyPI + uses: pypa/gh-action-pypi-publish@master + with: + user: __token__ + password: ${{ secrets.TEST_PYPI_API_TOKEN}} + repository_url: https://test.pypi.org/legacy/ + skip_existing: true diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml new file mode 100644 index 0000000..6060909 --- /dev/null +++ b/.github/workflows/release.yml @@ -0,0 +1,89 @@ +# Publish package on main branch if it's tagged with 'v*' + +name: release & publish workflow + +# Controls when the action will run. +on: + # Triggers the workflow on push events but only for the master branch + push: + tags: + - 'v*' + + # Allows you to run this workflow manually from the Actions tab + workflow_dispatch: + +# A workflow run is made up of one or more jobs that can run sequentially or in parallel +jobs: + # This workflow contains a single job called "release" + release: + name: Create Release + runs-on: ubuntu-20.04 + + strategy: + matrix: + python-versions: [3.8] + + # Steps represent a sequence of tasks that will be executed as part of the job + steps: + - name: Get version from tag + id: tag_name + run: | + echo ::set-output name=current_version::${GITHUB_REF#refs/tags/v} + shell: bash + + # Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it + - uses: actions/checkout@v2 + + - name: Get Changelog Entry + id: changelog_reader + uses: mindsers/changelog-reader-action@v2 + with: + validation_depth: 10 + version: ${{ steps.tag_name.outputs.current_version }} + path: ./CHANGELOG.md + + - uses: actions/setup-python@v2 + with: + python-version: ${{ matrix.python-versions }} + + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install poetry + + - name: build documentation + run: | + poetry install -E doc + poetry run mkdocs build + + - name: publish documentation + uses: peaceiris/actions-gh-pages@v3 + with: + personal_token: ${{ secrets.PERSONAL_TOKEN }} + publish_dir: ./site + + - name: Build wheels and source tarball + run: >- + poetry build + + - name: show temporary files + run: >- + ls -l + + - name: create github release + id: create_release + uses: softprops/action-gh-release@v1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + with: + body: ${{ steps.changelog_reader.outputs.changes }} + files: dist/*.whl + draft: false + prerelease: false + + - name: publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + user: __token__ + password: ${{ secrets.PYPI_API_TOKEN }} + skip_existing: true diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..e1f3e35 --- /dev/null +++ b/.gitignore @@ -0,0 +1,122 @@ +# macOS +.DS_Store +.AppleDouble +.LSOverride + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +env/ +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +*.egg-info/ +.installed.cfg +*.egg + +# 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/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# pyenv +.python-version + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# dotenv +.env + +# virtualenv +.venv +venv/ +ENV/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ + +# IDE settings +.vscode/ +.idea/ + +# mkdocs build dir +site/ +test_grid_visualize.png +*.png +*.jpg +*.jpeg +public/ + +a.out +*.o diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml new file mode 100644 index 0000000..d528f52 --- /dev/null +++ b/.gitlab-ci.yml @@ -0,0 +1,12 @@ +# The Docker image that will be used to build your app +image: ubuntu +pages: + script: "true" + artifacts: + paths: + # The folder that contains the files to be exposed at the Page URL + - public + rules: + # This ensures that only pushes to the default branch will trigger + # a pages deploy + - if: $CI_COMMIT_REF_NAME == "pages" diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..3f744b5 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,34 @@ +repos: + - repo: https://github.com/Lucas-C/pre-commit-hooks + rev: v1.1.9 + hooks: + - id: forbid-crlf + - id: remove-crlf + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v3.4.0 + hooks: + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-merge-conflict + - id: check-yaml + args: [ --unsafe ] + - repo: https://github.com/pre-commit/mirrors-isort + rev: v5.8.0 + hooks: + - id: isort + args: [ "--filter-files" ] + - repo: https://github.com/psf/black + rev: 22.10.0 + hooks: + - id: black + additional_dependencies: ["click==8.0.4"] + - repo: https://github.com/pycqa/flake8 + rev: 7.0.0 + hooks: + - id: flake8 + exclude: conf.py + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.8.0 + hooks: + - id: mypy + exclude: tests/ diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..a3343c9 --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,5 @@ +# Changelog + +## v2024.5.2 + +- First publicly available release diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..b72f172 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,76 @@ +# Contributing + +Contributions are welcome, and they are greatly appreciated! Every little bit +helps, and credit will always be given. + +## Developer Certificate of Origin + +``` +Developer Certificate of Origin +Version 1.1 + +Copyright (C) 2004, 2006 The Linux Foundation and its contributors. + +Everyone is permitted to copy and distribute verbatim copies of this +license document, but changing it is not allowed. + + +Developer's Certificate of Origin 1.1 + +By making a contribution to this project, I certify that: + +(a) The contribution was created in whole or in part by me and I + have the right to submit it under the open source license + indicated in the file; or + +(b) The contribution is based upon previous work that, to the best + of my knowledge, is covered under an appropriate open source + license and I have the right under that license to submit that + work with modifications, whether created in whole or in part + by me, under the same open source license (unless I am + permitted to submit under a different license), as indicated + in the file; or + +(c) The contribution was provided directly to me by some other + person who certified (a), (b) or (c) and I have not modified + it. + +(d) I understand and agree that this project and the contribution + are public and that a record of the contribution (including all + personal information I submit with it, including my sign-off) is + maintained indefinitely and may be redistributed consistent with + this project or the open source license(s) involved. +``` + +## Getting started + +1. Install the pre-commit hooks: `pre-commit install` +1. Checkout a feature branch `git checkout -b +1. When writing code, group changes logically into commits. Messy commits are + usually a result of excessive multitasking. If you work on one thing at a + time, then your commits will be cleaner. + 1. Name each commit "[imperative verb] [subject]" (e.g. "add earth2grid.some_modules"). Make sure it fits onto one line. + 1. Please provide context in the commit message. Start by explaining the + previous state of the code and why this needed changing. Focus on motivating + rather than explaining the changeset. + 1. run the test suite: `make unittest` + 1. run the docs: `make docs` +1. push the code to the repo `git push -u origin your-branch-name` and open an MR +1. ask for one reviewer. + +## Tips + + +## Deploying + +A reminder for the maintainers on how to deploy. +Make sure all your changes are committed (including an entry in CHANGELOG.md). +Then run: + +``` +$ poetry run bump2version patch # possible: major / minor / patch +$ git push +$ git push --tags +``` + +GitHub Actions will then deploy to PyPI if tests pass. diff --git a/LICENSE.txt b/LICENSE.txt new file mode 100644 index 0000000..f61f492 --- /dev/null +++ b/LICENSE.txt @@ -0,0 +1,201 @@ +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 2024 NVIDIA Corporation + + 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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..8950fb7 --- /dev/null +++ b/README.md @@ -0,0 +1,75 @@ +# Earth2 Grid Utilities + + + + +[![pypi](https://img.shields.io/pypi/v/earth2-grid.svg)](https://pypi.org/project/earth2-grid/) +[![python](https://img.shields.io/pypi/pyversions/earth2-grid.svg)](https://pypi.org/project/earth2-grid/) +[![Build Status](https://github.com/waynerv/earth2-grid/actions/workflows/dev.yml/badge.svg)](https://github.com/waynerv/earth2-grid/actions/workflows/dev.yml) +[![codecov](https://codecov.io/gh/waynerv/earth2-grid/branch/main/graphs/badge.svg)](https://codecov.io/github/waynerv/earth2-grid) + + + + +Utilities for working geographic data defined on various grids. + +Features: +- regridding + +Grids currently supported: +- regular lat lon +- HealPIX + +Under construction: +- exporting meshes to visualization software (e.g. VTK) +- neural network primitives for different grids: + - convolutional layers + - up/downsampling +- staggered grids + + +* Documentation: +* GitHub: +* PyPI: +* Free software: BSD-3-Clause + +## Install + + +``` +# Install the CUDA healpix padding library +git clone https://gitlab-master.nvidia.com/tkurth/healpixpad-pytorch +git clone ssh://git@gitlab-master.nvidia.com:12051/earth-2/earth2-grid.git +pip install --no-build-isolation healpixpad-pytorch earth2-grid +``` + +## Example + +``` +>>> import earth2grid +... # level is the resolution +... level = 6 +... hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.PixelOrder.XY) +... src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64) +... z_torch = torch.as_tensor(z) +... z_torch = torch.as_tensor(z) +>>> regrid = earth2grid.get_regridder(src, hpx) +>>> z_hpx = regrid(z_torch) +>>> z_hpx.shape +torch.Size([49152]) +>>> nside = 2**level +... reshaped = z_hpx.reshape(12, nside, nside) +... lat_r = hpx.lat.reshape(12, nside, nside) +... lon_r = hpx.lon.reshape(12, nside, nside) +>>> reshaped.shape +torch.Size([12, 64, 64]) +``` + + +## Features + +* TODO + +## Credits + +This package was created with [Cookiecutter](https://github.com/audreyr/cookiecutter) and the [waynerv/cookiecutter-pypackage](https://github.com/waynerv/cookiecutter-pypackage) project template. diff --git a/docs/.gitignore b/docs/.gitignore new file mode 100644 index 0000000..015219a --- /dev/null +++ b/docs/.gitignore @@ -0,0 +1,2 @@ +auto_examples/ +sg_execution_times.rst diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000..6570176 --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = . +BUILDDIR = _build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @PYTHONPATH=$(shell pwd):$(PYTHONPATH) $(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/api.rst b/docs/api.rst new file mode 100644 index 0000000..a272fc2 --- /dev/null +++ b/docs/api.rst @@ -0,0 +1,25 @@ +API +=== + +Grids +----- + +.. autoclass:: earth2grid.healpix.Grid + :members: + :show-inheritance: + +.. autoclass:: earth2grid.latlon.LatLonGrid + :members: + :show-inheritance: + +.. autofunction:: earth2grid.latlon.equiangular_lat_lon_grid + +Regridding +---------- + +.. autofunction:: earth2grid.get_regridder + +Other utilities +--------------- + +.. autofunction:: earth2grid.healpix.pad diff --git a/docs/changelog.md b/docs/changelog.md new file mode 120000 index 0000000..04c99a5 --- /dev/null +++ b/docs/changelog.md @@ -0,0 +1 @@ +../CHANGELOG.md \ No newline at end of file diff --git a/docs/conf.py b/docs/conf.py new file mode 100644 index 0000000..3196d06 --- /dev/null +++ b/docs/conf.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +# Configuration file for the Sphinx documentation builder. +# +# For the full list of built-in configuration values, see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Project information ----------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information + +project = 'Earth2-Grid' +copyright = '2024, NVIDIA Research' +author = 'NVIDIA Research' + +# -- General configuration --------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration + +extensions = [ + 'myst_parser', + 'sphinx.ext.autodoc', + 'sphinx.ext.napoleon', + 'sphinx_gallery.gen_gallery', +] + +templates_path = ['_templates'] +exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store'] + +# sphinx + +import os + +# pyvista +import pyvista + +pyvista.BUILDING_GALLERY = True + +import pyvista +from image_scraper import PNGScraper +from pyvista.core.errors import PyVistaDeprecationWarning +from pyvista.core.utilities.docs import linkcode_resolve, pv_html_page_context # noqa: F401 +from pyvista.plotting.utilities.sphinx_gallery import DynamicScraper + +# Manage errors +pyvista.set_error_output_file("errors.txt") +# Ensure that offscreen rendering is used for docs generation +pyvista.OFF_SCREEN = True # Not necessary - simply an insurance policy +# Preferred plotting style for documentation +pyvista.set_plot_theme("document") +pyvista.global_theme.window_size = [1024, 768] +pyvista.global_theme.font.size = 22 +pyvista.global_theme.font.label_size = 22 +pyvista.global_theme.font.title_size = 22 +pyvista.global_theme.return_cpos = False +pyvista.set_jupyter_backend(None) +# Save figures in specified directory +pyvista.FIGURE_PATH = os.path.join(os.path.abspath("./images/"), "auto-generated/") +if not os.path.exists(pyvista.FIGURE_PATH): + os.makedirs(pyvista.FIGURE_PATH) + +# necessary when building the sphinx gallery +pyvista.BUILDING_GALLERY = True +os.environ['PYVISTA_BUILDING_GALLERY'] = 'true' + + +# sphinx gallery +sphinx_gallery_conf = { + 'examples_dirs': '../examples/sphinx_gallery/', # path to your example scripts + 'gallery_dirs': 'auto_examples', # path to where to save gallery generated output + 'filename_pattern': '', + "image_scrapers": (DynamicScraper(), "matplotlib", PNGScraper()), +} + +# -- Options for HTML output ------------------------------------------------- +# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output +html_static_path = ['_static'] diff --git a/docs/contributing.md b/docs/contributing.md new file mode 120000 index 0000000..44fcc63 --- /dev/null +++ b/docs/contributing.md @@ -0,0 +1 @@ +../CONTRIBUTING.md \ No newline at end of file diff --git a/docs/image_scraper.py b/docs/image_scraper.py new file mode 100644 index 0000000..6ca0227 --- /dev/null +++ b/docs/image_scraper.py @@ -0,0 +1,59 @@ +# Copyright (c) 2015, Óscar Nájera +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# * Redistributions of source code must retain the above copyright notice, this +# list of conditions and the following disclaimer. +# +# * Redistributions in binary form must reproduce the above copyright notice, +# this list of conditions and the following disclaimer in the documentation +# and/or other materials provided with the distribution. +# +# * Neither the name of sphinx-gallery nor the names of its +# contributors may be used to endorse or promote products derived from +# this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# from https://sphinx-gallery.github.io/stable/advanced.html#example-2-detecting-image-files-on-disk +import os +import shutil +from glob import glob + +from sphinx_gallery.scrapers import figure_rst + + +class PNGScraper(object): + def __init__(self): + self.seen = set() + + def __repr__(self): + return 'PNGScraper' + + def __call__(self, block, block_vars, gallery_conf): + # Find all PNG files in the directory of this example. + path_current_example = os.path.dirname(block_vars['src_file']) + pngs = sorted(glob(os.path.join(path_current_example, '*.png'))) + + # Iterate through PNGs, copy them to the sphinx-gallery output directory + image_names = list() + image_path_iterator = block_vars['image_path_iterator'] + for png in pngs: + if png not in self.seen: + self.seen |= set(png) + this_image_path = image_path_iterator.next() + image_names.append(this_image_path) + shutil.move(png, this_image_path) + # Use the `figure_rst` helper function to generate reST for image files + return figure_rst(image_names, gallery_conf['src_dir']) diff --git a/docs/img/image.jpg b/docs/img/image.jpg new file mode 100644 index 0000000..3efeef3 Binary files /dev/null and b/docs/img/image.jpg differ diff --git a/docs/index.rst b/docs/index.rst new file mode 100644 index 0000000..a76e3e9 --- /dev/null +++ b/docs/index.rst @@ -0,0 +1,28 @@ +.. Earth2-Grid documentation master file, created by + sphinx-quickstart on Wed Feb 7 08:46:40 2024. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +Welcome to Earth2-Grid's documentation! +======================================= + +.. toctree:: + :maxdepth: 1 + :caption: Contents: + + installation.md + usage.md + api + changelog.md + contributing.md + auto_examples/index + + + + +Indices and tables +================== + +* :ref:`genindex` +* :ref:`modindex` +* :ref:`search` diff --git a/docs/installation.md b/docs/installation.md new file mode 100644 index 0000000..d1eda9d --- /dev/null +++ b/docs/installation.md @@ -0,0 +1,43 @@ +# Installation + +## Stable release + +To install Earth2 Grid Utilities, run this command in your +terminal: + +``` console +$ pip install earth2-grid +``` + +This is the preferred method to install Earth2 Grid Utilities, as it will always install the most recent stable release. + +If you don't have [pip][] installed, this [Python installation guide][] +can guide you through the process. + +## From source + +The source for Earth2 Grid Utilities can be downloaded from +the [Github repo][]. + +You can either clone the public repository: + +``` console +$ git clone git://github.com/waynerv/earth2-grid +``` + +Or download the [tarball][]: + +``` console +$ curl -OJL https://github.com/waynerv/earth2-grid/tarball/master +``` + +Once you have a copy of the source, you can install it with: + +``` console +$ pip install . +``` + + [pip]: https://pip.pypa.io + [Python installation guide]: http://docs.python-guide.org/en/latest/starting/installation/ + [Github repo]: https://github.com/%7B%7B%20cookiecutter.github_username%20%7D%7D/%7B%7B%20cookiecutter.project_slug%20%7D%7D + [tarball]: https://github.com/%7B%7B%20cookiecutter.github_username%20%7D%7D/%7B%7B%20cookiecutter.project_slug%20%7D%7D/tarball/master diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 0000000..954237b --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=. +set BUILDDIR=_build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/docs/push_docs.sh b/docs/push_docs.sh new file mode 100755 index 0000000..251efc4 --- /dev/null +++ b/docs/push_docs.sh @@ -0,0 +1,14 @@ +#!/bin/bash + +originalBranch=$(git rev-parse --abbrev-ref HEAD) +# Function to switch back to the original branch +function switchBack { + echo "Switching back to $originalBranch" + git checkout $originalBranch +} +git checkout pages +trap switchBack EXIT +rm -r public/ && cp -r docs/_build/html public +git add -f public/ +git commit -a -m "update doc from $ref" +git push origin pages diff --git a/docs/usage.md b/docs/usage.md new file mode 100644 index 0000000..27b3493 --- /dev/null +++ b/docs/usage.md @@ -0,0 +1,18 @@ +# Usage + +To use Earth2 Grid Utilities in a project + +``` +import earth2grid +# level is the resolution +level = 6 +hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.PixelOrder.XY) +src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64) +z_torch = torch.as_tensor(z) +z_torch = torch.as_tensor(z) +regrid = earth2grid.get_regridder(src, hpx) +z_hpx = regrid(z_torch) +z_hpx.shape +nside = 2**level +reshaped = z_hpx.reshape(12, nside, nside) +``` diff --git a/earth2grid/__init__.py b/earth2grid/__init__.py new file mode 100644 index 0000000..3143b70 --- /dev/null +++ b/earth2grid/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +from earth2grid import base, healpix, latlon +from earth2grid._regrid import get_regridder + +__all__ = ["base", "healpix", "latlon", "get_regridder"] diff --git a/earth2grid/_regrid.py b/earth2grid/_regrid.py new file mode 100644 index 0000000..4230026 --- /dev/null +++ b/earth2grid/_regrid.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import einops +import netCDF4 as nc +import torch + +from earth2grid import base, healpix +from earth2grid.latlon import LatLonGrid + + +class TempestRegridder(torch.nn.Module): + def __init__(self, file_path): + super().__init__() + dataset = nc.Dataset(file_path) + self.lat = dataset["latc_b"][:] + self.lon = dataset["lonc_b"][:] + + i = dataset["row"][:] - 1 + j = dataset["col"][:] - 1 + M = dataset["S"][:] + + i = i.data + j = j.data + M = M.data + + self.M = torch.sparse_coo_tensor((i, j), M, [max(i) + 1, max(j) + 1]).float() + + def to(self, device): + self.M = self.M.to(device) + return self + + def forward(self, x): + xr = einops.rearrange(x, "b c x y -> b c (x y)") + yr = xr @ self.M.T + y = einops.rearrange(yr, "b c (x y) -> b c x y", x=self.lat.size, y=self.lon.size) + return y + + +class Identity(torch.nn.Module): + def forward(self, x): + return x + + +def get_regridder(src: base.Grid, dest: base.Grid) -> torch.nn.Module: + """Get a regridder from `src` to `dest`""" + if src == dest: + return Identity() + elif isinstance(src, LatLonGrid) and isinstance(dest, LatLonGrid): + return src.get_bilinear_regridder_to(dest.lat, dest.lon) + elif isinstance(src, LatLonGrid) and isinstance(dest, healpix.Grid): + return src.get_bilinear_regridder_to(dest.lat, dest.lon) + elif isinstance(src, healpix.Grid): + return src.get_bilinear_regridder_to(dest.lat, dest.lon) + elif isinstance(dest, healpix.Grid): + return src.get_healpix_regridder(dest) # type: ignore + + raise ValueError(src, dest, "not supported.") diff --git a/earth2grid/base.py b/earth2grid/base.py new file mode 100644 index 0000000..2f43451 --- /dev/null +++ b/earth2grid/base.py @@ -0,0 +1,47 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +from typing import Protocol + +import numpy as np + + +class Grid(Protocol): + """lat and lon should be broadcastable arrays""" + + @property + def lat(self): + pass + + @property + def lon(self): + pass + + @property + def shape(self) -> tuple[int, ...]: + pass + + def get_bilinear_regridder_to(self, lat: np.ndarray, lon: np.ndarray): + """Return a regridder from `self` to lat/lon. + + Args: + lat, lon: broadcastable arrays for the lat/lon + """ + raise NotImplementedError() + + def visualize(self, data): + pass + + def to_pyvista(self): + pass diff --git a/earth2grid/csrc/healpix_bare_wrapper.cpp b/earth2grid/csrc/healpix_bare_wrapper.cpp new file mode 100644 index 0000000..2cb0537 --- /dev/null +++ b/earth2grid/csrc/healpix_bare_wrapper.cpp @@ -0,0 +1,228 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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. + */ +// This NVIDIA code +#include +#include +#include "healpix_bare.c" +#include "interpolation.h" +#include + + +#define FILLNA(x) std::isnan(x) ? 0.0 : x + +int return_1(void) { + return 1; +} + +template +auto wrap(Func func) { + return [func](int nside, torch::Tensor input) { + auto output = torch::empty_like(input); + auto accessor = input.accessor(); + auto out_accessor = output.accessor(); + + for (int64_t i = 0; i < input.size(0); ++i) { + out_accessor[i] = func(nside, accessor[i]); // Specify nside appropriately + } + + return output; + }; +} + + +template +auto wrap_2hpd(Func func) { + + return [func](int nside, torch::Tensor input){ + auto options = torch::TensorOptions().dtype(torch::kInt64); + auto x = torch::empty_like(input, options); + auto y = torch::empty_like(input, options); + auto f = torch::empty_like(input, options.dtype(torch::kInt32)); + + auto output_options = torch::TensorOptions().dtype(torch::kInt64); + auto output = torch::empty({input.size(0), 3}, output_options); + auto out_accessor = output.accessor(); + auto accessor = input.accessor(); + + for (int64_t i = 0; i < input.size(0); ++i) { + auto hpd = func(nside, accessor[i]); + out_accessor[i][0] = hpd.f; + out_accessor[i][1] = hpd.y; + out_accessor[i][2] = hpd.x; + } + + return output; + }; +} + +torch::Tensor hpd2loc_wrapper(int nside, torch::Tensor input) { + auto accessor = input.accessor(); + + auto output_options = torch::TensorOptions().dtype(torch::kDouble); + auto output = torch::empty({input.size(0), 3}, output_options); + auto out_accessor = output.accessor(); + + int64_t x, y, f; + + for (int64_t i = 0; i < input.size(0); ++i) { + f = accessor[i][0]; + y = accessor[i][1]; + x = accessor[i][2]; + + t_hpd hpd {x, y, f}; + tloc loc = hpd2loc(nside, hpd); + out_accessor[i][0] = loc.z; + out_accessor[i][1] = loc.s; + out_accessor[i][2] = loc.phi; + } + + return output; +} + +torch::Tensor hpc2loc_wrapper(torch::Tensor x, torch::Tensor y, torch::Tensor f) { + auto accessor_f = f.accessor(); + auto accessor_x = x.accessor(); + auto accessor_y = y.accessor(); + + auto output_options = torch::TensorOptions().dtype(torch::kDouble); + auto output = torch::empty({f.size(0), 3}, output_options); + auto out_accessor = output.accessor(); + + + for (int64_t i = 0; i < x.size(0); ++i) { + t_hpc hpc {accessor_x[i], accessor_y[i], accessor_f[i]}; + tloc loc = hpc2loc(hpc); + out_accessor[i][0] = loc.z; + out_accessor[i][1] = loc.s; + out_accessor[i][2] = loc.phi; + } + + return output; +} + + + +torch::Tensor corners(int nside, torch::Tensor pix, bool nest) { + auto accessor = pix.accessor(); + + auto output_options = torch::TensorOptions().dtype(torch::kDouble); + auto output = torch::empty({pix.size(0), 3, 4}, output_options); + auto out_accessor = output.accessor(); + + + t_hpd hpd; + double n = nside; + + for (int64_t i = 0; i < pix.size(0); ++i) { + if (nest) { + hpd = nest2hpd(nside, accessor[i]); + } else { + hpd = ring2hpd(nside, accessor[i]); + } + + t_hpc hpc; + int offset = 0; + t_vec vec; + + hpc.x = static_cast(hpd.x) / n; + hpc.y = static_cast(hpd.y) / n; + hpc.f = hpd.f; + vec = loc2vec(hpc2loc(hpc)); + out_accessor[i][0][offset] = FILLNA(vec.x); + out_accessor[i][1][offset] = FILLNA(vec.y); + out_accessor[i][2][offset] = vec.z; + offset++; + + hpc.x = static_cast(hpd.x + 1) / n; + hpc.y = static_cast(hpd.y) / n; + vec = loc2vec(hpc2loc(hpc)); + out_accessor[i][0][offset] = FILLNA(vec.x); + out_accessor[i][1][offset] = FILLNA(vec.y); + out_accessor[i][2][offset] = vec.z; + offset++; + + hpc.x = static_cast(hpd.x + 1) / n; + hpc.y = static_cast(hpd.y + 1) / n; + vec = loc2vec(hpc2loc(hpc)); + out_accessor[i][0][offset] = FILLNA(vec.x); + out_accessor[i][1][offset] = FILLNA(vec.y); + out_accessor[i][2][offset] = vec.z; + offset++; + + hpc.x = static_cast(hpd.x) / n; + hpc.y = static_cast(hpd.y + 1) / n; + vec = loc2vec(hpc2loc(hpc)); + out_accessor[i][0][offset] = FILLNA(vec.x); + out_accessor[i][1][offset] = FILLNA(vec.y); + out_accessor[i][2][offset] = vec.z; + offset++; + } + return output; +} + +// these are the minimal routines +// nest2hpd +// ring2hpd +// hpd2loc +// loc2ang + +std::vector get_interp_weights(int nside, torch::Tensor lon, torch::Tensor lat) { + + // setup outputs + auto weight_options = torch::TensorOptions().dtype(torch::kDouble); + auto weight = torch::empty({4, lon.size(0)}, weight_options); + + auto pix_options = torch::TensorOptions().dtype(torch::kLong); + auto pix = torch::empty({4, lon.size(0)}, pix_options); + + + auto pix_a = pix.accessor(); + auto weight_a = weight.accessor(); + + auto lon_a = lon.accessor(); + auto lat_a = lat.accessor(); + const bool nest = false; + + { + // output information + std::array pix_i; + std::array wgt_i; + for (int64_t i = 0; i < lon.size(0); ++i) { + t_ang ptg = latlon2ang(lat_a[i], lon_a[i]); + interpolation_weights(ptg, pix_i, wgt_i, nside); + // TODO flip i and j to get better cache performance + for (int j= 0; j < 4; ++j){ + pix_a[j][i] = pix_i[j]; + weight_a[j][i] = wgt_i[j]; + } + }; + } + return std::vector{pix, weight}; +} + + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("return_1", &return_1, "First implementation."); + m.def("ring2nest", wrap(ring2nest), "Element-wise ring2nest conversion"); + m.def("nest2ring", wrap(nest2ring), "Element-wise nest2ring conversion"); + m.def("nest2hpd", wrap_2hpd(nest2hpd), "hpd is f, y ,x"); + m.def("ring2hpd", wrap_2hpd(ring2hpd), "hpd is f, y ,x"); + m.def("hpd2loc", &hpd2loc_wrapper, "loc is in z, s, phi"); + m.def("hpc2loc", &hpc2loc_wrapper, "hpc2loc(x, y, f) -> z, s, phi"); + m.def("corners", &corners, ""); + m.def("get_interp_weights", &get_interp_weights, ""); +}; diff --git a/earth2grid/csrc/interpolation.h b/earth2grid/csrc/interpolation.h new file mode 100644 index 0000000..55aa48f --- /dev/null +++ b/earth2grid/csrc/interpolation.h @@ -0,0 +1,167 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + * + * 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. + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "healpix_bare.h" + +void hpx_info(int64_t nside, int64_t &npface, int64_t &ncap, int64_t &npix, + double &fact1, double &fact2) { + npix = nside2npix(nside); + npface = npix / 12; + ncap = (npface - nside) << 1; + fact2 = 4. / npix; + fact1 = (nside << 1) * fact2; +} + +int64_t ring_above(double z, int64_t nside) { + double az = std::abs(z); + if (az <= 2. / 3.) + return int64_t(nside * (2 - 1.5 * z)); + int64_t iring = int64_t(nside * std::sqrt(3 * (1 - az))); + return (z > 0) ? iring : 4 * nside - iring - 1; +} + +void ring_info(int64_t ring, int64_t &startpix, int64_t &ringpix, double &theta, + bool &shifted, int64_t nside) { + int64_t northring = (ring > 2 * nside) ? 4 * nside - ring : ring; + int64_t npface, ncap, npix; + double fact1, fact2; + hpx_info(nside, npface, ncap, npix, fact1, fact2); + if (northring < nside) { + double tmp = northring * northring * fact2; + double costheta = 1 - tmp; + double sintheta = std::sqrt(tmp * (2 - tmp)); + theta = std::atan2(sintheta, costheta); + ringpix = 4 * northring; + shifted = true; + startpix = 2 * northring * (northring - 1); + } else { + theta = acos((2 * nside - northring) * fact1); + ringpix = 4 * nside; + shifted = ((northring - nside) & 1) == 0; + startpix = ncap + (northring - nside) * ringpix; + } + if (northring != ring) // southern hemisphere + { + theta = std::numbers::pi - theta; + startpix = npix - startpix - ringpix; + } +} + +// TODO (asubramaniam): switch from std::array to double* pointers to avoid +// copies in the batched case +template +void interpolation_weights(const t_ang &ptg, std::array &pix, + std::array &wgt, int64_t nside) { + assert((ptg.theta >= 0) && (ptg.theta <= std::numbers::pi)); + double z = std::cos(ptg.theta); + int64_t npix = nside2npix(nside); + + // Do everything in ring ordering first + // Can convert indices to nest ordering at the end if needed + int64_t ir1 = ring_above(z, nside); + int64_t ir2 = ir1 + 1; + double theta1, theta2, w1, tmp, dphi; + int64_t sp, nr; + bool shift; + int64_t i1, i2; + if (ir1 > 0) { + ring_info(ir1, sp, nr, theta1, shift, nside); + dphi = 2. * std::numbers::pi / nr; + tmp = (ptg.phi / dphi - .5 * shift); + i1 = (tmp < 0) ? int64_t(tmp) - 1 : int64_t(tmp); + w1 = (ptg.phi - (i1 + .5 * shift) * dphi) / dphi; + if (i1 < 0) { + i1 += nr; + } + i2 = i1 + 1; + if (i2 >= nr) { + i2 -= nr; + } + pix[0] = sp + i1; + pix[1] = sp + i2; + wgt[0] = 1 - w1; + wgt[1] = w1; + } + if (ir2 < (4 * nside)) { + ring_info(ir2, sp, nr, theta2, shift, nside); + dphi = 2. * std::numbers::pi / nr; + tmp = (ptg.phi / dphi - .5 * shift); + i1 = (tmp < 0) ? int64_t(tmp) - 1 : int64_t(tmp); + w1 = (ptg.phi - (i1 + .5 * shift) * dphi) / dphi; + if (i1 < 0) + i1 += nr; + i2 = i1 + 1; + if (i2 >= nr) + i2 -= nr; + pix[2] = sp + i1; + pix[3] = sp + i2; + wgt[2] = 1 - w1; + wgt[3] = w1; + } + + if (ir1 == 0) { + double wtheta = ptg.theta / theta2; + wgt[2] *= wtheta; + wgt[3] *= wtheta; + double fac = (1 - wtheta) * 0.25; + wgt[0] = fac; + wgt[1] = fac; + wgt[2] += fac; + wgt[3] += fac; + pix[0] = (pix[2] + 2) & 3; + pix[1] = (pix[3] + 2) & 3; + } else if (ir2 == 4 * nside) { + double wtheta = (ptg.theta - theta1) / (std::numbers::pi - theta1); + wgt[0] *= (1 - wtheta); + wgt[1] *= (1 - wtheta); + double fac = wtheta * 0.25; + wgt[0] += fac; + wgt[1] += fac; + wgt[2] = fac; + wgt[3] = fac; + pix[2] = ((pix[0] + 2) & 3) + npix - 4; + pix[3] = ((pix[1] + 2) & 3) + npix - 4; + } else { + double wtheta = (ptg.theta - theta1) / (theta2 - theta1); + wgt[0] *= (1 - wtheta); + wgt[1] *= (1 - wtheta); + wgt[2] *= wtheta; + wgt[3] *= wtheta; + } + + // Convert indices from ring to nest format if needed + if (NEST) + for (size_t m = 0; m < pix.size(); ++m) + pix[m] = ring2nest(pix[m], nside); +} + +double degrees2radians(double theta) { return theta * std::numbers::pi / 180.; } + +t_ang latlon2ang(double lat, double lon) { + double theta = std::numbers::pi / 2. - degrees2radians(lat); + double phi = degrees2radians(lon); + t_ang ang = {theta, phi}; + return ang; +} diff --git a/earth2grid/healpix.py b/earth2grid/healpix.py new file mode 100644 index 0000000..6c757be --- /dev/null +++ b/earth2grid/healpix.py @@ -0,0 +1,424 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +""" + +From this notebook: https://colab.research.google.com/drive/1MzTyeNFiy-7RNY6UtGKsmDavX5dk6epU + + +Healpy has two indexing conventions NEST and RING. But for convolutions we want +2D array indexing in row or column major order. Here are some vectorized +routines `nest2xy` and `x2nest` for going in between these conventions. The +previous code shared by Dale used string computations to handle these +operations, which was probably quite slow. Here we use vectorized bit-shifting. + +## XY orientation + +For array-like indexing can have a different origin and orientation. For +example, the default is the origin is S and the data arr[f, y, x] follows the +right hand rule. In other words, (x + 1, y) being counterclockwise from (x, y) +when looking down on the face. + +""" + +import math +from dataclasses import dataclass +from enum import Enum +from typing import Union + +import einops +import numpy as np +import torch + +from earth2grid import healpix_bare + +try: + import pyvista as pv +except ImportError: + pv = None + +from earth2grid import base +from earth2grid.third_party.zephyr.healpix import healpix_pad + +try: + import healpixpad +except ImportError: + healpixpad = None + +__all__ = ["pad", "PixelOrder", "XY", "Compass", "Grid", "HEALPIX_PAD_XY", "conv2d"] + + +def pad(x: torch.Tensor, padding: int) -> torch.Tensor: + """ + Pad each face consistently with its according neighbors in the HEALPix + + Args: + x: The input tensor of shape [N, F, H, W] + padding: the amount of padding + + Returns: + The padded tensor with shape [N, F, H+2*padding, W+2*padding] + + Examples: + + Ths example show to pad data described by a :py:class:`Grid` object. + + >>> grid = Grid(level=4, pixel_order=PixelOrder.RING) + >>> lon = torch.from_numpy(grid.lon) + >>> faces = grid.reorder(HEALPIX_PAD_XY, lon) + >>> faces = faces.view(1, 12, grid._nside(), grid._nside()) + >>> faces.shape + torch.Size([1, 12, 16, 16]) + >>> padded = pad(faces, padding=1) + >>> padded.shape + torch.Size([1, 12, 18, 18]) + + """ + if healpixpad is None or x.device.type != 'cuda': + return healpix_pad(x, padding) + else: + return healpixpad.HEALPixPadFunction.apply(x.unsqueeze(2), padding).squeeze(2) + + +class PixelOrder(Enum): + RING = 0 + NEST = 1 + + +class Compass(Enum): + """Cardinal directions in counter clockwise order""" + + S = 0 + E = 1 + N = 2 + W = 3 + + +@dataclass(frozen=True) +class XY: + """ + Assumes + - i = n * n * f + n * y + x + - the origin (x,y)=(0,0) is South + - if clockwise=False follows the hand rule: + + Space + | + | + | / y + | / + |/______ x + + (Thumb points towards Space, index finger towards x, middle finger towards y) + """ + + origin: Compass = Compass.S + clockwise: bool = False + + +PixelOrderT = Union[PixelOrder, XY] + +HEALPIX_PAD_XY = XY(origin=Compass.N, clockwise=True) + + +def _convert_xyindex(nside: int, src: XY, dest: XY, i): + if src.clockwise != dest.clockwise: + i = _flip_xy(nside, i) + + rotations = dest.origin.value - src.origin.value + i = _rotate_index(nside=nside, rotations=-rotations if dest.clockwise else rotations, i=i) + return i + + +class ApplyWeights(torch.nn.Module): + def __init__(self, pix: torch.Tensor, weight: torch.Tensor): + super().__init__() + + # the first dim is the 4 point stencil + n, *self.shape = pix.shape + + pix = pix.view(n, -1).T + weight = weight.view(n, -1).T + + self.register_buffer("index", pix) + self.register_buffer("weight", weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + *shape, npix = x.shape + x = x.view(-1, npix).T + interpolated = torch.nn.functional.embedding_bag(self.index, x, per_sample_weights=self.weight, mode="sum").T + return interpolated.view(shape + self.shape) + + +@dataclass +class Grid(base.Grid): + """A Healpix Grid + + Attrs: + level: 2^level = nside + pixel_order: the ordering convection of the data + """ + + level: int + pixel_order: PixelOrderT = PixelOrder.RING + + def __post_init__(self): + if self.level > ZOOM_LEVELS: + raise ValueError(f"`level` must be less than or equal to {ZOOM_LEVELS}") + + def _nside(self): + return 2**self.level + + def _npix(self): + return self._nside() ** 2 * 12 + + def _nest_ipix(self): + """convert to nested index number""" + i = torch.arange(self._npix()) + if isinstance(self.pixel_order, XY): + i_xy = _convert_xyindex(nside=self._nside(), src=self.pixel_order, dest=XY(), i=i) + i = xy2nest(self._nside(), i_xy) + elif self.pixel_order == PixelOrder.RING: + i = healpix_bare.ring2nest(self._nside(), i) + elif self.pixel_order == PixelOrder.NEST: + pass + else: + raise ValueError(self.pixel_order) + return i.numpy() + + def _nest2me(self, ipix: np.ndarray) -> np.ndarray: + """return the index in my PIXELORDER corresponding to ipix in NEST ordering""" + if isinstance(self.pixel_order, XY): + i_xy = nest2xy(self._nside(), ipix) + i_me = _convert_xyindex(nside=self._nside(), src=XY(), dest=self.pixel_order, i=i_xy) + elif self.pixel_order == PixelOrder.RING: + ipix_t = torch.from_numpy(ipix) + i_me = healpix_bare.nest2ring(self._nside(), ipix_t).numpy() + elif self.pixel_order == PixelOrder.NEST: + i_me = ipix + return i_me + + @property + def lat(self): + ipix = torch.from_numpy(self._nest_ipix()) + _, lat = healpix_bare.pix2ang(self._nside(), ipix, lonlat=True, nest=True) + return lat.numpy() + + @property + def lon(self): + ipix = torch.from_numpy(self._nest_ipix()) + lon, _ = healpix_bare.pix2ang(self._nside(), ipix, lonlat=True, nest=True) + return lon.numpy() + + @property + def shape(self) -> tuple[int, ...]: + return (self._npix(),) + + def visualize(self, map): + raise NotImplementedError() + + def to_pyvista(self): + if pv is None: + raise ImportError("Need to install pyvista") + + # Make grid + nside = 2**self.level + pix = self._nest_ipix() + points = healpix_bare.corners(nside, torch.from_numpy(pix), True).numpy() + out = einops.rearrange(points, "n d s -> (n s) d") + unique_points, inverse = np.unique(out, return_inverse=True, axis=0) + assert unique_points.ndim == 2 + assert unique_points.shape[1] == 3 + inverse = einops.rearrange(inverse, "(n s) -> n s", n=pix.size) + n, s = inverse.shape + cells = np.ones_like(inverse, shape=(n, s + 1)) + cells[:, 0] = s + cells[:, 1:] = inverse + celltypes = np.full(shape=(n,), fill_value=pv.CellType.QUAD) + grid = pv.UnstructuredGrid(cells, celltypes, unique_points) + return grid + + def get_bilinear_regridder_to(self, lat: np.ndarray, lon: np.ndarray): + """Get regridder to the specified lat and lon points""" + lat, lon = np.broadcast_arrays(lat, lon) + i_ring, weights = healpix_bare.get_interp_weights(self._nside(), torch.tensor(lon), torch.tensor(lat)) + i_nest = healpix_bare.ring2nest(self._nside(), i_ring.ravel()) + i_me = torch.from_numpy(self._nest2me(i_nest.numpy())).view(i_ring.shape) + return ApplyWeights(i_me, weights) + + def approximate_grid_length_meters(self): + return approx_grid_length_meters(self._nside()) + + def reorder(self, order: PixelOrderT, x: torch.Tensor) -> torch.Tensor: + """Rorder the pixels of ``x`` to have ``order``""" + output_grid = Grid(level=self.level, pixel_order=order) + i_nest = output_grid._nest_ipix() + i_me = self._nest2me(i_nest) + return x[..., i_me] + + def get_healpix_regridder(self, dest: "Grid"): + if self.level != dest.level: + return self.get_bilinear_regridder_to(dest.lat, dest.lon) + + def regridder(x: torch.Tensor) -> torch.Tensor: + return self.reorder(dest.pixel_order, x) + + return regridder + + def to_image(self, x: torch.Tensor, fill_value=torch.nan) -> torch.Tensor: + """Use the 45 degree rotated grid pixelation + i points to SE, j point to NE + """ + grid = [[6, 9, -1, -1, -1], [1, 5, 8, -1, -1], [-1, 0, 4, 11, -1], [-1, -1, 3, 7, 10], [-1, -1, -1, 2, 6]] + pixel_order = XY(origin=Compass.W, clockwise=True) + x = self.reorder(pixel_order, x) + nside = self._nside() + *shape, _ = x.shape + x = x.reshape((*shape, 12, nside, nside)) + output = torch.full((*shape, 5 * nside, 5 * nside), device=x.device, dtype=x.dtype, fill_value=fill_value) + + for j in range(len(grid)): + for i in range(len(grid[0])): + face = grid[j][i] + if face != -1: + output[j * nside : (j + 1) * nside, i * nside : (i + 1) * nside] = x[face] + return output + + +# nside = 2^ZOOM_LEVELS +ZOOM_LEVELS = 20 + + +def _extract_every_other_bit(binary_number): + result = 0 + shift_count = 0 + + for i in range(ZOOM_LEVELS): + # Check if the least significant bit is 1 + # Set the corresponding bit in the result + result |= (binary_number & 1) << shift_count + + # Shift to the next bit to check + binary_number = binary_number >> 2 + shift_count += 1 + + return result + + +def _flip_xy(nside: int, i): + n2 = nside * nside + f = i // n2 + y = (i % n2) // nside + x = i % nside + return n2 * f + nside * x + y + + +def _rotate_index(nside: int, rotations: int, i): + # Extract f, x, and y from i + # convention is arr[f, y, x] ... x is the fastest changing index + n2 = nside * nside + f = i // n2 + y = (i % n2) // nside + x = i % nside + + # Reduce k to its equivalent in the range [0, 3] + k = rotations % 4 + + assert 0 <= k < 4 + + # Apply the rotation based on k + if k == 1: # 90 degrees counterclockwise + new_x, new_y = -y - 1, x + elif k == 2: # 180 degrees + new_x, new_y = -x - 1, -y - 1 + elif k == 3: # 270 degrees counterclockwise + new_x, new_y = y, -x - 1 + else: # k == 0, no change + new_x, new_y = x, y + + # Adjust for negative indices + new_x = new_x % nside + new_y = new_y % nside + + # Recalculate the linear index with the rotated x and y + return n2 * f + nside * new_y + new_x + + +def nest2xy(nside, i): + """convert NEST to XY index""" + tile = i // nside**2 + j = i % (nside**2) + x = _extract_every_other_bit(j) + y = _extract_every_other_bit(j >> 1) + return tile * nside**2 + y * nside + x + + +def xy2nest(nside, i): + """convert XY index to NEST""" + tile = i // (nside**2) + y = (i % (nside**2)) // nside + x = i % nside + + result = 0 + for i in range(ZOOM_LEVELS): + # Extract the ith bit from the number + extracted_bit = (x >> i) & 1 + result |= extracted_bit << (2 * i) + + extracted_bit = (y >> i) & 1 + result |= extracted_bit << (2 * i + 1) + return result | (tile * nside**2) + + +def approx_grid_length_meters(nside): + r_m = 6378140 + area = 4 * np.pi * r_m**2 / (12 * nside**2) + return np.sqrt(area) + + +def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): + """conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor + + Applies a 2D convolution over an input image composed of several input + planes. + + This operator supports :ref:`TensorFloat32`. + + See :class:`~torch.nn.Conv2d` for details and output shape. + + """ + px, py = padding + assert px == py + + n, c, x, y = input.shape + npix = input.size(-1) + nside2 = npix // 12 + nside = int(math.sqrt(nside2)) + assert nside**2 * 12 == npix + + input = einops.rearrange(input, "n c () (f x y) -> (n c) f x y", f=12, x=nside) + input = pad(input, px) + input = einops.rearrange(input, "(n c) f x y -> n c f x y", c=c) + padding = (0, 0, 0) + padding = 'valid' + + if not isinstance(stride, int): + stride = stride + (1,) + + if not isinstance(dilation, int): + dilation = (1,) + dilation + + weight = weight.unsqueeze(-3) + out = torch.nn.functional.conv3d(input, weight, bias, stride, padding, dilation, groups) + return einops.rearrange(out, "n c f x y -> n c () (f x y)") diff --git a/earth2grid/healpix_bare.py b/earth2grid/healpix_bare.py new file mode 100644 index 0000000..7ec3a3b --- /dev/null +++ b/earth2grid/healpix_bare.py @@ -0,0 +1,79 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import torch + +from earth2grid import _healpix_bare +from earth2grid._healpix_bare import corners, hpc2loc, hpd2loc, nest2hpd, nest2ring, ring2hpd, ring2nest + +__all__ = [ + "pix2ang", + "ring2nest", + "nest2ring", + "hpc2loc", + "corners", +] + + +def pix2ang(nside, i, nest=False, lonlat=False): + if nest: + hpd = nest2hpd(nside, i) + else: + hpd = ring2hpd(nside, i) + loc = hpd2loc(nside, hpd) + lon, lat = _loc2ang(loc) + + if lonlat: + return torch.rad2deg(lon), 90 - torch.rad2deg(lat) + else: + return lat, lon + + +def _loc2ang(loc): + """ + static t_ang loc2ang(tloc loc) + { return (t_ang){atan2(loc.s,loc.z), loc.phi}; } + """ + z = loc[..., 0] + s = loc[..., 1] + phi = loc[..., 2] + return phi % (2 * torch.pi), torch.atan2(s, z) + + +def loc2vec(loc): + z = loc[..., 0] + s = loc[..., 1] + phi = loc[..., 2] + x = (s * torch.cos(phi),) + y = (s * torch.sin(phi),) + return x, y, z + + +def get_interp_weights(nside: int, lon: torch.Tensor, lat: torch.Tensor): + """ + + Args: + lon: longtiude in deg E. Shape (*) + lat: latitdue in deg E. Shape (*) + + Returns: + pix, weights: both shaped (4, *). pix is given in RING convention. + + """ + shape = lon.shape + lon = lon.double().cpu().flatten() + lat = lat.double().cpu().flatten() + pix, weights = _healpix_bare.get_interp_weights(nside, lon, lat) + + return pix.view(4, *shape), weights.view(4, *shape) diff --git a/earth2grid/latlon.py b/earth2grid/latlon.py new file mode 100644 index 0000000..8d55e7f --- /dev/null +++ b/earth2grid/latlon.py @@ -0,0 +1,200 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import numpy as np +import torch + +from earth2grid import base + +try: + import pyvista as pv +except ImportError: + pv = None + + +class BilinearInterpolator(torch.nn.Module): + """Bilinear interpolation for a non-uniform grid""" + + def __init__( + self, x_coords: torch.Tensor, y_coords: torch.Tensor, x_query: torch.Tensor, y_query: torch.Tensor + ) -> None: + """ + + Args: + x_coords (Tensor): X-coordinates of the input grid, shape [W]. Must be in increasing sorted order. + y_coords (Tensor): Y-coordinates of the input grid, shape [H]. Must be in increasing sorted order. + x_query (Tensor): X-coordinates for query points, shape [N]. + y_query (Tensor): Y-coordinates for query points, shape [N]. + """ + super().__init__() + + # Ensure input coordinates are float for interpolation + x_coords, y_coords = x_coords.float(), y_coords.float() + + if torch.any(x_coords[1:] < x_coords[:-1]): + raise ValueError("x_coords must be in non-decreasing order.") + + if torch.any(y_coords[1:] < y_coords[:-1]): + raise ValueError("y_coords must be in non-decreasing order.") + + # Find indices for the closest lower and upper bounds in x and y directions + x_l_idx = torch.searchsorted(x_coords, x_query, right=True) - 1 + x_u_idx = x_l_idx + 1 + y_l_idx = torch.searchsorted(y_coords, y_query, right=True) - 1 + y_u_idx = y_l_idx + 1 + + # Clip indices to ensure they are within the bounds of the input grid + x_l_idx = x_l_idx.clamp(0, x_coords.size(0) - 2) + x_u_idx = x_u_idx.clamp(1, x_coords.size(0) - 1) + y_l_idx = y_l_idx.clamp(0, y_coords.size(0) - 2) + y_u_idx = y_u_idx.clamp(1, y_coords.size(0) - 1) + + # Compute weights + x_l_weight = (x_coords[x_u_idx] - x_query) / (x_coords[x_u_idx] - x_coords[x_l_idx]) + x_u_weight = (x_query - x_coords[x_l_idx]) / (x_coords[x_u_idx] - x_coords[x_l_idx]) + y_l_weight = (y_coords[y_u_idx] - y_query) / (y_coords[y_u_idx] - y_coords[y_l_idx]) + y_u_weight = (y_query - y_coords[y_l_idx]) / (y_coords[y_u_idx] - y_coords[y_l_idx]) + weights = torch.stack( + [x_l_weight * y_l_weight, x_u_weight * y_l_weight, x_l_weight * y_u_weight, x_u_weight * y_u_weight], dim=-1 + ) + + self.register_buffer("weights", weights) + + stride = x_coords.size(-1) + index = torch.stack( + [ + x_l_idx + stride * y_l_idx, + x_u_idx + stride * y_l_idx, + x_l_idx + stride * y_u_idx, + x_u_idx + stride * y_u_idx, + ], + dim=-1, + ) + self.register_buffer("index", index) + + def forward(self, z: torch.Tensor): + """ + Interpolate the field + + Args: + z: shape [Y, X] + """ + *shape, y, x = z.shape + zrs = z.view(-1, y * x).T + # using embedding bag is 2x faster on cpu and 4x on gpu. + interpolated = torch.nn.functional.embedding_bag(self.index, zrs, per_sample_weights=self.weights, mode='sum') + interpolated = interpolated.T.view(*shape, self.weights.size(0)) + return interpolated + + +class LatLonGrid(base.Grid): + def __init__(self, lat: list[float], lon: list[float]): + """ + Args: + lat: center of lat cells + lon: center of lon cells + """ + self._lat = lat + self._lon = lon + + @property + def lat(self): + return np.array(self._lat)[:, None] + + @property + def lon(self): + return np.array(self._lon) + + @property + def shape(self): + return (len(self.lat), len(self.lon)) + + def get_bilinear_regridder_to(self, lat: np.ndarray, lon: np.ndarray): + """Get regridder to the specified lat and lon points""" + return _RegridFromLatLon(self, lat, lon) + + def _lonb(self): + edges = (self.lon[1:] + self.lon[:-1]) / 2 + d_left = self.lon[1] - self.lon[0] + d_right = self.lon[-1] - self.lon[-2] + return np.concatenate([self.lon[0:1] - d_left / 2, edges, self.lon[-1:] + d_right / 2]) + + def visualize(self, data): + raise NotImplementedError() + + def to_pyvista(self): + # TODO need to make lat the cell centers rather than boundaries + + if pv is None: + raise ImportError("Need to install pyvista") + + print(self._lonb()) + + lon, lat = np.meshgrid(self._lonb(), self.lat) + y = np.cos(np.deg2rad(lat)) * np.sin(np.deg2rad(lon)) + x = np.cos(np.deg2rad(lat)) * np.cos(np.deg2rad(lon)) + z = np.sin(np.deg2rad(lat)) + grid = pv.StructuredGrid(x, y, z) + return grid + + +class _RegridFromLatLon(torch.nn.Module): + """Regrid from lat-lon to unstructured grid with bilinear interpolation""" + + def __init__(self, src: LatLonGrid, lat: np.ndarray, lon: np.ndarray): + super().__init__() + + lat, lon = np.broadcast_arrays(lat, lon) + self.shape = lat.shape + + # TODO add device switching logic (maybe use torch registers for this + # info) + long = np.concatenate([src.lon.ravel(), [360]], axis=-1) + long_t = torch.from_numpy(long) + + # flip the order latg since bilinear only works with increasing coordinate values + lat_increasing = src.lat[1] > src.lat[0] + latg_t = torch.from_numpy(src.lat.ravel()) + lat_query = torch.from_numpy(lat.ravel()) + + if not lat_increasing: + lat_query = -lat_query + latg_t = -latg_t + + self._bilinear = BilinearInterpolator(long_t, latg_t, y_query=lat_query, x_query=torch.from_numpy(lon.ravel())) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # pad z in lon direction + # only works for a global grid + # TODO generalize this to local grids and add options for padding + x = torch.cat([x, x[..., 0:1]], axis=-1) + out = self._bilinear(x) + return out.view(out.shape[:-1] + self.shape) + + +def equiangular_lat_lon_grid(nlat: int, nlon: int, includes_south_pole: bool = True) -> LatLonGrid: + """Return a regular lat-lon grid + + Lat is ordered from 90 to -90. Includes -90 and only if if includes_south_pole is True. + Lon is ordered from 0 to 360. includes 0, but not 360. + + Args: + nlat: number of latitude points + nlon: number of longtidue points + includes_south_pole: if true the final ``nlat`` includes the south pole + + """ # noqa + lat = np.linspace(90, -90, nlat, endpoint=includes_south_pole) + lon = np.linspace(0, 360, nlon, endpoint=False) + return LatLonGrid(lat.tolist(), lon.tolist()) diff --git a/earth2grid/third_party/healpix_bare/LICENSE b/earth2grid/third_party/healpix_bare/LICENSE new file mode 100644 index 0000000..4f12e3a --- /dev/null +++ b/earth2grid/third_party/healpix_bare/LICENSE @@ -0,0 +1,29 @@ +Copyright (C) 1997-2019 Krzysztof M. Gorski, Eric Hivon, Martin Reinecke, + Benjamin D. Wandelt, Anthony J. Banday, + Matthias Bartelmann, + Reza Ansari & Kenneth M. Ganga + +Redistribution and use in source and binary forms, with or without modification, +are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, this + list of conditions and the following disclaimer in the documentation and/or + other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its contributors may + be used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON +ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/earth2grid/third_party/healpix_bare/NEWS b/earth2grid/third_party/healpix_bare/NEWS new file mode 100644 index 0000000..b0debc4 --- /dev/null +++ b/earth2grid/third_party/healpix_bare/NEWS @@ -0,0 +1,2 @@ +Version 1.0 (March 2019): + initial release diff --git a/earth2grid/third_party/healpix_bare/README b/earth2grid/third_party/healpix_bare/README new file mode 100644 index 0000000..d78d621 --- /dev/null +++ b/earth2grid/third_party/healpix_bare/README @@ -0,0 +1,42 @@ +This package implements a small subset of routines for working with the +HEALPix tesselation of the sphere. + +For information about HEALPix, see the publication by +K.M. Gorski et al., 2005, Ap.J., 622, p.759 +(http://adsabs.harvard.edu/abs/2005ApJ...622..759G) +More comprehensive software packages supporting HEALPix, as well as extensive +documentation and references, are available at +https://healpix.sourceforge.io/. + + +Compilation +=========== + +gcc -O2 healpix_bare.c -std=c99 test.c -lm +clang -O2 healpix_bare.c -std=c99 test.c -lm +icc -O2 healpix_bare.c -std=c99 test.c -lm + + +Notes +===== + +This package is designed to cover a small set of frequently used HEALPix-related +routines, which are implemented with a focus on robustness, accuracy and clarity. +They have good performance, but are not heavily tuned to keep the code simple - +the full-featured, GPL-licensed Healpix Fortran and C++ libraries (available +via the link above) will be more efficient in most situations. + +The motivation for this package is to give developers of BSD-licensed +Healpix-related codes a reliable starting point, so that they don't have to +re-implement the central algorithms (like conversions between pixel index and +angular coordinates), which are nontrivial to do correctly for all corner cases. + + +Acknowledgements +================ + +If you are using this code in your own packages, please consider citing +the original paper in your publications: + +K.M. Gorski et al., 2005, Ap.J., 622, p.759 +(http://adsabs.harvard.edu/abs/2005ApJ...622..759G) diff --git a/earth2grid/third_party/healpix_bare/healpix_bare.c b/earth2grid/third_party/healpix_bare/healpix_bare.c new file mode 100644 index 0000000..bb01aea --- /dev/null +++ b/earth2grid/third_party/healpix_bare/healpix_bare.c @@ -0,0 +1,310 @@ +/* ----------------------------------------------------------------------------- + * + * Copyright (C) 1997-2019 Krzysztof M. Gorski, Eric Hivon, Martin Reinecke, + * Benjamin D. Wandelt, Anthony J. Banday, + * Matthias Bartelmann, + * Reza Ansari & Kenneth M. Ganga + * + * Implementation of the Healpix bare bones C library + * + * Licensed under a 3-clause BSD style license - see LICENSE + * + * For more information on HEALPix and additional software packages, see + * https://healpix.sourceforge.io/ + * + * If you are using this code in your own packages, please consider citing + * the original paper in your publications: + * K.M. Gorski et al., 2005, Ap.J., 622, p.759 + * (http://adsabs.harvard.edu/abs/2005ApJ...622..759G) + * + *----------------------------------------------------------------------------*/ + +#include +#include "healpix_bare.h" + +const double pi = 3.141592653589793238462643383279502884197; + +static const int jrll[] = { 2,2,2,2,3,3,3,3,4,4,4,4 }; +static const int jpll[] = { 1,3,5,7,0,2,4,6,1,3,5,7 }; + +/* conversions between continuous coordinate systems */ + +typedef struct { double z, s, phi; } tloc; + +/*! A structure describing the continuous Healpix coordinate system. + \a f takes values in [0;11], \a x and \a y lie in [0.0; 1.0]. */ +typedef struct { double x, y; int32_t f; } t_hpc; + +static t_hpc loc2hpc (tloc loc) + { + double za = fabs(loc.z); + double x = loc.phi*(1./(2.*pi)); + if (x<0.) + x += (int64_t)x + 1.; + else if (x>=1.) + x -= (int64_t)x; + double tt = 4.*x; + + if (za<=2./3.) /* Equatorial region */ + { + double temp1 = 0.5+tt; // [0.5; 4.5) + double temp2 = loc.z*0.75; // [-0.5; +0.5] + double jp = temp1-temp2; /* index of ascending edge line */ // [0; 5) + double jm = temp1+temp2; /* index of descending edge line */ // [0; 5) + int ifp = (int)jp; /* in {0,4} */ + int ifm = (int)jm; + return (t_hpc){jm-ifm, 1+ifp - jp, + (ifp==ifm) ? (ifp|4) : ((ifp=4) ntt=3; + double tp = tt-ntt; // [0;1) + double tmp = loc.s/sqrt((1.+za)*(1./3.)); // FIXME optimize! + + double jp = tp*tmp; /* increasing edge line index */ + double jm = (1.0-tp)*tmp; /* decreasing edge line index */ + if (jp>1.) jp = 1.; /* for points too close to the boundary */ + if (jm>1.) jm = 1.; + return (loc.z >= 0) ? (t_hpc){1.-jm, 1.-jp, ntt} + : (t_hpc){jp, jm, ntt+8}; + } + +static tloc hpc2loc (t_hpc hpc) + { + double jr = jrll[hpc.f] - hpc.x - hpc.y; + if (jr<1.) + { + double tmp = jr*jr*(1./3.); + double z = 1. - tmp; + double s = sqrt(tmp*(2.-tmp)); + double phi = (pi*0.25)*(jpll[hpc.f] + (hpc.x-hpc.y)/jr); + return (tloc){z,s,phi}; + } + else if (jr>3.) + { + jr = 4.-jr; + double tmp = jr*jr*(1./3.); + double z = tmp - 1.; + double s = sqrt(tmp*(2.-tmp)); + double phi = (pi*0.25)*(jpll[hpc.f] + (hpc.x-hpc.y)/jr); + return (tloc){z,s,phi}; + } + else + { + double z = (2.-jr)*(2./3.); + double s = sqrt((1.+z)*(1.-z)); + double phi = (pi*0.25)*(jpll[hpc.f] + hpc.x - hpc.y); + return (tloc){z,s,phi}; + } + } + +static tloc ang2loc(t_ang ang) + { + double cth=cos(ang.theta), sth=sin(ang.theta); + if (sth<0) { sth=-sth; ang.phi+=pi; } + return (tloc){cth, sth, ang.phi}; + } + +static t_ang loc2ang(tloc loc) + { return (t_ang){atan2(loc.s,loc.z), loc.phi}; } + +static tloc vec2loc(t_vec vec) + { + double vlen=sqrt(vec.x*vec.x+vec.y*vec.y+vec.z*vec.z); + double cth = vec.z/vlen; + double sth = sqrt(vec.x*vec.x+vec.y*vec.y)/vlen; + return (tloc){cth,sth,atan2(vec.y,vec.x)}; + } + +static t_vec loc2vec(tloc loc) + { return (t_vec){loc.s*cos(loc.phi), loc.s*sin(loc.phi), loc.z}; } + +t_vec ang2vec(t_ang ang) + { return loc2vec(ang2loc(ang)); } + +t_ang vec2ang(t_vec vec) + { + return (t_ang) {atan2(sqrt(vec.x*vec.x+vec.y*vec.y),vec.z), + atan2(vec.y,vec.x)}; + } + +/* conversions between discrete coordinate systems */ + +static int64_t isqrt(int64_t v) + { + int64_t res = sqrt(v+0.5); + if (v<((int64_t)(1)<<50)) return res; + if (res*res>v) + --res; + else if ((res+1)*(res+1)<=v) + ++res; + return res; + } + +static int64_t spread_bits (int64_t v) + { + int64_t res = v & 0xffffffff; + res = (res^(res<<16)) & 0x0000ffff0000ffff; + res = (res^(res<< 8)) & 0x00ff00ff00ff00ff; + res = (res^(res<< 4)) & 0x0f0f0f0f0f0f0f0f; + res = (res^(res<< 2)) & 0x3333333333333333; + res = (res^(res<< 1)) & 0x5555555555555555; + return res; + } + +static int64_t compress_bits (int64_t v) + { + int64_t res = v & 0x5555555555555555; + res = (res^(res>> 1)) & 0x3333333333333333; + res = (res^(res>> 2)) & 0x0f0f0f0f0f0f0f0f; + res = (res^(res>> 4)) & 0x00ff00ff00ff00ff; + res = (res^(res>> 8)) & 0x0000ffff0000ffff; + res = (res^(res>>16)) & 0x00000000ffffffff; + return res; + } + +/*! A structure describing the discrete Healpix coordinate system. + \a f takes values in [0;11], \a x and \a y lie in [0; nside[. */ +typedef struct { int64_t x, y; int32_t f; } t_hpd; + +static int64_t hpd2nest (int64_t nside, t_hpd hpd) + { return (hpd.f*nside*nside) + spread_bits(hpd.x) + (spread_bits(hpd.y)<<1); } + +static t_hpd nest2hpd (int64_t nside, int64_t pix) + { + int64_t npface_=nside*nside, p2=pix&(npface_-1); + return (t_hpd){compress_bits(p2), compress_bits(p2>>1), pix/npface_}; + } + +static int64_t hpd2ring (int64_t nside_, t_hpd hpd) + { + int64_t nl4 = 4*nside_; + int64_t jr = (jrll[hpd.f]*nside_) - hpd.x - hpd.y - 1; + + if (jrnl4) ? jp-nl4 : ((jp<1) ? jp+nl4 : jp); + return 2*jr*(jr-1) + jp - 1; + } + else if (jr > 3*nside_) + { + jr = nl4-jr; + int64_t jp = (jpll[hpd.f]*jr + hpd.x - hpd.y + 1) / 2; + jp = (jp>nl4) ? jp-nl4 : ((jp<1) ? jp+nl4 : jp); + return 12*nside_*nside_ - 2*(jr+1)*jr + jp - 1; + } + else + { + int64_t jp = (jpll[hpd.f]*nside_ + hpd.x - hpd.y + 1 + ((jr-nside_)&1)) / 2; + jp = (jp>nl4) ? jp-nl4 : ((jp<1) ? jp+nl4 : jp); + return 2*nside_*(nside_-1) + (jr-nside_)*nl4 + jp - 1; + } + } + +static t_hpd ring2hpd (int64_t nside_, int64_t pix) + { + int64_t ncap_=2*nside_*(nside_-1); + int64_t npix_=12*nside_*nside_; + + if (pix>1; /* counted from North pole */ + int64_t iphi = (pix+1) - 2*iring*(iring-1); + int64_t face = (iphi-1)/iring; + int64_t irt = iring - (jrll[face]*nside_) + 1; + int64_t ipt = 2*iphi- jpll[face]*iring -1; + if (ipt>=2*nside_) ipt-=8*nside_; + return (t_hpd) {(ipt-irt)>>1, (-(ipt+irt))>>1, face}; + } + else if (pix<(npix_-ncap_)) /* Equatorial region */ + { + int64_t ip = pix - ncap_; + int64_t iring = (ip/(4*nside_)) + nside_; /* counted from North pole */ + int64_t iphi = (ip%(4*nside_)) + 1; + int64_t kshift = (iring+nside_)&1; + int64_t ire = iring-nside_+1; + int64_t irm = 2*nside_+2-ire; + int64_t ifm = (iphi - ire/2 + nside_ -1) / nside_; + int64_t ifp = (iphi - irm/2 + nside_ -1) / nside_; + int64_t face = (ifp==ifm) ? (ifp|4) : ((ifp=2*nside_) ipt-=8*nside_; + return (t_hpd) {(ipt-irt)>>1, (-(ipt+irt))>>1, face}; + } + else /* South Polar cap */ + { + int64_t ip = npix_ - pix; + int64_t iring = (1+isqrt(2*ip-1))>>1; /* counted from South pole */ + int64_t iphi = 4*iring + 1 - (ip - 2*iring*(iring-1)); + int64_t face=8+(iphi-1)/iring; + int64_t irt = 4*nside_ - iring - (jrll[face]*nside_) + 1; + int64_t ipt = 2*iphi- jpll[face]*iring -1; + if (ipt>=2*nside_) ipt-=8*nside_; + return (t_hpd) {(ipt-irt)>>1, (-(ipt+irt))>>1, face}; + } + } + +int64_t nest2ring(int64_t nside, int64_t ipnest) + { + if ((nside&(nside-1))!=0) return -1; + return hpd2ring (nside, nest2hpd (nside, ipnest)); + } + +int64_t ring2nest(int64_t nside, int64_t ipring) + { + if ((nside&(nside-1))!=0) return -1; + return hpd2nest (nside, ring2hpd (nside, ipring)); + } + +/* mixed conversions */ + +static t_hpd loc2hpd (int64_t nside_, tloc loc) + { + t_hpc tmp = loc2hpc(loc); + return (t_hpd){(tmp.x*nside_), (tmp.y*nside_), tmp.f}; + } + +static tloc hpd2loc (int64_t nside_, t_hpd hpd) + { + double xns = 1./nside_; + t_hpc tmp = (t_hpc){(hpd.x+0.5)*xns,(hpd.y+0.5)*xns,hpd.f}; + return hpc2loc(tmp); + } + +int64_t npix2nside(int64_t npix) + { + int64_t res = isqrt(npix/12); + return (res*res*12==npix) ? res : -1; + } + +int64_t nside2npix(int64_t nside) + { return 12*nside*nside; } + +double vec_angle(t_vec v1, t_vec v2) + { + t_vec cross = { v1.y*v2.z - v1.z*v2.y, + v1.z*v2.x - v1.x*v2.z, + v1.x*v2.y - v1.y*v2.x }; + double len_cross = sqrt(cross.x*cross.x + cross.y*cross.y + cross.z*cross.z); + double dot = v1.x*v2.x + v1.y*v2.y + v1.z*v2.z; + return atan2 (len_cross, dot); + } + +int64_t ang2ring(int64_t nside, t_ang ang) + { return hpd2ring(nside, loc2hpd(nside, ang2loc(ang))); } +int64_t ang2nest(int64_t nside, t_ang ang) + { return hpd2nest(nside, loc2hpd(nside, ang2loc(ang))); } +int64_t vec2ring(int64_t nside, t_vec vec) + { return hpd2ring(nside, loc2hpd(nside, vec2loc(vec))); } +int64_t vec2nest(int64_t nside, t_vec vec) + { return hpd2nest(nside, loc2hpd(nside, vec2loc(vec))); } +t_ang ring2ang(int64_t nside, int64_t ipix) + { return loc2ang(hpd2loc(nside, ring2hpd(nside, ipix))); } +t_ang nest2ang(int64_t nside, int64_t ipix) + { return loc2ang(hpd2loc(nside, nest2hpd(nside, ipix))); } +t_vec ring2vec(int64_t nside, int64_t ipix) + { return loc2vec(hpd2loc(nside, ring2hpd(nside, ipix))); } +t_vec nest2vec(int64_t nside, int64_t ipix) + { return loc2vec(hpd2loc(nside, nest2hpd(nside, ipix))); } diff --git a/earth2grid/third_party/healpix_bare/healpix_bare.h b/earth2grid/third_party/healpix_bare/healpix_bare.h new file mode 100644 index 0000000..2314171 --- /dev/null +++ b/earth2grid/third_party/healpix_bare/healpix_bare.h @@ -0,0 +1,118 @@ +/* ----------------------------------------------------------------------------- + * + * Copyright (C) 1997-2019 Krzysztof M. Gorski, Eric Hivon, Martin Reinecke, + * Benjamin D. Wandelt, Anthony J. Banday, + * Matthias Bartelmann, + * Reza Ansari & Kenneth M. Ganga + * + * Public interface of the Healpix bare bones C library + * + * Licensed under a 3-clause BSD style license - see LICENSE + * + * For more information on HEALPix and additional software packages, see + * https://healpix.sourceforge.io/ + * + * If you are using this code in your own packages, please consider citing + * the original paper in your publications: + * K.M. Gorski et al., 2005, Ap.J., 622, p.759 + * (http://adsabs.harvard.edu/abs/2005ApJ...622..759G) + * + *----------------------------------------------------------------------------*/ + +#ifndef HEALPIX_BARE_H +#define HEALPIX_BARE_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* This code is written in C99; you may have to adjust your compiler flags. */ + +/* Continuous coordinate systems */ + +/* Admissible values for theta (definition see below) + 0 <= theta <= pi + + Admissible values for phi (definition see below) + In principle unconstrained, but best accuracy is obtained for + -2*pi <= phi <= 2*pi */ + +/*! A structure describing a location on the sphere. \a Theta is the co-latitude + in radians (0 at the North Pole, increasing to pi at the South Pole. + \a Phi is the azimuth in radians. */ +typedef struct { double theta, phi; } t_ang; +/*! A structure describing a 3-vector with coordinates \a x, \a y and \a z.*/ +typedef struct { double x, y, z; } t_vec; + +/*! Returns a normalized 3-vector pointing in the same direction as \a ang. */ +t_vec ang2vec(t_ang ang); +/*! Returns a t_ang describing the same direction as the 3-vector \a vec. + \a vec need not be normalized. */ +t_ang vec2ang(t_vec vec); + + +/* Discrete coordinate systems */ + +/* Admissible values for nside parameters: + any integer power of 2 with 1 <= nside <= 1<<29 + + Admissible values for pixel indices: + 0 <= idx < 12*nside*nside */ + +/*! Returns the RING pixel index of pixel \a ipnest at resolution \a nside. + On error, returns -1. */ +int64_t nest2ring(int64_t nside, int64_t ipnest); +/*! Returns the NEST pixel index of pixel \a ipring at resolution \a nside. + On error, returns -1. */ +int64_t ring2nest(int64_t nside, int64_t ipring); + + +/* Conversions between continuous and discrete coordinate systems */ + +/*! Returns the pixel number in NEST scheme at resolution \a nside, + which contains the position \a ang. */ +int64_t ang2nest(int64_t nside, t_ang ang); +/*! Returns the pixel number in RING scheme at resolution \a nside, + which contains the position \a ang. */ +int64_t ang2ring(int64_t nside, t_ang ang); + +/*! Returns a t_ang corresponding to the angular position of the center of + pixel \a ipix in NEST scheme at resolution \a nside. */ +t_ang nest2ang(int64_t nside, int64_t ipix); +/*! Returns a t_ang corresponding to the angular position of the center of + pixel \a ipix in RING scheme at resolution \a nside. */ +t_ang ring2ang(int64_t nside, int64_t ipix); + +/*! Returns the pixel number in NEST scheme at resolution \a nside, + which contains the direction described the 3-vector \a vec. */ +int64_t vec2nest(int64_t nside, t_vec vec); +/*! Returns the pixel number in RING scheme at resolution \a nside, + which contains the direction described the 3-vector \a vec. */ +int64_t vec2ring(int64_t nside, t_vec vec); + +/*! Returns a normalized 3-vector pointing in the direction of the center + of pixel \a ipix in NEST scheme at resolution \a nside. */ +t_vec nest2vec(int64_t nside, int64_t ipix); +/*! Returns a normalized 3-vector pointing in the direction of the center + of pixel \a ipix in RING scheme at resolution \a nside. */ +t_vec ring2vec(int64_t nside, int64_t ipix); + + +/* Miscellaneous utility routines */ + +/*! Returns \a 12*nside*nside. */ +int64_t nside2npix(int64_t nside); +/*! Returns \a sqrt(npix/12) if this is an integer number, otherwise \a -1. */ +int64_t npix2nside(int64_t npix); + +/*! Returns the angle (in radians) between the vectors \a v1 and \a v2. + The result is accurate even for angles close to 0 and pi. */ +double vec_angle(t_vec v1, t_vec v2); + +#ifdef __cplusplus +} /* extern "C" */ +#endif + +#endif /* HEALPIX_REFERENCE_H */ diff --git a/earth2grid/third_party/healpix_bare/test.c b/earth2grid/third_party/healpix_bare/test.c new file mode 100644 index 0000000..048a68a --- /dev/null +++ b/earth2grid/third_party/healpix_bare/test.c @@ -0,0 +1,129 @@ +/* ----------------------------------------------------------------------------- + * + * Copyright (C) 1997-2019 Krzysztof M. Gorski, Eric Hivon, Martin Reinecke, + * Benjamin D. Wandelt, Anthony J. Banday, + * Matthias Bartelmann, + * Reza Ansari & Kenneth M. Ganga + * + * Test for the Healpix bare bones C library + * + * Licensed under a 3-clause BSD style license - see LICENSE + * + * For more information on HEALPix and additional software packages, see + * https://healpix.sourceforge.io/ + * + * If you are using this code in your own packages, please consider citing + * the original paper in your publications: + * K.M. Gorski et al., 2005, Ap.J., 622, p.759 + * (http://adsabs.harvard.edu/abs/2005ApJ...622..759G) + * + *----------------------------------------------------------------------------*/ + +#include +#include +#include +#include +#include + +#include "healpix_bare.h" + +static void error_check(int64_t nside, int64_t ipix, int64_t ip1) + { + if (ip1 != ipix) { + printf("Error0: %" PRId64 " %" PRId64 " %" PRId64 "\n", nside, ipix, ip1); + abort(); + } + } + +static void test2_helper1 (int64_t nside, int64_t dpix, int64_t epix) + { + /* Find the number of pixels in the full map */ + for (int64_t ipix = 0; ipix < epix; ipix +=dpix) + error_check(nside, ipix, ring2nest(nside, nest2ring(nside, ipix))); + for (int64_t ipix = 0; ipix < epix; ipix +=dpix) + error_check(nside, ipix, ang2nest(nside, nest2ang(nside, ipix))); + for (int64_t ipix = 0; ipix < epix; ipix +=dpix) + error_check(nside, ipix, vec2nest(nside, nest2vec(nside, ipix))); + } + +static void test2_helper2 (int64_t nside, int64_t dpix) + { + /* Find the number of pixels in the full map */ + int64_t npix = nside2npix(nside); + for (int64_t ipix = 0; ipix < npix; ipix +=dpix) + error_check(nside, ipix, ang2ring(nside, ring2ang(nside, ipix))); + for (int64_t ipix = 0; ipix < npix; ipix +=dpix) + error_check(nside, ipix, vec2ring(nside, ring2vec(nside, ipix))); + } + +static void test2(void) { + printf("Starting basic C Healpix pixel routines test\n"); + + for (int i=0; i<=29; ++i) + { + int64_t nside = 1LL<12*nside*nside) epix = 12*nside*nside; + printf("Nside: %" PRId64 "\n",nside); + test2_helper1(nside,nside*nside/123456+1, npix); + test2_helper1(nside, 1, epix); + test2_helper2(nside,nside*nside/123456+1); + } + + printf("test completed\n\n"); +} + +static void test3(void) { + printf("Starting nontrivial C Healpix pixel routines test\n"); + + int64_t nside = 1024; + int64_t dpix = 23; + + /* Find the number of pixels in the full map */ + int64_t npix = nside2npix(nside); + printf("Number of pixels in full map: %ld\n", npix); + + printf("dpix: %ld\n", dpix); + printf("Nest -> ang -> vec -> ang -> Ring -> Nest\n"); + for (int64_t ipix = 0; ipix < npix; ipix +=dpix) + error_check(nside, ipix, + ring2nest(nside, ang2ring(nside, vec2ang(ang2vec(nest2ang(nside, ipix)))))); + printf("Ring -> ang -> Nest -> Ring\n"); + for (int64_t ipix = 0; ipix < npix; ipix +=dpix) + error_check(nside, ipix, + nest2ring(nside,ang2nest(nside, ring2ang(nside, ipix)))); + + printf("Nest -> vec -> Ring -> Nest\n"); + for (int64_t ipix = 0; ipix < npix; ipix +=dpix) + error_check(nside, ipix, + ring2nest(nside,vec2ring(nside,nest2vec(nside, ipix)))); + printf("Ring -> vec -> Nest -> Ring\n"); + for (int64_t ipix = 0; ipix < npix; ipix +=dpix) + error_check(nside, ipix, + nest2ring(nside,vec2nest(nside,ring2vec(nside, ipix)))); + + printf("%" PRId64 "\n", nside); + printf("test completed\n\n"); +} + +static void test_ang(void) { + printf("Starting vector angle test\n"); + for (double i=0; i<10; i+=0.01) { + t_vec v1 = { i,10-i,sin(i) }, v2 = { 3-1.5*i, sqrt(i), i*0.3 }; + double dot = v1.x*v2.x + v1.y*v2.y + v1.z*v2.z, + l1 = sqrt(v1.x*v1.x+v1.y*v1.y+v1.z*v1.z), + l2 = sqrt(v2.x*v2.x+v2.y*v2.y+v2.z*v2.z); + double ang = acos(dot/(l1*l2)); + if (fabs(ang - vec_angle(v1,v2))>1e-13) + printf("error in vector angle test: %e %e\n", ang, ang-vec_angle(v1, v2)); + } + printf("test completed\n\n"); +} + +int main(void) { + test2(); + test3(); + test_ang(); + return 0; +} diff --git a/earth2grid/third_party/zephyr/LICENSE b/earth2grid/third_party/zephyr/LICENSE new file mode 100644 index 0000000..78dd216 --- /dev/null +++ b/earth2grid/third_party/zephyr/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2018 Jonathan Weyn + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/earth2grid/third_party/zephyr/healpix.py b/earth2grid/third_party/zephyr/healpix.py new file mode 100644 index 0000000..c322afc --- /dev/null +++ b/earth2grid/third_party/zephyr/healpix.py @@ -0,0 +1,281 @@ +#!/usr/bin/env python3 + +""" +This file contains padding and convolution classes to perform according operations on the twelve faces of the HEALPix. + + + HEALPix Face order 3D array representation + ----------------- +-------------------------- //\\ //\\ //\\ //\\ | | | | | +|| 0 | 1 | 2 | 3 || // \\// \\// \\// \\ |0 |1 |2 |3 | +|\\ //\\ //\\ //\\ //| /\\0 //\\1 //\\2 //\\3 // ----------------- +| \\// \\// \\// \\// | // \\// \\// \\// \\// | | | | | +|4//\\5 //\\6 //\\7 //\\4| \\4//\\5 //\\6 //\\7 //\\ |4 |5 |6 |7 | +|// \\// \\// \\// \\| \\/ \\// \\// \\// \\ ----------------- +|| 8 | 9 | 10 | 11 | \\8 //\\9 //\\10//\\11// | | | | | +-------------------------- \\// \\// \\// \\// |8 |9 |10 |11 | + ----------------- + "\\" are top and bottom, whereas + "//" are left and right borders + + +Details on the HEALPix can be found at https://iopscience.iop.org/article/10.1086/427976 + +""" + +import sys + +import torch +import torch as th + +sys.path.append('/home/disk/quicksilver/nacc/dlesm/HealPixPad') +have_healpixpad = False +try: + from healpixpad import HEALPixPad # noqa + + have_healpixpad = True +except ImportError: + print("Warning, cannot find healpixpad module") + have_healpixpad = False + + +def healpix_pad(x: torch.Tensor, padding: int, enable_nhwc: bool = False) -> torch.Tensor: + """ + Pad each face consistently with its according neighbors in the HEALPix (see ordering and neighborhoods above). + + :param data: The input tensor of shape [N, F, H, W] where each face is to be padded in its HPX context + :return: The padded tensor where each face's height and width are increased by 2*p + """ + if padding == 0: + return x + + padder = HEALPixPadding(padding, enable_nhwc) + return padder.pad(x) + + +class HEALPixPadding(th.nn.Module): + """ + Padding layer for data on a HEALPix sphere. The requirements for using this layer are as follows: + - The last three dimensions are (face=12, height, width) + - The first four indices in the faces dimension [0, 1, 2, 3] are the faces on the northern hemisphere + - The second four indices in the faces dimension [4, 5, 6, 7] are the faces on the equator + - The last four indices in the faces dimension [8, 9, 10, 11] are the faces on the southern hemisphere + + Orientation and arrangement of the HEALPix faces are outlined above. + """ + + def __init__(self, padding: int, enable_nhwc: bool = False): + """ + Constructor for a HEALPix padding layer. + + :param padding: The padding size + """ + super().__init__() + self.p = padding + self.d = [-2, -1] + self.ret_tl = th.zeros(1, 1, 1) + self.ret_br = th.zeros(1, 1, 1) + self.enable_nhwc = enable_nhwc + if not isinstance(padding, int) or padding < 1: + raise ValueError(f"invalid value for 'padding', expected int > 0 but got {padding}") + + def pad(self, data: th.Tensor) -> th.Tensor: + """ + Pad each face consistently with its according neighbors in the HEALPix (see ordering and neighborhoods above). + + :param data: The input tensor of shape [..., F, H, W] where each face is to be padded in its HPX context + :return: The padded tensor where each face's height and width are increased by 2*p + """ + # Extract the twelve faces (as views of the original tensors) + f00, f01, f02, f03, f04, f05, f06, f07, f08, f09, f10, f11 = [ + torch.squeeze(x, dim=1) for x in th.split(tensor=data, split_size_or_sections=1, dim=1) + ] + + # Assemble the four padded faces on the northern hemisphere + p00 = self.pn(c=f00, t=f01, tl=f02, l=f03, bl=f03, b=f04, br=f08, r=f05, tr=f01) + p01 = self.pn(c=f01, t=f02, tl=f03, l=f00, bl=f00, b=f05, br=f09, r=f06, tr=f02) + p02 = self.pn(c=f02, t=f03, tl=f00, l=f01, bl=f01, b=f06, br=f10, r=f07, tr=f03) + p03 = self.pn(c=f03, t=f00, tl=f01, l=f02, bl=f02, b=f07, br=f11, r=f04, tr=f00) + + # Assemble the four padded faces on the equator + p04 = self.pe(c=f04, t=f00, tl=self.tl(f00, f03), l=f03, bl=f07, b=f11, br=self.br(f11, f08), r=f08, tr=f05) + p05 = self.pe(c=f05, t=f01, tl=self.tl(f01, f00), l=f00, bl=f04, b=f08, br=self.br(f08, f09), r=f09, tr=f06) + p06 = self.pe(c=f06, t=f02, tl=self.tl(f02, f01), l=f01, bl=f05, b=f09, br=self.br(f09, f10), r=f10, tr=f07) + p07 = self.pe(c=f07, t=f03, tl=self.tl(f03, f02), l=f02, bl=f06, b=f10, br=self.br(f10, f11), r=f11, tr=f04) + + # Assemble the four padded faces on the southern hemisphere + p08 = self.ps(c=f08, t=f05, tl=f00, l=f04, bl=f11, b=f11, br=f10, r=f09, tr=f09) + p09 = self.ps(c=f09, t=f06, tl=f01, l=f05, bl=f08, b=f08, br=f11, r=f10, tr=f10) + p10 = self.ps(c=f10, t=f07, tl=f02, l=f06, bl=f09, b=f09, br=f08, r=f11, tr=f11) + p11 = self.ps(c=f11, t=f04, tl=f03, l=f07, bl=f10, b=f10, br=f09, r=f08, tr=f08) + + res = th.stack((p00, p01, p02, p03, p04, p05, p06, p07, p08, p09, p10, p11), dim=1) + + return res + + def pn( + self, + c: th.Tensor, + t: th.Tensor, + tl: th.Tensor, + l: th.Tensor, + bl: th.Tensor, + b: th.Tensor, + br: th.Tensor, + r: th.Tensor, + tr: th.Tensor, + ) -> th.Tensor: + """ + Applies padding to a northern hemisphere face c under consideration of its given neighbors. + + :param c: The central face and tensor that is subject for padding + :param t: The top neighboring face tensor + :param tl: The top left neighboring face tensor + :param l: The left neighboring face tensor + :param bl: The bottom left neighboring face tensor + :param b: The bottom neighboring face tensor + :param br: The bottom right neighboring face tensor + :param r: The right neighboring face tensor + :param tr: The top right neighboring face tensor + :return: The padded tensor p + """ + p = self.p # Padding size + d = self.d # Dimensions for rotations + + # Start with top and bottom to extend the height of the c tensor + c = th.cat((t.rot90(1, d)[..., -p:, :], c, b[..., :p, :]), dim=-2) + + # Construct the left and right pads including the corner faces + left = th.cat((tl.rot90(2, d)[..., -p:, -p:], l.rot90(-1, d)[..., -p:], bl[..., :p, -p:]), dim=-2) + right = th.cat((tr[..., -p:, :p], r[..., :p], br[..., :p, :p]), dim=-2) + + return th.cat((left, c, right), dim=-1) + + def pe( + self, + c: th.Tensor, + t: th.Tensor, + tl: th.Tensor, + l: th.Tensor, + bl: th.Tensor, + b: th.Tensor, + br: th.Tensor, + r: th.Tensor, + tr: th.Tensor, + ) -> th.Tensor: + """ + Applies padding to an equatorial face c under consideration of its given neighbors. + + :param c: The central face and tensor that is subject for padding + :param t: The top neighboring face tensor + :param tl: The top left neighboring face tensor + :param l: The left neighboring face tensor + :param bl: The bottom left neighboring face tensor + :param b: The bottom neighboring face tensor + :param br: The bottom right neighboring face tensor + :param r: The right neighboring face tensor + :param tr: The top right neighboring face tensor + :return: The padded tensor p + """ + p = self.p # Padding size + + # Start with top and bottom to extend the height of the c tensor + c = th.cat((t[..., -p:, :], c, b[..., :p, :]), dim=-2) + + # Construct the left and right pads including the corner faces + left = th.cat((tl[..., -p:, -p:], l[..., -p:], bl[..., :p, -p:]), dim=-2) + right = th.cat((tr[..., -p:, :p], r[..., :p], br[..., :p, :p]), dim=-2) + + return th.cat((left, c, right), dim=-1) + + def ps( + self, + c: th.Tensor, + t: th.Tensor, + tl: th.Tensor, + l: th.Tensor, + bl: th.Tensor, + b: th.Tensor, + br: th.Tensor, + r: th.Tensor, + tr: th.Tensor, + ) -> th.Tensor: + """ + Applies padding to a southern hemisphere face c under consideration of its given neighbors. + + :param c: The central face and tensor that is subject for padding + :param t: The top neighboring face tensor + :param tl: The top left neighboring face tensor + :param l: The left neighboring face tensor + :param bl: The bottom left neighboring face tensor + :param b: The bottom neighboring face tensor + :param br: The bottom right neighboring face tensor + :param r: The right neighboring face tensor + :param tr: The top right neighboring face tensor + :return: The padded tensor p + """ + p = self.p # Padding size + d = self.d # Dimensions for rotations + + # Start with top and bottom to extend the height of the c tensor + c = th.cat((t[..., -p:, :], c, b.rot90(1, d)[..., :p, :]), dim=-2) + + # Construct the left and right pads including the corner faces + left = th.cat((tl[..., -p:, -p:], l[..., -p:], bl[..., :p, -p:]), dim=-2) + right = th.cat((tr[..., -p:, :p], r.rot90(-1, d)[..., :p], br.rot90(2, d)[..., :p, :p]), dim=-2) + + return th.cat((left, c, right), dim=-1) + + def tl(self, t: th.Tensor, l: th.Tensor) -> th.Tensor: # noqa + """ + Assembles the top left corner of a center face in the cases where no according top left face is defined on the + HPX. + + :param t: The face above the center face + :param l: The face left of the center face + :return: The assembled top left corner (only the sub-part that is required for padding) + """ + # ret = th.zeros((*t.shape[:-2], self.p, self.p), dtype=t.dtype, device=t.device) + ret = th.zeros_like(t)[..., : self.p, : self.p] # super ugly but super fast + # tc = t[..., :self.p, :self.p] + # lc = l[..., :self.p, :self.p] + # td = torch.diag_embed(torch.diagonal(tc, dim1=-2, dim2=-1, offset=1), dim1=-2, dim2=-1) + # ld = torch.diag_embed(torch.diagonal(lc, dim1=-2, dim2=-1, offset=-1), dim1=-2, dim2=-1) + # ret2 = torch.triu(tc, diagonal = 1) + torch.tril(lc, diagonal = -1) + 0.5 * (td + ld) + + # Bottom left point + ret[..., -1, -1] = 0.5 * t[..., -1, 0] + 0.5 * l[..., 0, -1] + + # Remaining points + for i in range(1, self.p): + ret[..., -i - 1, -i:] = t[..., -i - 1, :i] # Filling top right above main diagonal + ret[..., -i:, -i - 1] = l[..., :i, -i - 1] # Filling bottom left below main diagonal + ret[..., -i - 1, -i - 1] = 0.5 * t[..., -i - 1, 0] + 0.5 * l[..., 0, -i - 1] # Diagonal + + # print("ALL GOOD", torch.allclose(ret, ret2), ret[0,0,...], ret2[0,0,...]) + # sys.exit(1) + + return ret + + def br(self, b: th.Tensor, r: th.Tensor) -> th.Tensor: + """ + Assembles the bottom right corner of a center face in the cases where no according bottom right face is defined + on the HPX. + + :param b: The face below the center face + :param r: The face right of the center face + :return: The assembled bottom right corner (only the sub-part that is required for padding) + """ + # ret = th.zeros((*b.shape[:-2], self.p, self.p), dtype=b.dtype, device=b.device) + ret = th.zeros_like(b)[..., : self.p, : self.p] + + # Top left point + ret[..., 0, 0] = 0.5 * b[..., 0, -1] + 0.5 * r[..., -1, 0] + + # Remaining points + for i in range(1, self.p): + ret[..., :i, i] = r[..., -i:, i] # Filling top right above main diagonal + ret[..., i, :i] = b[..., i, -i:] # Filling bottom left below main diagonal + ret[..., i, i] = 0.5 * b[..., i, -1] + 0.5 * r[..., -1, i] # Diagonal + + return ret diff --git a/earth2grid/third_party/zephyr/test_healpix_pad.py b/earth2grid/third_party/zephyr/test_healpix_pad.py new file mode 100644 index 0000000..70a3654 --- /dev/null +++ b/earth2grid/third_party/zephyr/test_healpix_pad.py @@ -0,0 +1,13 @@ +import torch + +from earth2grid.third_party.zephyr.healpix import healpix_pad + + +def test_healpix_pad(): + ntile = 12 + nside = 32 + padding = 1 + n = 3 + x = torch.ones([n, ntile, nside, nside]) + out = healpix_pad(x, padding=padding) + assert out.shape == (n, ntile, nside + padding * 2, nside + padding * 2) diff --git a/examples/sphinx_gallery/README.rst b/examples/sphinx_gallery/README.rst new file mode 100644 index 0000000..32e0929 --- /dev/null +++ b/examples/sphinx_gallery/README.rst @@ -0,0 +1,6 @@ +Examples +======== + +The examples are below + +The examples in misc/ are not included in the Sphinx Gallery. diff --git a/examples/sphinx_gallery/hpx2grid.py b/examples/sphinx_gallery/hpx2grid.py new file mode 100644 index 0000000..42655a2 --- /dev/null +++ b/examples/sphinx_gallery/hpx2grid.py @@ -0,0 +1,46 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +""" +HealPIX Image visualization +--------------------------- + +HealPIX maps can be viewed as a 2D image rotated by 45 deg. This is useful for +quick visualization with image viewers without distorting the native pixels of +the image. +""" + +import matplotlib.pyplot as plt +import numpy as np +import torch + +# %% +from matplotlib.colors import Normalize +from PIL import Image + +from earth2grid.healpix import Grid + +grid = Grid(level=8) +lat = torch.tensor(grid.lat) +lat_img = grid.to_image(lat) + +# Use Image to save at full resolution +normalizer = Normalize(vmin=np.nanmin(lat_img), vmax=np.nanmax(lat_img)) +array = normalizer(lat_img) +array = plt.cm.viridis(array) +array = (256 * array).astype("uint8") +# set transparency for nans +array[..., -1] = np.where(np.isnan(lat_img), 0, 255) +image = Image.fromarray(array) +image.save("hpx_grid.png") diff --git a/examples/sphinx_gallery/latlon_to_healpix.py b/examples/sphinx_gallery/latlon_to_healpix.py new file mode 100644 index 0000000..fa66776 --- /dev/null +++ b/examples/sphinx_gallery/latlon_to_healpix.py @@ -0,0 +1,66 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +""" +HealPIX regridding +------------------ + +In this notebook, I demonstrate bilinear regridding onto healpy grids in O(10) +ms. this is a 3 order of magnitude speed-up compared to what Dale has reported. + +Now, lets define a healpix grid with indexing in the XY convention. we convert +to NEST indexing in order to use the `healpy.pix2ang` to get the lat lon +coordinates. This operation is near instant. + +""" +# %% + +import matplotlib.pyplot as plt +import numpy as np +import torch + +import earth2grid + +# level is the resolution +level = 6 +hpx = earth2grid.healpix.Grid(level=level, pixel_order=earth2grid.healpix.XY()) +src = earth2grid.latlon.equiangular_lat_lon_grid(32, 64) +regrid = earth2grid.get_regridder(src, hpx) + + +z = np.cos(np.deg2rad(src.lat)) * np.cos(np.deg2rad(src.lon)) + + +z_torch = torch.as_tensor(z) +z_hpx = regrid(z_torch) + +fig, (a, b) = plt.subplots(2, 1) +a.pcolormesh(src.lon, src.lat, z) +a.set_title("Lat Lon Grid") + +b.scatter(hpx.lon, hpx.lat, c=z_hpx, s=0.1) +b.set_title("Healpix") + +# %% one tile +nside = 2**level +reshaped = z_hpx.reshape(12, nside, nside) +lat_r = hpx.lat.reshape(12, nside, nside) +lon_r = hpx.lon.reshape(12, nside, nside) + +tile = 11 +fig, axs = plt.subplots(3, 4, sharex=True, sharey=True) +axs = axs.ravel() + +for tile in range(12): + axs[tile].pcolormesh(lon_r[tile], lat_r[tile], reshaped[tile]) diff --git a/examples/sphinx_gallery/pyvista_grids.py b/examples/sphinx_gallery/pyvista_grids.py new file mode 100644 index 0000000..ce36d59 --- /dev/null +++ b/examples/sphinx_gallery/pyvista_grids.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +""" +Plot grids with PyVista +----------------------- + +""" +import pyvista as pv + +# %% +import earth2grid + + +def label(mesh, plotter, text): + """ + Add a label above a mesh in a PyVista plot. + + Parameters: + - mesh: The mesh to label. + - plotter: A PyVista plotter instance. + - text: The label text. + - color: The color of the text label. Default is 'white'. + """ + # Calculate the center of the mesh and the top Z-coordinate plus an offset + center = mesh.center + label_pos = [center[0], center[1], mesh.bounds[5] + 0.5] # Offset to place label above the mesh + + # Add the label using point labels for precise 3D positioning + plotter.add_point_labels( + [label_pos], [text], point_size=0, render_points_as_spheres=False, shape_opacity=0, font_size=20 + ) + + +grid = earth2grid.healpix.Grid(level=4) +hpx = grid.to_pyvista() +latlon = earth2grid.latlon.equiangular_lat_lon_grid(32, 64, includes_south_pole=False).to_pyvista() + + +pl = pv.Plotter() +mesh = hpx.translate([0, 2.5, 0]) +pl.add_mesh(mesh, show_edges=True) +label(mesh, pl, "HealPix") + +pl.add_mesh(latlon, show_edges=True) +label(latlon, pl, "Equiangular Lat/Lon") + +pl.camera.position = (5, 0, 5) +pl.show() diff --git a/makefile b/makefile new file mode 100644 index 0000000..f5cb2a1 --- /dev/null +++ b/makefile @@ -0,0 +1,39 @@ +sources = earth2grid + +.PHONY: test format lint unittest coverage pre-commit clean +test: format lint unittest + +.PHONY: license +license: + python tests/_license/header_check.py + +format: license + isort $(sources) tests + black $(sources) tests + +lint: license + pre-commit run --all-files + +unittest: + coverage run --source earth2grid/ -m pytest + coverage run --source earth2grid/ -a -m pytest --doctest-modules earth2grid/ -vv + +coverage: + coverage report + +pre-commit: + pre-commit run --all-files + +clean: + rm -rf .mypy_cache .pytest_cache + rm -rf *.egg-info + rm -rf .tox dist site + rm -rf coverage.xml .coverage + + +.PHONY: docs +docs: + $(MAKE) -C docs html + +push_docs: docs + docs/push_docs.sh diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..29ff5c1 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,89 @@ +[build-system] +requires = ["setuptools>=40.8.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "earth2-grid" +version = "2024.5.1" +description = "Utilities for working with geographic data defined on various grids." +readme = "README.md" +license = { file="LICENSE.txt" } +authors = [ + { name = "NVIDIA", email = "nbrenowitz@nvidia.com" } +] +classifiers = [ + "Development Status :: 2 - Pre-Alpha", + "Intended Audience :: Developers", + "License :: OSI Approved :: BSD License", + "Natural Language :: English", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", +] + +dependencies = [ + "einops>=0.7.0", + "netCDF4>=1.6.5", + "numpy>=1.23.3", + "pyvista>=0.43.2", + "torch>=2.0.1", +] + +[project.urls] +"Homepage" = "https://github.com/waynerv/earth2-grid" + + +[project.optional-dependencies] +test = [ + "pytest>=6.2.4", + "black>=21.5b2", + "isort>=5.8.0", + "mypy>=0.900", + "flake8>=3.9.2", + "flake8-docstrings>=1.6.0", + "pytest-cov>=2.12.0" +] +dev = [ + "tox>=3.20.1", + "pre-commit>=2.12.0", + "virtualenv>=20.2.2", + "pip>=20.3.1", + "twine>=3.3.0", + "toml>=0.10.2", + "bump2version>=1.0.1" +] +doc = [ + "sphinx", + "myst-parser", + "sphinx_gallery", +] + +[tool.black] +line-length = 120 +skip-string-normalization = true +target-version = ['py36', 'py37', 'py38'] +include = '\\.pyi?$' +exclude = ''' +/( + \\.eggs + | \\.git + | \\.hg + | \\.mypy_cache + | \\.tox + | \\.venv + | _build + | buck-out + | build + | dist +)/ +''' + +[tool.isort] +multi_line_output = 3 +include_trailing_comma = true +force_grid_wrap = 0 +use_parentheses = true +ensure_newline_before_comments = true +line_length = 120 +skip_gitignore = true diff --git a/setup.cfg b/setup.cfg new file mode 100644 index 0000000..51d396b --- /dev/null +++ b/setup.cfg @@ -0,0 +1,94 @@ +[flake8] +max-line-length = 120 +max-complexity = 18 +ignore = E203, E266, W503 +docstring-convention = google +per-file-ignores = __init__.py:F401 +exclude = .git, + __pycache__, + setup.py, + build, + dist, + docs, + releases, + .venv, + .tox, + .mypy_cache, + .pytest_cache, + .vscode, + .github, + # By default test codes will be linted. + # tests + +[mypy] +ignore_missing_imports = True + +[coverage:run] +# uncomment the following to omit files during running +#omit = +[coverage:report] +exclude_lines = + pragma: no cover + def __repr__ + if self.debug: + if settings.DEBUG + raise AssertionError + raise NotImplementedError + if 0: + if __name__ == .__main__.: + def main + +[tox:tox] +isolated_build = true +envlist = py36, py37, py38, py39, format, lint, build + +[gh-actions] +python = + 3.9: py39 + 3.8: py38, format, lint, build + 3.7: py37 + 3.6: py36 + +[testenv] +allowlist_externals = pytest +extras = + test +passenv = * +setenv = + PYTHONPATH = {toxinidir} + PYTHONWARNINGS = ignore +commands = + pytest --cov=earth2grid --cov-branch --cov-report=xml --cov-report=term-missing tests + +[testenv:format] +allowlist_externals = + isort + black +extras = + test +commands = + isort earth2grid + black earth2grid tests + +[testenv:lint] +allowlist_externals = + flake8 + mypy +extras = + test +commands = + flake8 earth2grid tests + mypy earth2grid tests + +[testenv:build] +allowlist_externals = + poetry + mkdocs + twine +extras = + doc + dev +commands = + poetry build + mkdocs build + twine check dist/* diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..f8c4e84 --- /dev/null +++ b/setup.py @@ -0,0 +1,69 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import os +import subprocess +from typing import List + +from setuptools import setup +from torch.utils import cpp_extension + + +def get_compiler(): + try: + # Try to get the compiler path from the CC environment variable + # If not set, it will default to gcc (which could be symlinked to clang or g++) + compiler = subprocess.check_output(["gcc", "--version"], universal_newlines=True) + + if "clang" in compiler: + return "clang" + elif "g++" in compiler or "gcc" in compiler: + return "gnu" + else: + return "unknown" + except Exception as e: + print(f"Error detecting compiler: {e}") + return "unknown" + + +compiler_type = get_compiler() +extra_compile_args: List[str] = ["-std=c++20"] + +if compiler_type == "clang": + print("Detected Clang compiler.") + # Additional settings or flags specific to Clang can be added here + extra_compile_args += ["-Wno-error=c++11-narrowing", "-Wno-c++11-narrowing"] +elif compiler_type == "gnu": + print("Detected GNU compiler.") + # Additional settings or flags specific to G++ can be added here +else: + print("Could not detect compiler or unknown compiler detected.") + + +src_files = [ + "earth2grid/csrc/healpix_bare_wrapper.cpp", +] + +setup( + name='earth2grid', + ext_modules=[ + cpp_extension.CppExtension( + 'earth2grid._healpix_bare', + src_files, + extra_compile_args=extra_compile_args, + include_dirs=[os.path.abspath("earth2grid/csrc"), os.path.abspath("earth2grid/third_party/healpix_bare")], + ) + ], + cmdclass={'build_ext': cpp_extension.BuildExtension}, +) diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 0000000..edda28f --- /dev/null +++ b/tests/__init__.py @@ -0,0 +1,15 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +"""Unit test package for earth2-grid.""" diff --git a/tests/_license/config.json b/tests/_license/config.json new file mode 100644 index 0000000..bc55e7c --- /dev/null +++ b/tests/_license/config.json @@ -0,0 +1,13 @@ +{ + "copyright_file": "header.txt", + "dir": "../../", + "exclude-dir": [ + "../../build/", + "../../docs/", + "../../earth2grid/third_party", + "." + ], + "include-ext": [ + ".py" + ] + } diff --git a/tests/_license/header.txt b/tests/_license/header.txt new file mode 100644 index 0000000..a08b2c2 --- /dev/null +++ b/tests/_license/header.txt @@ -0,0 +1,14 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. diff --git a/tests/_license/header_check.py b/tests/_license/header_check.py new file mode 100644 index 0000000..a2e12e0 --- /dev/null +++ b/tests/_license/header_check.py @@ -0,0 +1,141 @@ +# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. +# SPDX-FileCopyrightText: All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""A script to check that copyright headers exists""" + +import itertools +import json +import re +from datetime import datetime +from pathlib import Path + + +def get_top_comments(_data): + """ + Get all lines where comments should exist + """ + lines_to_extract = [] + for i, line in enumerate(_data): + # If empty line, skip + if line in ["", "\n", "", "\r", "\r\n"]: + continue + # If it is a comment line, we should get it + if line.startswith("#"): + lines_to_extract.append(i) + # Assume all copyright headers occur before any import or from statements + # and not enclosed in a comment block + elif "import" in line: + break + elif "from" in line: + break + + comments = [_data[line] for line in lines_to_extract] + return comments + + +def main(): + + with open(Path(__file__).parent.resolve() / Path("config.json")) as f: + config = json.loads(f.read()) + print("License check config:") + print(json.dumps(config, sort_keys=True, indent=4)) + + current_year = int(datetime.today().year) + starting_year = 2023 + python_header_path = Path(__file__).parent.resolve() / Path(config["copyright_file"]) + working_path = Path(__file__).parent.resolve() / Path(config["dir"]) + exts = config["include-ext"] + + with open(python_header_path, "r", encoding="utf-8") as original: + pyheader = original.read().split("\n") + pyheader_lines = len(pyheader) + + # Build list of files to check + exclude_paths = [(Path(__file__).parent / Path(path)).resolve().rglob("*") for path in config["exclude-dir"]] + all_exclude_paths = itertools.chain.from_iterable(exclude_paths) + exclude_filenames = [p for p in all_exclude_paths if p.suffix in exts] + filenames = [p for p in working_path.resolve().rglob("*") if p.suffix in exts] + filenames = [filename for filename in filenames if filename not in exclude_filenames] + problematic_files = [] + gpl_files = [] + + for filename in filenames: + with open(str(filename), "r", encoding="utf-8") as original: + data = original.readlines() + + data = get_top_comments(data) + if data and "# ignore_header_test" in data[0]: + continue + if len(data) < pyheader_lines - 1: + print(f"{filename} has less header lines than the copyright template") + problematic_files.append(filename) + continue + + found = False + for i, line in enumerate(data): + if re.search(re.compile("Copyright.*NVIDIA.*", re.IGNORECASE), line): + found = True + # Check 1st line manually + year_good = False + for year in range(starting_year, current_year + 1): + year_line = pyheader[0].format(CURRENT_YEAR=year) + if year_line in data[i]: + year_good = True + break + year_line_aff = year_line.split(".") + year_line_aff = year_line_aff[0] + " & AFFILIATES." + year_line_aff[1] + if year_line_aff in data[i]: + year_good = True + break + if not year_good: + problematic_files.append(filename) + print(f"{filename} had an error with the year") + break + # while "opyright" in data[i]: + # i += 1 + # for j in range(1, pyheader_lines): + # if pyheader[j] not in data[i + j - 1]: + # problematic_files.append(filename) + # print(f"{filename} missed the line: {pyheader[j]}") + # break + if found: + break + if not found: + print(f"{filename} did not match the regex: `Copyright.*NVIDIA.*`") + problematic_files.append(filename) + + # test if GPL license exists + for lines in data: + if "gpl" in lines.lower(): + gpl_files.append(filename) + break + + if len(problematic_files) > 0: + print("The following files that might not have a copyright header:") + for _file in problematic_files: + print(_file) + if len(gpl_files) > 0: + print("test_header.py found the following files that might have GPL copyright:") + for _file in gpl_files: + print(_file) + assert len(problematic_files) == 0, "header test failed!" + assert len(gpl_files) == 0, "found gpl license, header test failed!" + + print(f"Success: File headers of {len(filenames)} files look good!") + + +if __name__ == "__main__": + main() diff --git a/tests/_regtest_outputs/test_healpix_bare.test_boundaries.out b/tests/_regtest_outputs/test_healpix_bare.test_boundaries.out new file mode 100644 index 0000000..78eb696 --- /dev/null +++ b/tests/_regtest_outputs/test_healpix_bare.test_boundaries.out @@ -0,0 +1,12 @@ +7.07107e-01 +4.56399e-17 +0.00000e+00 +7.45356e-01 +7.07107e-01 +7.45356e-01 +0.00000e+00 +0.00000e+00 +0.00000e+00 +6.66667e-01 +1.00000e+00 +6.66667e-01 diff --git a/tests/_regtest_outputs/test_healpix_bare.test_hpc2loc.out b/tests/_regtest_outputs/test_healpix_bare.test_hpc2loc.out new file mode 100644 index 0000000..dfbaa9f --- /dev/null +++ b/tests/_regtest_outputs/test_healpix_bare.test_hpc2loc.out @@ -0,0 +1,3 @@ +7.07107e-01 +7.07107e-01 +0.00000e+00 diff --git a/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[False-False].out b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[False-False].out new file mode 100644 index 0000000..0332266 --- /dev/null +++ b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[False-False].out @@ -0,0 +1,98 @@ +x +4.11138e-01 +4.11138e-01 +4.11138e-01 +4.11138e-01 +8.41069e-01 +8.41069e-01 +8.41069e-01 +8.41069e-01 +8.41069e-01 +8.41069e-01 +8.41069e-01 +8.41069e-01 +1.23096e+00 +1.23096e+00 +1.23096e+00 +1.23096e+00 +1.23096e+00 +1.23096e+00 +1.23096e+00 +1.23096e+00 +1.57080e+00 +1.57080e+00 +1.57080e+00 +1.57080e+00 +1.57080e+00 +1.57080e+00 +1.57080e+00 +1.57080e+00 +1.91063e+00 +1.91063e+00 +1.91063e+00 +1.91063e+00 +1.91063e+00 +1.91063e+00 +1.91063e+00 +1.91063e+00 +2.30052e+00 +2.30052e+00 +2.30052e+00 +2.30052e+00 +2.30052e+00 +2.30052e+00 +2.30052e+00 +2.30052e+00 +2.73045e+00 +2.73045e+00 +2.73045e+00 +2.73045e+00 +y +7.85398e-01 +2.35619e+00 +3.92699e+00 +5.49779e+00 +3.92699e-01 +1.17810e+00 +1.96350e+00 +2.74889e+00 +3.53429e+00 +4.31969e+00 +5.10509e+00 +5.89049e+00 +0.00000e+00 +7.85398e-01 +1.57080e+00 +2.35619e+00 +3.14159e+00 +3.92699e+00 +4.71239e+00 +5.49779e+00 +3.92699e-01 +1.17810e+00 +1.96350e+00 +2.74889e+00 +3.53429e+00 +4.31969e+00 +5.10509e+00 +5.89049e+00 +0.00000e+00 +7.85398e-01 +1.57080e+00 +2.35619e+00 +3.14159e+00 +3.92699e+00 +4.71239e+00 +5.49779e+00 +3.92699e-01 +1.17810e+00 +1.96350e+00 +2.74889e+00 +3.53429e+00 +4.31969e+00 +5.10509e+00 +5.89049e+00 +7.85398e-01 +2.35619e+00 +3.92699e+00 +5.49779e+00 diff --git a/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[False-True].out b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[False-True].out new file mode 100644 index 0000000..0c91d2c --- /dev/null +++ b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[False-True].out @@ -0,0 +1,98 @@ +x +1.23096e+00 +8.41069e-01 +8.41069e-01 +4.11138e-01 +1.23096e+00 +8.41069e-01 +8.41069e-01 +4.11138e-01 +1.23096e+00 +8.41069e-01 +8.41069e-01 +4.11138e-01 +1.23096e+00 +8.41069e-01 +8.41069e-01 +4.11138e-01 +1.91063e+00 +1.57080e+00 +1.57080e+00 +1.23096e+00 +1.91063e+00 +1.57080e+00 +1.57080e+00 +1.23096e+00 +1.91063e+00 +1.57080e+00 +1.57080e+00 +1.23096e+00 +1.91063e+00 +1.57080e+00 +1.57080e+00 +1.23096e+00 +2.73045e+00 +2.30052e+00 +2.30052e+00 +1.91063e+00 +2.73045e+00 +2.30052e+00 +2.30052e+00 +1.91063e+00 +2.73045e+00 +2.30052e+00 +2.30052e+00 +1.91063e+00 +2.73045e+00 +2.30052e+00 +2.30052e+00 +1.91063e+00 +y +7.85398e-01 +1.17810e+00 +3.92699e-01 +7.85398e-01 +2.35619e+00 +2.74889e+00 +1.96350e+00 +2.35619e+00 +3.92699e+00 +4.31969e+00 +3.53429e+00 +3.92699e+00 +5.49779e+00 +5.89049e+00 +5.10509e+00 +5.49779e+00 +0.00000e+00 +3.92699e-01 +5.89049e+00 +0.00000e+00 +1.57080e+00 +1.96350e+00 +1.17810e+00 +1.57080e+00 +3.14159e+00 +3.53429e+00 +2.74889e+00 +3.14159e+00 +4.71239e+00 +5.10509e+00 +4.31969e+00 +4.71239e+00 +7.85398e-01 +1.17810e+00 +3.92699e-01 +7.85398e-01 +2.35619e+00 +2.74889e+00 +1.96350e+00 +2.35619e+00 +3.92699e+00 +4.31969e+00 +3.53429e+00 +3.92699e+00 +5.49779e+00 +5.89049e+00 +5.10509e+00 +5.49779e+00 diff --git a/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[True-False].out b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[True-False].out new file mode 100644 index 0000000..925fbfa --- /dev/null +++ b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[True-False].out @@ -0,0 +1,98 @@ +x +4.50000e+01 +1.35000e+02 +2.25000e+02 +3.15000e+02 +2.25000e+01 +6.75000e+01 +1.12500e+02 +1.57500e+02 +2.02500e+02 +2.47500e+02 +2.92500e+02 +3.37500e+02 +0.00000e+00 +4.50000e+01 +9.00000e+01 +1.35000e+02 +1.80000e+02 +2.25000e+02 +2.70000e+02 +3.15000e+02 +2.25000e+01 +6.75000e+01 +1.12500e+02 +1.57500e+02 +2.02500e+02 +2.47500e+02 +2.92500e+02 +3.37500e+02 +0.00000e+00 +4.50000e+01 +9.00000e+01 +1.35000e+02 +1.80000e+02 +2.25000e+02 +2.70000e+02 +3.15000e+02 +2.25000e+01 +6.75000e+01 +1.12500e+02 +1.57500e+02 +2.02500e+02 +2.47500e+02 +2.92500e+02 +3.37500e+02 +4.50000e+01 +1.35000e+02 +2.25000e+02 +3.15000e+02 +y +6.64435e+01 +6.64435e+01 +6.64435e+01 +6.64435e+01 +4.18103e+01 +4.18103e+01 +4.18103e+01 +4.18103e+01 +4.18103e+01 +4.18103e+01 +4.18103e+01 +4.18103e+01 +1.94712e+01 +1.94712e+01 +1.94712e+01 +1.94712e+01 +1.94712e+01 +1.94712e+01 +1.94712e+01 +1.94712e+01 +0.00000e+00 +0.00000e+00 +0.00000e+00 +0.00000e+00 +0.00000e+00 +0.00000e+00 +0.00000e+00 +0.00000e+00 +-1.94712e+01 +-1.94712e+01 +-1.94712e+01 +-1.94712e+01 +-1.94712e+01 +-1.94712e+01 +-1.94712e+01 +-1.94712e+01 +-4.18103e+01 +-4.18103e+01 +-4.18103e+01 +-4.18103e+01 +-4.18103e+01 +-4.18103e+01 +-4.18103e+01 +-4.18103e+01 +-6.64435e+01 +-6.64435e+01 +-6.64435e+01 +-6.64435e+01 diff --git a/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[True-True].out b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[True-True].out new file mode 100644 index 0000000..a19f394 --- /dev/null +++ b/tests/_regtest_outputs/test_healpix_bare.test_pix2ang[True-True].out @@ -0,0 +1,98 @@ +x +4.50000e+01 +6.75000e+01 +2.25000e+01 +4.50000e+01 +1.35000e+02 +1.57500e+02 +1.12500e+02 +1.35000e+02 +2.25000e+02 +2.47500e+02 +2.02500e+02 +2.25000e+02 +3.15000e+02 +3.37500e+02 +2.92500e+02 +3.15000e+02 +0.00000e+00 +2.25000e+01 +3.37500e+02 +0.00000e+00 +9.00000e+01 +1.12500e+02 +6.75000e+01 +9.00000e+01 +1.80000e+02 +2.02500e+02 +1.57500e+02 +1.80000e+02 +2.70000e+02 +2.92500e+02 +2.47500e+02 +2.70000e+02 +4.50000e+01 +6.75000e+01 +2.25000e+01 +4.50000e+01 +1.35000e+02 +1.57500e+02 +1.12500e+02 +1.35000e+02 +2.25000e+02 +2.47500e+02 +2.02500e+02 +2.25000e+02 +3.15000e+02 +3.37500e+02 +2.92500e+02 +3.15000e+02 +y +1.94712e+01 +4.18103e+01 +4.18103e+01 +6.64435e+01 +1.94712e+01 +4.18103e+01 +4.18103e+01 +6.64435e+01 +1.94712e+01 +4.18103e+01 +4.18103e+01 +6.64435e+01 +1.94712e+01 +4.18103e+01 +4.18103e+01 +6.64435e+01 +-1.94712e+01 +0.00000e+00 +0.00000e+00 +1.94712e+01 +-1.94712e+01 +0.00000e+00 +0.00000e+00 +1.94712e+01 +-1.94712e+01 +0.00000e+00 +0.00000e+00 +1.94712e+01 +-1.94712e+01 +0.00000e+00 +0.00000e+00 +1.94712e+01 +-6.64435e+01 +-4.18103e+01 +-4.18103e+01 +-1.94712e+01 +-6.64435e+01 +-4.18103e+01 +-4.18103e+01 +-1.94712e+01 +-6.64435e+01 +-4.18103e+01 +-4.18103e+01 +-1.94712e+01 +-6.64435e+01 +-4.18103e+01 +-4.18103e+01 +-1.94712e+01 diff --git a/tests/test_earth2grid.py b/tests/test_earth2grid.py new file mode 100644 index 0000000..f25487d --- /dev/null +++ b/tests/test_earth2grid.py @@ -0,0 +1,35 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. + +"""Tests for `earth2grid` package.""" + +import pytest + + +@pytest.fixture +def response(): + """Sample pytest fixture. + + See more at: http://doc.pytest.org/en/latest/fixture.html + """ + # import requests + # return requests.get('https://github.com/audreyr/cookiecutter-pypackage') + + +def test_content(response): + """Sample pytest test function with the pytest fixture as an argument.""" + # from bs4 import BeautifulSoup + # assert 'GitHub' in BeautifulSoup(response.content).title.string + del response diff --git a/tests/test_healpix.py b/tests/test_healpix.py new file mode 100644 index 0000000..d7d0e63 --- /dev/null +++ b/tests/test_healpix.py @@ -0,0 +1,145 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import matplotlib.pyplot as plt +import numpy as np +import pytest +import torch + +from earth2grid import get_regridder, healpix + + +@pytest.mark.xfail +def test_grid_visualize(): + grid = healpix.Grid(level=4, pixel_order=healpix.XY()) + z = np.cos(10 * np.deg2rad(grid.lat)) + grid.visualize(z) + plt.savefig("test_grid_visualize.png") + + +@pytest.mark.parametrize("origin", list(healpix.Compass)) +def test_grid_healpix_orientations(tmp_path, origin): + + nest_grid = healpix.Grid(level=4, pixel_order=healpix.PixelOrder.NEST) + grid = healpix.Grid(level=4, pixel_order=healpix.XY(origin=origin)) + + nest_lat = nest_grid.lat.reshape([12, -1]) + lat = grid.lat.reshape([12, -1]) + + for i in range(12): + assert set(nest_lat[i]) == set(lat[i]) + + +@pytest.mark.parametrize("rot", range(4)) +def test_rotate_index_same_values(tmp_path, rot): + n = 8 + i = np.arange(12 * n * n) + i_rot = healpix._rotate_index(n, rot, i=i) + i = i.reshape(12, -1) + i_rot = i_rot.reshape(12, -1) + for f in range(12): + assert set(i[f]) == set(i_rot[f]) + + +@pytest.mark.parametrize("rot", range(4)) +def test_rotate_index(rot): + n = 32 + i = np.arange(12 * n * n) + i_rot = healpix._rotate_index(n, rot, i=i) + i_back = healpix._rotate_index(n, 4 - rot, i=i_rot) + np.testing.assert_array_equal(i_back, i) + + +@pytest.mark.parametrize("origin", list(healpix.Compass)) +@pytest.mark.parametrize("clockwise", [True, False]) +def test_reorder(tmp_path, origin, clockwise): + src_grid = healpix.Grid(level=4, pixel_order=healpix.XY(origin=origin, clockwise=clockwise)) + dest_grid = healpix.Grid(level=4, pixel_order=healpix.PixelOrder.NEST) + + z = np.cos(np.deg2rad(src_grid.lat)) * np.cos(np.deg2rad(src_grid.lon)) + z = torch.from_numpy(z) + z_reorder = src_grid.reorder(dest_grid.pixel_order, z) + z_roundtrip = dest_grid.reorder(src_grid.pixel_order, z_reorder) + np.testing.assert_array_equal(z, z_roundtrip) + + +def get_devices(): + devices = [torch.device("cpu")] + if torch.cuda.is_available(): + devices += [torch.device("cuda")] + return devices + + +@pytest.mark.parametrize("origin", list(healpix.Compass)) +@pytest.mark.parametrize("clockwise", [True, False]) +@pytest.mark.parametrize("padding", [0, 1, 2]) +@pytest.mark.parametrize("device", get_devices()) +def test_grid_healpix_pad(tmp_path, origin, clockwise, padding, device): + grid = healpix.Grid(level=4, pixel_order=healpix.XY(origin=origin, clockwise=clockwise)) + hpx_pad_grid = healpix.Grid(level=4, pixel_order=healpix.HEALPIX_PAD_XY) + z = np.cos(np.deg2rad(grid.lat)) * np.cos(np.deg2rad(grid.lon)) + z = torch.from_numpy(z) + regrid = get_regridder(grid, hpx_pad_grid) + z_hpx_pad = regrid(z) + + n = grid._nside() + z = z.view(-1, 12, n, n) + z_hpx_pad = z_hpx_pad.view(-1, 12, n, n) + + padded = healpix.pad(z_hpx_pad.to(device), padding).cpu() + + def grad_abs(z): + fx, fy = np.gradient(z, axis=(-1, -2)) + return np.mean(np.abs(fx)) + np.mean(np.abs(fy)) + + # the padded dtile should not vary much more than the non-padded tile + sigma_padded = grad_abs(padded) + sigma = grad_abs(z) + + if sigma_padded > sigma * 1.1: + + fig, axs = plt.subplots(3, 4) + axs = axs.ravel() + for i in range(12): + ax = axs[i] + ax.pcolormesh(padded[0, i]) + output_path = tmp_path / "output.png" + fig.savefig(output_path.as_posix()) + + raise ValueError( + f"The gradient of the padded data {sigma_padded} is too large. " + f"Examine the padding in the image at {output_path}." + ) + + +def test_to_image(): + grid = healpix.Grid(level=4) + lat = torch.tensor(grid.lat) + lat_img = grid.to_image(lat) + n = 2**grid.level + assert lat_img.shape == (5 * n, 5 * n) + + +def test_conv2d(): + f = 12 + nside = 16 + npix = f * nside * nside + cin = 3 + cout = 4 + n = 1 + + x = torch.ones(n, cin, 1, npix) + weight = torch.zeros(cout, cin, 3, 3) + out = healpix.conv2d(x, weight, padding=(1, 1)) + assert out.shape == (n, cout, 1, npix) diff --git a/tests/test_healpix_bare.py b/tests/test_healpix_bare.py new file mode 100644 index 0000000..c5b9f08 --- /dev/null +++ b/tests/test_healpix_bare.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import numpy +import numpy as np +import pytest +import torch + +import earth2grid.healpix_bare + + +def test_ring2nest(): + n = 8 + i = torch.arange(n * n * 12) + + j = earth2grid.healpix_bare.ring2nest(n, i) + i_round = earth2grid.healpix_bare.nest2ring(n, j) + numpy.testing.assert_array_equal(i, i_round) + + +@pytest.mark.parametrize("nest", [True, False]) +@pytest.mark.parametrize("lonlat", [True, False]) +def test_pix2ang(nest, lonlat, regtest): + n = 2 + i = torch.arange(n * n * 12) + + x, y = earth2grid.healpix_bare.pix2ang(n, i, nest=nest, lonlat=lonlat) + print("x", file=regtest) + np.savetxt(regtest, x, fmt="%.5e") + + print("y", file=regtest) + np.savetxt(regtest, y, fmt="%.5e") + + +def savetxt(file, array): + np.savetxt(file, array, fmt="%.5e") + + +def test_hpc2loc(regtest): + x = torch.tensor([0.0]).double() + y = torch.tensor([0.0]).double() + f = torch.tensor([0]) + + loc = earth2grid.healpix_bare.hpc2loc(x, y, f) + vec = earth2grid.healpix_bare.loc2vec(loc) + for array in vec: + savetxt(regtest, array) + + +def test_boundaries(regtest): + ipix = torch.tensor([0]) + boundaries = earth2grid.healpix_bare.corners(1, ipix, False) + assert not torch.any(torch.isnan(boundaries)), boundaries + assert boundaries.shape == (1, 3, 4) + savetxt(regtest, boundaries.flatten()) + + +def test_get_interp_weights_vector(): + lon = torch.tensor([23, 84, -23]).float() + lat = torch.tensor([0, 12, 67]).float() + pix, weights = earth2grid.healpix_bare.get_interp_weights(8, lon, lat) + assert pix.device == lon.device + assert pix.shape == (4, 3) + assert weights.shape == (4, 3) + + +def test_get_interp_weights_vector_interp_y(): + nside = 16 + inpix = torch.tensor([0, 1, 5, 6]) + + lon, lat = earth2grid.healpix_bare.pix2ang(nside, inpix, lonlat=True) + ay = 0.8 + + lonc = (lon[0] + lon[1]) / 2 + latc = lat[0] * ay + lat[3] * (1 - ay) + + pix, weights = earth2grid.healpix_bare.get_interp_weights(nside, lonc.unsqueeze(0), latc.unsqueeze(0)) + + assert torch.all(pix == inpix[:, None]) + expected_weights = torch.tensor([ay / 2, ay / 2, (1 - ay) / 2, (1 - ay) / 2]).double()[:, None] + assert torch.allclose(weights, expected_weights) diff --git a/tests/test_latlon.py b/tests/test_latlon.py new file mode 100644 index 0000000..2b20940 --- /dev/null +++ b/tests/test_latlon.py @@ -0,0 +1,16 @@ +import torch + +from earth2grid.latlon import equiangular_lat_lon_grid + + +def test_lat_lon_bilinear_regrid_to(): + src = equiangular_lat_lon_grid(15, 30) + dest = equiangular_lat_lon_grid(30, 60) + regrid = src.get_bilinear_regridder_to(dest.lat, dest.lon) + + regrid.float() + lat = torch.broadcast_to(torch.tensor(src.lat), src.shape) + z = torch.tensor(lat).float() + + out = regrid(z) + assert out.shape == dest.shape diff --git a/tests/test_regrid.py b/tests/test_regrid.py new file mode 100644 index 0000000..56e47d5 --- /dev/null +++ b/tests/test_regrid.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# 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. +import unittest + +import matplotlib.pyplot as plt +import numpy as np +import pytest +import torch + +import earth2grid +from earth2grid.latlon import BilinearInterpolator + + +@pytest.mark.parametrize("with_channels", [True, False]) +def test_latlon_regridder(with_channels, tmp_path): + nlat = 30 + nlon = 60 + + src = earth2grid.healpix.Grid(level=6, pixel_order=earth2grid.healpix.XY()) + + lat = np.linspace(-90, 90, nlat + 2)[1:-1] + lon = np.linspace(0, 360, nlon) + dest = earth2grid.latlon.LatLonGrid(lat, lon) + regridder = earth2grid.get_regridder(src, dest) + + z = np.cos(10 * np.deg2rad(src.lat)) + z = torch.from_numpy(z) + if with_channels: + z = torch.stack([z, 0 * z]) + + out = regridder(z) + + if out.ndim == 3: + out = out[0] + + assert out.shape[-2:] == (nlat, nlon) + + expected = np.cos(10 * np.deg2rad(lat))[:, None] + diff = np.mean(np.abs(out.numpy() - expected)) + if diff > 1e-3 * 90 / nlat: + plt.figure() + if with_channels: + out = out[0] + plt.pcolormesh(lon, lat, out - expected) + plt.title("regridded - expected") + plt.colorbar() + image_path = tmp_path / "test_latlon_regridder.png" + plt.savefig(image_path) + raise ValueError(f"{diff} too big. See {image_path}.") + + +@pytest.mark.parametrize("with_channels", [True, False]) +def test_healpix_to_lat_lon(with_channels): + dest = earth2grid.healpix.Grid(level=6, pixel_order=earth2grid.healpix.XY()) + src = earth2grid.latlon.equiangular_lat_lon_grid(33, 64) + regrid = earth2grid.get_regridder(src, dest) + + def f(lat, lon): + lat = torch.from_numpy(lat) + lon = torch.from_numpy(lon) + return torch.cos(torch.deg2rad(lat)) * torch.sin(2 * torch.deg2rad(lon)) + + z = f(src.lat, src.lon) + if with_channels: + z = z[None] + z_regridded = regrid(z) + expected = f(dest.lat, dest.lon) + assert torch.allclose(z_regridded, expected, rtol=0.01) + + +@pytest.mark.parametrize("with_channels", [True, False]) +@pytest.mark.parametrize("level", [4, 5]) +def test_healpix_to_healpix(with_channels, level): + """Test finer healpix interpolation""" + src = earth2grid.healpix.Grid(level=4) + dest = earth2grid.healpix.Grid(level=level) + regrid = earth2grid.get_regridder(src, dest) + + def f(lat, lon): + lat = torch.from_numpy(lat) + lon = torch.from_numpy(lon) + return torch.cos(torch.deg2rad(lat)) * torch.sin(2 * torch.deg2rad(lon)) + + z = f(src.lat, src.lon) + if with_channels: + z = z[None] + z_regridded = regrid(z) + expected = f(dest.lat, dest.lon) + max_diff = torch.max(z_regridded - expected) + print(max_diff) + assert torch.allclose(z_regridded, expected, atol=0.03) + + +@pytest.mark.parametrize("reverse", [True, False]) +def test_latlon_to_latlon(reverse): + nlat = 30 + nlon = 60 + lon = np.linspace(0, 360, nlon) + lat = np.linspace(-90, 90, nlat) + src = earth2grid.latlon.LatLonGrid(lat[::-1] if reverse else lat, lon) + dest = earth2grid.latlon.LatLonGrid(lat, lon) + regrid = earth2grid.get_regridder(src, dest) + regrid.float() + + z = torch.zeros(src.shape) + regrid(z) + + +class TestBilinearInterpolateNonUniform(unittest.TestCase): + def test_interpolation(self): + # Setup + H, W = 5, 5 # Input tensor height and width + input_tensor = torch.arange(1.0, H * W + 1).view(H, W) # Example 2D tensor + x_coords = torch.linspace(-1, 1, steps=W) # Example non-uniform x-coordinates + y_coords = torch.linspace(-1, 1, steps=H) # Example non-uniform y-coordinates + x_query = torch.tensor([0.0]) # Query x-coordinates at the center + y_query = torch.tensor([0.0]) # Query y-coordinates at the center + + # Expected value at the center of a linearly spaced grid + expected = torch.tensor([(H * W + 1) / 2]) + + # Execute + interpolator = BilinearInterpolator(x_coords, y_coords, x_query, y_query) + result = interpolator(input_tensor) + + # Verify + self.assertTrue(torch.allclose(result, expected), "The interpolated value does not match the expected value.") + + def test_raises_error_when_coordinates_not_increasing_x(self): + x_coords = torch.linspace(1, -1, steps=32) # Example non-uniform x-coordinates + y_coords = torch.linspace(-1, 1, steps=32) # Example non-uniform y-coordinates + with self.assertRaises(ValueError): + BilinearInterpolator(x_coords, y_coords, [0], [0]) + + def test_raises_error_when_coordinates_not_increasing_y(self): + x_coords = torch.linspace(-1, 1, steps=32) # Example non-uniform x-coordinates + y_coords = torch.linspace(1, -1, steps=32) # Example non-uniform y-coordinates + with self.assertRaises(ValueError): + BilinearInterpolator(x_coords, y_coords, [0], [0]) + + def test_interpolation_func(self): + # Setup + H, W = 32, 32 # Input tensor height and width + + def func(x, y): + + return 10 * x + 5 * y**2 + 4 + + x_coords = torch.linspace(-1, 1, steps=W) # Example non-uniform x-coordinates + y_coords = torch.linspace(-1, 1, steps=H) # Example non-uniform y-coordinates + x_query = torch.tensor([0.0, 0.5, 0.25]) # Query x-coordinates at the center + y_query = torch.tensor([0.0, 0.0, -0.4]) # Query y-coordinates at the center + + input_tensor = func(x_coords, y_coords[:, None]) + + # Expected value at the center of a linearly spaced grid + expected = func(x_query, y_query) + + # Execute + interpolator = BilinearInterpolator(x_coords, y_coords, x_query, y_query) + result = interpolator(input_tensor) + + # Verify + self.assertTrue(torch.allclose(result, expected, rtol=0.01)) + + if torch.cuda.is_available() and torch.cuda.device_count() > 0: + interpolator.cuda() + interpolator(input_tensor.cuda())