Metadata-Version: 2.4
Name: healjax
Version: 0.2.3
Summary: Healpix routines for JAX.
Author: ghcollin, AaronParsons
License: MIT
Project-URL: Homepage, https://github.com/EIGSEP/healjax
Keywords: jax,healpix,healpy
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Information Technology
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE.md
Requires-Dist: numpy
Requires-Dist: jax
Provides-Extra: maps
Requires-Dist: healpy; extra == "maps"
Requires-Dist: astropy; extra == "maps"
Requires-Dist: scipy; extra == "maps"
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: healpy; extra == "test"
Requires-Dist: astropy; extra == "test"
Requires-Dist: scipy; extra == "test"
Requires-Dist: astropy-healpix; extra == "test"
Dynamic: license-file

# Healjax

Some healpix routines which run with JAX. 

These routines are ports of the BSD licensed healpix library https://github.com/astropy/astropy-healpix

# Installation

```
pip install healjax
```

The low-level pixel functions need only `numpy` and `jax`. The `healjax.maps`
object layer additionally uses `healpy` (scheme conversion, spherical harmonic
transforms), `astropy` (FITS I/O) and `scipy` (spherical harmonic fitting);
install those with the `maps` extra:

```
pip install 'healjax[maps]
```

These are imported lazily, so the error only appears if you call a feature
that needs them.

# Package layout

 - `healjax` — the low-level JAX pixel functions documented below
   (`ang2pix`, `vec2pix`, `get_interp_weights`, …).
 - `healjax.maps` — the object layer: `HealpixBase`, `Alm`, `HealpixMap` and
   `HPM`, plus spherical harmonic fitting (`sph_fit`, `fit_alms_from_maps`,
   `alms_to_filled_maps`).
 - `healjax.coord` — coordinate transforms: `rot_m`, `xyz2thphi`,
   `thphi2xyz`, `eq2top_m`, and the usual eq/topocentric helpers. Each
   dispatches on its input, returning NumPy for NumPy and JAX for JAX.
 - `healjax.interp` — `interpolate_map` and `rotate_interpolate_and_sum`,
   jitted with `nside` as a static argument.

The most common symbols (`HPM`, `rot_m`, `interpolate_map`, …) are re-exported
at the package root.

## `HPM`

`HPM` is the main user-facing map container: a `HealpixMap` whose coordinate
lookup and interpolation run through the JAX kernels. It replaces both
`aipy.healpix.HealpixMap` and the per-package `HPM` copies that used to live in
`eigsep_sim` and `eigsep_data`, and carries no aipy dependency.

```python
import numpy as np
import healjax

hpm = healjax.HPM(nside=64, interp=True)
hpm.set_map(sky_map)              # (npix,) or (npix, nfreq)
vals = hpm[theta, phi]            # interpolated read
vals = hpm[x, y, z]               # same, Cartesian
data = hpm.rotate_interpolate_and_sum(sky, crds, rot_ms)
```

Only the RING scheme is supported on the JAX interpolation path; interpolating
a NEST-ordered `HPM` raises `NotImplementedError`.

# Some notes

All functions take a `scheme` argument that can be one of
 - `'ring'` The ring indexing healpix convention.
 - `'nest'` The nested indexing healpix convetion.
 - `'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.

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.

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.

 - `vec2pix(scheme, nside, x, y, z, out_dtype=None) -> hp` Takes an xyz vector and returns the healpix bin that it lies in.
 - `ang2pix_radec(scheme, nside, ra, dec, out_dtype=None) -> hp` Takes a right ascension and declination and returns a healpix bin.
 - `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.

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.

 - `pix2vec(scheme, nside, hp, dx=None, dy=None) -> (x, y, z)` Takes a healpix index and returns the xyz vector that it corresponds to.
 - `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.
 - `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.
 - `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.

There are in addition two functions for finding neighbouring healpix bins.
 - `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.
 - `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.

The following functions can be used to convert between coordinate systems:
 - `ang2vec_radec(ra, dec) -> (x, y, z)` converts from right ascension and declination to the x, y, z vector system.
 - `ang2vec(theta, phi) -> (x, y, z)` converts from the healpy colatitude and longitude system to the x, y, z vector system.
 - `vec2ang_radec(x, y, z) -> (ra, dec)` converts from the x, y, z vector system to right ascension and declination.
 - `vec2ang(x, y, z) -> (theta, phi)` converts from the x, y, z vector system to the healpy colatitude longitude system, returning theta, phi.

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.
 - `scheme2bighpxy(scheme, nside, hp_idx) -> (bighp, x, y)` converts from the given scheme (including 'xy') to the internal big healpix, x, y system.
 - `bighpxy2scheme(scheme, nside, bighp, x, y) -> hp` converts from the given big healpix face number, x and y to the provided scheme (including 'xy').

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.
 - `nside2npix(nside) -> npix`
 - `npix2nside(npix) -> nside`
 - `get_nside(map) -> nside` returns the nside for a provided healpix array.

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.

# Compatibility

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. 

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.

The neighbours and convert map functions have been tested to agree with astropy-healpix.

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.
Tests pass on ARM and x86 CPUs.
