https://github.com/wcxve/xspex

JAX interface for XSPEC spectral models.

https://github.com/wcxve/xspex

Science Score: 26.0%

This score indicates how likely this project is to be science-related based on various indicators:

  • CITATION.cff file
  • codemeta.json file
    Found codemeta.json file
  • .zenodo.json file
    Found .zenodo.json file
  • DOI references
  • Academic publication links
  • Academic email domains
  • Institutional organization owner
  • JOSS paper metadata
  • Scientific vocabulary similarity
    Low similarity (14.2%) to scientific vocabulary

Keywords

astronomy astrophysics jax python spectral-analysis xla xspec
Last synced: 6 months ago · JSON representation

Repository

JAX interface for XSPEC spectral models.

Basic Info
  • Host: GitHub
  • Owner: wcxve
  • License: gpl-3.0
  • Language: Python
  • Default Branch: main
  • Homepage:
  • Size: 288 KB
Statistics
  • Stars: 4
  • Watchers: 2
  • Forks: 1
  • Open Issues: 3
  • Releases: 6
Topics
astronomy astrophysics jax python spectral-analysis xla xspec
Created almost 2 years ago · Last pushed 6 months ago
Metadata Files
Readme License

README.md

xspex

JAX Interface for XSPEC Spectral Models.

PyPI - Python Version PyPI - Version License: GPL v3
Coverage Status

Installation

NOTE: Before installation, HEASoft & XSPEC v12.12.1+ are required to be installed on your system. You can download from here, or install from conda.

Once the HEADAS environment has been initialized, xspex can be installed directly from PyPI using:

console pip install xspex

Examples

Basic Usage

```python import os

os.environ['XLAFLAGS'] = '--xlaforcehostplatformdevicecount=2'

import jax import jax.numpy as jnp import numpy as np import xspex as xx

Double precision is required for XSPEC models

jax.config.update('jaxenablex64', True)

Get APEC model function

fn, info = xx.get_model('apec')

Define parameters and energy grid

params = jnp.array([1.0, 1.0, 0.0]) egrid = jnp.linspace(0.1, 0.2, 6)

Evaluate the model function

value = fn(params, egrid) print(value)

output:

[1.29398149 0.43733974 0.11998129 0.08725635 0.07757024]

```

JAX Transformations

xspex provides JAX automatic differentiation support for XSPEC models through finite difference approximation. This allows seamless integration with JAX's transformations like grad, jacfwd, jacrev, etc.

Computing Gradients

```python

Get gradient function with respect to parameters

grad_fn = jax.grad(lambda p, e: jnp.sum(fn(p, e)))

Compute gradient, note that the abundance and redshift are fixed by default

grad = grad_fn(params, egrid) print(grad)

output:

[-3.1665168 0. 0. ]

```

Computing Jacobian

```python

Get Jacobian function

jac_fn = jax.jacfwd(lambda p, e: fn(p, e)) # or jax.jacrev, jax.jacobian

Compute Jacobian matrix

jacobian = jac_fn(params, egrid) print(jacobian)

output:

[[-2.01717805 -0. -0. ]

[-1.05626962 -0. -0. ]

[-0.03252301 -0. -0. ]

[-0.02018553 -0. -0. ]

[-0.0403606 -0. -0. ]]

```

Vectorization with vmap

```python

Create multiple parameter sets

param_sets = jnp.array([ [0.5, 1.0, 0.0], [1.0, 1.0, 0.0], [2.0, 1.0, 0.0], ])

Vectorize the function

vmappedfn = jax.vmap(fn, inaxes=(0, None)) results = vmappedfn(paramsets, egrid) print(results)

output:

[[0.52477309 0.56379027 0.13421626 0.11663016 0.17570166]

[1.29398149 0.43733974 0.11998129 0.08725635 0.07757024]

[0.18578006 0.10865312 0.08878279 0.07398064 0.06353859]]

```

Parallel evaluation with pmap

```python

Replicate parameters across devices

param_sets = jnp.array([ [1.0, 1.0, 0.0], [2.0, 1.0, 0.0], ])

pmappedfn = jax.pmap(fn, inaxes=(0, None)) results = pmappedfn(paramsets, egrid) print(results)

output:

[[1.29398149 0.43733974 0.11998129 0.08725635 0.07757024]

[0.18578006 0.10865312 0.08878279 0.07398064 0.06353859]]

```

Custom Finite Difference Automatic Differentiation

```python

Create model with custom finite difference settings

fn, info = xx.getmodel('powerlaw') fn2 = xx.definefdjvp( # see the docstring for more details fn, info, delta=1e-6, # Custom step size (relative to parameter value) method='central', # 'central' or 'forward' finite differences fixed=None # Optional: specify which parameters to keep fixed ) ```

Owner

  • Name: W.-C. Xue
  • Login: wcxve
  • Kind: user
  • Company: Institute of High Energy Physics

Ph. D. student of High Energy Astrophysics

GitHub Events

Total
  • Create event: 29
  • Release event: 8
  • Issues event: 3
  • Watch event: 3
  • Delete event: 25
  • Issue comment event: 54
  • Push event: 123
  • Pull request review comment event: 66
  • Pull request review event: 68
  • Pull request event: 48
  • Fork event: 1
Last Year
  • Create event: 29
  • Release event: 8
  • Issues event: 3
  • Watch event: 3
  • Delete event: 25
  • Issue comment event: 54
  • Push event: 123
  • Pull request review comment event: 66
  • Pull request review event: 68
  • Pull request event: 48
  • Fork event: 1

Issues and Pull Requests

Last synced: 6 months ago

All Time
  • Total issues: 2
  • Total pull requests: 26
  • Average time to close issues: about 1 year
  • Average time to close pull requests: about 15 hours
  • Total issue authors: 1
  • Total pull request authors: 3
  • Average comments per issue: 0.0
  • Average comments per pull request: 1.54
  • Merged pull requests: 17
  • Bot issues: 0
  • Bot pull requests: 14
Past Year
  • Issues: 1
  • Pull requests: 25
  • Average time to close issues: 10 months
  • Average time to close pull requests: about 15 hours
  • Issue authors: 1
  • Pull request authors: 3
  • Average comments per issue: 0.0
  • Average comments per pull request: 1.6
  • Merged pull requests: 16
  • Bot issues: 0
  • Bot pull requests: 14
Top Authors
Issue Authors
  • wcxve (2)
Pull Request Authors
  • wcxve (12)
  • pre-commit-ci[bot] (7)
  • dependabot[bot] (7)
Top Labels
Issue Labels
bug (1)
Pull Request Labels
dependencies (7) python (4) github_actions (3)

Packages

  • Total packages: 1
  • Total downloads:
    • pypi 63 last-month
  • Total dependent packages: 1
  • Total dependent repositories: 0
  • Total versions: 8
  • Total maintainers: 1
pypi.org: xspex

Access Xspec models and corresponding JAX/XLA ops.

  • Versions: 8
  • Dependent Packages: 1
  • Dependent Repositories: 0
  • Downloads: 63 Last month
Rankings
Dependent packages count: 9.5%
Average: 36.2%
Dependent repos count: 62.9%
Maintainers (1)
Last synced: 6 months ago

Dependencies

pyproject.toml pypi
setup.py pypi
.github/workflows/ci.yml actions
  • actions/checkout v4 composite
  • actions/download-artifact v4 composite
  • actions/upload-artifact v4 composite
  • astral-sh/setup-uv v6 composite
  • codecov/codecov-action v5 composite
  • conda-incubator/setup-miniconda v3 composite
  • pypa/gh-action-pypi-publish v1.12.4 composite