© The Authors 2024

1 Introduction

Spectral fitting is one of the cornerstones of X-ray astrophysics and has been widely used and developed since the 1990s as the number of X-ray instruments has increased. An X-ray spectrum, represented by a number of counts in energy bins, can be reduced to meaningful parameters using an appropriate physical model. Constraining the plausible values of these parameters given an X-ray spectrum is usually achieved using publicly available software, such as xspec (Arnaud 1996), SPEX (Kaastra et al. 1996), sherpa (Freeman et al. 2001; Doe et al. 2007), or ISIS (Houck & Denicola 2000).

The typical approach is to define a fit statistic (e.g. χ² or C-stat; see Cash 1979) and optimise it using gradient-free minimisation algorithms. This type of approach is efficient if the initial parameter estimates are close to their optimal values, but is highly dependent on the dimensionality and smoothness of the parameter space. In general, there is no guarantee that these algorithms will not get stuck in local minima, which would ultimately lead to incorrect physical interpretations. Moreover, deriving a meaningful error budget in this situation is not trivial and often requires an exhaustive evaluation of the fit statistic in an N-dimensional lattice in the neighbourhood of the best fit, which is not efficient as the number of parameters, N, increases.

On the other hand, Bayesian approaches provide a sample of parameters that are distributed according to the likelihood. The sample of posterior parameters carries information about the dispersion and is less likely to get stuck in a local minimum (e.g. Bonson & Gallo 2016; Choudhury et al. 2017; Barret & Cappi 2019). This type of approach requires orders of magnitude more computation time than direct minimisation. Posterior samples can also be generated using Markov chain Monte Carlo (MCMC) approaches, or other algorithms such as the nested sampling (Ns) implemented in BXA (Buchner et al. 2014; Buchner 2021), which can efficiently check their convergence. We recently proposed a new approach using neural networks to directly learn the posterior distribution, given a set of simulated spectra (see SIXSA, Barret & Dupourqué 2024). This simulation-based inference approach showed a similar performance compared to the existing solutions, while being much faster.

As a more general trend, we have seen major advances in the field of machine learning over the last decade, mainly due to the use of back-propagation algorithms (Rumelhart et al. 1986). Today’s neural networks can be parameterised with up to billions of parameters, and finding the best parameters to achieve a given task is unthinkable using classical minimisation algorithms. The back-propagation algorithm allows the gradient of any function to be estimated as long as it can be represented by a computational graph of differentiable components. This additional gradient information allows the use of powerful minimisation algorithms, such as the adam algorithm (Kingma & Ba 2017), which can find optimal values for parameters in high-dimensional spaces. Modern frameworks such as tensorflow (Abadi et al. 2015a,b), PyTorch (Paszke et al. 2019) and JAX (Bradbury et al. 2018) can provide this gradient information automatically. In particular, JAX is already used in the astrophysics community to facilitate the inference of cosmological parameters (e.g. Campagne et al. 2023; Piras & Mancini 2023; Balkenhol et al. 2024) by providing differentiable likelihood functions.

In this article we present jaxspec, a JAX-based X-ray spectral fitting package written in pure Python. Using differentiable programming, gradient-based MCMC samplers and variational inference (VI) algorithms, jaxspec quickly and efficiently provides distributions for the spectral model parameters. We describe the design of jaxspec and we benchmark it against existing packages to demonstrate the benefits of just-in-time compilation and auto-differentiability. By studying Hitomi/Soft X-ray Spectrometer (SXS) observations of the core of the Perseus cluster and constraining the gas motions using a Bayesian hierarchical model, we demonstrate the feasibility of VI for high-resolution X-ray spectral fitting. Finally, we discuss the philosophy behind the development of jaxspec and our desire to provide software that is modern, flexible, open source, and driven by the community.

2 Description of jaxspec

We have designed jaxspec as a software package written entirely in Python. By leveraging key packages from the data science ecosystem, the source code itself is greatly simplified compared to other existing solutions and some sustainability is ensured over time. In particular, jaxspec relies on two widely used and maintained dependences. The first is JAX (Bradbury et al. 2018), a library similar to tensorflow and pytorch for writing automatically differentiable code that can be compiled at runtime. JAX is widely used in the machine learning community and allows code to be executed transparently on core, graphical or tensor process units (CPUs, GPUs and TPUs, respectively) to achieve significant speed-up. The second is numpyro (Phan et al. 2019), a probabilistic programming library that implements the construction of hierarchical Bayesian models and the tools needed to perform the associated inference. In particular, numpyro includes a JAX implementation of the no U-turn sampler (Hoffman & Gelman 2014) and an interface for applying VI methods. We are distributing jaxspec to the community as a Python package that can be downloaded and installed directly from the pypi portal. The codebase is open source, publicly available in a GitHub repository¹ and documented².

Spectral models in jaxspec can be built in the same way as one would expect from xspec, namely, by organically nesting model components. We have already implemented some commonly used and analytically computable model components³ (e.g. blackbody spectra and power-law spectra) and tabulated absorption models, such as the Tübingen-Boulder interstellar medium absorption model tbabs (Wilms et al. 2000). jaxspec does not (yet) implement many of the models available in xspec, and porting legacy Fortran / C models to pure JAX is neither a quick nor an easy task. However, models that are interpolated on rectangular grids are readily implementable in jaxspec, such as the nsatmos neutron star atmosphere model (Heinke et al. 2006). As described in Sect. 5, we are counting on the community to update, maintain and share models, and we will do our best to help model developers integrate their work into this codebase.

Each spectral model is parameterised by a subset of the parameters θ. These models are folded into the instrumental response using matrix multiplications, which allow the expected number of counts in each instrumental bin, λ_i, to be estimated given a set of parameters and for a particular observing setup. Finding the best set of parameters to describe the observed data requires a metric to define how good a model is. In the case of X-ray spectral fitting, the detection of photons within instrumental bins can be described as a counting process, which means it can be modelled as a Poisson random variable. The corresponding log-likelihood⁴ for the expected photons, λ(θ), according to the model and the observed values, c, is written as $$\log ℒ(c\mid \theta )={\displaystyle \sum _i{c}_i}\log {\lambda }_i(\theta )-\log \text{Γ}\left(c_i+1\right)-{\lambda }_i(\theta )$$(1)

where ${\displaystyle {\sum }_i}$ denotes the summation over the instrument bins. This Poisson likelihood is directly related to the well-known Cash statistic (Cash 1979), which is actually a transformation of the Poisson likelihood that corresponds to the χ² statistic at high counts.jaxspec implement only on the Poisson likelihood, as the χ² likelihood has historically been used to study high-count spectra and to simplify error calculations, and is biased at lower counts (see, e.g. Humphrey et al. 2009; Kaastra 2017; Buchner & Boorman 2023). Finding the best set of parameters that maximises this likelihood and its associated variance can be achieved using Bayesian inference approaches.

3 Benchmarking Bayesian inference in jaxspec

Solving an inference problem with Bayesian approaches requires inverting this likelihood into a posterior distribution of θ constrained to the observed number of counts c. This requires defining prior distributions for θ that encapsulate our (lack of) knowledge about the expected values of our model parameters. Once the Bayesian inference problem is set up, the easiest way to solve it is to sample the parameter values {θ}_i from the posterior distribution. There are many methods for doing this, such as MCMC or NS.

3.1 Quality of the MCMC samples

Markov chain Monte Carlo sampling has been one of the most widely used approaches to obtaining parameter samples that are distributed according to the posterior distribution of a Bayesian inference problem since the Metropolis-Hastings algorithm was proposed (Hastings 1970). MCMC algorithms generate chains of samples by moving gradually through the parameter space, making small, local changes at each step. They are designed so that the stationary distribution of these chains is the posterior distribution of the parameters. As such, reaching the stationary distribution requires a sufficiently large number of iterations, which is not trivial to quantify. Many solutions based on MCMC sampling can provide posterior samples. The affine invariant ensemble sampler (AIES, Goodman & Weare 2010) uses multiple chains that are not independent (often referred to as walkers) to shift the proposal distribution towards optimal regions of the parameter space. This is the sampler implemented in the well-known emcee package (Foreman-Mackey et al. 2013), and also the default sampler implemented in xspec. While the traditional MCMC sampling algorithm requires numerous steps to achieve stationarity, gradient-based samplers can significantly reduce the number of steps required. Hamiltonian Monte Carlo (e.g. Betancourt & Girolami 2013) uses the gradient of the likelihood to perform ‘energy saving’ steps in the parameter space, allowing much better exploration and faster convergence. In particular, NUTS (Hoffman & Gelnian 2014) uses this dynamic to adoptively explore the likelihood, setting the step size depending on the local curvature. Since jaxspec provides auto-differentiable likelihoods, we set the numpyro implementation of NUTS as the default sampler, and compared the numpyro and xspec implementations of AIES.

To compare the performance of jaxspec with existing solutions, we chose to solve a similar inference problem using the different methods. The task was to obtain the posterior distribution of the five parameters of an absorbed dual-component model (tbabs*(powerlaw+blackbodyrad)) fitted to an observation of NGC 7793 ULX-4. Ultra-luminous X-ray sources (ULXs, e.g., Kaaret et al. 2017) are highly accreting stellar-mass binary systems, in which the precise physical processes at play are still largely debated. In particular, some ULXs display a two-component spectral emission in X-rays (e.g., Sutton et al. 2013), the first component (thermal) generally being interpreted as an accretion disk and the second (power-law-like) as either a hot accretion column or a Compton-scattering outflow (e.g., Gúrpide et al. 2021). The relative temporal evolution of one component with respect to the other is generally different in ULXs compared to standard X-ray binaries (e.g., Koliopanos et al. 2017). As such, precisely and carefully constraining the parameters from both components can shed some light on the physics of these objects. We used the XMM-Newton observations of NGC 7793 ULX-4 as reduced in Quintin et al. (2021), and focused on the spectrum produced by the EPIC-PN instrument, with no background. All the calculations described in this section are performed on a single core in order to obtain comparable results between the different methods used, but it is worth noting that they all benefit from being performed on multiple cores.

For this benchmark, we ran MCMC chains using AIES and NUTS and measured how quickly the chain converged to a stationary distribution using the $\widehat{R}$ statistic (Vehtari et al. 2021) computed for an increasing number of samples. The $\widehat{R}$ statistic compares the intra- and inter-chain variance to assess their individual convergence. AIES was run with 2⁵ = 32 walkers, as suggested by the number of parameters in our model, while NUTS was run with four chains. We ran 10³ burn-in steps and 10⁵ sampling steps, and computed the $\widehat{R}$ statistic for an increasing number of samples, with a step of 1000. The $\widehat{R}$ statistic was computed using only the best four chains of each sampler, defined in terms of the best likelihood on average. We show in the left panel of Fig. 1 that the $\widehat{R}$ decreases as the number of steps increases, as the chains begin to converge. The usual convergence threshold of $\widehat{R}<1.01$ is reached well within the first 1000 steps with NUTS, ~5000 steps with AIES in jaxspec, and 25 000 steps with AIES in xspec. The right panel of Fig. 1 shows the individual samples of a single chain using the norm_BB parameter. As expected, NUTS is much better at providing uncorrelated samples, while AIES and especially the xspec implementation provide samples with visible correlations, meaning that it takes a large number of steps and additional thinning to get a meaningful sample of parameters from the posterior distribution.

Fig. 1 Comparison of the sampling quality using jaxspec and xspec. Left: $\widehat{R}$ statistic computed for 100 runs of 105 steps of NUTS and AIES using the jaxpsec and xspec implementations. The usual convergence threshold of $\widehat{R}<1.01$ is shown in black. Right: value of the parameter normBB for 104 steps of a single chain in the sample using NUTS and AIES with jaxpsec and xspec implementations. The solid line represents a 100-step moving average of the parameter value.

3.2 Comparison of inference times for various methods

Another common approach to solving Bayesian inference problems is to use NS (see e.g. Büchner 2023), which is a powerful alternative to MCMC sampling. NS relies on a population of points in the parameter space that evolves according to their likelihood at each step. As the number of steps increases, the points with the lowest likelihood are excluded and new points are drawn according to the distribution of the remaining points. Such evolutionary approaches shrink the volume occupied by the samples until there is no further improvement in likelihood. Unlike MCMC approaches, NS provides a stopping criterion and can explore a much larger part of the parameter space, providing robust estimates for the posterior distribution, especially in multi-mode cases. A popular implementation of NS used in X-ray spectral fitting is BXA (Buchner et al. 2014), which interfaces the NS algorithm of ultranest (Büchner 2021) with xspec or sherpa.

We used different inference algorithms and compared their results and total inference time. We used the previously introduced NUTS and AIES samplers implemented in j axspec and xpsec. We also compared the existing BXA sampler (with xspec interface) to the implementation of phantom-powered nested sampling (PPNS, Albert 2023) provided in jaxns. We used the same observational setup as described in the previous subsection and fitted the same five-parameter model. We ran all algorithms with default hyperparameters for sampling and termination, and used the previously determined number of steps to convergence for the MCMC sampler.

The time for 100 runs of each inference approach is plotted in Fig. 2. For this particular case, NUTS is an order of magnitude faster than the other approaches, AIES in jaxspec is slightly faster than other approaches, and PPNS in jaxspec is the equivalent to NS in BXA in terms of speed. For a well-behaved inference problem, NUTS can generate stable posterior distributions in less than a minute. In Fig. 3, we show the results of a single run of each algorithm. The posterior distributions are in good agreement. We note that the AIES implementation in numpyro, which is used in jaxspec, produces marginal distributions that are tighter than the other algorithms, due to differences in the underlying implementations. In a nutshell, the xspec implementation uses the ‘stretch’ move from Goodman & Weare (2010) while j axspec uses the ‘differential evolution’ move from Nelson et al. (2014). The second converges faster but struggles to explore the tail of the distribution, slightly underestimating the spread of the parameters. Increasing the number of walkers should solve this. In either case, the user can choose to use the stretch move or a combination of the two moves in numpyro to achieve the desired behaviour. A composite model like the one presented in this benchmark is a standard use case in the community, meaning this benchmark is representative of a typical inference case for accreting objects. As it is easy to choose a specific sampler in jaxspec, the user can try different approaches to solving a single problem, leading to robust results in the end.

Fig. 2 Time per run and standard deviation for NUTS, AIES and NS in both jaxspec, BXA and xspec, estimated for 100 runs of each algorithm on a five-parameter, three-component inference problem.

4 High-resolution spectroscopy with jaxspec

For the sake of demonstrating the capabilities of jaxspec, we considered the high-resolution X-ray spectrum from the Hitomi/SXS observations of the Perseus core, obtained with the exquisite spectral resolution of ~5eV, which enabled the first high-precision direct measurement of the motions of the intracluster gas (Hitomi Collaboration 2016, 2018). By studying the Fe xxv Heα, Heβ and Heγ complexes, the bulk and velocity dispersions have been constrained to ~50 km/s and ~200 km/s, respectively, in the central regions of the cluster. The high spectral resolution results in extremely dense spectra with large response matrices, which makes direct minimisation more difficult and the general inference more computationally expensive within existing frameworks. Here we demonstrate the use of jaxspec for accurately modelling the observations of the Perseus core with Hitomi. To achieve this, we used the publicly available data, bearing in mind that these data are subject to calibration problems that limit the scientific scope of the results. We used the pipeline products from ObsIds 40, 50 and 60. The spectra weree optimally binned using the algorithm from Kaastra & Bleeker (2016). As in Hitomi Collaboration (2016), we focused on modelling the Fe XXV Hea complex with Gaussian emission lines of known energies in the [6.49, 6.61] keV band. The measurement of the centroid shift and the broadening of these lines provide information about the dynamics of the intra-cluster medium along the line of sight, and is a direct measure of the turbulent motions. In Appendix B, we define a spectral model consisting of nine Gaussian emission lines with known energies to assess the centroid shift and broadening resulting from the gas motion. We used the cluster redshift determination from Hitomi Collaboration (2018) and assumed a Gaussian prior such as z ~ 𝒩(0.017284,0.00005). All the model parameters were constrained using the three observations simultaneously.

As a novelty, we solved this Bayesian inference problem using a VI approach. This family of methods takes a different approach from the MCMC family. Instead of sampling directly from the posterior distribution, it analytically computes an approximation to the posterior density. It requires the specification of a family of distributions, q(θ), called the variational distribution, which is fitted to the true posterior distribution. In practice, this is equivalent to saying that the posterior distribution would be well approximated by, for example a Gaussian distribution for each variable, a multivariate Gaussian for all parameters with a mean and a correlation matrix to be determined, or any analytically computable and parameterised probability density. Selecting a surrogate posterior distribution from a set of known densities transforms the inference problem into an optimisation problem in which we try to recover the best parameters of the surrogate distribution. This in turn requires defining a distance between the variational distribution and the true posterior, minimised by adjusting the variational distribution parameters. One commonly used measure between two distributions in machine learning studies is the Kullback-Leibler divergence (KL divergence; Kullback & Leibler 1951). It shows how much information is lost when a model distribution is used to approximate the true underlying distribution. For the optimisation process in VI, the commonly used metric is the Evidence Lower Bound, which is directly derived from the KL divergence between the variational distribution and the posterior distribution. For further insight, we direct the interested reader to the review by Blei et al. (2017). The ELBO is defined as the difference between the expectation (with respect to the variational distribution, q(θ)) of the logarithm of the joint distribution, p (θ, X), and the logarithm of the variational distribution, $$\mathrm{ELBO}(q(\theta ))={\mathbb{𝔼}}_{q(\theta )}\{\log p(\theta ,X)\}-{\mathbb{𝔼}}_{q(\theta )}\{\log q(\theta )\},$$(2)

while the best surrogate posterior distribution, q*(θ), is the one that maximises the ELBO, $$q^{*}(\theta )=\underset{q(\theta )}{\mathrm{argmin}}\{-\mathrm{ELBO}(q(\theta ))\}.$$(3)

We aim to demonstrate the feasibility and effectiveness of VI for X-ray spectral fitting. To our knowledge, this is the first time that such methods have been used for X-ray spectroscopy in astrophysics, with only a handful of works on VI in astrophysics in general (e.g. Gunapati et al. 2022). It has to be noted that classical approaches can also work here. We assumed a multivariate Gaussian distribution as the variational distribution for these observations. This optimisation problem takes particular advantage of the differentiability of its metric, provided here by the auto-differentiation in JAX. This makes it possible to mobilise first-order optimisers (e.g. adam, Kingma & Ba 2017) using gradient information to find a satisfactory minimum for the problem. These algorithms require the user to set hyperparameters, such as the learning rate, which defines the length of each step in the parameter space, or the total number of iterations. For this particular problem, we used the adamw optimiser (Loshchilov & Hutter 2019) with default parameters and a learning rate of 5⁻⁴ for 10⁵ iterations. The gradient descent can be run on a GPU and takes ~7 minutes on a Nvidia A2. In Fig. 4 we show the ELBO evolution for the 10⁵ adamw steps, the resulting redshift, bulk motion and velocity dispersion in the observer frame, and the distribution of the posterior folded spectra. The stationarity of ELBO, coupled with the small residuals on the observed spectra, indicates a qualitative fit. The line-of-sight bulk motion of 51 ± 15 km/s and the velocity dispersion of 196 ± 4 km/s are compatible with those measured by Hitomi Collaboration (2016) in the central region of Perseus. However, we stress that our results are not directly comparable, as our analysis is averaged over all Perseus observations and all SXS pixels, and a more exhaustive treatment that includes systematic uncertainties and calibration issues is beyond the scope of this paper. We also show that the posterior has the expected behaviour, as no insight into the cluster redshift is gained due to a degeneracy with the bulk motion captured by our variational distribution. This confirms the efficiency of VI methods for X-ray spectral fitting.

Fig. 3 Posterior distribution of tbabs*(powerlaw+blackbodyrad) parameters from this benchmark, drawn with the 104 last steps of a run of NUTS and AIES using the jaxpsec and xspec implementations, and as provided by the NS algorithms implemented in BXA and jaxspec.

5 Enabling a community-driven package

As the X-ray spectral fitting ecosystem is already relatively crowded, we feel it is important to justify our efforts to add a new library to it. Beyond the efficiency of JAX and the complexity of scientific use cases enabled by the implementation of state-of-the-art sampling algorithms, the key strengths of jaxspec are its ease of installation and use. It is worth noting that our long-term goal is to exploit the scientific data produced by high-resolution spectrometers such as XR/SM/Resolve and the future newAthena/X-ray Integral Field Unit (XIFU; Barret et al. 2023), which would greatly benefit from modern and fast spectral fitting software.

Fig. 4 Variational inference results for the core of Perseus as seen by Hitomi. (a) ELBO metric during variational distribution optimisation. (b) Redshift, bulk motion (vbulk), and velocity dispersion (σv) posterior distributions obtained for the central regions of the Perseus cluster as seen by Hitomi/SXS, superimposed on the prior distribution for the redshift and the previous results from (Hitomi Collaboration 2016) in the central region. (c) Count spectra of the observations used in this spectral fit, superimposed on the folded model for each observation, and the corresponding residuals in σ.

5.1 From the casual user’s point of view

Our objective is to provide users with a straightforward software solution for spectral inference that is both readily deployable and intuitive to use. As jaxspec is not dependent on xspec or any equivalent, the installation process is as straightforward as that of any Python package, requiring only the pip command. The software is fully documented and provides examples of inference. As detailed in Appendix A, the syntax for building models and fitting them to observations is kept as simple as possible. For scientists using other tools, the interface we propose should be familiar. The default setups for inference are fast and efficient at providing reliable scientific results. Nevertheless, as jaxspec exposes its likelihoods, the user can extend its use beyond what is provided in the package.

5.2 From the contributor’s point of view

The jaxspec codebaserepresents a fresh start from the existing frameworks. It greatly reduces the code complexity when compared to existing alternatives. Relying on the Python ecosystem for data science reduces the need for maintaining heavy boilerplates. The continuous integration framework, with automated tests, code linters, and automated documentation, allows updates and patches to be deployed easily. The package is open source and hosted on GitHub, and individual contributions are encouraged. This facilitates a collaborative environment in which users can raise issues and directly contribute to the codebase by correcting and extending source code. In particular, it is possible to extend the spectral models, as was previously possible with xspec and other similar software. Analytical models and models interpolated on regular grids can be readily implemented.

6 Conclusions

In this paper, we have presented jaxspec, a pure Python and JAX software for X-ray spectral fitting. Since JAX is differentiable and compilable, powerful samplers such as NUTS can be used to solve the inference problems quickly and reliably, running seamlessly on CPUs or GPUs. We show that it outperforms existing software by an order of magnitude in execution time on a five-parameter, three-component model, while providing robust posterior distributions. The compilable and auto-differentiable likelihood provided by jaxspec allows the use of VI, which can provide accurate results for high-dimensional models, as demonstrated on high-energy-resolution Hitomi/SXS spectra. The likelihood functions built into jaxspec are exposed so that it can be interfaced with external packages such as ultranest, giving users many options.

The current limitation of jaxspec is the library of spectral models that can be fitted, as models must be translated in pure JAX. Regularly tabulated models are already easy to port. We also plan to integrate existing features from other software, such as direct minimisation, multi-instrument gain and shift, convolutional spectral components, and a restriction prior as implemented in SIXSA, and to systematically implement xspec models, which could be used with non-differentiable methods. We also plan to explore the possibility of using neural networks to create surrogate models to port existing models to our codebase. If successful, this approach would have the advantage of providing naturally differentiable codes that are fast to run on CPUs and GPUs, at the cost of controlled error in model evaluation (see e.g. Varma et al. 2019a,b for gravitational waves or Euclid Collaboration 2021 for cosmology applications). As we have taken great care in building this software using good software engineering practices, with a contribution policy, automated testing, and code quality linters, we hope that it will evolve according to the needs of the community. We encourage individual contributions to the codebase and hope to provide astrophysicists with a simple yet reliable tool for X-ray spectral fitting.

Acknowledgements

S.D. and D.B. would like to thank Florent Castellani, and Alexeï Molin for their early contributions to this project, and Fabio Acero and Johannes Buchner for the fruitful discussions about spectral fitting. S.D. and D.B. thank the Centre National d’Etudes Spatiales for financial support. All the calculations involving GPUs were carried out on the SSPCloud (Comte et al. 2022), a platform freely provided to French academics. This work made use of various open-source packages such as numpy (Harris et al. 2020) matplotlib (Hunter 2007), astropy (Astropy Collaboration 2013, 2018, 2022), ChainConsumer (Hinton 2016), cmasher (Velden 2020), haiku (Hennigan et al. 2020).

Appendix A Basic cookbook

The user can install jaxspec with a simple call to pip.

pip install jaxspec

We recommend the user start from a fresh Python 3.10 or 3.11 environment. To build a spectral model, additive and multiplicative components can be tied together as follows :

from jaxspec.model.addi tive import Powerlaw
from jaxspec.model.addi tive import Blackbodyrad
from jaxspec.model.multiplicative import Tbabs
model = Tbabs() * (Powerlaw() + Blackbodyrad())

If the data used are compliant with the OGIP standard, they can be loaded as follows :

from jaxspec.data import ObsConfiguration
obs = ObsConfiguration.from_pha_file(
‘path/to/data.pha’,
low_energy=0.5, high_energy=10
# RMF and ARF can be explicitly stated
#rmf_path=‘path/to/data.rmf’,
#arf_path=‘path/to/data.arf’,
)

Then, the data can be fitted by providing prior distributions.

from jaxspec.fit import MCMCFitter
from numpyro.distributions import Uniform,
LogUniform
prior = {
“blackbodyrad_l_kT”: Uniform (0, 3),
“blackbodyrad_l_norm”: LogUniform(1e-3, 1),
“ powerlaw_1_alpha”: Uni form(0, 3) ,
“powerlaw_1_norm”: LogUniform(1e-5, 1e-1),
“tbabs_1_N_H” : Uniform(0, 1)
}
result = MCMCFitter(model, prior, obs).fit()

From this result, one can plot the posterior distributions, and perform posterior predictive checks.

result.plot_corner()
result.plot_ppc()

We refer the interested user to the documentation of jaxspec, which will evolve, unlike what is presented in this article. It can be found at http://jaxspec.rtfd.io/.

Appendix B Perseus model

The unfolded spectrum is modelled by nine Gaussian emission lines with a common broadening ${\displaystyle \sum ={\sigma }_v/c}$ and energy shift δ = v_bulk/c. For each of these lines, we fitted a free normalisation and an additional intrinsic energy perturbation, which should be much smaller than the total shift.

$$E_{L,i}=\frac{E_{\text{tab},i}}{(1+z)(1+\delta )}+\text{Δ}{E}_{\text{noise},i}$$(B.1)

where E_tab,i is the expected energy of line i which can be found in the appendix of Hitomi Collaboration (2016). This spectrum is then folded into Hitomi responses for the three observations; for two of them we added a gain and shift parameter. An additional count-rate baseline was added for each instrument to model the continuum. The prior used for each parameter and the posterior values are displayed in Table B.1. We also used the automatic re-parametrisation provided by numpyro to rescale and de-center the Gaussian prior distributions, and represent the loguniform prior in a flat space, which regularises the likelihood and facilitates the gradient descent.

References

All Tables

All Figures

Fig. 1 Comparison of the sampling quality using jaxspec and xspec. Left: $\widehat{R}$ statistic computed for 100 runs of 105 steps of NUTS and AIES using the jaxpsec and xspec implementations. The usual convergence threshold of $\widehat{R}<1.01$ is shown in black. Right: value of the parameter normBB for 104 steps of a single chain in the sample using NUTS and AIES with jaxpsec and xspec implementations. The solid line represents a 100-step moving average of the parameter value.

In the text

	Fig. 2 Time per run and standard deviation for NUTS, AIES and NS in both jaxspec, BXA and xspec, estimated for 100 runs of each algorithm on a five-parameter, three-component inference problem.
In the text

	Fig. 3 Posterior distribution of tbabs*(powerlaw+blackbodyrad) parameters from this benchmark, drawn with the 104 last steps of a run of NUTS and AIES using the jaxpsec and xspec implementations, and as provided by the NS algorithms implemented in BXA and jaxspec.
In the text

Fig. 4 Variational inference results for the core of Perseus as seen by Hitomi. (a) ELBO metric during variational distribution optimisation. (b) Redshift, bulk motion (vbulk), and velocity dispersion (σv) posterior distributions obtained for the central regions of the Perseus cluster as seen by Hitomi/SXS, superimposed on the prior distribution for the redshift and the previous results from (Hitomi Collaboration 2016) in the central region. (c) Count spectra of the observations used in this spectral fit, superimposed on the folded model for each observation, and the corresponding residuals in σ.

In the text