healjax 0.2.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- healjax-0.2.3/LICENSE.md +37 -0
- healjax-0.2.3/PKG-INFO +148 -0
- healjax-0.2.3/README.md +114 -0
- healjax-0.2.3/pyproject.toml +42 -0
- healjax-0.2.3/setup.cfg +4 -0
- healjax-0.2.3/src/healjax/__init__.py +33 -0
- healjax-0.2.3/src/healjax/coord/__init__.py +29 -0
- healjax-0.2.3/src/healjax/coord/transforms.py +317 -0
- healjax-0.2.3/src/healjax/healjax.py +1018 -0
- healjax-0.2.3/src/healjax/interp/__init__.py +5 -0
- healjax-0.2.3/src/healjax/interp/interpolation.py +85 -0
- healjax-0.2.3/src/healjax/maps/__init__.py +35 -0
- healjax-0.2.3/src/healjax/maps/_optional.py +55 -0
- healjax-0.2.3/src/healjax/maps/alm.py +375 -0
- healjax-0.2.3/src/healjax/maps/base.py +136 -0
- healjax-0.2.3/src/healjax/maps/healpix_map.py +173 -0
- healjax-0.2.3/src/healjax/maps/hpm.py +129 -0
- healjax-0.2.3/src/healjax.egg-info/PKG-INFO +148 -0
- healjax-0.2.3/src/healjax.egg-info/SOURCES.txt +24 -0
- healjax-0.2.3/src/healjax.egg-info/dependency_links.txt +1 -0
- healjax-0.2.3/src/healjax.egg-info/requires.txt +14 -0
- healjax-0.2.3/src/healjax.egg-info/top_level.txt +1 -0
- healjax-0.2.3/tests/test_coord.py +345 -0
- healjax-0.2.3/tests/test_interp.py +198 -0
- healjax-0.2.3/tests/test_maps.py +725 -0
- healjax-0.2.3/tests/tests.py +672 -0
healjax-0.2.3/LICENSE.md
ADDED
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2023 ghcollin
|
|
4
|
+
|
|
5
|
+
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:
|
|
6
|
+
|
|
7
|
+
The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
|
|
8
|
+
|
|
9
|
+
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.
|
|
10
|
+
|
|
11
|
+
This software was derived from the astropy-healpix library, that is covered under the following license:
|
|
12
|
+
|
|
13
|
+
Copyright (c) 2016-2018, Astropy Developers
|
|
14
|
+
All rights reserved.
|
|
15
|
+
|
|
16
|
+
Redistribution and use in source and binary forms, with or without modification,
|
|
17
|
+
are permitted provided that the following conditions are met:
|
|
18
|
+
|
|
19
|
+
* Redistributions of source code must retain the above copyright notice, this
|
|
20
|
+
list of conditions and the following disclaimer.
|
|
21
|
+
* Redistributions in binary form must reproduce the above copyright notice, this
|
|
22
|
+
list of conditions and the following disclaimer in the documentation and/or
|
|
23
|
+
other materials provided with the distribution.
|
|
24
|
+
* Neither the name of the Astropy Team nor the names of its contributors may be
|
|
25
|
+
used to endorse or promote products derived from this software without
|
|
26
|
+
specific prior written permission.
|
|
27
|
+
|
|
28
|
+
THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
|
|
29
|
+
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
|
|
30
|
+
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
|
|
31
|
+
DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
|
|
32
|
+
ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
|
|
33
|
+
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
|
|
34
|
+
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
|
|
35
|
+
ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
|
|
36
|
+
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
|
|
37
|
+
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
|
healjax-0.2.3/PKG-INFO
ADDED
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: healjax
|
|
3
|
+
Version: 0.2.3
|
|
4
|
+
Summary: Healpix routines for JAX.
|
|
5
|
+
Author: ghcollin, AaronParsons
|
|
6
|
+
License: MIT
|
|
7
|
+
Project-URL: Homepage, https://github.com/EIGSEP/healjax
|
|
8
|
+
Keywords: jax,healpix,healpy
|
|
9
|
+
Classifier: Development Status :: 4 - Beta
|
|
10
|
+
Classifier: Intended Audience :: Developers
|
|
11
|
+
Classifier: Intended Audience :: Information Technology
|
|
12
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.14
|
|
17
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
18
|
+
Requires-Python: >=3.11
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
License-File: LICENSE.md
|
|
21
|
+
Requires-Dist: numpy
|
|
22
|
+
Requires-Dist: jax
|
|
23
|
+
Provides-Extra: maps
|
|
24
|
+
Requires-Dist: healpy; extra == "maps"
|
|
25
|
+
Requires-Dist: astropy; extra == "maps"
|
|
26
|
+
Requires-Dist: scipy; extra == "maps"
|
|
27
|
+
Provides-Extra: test
|
|
28
|
+
Requires-Dist: pytest; extra == "test"
|
|
29
|
+
Requires-Dist: healpy; extra == "test"
|
|
30
|
+
Requires-Dist: astropy; extra == "test"
|
|
31
|
+
Requires-Dist: scipy; extra == "test"
|
|
32
|
+
Requires-Dist: astropy-healpix; extra == "test"
|
|
33
|
+
Dynamic: license-file
|
|
34
|
+
|
|
35
|
+
# Healjax
|
|
36
|
+
|
|
37
|
+
Some healpix routines which run with JAX.
|
|
38
|
+
|
|
39
|
+
These routines are ports of the BSD licensed healpix library https://github.com/astropy/astropy-healpix
|
|
40
|
+
|
|
41
|
+
# Installation
|
|
42
|
+
|
|
43
|
+
```
|
|
44
|
+
pip install healjax
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
The low-level pixel functions need only `numpy` and `jax`. The `healjax.maps`
|
|
48
|
+
object layer additionally uses `healpy` (scheme conversion, spherical harmonic
|
|
49
|
+
transforms), `astropy` (FITS I/O) and `scipy` (spherical harmonic fitting);
|
|
50
|
+
install those with the `maps` extra:
|
|
51
|
+
|
|
52
|
+
```
|
|
53
|
+
pip install 'healjax[maps]
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
These are imported lazily, so the error only appears if you call a feature
|
|
57
|
+
that needs them.
|
|
58
|
+
|
|
59
|
+
# Package layout
|
|
60
|
+
|
|
61
|
+
- `healjax` — the low-level JAX pixel functions documented below
|
|
62
|
+
(`ang2pix`, `vec2pix`, `get_interp_weights`, …).
|
|
63
|
+
- `healjax.maps` — the object layer: `HealpixBase`, `Alm`, `HealpixMap` and
|
|
64
|
+
`HPM`, plus spherical harmonic fitting (`sph_fit`, `fit_alms_from_maps`,
|
|
65
|
+
`alms_to_filled_maps`).
|
|
66
|
+
- `healjax.coord` — coordinate transforms: `rot_m`, `xyz2thphi`,
|
|
67
|
+
`thphi2xyz`, `eq2top_m`, and the usual eq/topocentric helpers. Each
|
|
68
|
+
dispatches on its input, returning NumPy for NumPy and JAX for JAX.
|
|
69
|
+
- `healjax.interp` — `interpolate_map` and `rotate_interpolate_and_sum`,
|
|
70
|
+
jitted with `nside` as a static argument.
|
|
71
|
+
|
|
72
|
+
The most common symbols (`HPM`, `rot_m`, `interpolate_map`, …) are re-exported
|
|
73
|
+
at the package root.
|
|
74
|
+
|
|
75
|
+
## `HPM`
|
|
76
|
+
|
|
77
|
+
`HPM` is the main user-facing map container: a `HealpixMap` whose coordinate
|
|
78
|
+
lookup and interpolation run through the JAX kernels. It replaces both
|
|
79
|
+
`aipy.healpix.HealpixMap` and the per-package `HPM` copies that used to live in
|
|
80
|
+
`eigsep_sim` and `eigsep_data`, and carries no aipy dependency.
|
|
81
|
+
|
|
82
|
+
```python
|
|
83
|
+
import numpy as np
|
|
84
|
+
import healjax
|
|
85
|
+
|
|
86
|
+
hpm = healjax.HPM(nside=64, interp=True)
|
|
87
|
+
hpm.set_map(sky_map) # (npix,) or (npix, nfreq)
|
|
88
|
+
vals = hpm[theta, phi] # interpolated read
|
|
89
|
+
vals = hpm[x, y, z] # same, Cartesian
|
|
90
|
+
data = hpm.rotate_interpolate_and_sum(sky, crds, rot_ms)
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Only the RING scheme is supported on the JAX interpolation path; interpolating
|
|
94
|
+
a NEST-ordered `HPM` raises `NotImplementedError`.
|
|
95
|
+
|
|
96
|
+
# Some notes
|
|
97
|
+
|
|
98
|
+
All functions take a `scheme` argument that can be one of
|
|
99
|
+
- `'ring'` The ring indexing healpix convention.
|
|
100
|
+
- `'nest'` The nested indexing healpix convetion.
|
|
101
|
+
- `'xy'` The xy convention used in the astropy healpix library. All calculations are done in this scheme, the other schemes are provided by transforming to and from xy.
|
|
102
|
+
|
|
103
|
+
All functions also take an `nside` argument that should be an integer. In typical JAX style, these functions are only defined to operate on single bin indices/vectors/angle-pairs, use `jax.vmap` to operate on arrays.
|
|
104
|
+
|
|
105
|
+
The following functions are used to convert to healpix bin indices. They all take an `out_dtype` option that defaults to `int`. This option can be used if you wish to perform calculations in a smaller dtype than the default JAX integer type.
|
|
106
|
+
|
|
107
|
+
- `vec2pix(scheme, nside, x, y, z, out_dtype=None) -> hp` Takes an xyz vector and returns the healpix bin that it lies in.
|
|
108
|
+
- `ang2pix_radec(scheme, nside, ra, dec, out_dtype=None) -> hp` Takes a right ascension and declination and returns a healpix bin.
|
|
109
|
+
- `ang2pix(scheme, nside, theta, phi, out_dtype=None) -> hp` Takes a theta and phi pair of angles in the healpy convention and returns a healpix bin.
|
|
110
|
+
|
|
111
|
+
The following functions are used to convert from healpix bin indices. They all take a `dx` and `dy` option that is used to offset the output within the requested bin. The default for both is 0.5, which returns a vector/angle-pair that lies in the center of the bin.
|
|
112
|
+
|
|
113
|
+
- `pix2vec(scheme, nside, hp, dx=None, dy=None) -> (x, y, z)` Takes a healpix index and returns the xyz vector that it corresponds to.
|
|
114
|
+
- `pix2ang_radec(scheme, nside, hp, dx=None, dy=None) -> (ra, dec)` Takes a healpix index and returns the right ascension and declination that it corresponds to.
|
|
115
|
+
- `pix2ang_colatlong(scheme, nside, hp, dx=None, dy=None) -> (theta, phi)` Takes a healpix index and returns theta and phi in the healpy colatitude longitude convention.
|
|
116
|
+
- `pix2ang(scheme, nside, hp, dx=None, dy=None) -> (theta, phi)` Same as `pix2ang_colonglat` but the conversion is done through xyz vectors as in the astropy healpix library.
|
|
117
|
+
|
|
118
|
+
There are in addition two functions for finding neighbouring healpix bins.
|
|
119
|
+
- `get_neighbours(scheme, nside, hp) -> hp[8]` This is intended to have the same API as the healpy neighbour function (as tested against the astropy healpix library). It returns 8 neighbours, with -1 denoting a non-existent neighbour.
|
|
120
|
+
- `get_patch(scheme, nside, hp) -> hp[3, 3]` This is used to implement `get_neighbours` and returns a 3x3 array of the neighbours and the input index (located in the central element). A -1 denotes that the neighbour in that position does not exist.
|
|
121
|
+
|
|
122
|
+
The following functions can be used to convert between coordinate systems:
|
|
123
|
+
- `ang2vec_radec(ra, dec) -> (x, y, z)` converts from right ascension and declination to the x, y, z vector system.
|
|
124
|
+
- `ang2vec(theta, phi) -> (x, y, z)` converts from the healpy colatitude and longitude system to the x, y, z vector system.
|
|
125
|
+
- `vec2ang_radec(x, y, z) -> (ra, dec)` converts from the x, y, z vector system to right ascension and declination.
|
|
126
|
+
- `vec2ang(x, y, z) -> (theta, phi)` converts from the x, y, z vector system to the healpy colatitude longitude system, returning theta, phi.
|
|
127
|
+
|
|
128
|
+
This library follows astropy healpix in using an internal pixel addressing system that has three components: a 'big' healpix index that determines which of the twelve healpix faces the pixel in located in, as well as an x and y coordinate that addresses the pixel within this face using a cartesian coordinate chart. This should not be confused with the 'xy' healpix indexing scheme that packs these three numbers into a single healpix index. Nevertheless, knowing the x, y coordinates of a pixel can be useful, and so these functions are defined to retrieve them.
|
|
129
|
+
- `scheme2bighpxy(scheme, nside, hp_idx) -> (bighp, x, y)` converts from the given scheme (including 'xy') to the internal big healpix, x, y system.
|
|
130
|
+
- `bighpxy2scheme(scheme, nside, bighp, x, y) -> hp` converts from the given big healpix face number, x and y to the provided scheme (including 'xy').
|
|
131
|
+
|
|
132
|
+
There is also the usual nside and npix utility functions. Note that these are all numpy functions, they cannot be used for runtime JAX and are intended for non-traced computations such as array dimensions.
|
|
133
|
+
- `nside2npix(nside) -> npix`
|
|
134
|
+
- `npix2nside(npix) -> nside`
|
|
135
|
+
- `get_nside(map) -> nside` returns the nside for a provided healpix array.
|
|
136
|
+
|
|
137
|
+
Finally, the `convert_map(in_scheme, out_scheme, map) -> converted_map` function can be used to convert a map to and from the various schemes detailed above.
|
|
138
|
+
|
|
139
|
+
# Compatibility
|
|
140
|
+
|
|
141
|
+
Angle-pair to healpix functions tested to agree with astropy-healpix up to 1e-15 away from boundary in 64bit mode. Tests in 32bit mode predictably fail, as astropy-healpix computes in 64bit mode, meaning that we can't expect the binning results to be reproduced by 32bit calculations.
|
|
142
|
+
|
|
143
|
+
Healpix to angle-pair functions current exhibit a relative error that grows with the nside. For nsides up to 1024, this stays within a factor of 1e3 times to 1e7 times the floating point epsilon (this is still a relative error of less than 1e-9). This may be due to an implementation error that I haven't been able to track down.
|
|
144
|
+
|
|
145
|
+
The neighbours and convert map functions have been tested to agree with astropy-healpix.
|
|
146
|
+
|
|
147
|
+
Currently, there is an issue with running this on NVIDIA GPUs. Vmapping over certain, medium-sized arrays of inputs can cause JAX to crash. If you run into this, you can try increasing or decreasing the size of the array you vmap over to work around it. This bug has been reported to the JAX team.
|
|
148
|
+
Tests pass on ARM and x86 CPUs.
|
healjax-0.2.3/README.md
ADDED
|
@@ -0,0 +1,114 @@
|
|
|
1
|
+
# Healjax
|
|
2
|
+
|
|
3
|
+
Some healpix routines which run with JAX.
|
|
4
|
+
|
|
5
|
+
These routines are ports of the BSD licensed healpix library https://github.com/astropy/astropy-healpix
|
|
6
|
+
|
|
7
|
+
# Installation
|
|
8
|
+
|
|
9
|
+
```
|
|
10
|
+
pip install healjax
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
The low-level pixel functions need only `numpy` and `jax`. The `healjax.maps`
|
|
14
|
+
object layer additionally uses `healpy` (scheme conversion, spherical harmonic
|
|
15
|
+
transforms), `astropy` (FITS I/O) and `scipy` (spherical harmonic fitting);
|
|
16
|
+
install those with the `maps` extra:
|
|
17
|
+
|
|
18
|
+
```
|
|
19
|
+
pip install 'healjax[maps]
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
These are imported lazily, so the error only appears if you call a feature
|
|
23
|
+
that needs them.
|
|
24
|
+
|
|
25
|
+
# Package layout
|
|
26
|
+
|
|
27
|
+
- `healjax` — the low-level JAX pixel functions documented below
|
|
28
|
+
(`ang2pix`, `vec2pix`, `get_interp_weights`, …).
|
|
29
|
+
- `healjax.maps` — the object layer: `HealpixBase`, `Alm`, `HealpixMap` and
|
|
30
|
+
`HPM`, plus spherical harmonic fitting (`sph_fit`, `fit_alms_from_maps`,
|
|
31
|
+
`alms_to_filled_maps`).
|
|
32
|
+
- `healjax.coord` — coordinate transforms: `rot_m`, `xyz2thphi`,
|
|
33
|
+
`thphi2xyz`, `eq2top_m`, and the usual eq/topocentric helpers. Each
|
|
34
|
+
dispatches on its input, returning NumPy for NumPy and JAX for JAX.
|
|
35
|
+
- `healjax.interp` — `interpolate_map` and `rotate_interpolate_and_sum`,
|
|
36
|
+
jitted with `nside` as a static argument.
|
|
37
|
+
|
|
38
|
+
The most common symbols (`HPM`, `rot_m`, `interpolate_map`, …) are re-exported
|
|
39
|
+
at the package root.
|
|
40
|
+
|
|
41
|
+
## `HPM`
|
|
42
|
+
|
|
43
|
+
`HPM` is the main user-facing map container: a `HealpixMap` whose coordinate
|
|
44
|
+
lookup and interpolation run through the JAX kernels. It replaces both
|
|
45
|
+
`aipy.healpix.HealpixMap` and the per-package `HPM` copies that used to live in
|
|
46
|
+
`eigsep_sim` and `eigsep_data`, and carries no aipy dependency.
|
|
47
|
+
|
|
48
|
+
```python
|
|
49
|
+
import numpy as np
|
|
50
|
+
import healjax
|
|
51
|
+
|
|
52
|
+
hpm = healjax.HPM(nside=64, interp=True)
|
|
53
|
+
hpm.set_map(sky_map) # (npix,) or (npix, nfreq)
|
|
54
|
+
vals = hpm[theta, phi] # interpolated read
|
|
55
|
+
vals = hpm[x, y, z] # same, Cartesian
|
|
56
|
+
data = hpm.rotate_interpolate_and_sum(sky, crds, rot_ms)
|
|
57
|
+
```
|
|
58
|
+
|
|
59
|
+
Only the RING scheme is supported on the JAX interpolation path; interpolating
|
|
60
|
+
a NEST-ordered `HPM` raises `NotImplementedError`.
|
|
61
|
+
|
|
62
|
+
# Some notes
|
|
63
|
+
|
|
64
|
+
All functions take a `scheme` argument that can be one of
|
|
65
|
+
- `'ring'` The ring indexing healpix convention.
|
|
66
|
+
- `'nest'` The nested indexing healpix convetion.
|
|
67
|
+
- `'xy'` The xy convention used in the astropy healpix library. All calculations are done in this scheme, the other schemes are provided by transforming to and from xy.
|
|
68
|
+
|
|
69
|
+
All functions also take an `nside` argument that should be an integer. In typical JAX style, these functions are only defined to operate on single bin indices/vectors/angle-pairs, use `jax.vmap` to operate on arrays.
|
|
70
|
+
|
|
71
|
+
The following functions are used to convert to healpix bin indices. They all take an `out_dtype` option that defaults to `int`. This option can be used if you wish to perform calculations in a smaller dtype than the default JAX integer type.
|
|
72
|
+
|
|
73
|
+
- `vec2pix(scheme, nside, x, y, z, out_dtype=None) -> hp` Takes an xyz vector and returns the healpix bin that it lies in.
|
|
74
|
+
- `ang2pix_radec(scheme, nside, ra, dec, out_dtype=None) -> hp` Takes a right ascension and declination and returns a healpix bin.
|
|
75
|
+
- `ang2pix(scheme, nside, theta, phi, out_dtype=None) -> hp` Takes a theta and phi pair of angles in the healpy convention and returns a healpix bin.
|
|
76
|
+
|
|
77
|
+
The following functions are used to convert from healpix bin indices. They all take a `dx` and `dy` option that is used to offset the output within the requested bin. The default for both is 0.5, which returns a vector/angle-pair that lies in the center of the bin.
|
|
78
|
+
|
|
79
|
+
- `pix2vec(scheme, nside, hp, dx=None, dy=None) -> (x, y, z)` Takes a healpix index and returns the xyz vector that it corresponds to.
|
|
80
|
+
- `pix2ang_radec(scheme, nside, hp, dx=None, dy=None) -> (ra, dec)` Takes a healpix index and returns the right ascension and declination that it corresponds to.
|
|
81
|
+
- `pix2ang_colatlong(scheme, nside, hp, dx=None, dy=None) -> (theta, phi)` Takes a healpix index and returns theta and phi in the healpy colatitude longitude convention.
|
|
82
|
+
- `pix2ang(scheme, nside, hp, dx=None, dy=None) -> (theta, phi)` Same as `pix2ang_colonglat` but the conversion is done through xyz vectors as in the astropy healpix library.
|
|
83
|
+
|
|
84
|
+
There are in addition two functions for finding neighbouring healpix bins.
|
|
85
|
+
- `get_neighbours(scheme, nside, hp) -> hp[8]` This is intended to have the same API as the healpy neighbour function (as tested against the astropy healpix library). It returns 8 neighbours, with -1 denoting a non-existent neighbour.
|
|
86
|
+
- `get_patch(scheme, nside, hp) -> hp[3, 3]` This is used to implement `get_neighbours` and returns a 3x3 array of the neighbours and the input index (located in the central element). A -1 denotes that the neighbour in that position does not exist.
|
|
87
|
+
|
|
88
|
+
The following functions can be used to convert between coordinate systems:
|
|
89
|
+
- `ang2vec_radec(ra, dec) -> (x, y, z)` converts from right ascension and declination to the x, y, z vector system.
|
|
90
|
+
- `ang2vec(theta, phi) -> (x, y, z)` converts from the healpy colatitude and longitude system to the x, y, z vector system.
|
|
91
|
+
- `vec2ang_radec(x, y, z) -> (ra, dec)` converts from the x, y, z vector system to right ascension and declination.
|
|
92
|
+
- `vec2ang(x, y, z) -> (theta, phi)` converts from the x, y, z vector system to the healpy colatitude longitude system, returning theta, phi.
|
|
93
|
+
|
|
94
|
+
This library follows astropy healpix in using an internal pixel addressing system that has three components: a 'big' healpix index that determines which of the twelve healpix faces the pixel in located in, as well as an x and y coordinate that addresses the pixel within this face using a cartesian coordinate chart. This should not be confused with the 'xy' healpix indexing scheme that packs these three numbers into a single healpix index. Nevertheless, knowing the x, y coordinates of a pixel can be useful, and so these functions are defined to retrieve them.
|
|
95
|
+
- `scheme2bighpxy(scheme, nside, hp_idx) -> (bighp, x, y)` converts from the given scheme (including 'xy') to the internal big healpix, x, y system.
|
|
96
|
+
- `bighpxy2scheme(scheme, nside, bighp, x, y) -> hp` converts from the given big healpix face number, x and y to the provided scheme (including 'xy').
|
|
97
|
+
|
|
98
|
+
There is also the usual nside and npix utility functions. Note that these are all numpy functions, they cannot be used for runtime JAX and are intended for non-traced computations such as array dimensions.
|
|
99
|
+
- `nside2npix(nside) -> npix`
|
|
100
|
+
- `npix2nside(npix) -> nside`
|
|
101
|
+
- `get_nside(map) -> nside` returns the nside for a provided healpix array.
|
|
102
|
+
|
|
103
|
+
Finally, the `convert_map(in_scheme, out_scheme, map) -> converted_map` function can be used to convert a map to and from the various schemes detailed above.
|
|
104
|
+
|
|
105
|
+
# Compatibility
|
|
106
|
+
|
|
107
|
+
Angle-pair to healpix functions tested to agree with astropy-healpix up to 1e-15 away from boundary in 64bit mode. Tests in 32bit mode predictably fail, as astropy-healpix computes in 64bit mode, meaning that we can't expect the binning results to be reproduced by 32bit calculations.
|
|
108
|
+
|
|
109
|
+
Healpix to angle-pair functions current exhibit a relative error that grows with the nside. For nsides up to 1024, this stays within a factor of 1e3 times to 1e7 times the floating point epsilon (this is still a relative error of less than 1e-9). This may be due to an implementation error that I haven't been able to track down.
|
|
110
|
+
|
|
111
|
+
The neighbours and convert map functions have been tested to agree with astropy-healpix.
|
|
112
|
+
|
|
113
|
+
Currently, there is an issue with running this on NVIDIA GPUs. Vmapping over certain, medium-sized arrays of inputs can cause JAX to crash. If you run into this, you can try increasing or decreasing the size of the array you vmap over to work around it. This bug has been reported to the JAX team.
|
|
114
|
+
Tests pass on ARM and x86 CPUs.
|
|
@@ -0,0 +1,42 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["setuptools>=61"]
|
|
3
|
+
build-backend = "setuptools.build_meta"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "healjax"
|
|
7
|
+
version = "0.2.3"
|
|
8
|
+
description = "Healpix routines for JAX."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
license = { text = "MIT" }
|
|
11
|
+
authors = [{ name = "ghcollin" }, { name = "AaronParsons" }]
|
|
12
|
+
keywords = ["jax", "healpix", "healpy"]
|
|
13
|
+
classifiers = [
|
|
14
|
+
"Development Status :: 4 - Beta",
|
|
15
|
+
"Intended Audience :: Developers",
|
|
16
|
+
"Intended Audience :: Information Technology",
|
|
17
|
+
"License :: OSI Approved :: MIT License",
|
|
18
|
+
"Programming Language :: Python :: 3.11",
|
|
19
|
+
"Programming Language :: Python :: 3.12",
|
|
20
|
+
"Programming Language :: Python :: 3.13",
|
|
21
|
+
"Programming Language :: Python :: 3.14",
|
|
22
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
23
|
+
]
|
|
24
|
+
requires-python = ">=3.11"
|
|
25
|
+
dependencies = ["numpy", "jax"]
|
|
26
|
+
|
|
27
|
+
[project.optional-dependencies]
|
|
28
|
+
# healjax.maps: pixel-scheme conversion and spherical harmonics (healpy),
|
|
29
|
+
# FITS I/O (astropy), and spherical-harmonic fitting (scipy).
|
|
30
|
+
maps = ["healpy", "astropy", "scipy"]
|
|
31
|
+
test = ["pytest", "healpy", "astropy", "scipy", "astropy-healpix"]
|
|
32
|
+
|
|
33
|
+
[project.urls]
|
|
34
|
+
Homepage = "https://github.com/EIGSEP/healjax"
|
|
35
|
+
|
|
36
|
+
[tool.setuptools.packages.find]
|
|
37
|
+
where = ["src"]
|
|
38
|
+
|
|
39
|
+
[tool.pytest.ini_options]
|
|
40
|
+
testpaths = ["tests"]
|
|
41
|
+
# tests/tests.py is the original upstream unittest suite.
|
|
42
|
+
python_files = ["test_*.py", "tests.py"]
|
healjax-0.2.3/setup.cfg
ADDED
|
@@ -0,0 +1,33 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Copyright (c) 2023 ghcollin
|
|
3
|
+
|
|
4
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
5
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
6
|
+
in the Software without restriction, including without limitation the rights
|
|
7
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
8
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
9
|
+
furnished to do so, subject to the following conditions:
|
|
10
|
+
|
|
11
|
+
The above copyright notice and this permission notice shall be included in all
|
|
12
|
+
copies or substantial portions of the Software.
|
|
13
|
+
|
|
14
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
15
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
16
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
17
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
18
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
19
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
20
|
+
SOFTWARE.
|
|
21
|
+
"""
|
|
22
|
+
|
|
23
|
+
from .healjax import vec2pix, ang2pix_radec, ang2pix, pix2vec, pix2ang_radec, pix2ang_colatlong, pix2ang, get_patch, get_neighbours, convert_map, bighpxy2scheme, scheme2bighpxy, ang2vec_radec, ang2vec, vec2ang_radec, vec2ang, nside2npix, npix2nside, get_nside, ring_above, get_ring_info2
|
|
24
|
+
from .healjax import get_interpol as get_interp_weights
|
|
25
|
+
from .healjax import precompute_ring_info, get_interpol_precomp as get_interp_weights_precomp
|
|
26
|
+
from .healjax import FLOAT_TYPE, INT_TYPE
|
|
27
|
+
|
|
28
|
+
# Subpackages layered on top of the pixel functions above. These are imported
|
|
29
|
+
# after `.healjax` so that their `from ..healjax import ...` lines resolve.
|
|
30
|
+
from . import coord, interp, maps
|
|
31
|
+
from .coord import eq2top_m, rot_m, thphi2xyz, xyz2thphi
|
|
32
|
+
from .interp import interpolate_map, rotate_interpolate_and_sum
|
|
33
|
+
from .maps import HPM, Alm, HealpixBase, HealpixMap
|
|
@@ -0,0 +1,29 @@
|
|
|
1
|
+
"""Coordinate transforms shared across the EIGSEP packages."""
|
|
2
|
+
|
|
3
|
+
from .transforms import (
|
|
4
|
+
rot_m,
|
|
5
|
+
xyz2thphi,
|
|
6
|
+
thphi2xyz,
|
|
7
|
+
eq2top_m,
|
|
8
|
+
top2eq_m,
|
|
9
|
+
eq2radec,
|
|
10
|
+
radec2eq,
|
|
11
|
+
latlong2xyz,
|
|
12
|
+
top2azalt,
|
|
13
|
+
azalt2top,
|
|
14
|
+
angles_to_coord,
|
|
15
|
+
)
|
|
16
|
+
|
|
17
|
+
__all__ = [
|
|
18
|
+
"rot_m",
|
|
19
|
+
"xyz2thphi",
|
|
20
|
+
"thphi2xyz",
|
|
21
|
+
"eq2top_m",
|
|
22
|
+
"top2eq_m",
|
|
23
|
+
"eq2radec",
|
|
24
|
+
"radec2eq",
|
|
25
|
+
"latlong2xyz",
|
|
26
|
+
"top2azalt",
|
|
27
|
+
"azalt2top",
|
|
28
|
+
"angles_to_coord",
|
|
29
|
+
]
|