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

jaxspec/{fit.py → fit/_bayesian_model.py}

@@ -1,54 +1,53 @@
 import operator
-import warnings
 
-from abc import ABC, abstractmethod
 from collections.abc import Callable
 from functools import cached_property
-from typing import Literal
+from typing import Any
 
-import arviz as az
 import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
 import numpyro
 
-from jax import random
-from jax.experimental import mesh_utils
+from flax import nnx
 from jax.random import PRNGKey
-from jax.sharding import PositionalSharding
-from numpyro.contrib.nested_sampling import NestedSampler
+from jax.sharding import NamedSharding, PartitionSpec
 from numpyro.distributions import Poisson, TransformedDistribution
-from numpyro.infer import AIES, ESS, MCMC, NUTS, Predictive
+from numpyro.infer import Predictive
 from numpyro.infer.inspect import get_model_relations
 from numpyro.infer.reparam import TransformReparam
 from numpyro.infer.util import log_density
 
-from ._fit._build_model import build_prior, forward_model
-from .analysis._plot import (
+from ..analysis._plot import (
     _error_bars_for_observed_data,
     _plot_binned_samples_with_error,
     _plot_poisson_data_with_error,
 )
-from .analysis.results import FitResult
-from .data import ObsConfiguration
-from .model.abc import SpectralModel
-from .model.background import BackgroundModel
-from .util.typing import PriorDictType
+from ..data import ObsConfiguration
+from ..model.abc import SpectralModel
+from ..model.background import BackgroundModel
+from ..model.instrument import InstrumentModel
+from ..util.typing import PriorDictType
+from ._build_model import build_prior, forward_model
 
 
-class BayesianModel:
+class BayesianModel(nnx.Module):
     """
     Base class for a Bayesian model. This class contains the necessary methods to build a model, sample from the prior
     and compute the log-likelihood and posterior probability.
     """
 
+    settings: dict[str, Any]
+
     def __init__(
         self,
         model: SpectralModel,
         prior_distributions: PriorDictType | Callable,
         observations: ObsConfiguration | list[ObsConfiguration] | dict[str, ObsConfiguration],
         background_model: BackgroundModel = None,
+        instrument_model: InstrumentModel = None,
         sparsify_matrix: bool = False,
+        n_points: int = 2,
     ):
         """
         Build a Bayesian model for a given spectral model and observations.
@@ -59,74 +58,36 @@ class BayesianModel:
                 callable function that returns parameter samples.
             observations: the observations to fit the model to.
             background_model: the background model to fit.
+            instrument_model: the instrument model to fit.
             sparsify_matrix: whether to sparsify the transfer matrix.
         """
-        self.model = model
+
+        self.spectral_model = model
         self._observations = observations
         self.background_model = background_model
-        self.sparse = sparsify_matrix
+        self.instrument_model = instrument_model
+        self.settings = {"sparse": sparsify_matrix}
 
         if not callable(prior_distributions):
-            # Validate the entry with pydantic
-            # prior = PriorDictModel.from_dict(prior_distributions).
 
             def prior_distributions_func():
                 return build_prior(
-                    prior_distributions, expand_shape=(len(self.observation_container),)
+                    prior_distributions,
+                    expand_shape=(len(self._observation_container),),
+                    prefix="mod/~/",
                 )
 
         else:
             prior_distributions_func = prior_distributions
 
         self.prior_distributions_func = prior_distributions_func
-        self.init_params = self.prior_samples()
-
-        # Check the priors are suited for the observations
-        split_parameters = [
-            (param, shape[-1])
-            for param, shape in jax.tree.map(lambda x: x.shape, self.init_params).items()
-            if (len(shape) > 1)
-            and not param.startswith("_")
-            and not param.startswith("bkg")  # hardcoded for subtracted background
-        ]
-
-        for parameter, proposed_number_of_obs in split_parameters:
-            if proposed_number_of_obs != len(self.observation_container):
-                raise ValueError(
-                    f"Invalid splitting in the prior distribution. "
-                    f"Expected {len(self.observation_container)} but got {proposed_number_of_obs} for {parameter}"
-                )
-
-    @cached_property
-    def observation_container(self) -> dict[str, ObsConfiguration]:
-        """
-        The observations used in the fit as a dictionary of observations.
-        """
-
-        if isinstance(self._observations, dict):
-            return self._observations
-
-        elif isinstance(self._observations, list):
-            return {f"data_{i}": obs for i, obs in enumerate(self._observations)}
-
-        elif isinstance(self._observations, ObsConfiguration):
-            return {"data": self._observations}
-
-        else:
-            raise ValueError(f"Invalid type for observations : {type(self._observations)}")
-
-    @cached_property
-    def numpyro_model(self) -> Callable:
-        """
-        Build the numpyro model using the observed data, the prior distributions and the spectral model.
-        """
 
         def numpyro_model(observed=True):
             # Instantiate and register the parameters of the spectral model and the background
             prior_params = self.prior_distributions_func()
 
             # Iterate over all the observations in our container and build a single numpyro model for each observation
-            for i, (name, observation) in enumerate(self.observation_container.items()):
+            for i, (name, observation) in enumerate(self._observation_container.items()):
                 # Check that we can indeed fit a background
                 if (getattr(observation, "folded_background", None) is not None) and (
                     self.background_model is not None
@@ -150,23 +111,70 @@ class BayesianModel:
                 # They can be identical or different for each observation
                 params = jax.tree.map(lambda x: x[i], prior_params)
 
+                if self.instrument_model is not None:
+                    gain, shift = self.instrument_model.get_gain_and_shift_model(name)
+                else:
+                    gain, shift = None, None
+
                 # Forward model the observation and get the associated countrate
                 obs_model = jax.jit(
-                    lambda par: forward_model(self.model, par, observation, sparse=self.sparse)
+                    lambda par: forward_model(
+                        self.spectral_model,
+                        par,
+                        observation,
+                        sparse=self.settings.get("sparse", False),
+                        gain=gain,
+                        shift=shift,
+                        n_points=n_points,
+                    )
                 )
+
                 obs_countrate = obs_model(params)
 
                 # Register the observation as an observed site
-                with numpyro.plate("obs_plate_" + name, len(observation.folded_counts)):
+                with numpyro.plate("obs_plate/~/" + name, len(observation.folded_counts)):
                     numpyro.sample(
-                        "obs_" + name,
-                        Poisson(
-                            obs_countrate + bkg_countrate
-                        ),  # / observation.folded_backratio.data
+                        "obs/~/" + name,
+                        Poisson(obs_countrate + bkg_countrate / observation.folded_backratio.data),
                         obs=observation.folded_counts.data if observed else None,
                     )
 
-        return numpyro_model
+        self.numpyro_model = numpyro_model
+        self._init_params = self.prior_samples()
+        # Check the priors are suited for the observations
+        split_parameters = [
+            (param, shape[-1])
+            for param, shape in jax.tree.map(lambda x: x.shape, self._init_params).items()
+            if (len(shape) > 1)
+            and not param.startswith("_")
+            and not param.startswith("bkg")  # hardcoded for subtracted background
+            and not param.startswith("ins")
+        ]
+
+        for parameter, proposed_number_of_obs in split_parameters:
+            if proposed_number_of_obs != len(self._observation_container):
+                raise ValueError(
+                    f"Invalid splitting in the prior distribution. "
+                    f"Expected {len(self._observation_container)} but got {proposed_number_of_obs} for {parameter}"
+                )
+
+    @cached_property
+    def _observation_container(self) -> dict[str, ObsConfiguration]:
+        """
+        The observations used in the fit as a dictionary of observations.
+        """
+
+        if isinstance(self._observations, dict):
+            return self._observations
+
+        elif isinstance(self._observations, list):
+            return {f"data_{i}": obs for i, obs in enumerate(self._observations)}
+
+        elif isinstance(self._observations, ObsConfiguration):
+            return {"data": self._observations}
+
+        else:
+            raise ValueError(f"Invalid type for observations : {type(self._observations)}")
 
     @cached_property
     def transformed_numpyro_model(self) -> Callable:
@@ -315,7 +323,8 @@ class BayesianModel:
         """
         key_prior, key_posterior = jax.random.split(key, 2)
         n_devices = len(jax.local_devices())
-        sharding = PositionalSharding(mesh_utils.create_device_mesh((n_devices,)))
+        mesh = jax.make_mesh((n_devices,), ("batch",))
+        sharding = NamedSharding(mesh, PartitionSpec("batch"))
 
         # Sample from prior and correct if the number of samples is not a multiple of the number of devices
         if num_samples % n_devices != 0:
@@ -327,7 +336,7 @@ class BayesianModel:
         sharded_parameters = jax.device_put(prior_params, sharding)
         posterior_observations = self.mock_observations(sharded_parameters, key=key_posterior)
 
-        for key, value in self.observation_container.items():
+        for key, value in self._observation_container.items():
             fig, ax = plt.subplots(
                 nrows=2, ncols=1, sharex=True, figsize=(5, 6), height_ratios=[3, 1]
             )
@@ -349,7 +358,7 @@ class BayesianModel:
             )
 
             prior_plot = _plot_binned_samples_with_error(
-                ax[0], value.out_energies, posterior_observations["obs_" + key], n_sigmas=3
+                ax[0], value.out_energies, posterior_observations["obs/~/" + key], n_sigmas=3
             )
 
             legend_plots.append((true_data_plot,))
@@ -358,7 +367,7 @@ class BayesianModel:
             legend_labels.append("Prior Predictive")
 
             # rank = np.vstack((posterior_observations["obs_" + key], value.folded_counts.values)).argsort(axis=0)[-1] / (num_samples) * 100
-            counts = posterior_observations["obs_" + key]
+            counts = posterior_observations["obs/~/" + key]
             observed = value.folded_counts.values
 
             num_samples = counts.shape[0]
@@ -387,264 +396,3 @@ class BayesianModel:
             plt.suptitle(f"Prior Predictive coverage for {key}")
             plt.tight_layout()
             plt.show()
-
-
-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}")
-            else:
-                predictive_parameters.append(f"obs_{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=True
-        )
-
-        return FitResult(
-            self,
-            inference_data,
-            background_model=self.background_model,
-        )
-
-
-class NSFitter(BayesianModelFitter):
-    r"""
-    A class to fit a model to a given set of observation using the Nested Sampling algorithm. This class uses the
-    [`DefaultNestedSampler`][jaxns.DefaultNestedSampler] from [`jaxns`](https://jaxns.readthedocs.io/en/latest/) which
-    implements the [Phantom-Powered Nested Sampling](https://arxiv.org/abs/2312.11330) algorithm.
-
-    !!! info
-        Ensure large prior volume is covered by the prior distributions to ensure the algorithm yield proper results.
-
-    """
-
-    def fit(
-        self,
-        rng_key: int = 0,
-        num_samples: int = 1000,
-        num_live_points: int = 1000,
-        plot_diagnostics=False,
-        termination_kwargs: dict | None = None,
-        verbose=True,
-    ) -> FitResult:
-        """
-        Fit the model to the data using the Phantom-Powered nested sampling algorithm.
-
-        Parameters:
-            rng_key: the random key used to initialize the sampler.
-            num_samples: the number of samples to draw.
-            num_live_points: the number of live points to use at the start of the NS algorithm.
-            plot_diagnostics: whether to plot the diagnostics of the NS algorithm.
-            termination_kwargs: additional arguments to pass to the termination criterion of the NS algorithm.
-            verbose: whether to print the progress of the NS algorithm.
-
-        Returns:
-            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
-        """
-
-        bayesian_model = self.transformed_numpyro_model
-        keys = random.split(random.PRNGKey(rng_key), 4)
-
-        ns = NestedSampler(
-            bayesian_model,
-            constructor_kwargs=dict(
-                verbose=verbose,
-                difficult_model=True,
-                max_samples=1e5,
-                parameter_estimation=True,
-                gradient_guided=True,
-                devices=jax.devices(),
-                # init_efficiency_threshold=0.01,
-                num_live_points=num_live_points,
-            ),
-            termination_kwargs=termination_kwargs if termination_kwargs else dict(),
-        )
-
-        ns.run(keys[0])
-
-        if plot_diagnostics:
-            ns.diagnostics()
-
-        posterior = ns.get_samples(keys[1], num_samples=num_samples)
-        inference_data = self.build_inference_data(
-            posterior, num_chains=1, use_transformed_model=True
-        )
-
-        return FitResult(
-            self,
-            inference_data,
-            background_model=self.background_model,
-        )

jaxspec/{_fit → fit}/_build_model.py

@@ -1,5 +1,7 @@
+from collections.abc import Callable
 from typing import TYPE_CHECKING
 
+import jax
 import jax.numpy as jnp
 import numpy as np
 import numpyro
@@ -9,9 +11,15 @@ from jax.typing import ArrayLike
 from numpyro.distributions import Distribution
 
 if TYPE_CHECKING:
-    from ..data import ObsConfiguration
-    from ..model.abc import SpectralModel
-    from ..util.typing import PriorDictType
+    from jaxspec.data import ObsConfiguration
+    from jaxspec.model.abc import SpectralModel
+    from jaxspec.util.typing import PriorDictType
+
+
+class TiedParameter:
+    def __init__(self, tied_to, func):
+        self.tied_to = tied_to
+        self.func = func
 
 
 def forward_model(
@@ -19,6 +27,10 @@ def forward_model(
     parameters,
     obs_configuration: "ObsConfiguration",
     sparse=False,
+    gain: Callable | None = None,
+    shift: Callable | None = None,
+    split_branches: bool = False,
+    n_points: int | None = 2,
 ):
     energies = np.asarray(obs_configuration.in_energies)
 
@@ -31,10 +43,24 @@ def forward_model(
     else:
         transfer_matrix = np.asarray(obs_configuration.transfer_matrix.data.todense())
 
-    expected_counts = transfer_matrix @ model.photon_flux(parameters, *energies)
+    energies = shift(energies) if shift is not None else energies
+    energies = jnp.clip(energies, min=1e-6)  # Ensure shifted energies remain positive
+    factor = gain(energies) if gain is not None else 1.0
+    factor = jnp.clip(factor, min=0.0)  # Ensure the gain is positive to avoid NaNs
 
-    # The result is clipped at 1e-6 to avoid 0 round-off and diverging likelihoods
-    return jnp.clip(expected_counts, a_min=1e-6)
+    if not split_branches:
+        expected_counts = transfer_matrix @ (
+            model.photon_flux(parameters, *energies, n_points=n_points) * factor
+        )
+        return jnp.clip(expected_counts, min=1e-6)  # Ensure the expected counts are positive
+
+    else:
+        model_flux = model.photon_flux(
+            parameters, *energies, split_branches=True, n_points=n_points
+        )
+        return jax.tree.map(
+            lambda f: jnp.clip(transfer_matrix @ (f * factor), min=1e-6), model_flux
+        )
 
 
 def build_prior(prior: "PriorDictType", expand_shape: tuple = (), prefix=""):
@@ -43,15 +69,20 @@ def build_prior(prior: "PriorDictType", expand_shape: tuple = (), prefix=""):
     Must be used within a numpyro model.
     """
     parameters = {}
+    params_to_tie = {}
 
     for key, value in prior.items():
         # Split the key to extract the module name and parameter name
         module_name, param_name = key.rsplit("_", 1)
+
         if isinstance(value, Distribution):
             parameters[key] = jnp.ones(expand_shape) * numpyro.sample(
                 f"{prefix}{module_name}_{param_name}", value
             )
 
+        elif isinstance(value, TiedParameter):
+            params_to_tie[key] = value
+
         elif isinstance(value, ArrayLike):
             parameters[key] = jnp.ones(expand_shape) * value
 
@@ -60,4 +91,9 @@ def build_prior(prior: "PriorDictType", expand_shape: tuple = (), prefix=""):
                 f"Invalid prior type {type(value)} for parameter {prefix}{module_name}_{param_name} : {value}"
             )
 
+    for key, value in params_to_tie.items():
+        func_to_apply = value.func
+        tied_to = value.tied_to
+        parameters[key] = func_to_apply(parameters[tied_to])
+
     return parameters