jaxspec 0.0.6__py3-none-any.whl → 0.0.7__py3-none-any.whl

jaxspec/fit.py

@@ -1,67 +1,84 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Literal
+
+import arviz as az
 import haiku as hk
+import jax
 import jax.numpy as jnp
 import numpyro
-import arviz as az
-import jax
-from typing import Callable, TypeVar
-from abc import ABC, abstractmethod
+import optimistix as optx
+
 from jax import random
-from jax.tree_util import tree_map
+from jax.experimental.sparse import BCOO
 from jax.flatten_util import ravel_pytree
-from jax.experimental.sparse import BCSR
-from .analysis.results import FitResult
-from .model.abc import SpectralModel
-from .data import ObsConfiguration
-from .model.background import BackgroundModel
-from numpyro.infer import MCMC, NUTS, Predictive
-from numpyro.distributions import Distribution, TransformedDistribution
-from numpyro.distributions import Poisson
+from jax.random import PRNGKey
+from jax.tree_util import tree_map
 from jax.typing import ArrayLike
+from numpyro.contrib.nested_sampling import NestedSampler
+from numpyro.distributions import Distribution, Poisson, TransformedDistribution
+from numpyro.infer import MCMC, NUTS, Predictive
+from numpyro.infer.initialization import init_to_value
 from numpyro.infer.reparam import TransformReparam
-from numpyro.infer.util import initialize_model
-from jax.random import PRNGKey
-import jaxopt
+from numpyro.infer.util import constrain_fn
+from scipy.stats import Covariance, multivariate_normal
 
-
-T = TypeVar("T")
-
-
-class HaikuDict(dict[str, dict[str, T]]): ...
+from .analysis.results import FitResult
+from .data import ObsConfiguration
+from .model.abc import SpectralModel
+from .model.background import BackgroundModel
+from .util import catchtime, sample_prior
+from .util.typing import PriorDictModel, PriorDictType
 
 
-def build_prior(prior: HaikuDict[Distribution | ArrayLike], expand_shape: tuple = ()):
+def build_prior(prior: PriorDictType, expand_shape: tuple = (), prefix=""):
+    """
+    Transform a dictionary of prior distributions into a dictionary of parameters sampled from the prior.
+    Must be used within a numpyro model.
+    """
     parameters = dict(hk.data_structures.to_haiku_dict(prior))
 
     for i, (m, n, sample) in enumerate(hk.data_structures.traverse(prior)):
         match sample:
             case Distribution():
-                parameters[m][n] = jnp.ones(expand_shape) * numpyro.sample(f"{m}_{n}", sample)
-                # parameters[m][n] = numpyro.sample(f"{m}_{n}", sample.expand(expand_shape)) build a free parameter for each obs
-            case float() | ArrayLike():
+                parameters[m][n] = jnp.ones(expand_shape) * numpyro.sample(
+                    f"{prefix}{m}_{n}", sample
+                )
+            case ArrayLike():
                 parameters[m][n] = jnp.ones(expand_shape) * sample
             case _:
-                raise ValueError(f"Invalid prior type {type(sample)} for parameter {m}_{n} : {sample}")
+                raise ValueError(
+                    f"Invalid prior type {type(sample)} for parameter {prefix}{m}_{n} : {sample}"
+                )
 
     return parameters
 
 
-def build_numpyro_model(
+def build_numpyro_model_for_single_obs(
     obs: ObsConfiguration,
     model: SpectralModel,
     background_model: BackgroundModel,
     name: str = "",
     sparse: bool = False,
 ) -> Callable:
-    def numpro_model(prior_params, observed=True):
+    """
+    Build a numpyro model for a given observation and spectral model.
+    """
+
+    def numpyro_model(prior_params, observed=True):
         # prior_params = build_prior(prior_distributions, name=name)
-        transformed_model = hk.without_apply_rng(hk.transform(lambda par: CountForwardModel(model, obs, sparse=sparse)(par)))
+        transformed_model = hk.without_apply_rng(
+            hk.transform(lambda par: CountForwardModel(model, obs, sparse=sparse)(par))
+        )
 
         if (getattr(obs, "folded_background", None) is not None) and (background_model is not None):
             bkg_countrate = background_model.numpyro_model(
-                obs.out_energies, obs.folded_background.data, name="bkg_" + name, observed=observed
+                obs, model, name="bkg_" + name, observed=observed
             )
         elif (getattr(obs, "folded_background", None) is None) and (background_model is not None):
-            raise ValueError("Trying to fit a background model but no background is linked to this observation")
+            raise ValueError(
+                "Trying to fit a background model but no background is linked to this observation"
+            )
 
         else:
             bkg_countrate = 0.0
@@ -73,14 +90,16 @@ def build_numpyro_model(
         with numpyro.plate("obs_plate_" + name, len(obs.folded_counts)):
             numpyro.sample(
                 "obs_" + name,
-                Poisson(countrate + bkg_countrate / obs.folded_backratio.data),
+                Poisson(countrate + bkg_countrate * obs.folded_backratio.data),
                 obs=obs.folded_counts.data if observed else None,
             )
 
-    return numpro_model
+    return numpyro_model
 
 
-def filter_inference_data(inference_data, observation_container, background_model=None) -> az.InferenceData:
+def filter_inference_data(
+    inference_data, observation_container, background_model=None
+) -> az.InferenceData:
     predictive_parameters = []
 
     for key, value in observation_container.items():
@@ -109,8 +128,12 @@ class CountForwardModel(hk.Module):
         self.model = model
         self.energies = jnp.asarray(folding.in_energies)
 
-        if sparse:  # folding.transfer_matrix.data.density > 0.015 is a good criterion to consider sparsify
-            self.transfer_matrix = BCSR.from_scipy_sparse(folding.transfer_matrix.data.to_scipy_sparse().tocsr())  #
+        if (
+            sparse
+        ):  # folding.transfer_matrix.data.density > 0.015 is a good criterion to consider sparsify
+            self.transfer_matrix = BCOO.from_scipy_sparse(
+                folding.transfer_matrix.data.to_scipy_sparse().tocsr()
+            )
 
         else:
             self.transfer_matrix = jnp.asarray(folding.transfer_matrix.data.todense())
@@ -140,7 +163,8 @@ class ModelFitter(ABC):
         """
         Initialize the fitter.
 
-        Parameters:
+        Parameters
+        ----------
             model: the spectral model to fit.
             observations: the observations to fit the model to.
             background_model: the background model to fit.
@@ -170,30 +194,61 @@ class ModelFitter(ABC):
         else:
             raise ValueError(f"Invalid type for observations : {type(self._observations)}")
 
-    def numpyro_model(self, prior_distributions: HaikuDict[Distribution]) -> Callable:
+    def numpyro_model(self, prior_distributions: PriorDictType | Callable) -> Callable:
         """
         Build the numpyro model using the observed data, the prior distributions and the spectral model.
 
-        Parameters:
+        Parameters
+        ----------
             prior_distributions: a nested dictionary containing the prior distributions for the model parameters.
 
         Returns:
+        -------
             A model function that can be used with numpyro.
         """
 
+        if not callable(prior_distributions):
+            # Validate the entry with pydantic
+            prior_distributions = PriorDictModel(nested_dict=prior_distributions).nested_dict
+
+            def prior_distributions_func():
+                return build_prior(
+                    prior_distributions, expand_shape=(len(self._observation_container),)
+                )
+
+        else:
+            prior_distributions_func = prior_distributions
+
         def model(observed=True):
-            prior_params = build_prior(prior_distributions, expand_shape=(len(self._observation_container),))
+            prior_params = prior_distributions_func()
 
+            # Iterate over all the observations in our container and build a single numpyro model for each observation
             for i, (key, observation) in enumerate(self._observation_container.items()):
+                # We expect that prior_params contains an array of parameters for each observation
+                # They can be identical or different for each observation
                 params = tree_map(lambda x: x[i], prior_params)
 
-                obs_model = build_numpyro_model(observation, self.model, self.background_model, name=key, sparse=self.sparse)
+                obs_model = build_numpyro_model_for_single_obs(
+                    observation, self.model, self.background_model, name=key, sparse=self.sparse
+                )
+
                 obs_model(params, observed=observed)
 
         return model
 
+    def transformed_numpyro_model(self, prior_distributions: PriorDictType) -> Callable:
+        transform_dict = {}
+
+        for m, n, val in hk.data_structures.traverse(prior_distributions):
+            if isinstance(val, TransformedDistribution):
+                transform_dict[f"{m}_{n}"] = TransformReparam()
+
+        return numpyro.handlers.reparam(
+            self.numpyro_model(prior_distributions), config=transform_dict
+        )
+
     @abstractmethod
-    def fit(self, prior_distributions: HaikuDict[Distribution], **kwargs) -> FitResult: ...
+    def fit(self, prior_distributions: PriorDictType, **kwargs) -> FitResult: ...
 
 
 class BayesianFitter(ModelFitter):
@@ -204,20 +259,22 @@ class BayesianFitter(ModelFitter):
 
     def fit(
         self,
-        prior_distributions: HaikuDict[Distribution],
+        prior_distributions: PriorDictType,
         rng_key: int = 0,
-        num_chains: int = 4,
+        num_chains: int = len(jax.devices()),
         num_warmup: int = 1000,
         num_samples: int = 1000,
         max_tree_depth: int = 10,
         target_accept_prob: float = 0.8,
         dense_mass: bool = False,
+        kernel_kwargs: dict = {},
         mcmc_kwargs: dict = {},
     ) -> FitResult:
         """
         Fit the model to the data using NUTS sampler from numpyro.
 
-        Parameters:
+        Parameters
+        ----------
             prior_distributions: a nested dictionary containing the prior distributions for the model parameters.
             rng_key: the random key used to initialize the sampler.
             num_chains: the number of chains to run.
@@ -229,14 +286,12 @@ class BayesianFitter(ModelFitter):
             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.
         """
 
-        transform_dict = {}
-
-        for m, n, val in hk.data_structures.traverse(prior_distributions):
-            if isinstance(val, TransformedDistribution):
-                transform_dict[f"{m}_{n}"] = TransformReparam()
+        bayesian_model = self.transformed_numpyro_model(prior_distributions)
+        # bayesian_model = self.numpyro_model(prior_distributions)
 
         chain_kwargs = {
             "num_warmup": num_warmup,
@@ -244,20 +299,29 @@ class BayesianFitter(ModelFitter):
             "num_chains": num_chains,
         }
 
-        bayesian_model = numpyro.handlers.reparam(self.numpyro_model(prior_distributions), config=transform_dict)
-
-        kernel = NUTS(bayesian_model, max_tree_depth=max_tree_depth, target_accept_prob=target_accept_prob, dense_mass=dense_mass)
+        kernel = NUTS(
+            bayesian_model,
+            max_tree_depth=max_tree_depth,
+            target_accept_prob=target_accept_prob,
+            dense_mass=dense_mass,
+            **kernel_kwargs,
+        )
 
         mcmc = MCMC(kernel, **(chain_kwargs | mcmc_kwargs))
-
         keys = random.split(random.PRNGKey(rng_key), 3)
 
         mcmc.run(keys[0])
-        posterior_predictive = Predictive(bayesian_model, mcmc.get_samples())(keys[1], observed=False)
+        posterior_predictive = Predictive(bayesian_model, mcmc.get_samples())(
+            keys[1], observed=False
+        )
         prior = Predictive(bayesian_model, num_samples=num_samples)(keys[2], observed=False)
-        inference_data = az.from_numpyro(mcmc, prior=prior, posterior_predictive=posterior_predictive)
+        inference_data = az.from_numpyro(
+            mcmc, prior=prior, posterior_predictive=posterior_predictive
+        )
 
-        inference_data = filter_inference_data(inference_data, self._observation_container, self.background_model)
+        inference_data = filter_inference_data(
+            inference_data, self._observation_container, self.background_model
+        )
 
         return FitResult(
             self.model,
@@ -272,61 +336,242 @@ class MinimizationFitter(ModelFitter):
     """
     A class to fit a model to a given set of observation using a minimization algorithm. This class uses the L-BFGS
     algorithm from jaxopt to perform the minimization on the model parameters. The uncertainties are computed using the
-    Hessian of the log-likelihood, assuming that it is a multivariate Gaussian in the unbounded space defined by
+    Hessian of the log-log_likelihood, assuming that it is a multivariate Gaussian in the unbounded space defined by
     numpyro.
     """
 
     def fit(
         self,
-        prior_distributions: HaikuDict[Distribution],
+        prior_distributions: PriorDictType,
         rng_key: int = 0,
-        num_iter_max: int = 10_000,
+        num_iter_max: int = 100_000,
         num_samples: int = 1_000,
+        solver: Literal["bfgs", "levenberg_marquardt"] = "bfgs",
+        init_params=None,
+        refine_first_guess=True,
     ) -> FitResult:
         """
         Fit the model to the data using L-BFGS algorithm.
 
-        Parameters:
+        Parameters
+        ----------
             prior_distributions: a nested dictionary containing the prior distributions for the model parameters.
             rng_key: the random key used to initialize the sampler.
             num_iter_max: the maximum number of iteration in the minimization algorithm.
             num_samples: the number of sample to draw from the best-fit covariance.
 
         Returns:
+        -------
             A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
         """
 
         bayesian_model = self.numpyro_model(prior_distributions)
+        keys = jax.random.split(PRNGKey(rng_key), 4)
 
-        param_info, potential_fn, postprocess_fn, *_ = initialize_model(
-            PRNGKey(0),
-            bayesian_model,
-            model_args=tuple(),
-            dynamic_args=True,  # <- this is important!
-        )
+        if init_params is not None:
+            # We initialize the parameters by randomly sampling from the prior
+            local_keys = jax.random.split(keys[0], 2)
+            starting_value = sample_prior(
+                prior_distributions, key=local_keys[0], flat_parameters=True
+            )
+
+            # We update the starting value with the provided init_params
+            for m, n, val in hk.data_structures.traverse(init_params):
+                if f"{m}_{n}" in starting_value.keys():
+                    starting_value[f"{m}_{n}"] = val
+
+            init_params, _ = numpyro.infer.util.find_valid_initial_params(
+                local_keys[1], bayesian_model, init_strategy=init_to_value(values=starting_value)
+            )
+
+        else:
+            init_params, _ = numpyro.infer.util.find_valid_initial_params(keys[0], bayesian_model)
+
+        init_params = init_params[0]
 
-        # get negative log-density from the potential function
         @jax.jit
-        def nll_fn(position):
-            func = potential_fn()
-            return func(position)
+        def nll(unconstrained_params, _):
+            constrained_params = constrain_fn(
+                bayesian_model, tuple(), dict(observed=True), unconstrained_params
+            )
 
-        solver = jaxopt.LBFGS(fun=nll_fn, maxiter=10_000)
-        params, state = solver.run(param_info.z)
-        keys = random.split(random.PRNGKey(rng_key), 3)
+            log_likelihood = numpyro.infer.util.log_likelihood(
+                model=bayesian_model, posterior_samples=constrained_params
+            )
+
+            # We solve a least square problem, this function ensure that the total residual is indeed the nll
+            return jax.tree.map(lambda x: jnp.sqrt(-x), log_likelihood)
+
+        """
+        if refine_first_guess:
+            with catchtime("Refine_first"):
+                solution = optx.least_squares(
+                    nll,
+                    optx.BestSoFarMinimiser(optx.OptaxMinimiser(optax.adam(1e-4), 1e-6, 1e-6)),
+                    init_params,
+                    max_steps=1000,
+                    throw=False
+                )
+            init_params = solution.value
+        """
 
+        if solver == "bfgs":
+            solver = optx.BestSoFarMinimiser(optx.BFGS(1e-6, 1e-6))
+        elif solver == "levenberg_marquardt":
+            solver = optx.BestSoFarLeastSquares(optx.LevenbergMarquardt(1e-6, 1e-6))
+        else:
+            raise NotImplementedError(f"{solver} is not implemented")
+
+        with catchtime("Minimization"):
+            solution = optx.least_squares(
+                nll,
+                solver,
+                init_params,
+                max_steps=num_iter_max,
+            )
+
+        params = solution.value
         value_flat, unflatten_fun = ravel_pytree(params)
-        covariance = jnp.linalg.inv(jax.hessian(lambda p: nll_fn(unflatten_fun(p)))(value_flat))
 
-        samples_flat = jax.random.multivariate_normal(keys[0], value_flat, covariance, shape=(num_samples,))
-        samples = jax.vmap(unflatten_fun)(samples_flat.block_until_ready())
-        posterior_samples = postprocess_fn()(samples)
+        with catchtime("Compute error"):
+            precision = jax.hessian(
+                lambda p: jnp.sum(ravel_pytree(nll(unflatten_fun(p), None))[0] ** 2)
+            )(value_flat)
 
-        posterior_predictive = Predictive(bayesian_model, posterior_samples)(keys[1], observed=False)
-        prior = Predictive(bayesian_model, num_samples=num_samples)(keys[2], observed=False)
+            cov = Covariance.from_precision(precision)
+
+            samples_flat = multivariate_normal.rvs(mean=value_flat, cov=cov, size=num_samples)
+
+        samples = jax.vmap(unflatten_fun)(samples_flat)
+        posterior_samples = jax.jit(
+            jax.vmap(lambda p: constrain_fn(bayesian_model, tuple(), dict(observed=True), p))
+        )(samples)
+
+        with catchtime("Posterior"):
+            posterior_predictive = Predictive(bayesian_model, posterior_samples)(
+                keys[2], observed=False
+            )
+            prior = Predictive(bayesian_model, num_samples=num_samples)(keys[3], observed=False)
+            log_likelihood = numpyro.infer.log_likelihood(bayesian_model, posterior_samples)
+
+        def sanitize_chain(chain):
+            """
+            reshape the samples so that it is arviz compliant with an extra starting dimension
+            """
+            return tree_map(lambda x: x[None, ...], chain)
+
+        # We export the observed values to the inference_data
+        seeded_model = numpyro.handlers.substitute(
+            numpyro.handlers.seed(bayesian_model, jax.random.PRNGKey(0)),
+            substitute_fn=numpyro.infer.init_to_sample,
+        )
+        trace = numpyro.handlers.trace(seeded_model).get_trace()
+        observations = {
+            name: site["value"]
+            for name, site in trace.items()
+            if site["type"] == "sample" and site["is_observed"]
+        }
+
+        with catchtime("InferenceData wrapping"):
+            inference_data = az.from_dict(
+                sanitize_chain(posterior_samples),
+                prior=sanitize_chain(prior),
+                posterior_predictive=sanitize_chain(posterior_predictive),
+                log_likelihood=sanitize_chain(log_likelihood),
+                observed_data=observations,
+            )
+
+        inference_data = filter_inference_data(
+            inference_data, self._observation_container, self.background_model
+        )
+
+        return FitResult(
+            self.model,
+            self._observation_container,
+            inference_data,
+            self.model.params,
+            background_model=self.background_model,
+        )
+
+
+class NestedSamplingFitter(ModelFitter):
+    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.
+    Add Citation to jaxns
+    """
+
+    def fit(
+        self,
+        prior_distributions: PriorDictType,
+        rng_key: int = 0,
+        num_parallel_workers: int = len(jax.devices()),
+        num_samples: int = 1000,
+        plot_diagnostics=False,
+        verbose=True,
+    ) -> FitResult:
+        """
+        Fit the model to the data using the Phantom-Powered nested sampling algorithm.
+
+        Parameters:
+            prior_distributions: a nested dictionary containing the prior distributions for the model parameters.
+            rng_key: the random key used to initialize the sampler.
+            num_samples: the number of samples to draw.
+
+        Returns:
+            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
+        """
+
+        bayesian_model = self.transformed_numpyro_model(prior_distributions)
+        keys = random.split(random.PRNGKey(rng_key), 4)
+
+        ns = NestedSampler(
+            bayesian_model,
+            constructor_kwargs=dict(
+                num_parallel_workers=num_parallel_workers,
+                verbose=verbose,
+                difficult_model=True,
+                # max_samples=1e6,
+                # num_live_points=10_000,
+                # init_efficiency_threshold=0.5,
+                parameter_estimation=True,
+            ),
+            termination_kwargs=dict(dlogZ=1e-2),
+        )
+
+        ns.run(keys[0])
+
+        self.ns = ns
+
+        if plot_diagnostics:
+            ns.diagnostics()
+
+        posterior_samples = ns.get_samples(keys[1], num_samples=num_samples * num_parallel_workers)
         log_likelihood = numpyro.infer.log_likelihood(bayesian_model, posterior_samples)
+        posterior_predictive = Predictive(bayesian_model, posterior_samples)(
+            keys[2], observed=False
+        )
+
+        prior = Predictive(bayesian_model, num_samples=num_samples * num_parallel_workers)(
+            keys[3], observed=False
+        )
+
+        seeded_model = numpyro.handlers.substitute(
+            numpyro.handlers.seed(bayesian_model, jax.random.PRNGKey(0)),
+            substitute_fn=numpyro.infer.init_to_sample,
+        )
+        trace = numpyro.handlers.trace(seeded_model).get_trace()
+        observations = {
+            name: site["value"]
+            for name, site in trace.items()
+            if site["type"] == "sample" and site["is_observed"]
+        }
 
         def sanitize_chain(chain):
+            """
+            reshape the samples so that it is arviz compliant with an extra starting dimension
+            """
             return tree_map(lambda x: x[None, ...], chain)
 
         inference_data = az.from_dict(
@@ -334,9 +579,12 @@ class MinimizationFitter(ModelFitter):
             prior=sanitize_chain(prior),
             posterior_predictive=sanitize_chain(posterior_predictive),
             log_likelihood=sanitize_chain(log_likelihood),
+            observed_data=observations,
         )
 
-        inference_data = filter_inference_data(inference_data, self._observation_container, self.background_model)
+        inference_data = filter_inference_data(
+            inference_data, self._observation_container, self.background_model
+        )
 
         return FitResult(
             self.model,

jaxspec/model/__init__.py

@@ -1 +0,0 @@
-