jaxspec 0.2.2.dev0__py3-none-any.whl → 0.3.0__py3-none-any.whl

jaxspec/fit/_fitter.py

@@ -0,0 +1,255 @@
+import warnings
+
+from abc import ABC, abstractmethod
+from typing import Literal
+
+import arviz as az
+import jax
+import matplotlib.pyplot as plt
+import numpyro
+
+from jax import random
+from jax.random import PRNGKey
+from numpyro.infer import AIES, ESS, MCMC, NUTS, SVI, Predictive, Trace_ELBO
+from numpyro.infer.autoguide import AutoMultivariateNormal
+
+from ..analysis.results import FitResult
+from ._bayesian_model import BayesianModel
+
+
+class BayesianModelFitter(BayesianModel, ABC):
+    def build_inference_data(
+        self,
+        posterior_samples,
+        num_chains: int = 1,
+        num_predictive_samples: int = 1000,
+        key: PRNGKey = PRNGKey(42),
+        use_transformed_model: bool = False,
+        filter_inference_data: bool = True,
+    ) -> az.InferenceData:
+        """
+        Build an [InferenceData][arviz.InferenceData] object from posterior samples.
+
+        Parameters:
+            posterior_samples: the samples from the posterior distribution.
+            num_chains: the number of chains used to sample the posterior.
+            num_predictive_samples: the number of samples to draw from the prior.
+            key: the random key used to initialize the sampler.
+            use_transformed_model: whether to use the transformed model to build the InferenceData.
+            filter_inference_data: whether to filter the InferenceData to keep only the relevant parameters.
+        """
+
+        numpyro_model = (
+            self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
+        )
+
+        keys = random.split(key, 3)
+
+        posterior_predictive = Predictive(numpyro_model, posterior_samples)(keys[0], observed=False)
+
+        prior = Predictive(numpyro_model, num_samples=num_predictive_samples * num_chains)(
+            keys[1], observed=False
+        )
+
+        log_likelihood = numpyro.infer.log_likelihood(numpyro_model, posterior_samples)
+
+        seeded_model = numpyro.handlers.substitute(
+            numpyro.handlers.seed(numpyro_model, keys[3]),
+            substitute_fn=numpyro.infer.init_to_sample,
+        )
+
+        observations = {
+            name: site["value"]
+            for name, site in numpyro.handlers.trace(seeded_model).get_trace().items()
+            if site["type"] == "sample" and site["is_observed"]
+        }
+
+        def reshape_first_dimension(arr):
+            new_dim = arr.shape[0] // num_chains
+            new_shape = (num_chains, new_dim) + arr.shape[1:]
+            reshaped_array = arr.reshape(new_shape)
+
+            return reshaped_array
+
+        posterior_samples = {
+            key: reshape_first_dimension(value) for key, value in posterior_samples.items()
+        }
+        prior = {key: value[None, :] for key, value in prior.items()}
+        posterior_predictive = {
+            key: reshape_first_dimension(value) for key, value in posterior_predictive.items()
+        }
+        log_likelihood = {
+            key: reshape_first_dimension(value) for key, value in log_likelihood.items()
+        }
+
+        inference_data = az.from_dict(
+            posterior_samples,
+            prior=prior,
+            posterior_predictive=posterior_predictive,
+            log_likelihood=log_likelihood,
+            observed_data=observations,
+        )
+
+        return (
+            self.filter_inference_data(inference_data) if filter_inference_data else inference_data
+        )
+
+    def filter_inference_data(
+        self,
+        inference_data: az.InferenceData,
+    ) -> az.InferenceData:
+        """
+        Filter the inference data to keep only the relevant parameters for the observations.
+
+        - Removes predictive parameters from deterministic random variables (e.g. kernel of background GP)
+        - Removes parameters build from reparametrised variables (e.g. ending with `"_base"`)
+        """
+
+        predictive_parameters = []
+
+        for key, value in self._observation_container.items():
+            if self.background_model is not None:
+                predictive_parameters.append(f"obs/~/{key}")
+                predictive_parameters.append(f"bkg/~/{key}")
+            #                predictive_parameters.append(f"ins/~/{key}")
+            else:
+                predictive_parameters.append(f"obs/~/{key}")
+        #                predictive_parameters.append(f"ins/~/{key}")
+
+        inference_data.posterior_predictive = inference_data.posterior_predictive[
+            predictive_parameters
+        ]
+
+        parameters = [
+            x
+            for x in inference_data.posterior.keys()
+            if not (x.endswith("_base") or x.startswith("_"))
+        ]
+        inference_data.posterior = inference_data.posterior[parameters]
+        inference_data.prior = inference_data.prior[parameters]
+
+        return inference_data
+
+    @abstractmethod
+    def fit(self, **kwargs) -> FitResult: ...
+
+
+class MCMCFitter(BayesianModelFitter):
+    """
+    A class to fit a model to a given set of observation using a Bayesian approach. This class uses samplers
+    from numpyro to perform the inference on the model parameters.
+    """
+
+    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 = {},
+        mcmc_kwargs: dict = {},
+    ) -> 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.
+        """
+
+        bayesian_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](bayesian_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,
+            background_model=self.background_model,
+        )
+
+
+class VIFitter(BayesianModelFitter):
+    def fit(
+        self,
+        rng_key: int = 0,
+        num_steps: int = 10_000,
+        optimizer=numpyro.optim.Adam(step_size=0.0005),
+        loss=Trace_ELBO(),
+        num_samples: int = 1000,
+        guide=None,
+        use_transformed_model: bool = True,
+        plot_diagnostics: bool = False,
+    ) -> FitResult:
+        bayesian_model = (
+            self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
+        )
+
+        if guide is None:
+            guide = AutoMultivariateNormal(bayesian_model)
+
+        svi = SVI(bayesian_model, guide, optimizer, loss=loss)
+
+        keys = random.split(random.PRNGKey(rng_key), 3)
+        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,
+            background_model=self.background_model,
+        )

jaxspec/model/abc.py

@@ -11,8 +11,7 @@ import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
 import networkx as nx
-
-from simpleeval import simple_eval
+import numpy as np
 
 from jaxspec.util.typing import PriorDictType
 
@@ -31,14 +30,14 @@ def set_parameters(params: PriorDictType, state: nnx.State) -> nnx.State:
         A spectral model with the newly set parameters.
     """
 
-    state_dict = state.to_pure_dict()  # haiku-like 2 level dictionary
+    state_dict = nnx.to_pure_dict(state)  # haiku-like 2 level dictionary
 
     for key, value in params.items():
         # Split the key to extract the module name and parameter name
         module_name, param_name = key.rsplit("_", 1)
-        state_dict["modules"][module_name][param_name] = value
+        state_dict["components"][module_name][param_name] = value
 
-    state.replace_by_pure_dict(state_dict)
+    nnx.replace_by_pure_dict(state, state_dict)
 
     return state
 
@@ -71,52 +70,16 @@ class Composable(ABC):
 
 
 class SpectralModel(nnx.Module, Composable):
-    # graph: nx.DiGraph = eqx.field(static=True)
-    # modules: dict
+    _graph: nx.DiGraph
 
     def __init__(self, graph: nx.DiGraph):
-        self.graph = graph
-        self.modules = {}
+        self._graph = graph
+        self.components = {}
+        self._energy_grid = np.geomspace(0.1, 50, 1000, dtype=np.float64)
 
-        for node, data in self.graph.nodes(data=True):
+        for node, data in self._graph.nodes(data=True):
             if "component" in data["type"]:
-                self.modules[data["name"]] = data["component"]  # (**data['kwargs'])
-
-    @classmethod
-    def from_string(cls, string: str) -> SpectralModel:
-        """
-        This constructor enable to build a model from a string. The string should be a valid python expression, with
-        the following constraints :
-
-        * The model components should be defined in the jaxspec.model.list module
-        * The model components should be separated by a * or a + (no convolution yet)
-        * The model components should be written with their parameters in parentheses
-
-        Parameters:
-            string : The string to parse
-
-        Examples:
-            An absorbed model with a powerlaw and a blackbody:
-
-            >>> model = SpectralModel.from_string("Tbabs()*(Powerlaw() + Blackbody())")
-        """
-
-        from .list import model_components
-
-        return simple_eval(string, functions=model_components)
-
-    def to_string(self) -> str:
-        """
-        This method return the string representation of the model.
-
-        Examples:
-            Build a model from a string and convert it back to a string:
-
-            >>> model = SpectralModel.from_string("Tbabs()*(Powerlaw() + Blackbody())")
-            >>> model.to_string()
-            "Tbabs()*(Powerlaw() + Blackbody())"
-        """
-        return str(self)
+                self.components[data["name"]] = data["component"]  # (**data['kwargs'])
 
     """
     def __str__(self) -> str:
@@ -153,7 +116,7 @@ class SpectralModel(nnx.Module, Composable):
         """
 
         composed_graph = compose(
-            self.graph, other.graph, operation=operation, operation_func=operation_func
+            self._graph, other._graph, operation=operation, operation_func=operation_func
         )
 
         return SpectralModel(composed_graph)
@@ -161,7 +124,7 @@ class SpectralModel(nnx.Module, Composable):
     @classmethod
     def from_component(cls, component):
         node_id = str(uuid4())
-        graph = nx.DiGraph()
+        _graph = nx.DiGraph()
 
         node_properties = {
             "type": f"{component.type}_component",
@@ -170,26 +133,26 @@ class SpectralModel(nnx.Module, Composable):
             "depth": 0,
         }
 
-        graph.add_node(node_id, **node_properties)
-        graph.add_node("out", type="out", depth=1)
-        graph.add_edge(node_id, "out")
+        _graph.add_node(node_id, **node_properties)
+        _graph.add_node("out", type="out", depth=1)
+        _graph.add_edge(node_id, "out")
 
-        return cls(graph)
+        return cls(_graph)
 
     def _find_multiplicative_components(self, node_id):
         """
         Recursively finds all the multiplicative components connected to the node with the given ID.
         """
-        node = self.graph.nodes[node_id]
+        node = self._graph.nodes[node_id]
         multiplicative_nodes = []
 
         if node.get("type") == "mul_operation":
             # Recursively find all the multiplicative components using the predecessors
-            predecessors = self.graph.pred[node_id]
+            predecessors = self._graph.pred[node_id]
             for node_id in predecessors:
-                if "multiplicative_component" == self.graph.nodes[node_id].get("type"):
+                if "multiplicative_component" == self._graph.nodes[node_id].get("type"):
                     multiplicative_nodes.append(node_id)
-                elif "mul_operation" == self.graph.nodes[node_id].get("type"):
+                elif "mul_operation" == self._graph.nodes[node_id].get("type"):
                     multiplicative_nodes.extend(self._find_multiplicative_components(node_id))
 
         return multiplicative_nodes
@@ -198,8 +161,8 @@ class SpectralModel(nnx.Module, Composable):
     def root_nodes(self) -> list[str]:
         return [
             node_id
-            for node_id, in_degree in self.graph.in_degree(self.graph.nodes)
-            if in_degree == 0 and ("additive" in self.graph.nodes[node_id].get("type"))
+            for node_id, in_degree in self._graph.in_degree(self._graph.nodes)
+            if in_degree == 0 and ("additive" in self._graph.nodes[node_id].get("type"))
         ]
 
     @property
@@ -207,8 +170,8 @@ class SpectralModel(nnx.Module, Composable):
         branches = []
 
         for root_node_id in self.root_nodes:
-            root_node_name = self.graph.nodes[root_node_id].get("name")
-            path = nx.shortest_path(self.graph, source=root_node_id, target="out")
+            root_node_name = self._graph.nodes[root_node_id].get("name")
+            path = nx.shortest_path(self._graph, source=root_node_id, target="out")
             multiplicative_components = []
 
             # Search all multiplicative components connected to this node
@@ -218,10 +181,12 @@ class SpectralModel(nnx.Module, Composable):
                     [node_id for node_id in self._find_multiplicative_components(node_id)]
                 )
 
+            multiplicative_components = set(multiplicative_components)
+
             branch = ""
 
             for multiplicative_node_id in multiplicative_components:
-                multiplicative_node_name = self.graph.nodes[multiplicative_node_id].get("name")
+                multiplicative_node_name = self._graph.nodes[multiplicative_node_id].get("name")
                 branch += f"{multiplicative_node_name}*"
 
             branch += f"{root_node_name}"
@@ -233,12 +198,12 @@ class SpectralModel(nnx.Module, Composable):
         continuum = {}
 
         ## Evaluate the expected contribution for each component
-        for node_id in nx.dag.topological_sort(self.graph):
-            node = self.graph.nodes[node_id]
+        for node_id in nx.dag.topological_sort(self._graph):
+            node = self._graph.nodes[node_id]
 
             if node["type"] == "additive_component":
                 node_name = node["name"]
-                runtime_modules = self.modules[node_name]
+                runtime_modules = self.components[node_name]
 
                 if not energy_flux:
                     continuum[node_name] = runtime_modules._photon_flux(
@@ -252,8 +217,8 @@ class SpectralModel(nnx.Module, Composable):
 
             elif node["type"] == "multiplicative_component":
                 node_name = node["name"]
-                runtime_modules = self.modules[node_name]
-                continuum[node_name] = runtime_modules.factor((e_low + e_high) / 2)
+                runtime_modules = self.components[node_name]
+                continuum[node_name] = runtime_modules._factor(e_low, e_high, n_points=n_points)
 
             else:
                 pass
@@ -264,10 +229,10 @@ class SpectralModel(nnx.Module, Composable):
         branches = {}
 
         for root_node_id in root_nodes:
-            root_node_name = self.graph.nodes[root_node_id].get("name")
+            root_node_name = self._graph.nodes[root_node_id].get("name")
             root_continuum = continuum[root_node_name]
 
-            path = nx.shortest_path(self.graph, source=root_node_id, target="out")
+            path = nx.shortest_path(self._graph, source=root_node_id, target="out")
             multiplicative_components = []
 
             # Search all multiplicative components connected to this node
@@ -279,8 +244,8 @@ class SpectralModel(nnx.Module, Composable):
 
             branch = ""
 
-            for multiplicative_node_id in multiplicative_components:
-                multiplicative_node_name = self.graph.nodes[multiplicative_node_id].get("name")
+            for multiplicative_node_id in set(multiplicative_components):
+                multiplicative_node_name = self._graph.nodes[multiplicative_node_id].get("name")
                 root_continuum *= continuum[multiplicative_node_name]
                 branch += f"{multiplicative_node_name}*"
 
@@ -302,7 +267,7 @@ class SpectralModel(nnx.Module, Composable):
         Returns:
             A string containing the mermaid representation of the model.
         """
-        return export_to_mermaid(self.graph, file)
+        return export_to_mermaid(self._graph, file)
 
     @partial(jax.jit, static_argnums=0, static_argnames=("n_points", "split_branches"))
     def photon_flux(self, params, e_low, e_high, n_points=2, split_branches=False):
@@ -327,12 +292,12 @@ class SpectralModel(nnx.Module, Composable):
             instead.
         """
 
-        graphdef, state = nnx.split(self)
-        state = set_parameters(params, state)
+        graphdef, parameters, tables = nnx.split(self, nnx.Param, ...)
+        parameters = set_parameters(params, parameters)
 
-        return nnx.call((graphdef, state)).turbo_flux(
+        return nnx.merge(graphdef, parameters, tables).turbo_flux(
             e_low, e_high, n_points=n_points, return_branches=split_branches
-        )[0]
+        )
 
     @partial(jax.jit, static_argnums=0, static_argnames="n_points")
     def energy_flux(self, params, e_low, e_high, n_points=2):
@@ -357,12 +322,12 @@ class SpectralModel(nnx.Module, Composable):
             instead.
         """
 
-        graphdef, state = nnx.split(self)
-        state = set_parameters(params, state)
+        graphdef, parameters, tables = nnx.split(self, nnx.Param, ...)
+        parameters = set_parameters(params, parameters)
 
-        return nnx.call((graphdef, state)).turbo_flux(
+        return nnx.merge(graphdef, parameters, tables).turbo_flux(
             e_low, e_high, n_points=n_points, energy_flux=True
-        )[0]
+        )
 
 
 class ModelComponent(nnx.Module, Composable, ABC):
@@ -427,6 +392,13 @@ class AdditiveComponent(ModelComponent):
 class MultiplicativeComponent(ModelComponent):
     type = "multiplicative"
 
+    def _factor(self, e_low, e_high, n_points=2):
+        energy = jnp.linspace(e_low, e_high, n_points, axis=-1)
+        factor = self.factor(energy)
+
+        return jsp.integrate.trapezoid(factor * energy, jnp.log(energy), axis=-1) / (e_high - e_low)
+        # return jnp.mean(factor, axis = -1)
+
     def factor(self, energy):
         """
         Absorption factor applied for a given energy

jaxspec/model/additive.py

@@ -166,19 +166,28 @@ class Gauss(AdditiveComponent):
         self.sigma = nnx.Param(1e-2)
         self.norm = nnx.Param(1.0)
 
+    def continuum(self, energy):
+        return (
+            self.norm
+            * jsp.stats.norm.pdf(energy, loc=jnp.asarray(self.El), scale=jnp.asarray(self.sigma))
+            / (1 - jsp.special.erf(-self.El / (self.sigma * jnp.sqrt(2))))
+        )
+
+    """
     def integrated_continuum(self, e_low, e_high):
         return self.norm * (
             jsp.stats.norm.cdf(
                 e_high,
-                loc=jnp.asarray(self.El, dtype=jnp.float64),
-                scale=jnp.asarray(self.sigma, dtype=jnp.float64),
+                loc=jnp.asarray(self.El),
+                scale=jnp.asarray(self.sigma),
             )
             - jsp.stats.norm.cdf(
                 e_low,
-                loc=jnp.asarray(self.El, dtype=jnp.float64),
-                scale=jnp.asarray(self.sigma, dtype=jnp.float64),
-            )
+                loc=jnp.asarray(self.El),
+                scale=jnp.asarray(self.sigma),
+            ) #/  (1 - jsp.special.erf(- self.El / (self.sigma * jnp.sqrt(2))))
         )
+    """
 
 
 class Cutoffpl(AdditiveComponent):

jaxspec/model/background.py

@@ -5,15 +5,15 @@ import jax.numpy as jnp
 import numpyro
 import numpyro.distributions as dist
 
+from flax import nnx
 from jax.scipy.integrate import trapezoid
 from numpyro.distributions import Poisson
 from tinygp import GaussianProcess, kernels
 
-from .._fit._build_model import build_prior, forward_model
 from .abc import SpectralModel
 
 
-class BackgroundModel(ABC):
+class BackgroundModel(ABC, nnx.Module):
     """
     Handles the background modelling in our spectra. This is handled in a separate class for now
     since backgrounds can be phenomenological models fitted directly on the folded spectrum. This is not the case for
@@ -42,7 +42,7 @@ class SubtractedBackground(BackgroundModel):
 
     def numpyro_model(self, observation, name: str = "", observed=True):
         _, observed_counts = observation.out_energies, observation.folded_background.data
-        numpyro.deterministic(f"bkg_{name}", observed_counts)
+        numpyro.deterministic(f"bkg/~/{name}", observed_counts)
 
         return observed_counts
 
@@ -61,11 +61,11 @@ class BackgroundWithError(BackgroundModel):
         _, observed_counts = obs.out_energies, obs.folded_background.data
         alpha = observed_counts + 1
         beta = 1
-        countrate = numpyro.sample(f"_bkg_{name}_countrate", dist.Gamma(alpha, rate=beta))
+        countrate = numpyro.sample(f"bkg/~/_{name}_countrate", dist.Gamma(alpha, rate=beta))
 
-        with numpyro.plate(f"bkg_{name}_plate", len(observed_counts)):
+        with numpyro.plate(f"bkg/~/{name}_plate", len(observed_counts)):
             numpyro.sample(
-                f"bkg_{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
+                f"bkg/~/{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
             )
 
         return countrate
@@ -111,27 +111,28 @@ class GaussianProcessBackground(BackgroundModel):
 
         # The parameters of the GP model
         mean = numpyro.sample(
-            f"_bkg_{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0)
+            f"bkg/~/_{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0)
         )
-        sigma = numpyro.sample(f"_bkg_{name}_sigma", dist.HalfNormal(3.0))
-        rho = numpyro.sample(f"_bkg_{name}_rho", dist.HalfNormal(10.0))
+        sigma = numpyro.sample(f"bkg/~/_{name}_sigma", dist.HalfNormal(3.0))
+        rho = numpyro.sample(f"bkg/~/_{name}_rho", dist.HalfNormal(10.0))
 
         # Set up the kernel and GP objects
         kernel = sigma**2 * self.kernel(rho)
         nodes = jnp.linspace(0, 1, self.n_nodes)
         gp = GaussianProcess(kernel, nodes, diag=1e-5 * jnp.ones_like(nodes), mean=mean)
 
-        log_rate = numpyro.sample(f"_bkg_{name}_log_rate_nodes", gp.numpyro_dist())
+        log_rate = numpyro.sample(f"bkg/~/_{name}_log_rate_nodes", gp.numpyro_dist())
 
         interp_count_rate = jnp.exp(
             jnp.interp(energy, nodes * (self.e_max - self.e_min) + self.e_min, log_rate)
         )
+
         count_rate = trapezoid(interp_count_rate, energy, axis=0)
 
         # Finally, our observation model is Poisson
-        with numpyro.plate("bkg_plate_" + name, len(observed_counts)):
+        with numpyro.plate("bkg/~/plate_" + name, len(observed_counts)):
             numpyro.sample(
-                f"bkg_{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
+                f"bkg/~/{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
             )
 
         return count_rate
@@ -144,15 +145,17 @@ class SpectralModelBackground(BackgroundModel):
         self.sparse = sparse
 
     def numpyro_model(self, observation, name: str = "", observed=True):
+        from jaxspec.fit._build_model import build_prior, forward_model
+
         params = build_prior(self.prior, prefix=f"_bkg_{name}_")
         bkg_model = jax.jit(
             lambda par: forward_model(self.spectral_model, par, observation, sparse=self.sparse)
         )
         bkg_countrate = bkg_model(params)
 
-        with numpyro.plate("bkg_plate_" + name, len(observation.folded_background)):
+        with numpyro.plate("bkg/~/plate_" + name, len(observation.folded_background)):
             numpyro.sample(
-                "bkg_" + name,
+                "bkg/~/" + name,
                 Poisson(bkg_countrate),
                 obs=observation.folded_background.data if observed else None,
             )