diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml index 8b3312a..349d67e 100644 --- a/.github/workflows/bench.yml +++ b/.github/workflows/bench.yml @@ -13,10 +13,8 @@ jobs: runs-on: ${{ matrix.os }} strategy: matrix: - #os: [windows-latest, macOS-latest, ubuntu-latest] - #python-version: [3.8, 3.9, 3.10] os: [ubuntu-latest] - python-version: [3.8] + python-version: [3.12] steps: - name: "Set up Python ${{ matrix.python-version }} on ${{ matrix.os }}" uses: actions/setup-python@v4 diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 512b445..a032405 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -45,8 +45,7 @@ jobs: maturin build --release --compatibility manylinux2014 maturin develop --release pip install -r requirements-dev.txt - python -m pytest python/cdshealpix/tests/test_nested_healpix.py - python -m pytest python/cdshealpix/tests/test_ring_healpix.py + python -m pytest -v deactivate rm -r cdshealpixenv done @@ -96,10 +95,8 @@ jobs: python-version: ['3.8', '3.9', '3.10', '3.11', '3.12'] steps: # Checkout the project - - name: "Checkout branch ${{ github.head_ref }}" + - name: "Checkout branch" uses: actions/checkout@v4 - with: - ref: ${{ github.head_ref }} # Test Rust code #- name: Run rust wrapper tests # run: cargo test --verbose -- --nocapture @@ -125,10 +122,7 @@ jobs: # Install dependencies pip install -r requirements-dev.txt # Run tests - cd python - python -m pytest -v -s cdshealpix/tests/test_nested_healpix.py - python -m pytest -v -s cdshealpix/tests/test_ring_healpix.py - cd .. + python -m pytest -v # Clean #pip freeze > requirements-uninstall.txt #pip uninstall -r requirements-uninstall.txt -y @@ -142,10 +136,8 @@ jobs: python-version: ['3.8', '3.9', '3.10', '3.11', '3.12'] steps: # Checkout the project - - name: "Checkout branch ${{ github.head_ref }}" + - name: "Checkout branch" uses: actions/checkout@v4 - with: - ref: ${{ github.head_ref }} # Test Rust code #- name: Run rust wrapper tests # run: cargo test --verbose -- --nocapture @@ -170,20 +162,15 @@ jobs: # Install dependencies pip install -r requirements-dev.txt # Run tests - cd python - python -m pytest -v "cdshealpix\tests\test_nested_healpix.py" - python -m pytest -v "cdshealpix\tests\test_ring_healpix.py" - cd .. + python -m pytest -v deactivate - # Build the doc and run the tests in the doc (only for python 3.9 on ubuntu) + # Build the doc and run the tests in the doc (only for python 3.11 on ubuntu) test-doc: runs-on: ubuntu-latest steps: - - name: "Checkout branch ${{ github.head_ref }}" + - name: "Checkout branch" uses: actions/checkout@v4 - with: - ref: ${{ github.head_ref }} - name: "Set up Python 3.11 on Ubuntu" uses: actions/setup-python@v4 with: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 8368a56..6428619 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -13,7 +13,6 @@ repos: rev: 'v1.10.0' hooks: - id: python-no-eval - - id: rst-backticks # For python files - repo: https://github.com/psf/black # Code style diff --git a/CHANGELOG.md b/CHANGELOG.md index cf64426..57c0f39 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ * new features `box_search` and `zone_search` in `cdshealpix.nested` * `cdshaelpix.nested.vertices` can now take depth as a `numpy.ndarray` instead of only accepting a single depth +* new module `skymap` added. +* read/write, and plot nested all-sky skymaps in the implicit scheme from fits files with + `Skymap.from_fits`, `Skymap.from_array`, `Skymap.quick_plot`, and `Skymap.to_fits` ### Fixed @@ -18,6 +21,10 @@ * `nested.healpix_to_lonlat`, failed into rust panic for `dx=1` or `dy=1`. This is now indicated in the documentation and is catched in a `ValueError` on the python side. +### Changed + +* `matplotlib` is now an optional dependency, to plot previews of skymaps. + ## 0.6.5 Released 2023-11-28 diff --git a/Cargo.toml b/Cargo.toml index 4e01d57..885501b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,18 +20,19 @@ name = "cdshealpix" crate-type = ["cdylib"] [dependencies] -healpix = { package = "cdshealpix", version = "0.6" } +healpix = { rev="5ab172b4d1d206c973f9f170b14509f4982e0090", package = "cdshealpix", git = 'https://github.com/cds-astro/cds-healpix-rust' } +mapproj = "0.3.0" rayon = "1.10" [dependencies.numpy] -version = "0.20" +version = "0.22" [dependencies.pyo3] -version = "0.20" +version = "0.22" features = ["extension-module"] [dependencies.ndarray] -version = "0.15" +version = "0.16" default-features = false # do not include the default features, and optionally # cherry-pick individual features features = ["rayon"] diff --git a/docs/api.rst b/docs/api.rst index b899e5f..12f01ed 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -87,4 +87,14 @@ While in the nested scheme, ``nside`` is a power of two vertices vertices_skycoord +cdshealpix.skymap +~~~~~~~~~~~~~~~~~ + +This module provides a minimal interface to interact with Skymaps, as defined in the +`data format for gamma ray astronomy specification +`_. + +.. autoclass:: cdshealpix.skymap.Skymap + :members: + .. _cdshealpix: https://github.com/cds-astro/cds-healpix-python diff --git a/docs/conf.py b/docs/conf.py index 1959d34..305afd9 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -65,6 +65,8 @@ ] default_role = "py:obj" numpydoc_class_members_toctree = False +numpydoc_show_class_members = False +autosummary_generate = False # Add any paths that contain templates here, relative to this directory. templates_path = ["_templates"] bibtex_bibfiles = ["references.bib"] diff --git a/docs/examples/examples.rst b/docs/examples/examples.rst index c26a06f..b45e8f3 100644 --- a/docs/examples/examples.rst +++ b/docs/examples/examples.rst @@ -42,11 +42,23 @@ Zone search ----------- In this example, we get the ``ipix`` and ``depth`` in a zone and plot them by combining -:external:ref:`~cdshealpix.nested.vertices`_ with :external:ref:`~matplotlib.path.Polygon`_ +`cdshealpix.nested.vertices` with `matplotlib.path.Polygon` .. plot:: examples/search_methods/zone_search.py :include-source: +Skymaps +======= + +The skymap sub-module allows to manipulate easily all-sky skymaps in the nested ordering +and implicit schema. +The class can be instantiated either from a fits file, with `Skymap.from_fits`, or +directly with a numpy `numpy.array` containing the values associated to each HEALPix +pixel. + +.. plot:: examples/skymaps/skymaps.py + :include-source: + Notebook examples ================= diff --git a/docs/examples/skymaps/skymap.fits b/docs/examples/skymaps/skymap.fits new file mode 100644 index 0000000..db19a0f Binary files /dev/null and b/docs/examples/skymaps/skymap.fits differ diff --git a/docs/examples/skymaps/skymaps.py b/docs/examples/skymaps/skymaps.py new file mode 100644 index 0000000..42e532d --- /dev/null +++ b/docs/examples/skymaps/skymaps.py @@ -0,0 +1,7 @@ +"""Read and plots a quick preview of a skymap in a FITS file.""" + +from cdshealpix.skymap import Skymap + +skymap = Skymap.from_fits("skymap.fits") +print(skymap.depth) +skymap.quick_plot() diff --git a/pyproject.toml b/pyproject.toml index 9cdf3bc..8adaf29 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -2,18 +2,31 @@ [project] name = "cdshealpix" requires-python = ">=3.8" -# https://numpy.org/neps/nep-0029-deprecation_policy.html -# https://docs.astropy.org/en/stable/changelog.html dependencies = [ "astropy<5.3; python_version == '3.8'", "astropy; python_version > '3.8'" ] classifiers = [ "Programming Language :: Rust", + "Programming Language :: Python :: 3.8", + "Programming Language :: Python :: 3.9", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", ] +description = "A healpix manipulation library." +readme = "README.md" +keywords = ["astronomy", "healpix"] + +[project.optional-dependencies] +plot = ["matplotlib"] [project.urls] repository = "https://github.com/cds-astro/cds-healpix-python" +documentation = "https://cds-astro.github.io/cds-healpix-python/" +issues = "https://github.com/cds-astro/cds-healpix-python/issues" +changelog = "https://github.com/cds-astro/cds-healpix-python/blob/master/CHANGELOG.md" + # Build a cdshealpix-x.x.x.tar.gz containing sources (from maturin). [build-system] @@ -44,7 +57,7 @@ output-format = "grouped" show-source = true ignore-init-module-imports = true target-version = "py37" -extend-select = ["SIM", "D", "UP", "N", "S", "B", "A", "C4", "ICN", "RET", "ARG", "PD", "PGH", "RUF"] +extend-select = ["SIM", "D", "UP", "N", "S", "B", "A", "C4", "ICN", "RET", "ARG", "PGH", "RUF"] extend-ignore = ["E501"] # E501: line length (done by black in our case) exclude = ["conf.py"] @@ -67,6 +80,6 @@ max-string-length = 20 convention = "numpy" # Accepts: "google", "numpy", or "pep257" [tool.pytest.ini_options] -addopts = "--doctest-modules" +addopts = "--doctest-modules --ignore-glob=python/cdshealpix/tests/test_bench*" doctest_optionflags = "NORMALIZE_WHITESPACE" testpaths = "python" \ No newline at end of file diff --git a/python/cdshealpix/skymap/__init__.py b/python/cdshealpix/skymap/__init__.py new file mode 100644 index 0000000..0421c99 --- /dev/null +++ b/python/cdshealpix/skymap/__init__.py @@ -0,0 +1 @@ +from .skymap import Skymap # noqa: F401 diff --git a/python/cdshealpix/skymap/skymap.py b/python/cdshealpix/skymap/skymap.py new file mode 100644 index 0000000..bf397d4 --- /dev/null +++ b/python/cdshealpix/skymap/skymap.py @@ -0,0 +1,159 @@ +"""Manipulation of skymaps. + +SkyMaps are described in _ +This sub-module supports skymaps in the nested scheme, and in the implicit format. +The coordinates system should be 'CEL'. +""" +from .. import cdshealpix + +from pathlib import Path +from typing import Union + +try: + import matplotlib.pyplot as plt +except ImportError: + _matplotlib_missing = True +import numpy as np + + +class Skymap: + """A Skymap, containing values to associate to healpix cells.""" + + def __init__(self, values): + self.values = values + + @property + def depth(self): + """The depth of the skymap. + + Avoids the costly log calculation. + + Returns + ------- + `int` + The depth of the skymap. + + Examples + -------- + >>> from cdshealpix.skymap import Skymap + >>> map = Skymap.from_array([0]*12) + >>> map.depth + 0 + """ + return cdshealpix.depth_skymap(self.values) + + @classmethod + def from_fits(cls, path: Union[str, Path]): + """Read a skymap in the nested schema from a FITS file. + + This reader supports files which are: + + - all sky maps + - in the nested scheme + - and the implicit format + + Parameters + ---------- + path : str, `pathlib.Path` + The file's path. + + Returns + ------- + `Skymap` + A skymap. Its values are in a numpy array which data type in inferred from + the FITS header. + """ + return cls(cdshealpix.read_skymap(str(path))) + + @classmethod + def from_array(cls, values): + """Create a skymap from an array. + + Parameters + ---------- + values : `numpy.array` + An array-like object. It should be one-dimensional, and its length should be + the number of cells in a HEALPix order. + It should be in the nested ordering (not tested). + + Returns + ------- + `Skymap` + A skymap object. + + Examples + -------- + >>> from cdshealpix.skymap import Skymap + >>> import numpy as np + >>> skymap =Skymap.from_array(np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], dtype=np.uint8)) + """ + # only makes a copy if it was not C-contiguous in the first place + values = np.ascontiguousarray(values) + if values.ndim != 1: + raise ValueError( + "Skymap values should be one-dimensional. Got an array of " + f"shape {values.shape}." + ) + n = int(len(values) / 12) + # test if it is a power of two (1000 & 0111 = 0000) + if n & (n - 1) != 0 or n == 0: + raise ValueError( + "The length of values should be a valid number of cells in " + "a given HEALPix order, i.e something like 12, 48, 192... " + f"Got '{len(values)}'." + ) + + if values.dtype not in ( + np.float64, + np.float32, + np.int64, + np.int32, + np.int16, + np.uint8, + float, + int, + ): + raise ValueError( + f"The accepted types are f64, f32, i64, i32, u8. Got '{values.dtype}'." + ) + return cls(values) + + def to_fits(self, path): + """Write a Skymap in a fits file. + + Parameters + ---------- + path : `str`, `pathlib.Path` + The file's path. + """ + cdshealpix.write_skymap(self.values, str(path)) + + def quick_plot(self, *, size=256, convert_to_gal=True, path=None): + """Preview a skymap in the Mollweide projection. + + Parameters + ---------- + size : `int`, optional + The size of the plot in the y-axis in pixels. + It fixes the resolution of the image. By default 256 + convert_to_gal : `bool`, optional + Should the image be converted into a galactic frame? by default True + path : `str` or `pathlib.Path`, optional + If different from none, the image will not only be displayed, but also saved + at the given location. By default None + """ + if _matplotlib_missing: + raise ModuleNotFoundError( + "matplotlib is mandatory to use 'quick_plot'. " + "See https://matplotlib.org/ for installation " + "instructions." + ) + img = cdshealpix.pixels_skymap(self.values, size, convert_to_gal) + fig = plt.imshow(img) + plt.axis("off") + fig.axes.get_xaxis().set_visible(False) + fig.axes.get_yaxis().set_visible(False) + if path: + plt.savefig(path, bbox_inches="tight", pad_inches=0, transparent=True) + plt.show() diff --git a/python/cdshealpix/tests/resources/skymap.fits b/python/cdshealpix/tests/resources/skymap.fits new file mode 100644 index 0000000..db19a0f Binary files /dev/null and b/python/cdshealpix/tests/resources/skymap.fits differ diff --git a/python/cdshealpix/tests/test_skymaps.py b/python/cdshealpix/tests/test_skymaps.py new file mode 100644 index 0000000..b24844d --- /dev/null +++ b/python/cdshealpix/tests/test_skymaps.py @@ -0,0 +1,45 @@ +from pathlib import Path +from tempfile import NamedTemporaryFile + +import numpy as np +import pytest + +from ..skymap import Skymap +from .. import cdshealpix + +path_to_test_skymap = Path(__file__).parent.resolve() / "resources" / "skymap.fits" + + +def test_read(): + values = Skymap.from_fits(path_to_test_skymap).values + assert values.dtype == np.int32 + assert len(values) == 49152 + + +def test_read_write_read_conservation(): + skymap = Skymap.from_fits(path_to_test_skymap) + with NamedTemporaryFile() as fp: + skymap.to_fits(fp.name) + skymap2 = Skymap.from_fits(fp.name) + assert all(skymap.values == skymap2.values) + + +def test_plot(): + skymap = Skymap.from_array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]) + img = cdshealpix.pixels_skymap(skymap.values, 256, True) + assert img.shape == (256, 512, 4) + + +def test_depth(): + n = 12 + skymap = Skymap.from_array(np.zeros(12 * 4**n)) + assert skymap.depth == n + + +def test_from_array(): + with pytest.raises(ValueError, match="The length of values should be*"): + Skymap.from_array(np.zeros(3)) + with pytest.raises(ValueError, match="The accepted types are*"): + Skymap.from_array(["a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l"]) + with pytest.raises(ValueError, match="Skymap values should be one-dimensional*"): + Skymap.from_array([[1, 2, 3], [1, 2, 3]]) diff --git a/requirements-dev.txt b/requirements-dev.txt index 0ece12f..f8da8a5 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,5 +1,6 @@ astropy<5.3; python_version == '3.8' astropy; python_version > '3.8' +astropy_healpix +matplotlib pre-commit==2.21.* pytest -astropy_healpix \ No newline at end of file diff --git a/src/lib.rs b/src/lib.rs index 0f9b2d4..068f310 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,18 +1,15 @@ -#[cfg(feature = "rayon")] -extern crate healpix; - -extern crate ndarray; - -extern crate numpy; -extern crate pyo3; -extern crate rayon; - use ndarray::{Array1, Zip}; -use numpy::{IntoPyArray, PyArray1, PyArrayDyn}; -use pyo3::prelude::{pymodule, Py, PyModule, PyResult, Python}; +use numpy::{IntoPyArray, PyArray1, PyArrayDyn, PyArrayMethods, PyReadonlyArrayDyn}; +use pyo3::{ + prelude::{pymodule, Bound, PyModule, PyResult, Python}, + types::PyModuleMethods, + wrap_pyfunction, +}; use healpix::compass_point::{Cardinal, MainWind, Ordinal}; +mod skymap_functions; + /// This uses rust-numpy for numpy interoperability between /// Python and Rust. /// PyArrayDyn rust-numpy array types are converted to ndarray @@ -23,16 +20,25 @@ use healpix::compass_point::{Cardinal, MainWind, Ordinal}; /// operate on them element-wisely. This is done in parallel using the /// ndarray-parallel crate that offers the par_for_each method on zipped arrays. #[pymodule] -fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { - // wrapper of to_ring and from_ring +fn cdshealpix(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { + // add skymap pyfunctions here + m.add_function(wrap_pyfunction!(skymap_functions::read_skymap, m)?) + .unwrap(); + m.add_function(wrap_pyfunction!(skymap_functions::write_skymap, m)?) + .unwrap(); + m.add_function(wrap_pyfunction!(skymap_functions::pixels_skymap, m)?) + .unwrap(); + m.add_function(wrap_pyfunction!(skymap_functions::depth_skymap, m)?) + .unwrap(); + // wrapper of to_ring and from_ring #[pyfn(m)] #[pyo3(name = "to_ring")] - unsafe fn to_ring( + unsafe fn to_ring<'a>( _py: Python, depth: u8, - ipix: &PyArrayDyn, - ipix_ring: &PyArrayDyn, + ipix: &Bound<'a, PyArrayDyn>, + ipix_ring: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let ipix = ipix.as_array(); @@ -65,11 +71,11 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { } #[pyfn(m)] - unsafe fn from_ring( + unsafe fn from_ring<'a>( _py: Python, depth: u8, - ipix_ring: &PyArrayDyn, - ipix: &PyArrayDyn, + ipix_ring: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let ipix_ring = ipix_ring.as_array(); @@ -104,14 +110,14 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// wrapper of `lonlat_to_healpix` #[pyfn(m)] - unsafe fn lonlat_to_healpix( + unsafe fn lonlat_to_healpix<'a>( _py: Python, - depth: &PyArrayDyn, - lon: &PyArrayDyn, - lat: &PyArrayDyn, - ipix: &PyArrayDyn, - dx: &PyArrayDyn, - dy: &PyArrayDyn, + depth: &Bound<'a, PyArrayDyn>, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, + dx: &Bound<'a, PyArrayDyn>, + dy: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let lon = lon.as_array(); @@ -161,14 +167,14 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { } #[pyfn(m)] - unsafe fn lonlat_to_healpix_ring( + unsafe fn lonlat_to_healpix_ring<'a>( _py: Python, - nside: &PyArrayDyn, - lon: &PyArrayDyn, - lat: &PyArrayDyn, - ipix: &PyArrayDyn, - dx: &PyArrayDyn, - dy: &PyArrayDyn, + nside: &Bound<'a, PyArrayDyn>, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, + dx: &Bound<'a, PyArrayDyn>, + dy: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let lon = lon.as_array(); @@ -218,14 +224,14 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// wrapper of `healpix_to_lonlat` #[pyfn(m)] - unsafe fn healpix_to_lonlat( + unsafe fn healpix_to_lonlat<'a>( _py: Python, - depth: &PyArrayDyn, - ipix: &PyArrayDyn, + depth: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, dx: f64, dy: f64, - lon: &PyArrayDyn, - lat: &PyArrayDyn, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let mut lon = lon.as_array_mut(); @@ -266,14 +272,14 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { } #[pyfn(m)] - unsafe fn healpix_to_lonlat_ring( + unsafe fn healpix_to_lonlat_ring<'a>( _py: Python, - nside: &PyArrayDyn, - ipix: &PyArrayDyn, + nside: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, dx: f64, dy: f64, - lon: &PyArrayDyn, - lat: &PyArrayDyn, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let mut lon = lon.as_array_mut(); @@ -315,12 +321,12 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// wrapper of `healpix_to_xy` #[pyfn(m)] - unsafe fn healpix_to_xy( + unsafe fn healpix_to_xy<'a>( _py: Python, - ipix: &PyArrayDyn, - depth: &PyArrayDyn, - x: &PyArrayDyn, - y: &PyArrayDyn, + ipix: &Bound<'a, PyArrayDyn>, + depth: &Bound<'a, PyArrayDyn>, + x: &Bound<'a, PyArrayDyn>, + y: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let mut x = x.as_array_mut(); @@ -363,12 +369,12 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { } #[pyfn(m)] - unsafe fn healpix_to_xy_ring( + unsafe fn healpix_to_xy_ring<'a>( _py: Python, - nside: &PyArrayDyn, - ipix: &PyArrayDyn, - x: &PyArrayDyn, - y: &PyArrayDyn, + nside: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, + x: &Bound<'a, PyArrayDyn>, + y: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let mut x = x.as_array_mut(); @@ -410,12 +416,12 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// wrapper of `lonlat_to_xy` #[pyfn(m)] - unsafe fn lonlat_to_xy( + unsafe fn lonlat_to_xy<'a>( _py: Python, - lon: &PyArrayDyn, - lat: &PyArrayDyn, - x: &PyArrayDyn, - y: &PyArrayDyn, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, + x: &Bound<'a, PyArrayDyn>, + y: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let mut x = x.as_array_mut(); @@ -457,12 +463,12 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// wrapper of `xy_to_lonlat` #[pyfn(m)] - unsafe fn xy_to_lonlat( + unsafe fn xy_to_lonlat<'a>( _py: Python, - x: &PyArrayDyn, - y: &PyArrayDyn, - lon: &PyArrayDyn, - lat: &PyArrayDyn, + x: &Bound<'a, PyArrayDyn>, + y: &Bound<'a, PyArrayDyn>, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let x = x.as_array(); @@ -504,13 +510,13 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// wrapper of `vertices` #[pyfn(m)] - unsafe fn vertices( + unsafe fn vertices<'a>( _py: Python, - depth: &PyArrayDyn, - ipix: &PyArrayDyn, + depth: &Bound<'a, PyArrayDyn>, + ipix: &Bound<'a, PyArrayDyn>, step: usize, - lon: &PyArrayDyn, - lat: &PyArrayDyn, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let depth = depth.as_array(); @@ -605,13 +611,13 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { } #[pyfn(m)] - unsafe fn vertices_ring( + unsafe fn vertices_ring<'a>( _py: Python, nside: u32, - ipix: &PyArrayDyn, + ipix: &Bound<'a, PyArrayDyn>, _step: usize, - lon: &PyArrayDyn, - lat: &PyArrayDyn, + lon: &Bound<'a, PyArrayDyn>, + lat: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let ipix = ipix.as_array(); @@ -672,11 +678,11 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// The given array must be of size 9 /// `[S, SE, E, SW, C, NE, W, NW, N]` #[pyfn(m)] - unsafe fn neighbours( + unsafe fn neighbours<'a>( _py: Python, depth: u8, - ipix: &PyArrayDyn, - neighbours: &PyArrayDyn, + ipix: &Bound<'a, PyArrayDyn>, + neighbours: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let ipix = ipix.as_array(); @@ -761,14 +767,18 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// Cone search #[pyfn(m)] fn cone_search( - py: Python, + py: Python<'_>, depth: u8, delta_depth: u8, lon: f64, lat: f64, radius: f64, flat: bool, - ) -> (Py>, Py>, Py>) { + ) -> ( + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + ) { let bmoc = healpix::nested::cone_coverage_approx_custom(depth, delta_depth, lon, lat, radius); let (ipix, depth, fully_covered) = if flat { @@ -778,16 +788,16 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { }; ( - ipix.into_pyarray(py).to_owned(), - depth.into_pyarray(py).to_owned(), - fully_covered.into_pyarray(py).to_owned(), + ipix.into_pyarray_bound(py), + depth.into_pyarray_bound(py), + fully_covered.into_pyarray_bound(py), ) } /// Elliptical cone search #[pyfn(m)] fn elliptical_cone_search( - py: Python, + py: Python<'_>, depth: u8, delta_depth: u8, lon: f64, @@ -796,7 +806,11 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { b: f64, pa: f64, flat: bool, - ) -> (Py>, Py>, Py>) { + ) -> ( + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + ) { let bmoc = healpix::nested::elliptical_cone_coverage_custom(depth, delta_depth, lon, lat, a, b, pa); @@ -807,21 +821,25 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { }; ( - ipix.into_pyarray(py).to_owned(), - depth.into_pyarray(py).to_owned(), - fully_covered.into_pyarray(py).to_owned(), + ipix.into_pyarray_bound(py), + depth.into_pyarray_bound(py), + fully_covered.into_pyarray_bound(py), ) } /// Polygon search #[pyfn(m)] - unsafe fn polygon_search( - py: Python, + unsafe fn polygon_search<'a>( + py: Python<'a>, depth: u8, - lon: &PyArrayDyn, - lat: &PyArrayDyn, + lon: PyReadonlyArrayDyn<'a, f64>, + lat: PyReadonlyArrayDyn<'a, f64>, flat: bool, - ) -> (Py>, Py>, Py>) { + ) -> ( + Bound<'a, PyArray1>, + Bound<'a, PyArray1>, + Bound<'a, PyArray1>, + ) { let lon = lon.as_array(); let lat = lat.as_array(); @@ -842,9 +860,9 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { }; ( - ipix.into_pyarray(py).to_owned(), - depth.into_pyarray(py).to_owned(), - fully_covered.into_pyarray(py).to_owned(), + ipix.into_pyarray_bound(py), + depth.into_pyarray_bound(py), + fully_covered.into_pyarray_bound(py), ) } @@ -861,7 +879,7 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// * ``pa`` -rotation angle in degrees #[pyfn(m)] fn box_search( - py: Python, + py: Python<'_>, depth: u8, lon: f64, lat: f64, @@ -869,7 +887,11 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { b: f64, pa: f64, flat: bool, - ) -> (Py>, Py>, Py>) { + ) -> ( + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + ) { let bmoc = healpix::nested::box_coverage(depth, lon, lat, a, b, pa); let (ipix, depth, fully_covered) = if flat { @@ -879,9 +901,9 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { }; ( - ipix.into_pyarray(py).to_owned(), - depth.into_pyarray(py).to_owned(), - fully_covered.into_pyarray(py).to_owned(), + ipix.into_pyarray_bound(py), + depth.into_pyarray_bound(py), + fully_covered.into_pyarray_bound(py), ) } @@ -897,14 +919,18 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { /// * ``lat_max`` - east north corner latitude #[pyfn(m)] fn zone_search( - py: Python, + py: Python<'_>, depth: u8, lon_min: f64, lat_min: f64, lon_max: f64, lat_max: f64, flat: bool, - ) -> (Py>, Py>, Py>) { + ) -> ( + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + Bound<'_, PyArray1>, + ) { let bmoc = healpix::nested::zone_coverage(depth, lon_min, lat_min, lon_max, lat_max); let (ipix, depth, fully_covered) = if flat { @@ -914,20 +940,20 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { }; ( - ipix.into_pyarray(py).to_owned(), - depth.into_pyarray(py).to_owned(), - fully_covered.into_pyarray(py).to_owned(), + ipix.into_pyarray_bound(py), + depth.into_pyarray_bound(py), + fully_covered.into_pyarray_bound(py), ) } #[pyfn(m)] - unsafe fn external_neighbours( + unsafe fn external_neighbours<'a>( _py: Python, depth: u8, delta_depth: u8, - ipix: &PyArrayDyn, - corners: &PyArrayDyn, - edges: &PyArrayDyn, + ipix: &Bound<'a, PyArrayDyn>, + corners: &Bound<'a, PyArrayDyn>, + edges: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let ipix = ipix.as_array(); @@ -1045,20 +1071,19 @@ fn cdshealpix(_py: Python, m: &PyModule) -> PyResult<()> { // Bilinear interpolation // //////////////////////////// #[pyfn(m)] - unsafe fn bilinear_interpolation( - _py: Python, + fn bilinear_interpolation<'a>( depth: u8, - lon: &PyArrayDyn, - lat: &PyArrayDyn, - ipix: &PyArrayDyn, - weights: &PyArrayDyn, + lon: PyReadonlyArrayDyn<'a, f64>, + lat: PyReadonlyArrayDyn<'a, f64>, + ipix: &Bound<'a, PyArrayDyn>, + weights: &Bound<'a, PyArrayDyn>, nthreads: u16, ) -> PyResult<()> { let lon = lon.as_array(); let lat = lat.as_array(); - let mut ipix = ipix.as_array_mut(); - let mut weights = weights.as_array_mut(); + let mut ipix = unsafe { ipix.as_array_mut() }; + let mut weights = unsafe { weights.as_array_mut() }; let layer = healpix::nested::get(depth); #[cfg(not(target_arch = "wasm32"))] diff --git a/src/skymap_functions.rs b/src/skymap_functions.rs new file mode 100644 index 0000000..29c37a4 --- /dev/null +++ b/src/skymap_functions.rs @@ -0,0 +1,232 @@ +extern crate healpix; +extern crate mapproj; + +use std::fs::File; +use std::io::BufWriter; + +use numpy::{ + IntoPyArray, Ix3, NotContiguousError, PyArray1, PyArray3, PyArrayMethods, PyReadonlyArray1, +}; +use pyo3::{ + exceptions::{PyIOError, PyValueError}, + prelude::*, + types::{PyAny, PyModule}, + Bound, PyErr, PyResult, +}; + +use healpix::nested::map::skymap::SkyMapValue; +use healpix::{ + depth_from_n_hash_unsafe, + nested::map::{ + fits::write::write_implicit_skymap_fits, + img::{to_skymap_img_default, PosConversion, Val}, + skymap::{ImplicitSkyMapArrayRef, SkyMap, SkyMapEnum}, + }, +}; + +#[pyfunction] +#[pyo3(pass_module)] +pub fn read_skymap<'py>( + module: &Bound<'py, PyModule>, + path: String, +) -> PyResult> { + SkyMapEnum::from_fits_file(&path) + .map_err(|err| PyIOError::new_err(err.to_string())) + .map(|sky_map_enum| match sky_map_enum { + SkyMapEnum::ImplicitU64U8(s) => s + .values() + .copied() + .collect::>() + .into_pyarray_bound(module.py()) + .into_any(), + SkyMapEnum::ImplicitU64I16(s) => s + .values() + .copied() + .collect::>() + .into_pyarray_bound(module.py()) + .into_any(), + SkyMapEnum::ImplicitU64I32(s) => s + .values() + .copied() + .collect::>() + .into_pyarray_bound(module.py()) + .into_any(), + SkyMapEnum::ImplicitU64I64(s) => s + .values() + .copied() + .collect::>() + .into_pyarray_bound(module.py()) + .into_any(), + SkyMapEnum::ImplicitU64F32(s) => s + .values() + .copied() + .collect::>() + .into_pyarray_bound(module.py()) + .into_any(), + SkyMapEnum::ImplicitU64F64(s) => s + .values() + .copied() + .collect::>() + .into_pyarray_bound(module.py()) + .into_any(), + }) +} + +// we define an enum for the supported numpy dtypes +#[derive(FromPyObject)] +pub enum SupportedArray<'py> { + F64(PyReadonlyArray1<'py, f64>), + I64(PyReadonlyArray1<'py, i64>), + F32(PyReadonlyArray1<'py, f32>), + I32(PyReadonlyArray1<'py, i32>), + I16(PyReadonlyArray1<'py, i16>), + U8(PyReadonlyArray1<'py, u8>), +} +impl<'py> SupportedArray<'py> { + fn n_hash(&self) -> u64 { + let n = match self { + SupportedArray::F64(values) => values.as_array().shape()[0], + SupportedArray::I64(values) => values.as_array().shape()[0], + SupportedArray::F32(values) => values.as_array().shape()[0], + SupportedArray::I32(values) => values.as_array().shape()[0], + SupportedArray::I16(values) => values.as_array().shape()[0], + SupportedArray::U8(values) => values.as_array().shape()[0], + }; + n as u64 + } +} + +#[pyfunction] +pub fn write_skymap(values: SupportedArray<'_>, path: String) -> Result<(), PyErr> { + let writer = + BufWriter::new(File::create(path).map_err(|err| PyIOError::new_err(err.to_string()))?); + match values { + SupportedArray::F64(values) => write_skymap_gen(writer, values.as_slice()), + SupportedArray::I64(values) => write_skymap_gen(writer, values.as_slice()), + SupportedArray::F32(values) => write_skymap_gen(writer, values.as_slice()), + SupportedArray::I32(values) => write_skymap_gen(writer, values.as_slice()), + SupportedArray::I16(values) => write_skymap_gen(writer, values.as_slice()), + SupportedArray::U8(values) => write_skymap_gen(writer, values.as_slice()), + } +} +fn write_skymap_gen( + writer: BufWriter, + as_slice_res: Result<&[T], NotContiguousError>, +) -> Result<(), PyErr> { + as_slice_res.map_err(move |e| e.into()).and_then(|slice| { + write_implicit_skymap_fits(writer, slice).map_err(|err| PyIOError::new_err(err.to_string())) + }) +} + +#[pyfunction] +#[pyo3(pass_module)] +pub fn pixels_skymap<'py>( + module: &Bound<'py, PyModule>, + values: SupportedArray<'py>, + image_size: u16, + convert_to_gal: bool, +) -> PyResult>> { + let n_hash = values.n_hash(); + let depth = depth_from_n_hash_unsafe(n_hash); + // we have to use https://github.com/cds-astro/cds-healpix-rust/blob/847ae35945708efb6b949c3d15b3726ab7adeb2f/src/nested/map/img.rs#L391 + match values { + SupportedArray::F64(values) => values.as_slice().map_err(|e| e.into()).and_then(|v| { + skymap_ref_to_img( + &ImplicitSkyMapArrayRef::<'_, u64, f64>::new(depth, v), + image_size, + convert_to_gal, + module.py(), + ) + }), + SupportedArray::I64(values) => values.as_slice().map_err(|e| e.into()).and_then(|v| { + skymap_ref_to_img( + &ImplicitSkyMapArrayRef::<'_, u64, i64>::new(depth, v), + image_size, + convert_to_gal, + module.py(), + ) + }), + SupportedArray::F32(values) => values.as_slice().map_err(|e| e.into()).and_then(|v| { + skymap_ref_to_img( + &ImplicitSkyMapArrayRef::<'_, u64, f32>::new(depth, v), + image_size, + convert_to_gal, + module.py(), + ) + }), + SupportedArray::I32(values) => values.as_slice().map_err(|e| e.into()).and_then(|v| { + skymap_ref_to_img( + &ImplicitSkyMapArrayRef::<'_, u64, i32>::new(depth, v), + image_size, + convert_to_gal, + module.py(), + ) + }), + SupportedArray::I16(values) => values.as_slice().map_err(|e| e.into()).and_then(|v| { + skymap_ref_to_img( + &ImplicitSkyMapArrayRef::<'_, u64, i16>::new(depth, v), + image_size, + convert_to_gal, + module.py(), + ) + }), + SupportedArray::U8(values) => values.as_slice().map_err(|e| e.into()).and_then(|v| { + skymap_ref_to_img( + &ImplicitSkyMapArrayRef::<'_, u64, u8>::new(depth, v), + image_size, + convert_to_gal, + module.py(), + ) + }), + } +} + +fn skymap_ref_to_img<'py, 'a, S>( + skymap: &'a S, + image_size: u16, + convert_to_gal: bool, + py: Python<'py>, +) -> PyResult>> +where + S: SkyMap<'a> + 'a, + S::ValueType: Val, +{ + if convert_to_gal { + let vec = to_skymap_img_default( + skymap, + (image_size << 1, image_size), + None, + None, + Some(PosConversion::EqMap2GalImg), + None, + None, + ) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + PyArray1::from_slice_bound(py, vec.as_slice()).reshape(Ix3( + image_size as usize, + (image_size << 1) as usize, + 4_usize, + )) + } else { + let vec = to_skymap_img_default( + skymap, + (image_size << 1, image_size), + None, + None, + None, + None, + None, + ) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + PyArray1::from_slice_bound(py, vec.as_slice()).reshape(Ix3( + image_size as usize, + (image_size << 1) as usize, + 4_usize, + )) + } +} + +#[pyfunction] +pub fn depth_skymap(values: SupportedArray) -> u8 { + depth_from_n_hash_unsafe(values.n_hash()) +}