Fitting

BayesianModel

Bayesian spectral model that composes a deterministic forward model with numpyro priors and Poisson likelihoods.

Parameters:

Name	Type	Description	Default
model	SpectralModel | ModelComponent	The spectral model to fit (cloned per observation inside the internal :class:~jaxspec.fit._forward_model.ForwardModel). A single bare component (e.g. Powerlaw()) is accepted and auto-wrapped via :meth:~jaxspec.model.abc.SpectralModel.from_component.	required
prior	dict | Callable	Either a unified prior dict using the [obs] / [*] scoping syntax (see module docs), or a factory callable () -> ((leaf_path, shape) -> Distribution). The factory form runs inside the numpyro trace, letting it sample shared / hierarchical params before returning the leaf callable.	required
observations	ObsConfiguration | list[ObsConfiguration] | dict[str, ObsConfiguration]	One or more observation configurations.	required
background_model	BackgroundModel | dict[str, BackgroundModel | None] | None	None, a singleton BackgroundModel, or a {obs_name: BackgroundModel | None} dict.	None
instrument_model	dict[str, InstrumentModel | None] | None	None, or a {obs_name: InstrumentModel | None} dict. None entries (and observations omitted from the dict) apply the identity fold (no instrument calibration).	None
sparsify_matrix	bool	Whether to use sparse transfer matrices.	False
n_points	int	Number of quadrature points per energy bin.	2

Source code in src/jaxspec/fit/_bayesian_model.py

class BayesianModel:
    """
    Bayesian spectral model that composes a deterministic forward model with
    numpyro priors and Poisson likelihoods.

    Parameters:
        model: The spectral model to fit (cloned per observation inside the
            internal :class:`~jaxspec.fit._forward_model.ForwardModel`). A single
            bare component (e.g. ``Powerlaw()``) is accepted and auto-wrapped via
            :meth:`~jaxspec.model.abc.SpectralModel.from_component`.
        prior: Either a unified prior dict using the ``[obs]`` / ``[*]``
            scoping syntax (see module docs), or a factory callable
            ``() -> ((leaf_path, shape) -> Distribution)``. The factory form
            runs inside the numpyro trace, letting it sample shared /
            hierarchical params before returning the leaf callable.
        observations: One or more observation configurations.
        background_model: ``None``, a singleton ``BackgroundModel``, or a
            ``{obs_name: BackgroundModel | None}`` dict.
        instrument_model: ``None``, or a ``{obs_name: InstrumentModel | None}``
            dict. ``None`` entries (and observations omitted from the dict)
            apply the identity fold (no instrument calibration).
        sparsify_matrix: Whether to use sparse transfer matrices.
        n_points: Number of quadrature points per energy bin.
    """

    def __init__(
        self,
        model: SpectralModel | ModelComponent,
        prior: dict | Callable,
        observations: ObsConfiguration | list[ObsConfiguration] | dict[str, ObsConfiguration],
        background_model: BackgroundModel | dict[str, BackgroundModel | None] | None = None,
        instrument_model: dict[str, InstrumentModel | None] | None = None,
        sparsify_matrix: bool = False,
        n_points: int = 2,
        energy_grid: ArrayLike | None = None,
    ):
        self.forward_model = ForwardModel(
            model,
            observations,
            background_model=background_model,
            instrument_model=instrument_model,
            sparsify_matrix=sparsify_matrix,
            n_points=n_points,
            energy_grid=energy_grid,
        )

        self._user_prior = prior
        self._effective_prior = self._build_prior_dict()
        self._validate_prior_dict()

        # Tell the ForwardModel whether the eval-once-then-fold fast path is
        # safe for this fit. Only safe when an explicit energy_grid is set AND
        # every spectral prior entry is shared (no [obs] / [*] scopes), since
        # any per-obs spectral param means each replica produces a different
        # flux on the grid. Callable priors are conservatively treated as
        # non-shared (we can't statically inspect them).
        self.forward_model.settings["spectrum_shared"] = self._spectrum_is_shared()

    @property
    def spectral_model(self) -> SpectralModel:
        """A representative spectral model replica (PN's, or whichever obs is first).

        All per-obs replicas share the same structure; this is provided for
        callers (e.g. :class:`~jaxspec.analysis.results.FitResult`) that want
        a single ``SpectralModel`` to introspect topology or compute fluxes.
        Per-obs *parameter values* live on the bound replicas after sampling.
        """
        return next(iter(self.forward_model.spectrum.values()))

    @property
    def spectrum(self) -> dict[str, SpectralModel]:
        """Per-obs replicas of the spectral model held on the ForwardModel."""
        return self.forward_model.spectrum

    @property
    def instrument(self) -> dict[str, InstrumentModel]:
        return self.forward_model.instrument

    @property
    def background(self) -> dict[str, BackgroundModel]:
        return self.forward_model.background

    @property
    def settings(self) -> dict[str, Any]:
        return self.forward_model.settings

    # ----- Prior dict validation + default-merge -----

    def _build_prior_dict(self) -> dict | Callable:
        """Merge per-obs background and instrument defaults into the user prior
        (user wins). For callable priors, pass through unchanged."""
        if not isinstance(self._user_prior, dict):
            return self._user_prior

        prior = dict(self._user_prior)
        for modules in (self.forward_model.background, self.forward_model.instrument):
            for obs_name, module in modules.items():
                obs = self.forward_model.observations[obs_name]
                for key, value in module.default_prior(obs, obs_name).items():
                    prior.setdefault(key, value)
        return prior

    def _applicable_obs(self, prefix: str) -> set[str]:
        return set(_prefix_to_obs_names(self.forward_model).get(prefix, []))

    def _validate_prior_dict(self) -> None:
        """Validate the (effective) prior dict against the forward model.

        Callable priors are not validated structurally — the user is on the hook
        for their callable's correctness.
        """
        if not isinstance(self._effective_prior, dict):
            return
        leaves = _enumerate_leaves(self.forward_model)
        applicable = {prefix: self._applicable_obs(prefix) for prefix in _KNOWN_PREFIXES}
        for raw_key, value in self._effective_prior.items():
            self._validate_prior_entry(raw_key, value, leaves, applicable)

    def _validate_prior_entry(self, raw_key, value, leaves, applicable_by_prefix) -> None:
        # Catch flat keys like "tbabs_1_nh" before parse_prior_key — the
        # regex would accept them but downstream errors would be cryptic.
        if "." not in raw_key.split("[", 1)[0]:
            raise ValueError(
                f"Prior key {raw_key!r} has no module prefix. Expected a "
                f"dotted path like 'spectrum.<component>.<param>' (or "
                f"'instrument.<...>' / 'background.<...>'), optionally with "
                f"an [obs] or [*] suffix."
            )

        path, scope = parse_prior_key(raw_key)
        prefix = path.split(".", 1)[0]

        if prefix not in _KNOWN_PREFIXES:
            raise ValueError(
                f"Prior key {raw_key!r} starts with unknown module {prefix!r}. "
                f"The first dotted segment must be one of {_KNOWN_PREFIXES}. "
                f"Did you mean 'spectrum.{path}'?"
            )

        applicable = applicable_by_prefix[prefix]
        if not applicable:
            hint = (
                "Did you forget to pass instrument_model= to the fitter?"
                if prefix == "instrument"
                else "Did you forget to pass background_model= to the fitter?"
                if prefix == "background"
                else ""
            )
            raise ValueError(
                f"Prior key {raw_key!r} has prefix {prefix!r} but no observations "
                f"are attached to the {prefix!r} model. {hint}".rstrip()
            )
        if scope is not None and scope != "*" and scope not in applicable:
            raise ValueError(
                f"Prior key {raw_key!r} references observation {scope!r} which is "
                f"not in the {prefix!r} applicable set {sorted(applicable)}."
            )
        # Strict leaf-existence check: a key that resolves to zero leaves is a
        # typo'd parameter path — surface it at build time, not silently drop it.
        if not _resolve_targets(path, scope, leaves, applicable_by_prefix):
            raise KeyError(_unmatched_key_message(path, scope, leaves))
        if isinstance(value, dist.Distribution | TiedParameter):
            return
        try:
            jnp.asarray(value)
        except Exception as exc:
            raise TypeError(f"Invalid fixed prior value for {raw_key!r}: {value!r}") from exc

    # ----- numpyro model wiring -----

    def numpyro_model(self, observed: bool = True):
        """Sample the prior, evaluate the forward model, register likelihoods.

        Thin wrapper around :meth:`ForwardModel.evaluate`: this method owns
        only the numpyro-specific concerns (sample sites + Poisson
        likelihoods on the observed counts). The deterministic forward pass
        — spectral evaluation, instrument folding, background — lives on
        the forward model and is reused by ``fakeit`` and posterior-predictive
        checks.
        """
        inputs = self._sample_inputs()

        # Clone the forward_model per call so each evaluate sees a fresh tree
        # (the original module's Variables would otherwise accumulate tracers
        # across MCMC's repeated traces, surfacing as UnexpectedTracerError).
        # ``evaluate`` itself does NOT clone — that would break ``jax.vmap``.
        fresh_forward = nnx.clone(self.forward_model)
        predictions = fresh_forward.evaluate(inputs, missing_key_style="prior")

        fm = self.forward_model
        for obs_name, obs in fm.observations.items():
            source_flux = predictions[obs_name]["source"]
            bkg_rate = predictions[obs_name]["background"]
            bg = fm.background.get(obs_name)

            if bkg_rate is not None:
                if getattr(obs, "folded_background", None) is None:
                    raise ValueError(
                        "Trying to fit a background model but no background is "
                        "linked to this observation"
                    )
                bkg_in_obs = bkg_rate * obs.folded_backratio.data
                total = source_flux + bkg_in_obs

                if bg.is_stochastic:
                    with numpyro.plate(
                        f"observed_background_plate.{obs_name}", len(obs.folded_background)
                    ):
                        numpyro.sample(
                            f"observed_background.{obs_name}",
                            Poisson(bkg_rate),
                            obs=obs.folded_background.data if observed else None,
                        )
                else:
                    numpyro.deterministic(f"observed_background.{obs_name}", bkg_rate)
            else:
                total = source_flux

            with numpyro.plate(f"observed_plate.{obs_name}", len(obs.folded_counts)):
                numpyro.sample(
                    f"observed.{obs_name}",
                    Poisson(total),
                    obs=obs.folded_counts.data if observed else None,
                )

    # ----- Prior sampling -----

    def _sample_inputs(self) -> dict[str, Any]:
        """Sample the (effective) prior into the leaf-path inputs dict that
        :meth:`ForwardModel.evaluate` consumes.

        Creates the per-leaf numpyro sample sites along the way. Thin wrapper
        around :func:`~jaxspec.fit._prior_resolution.sample_prior` that
        provides the prefix → applicable-obs table built from
        :meth:`_applicable_obs`.
        """
        applicable = {prefix: self._applicable_obs(prefix) for prefix in _KNOWN_PREFIXES}
        return sample_prior(self.forward_model, self._effective_prior, applicable)

    def _spectrum_is_shared(self) -> bool:
        """Whether every spectral prior entry is shared across obs.

        Used at construction time to set :attr:`ForwardModel.settings`'s
        ``"spectrum_shared"`` flag, which lets :meth:`ForwardModel.evaluate`
        evaluate the spectrum **once** when a user energy grid is set
        (otherwise each obs's per-obs replica must be evaluated separately,
        e.g. when any spectral param has a ``[*]`` / ``[obs]`` scope).

        Conservative: callable priors return ``False`` since we cannot
        statically inspect them.
        """
        if not isinstance(self._effective_prior, dict):
            return False
        for raw_key in self._effective_prior:
            path, scope = parse_prior_key(raw_key)
            if path.startswith("spectrum.") and scope is not None:
                return False
        return True

    # ----- Cached properties for fitter machinery -----

    @cached_property
    def transformed_numpyro_model(self) -> Callable:
        transform_dict = {}

        relations = get_model_relations(self.numpyro_model)
        distributions = {
            parameter: getattr(numpyro.distributions, value, None)
            for parameter, value in relations["sample_dist"].items()
        }

        for parameter, distribution in distributions.items():
            if isinstance(distribution, TransformedDistribution):
                transform_dict[parameter] = TransformReparam()

        return numpyro.handlers.reparam(self.numpyro_model, config=transform_dict)

    @cached_property
    def log_likelihood_per_obs(self) -> Callable:
        """Build the log likelihood function for each bin in each observation."""

        @jax.jit
        def log_likelihood_per_obs(constrained_params):
            log_likelihood = numpyro.infer.util.log_likelihood(
                model=self.numpyro_model, posterior_samples=constrained_params
            )
            return jax.tree.map(lambda x: jnp.where(jnp.isnan(x), -jnp.inf, x), log_likelihood)

        return log_likelihood_per_obs

    @cached_property
    def log_likelihood(self) -> Callable:
        """Build the total log likelihood function."""

        @jax.jit
        def log_likelihood(constrained_params):
            log_likelihood = self.log_likelihood_per_obs(constrained_params)
            return jax.tree.reduce(operator.add, jax.tree.map(jnp.sum, log_likelihood))

        return log_likelihood

    @cached_property
    def log_posterior_prob(self) -> Callable:
        """Build the posterior probability function.

        Enables distribution validation only during this trace so that
        out-of-support parameter values produce ``-inf`` log-probabilities
        (instead of silently bogus finite values), letting external samplers
        such as AIES / ESS / MH correctly reject them.

        We don't use ``numpyro.validation_enabled()`` because that context
        manager has an upstream bug: it saves ``_VALIDATION_ENABLED`` (a
        module-global) but restores via ``enable_validation`` which writes
        to *both* that global *and* ``Distribution._validate_args`` — and
        those two can be out of sync (they are at fresh import). The
        manual save/restore below targets ``Distribution._validate_args``
        directly, which is the attribute every ``Distribution`` instance
        actually reads at construction time.
        """

        @jax.jit
        def log_posterior_prob(constrained_params):
            with numpyro.validation_enabled(True):
                log_posterior_prob, _ = log_density(
                    self.numpyro_model, (), dict(observed=True), constrained_params
                )

            return jnp.where(jnp.isnan(log_posterior_prob), -jnp.inf, log_posterior_prob)

        return log_posterior_prob

    @cached_property
    def parameter_names(self) -> list[str]:
        """List of parameter names for the model."""
        relations = get_model_relations(self.numpyro_model)
        all_sites = relations["sample_sample"].keys()
        observed_sites = relations["observed"]
        return sorted(site for site in all_sites if site not in observed_sites)

    @cached_property
    def observation_names(self) -> list[str]:
        """List of the observations."""
        relations = get_model_relations(self.numpyro_model)
        all_sites = relations["sample_sample"].keys()
        observed_sites = relations["observed"]
        return sorted(site for site in all_sites if site in observed_sites)

    def array_to_dict(self, theta):
        return {name: theta[i] for i, name in enumerate(self.parameter_names)}

    def dict_to_array(self, dict_of_params):
        theta = jnp.zeros(len(self.parameter_names))
        for index, parameter_key in enumerate(self.parameter_names):
            theta = theta.at[index].set(dict_of_params[parameter_key])
        return theta

    def prior_samples(self, key: Array = rng_key(0), num_samples: int = 100):
        """Sample from the prior distribution."""

        @jax.jit
        def prior_sample(key):
            return Predictive(
                self.numpyro_model, return_sites=self.parameter_names, num_samples=num_samples
            )(key, observed=False)

        return prior_sample(key)

    def mock_observations(self, parameters, key: Array = rng_key(0)):
        @jax.jit
        def fakeit(key, parameters):
            return Predictive(
                self.numpyro_model,
                return_sites=self.observation_names,
                posterior_samples=parameters,
            )(key, observed=False)

        return fakeit(key, parameters)

    def prior_predictive_coverage(
        self,
        key: Array = rng_key(0),
        num_samples: int = 1000,
        min_counts: int | None = None,
        grouping: int | None = None,
    ):
        """Check if the prior distribution includes the observed data."""
        if min_counts is not None and grouping is not None:
            raise ValueError("min_counts and grouping are mutually exclusive")

        key_prior, key_posterior = jax.random.split(key, 2)
        prior_params = self.prior_samples(key=key_prior, num_samples=num_samples)
        posterior_observations = self.mock_observations(prior_params, key=key_posterior)

        for key, value in self.forward_model.observations.items():
            fig, ax = plt.subplots(
                nrows=2, ncols=1, sharex=True, figsize=(5, 6), height_ratios=[3, 1]
            )

            legend_plots = []
            legend_labels = []

            observed = value.folded_counts.values
            counts = np.asarray(posterior_observations[f"observed.{key}"])
            out_energies = value.out_energies

            bin_ids = _compute_bin_ids(observed, min_counts, grouping)

            if bin_ids is not None:
                observed = rebin_counts(observed, bin_ids)
                counts = rebin_counts(counts, bin_ids)
                out_energies = _rebin_xbins(out_energies, bin_ids)

            y_observed, y_observed_low, y_observed_high = _error_bars_for_observed_data(
                observed, 1.0, "ct"
            )

            true_data_plot = _plot_poisson_data_with_error(
                ax[0],
                out_energies,
                y_observed.value,
                y_observed_low.value,
                y_observed_high.value,
                alpha=0.7,
            )

            prior_plot = _plot_binned_samples_with_error(ax[0], out_energies, counts, n_sigmas=3)

            legend_plots.append((true_data_plot,))
            legend_labels.append("Observed")
            legend_plots += prior_plot
            legend_labels.append("Prior Predictive")

            num_samples = counts.shape[0]

            less_than_obs = (counts < observed).sum(axis=0)
            equal_to_obs = (counts == observed).sum(axis=0)

            rank = (less_than_obs + 0.5 * equal_to_obs) / num_samples * 100

            ax[1].stairs(rank, edges=[*list(out_energies[0]), out_energies[1][-1]])

            ax[1].plot(
                (out_energies.min(), out_energies.max()),
                (50, 50),
                color="black",
                linestyle="--",
            )

            ax[1].set_xlabel("Energy (keV)")
            ax[0].set_ylabel("Counts")
            ax[1].set_ylabel("Rank (%)")
            ax[1].set_ylim(0, 100)
            ax[0].set_xlim(out_energies.min(), out_energies.max())
            ax[0].loglog()
            ax[0].legend(legend_plots, legend_labels)
            plt.suptitle(f"Prior Predictive coverage for {key}")
            plt.tight_layout()
            plt.show()

log_likelihood cached property

Build the total log likelihood function.

log_likelihood_per_obs cached property

Build the log likelihood function for each bin in each observation.

log_posterior_prob cached property

Build the posterior probability function.

Enables distribution validation only during this trace so that out-of-support parameter values produce -inf log-probabilities (instead of silently bogus finite values), letting external samplers such as AIES / ESS / MH correctly reject them.

We don't use numpyro.validation_enabled() because that context manager has an upstream bug: it saves _VALIDATION_ENABLED (a module-global) but restores via enable_validation which writes to both that global and Distribution._validate_args — and those two can be out of sync (they are at fresh import). The manual save/restore below targets Distribution._validate_args directly, which is the attribute every Distribution instance actually reads at construction time.

observation_names cached property

List of the observations.

parameter_names cached property

List of parameter names for the model.

spectral_model property

A representative spectral model replica (PN's, or whichever obs is first).

All per-obs replicas share the same structure; this is provided for callers (e.g. :class:~jaxspec.analysis.results.FitResult) that want a single SpectralModel to introspect topology or compute fluxes. Per-obs parameter values live on the bound replicas after sampling.

spectrum property

Per-obs replicas of the spectral model held on the ForwardModel.

numpyro_model(observed=True)

Sample the prior, evaluate the forward model, register likelihoods.

Thin wrapper around :meth:ForwardModel.evaluate: this method owns only the numpyro-specific concerns (sample sites + Poisson likelihoods on the observed counts). The deterministic forward pass — spectral evaluation, instrument folding, background — lives on the forward model and is reused by fakeit and posterior-predictive checks.

Source code in src/jaxspec/fit/_bayesian_model.py

def numpyro_model(self, observed: bool = True):
    """Sample the prior, evaluate the forward model, register likelihoods.

    Thin wrapper around :meth:`ForwardModel.evaluate`: this method owns
    only the numpyro-specific concerns (sample sites + Poisson
    likelihoods on the observed counts). The deterministic forward pass
    — spectral evaluation, instrument folding, background — lives on
    the forward model and is reused by ``fakeit`` and posterior-predictive
    checks.
    """
    inputs = self._sample_inputs()

    # Clone the forward_model per call so each evaluate sees a fresh tree
    # (the original module's Variables would otherwise accumulate tracers
    # across MCMC's repeated traces, surfacing as UnexpectedTracerError).
    # ``evaluate`` itself does NOT clone — that would break ``jax.vmap``.
    fresh_forward = nnx.clone(self.forward_model)
    predictions = fresh_forward.evaluate(inputs, missing_key_style="prior")

    fm = self.forward_model
    for obs_name, obs in fm.observations.items():
        source_flux = predictions[obs_name]["source"]
        bkg_rate = predictions[obs_name]["background"]
        bg = fm.background.get(obs_name)

        if bkg_rate is not None:
            if getattr(obs, "folded_background", None) is None:
                raise ValueError(
                    "Trying to fit a background model but no background is "
                    "linked to this observation"
                )
            bkg_in_obs = bkg_rate * obs.folded_backratio.data
            total = source_flux + bkg_in_obs

            if bg.is_stochastic:
                with numpyro.plate(
                    f"observed_background_plate.{obs_name}", len(obs.folded_background)
                ):
                    numpyro.sample(
                        f"observed_background.{obs_name}",
                        Poisson(bkg_rate),
                        obs=obs.folded_background.data if observed else None,
                    )
            else:
                numpyro.deterministic(f"observed_background.{obs_name}", bkg_rate)
        else:
            total = source_flux

        with numpyro.plate(f"observed_plate.{obs_name}", len(obs.folded_counts)):
            numpyro.sample(
                f"observed.{obs_name}",
                Poisson(total),
                obs=obs.folded_counts.data if observed else None,
            )

prior_predictive_coverage(key=rng_key(0), num_samples=1000, min_counts=None, grouping=None)

Check if the prior distribution includes the observed data.

Source code in src/jaxspec/fit/_bayesian_model.py

def prior_predictive_coverage(
    self,
    key: Array = rng_key(0),
    num_samples: int = 1000,
    min_counts: int | None = None,
    grouping: int | None = None,
):
    """Check if the prior distribution includes the observed data."""
    if min_counts is not None and grouping is not None:
        raise ValueError("min_counts and grouping are mutually exclusive")

    key_prior, key_posterior = jax.random.split(key, 2)
    prior_params = self.prior_samples(key=key_prior, num_samples=num_samples)
    posterior_observations = self.mock_observations(prior_params, key=key_posterior)

    for key, value in self.forward_model.observations.items():
        fig, ax = plt.subplots(
            nrows=2, ncols=1, sharex=True, figsize=(5, 6), height_ratios=[3, 1]
        )

        legend_plots = []
        legend_labels = []

        observed = value.folded_counts.values
        counts = np.asarray(posterior_observations[f"observed.{key}"])
        out_energies = value.out_energies

        bin_ids = _compute_bin_ids(observed, min_counts, grouping)

        if bin_ids is not None:
            observed = rebin_counts(observed, bin_ids)
            counts = rebin_counts(counts, bin_ids)
            out_energies = _rebin_xbins(out_energies, bin_ids)

        y_observed, y_observed_low, y_observed_high = _error_bars_for_observed_data(
            observed, 1.0, "ct"
        )

        true_data_plot = _plot_poisson_data_with_error(
            ax[0],
            out_energies,
            y_observed.value,
            y_observed_low.value,
            y_observed_high.value,
            alpha=0.7,
        )

        prior_plot = _plot_binned_samples_with_error(ax[0], out_energies, counts, n_sigmas=3)

        legend_plots.append((true_data_plot,))
        legend_labels.append("Observed")
        legend_plots += prior_plot
        legend_labels.append("Prior Predictive")

        num_samples = counts.shape[0]

        less_than_obs = (counts < observed).sum(axis=0)
        equal_to_obs = (counts == observed).sum(axis=0)

        rank = (less_than_obs + 0.5 * equal_to_obs) / num_samples * 100

        ax[1].stairs(rank, edges=[*list(out_energies[0]), out_energies[1][-1]])

        ax[1].plot(
            (out_energies.min(), out_energies.max()),
            (50, 50),
            color="black",
            linestyle="--",
        )

        ax[1].set_xlabel("Energy (keV)")
        ax[0].set_ylabel("Counts")
        ax[1].set_ylabel("Rank (%)")
        ax[1].set_ylim(0, 100)
        ax[0].set_xlim(out_energies.min(), out_energies.max())
        ax[0].loglog()
        ax[0].legend(legend_plots, legend_labels)
        plt.suptitle(f"Prior Predictive coverage for {key}")
        plt.tight_layout()
        plt.show()

prior_samples(key=rng_key(0), num_samples=100)

Sample from the prior distribution.

Source code in src/jaxspec/fit/_bayesian_model.py

def prior_samples(self, key: Array = rng_key(0), num_samples: int = 100):
    """Sample from the prior distribution."""

    @jax.jit
    def prior_sample(key):
        return Predictive(
            self.numpyro_model, return_sites=self.parameter_names, num_samples=num_samples
        )(key, observed=False)

    return prior_sample(key)

ForwardModel

Bases: HideUnderscoreMixin, Module

Parametric nnx tree + per-obs caches + a deterministic evaluate.

Only parameters and parametric submodules live on the nnx tree. Non-parametric state (xarray observations, response caches, settings) is held off the nnx tree on _aux — these Python objects aren't pytree-friendly and don't belong in nnx's Variable tracking.

Parameters live as nnx.Param leaves under three dict-of-modules attributes:

• spectrum: {obs_name: SpectralModel} — one cloned replica per observation, so per-obs spectral params become natural nnx leaves at spectrum.<obs>.<path>.
• instrument: {obs_name: InstrumentModel} — only observations with a non-None entry in the user's instrument_model arg.
• background: {obs_name: BackgroundModel} — singleton expanded to per-obs clones, or the per-obs dict as supplied; None entries dropped.

Parameters:

Name	Type	Description	Default
spectral_model	SpectralModel | ModelComponent	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][jaxspec.model.abc.SpectralModel.from_component].	required
observations	ObsConfiguration | list | dict	One or more observation configurations. Accepts a single ObsConfiguration, a list (auto-named data_0, data_1, ...), or a {name: obs} dict.	required
background_model	BackgroundModel | dict[str, BackgroundModel | None] | None	None, a singleton BackgroundModel (applied to every observation as a clone), or a {obs_name: BackgroundModel | None} dict for per-obs heterogeneous backgrounds.	None
instrument_model	dict[str, InstrumentModel | None] | None	None, or a {obs_name: InstrumentModel | None} dict. None entries (and observations missing from the dict) apply the identity fold.	None
sparsify_matrix	bool	Whether to store transfer matrices as sparse BCOO.	False
n_points	int	Number of quadrature points per energy bin for the flux integration.	2
energy_grid	ArrayLike | None	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.	None

Source code in src/jaxspec/fit/_forward_model.py

class ForwardModel(HideUnderscoreMixin, nnx.Module):
    """Parametric nnx tree + per-obs caches + a deterministic ``evaluate``.

    Only parameters and parametric submodules live on the nnx tree.
    Non-parametric state (xarray observations, response caches, settings) is
    held off the nnx tree on ``_aux`` — these Python objects aren't
    pytree-friendly and don't belong in nnx's Variable tracking.

    Parameters live as ``nnx.Param`` leaves under three dict-of-modules attributes:

    - ``spectrum``: ``{obs_name: SpectralModel}`` — one cloned replica per
      observation, so per-obs spectral params become natural nnx leaves at
      ``spectrum.<obs>.<path>``.
    - ``instrument``: ``{obs_name: InstrumentModel}`` — only observations
      with a non-``None`` entry in the user's ``instrument_model`` arg.
    - ``background``: ``{obs_name: BackgroundModel}`` — singleton expanded
      to per-obs clones, or the per-obs dict as supplied; ``None`` entries
      dropped.

    Parameters:
        spectral_model: 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`][jaxspec.model.abc.SpectralModel.from_component].
        observations: One or more observation configurations. Accepts a single
            [`ObsConfiguration`][jaxspec.data.obsconf.ObsConfiguration], a list
            (auto-named ``data_0``, ``data_1``, ...), or a ``{name: obs}`` dict.
        background_model: ``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: ``None``, or a ``{obs_name: InstrumentModel | None}``
            dict. ``None`` entries (and observations missing from the dict)
            apply the identity fold.
        sparsify_matrix: Whether to store transfer matrices as sparse BCOO.
        n_points: Number of quadrature points per energy bin for the flux
            integration.
        energy_grid: 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`][jaxspec.model.instrument.InstrumentModel.fold].
            When ``None`` (default) each observation evaluates the spectrum on
            its own native grid.
    """

    def __init__(
        self,
        spectral_model: SpectralModel | ModelComponent,
        observations: ObsConfiguration | list | dict,
        background_model: BackgroundModel | dict[str, BackgroundModel | None] | None = None,
        instrument_model: dict[str, InstrumentModel | None] | None = None,
        sparsify_matrix: bool = False,
        n_points: int = 2,
        energy_grid: ArrayLike | None = None,
    ):
        # Accept a bare component (e.g. ``Powerlaw()``) where a SpectralModel is
        # expected — wrap it so it gains flux_func / branch topology and its params
        # nest under the conventional ``<component>_1`` name (e.g. ``powerlaw_1``).
        if isinstance(spectral_model, ModelComponent):
            spectral_model = SpectralModel.from_component(spectral_model)

        obs_dict = _normalise_observations(observations)
        obs_names = list(obs_dict)

        # Catch typos in user-supplied per-obs dicts before normalisation
        # silently drops the misspelled entries.
        if isinstance(instrument_model, dict):
            _validate_obs_keys(instrument_model, obs_names, model_kind="instrument_model")
        if isinstance(background_model, dict):
            _validate_obs_keys(background_model, obs_names, model_kind="background_model")

        instrument_dict = _normalise_instrument(instrument_model)
        background_dict = _normalise_background(background_model, obs_names)

        self.spectrum = nnx.data({name: nnx.clone(spectral_model) for name in obs_dict})
        self.instrument = nnx.data(instrument_dict)
        self.background = nnx.data(background_dict)

        validated_grid = _validate_energy_grid(energy_grid) if energy_grid is not None else None

        # Non-parametric state lives OFF the nnx tree (plain attributes on
        # ForwardModel itself would still be tracked; stash them on the orchestrator).
        # These are exposed to the BayesianModel via the public accessors below.
        #
        # ``spectrum_shared`` is a hint set by :class:`BayesianModel` after it
        # inspects the prior dict; ``evaluate`` reads it to decide whether a
        # user-supplied ``energy_grid`` enables the eval-once-then-fold fast
        # path. Defaults to ``False`` so direct ForwardModel use is correct
        # by default (no fast path; one spectral eval per obs).
        self._aux = _ForwardModelAux(
            observations=obs_dict,
            caches={
                name: _build_obs_cache(obs, self.instrument.get(name), sparse=sparsify_matrix)
                for name, obs in obs_dict.items()
            },
            settings={
                "sparse": sparsify_matrix,
                "n_points": n_points,
                "energy_grid": validated_grid,
                "spectrum_shared": False,
            },
        )

        # Background models with caches (e.g. SpectralModelBackground transfer matrix,
        # BackgroundWithError per-bin shape) need their per-obs cache before any
        # JAX trace runs over their __call__.
        for name, bg in self.background.items():
            bg._set_obs_cache(obs_dict[name], sparse=sparsify_matrix)

    # ----- Non-parametric state accessors (read-through to self._aux) -----

    @property
    def observations(self) -> dict[str, ObsConfiguration]:
        """The name-keyed observation configurations, ``{obs_name: ObsConfiguration}``."""
        return self._aux.observations

    @property
    def settings(self) -> dict[str, Any]:
        """Evaluation settings (``sparse``, ``n_points``, ``energy_grid``, ``spectrum_shared``)."""
        return self._aux.settings

    @property
    def obs_caches(self) -> dict[str, dict[str, Any]]:
        """Per-observation JAX-typed response caches (transfer matrix and components)."""
        return self._aux.caches

    # ----- Unified evaluation entry point -----

    def evaluate(
        self,
        inputs: dict[str, ArrayLike],
        *,
        split_branches: bool = False,
        with_background: bool = True,
        missing_key_style: str = "inputs",
    ) -> dict[str, dict[str, Any]]:
        """Bind ``inputs`` and run the per-observation forward pass.

        ``inputs`` is a flat ``{leaf_path: value}`` dict keyed by nnx leaf
        paths (e.g. ``"spectrum.PN.powerlaw_1.norm"``,
        ``"instrument.MOS1.gain.factor"``, ``"background.PN.countrate"``).
        Every ``nnx.Param`` leaf of the tree must be covered; a miss raises a
        ``KeyError`` with the rich ``_missing_prior_message``.

        The method is **deterministic** — no numpyro sites are created here.
        Callers that need numpyro sampling build the inputs dict via
        ``sample_prior`` first (which is what
        [`BayesianModel.numpyro_model`][jaxspec.fit.BayesianModel.numpyro_model]
        does), then call ``evaluate``. Non-sampling callers (``fakeit``,
        posterior-predictive checks) build the inputs dict from concrete values
        and call ``evaluate`` directly, vmapping over batch dimensions.

        When ``settings["energy_grid"]`` is set, the spectral model is
        evaluated over that grid and redistributed onto each obs's native
        grid by [`InstrumentModel.fold`][jaxspec.model.instrument.InstrumentModel.fold].
        With ``settings["spectrum_shared"]`` additionally ``True`` the grid
        evaluation happens **once** and is broadcast to every obs (the
        BayesianModel sets this flag when no per-obs spectral prior is
        present). When ``energy_grid`` is ``None`` each obs evaluates the
        spectrum on its own (instrument-shifted) native grid.

        Parameters:
            inputs: Flat leaf-path → value dict covering every
                ``nnx.Param`` leaf of the tree.
            split_branches: If ``True``, the per-obs ``"source"`` entry is a
                ``{branch_name: folded_flux}`` pytree instead of a summed
                array. Used by posterior-predictive checks that overlay each
                component.
            with_background: If ``False``, skip the background evaluation and
                set each ``"background"`` entry to ``None``. Used by the
                source-only component overlay to avoid computing a rate it
                discards.
            missing_key_style: Internal selector for the "missing leaf" error
                wording. ``"inputs"`` (default) points direct callers at the
                resolved leaf path; the fitter passes ``"prior"`` to suggest the
                prior-dict key forms.

        Returns:
            ``{obs_name: {"source": folded_flux | {branch: folded_flux},
            "background": background_rate | None}}``. The background entry
            is ``None`` for obs without a background model.
        """
        from ._prior_resolution import bind_inputs

        bound = bind_inputs(self, inputs, missing_key_style=missing_key_style)
        settings = self.settings
        n_points = settings["n_points"]
        energy_grid = settings["energy_grid"]
        spectrum_shared = settings.get("spectrum_shared", False)

        # Slice the user grid once — it is identical for every obs. Fast path:
        # when the spectrum is also shared across every obs, evaluate the
        # spectral model once and reuse the result.
        shared_flux = None
        eval_energies = None
        if energy_grid is not None:
            e_low = energy_grid[:-1]
            e_high = energy_grid[1:]
            eval_energies = jnp.stack([e_low, e_high])
            if spectrum_shared:
                any_replica = next(iter(bound.spectrum.values()))
                shared_flux = any_replica.flux_func(
                    e_low, e_high, n_points=n_points, return_branches=split_branches
                )

        predictions: dict[str, dict[str, Any]] = {}
        for obs_name, obs in self.observations.items():
            cache = self.obs_caches[obs_name]
            inst = bound.instrument.get(obs_name, _IDENTITY_INSTRUMENT)

            if eval_energies is not None:
                if shared_flux is not None:
                    flux = shared_flux
                else:
                    spec = bound.spectrum[obs_name]
                    flux = spec.flux_func(
                        e_low, e_high, n_points=n_points, return_branches=split_branches
                    )
                # Shift is applied *inside* fold when eval_energies is provided
                # (see InstrumentModel.fold's contract).
                source_flux = inst.fold(flux, cache, eval_energies=eval_energies)
            else:
                spec = bound.spectrum[obs_name]
                shifted = inst.apply_shift(cache["in_energies"])
                flux = spec.flux_func(
                    shifted[0], shifted[1], n_points=n_points, return_branches=split_branches
                )
                source_flux = inst.fold(flux, cache)

            if with_background:
                bg = bound.background.get(obs_name)
                bg_rate = bg(obs) if bg is not None else None
            else:
                bg_rate = None
            predictions[obs_name] = {"source": source_flux, "background": bg_rate}

        return predictions

obs_caches property

Per-observation JAX-typed response caches (transfer matrix and components).

observations property

The name-keyed observation configurations, {obs_name: ObsConfiguration}.

settings property

Evaluation settings (sparse, n_points, energy_grid, spectrum_shared).

evaluate(inputs, *, split_branches=False, with_background=True, missing_key_style='inputs')

Bind inputs and run the per-observation forward pass.

inputs is a flat {leaf_path: value} dict keyed by nnx leaf paths (e.g. "spectrum.PN.powerlaw_1.norm", "instrument.MOS1.gain.factor", "background.PN.countrate"). Every nnx.Param leaf of the tree must be covered; a miss raises a KeyError with the rich _missing_prior_message.

The method is deterministic — no numpyro sites are created here. Callers that need numpyro sampling build the inputs dict via sample_prior first (which is what BayesianModel.numpyro_model does), then call evaluate. Non-sampling callers (fakeit, posterior-predictive checks) build the inputs dict from concrete values and call evaluate directly, vmapping over batch dimensions.

When settings["energy_grid"] is set, the spectral model is evaluated over that grid and redistributed onto each obs's native grid by InstrumentModel.fold. With settings["spectrum_shared"] additionally True the grid evaluation happens once and is broadcast to every obs (the BayesianModel sets this flag when no per-obs spectral prior is present). When energy_grid is None each obs evaluates the spectrum on its own (instrument-shifted) native grid.

Parameters:

Name	Type	Description	Default
inputs	dict[str, ArrayLike]	Flat leaf-path → value dict covering every nnx.Param leaf of the tree.	required
split_branches	bool	If True, the per-obs "source" entry is a {branch_name: folded_flux} pytree instead of a summed array. Used by posterior-predictive checks that overlay each component.	False
with_background	bool	If False, skip the background evaluation and set each "background" entry to None. Used by the source-only component overlay to avoid computing a rate it discards.	True
missing_key_style	str	Internal selector for the "missing leaf" error wording. "inputs" (default) points direct callers at the resolved leaf path; the fitter passes "prior" to suggest the prior-dict key forms.	'inputs'

Returns:

Type	Description
dict[str, dict[str, Any]]	``{obs_name: {"source": folded_flux | {branch: folded_flux},
dict[str, dict[str, Any]]	"background": background_rate | None}}``. The background entry
dict[str, dict[str, Any]]	is None for obs without a background model.

Source code in src/jaxspec/fit/_forward_model.py

def evaluate(
    self,
    inputs: dict[str, ArrayLike],
    *,
    split_branches: bool = False,
    with_background: bool = True,
    missing_key_style: str = "inputs",
) -> dict[str, dict[str, Any]]:
    """Bind ``inputs`` and run the per-observation forward pass.

    ``inputs`` is a flat ``{leaf_path: value}`` dict keyed by nnx leaf
    paths (e.g. ``"spectrum.PN.powerlaw_1.norm"``,
    ``"instrument.MOS1.gain.factor"``, ``"background.PN.countrate"``).
    Every ``nnx.Param`` leaf of the tree must be covered; a miss raises a
    ``KeyError`` with the rich ``_missing_prior_message``.

    The method is **deterministic** — no numpyro sites are created here.
    Callers that need numpyro sampling build the inputs dict via
    ``sample_prior`` first (which is what
    [`BayesianModel.numpyro_model`][jaxspec.fit.BayesianModel.numpyro_model]
    does), then call ``evaluate``. Non-sampling callers (``fakeit``,
    posterior-predictive checks) build the inputs dict from concrete values
    and call ``evaluate`` directly, vmapping over batch dimensions.

    When ``settings["energy_grid"]`` is set, the spectral model is
    evaluated over that grid and redistributed onto each obs's native
    grid by [`InstrumentModel.fold`][jaxspec.model.instrument.InstrumentModel.fold].
    With ``settings["spectrum_shared"]`` additionally ``True`` the grid
    evaluation happens **once** and is broadcast to every obs (the
    BayesianModel sets this flag when no per-obs spectral prior is
    present). When ``energy_grid`` is ``None`` each obs evaluates the
    spectrum on its own (instrument-shifted) native grid.

    Parameters:
        inputs: Flat leaf-path → value dict covering every
            ``nnx.Param`` leaf of the tree.
        split_branches: If ``True``, the per-obs ``"source"`` entry is a
            ``{branch_name: folded_flux}`` pytree instead of a summed
            array. Used by posterior-predictive checks that overlay each
            component.
        with_background: If ``False``, skip the background evaluation and
            set each ``"background"`` entry to ``None``. Used by the
            source-only component overlay to avoid computing a rate it
            discards.
        missing_key_style: Internal selector for the "missing leaf" error
            wording. ``"inputs"`` (default) points direct callers at the
            resolved leaf path; the fitter passes ``"prior"`` to suggest the
            prior-dict key forms.

    Returns:
        ``{obs_name: {"source": folded_flux | {branch: folded_flux},
        "background": background_rate | None}}``. The background entry
        is ``None`` for obs without a background model.
    """
    from ._prior_resolution import bind_inputs

    bound = bind_inputs(self, inputs, missing_key_style=missing_key_style)
    settings = self.settings
    n_points = settings["n_points"]
    energy_grid = settings["energy_grid"]
    spectrum_shared = settings.get("spectrum_shared", False)

    # Slice the user grid once — it is identical for every obs. Fast path:
    # when the spectrum is also shared across every obs, evaluate the
    # spectral model once and reuse the result.
    shared_flux = None
    eval_energies = None
    if energy_grid is not None:
        e_low = energy_grid[:-1]
        e_high = energy_grid[1:]
        eval_energies = jnp.stack([e_low, e_high])
        if spectrum_shared:
            any_replica = next(iter(bound.spectrum.values()))
            shared_flux = any_replica.flux_func(
                e_low, e_high, n_points=n_points, return_branches=split_branches
            )

    predictions: dict[str, dict[str, Any]] = {}
    for obs_name, obs in self.observations.items():
        cache = self.obs_caches[obs_name]
        inst = bound.instrument.get(obs_name, _IDENTITY_INSTRUMENT)

        if eval_energies is not None:
            if shared_flux is not None:
                flux = shared_flux
            else:
                spec = bound.spectrum[obs_name]
                flux = spec.flux_func(
                    e_low, e_high, n_points=n_points, return_branches=split_branches
                )
            # Shift is applied *inside* fold when eval_energies is provided
            # (see InstrumentModel.fold's contract).
            source_flux = inst.fold(flux, cache, eval_energies=eval_energies)
        else:
            spec = bound.spectrum[obs_name]
            shifted = inst.apply_shift(cache["in_energies"])
            flux = spec.flux_func(
                shifted[0], shifted[1], n_points=n_points, return_branches=split_branches
            )
            source_flux = inst.fold(flux, cache)

        if with_background:
            bg = bound.background.get(obs_name)
            bg_rate = bg(obs) if bg is not None else None
        else:
            bg_rate = None
        predictions[obs_name] = {"source": source_flux, "background": bg_rate}

    return predictions

Fitter classes

MCMCFitter

Bases: BayesianModelFitter

Fit a spectral model via MCMC sampling (NUTS, AIES, or ESS).

Inherits from :class:BayesianModel and accepts the same constructor arguments (spectral model, prior dict with dotted-path keys, observations, optional background/instrument models).

Source code in src/jaxspec/fit/_fitter/_mcmc.py

class MCMCFitter(BayesianModelFitter):
    """Fit a spectral model via MCMC sampling (NUTS, AIES, or ESS).

    Inherits from :class:`BayesianModel` and accepts the same constructor
    arguments (spectral model, prior dict with dotted-path keys, observations,
    optional background/instrument models).
    """

    kernel_dict = {
        "nuts": NUTS,
        "aies": AIES,
        "ess": ESS,
    }

    def fit(
        self,
        rng_key: int = 0,
        num_chains: int = len(jax.devices()),
        num_warmup: int = 1000,
        num_samples: int = 1000,
        sampler: Literal["nuts", "aies", "ess"] = "nuts",
        use_transformed_model: bool = True,
        kernel_kwargs: dict | None = None,
        mcmc_kwargs: dict | None = None,
    ) -> FitResult:
        """
        Fit the model to the data using a MCMC sampler from numpyro.

        Parameters:
            rng_key: the random key used to initialize the sampler.
            num_chains: the number of chains to run.
            num_warmup: the number of warmup steps.
            num_samples: the number of samples to draw.
            sampler: the sampler to use. Can be one of "nuts", "aies" or "ess".
            use_transformed_model: whether to use the transformed model to build the InferenceData.
            kernel_kwargs: additional arguments to pass to the kernel. See [`NUTS`][numpyro.infer.mcmc.MCMCKernel] for more details.
            mcmc_kwargs: additional arguments to pass to the MCMC sampler. See [`MCMC`][numpyro.infer.mcmc.MCMC] for more details.

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

        kernel_kwargs: dict = kernel_kwargs or {}
        mcmc_kwargs: dict = mcmc_kwargs or {}

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

        chain_kwargs = {
            "num_warmup": num_warmup,
            "num_samples": num_samples,
            "num_chains": num_chains,
        }

        kernel = self.kernel_dict[sampler](numpyro_model, **kernel_kwargs)

        mcmc_kwargs = chain_kwargs | mcmc_kwargs

        if sampler in ["aies", "ess"] and mcmc_kwargs.get("chain_method", None) != "vectorized":
            mcmc_kwargs["chain_method"] = "vectorized"
            warnings.warn("The chain_method is set to 'vectorized' for AIES and ESS samplers")

        mcmc = MCMC(kernel, **mcmc_kwargs)
        keys = random.split(random.PRNGKey(rng_key), 3)

        mcmc.run(keys[0])

        posterior = mcmc.get_samples()

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

        return FitResult(self, inference_data)

fit(rng_key=0, num_chains=len(jax.devices()), num_warmup=1000, num_samples=1000, sampler='nuts', use_transformed_model=True, kernel_kwargs=None, mcmc_kwargs=None)

Fit the model to the data using a MCMC sampler from numpyro.

Parameters:

Name	Type	Description	Default
rng_key	int	the random key used to initialize the sampler.	0
num_chains	int	the number of chains to run.	len(devices())
num_warmup	int	the number of warmup steps.	1000
num_samples	int	the number of samples to draw.	1000
sampler	Literal['nuts', 'aies', 'ess']	the sampler to use. Can be one of "nuts", "aies" or "ess".	'nuts'
use_transformed_model	bool	whether to use the transformed model to build the InferenceData.	True
kernel_kwargs	dict | None	additional arguments to pass to the kernel. See [NUTS][numpyro.infer.mcmc.MCMCKernel] for more details.	None
mcmc_kwargs	dict | None	additional arguments to pass to the MCMC sampler. See [MCMC][numpyro.infer.mcmc.MCMC] for more details.	None

Returns:

Type	Description
FitResult	A FitResult instance containing the results of the fit.

Source code in src/jaxspec/fit/_fitter/_mcmc.py

def fit(
    self,
    rng_key: int = 0,
    num_chains: int = len(jax.devices()),
    num_warmup: int = 1000,
    num_samples: int = 1000,
    sampler: Literal["nuts", "aies", "ess"] = "nuts",
    use_transformed_model: bool = True,
    kernel_kwargs: dict | None = None,
    mcmc_kwargs: dict | None = None,
) -> FitResult:
    """
    Fit the model to the data using a MCMC sampler from numpyro.

    Parameters:
        rng_key: the random key used to initialize the sampler.
        num_chains: the number of chains to run.
        num_warmup: the number of warmup steps.
        num_samples: the number of samples to draw.
        sampler: the sampler to use. Can be one of "nuts", "aies" or "ess".
        use_transformed_model: whether to use the transformed model to build the InferenceData.
        kernel_kwargs: additional arguments to pass to the kernel. See [`NUTS`][numpyro.infer.mcmc.MCMCKernel] for more details.
        mcmc_kwargs: additional arguments to pass to the MCMC sampler. See [`MCMC`][numpyro.infer.mcmc.MCMC] for more details.

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

    kernel_kwargs: dict = kernel_kwargs or {}
    mcmc_kwargs: dict = mcmc_kwargs or {}

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

    chain_kwargs = {
        "num_warmup": num_warmup,
        "num_samples": num_samples,
        "num_chains": num_chains,
    }

    kernel = self.kernel_dict[sampler](numpyro_model, **kernel_kwargs)

    mcmc_kwargs = chain_kwargs | mcmc_kwargs

    if sampler in ["aies", "ess"] and mcmc_kwargs.get("chain_method", None) != "vectorized":
        mcmc_kwargs["chain_method"] = "vectorized"
        warnings.warn("The chain_method is set to 'vectorized' for AIES and ESS samplers")

    mcmc = MCMC(kernel, **mcmc_kwargs)
    keys = random.split(random.PRNGKey(rng_key), 3)

    mcmc.run(keys[0])

    posterior = mcmc.get_samples()

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

    return FitResult(self, inference_data)

VIFitter

Bases: BayesianModelFitter

Source code in src/jaxspec/fit/_fitter/_vi.py

class VIFitter(BayesianModelFitter):
    def fit(
        self,
        rng_key: int = 0,
        num_steps: int = 10_000,
        optimizer: numpyro.optim._NumPyroOptim = numpyro.optim.Adam(step_size=0.0005),
        loss: numpyro.infer.elbo.ELBO = Trace_ELBO(),
        num_samples: int = 1000,
        guide: numpyro.infer.autoguide.AutoGuide | None = None,
        use_transformed_model: bool = True,
        plot_diagnostics: bool = False,
    ) -> 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.
            num_samples: the number of samples to draw.
            loss: the loss function to use.
            guide: the guide to use.
            use_transformed_model: whether to use the transformed model to build the InferenceData.
            plot_diagnostics: plot the loss during VI.

        Returns:
            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
        """
        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)
        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
        )

        return FitResult(self, inference_data)

fit(rng_key=0, num_steps=10000, optimizer=numpyro.optim.Adam(step_size=0.0005), loss=Trace_ELBO(), num_samples=1000, guide=None, use_transformed_model=True, plot_diagnostics=False)

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

Parameters:

Name	Type	Description	Default
rng_key	int	the random key used to initialize the sampler.	0
num_steps	int	the number of steps for VI.	10000
optimizer	_NumPyroOptim	the optimizer to use.	Adam(step_size=0.0005)
num_samples	int	the number of samples to draw.	1000
loss	ELBO	the loss function to use.	Trace_ELBO()
guide	AutoGuide | None	the guide to use.	None
use_transformed_model	bool	whether to use the transformed model to build the InferenceData.	True
plot_diagnostics	bool	plot the loss during VI.	False

Returns:

Type	Description
FitResult	A FitResult instance containing the results of the fit.

Source code in src/jaxspec/fit/_fitter/_vi.py

def fit(
    self,
    rng_key: int = 0,
    num_steps: int = 10_000,
    optimizer: numpyro.optim._NumPyroOptim = numpyro.optim.Adam(step_size=0.0005),
    loss: numpyro.infer.elbo.ELBO = Trace_ELBO(),
    num_samples: int = 1000,
    guide: numpyro.infer.autoguide.AutoGuide | None = None,
    use_transformed_model: bool = True,
    plot_diagnostics: bool = False,
) -> 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.
        num_samples: the number of samples to draw.
        loss: the loss function to use.
        guide: the guide to use.
        use_transformed_model: whether to use the transformed model to build the InferenceData.
        plot_diagnostics: plot the loss during VI.

    Returns:
        A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
    """
    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)
    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
    )

    return FitResult(self, inference_data)

Prior helpers

joint_prior_factory(components, joint_dist, *, name=None)

Sample a single multivariate site and return a per-leaf value lookup.

Use from inside a factory-callable prior to give several parameters a correlated joint draw (the dict form is strictly per-leaf, so joint sampling has to drop into the callable form). The returned lookup function maps a leaf path to the pre-sampled component value when the leaf matches one of components, and None otherwise — letting the caller chain it with structural defaults via or.

Parameters:

Name	Type	Description	Default
components	Sequence[str]	Ordered tuple of user-facing parameter keys (without [obs] suffix). The k-th component of joint_dist's sample is bound to components[k].	required
joint_dist	Distribution	A multivariate distribution whose event_shape[-1] equals len(components).	required
name	str | None	Optional numpyro site name; defaults to "+"-joined components.	None

Example::

def my_prior_factory():
    alpha_norm = joint_prior_factory(
        components=("spectrum.powerlaw_1.alpha", "spectrum.powerlaw_1.norm"),
        joint_dist=dist.MultivariateNormal(
            loc=jnp.array([2.0, 1e-4]),
            covariance_matrix=jnp.array([[0.5, 1e-5], [1e-5, 1e-8]]),
        ),
    )

    def prior(path, shape):
        return alpha_norm(path) or _structural_defaults(path, shape)
    return prior

fitter = MCMCFitter(spectral_model, observations, prior=my_prior_factory, ...)

The implementation samples joint_dist exactly once (under the supplied name), then returns the per-component scalar slice on match. The leaf callable contract in :func:~jaxspec.fit._bayesian_model._sample_leaves treats array-like returns as pre-sampled values: they are written straight to the leaf with no extra numpyro.sample site.

Source code in src/jaxspec/fit/_parameter.py

def joint_prior_factory(
    components: Sequence[str],
    joint_dist: dist.Distribution,
    *,
    name: str | None = None,
) -> Callable[[str], ArrayLike | None]:
    """Sample a single multivariate site and return a per-leaf value lookup.

    Use from inside a factory-callable prior to give several parameters a
    correlated joint draw (the dict form is strictly per-leaf, so joint
    sampling has to drop into the callable form). The returned lookup
    function maps a leaf path to the pre-sampled component value when the
    leaf matches one of ``components``, and ``None`` otherwise — letting the
    caller chain it with structural defaults via ``or``.

    Parameters:
        components: Ordered tuple of user-facing parameter keys (without
            ``[obs]`` suffix). The k-th component of ``joint_dist``'s sample
            is bound to ``components[k]``.
        joint_dist: A multivariate distribution whose ``event_shape[-1]``
            equals ``len(components)``.
        name: Optional numpyro site name; defaults to ``"+"``-joined components.

    Example::

        def my_prior_factory():
            alpha_norm = joint_prior_factory(
                components=("spectrum.powerlaw_1.alpha", "spectrum.powerlaw_1.norm"),
                joint_dist=dist.MultivariateNormal(
                    loc=jnp.array([2.0, 1e-4]),
                    covariance_matrix=jnp.array([[0.5, 1e-5], [1e-5, 1e-8]]),
                ),
            )

            def prior(path, shape):
                return alpha_norm(path) or _structural_defaults(path, shape)
            return prior

        fitter = MCMCFitter(spectral_model, observations, prior=my_prior_factory, ...)

    The implementation samples ``joint_dist`` exactly once (under the supplied
    ``name``), then returns the per-component scalar slice on match. The leaf
    callable contract in :func:`~jaxspec.fit._bayesian_model._sample_leaves`
    treats array-like returns as pre-sampled values: they are written straight
    to the leaf with no extra numpyro.sample site.
    """
    site_name = name or "+".join(components)
    sample = numpyro.sample(site_name, joint_dist)

    component_to_idx = {comp: i for i, comp in enumerate(components)}

    def lookup(leaf_path: str) -> ArrayLike | None:
        # The callable form of ``sample_prior`` calls the prior callable with
        # paths like ``"spectrum.<obs>.powerlaw_1.alpha"`` (the obs segment is
        # part of the nnx tree shape). Strip it to match user-facing paths.
        parts = leaf_path.split(".")
        if len(parts) >= 3:
            stripped = f"{parts[0]}.{'.'.join(parts[2:])}"
            if stripped in component_to_idx:
                return sample[..., component_to_idx[stripped]]
        if leaf_path in component_to_idx:
            return sample[..., component_to_idx[leaf_path]]
        return None

    return lookup

dict_prior(prior_dict)

Wrap a dict-form prior as a leaf callable.

Useful for hybrid setups: cover the common cases with dict-style entries and the special ones with custom callable logic. The returned function returns None on miss (so callers can chain with or) — unlike the framework's internal dict adapter which raises KeyError on miss.

Per the leaf-callable contract: the return value is either a :class:numpyro.distributions.Distribution (registered as a numpyro sample site by :func:~jaxspec.fit._bayesian_model._sample_leaves), a pre-sampled array (written straight to the leaf, no site), or None on miss.

Must be called from INSIDE the numpyro trace (typically from inside a factory-callable prior) since shared entries get sampled here.

Example::

def my_prior_factory():
    covered = dict_prior({
        "spectrum.powerlaw_1.alpha": dist.Uniform(0, 5),
        "spectrum.powerlaw_1.norm[*]": dist.LogUniform(1e-5, 1e-2),
    })

    def prior(path, shape):
        return covered(path, shape) or _structural_defaults(path, shape)
    return prior

The dict resolution order is the same as the framework's: explicit [obs] > [*] > shared > miss.

Source code in src/jaxspec/fit/_parameter.py

def dict_prior(
    prior_dict: dict,
) -> Callable[[str, tuple], dist.Distribution | ArrayLike | None]:
    """Wrap a dict-form prior as a leaf callable.

    Useful for hybrid setups: cover the common cases with dict-style entries
    and the special ones with custom callable logic. The returned function
    returns ``None`` on miss (so callers can chain with ``or``) — unlike the
    framework's internal dict adapter which raises ``KeyError`` on miss.

    Per the leaf-callable contract: the return value is either a
    :class:`numpyro.distributions.Distribution` (registered as a numpyro
    sample site by :func:`~jaxspec.fit._bayesian_model._sample_leaves`), a
    pre-sampled array (written straight to the leaf, no site), or ``None``
    on miss.

    Must be called from INSIDE the numpyro trace (typically from inside a
    factory-callable prior) since shared entries get sampled here.

    Example::

        def my_prior_factory():
            covered = dict_prior({
                "spectrum.powerlaw_1.alpha": dist.Uniform(0, 5),
                "spectrum.powerlaw_1.norm[*]": dist.LogUniform(1e-5, 1e-2),
            })

            def prior(path, shape):
                return covered(path, shape) or _structural_defaults(path, shape)
            return prior

    The dict resolution order is the same as the framework's: explicit
    ``[obs]`` > ``[*]`` > shared > miss.
    """
    from ._prior_resolution import _split_nnx_leaf, parse_prior_key

    # Pre-sample shared entries once.
    shared: dict[str, object] = {}
    for raw_key, value in prior_dict.items():
        path, scope = parse_prior_key(raw_key)
        if scope is None:
            if isinstance(value, dist.Distribution):
                shared[path] = numpyro.sample(path, value)
            else:
                shared[path] = jnp.asarray(value)

    def lookup(leaf_path: str, shape):
        try:
            base, obs = _split_nnx_leaf(leaf_path)
        except ValueError:
            return None
        for key in (f"{base}[{obs}]", f"{base}[*]"):
            if key in prior_dict:
                return _materialise_prior_value(prior_dict[key])
        if base in shared:
            return shared[base]
        return None

    return lookup