jaxspec

https://jaxspec.readthedocs.io/en/stable/
Admin
jaxspec is a pure Python library for statistical inference on X-ray spectra that allows users to...
Tokens:169,368
Snippets:754
Trust Score:9.6
Update:1 month ago
Show doc for...
Tokens:
Raw
### Setup and define spectral model and priors

Source: https://jaxspec.readthedocs.io/en/stable/examples/calibration

Initializes numpyro, imports jaxspec components, defines the spectral model Tbabs()*(Powerlaw()+Blackbodyrad()), loads observations from the example dataset, and sets up the prior dictionary for the spectral parameters.

```python
%%capture
# Hide the output of this cell
import numpyro

numpyro.enable_x64()
numpyro.set_platform("cpu")
numpyro.set_host_device_count(4)

import numpyro.distributions as dist
from jaxspec.fit import VIFitter
from jaxspec.data.util import load_example_obsconf
from jaxspec.model.additive import Blackbodyrad, Powerlaw
from jaxspec.model.multiplicative import Tbabs
from jaxspec.model.background import BackgroundWithError

spectral_model = Tbabs() * (Powerlaw() + Blackbodyrad())

obsconfs = load_example_obsconf("NGC7793_ULX4_ALL")

prior = {
    "spectrum.powerlaw_1.alpha": dist.Uniform(1, 3),
    "spectrum.powerlaw_1.norm": dist.LogUniform(1e-5, 1e-3),
    "spectrum.blackbodyrad_1.kT": dist.Uniform(0, 2),
    "spectrum.blackbodyrad_1.norm": dist.LogUniform(1e-2, 1),
    "spectrum.tbabs_1.nh": dist.Uniform(0, 1),
}
```

--------------------------------

### Install pre-commit hooks

Source: https://jaxspec.readthedocs.io/en/stable/contribute

Install pre-commit hooks so they run automatically before each commit.

```bash
uv run pre-commit install

```

--------------------------------

### Install uv

Source: https://jaxspec.readthedocs.io/en/stable/contribute

Install uv, the dependency manager used by jaxspec, inside the activated environment.

```bash
pip install uv

```

--------------------------------

### Set up MCMC fitter with jaxspec and numpyro

Source: https://jaxspec.readthedocs.io/en/stable/examples/work_with_arviz

Configures numpyro for CPU and x64, defines a spectral model (Tbabs * (Powerlaw + Blackbodyrad)), loads example observation data, sets priors, and creates an MCMCFitter instance. This setup is required before running the MCMC fit.

```python
%%capture 
# Hide the output of this cell
import numpyro

numpyro.enable_x64()
numpyro.set_platform("cpu")
numpyro.set_host_device_count(4)

import numpyro.distributions as dist
from jaxspec.fit import MCMCFitter
from jaxspec.data.util import load_example_obsconf
from jaxspec.model.additive import Blackbodyrad, Powerlaw
from jaxspec.model.multiplicative import Tbabs

spectral_model = Tbabs() * (Powerlaw() + Blackbodyrad())

obsconf = load_example_obsconf("NGC7793_ULX4_PN")

prior = {
    "spectrum.powerlaw_1.alpha": dist.Uniform(1, 3),
    "spectrum.powerlaw_1.norm": dist.LogUniform(1e-5, 1e-3),
    "spectrum.blackbodyrad_1.kT": dist.Uniform(0, 2),
    "spectrum.blackbodyrad_1.norm": dist.LogUniform(1e-2, 1),
    "spectrum.tbabs_1.nh": dist.Uniform(0, 1),
}

fitter = MCMCFitter(spectral_model, prior, obsconf)

```

--------------------------------

### Install jaxspec from source (non-editable)

Source: https://jaxspec.readthedocs.io/en/stable/frontpage/installation

Installs the latest development version directly from the GitHub repository without cloning. Use this if you do not need to edit the code.

```bash
pip install git+https://github.com/renecotyfanboy/jaxspec --upgrade

```

--------------------------------

### Install jaxspec from source (development)

Source: https://jaxspec.readthedocs.io/en/stable/frontpage/installation

Clones the jaxspec repository and installs it in editable mode for development. Use this if you plan to edit the source code.

```bash
git clone https://github.com/renecotyfanboy/jaxspec
cd jaxspec
pip install -e .

```

--------------------------------

### Sync dependencies with uv

Source: https://jaxspec.readthedocs.io/en/stable/contribute

Install jaxspec and all dependency groups exactly as recorded in uv.lock, matching CI behavior. Drop --frozen only when intentionally changing dependencies.

```bash
uv sync --all-groups --frozen

```

--------------------------------

### Install jaxspec with pip

Source: https://jaxspec.readthedocs.io/en/stable/frontpage/installation

Installs the latest stable version of jaxspec from PyPI using pip. Run inside an environment with a compatible Python version.

```bash
pip install jaxspec

```

--------------------------------

### Create ObsConfiguration from instrument and observation

Source: https://jaxspec.readthedocs.io/en/stable/examples/load_data

Builds a folding matrix from an instrument and observation, specifying the energy band from 0.5 to 8 keV. This is the primary method for defining an observational setup.

```python
from jaxspec.data import ObsConfiguration

folding_model = ObsConfiguration.from_instrument(
    instrument, observation, low_energy=0.5, high_energy=8
)
```

--------------------------------

### Load instrument from OGIP files

Source: https://jaxspec.readthedocs.io/en/stable/examples/load_data

Use Instrument.from_ogip_file to load an instrumental setup from a response matrix (.rmf) and ancillary response (.arf) file. Returns an Instrument object, a subclass of xarray.Dataset.

```python
from jaxspec.data import Instrument

instrument = Instrument.from_ogip_file("data/PN.rmf", arf_path="data/PN.arf")

```

--------------------------------

### Integrate with integration limits as parameters

Source: https://jaxspec.readthedocs.io/en/stable/references/integrate

Shows that the integration limits can be parameters. The example defines an erf function using integrate_interval with a variable upper limit.

```python
def erf(x):

    def integrand(t):
        return 2/jnp.sqrt(jnp.pi) * jnp.exp(-t**2)

    return integrate_interval(integrand)(0, x)

print(erf(1)) # 0.84270084
```

--------------------------------

### Initialize zeus sampler with jaxspec model

Source: https://jaxspec.readthedocs.io/en/stable/examples/external_samplers

Sets up a zeus ensemble sampler for a jaxspec Bayesian model. The starting points are keyed by parameter names and multiplied by random noise to create initial positions for each walker. The log-probability function is JIT-compiled and vectorized with jax.jit and jax.vmap.

```python
ndim, nwalkers = len(bayesian_model.parameter_names), 40 # `parameter_names` is sorted, so never hand-write a positional vector: key the # starting point by name and let the model order it. start = { "spectrum.powerlaw_1.alpha": 2.0, "spectrum.powerlaw_1.norm": 3e-4, "spectrum.blackbodyrad_1.kT": 0.7, "spectrum.blackbodyrad_1.norm": 0.2, "spectrum.tbabs_1.nh": 0.1, } p0 = np.asarray(bayesian_model.dict_to_array(start)) * np.random.normal( loc=1, scale=0.1, size=(nwalkers, ndim) ) @jax.jit @jax.vmap def zeus_log_prob(parameters: np.ndarray) -> jax.typing.ArrayLike: return bayesian_model.log_posterior_prob(bayesian_model.array_to_dict(parameters)) sampler = zeus.EnsembleSampler(nwalkers, ndim, zeus_log_prob, vectorize=True) sampler.run_mcmc(p0, 10_000, progress=True);
```

--------------------------------

### BayesianModel with InstrumentModel for PN, MOS1, MOS2

Source: https://jaxspec.readthedocs.io/en/stable/references/instrument

Shows how to pass an instrument_model dict to BayesianModel. The example sets PN to None and provides InstrumentModel instances with ConstantGain and ConstantShift for MOS1 and MOS2.

```python
BayesianModel(
    spectral_model, prior, observations,
    instrument_model={
        "PN": None,
        "MOS1": InstrumentModel(gain=ConstantGain(), shift=ConstantShift()),
        "MOS2": InstrumentModel(gain=ConstantGain(), shift=ConstantShift()),
    },
)
```

--------------------------------

### Run zeus MCMC with compiled posterior log-probability

Source: https://jaxspec.readthedocs.io/en/stable/examples/external_samplers

Sets up and runs the zeus sampler using a jitted and vmap'd posterior log-probability function. Builds starting points using dict_to_array and runs 10,000 steps with 40 walkers.

```python
ndim, nwalkers = len(bayesian_model.parameter_names), 40

# `parameter_names` is sorted, so never hand-write a positional vector: key the
# starting point by name and let the model order it.
start = {
    "spectrum.powerlaw_1.alpha": 2.0,
    "spectrum.powerlaw_1.norm": 3e-4,
    "spectrum.blackbodyrad_1.kT": 0.7,
    "spectrum.blackbodyrad_1.norm": 0.2,
    "spectrum.tbabs_1.nh": 0.1,
}
p0 = np.asarray(bayesian_model.dict_to_array(start)) * np.random.normal(
    loc=1, scale=0.1, size=(nwalkers, ndim)
)


@jax.jit
@jax.vmap
def zeus_log_prob(parameters: np.ndarray) -> jax.typing.ArrayLike:
    return bayesian_model.log_posterior_prob(bayesian_model.array_to_dict(parameters))


sampler = zeus.EnsembleSampler(nwalkers, ndim, zeus_log_prob, vectorize=True)

sampler.run_mcmc(p0, 10_000, progress=True);
```

--------------------------------

### joint_prior_factory usage example

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

Example of using joint_prior_factory inside a callable-prior factory. The returned lookup is used to retrieve joint samples for specified parameter paths, falling back to a default prior when the path is not listed.

```python
def prior_factory():
    joint = joint_prior_factory(
        ("spectrum.powerlaw_1.alpha", "spectrum.powerlaw_1.norm"),
        joint_dist,
    )

    def prior(path, shape):
        value = joint(path)
        return value if value is not None else default_prior(path, shape)

    return prior

```

--------------------------------

### SpectralModelBackground class implementation

Source: https://jaxspec.readthedocs.io/en/stable/references/background

Defines the SpectralModelBackground class with methods for path stripping, cache setup, and background flux computation. The user_path method removes the 'spectral_model.' prefix from prior keys, and _set_obs_cache builds a transfer-matrix cache. The __call__ method computes the background flux as the transfer matrix times the spectral model's photon flux, clipped to a minimum of 1e-6.

```python
class SpectralModelBackground(BackgroundModel):
    """Model the background as a spectral model folded through the instrument response.

    The inner ``spectral_model`` is evaluated against the observation's
    ``folded_background`` spectrum, using the same transfer matrix and energy
    grid as the source. Prior keys use dotted paths relative to the inner spectral
    model, provided in the unified prior dict. A composed model keeps its component
    names — ``SpectralModelBackground(Tbabs() * Powerlaw())`` gives
    ``"background.tbabs_1.nh"`` and ``"background.powerlaw_1.alpha"`` — whereas a bare
    component has no component segment at all, so
    ``SpectralModelBackground(Powerlaw())`` gives just ``"background.alpha"`` and
    ``"background.norm"``.

    Each per-obs instance carries its own transfer-matrix cache, populated by
    ``_set_obs_cache`` at [`ForwardModel`][jaxspec.fit._forward_model.ForwardModel]
    construction time.

    Parameters:
        spectral_model: The spectral model describing the background shape.
        sparse: Whether to use sparse transfer matrices for the background
            convolution.
    """

    def __init__(self, spectral_model, sparse: bool = False):
        self.spectral_model = spectral_model
        self.sparse = sparse
        self._tm = nnx.data(jnp.zeros((1, 1)))

    def user_path(self, nnx_path: str) -> str:
        """Strip the internal ``spectral_model.`` wrapper segment so user prior
        keys can be ``"background.powerlaw_1.alpha"`` instead of the verbose
        ``"background.spectral_model.powerlaw_1.alpha"``.

        The bare wrapper path maps to the empty string: when the background wraps a
        single component, that component *is* ``spectral_model``, so its own path is
        the background's.
        """
        prefix = "spectral_model."
        if nnx_path == prefix.rstrip("."):
            return ""
        return nnx_path[len(prefix) :] if nnx_path.startswith(prefix) else nnx_path

    def _set_obs_cache(self, observation, *, sparse: bool) -> None:
        """Build this instance's transfer-matrix cache for ``observation``.

        The bundled ``sparse`` flag from [`ForwardModel`][jaxspec.fit._forward_model.ForwardModel]
        takes precedence over the constructor flag so the background folding
        matches the source folding's storage choice.
        """
        self._tm = nnx.data(to_jax_matrix(observation.transfer_matrix.data, sparse=sparse))

    def __call__(self, observation):
        import numpy as np

        energies = np.asarray(observation.in_energies)
        flux = self.spectral_model.photon_flux(*energies)
        return jnp.clip(self._tm @ flux, min=1e-6)
```

--------------------------------

### VIFitter.fit

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

Fits the model to the data using a variational inference approach from numpyro. It runs SVI with configurable optimizer, loss, and guide, then returns a FitResult.

```APIDOC
## VIFitter.fit

### Description
Fits the model to the data using a variational inference approach from numpyro. It runs SVI with configurable optimizer, loss, and guide, then returns a FitResult.

### Method
fit

### Signature
VIFitter.fit(self, rng_key: int = 0, num_steps: int = 10_000, optimizer: numpyro.optim._NumPyroOptim | None = None, loss: numpyro.infer.elbo.ELBO | None = None, num_samples: int = 1000, guide: numpyro.infer.autoguide.AutoGuide | None = None, use_transformed_model: bool = True, plot_diagnostics: bool = False, svi_kwargs: dict | None = None) -> FitResult

### Parameters

#### rng_key
- **Type**: int
- **Default**: 0
- **Description**: The random key used to initialize the sampler.

#### num_steps
- **Type**: int
- **Default**: 10_000
- **Description**: The number of steps for VI.

#### optimizer
- **Type**: numpyro.optim._NumPyroOptim | None
- **Default**: None
- **Description**: The optimizer to use. Defaults to `Adam(step_size=0.0005)`.

#### loss
- **Type**: numpyro.infer.elbo.ELBO | None
- **Default**: None
- **Description**: The loss function to use. Defaults to `Trace_ELBO()`.

#### num_samples
- **Type**: int
- **Default**: 1000
- **Description**: The number of samples to draw.

#### guide
- **Type**: numpyro.infer.autoguide.AutoGuide | None
- **Default**: None
- **Description**: The guide to use. Defaults to an `AutoMultivariateNormal` guide.

#### use_transformed_model
- **Type**: bool
- **Default**: True
- **Description**: Whether to use the transformed model to build the InferenceData.

#### plot_diagnostics
- **Type**: bool
- **Default**: False
- **Description**: Plot the loss during VI.

#### svi_kwargs
- **Type**: dict | None
- **Default**: None
- **Description**: Additional arguments to pass to the SVI runner. See `SVI.run` for more details.

### Returns
- **FitResult**: A `FitResult` instance containing the results of the fit.
```

--------------------------------

### TiedParameter usage example

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

Defines a prior dictionary where the alpha parameter of powerlaw_2 is tied to powerlaw_1's alpha via a TiedParameter with a lambda transformation (0.5 * x).

```python
prior = {
    "spectrum.powerlaw_1.alpha": dist.Uniform(0, 5),
    "spectrum.powerlaw_2.alpha": TiedParameter(
        "spectrum.powerlaw_1.alpha", lambda x: 0.5 * x
    ),
}

```

--------------------------------

### VIFitter.fit method source code

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

The fit method performs variational inference using numpyro's SVI. It defaults to Adam optimizer with step size 0.0005, Trace_ELBO loss, and AutoMultivariateNormal guide. It supports optional diagnostics plotting and returns a FitResult.

```python
class VIFitter(BayesianModelFitter):
    def fit(
        self,
        rng_key: int = 0,
        num_steps: int = 10_000,
        optimizer: numpyro.optim._NumPyroOptim | None = None,
        loss: numpyro.infer.elbo.ELBO | None = None,
        num_samples: int = 1000,
        guide: numpyro.infer.autoguide.AutoGuide | None = None,
        use_transformed_model: bool = True,
        plot_diagnostics: bool = False,
        svi_kwargs: dict | None = None,
    ) -> FitResult:
        """
        Fit the model to the data using a variational inference approach from numpyro.

        Parameters:
            rng_key: the random key used to initialize the sampler.
            num_steps: the number of steps for VI.
            optimizer: the optimizer to use. Defaults to `Adam(step_size=0.0005)`.
            num_samples: the number of samples to draw.
            loss: the loss function to use. Defaults to `Trace_ELBO()`.
            guide: the guide to use. Defaults to an `AutoMultivariateNormal` guide.
            use_transformed_model: whether to use the transformed model to build the InferenceData.
            plot_diagnostics: plot the loss during VI.
            svi_kwargs: additional arguments to pass to the SVI runner. See [`SVI.run`][numpyro.infer.svi.SVI.run] for more details.

        Returns:
            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
        """

        svi_kwargs: dict = svi_kwargs or {}
        optimizer = numpyro.optim.Adam(step_size=0.0005) if optimizer is None else optimizer
        loss = Trace_ELBO() if loss is None else loss

        numpyro_model = (
            self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
        )

        if guide is None:
            guide = AutoMultivariateNormal(numpyro_model)

        svi = SVI(numpyro_model, guide, optimizer, loss=loss)

        keys = random.split(random.PRNGKey(rng_key), 2)
        svi_result = svi.run(keys[0], num_steps, **svi_kwargs)
        params = svi_result.params

        if plot_diagnostics:
            plt.plot(svi_result.losses)
            plt.xlabel("Steps")
            plt.ylabel("ELBO loss")
            plt.semilogy()

        predictive = Predictive(guide, params=params, num_samples=num_samples)
        posterior = predictive(keys[1])

        inference_data = self.build_inference_data(
            posterior, num_chains=1, use_transformed_model=use_transformed_model
        )

        from ...analysis.results import FitResult

        return FitResult(self, inference_data)
```

--------------------------------

### NSFitter class for Nested Sampling fitting with jaxns

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

Defines the NSFitter class that fits a model using the Nested Sampling algorithm via jaxns and the numpyro NestedSampler wrapper. It requires jaxns to be installed, otherwise raises ImportError. The fit method renames sample sites with invalid Python identifiers (e.g., containing dots) before sampling and maps them back afterwards, and stores log_Z, ESS, and other diagnostics in the resulting inference data.

```python
class NSFitter(BayesianModelFitter):
    r"""
    Fit a model using the Nested Sampling algorithm via
    [`jaxns`](https://jaxns.readthedocs.io/en/latest/) through the
    [`numpyro.contrib.nested_sampling.NestedSampler`][] wrapper.

    Because `jaxns` requires sample-site names to be valid Python identifiers,
    this fitter transparently renames sites that contain characters such as
    ``.`` (e.g. ``forward.spectrum.MOS1.powerlaw_1.alpha``) before passing
    the model to the nested sampler, and maps them back afterwards.

    !!! info
        Ensure the prior distributions cover a large enough volume for the
        algorithm to yield proper results.
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        try:
            if importlib.util.find_spec("jaxns") is None:
                raise ModuleNotFoundError
        except (ModuleNotFoundError, ValueError):
            raise ImportError("jaxns is not installed. Please install it with `pip install jaxns`.")

    def fit(
        self,
        rng_key: int = 0,
        num_samples: int = 1000,
        num_live_points: int = 1000,
        plot_diagnostics: bool = False,
        termination_kwargs: dict | None = None,
        constructor_kwargs: dict | None = None,
        verbose: bool = True,
        use_transformed_model: bool = True,
    ) -> FitResult:
        """
        Fit the model to the data using the Phantom-Powered nested sampling algorithm.

        Parameters:
            rng_key: the random key used to initialize the sampler.
            num_samples: the number of posterior samples to draw.
            num_live_points: the number of live points used by the NS algorithm.
            plot_diagnostics: whether to plot the diagnostics of the NS algorithm.
            termination_kwargs: keyword arguments forwarded to the jaxns termination
                criterion.
            constructor_kwargs: extra keyword arguments forwarded to
                ``jaxns.DefaultNestedSampler``.
            verbose: whether to print progress.
            use_transformed_model: whether to use the reparameterised model.

        Returns:
            A [`FitResult`][jaxspec.analysis.results.FitResult] instance.
        """
        from numpyro.contrib.nested_sampling import NestedSampler

        numpyro_model = (
            self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
        )

        all_site_names = _collect_site_names(numpyro_model)
        names_needing_rename = [n for n in all_site_names if n != _sanitise_name(n)]

        if names_needing_rename:
            original_to_safe, safe_to_original = _build_name_mapping(all_site_names)
            wrapped_model = _RenameSites(numpyro_model, original_to_safe)
        else:
            wrapped_model = numpyro_model
            safe_to_original = {}

        keys = random.split(random.PRNGKey(rng_key), 4)

        # A single device supports likelihood arrays with any number of bins.
        default_constructor = dict(
            verbose=verbose,
            difficult_model=True,
            max_samples=1e5,
            parameter_estimation=True,
            gradient_guided=False,
            devices=jax.devices()[:1],
            num_live_points=num_live_points,
        )
        if constructor_kwargs:
            default_constructor.update(constructor_kwargs)

        ns = NestedSampler(
            wrapped_model,
            constructor_kwargs=default_constructor,
            termination_kwargs=termination_kwargs or {},
        )

        ns.run(keys[0])

        if plot_diagnostics:
            ns.diagnostics()

        posterior = ns.get_samples(keys[1], num_samples=num_samples)

        if safe_to_original:
            posterior = {safe_to_original.get(k, k): v for k, v in posterior.items()}

        inference_data = self.build_inference_data(
            posterior, num_chains=1, use_transformed_model=use_transformed_model
        )

        results = ns._results
        if results is not None:
            inference_data.attrs["log_Z_mean"] = float(results.log_Z_mean)
            inference_data.attrs["log_Z_uncert"] = float(results.log_Z_uncert)
            inference_data.attrs["ESS"] = float(results.ESS)
            inference_data.attrs["H_mean"] = float(results.H_mean)
            inference_data.attrs["total_num_samples"] = int(results.total_num_samples)
            inference_data.attrs["total_num_likelihood_evaluations"] = int(
                results.total_num_likelihood_evaluations
            )

        from ...analysis.results import FitResult

        return FitResult(self, inference_data)

```

--------------------------------

### Serve documentation locally

Source: https://jaxspec.readthedocs.io/en/stable/contribute

Serve the documentation locally with mkdocs for development.

```bash
uv run mkdocs serve

```

--------------------------------

### Leaf callable structural prior example

Source: https://jaxspec.readthedocs.io/en/stable/examples/flexible_priors

Defines a leaf callable that returns a distribution based on the parameter path suffix, used as the `prior` argument in `MCMCFitter`. Each parameter gets its own draw per observation, increasing the number of free parameters compared to a shared dict.

```python
import numpyro.distributions as dist
import jax.numpy as jnp

def structural_prior(path, shape):
    if path.endswith(".nh"):           return dist.Uniform(0, 1)
    if path.endswith(".alpha"):        return dist.Uniform(0, 5)
    if path.endswith(".norm"):         return dist.LogUniform(1e-5, 1e-2)
    if path.endswith(".kT"):           return dist.Uniform(0, 5)
    if path.endswith(".gain.factor"):  return dist.Uniform(0.5, 1.5)
    if path.endswith(".shift.offset"): return dist.Uniform(-0.3, 0.3)
    if path.endswith(".countrate"):    return dist.Gamma(jnp.ones(shape), rate=1.0)
    raise KeyError(f"No prior defined for {path}")

fitter = MCMCFitter(spectral_model, structural_prior, observations)
```

--------------------------------

### ForwardModel.__init__

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

Constructs a ForwardModel instance by cloning the spectral model per observation, normalizing observation, background, and instrument configurations, and building per-observation caches.

```APIDOC
## ForwardModel.__init__

### Description
Initializes a ForwardModel with a spectral model, observations, optional background and instrument models, and evaluation settings.

### Method
Constructor

### Parameters
- **spectral_model** (SpectralModel | ModelComponent) - Required - The spectral model template; cloned per observation. A single bare component (e.g. `Powerlaw()`) is accepted too and is auto-wrapped via `SpectralModel.from_component`.
- **observations** (ObsConfiguration | list | dict) - Required - One or more observation configurations. Accepts a single `ObsConfiguration`, a list (auto-named `data_0`, `data_1`, ...), or a `{name: obs}` dict.
- **background_model** (BackgroundModel | dict[str, BackgroundModel | None] | None) - Optional - `None`, a singleton `BackgroundModel` (applied to every observation as a clone), or a `{obs_name: BackgroundModel | None}` dict for per-obs heterogeneous backgrounds.
- **instrument_model** (dict[str, InstrumentModel | None] | None) - Optional - `None`, or a `{obs_name: InstrumentModel | None}` dict. `None` entries (and observations missing from the dict) apply the identity fold.
- **sparsify_matrix** (bool) - Optional - Whether to store transfer matrices as sparse BCOO.
- **n_points** (int) - Optional - Number of quadrature points per energy bin for the flux integration.
- **energy_grid** (ArrayLike | None) - Optional - Optional shared 1-D array of energy bin edges (keV, strictly increasing) on which to evaluate the spectral model before redistributing onto each observation's native grid via `InstrumentModel.fold`. When `None` (default) each observation evaluates the spectrum on its own native grid.

### Response
No return value.
```

--------------------------------

### Import packages and configure JAX for external samplers

Source: https://jaxspec.readthedocs.io/en/stable/examples/external_samplers

Sets up JAX with x64 precision, host device count, and CPU platform. Imports required libraries including numpyro, arviz, zeus, jax, and jaxspec modules.

```python
%%capture 
# Hide the output of this cell
import numpyro

numpyro.enable_x64()
numpyro.set_host_device_count(4)
numpyro.set_platform("cpu")

import arviz as az
import numpyro.distributions as dist
import matplotlib.pyplot as plt
import zeus
import jax
import numpy as np
from jaxspec.data.util import load_example_obsconf
from jaxspec.fit import BayesianModel, MCMCFitter
from jaxspec.model.additive import Powerlaw, Blackbodyrad
from jaxspec.model.multiplicative import Tbabs

```

--------------------------------

### Install dill for saving inference results

Source: https://jaxspec.readthedocs.io/en/stable/faq/cookbook

Installs the dill package via pip, which is required to serialise and un-serialise jaxspec fit results.

```bash
pip install dill
```

--------------------------------

### Run zeus sampler and show progress output

Source: https://jaxspec.readthedocs.io/en/stable/examples/external_samplers

Output from running the zeus sampler for 10,000 steps with 40 walkers. The progress bar shows the sampling completed in about 5 minutes and 17 seconds.

```text
Initialising ensemble of 40 walkers...
Sampling progress : 100%|██████████| 10000/10000 [05:17<00:00, 31.47it/s]
```

--------------------------------

### Run fakeit with mock ObsConfiguration

Source: https://jaxspec.readthedocs.io/en/stable/examples/fakeits

Uses the mock ObsConfiguration with fakeit_for_multiple_parameters to generate spectra.

```python
spectra = fakeit_for_multiple_parameters(obsconf, model, parameters)

```

--------------------------------

### Clone jaxspec repository

Source: https://jaxspec.readthedocs.io/en/stable/contribute

Clone the jaxspec repository and navigate into its directory to start contributing.

```bash
git clone https://github.com/renecotyfanboy/jaxspec
cd jaxspec

```

--------------------------------

### from_pha_file classmethod: Build observation configuration from PHA file

Source: https://jaxspec.readthedocs.io/en/stable/references/data

Class method that constructs an ObsConfiguration from a PHA file, optionally using RMF, ARF, and background files. It uses data_path_finder to locate default paths if not provided, and sets energy bounds with defaults low_energy=1e-20 and high_energy=1e20. The method creates an Instrument from OGIP files and an Observation from the PHA file, then returns the configuration via from_instrument.

```python
@classmethod
def from_pha_file(
    cls,
    pha_path,
    rmf_path: str | None = None,
    arf_path: str | None = None,
    bkg_path: str | None = None,
    low_energy: float = 1e-20,
    high_energy: float = 1e20,
):
    r"""
    Build the observation configuration from a PHA file.

    Parameters:
        pha_path: The path to the PHA file.
        rmf_path: The path to the RMF file.
        arf_path: The path to the ARF file.
        bkg_path: The path to the background file.
        low_energy: The lower bound of the energy range to consider.
        high_energy: The upper bound of the energy range to consider.
    """

    from .util import data_path_finder

    arf_path_default, rmf_path_default, bkg_path_default = data_path_finder(
        pha_path,
        require_arf=arf_path is None,
        require_rmf=rmf_path is None,
        require_bkg=bkg_path is None,
    )

    arf_path = arf_path_default if arf_path is None else arf_path
    rmf_path = rmf_path_default if rmf_path is None else rmf_path
    bkg_path = bkg_path_default if bkg_path is None else bkg_path

    instrument = Instrument.from_ogip_file(
        rmf_path, arf_path=arf_path if arf_path != "" else None
    )
    observation = Observation.from_pha_file(pha_path, bkg_path=bkg_path)

    return cls.from_instrument(
        instrument, observation, low_energy=low_energy, high_energy=high_energy
    )

```

--------------------------------

### FitResult.spectrum_parameters

Source: https://jaxspec.readthedocs.io/en/stable/references/results

Property that returns a subset of `input_parameters` belonging to the spectral model (keys starting with 'spectrum.').

```APIDOC
## FitResult.spectrum_parameters

### Description
Subset of `input_parameters` belonging to the spectral model.

### Property
`spectrum_parameters` -> dict[str, ArrayLike]

### Returns
- **dict[str, ArrayLike]** - Dictionary containing only the spectral model parameters.

```

--------------------------------

### Plot the effective area

Source: https://jaxspec.readthedocs.io/en/stable/examples/load_data

Plot the effective area of the instrumental setup using matplotlib. Requires importing matplotlib.pyplot and showing the plot afterwards.

```python
import matplotlib.pyplot as plt

instrument.plot_area()
plt.show()

```

--------------------------------

### Plot the redistribution matrix

Source: https://jaxspec.readthedocs.io/en/stable/examples/load_data

Plot the redistribution matrix of the instrumental setup using matplotlib. Requires importing matplotlib.pyplot and showing the plot afterwards.

```python
import matplotlib.pyplot as plt

instrument.plot_redistribution()
plt.show()

```

--------------------------------

### APEC constructor with parameter defaults and validation

Source: https://jaxspec.readthedocs.io/en/stable/references/additive

Constructor initializes plasma parameters, optionally adding velocity and abundance parameters based on configuration. Validates 'abundances' and 'k_window' arguments, raises ValueError for unsupported values. Uses runtime or restricted table depending on 'energy_band'.

```python
def __init__(
        self,
        broadening: bool = False,
        abundances: Literal["fixed", "free", "all"] = "fixed",
        abundance_table: str = "angr",
        energy_band: tuple[float, float] | None = None,
        k_window: int | Literal["auto"] = "auto",
        line_chunks: int | None = None,
        z_max: float = 0.3,
    ):
        self.kT = nnx.Param(1.0)
        self.redshift = nnx.Param(0.0)
        self.norm = nnx.Param(1.0)

        if broadening:
            self.velocity = nnx.Param(0.0)

        if abundances == "fixed":
            self.abund = nnx.Param(1.0)
        elif abundances == "free":
            for symbol in VAPEC_SYMBOLS:
                setattr(self, symbol, nnx.Param(1.0))
        elif abundances == "all":
            for symbol in ELEMENT_SYMBOLS:
                setattr(self, symbol, nnx.Param(1.0))
        else:
            raise ValueError(f"abundances must be 'fixed', 'free' or 'all', got {abundances!r}")

        if isinstance(k_window, str):
            if k_window != "auto":
                raise ValueError(f"k_window must be a positive integer or 'auto', got {k_window!r}")
            self._k_window = None
        else:
            self._k_window = int(k_window)
            if self._k_window < 1:
                raise ValueError(f"k_window must be a positive integer or 'auto', got {k_window!r}")

        self._broadening = bool(broadening)
        self._abundances = abundances
        self._line_chunks = None if line_chunks is None else max(int(line_chunks), 1)
        self._k_window_checked = False
        # Grid of the last automatic window sizing (bins, first edge, last edge) and its result.
        self._auto_n_bins = -1
        self._auto_e_min = 0.0
        self._auto_e_max = 0.0
        self._auto_k = DEFAULT_K_WINDOW

        self._tables = _ApecTables(
            restricted_table(*energy_band, z_max=z_max)
            if energy_band is not None
            else runtime_table()
        )

        if abundance_table == "angr":
            self._ab_ratio = np.ones(len(ELEMENT_SYMBOLS))
        else:
            from ...util.abundance import abundance_table as abundance_df
```

--------------------------------

### fit

Source: https://jaxspec.readthedocs.io/en/stable/references/fitting

Fit the model to the data using a variational inference approach from numpyro. This method supports customization of optimizer, loss, guide, and sampling parameters.

```APIDOC
## `fit(rng_key=0, num_steps=10000, optimizer=None, loss=None, num_samples=1000, guide=None, use_transformed_model=True, plot_diagnostics=False, svi_kwargs=None)`

### Description
Fit the model to the data using a variational inference approach from numpyro.

### Method
Python method call on a fitter instance.

### Parameters
#### rng_key (int) - Optional
- the random key used to initialize the sampler.

#### num_steps (int) - Optional
- the number of steps for VI.

#### optimizer (_NumPyroOptim | None) - Optional
- the optimizer to use. Defaults to `Adam(step_size=0.0005)`.

#### num_samples (int) - Optional
- the number of samples to draw.

#### loss (ELBO | None) - Optional
- the loss function to use. Defaults to `Trace_ELBO()`.

#### guide (AutoGuide | None) - Optional
- the guide to use. Defaults to an `AutoMultivariateNormal` guide.

#### use_transformed_model (bool) - Optional
- whether to use the transformed model to build the InferenceData.

#### plot_diagnostics (bool) - Optional
- plot the loss during VI.

#### svi_kwargs (dict | None) - Optional
- additional arguments to pass to the SVI runner.

### Returns
- `FitResult` - A `FitResult` instance containing the results of the fit.

### Example (Python)
```python
fitter.fit(rng_key=0, num_steps=10000, num_samples=1000)
```
```

--------------------------------

### Integrate over the positive real line with integrate_positive

Source: https://jaxspec.readthedocs.io/en/stable/references/integrate

Computes the integral of a function over [0, ∞) using tanh-sinh quadrature. The example computes the gamma function value at 1/2.

```python
gamma = integrate_positive(lambda t, z: t**(z-1) * jnp.exp(-t))
print(gamma(1/2)) # 1.7716383
```

--------------------------------

### Create mock ObsConfiguration from instrument

Source: https://jaxspec.readthedocs.io/en/stable/examples/fakeits

Creates a mock ObsConfiguration from an OGIP instrument file (RMF and ARF) with a specified exposure, for use when no observation file is available.

```python
from jaxspec.data import ObsConfiguration, Instrument

instrument = Instrument.from_ogip_file(
    "instrument.rmf",
    arf_path="instrument.arf"
)

obsconf = ObsConfiguration.mock_from_instrument(
    instrument,
    exposure=1e5,
)

```

--------------------------------

### Integrate a function over an interval with integrate_interval

Source: https://jaxspec.readthedocs.io/en/stable/references/integrate

Computes the integral of a function over [a,b] using tanh-sinh quadrature. The example integrates 1/(1+x^2) from 0 to 1 to approximate pi.

```python
pi = 4*integrate_interval(lambda x: 1/(1+x**2))(0, 1)
print(pi) # 3.1415927
```
Tokens:
Raw