jax-healpy 0.2.1__tar.gz

jax_healpy-0.2.1/.github/workflows/release.yml

@@ -0,0 +1,34 @@
+name: Release to PyPI
+
+on:
+  release:
+    types: [created]
+
+jobs:
+  publish:
+    name: Publish package to PyPI
+    runs-on: ubuntu-latest
+
+    permissions:
+      id-token: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.x'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build twine
+
+      - name: Build package
+        run: |
+          python -m build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

jax_healpy-0.2.1/.gitignore

@@ -0,0 +1,10 @@
+.coverage
+MANIFEST
+build
+dist
+__pycache__
+*.egg-info
+/.idea/
+/.vscode/
+/venv*/
+/benchmarks/results/

jax_healpy-0.2.1/.pre-commit-config.yaml

@@ -0,0 +1,27 @@
+repos:
+- repo: https://github.com/hadialqattan/pycln
+  rev: "v2.4.0"
+  hooks:
+  - id: pycln
+    args:
+    - --all
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.8.0
+  hooks:
+    - id: ruff
+    - id: ruff-format
+
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: 'v5.0.0'
+  hooks:
+  - id: trailing-whitespace
+  - id: end-of-file-fixer
+  - id: check-yaml
+  - id: check-merge-conflict
+
+- repo: https://github.com/PyCQA/bandit
+  rev: '1.8.0'
+  hooks:
+  - id: bandit
+    files: ^jax_healpy/

jax_healpy-0.2.1/PKG-INFO

@@ -0,0 +1,66 @@
+Metadata-Version: 2.2
+Name: jax-healpy
+Version: 0.2.1
+Summary: Healpix JAX implementation.
+Author-email: Pierre Chanial <pierre.chanial@gmail.com>, Simon Biquard <sbiquard@gmail.com>
+Maintainer-email: Pierre Chanial <pierre.chanial@gmail.com>
+Project-URL: homepage, https://pchanial.github.io/jax-healpy
+Project-URL: repository, https://github.com/pchanial/jax-healpy
+Keywords: scientific computing
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+Requires-Dist: jax
+Requires-Dist: jaxtyping
+Provides-Extra: dev
+Requires-Dist: healpy; extra == "dev"
+Requires-Dist: matplotlib; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pandas; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-mock; extra == "dev"
+Requires-Dist: setuptools_scm; extra == "dev"
+Requires-Dist: typer; extra == "dev"
+Requires-Dist: PyYAML; extra == "dev"
+Provides-Extra: recommended
+Requires-Dist: s2fft; extra == "recommended"
+
+# Healpy with JAX
+
+This project intends to assess the interest of implementing healpy functions using JAX.
+
+_WARNING: BETA STAGE!!!_
+
+## Installation
+
+First install JAX following the [documentation](https://jax.readthedocs.io/en/latest/installation.html).
+
+Then install the package with:
+
+```bash
+pip install jax-healpy
+```
+
+To use the spherical harmonics functions,
+you will need [s2fft](https://astro-informatics.github.io/s2fft/),
+part of the recommended dependencies:
+
+```bash
+pip install jax-healpy[recommended]
+```
+
+## Benchmarks
+
+Execution time measured on the [Jean Zay supercomputer](http://www.idris.fr/jean-zay/cpu/jean-zay-cpu-hw.html).
+
+- CPU: Intel(R) Xeon(R) Gold 2648 @ 2.50GHz
+- GPU: NVIDIA Tesla V100-SXM2-16GB
+
+![Benchmark](/docs/benchmarks/chart-darkbackground-n10000000.png)

jax_healpy-0.2.1/README.md

@@ -0,0 +1,32 @@
+# Healpy with JAX
+
+This project intends to assess the interest of implementing healpy functions using JAX.
+
+_WARNING: BETA STAGE!!!_
+
+## Installation
+
+First install JAX following the [documentation](https://jax.readthedocs.io/en/latest/installation.html).
+
+Then install the package with:
+
+```bash
+pip install jax-healpy
+```
+
+To use the spherical harmonics functions,
+you will need [s2fft](https://astro-informatics.github.io/s2fft/),
+part of the recommended dependencies:
+
+```bash
+pip install jax-healpy[recommended]
+```
+
+## Benchmarks
+
+Execution time measured on the [Jean Zay supercomputer](http://www.idris.fr/jean-zay/cpu/jean-zay-cpu-hw.html).
+
+- CPU: Intel(R) Xeon(R) Gold 2648 @ 2.50GHz
+- GPU: NVIDIA Tesla V100-SXM2-16GB
+
+![Benchmark](/docs/benchmarks/chart-darkbackground-n10000000.png)

jax_healpy-0.2.1/benchmarks/__main__.py

@@ -0,0 +1,286 @@
+import os
+import platform
+import re
+import subprocess
+import timeit
+from collections.abc import Callable
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from pprint import pprint
+from typing import Any, Iterable, Literal
+
+import healpy as hp
+import jax
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import typer
+import yaml
+from jaxtyping import ArrayLike
+from matplotlib.ticker import ScalarFormatter
+
+BENCH_PATH = Path(__file__).parent / 'results'
+BENCHMARKED_FUNCS = [
+    'ang2vec',
+    'vec2ang',
+    'ang2pix',
+    'pix2ang',
+    'vec2pix',
+    'pix2vec',
+    'ring2nest',
+    'nest2ring',
+    'pix2xyf',
+    'xyf2pix',
+    'reorder',
+]
+CHART_PATH_NAME = 'chart-{style}-n{n}.png'
+
+# TODO: use those when typer supports Literals
+LibraryType = Literal['healpy', 'jax-healpy']
+PrecisionType = Literal['32', '64']
+
+app = typer.Typer()
+
+
+@dataclass(frozen=True)
+class BenchmarkResult:
+    library: str
+    version: str
+    processor: str
+    processor_type: str
+    precision: int
+    n: int
+    execution_times_s: dict[str, float]
+
+
+def bench_it(
+    library: str,
+    func_name: str,
+    nside: int,
+    n: int,
+    precision: str,
+    rng: np.random.Generator,
+) -> float:
+    if precision == '32':
+        dtype = np.dtype('float32')
+    elif precision == '64':
+        dtype = np.dtype('float64')
+    else:
+        raise ValueError(f'Invalid precision {precision}')
+
+    args = _get_args(library, func_name, nside, n, dtype, rng)
+    func = _get_func(library, func_name, *args)
+    with jax.experimental.enable_x64(precision == '64'):
+        return time_it(func)
+
+
+def _get_args(
+    library: str,
+    func_name: str,
+    nside: int,
+    n: int,
+    dtype: np.dtype,
+    rng: np.random.Generator,
+) -> tuple[ArrayLike]:
+    if func_name in {'ang2vec', 'ang2pix'}:
+        theta = rng.uniform(0, np.pi, size=n).astype(dtype)
+        phi = rng.uniform(0, 2 * np.pi, size=n).astype(dtype)
+        if func_name == 'ang2vec':
+            args = (theta, phi)
+        else:
+            args = (nside, theta, phi)
+
+    elif func_name in {'vec2ang', 'vec2pix'}:
+        vec = rng.uniform(-np.pi, np.pi, size=(3, n))
+        vec /= np.sqrt(np.sum(vec**2, axis=0))
+        vec = vec.astype(dtype)
+        if func_name == 'vec2ang':
+            args = (vec.T.copy(),)
+        else:
+            args = (nside, vec[0], vec[1], vec[2])
+
+    elif func_name in {'pix2ang', 'pix2vec', 'pix2xyf', 'ring2nest', 'nest2ring'}:
+        pixels = rng.uniform(0, hp.nside2npix(nside), size=n).astype(int)
+        args = (nside, pixels)
+
+    elif func_name == 'xyf2pix':
+        x = rng.uniform(0, nside, size=n).astype(int)
+        y = rng.uniform(0, nside, size=n).astype(int)
+        f = rng.uniform(0, 12, size=n).astype(int)
+        args = (nside, x, y, f)
+
+    elif func_name == 'reorder':
+        npix = hp.nside2npix(nside)
+        map_in = rng.uniform(size=npix)
+        args = (map_in,)
+
+    else:
+        raise NotImplementedError
+
+    if library == 'jax-healpy':
+        args = tuple(jax.device_put(_) if isinstance(_, np.ndarray) else _ for _ in args)
+
+    return args
+
+
+def _get_func(library: str, func_name: str, *args: Any):
+    if library == 'jax-healpy':
+        import jax_healpy as module
+    else:
+        import healpy as module
+
+    func = getattr(module, func_name)
+    if library == 'healpy':
+        if func_name == 'reorder':
+            func_ = lambda: func(*args, r2n=True)  # noqa: E731
+        else:
+            func_ = lambda: func(*args)  # noqa: E731
+    else:
+        if func_name in {'pix2ang', 'vec2ang'}:
+
+            def func_() -> None:
+                theta, phi = func(*args)
+                theta.block_until_ready()
+                phi.block_until_ready()
+
+        elif func_name == 'pix2xyf':
+
+            def func_() -> None:
+                x, y, f = func(*args)
+                x.block_until_ready()
+                y.block_until_ready()
+                f.block_until_ready()
+
+        elif func_name == 'reorder':
+
+            def func_() -> None:
+                func(*args, r2n=True).block_until_ready()
+
+        else:
+
+            def func_() -> None:
+                func(*args).block_until_ready()
+
+        func_()  # discard first call, which includes compilation
+
+    return func_
+
+
+def time_it(func: Callable[[], None]) -> float:
+    timer = timeit.Timer(func)
+    number, _ = timer.autorange()
+    execution_time = min(timer.repeat(number=number)) / number
+    return execution_time
+
+
+@app.command()
+def run(
+    library: str,
+    nside: int = 512,
+    n: int = 10_000_000,
+    precision: str = '64',
+) -> None:
+    if library == 'jax-healpy':
+        version = f'jax({jax.__version__})'
+        device = jax.devices()[0]
+        processor = device.device_kind
+        if processor == 'cpu':
+            processor = _get_cpu_name()
+            processor_type = 'cpu'
+        else:
+            processor_type = 'gpu'
+    elif library == 'healpy':
+        version = hp.__version__
+        processor = _get_cpu_name()
+        processor_type = 'cpu'
+    else:
+        raise ValueError(f'Invalid library {library}')
+
+    rng = np.random.default_rng(0)
+    execution_times = {}
+
+    print(f'Running {library}...')
+    for func_name in BENCHMARKED_FUNCS:
+        execution_time = bench_it(library, func_name, nside, n, precision, rng)
+        execution_times[func_name] = execution_time
+
+    result = BenchmarkResult(
+        library=library,
+        version=version,
+        processor=processor,
+        processor_type=processor_type,
+        precision=int(precision),
+        n=n,
+        execution_times_s=execution_times,
+    )
+    pprint(result)
+
+    chars = '[] '
+    formatted_processor = processor.lower().rstrip(chars).replace('(c)', '').replace('(tm)', '')
+    filename = f'{library}-{formatted_processor}-precision{precision}-n{n}.yaml'
+    for char in chars:
+        filename = filename.replace(char, '-')
+    BENCH_PATH.mkdir(parents=True, exist_ok=True)
+    with open(BENCH_PATH / filename, 'w') as f:
+        yaml.dump(asdict(result), f)
+
+
+def _get_cpu_name():
+    if platform.system() == 'Darwin':
+        os.environ['PATH'] = os.environ['PATH'] + os.pathsep + '/usr/sbin'
+        command = 'sysctl -n machdep.cpu.brand_string'
+        return subprocess.check_output(command).strip()
+    elif platform.system() == 'Linux':
+        command = 'cat /proc/cpuinfo'
+        all_info = subprocess.check_output(command, shell=True).decode().strip()
+        for line in all_info.split('\n'):
+            if 'model name' in line:
+                return re.sub('.*model name.*:', '', line, 1).strip()
+    return platform.processor()
+
+
+@app.command()
+def collect():
+    files = BENCH_PATH.glob('*-n*.yaml')
+    ns = {}
+    for file in files:
+        result_as_dict = yaml.safe_load(file.read_text())
+        result = BenchmarkResult(**result_as_dict)
+        ns.setdefault(result.n, []).append(result)
+
+    for n, results in ns.items():
+        results = sorted(
+            results,
+            key=lambda _: sum(_.execution_times_s.values()) / len(BENCHMARKED_FUNCS),
+        )
+        export_chart(n, results, 'default')
+        export_chart(n, results, 'dark_background')
+
+
+def export_chart(n: int, results: Iterable[BenchmarkResult], style: str):
+    chart_path = BENCH_PATH / CHART_PATH_NAME.format(style=style.replace('_', ''), n=n)
+    plt.style.use(style)
+    # I've used pandas because I was not able to do what I wanted with matplotlib alone
+    df = pd.DataFrame(
+        [
+            [f'{res.library}[{res.processor_type}{res.precision}]']
+            + [res.execution_times_s[_] for _ in BENCHMARKED_FUNCS]
+            for res in results
+        ],
+        columns=['library'] + BENCHMARKED_FUNCS,
+    )
+    for func in BENCHMARKED_FUNCS:
+        df[func] /= df[func][0]
+    pprint(df)
+
+    ax = df.plot.bar(x='library', title='Benchmark (lower is better)', rot=60)
+    ax.legend(framealpha=0, loc='upper left')
+    plt.yscale('log')
+    ax.yaxis.set_major_formatter(ScalarFormatter())
+    plt.xlabel('')
+    plt.tight_layout()
+    plt.savefig(chart_path, transparent=True)
+
+
+if __name__ == '__main__':
+    app()

jax_healpy-0.2.1/benchmarks/script_sphtfunc.py

@@ -0,0 +1,53 @@
+import timeit
+from pathlib import Path
+from typing import Callable
+
+import healpy as hp
+import jax
+import numpy as np
+
+import jax_healpy as jhp
+
+data_path = Path(jhp.__file__).parent.parent / 'tests/data'
+
+
+def cla(data_path: Path) -> np.ndarray:
+    return hp.read_cl(data_path / 'cl_wmap_band_iqumap_r9_7yr_W_v4_udgraded32_II_lmax64_rmmono_3iter.fits')
+
+
+nside = 256
+lmax = 2 * nside
+fwhm_deg = 7.0
+
+
+orig = hp.synfast(
+    cla(data_path),
+    nside,
+    lmax=lmax,
+    pixwin=False,
+    fwhm=np.radians(fwhm_deg),
+    new=False,
+)
+orig_d = jax.device_put(orig)
+
+
+def map2alm_hp():
+    hp.map2alm(orig, iter=0, lmax=2 * nside - 1, use_weights=False)
+
+
+def map2alm_jhp():
+    jhp.map2alm(orig_d).block_until_ready()
+
+
+def time_it(func: Callable[[], None]) -> float:
+    timer = timeit.Timer(func)
+    number, _ = timer.autorange()
+    execution_time = min(timer.repeat(number=number)) / number
+    return execution_time
+
+
+print(time_it(map2alm_hp))
+# 0.0038599122400046326
+
+print(time_it(map2alm_jhp))
+# 0.437624677999338

jax_healpy-0.2.1/jax_healpy/__init__.py

@@ -0,0 +1,75 @@
+from jax import config as _config
+
+from .pixelfunc import (
+    UNSEEN,
+    ang2pix,
+    ang2vec,
+    isnpixok,
+    isnsideok,
+    maptype,
+    nest2ring,
+    npix2nside,
+    npix2order,
+    nside2npix,
+    nside2order,
+    nside2pixarea,
+    nside2resol,
+    order2npix,
+    order2nside,
+    pix2ang,
+    pix2vec,
+    pix2xyf,
+    reorder,
+    ring2nest,
+    vec2ang,
+    vec2pix,
+    xyf2pix,
+)
+from .sphtfunc import alm2map, map2alm
+
+__all__ = [
+    'UNSEEN',
+    'pix2ang',
+    'ang2pix',
+    'pix2xyf',
+    'xyf2pix',
+    'pix2vec',
+    'vec2pix',
+    'ang2vec',
+    'vec2ang',
+    # 'get_interp_weights',
+    # 'get_interp_val',
+    # 'get_all_neighbours',
+    # 'max_pixrad',
+    'nest2ring',
+    'ring2nest',
+    'reorder',
+    # 'ud_grade',
+    # 'UNSEEN',
+    # 'mask_good',
+    # 'mask_bad',
+    # 'ma',
+    # 'fit_dipole',
+    # 'remove_dipole',
+    # 'fit_monopole',
+    # 'remove_monopole',
+    'nside2npix',
+    'npix2nside',
+    'nside2order',
+    'order2nside',
+    'order2npix',
+    'npix2order',
+    'nside2resol',
+    'nside2pixarea',
+    'isnsideok',
+    'isnpixok',
+    # 'get_map_size',
+    # 'get_min_valid_nside',
+    # 'get_nside',
+    'maptype',
+    # 'ma_to_array',
+    'alm2map',
+    'map2alm',
+]
+
+_config.update('jax_enable_x64', True)