jaxspec 0.0.6__py3-none-any.whl → 0.0.8__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.inspect import get_model_relations
 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
+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) * sample
-            case _:
-                raise ValueError(f"Invalid prior type {type(sample)} for parameter {m}_{n} : {sample}")
+        if isinstance(sample, Distribution):
+            parameters[m][n] = jnp.ones(expand_shape) * numpyro.sample(f"{prefix}{m}_{n}", sample)
+
+        elif isinstance(sample, ArrayLike):
+            parameters[m][n] = jnp.ones(expand_shape) * sample
+
+        else:
+            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
@@ -77,10 +94,12 @@ def build_numpyro_model(
                 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())
@@ -125,14 +148,15 @@ class CountForwardModel(hk.Module):
         return jnp.clip(expected_counts, a_min=1e-6)
 
 
-class ModelFitter(ABC):
+class BayesianModel:
     """
-    Abstract class to fit a model to a given set of observation.
+    Class to fit a model to a given set of observation.
     """
 
     def __init__(
         self,
         model: SpectralModel,
+        prior_distributions: PriorDictType | Callable,
         observations: ObsConfiguration | list[ObsConfiguration] | dict[str, ObsConfiguration],
         background_model: BackgroundModel = None,
         sparsify_matrix: bool = False,
@@ -142,6 +166,8 @@ class ModelFitter(ABC):
 
         Parameters:
             model: the spectral model to fit.
+            prior_distributions: a nested dictionary containing the prior distributions for the model parameters, or a
+                callable function that returns parameter samples.
             observations: the observations to fit the model to.
             background_model: the background model to fit.
             sparsify_matrix: whether to sparsify the transfer matrix.
@@ -152,8 +178,20 @@ class ModelFitter(ABC):
         self.pars = tree_map(lambda x: jnp.float64(x), self.model.params)
         self.sparse = sparsify_matrix
 
+        if not callable(prior_distributions):
+            # Validate the entry with pydantic
+            prior = PriorDictModel(nested_dict=prior_distributions).nested_dict
+
+            def prior_distributions_func():
+                return build_prior(prior, expand_shape=(len(self.observation_container),))
+
+        else:
+            prior_distributions_func = prior_distributions
+
+        self.prior_distributions_func = prior_distributions_func
+
     @property
-    def _observation_container(self) -> dict[str, ObsConfiguration]:
+    def observation_container(self) -> dict[str, ObsConfiguration]:
         """
         The observations used in the fit as a dictionary of observations.
         """
@@ -170,33 +208,56 @@ class ModelFitter(ABC):
         else:
             raise ValueError(f"Invalid type for observations : {type(self._observations)}")
 
-    def numpyro_model(self, prior_distributions: HaikuDict[Distribution]) -> Callable:
+    @property
+    def numpyro_model(self) -> Callable:
         """
         Build the numpyro model using the observed data, the prior distributions and the spectral model.
 
-        Parameters:
-            prior_distributions: a nested dictionary containing the prior distributions for the model parameters.
-
         Returns:
+        -------
             A model function that can be used with numpyro.
         """
 
         def model(observed=True):
-            prior_params = build_prior(prior_distributions, expand_shape=(len(self._observation_container),))
+            prior_params = self.prior_distributions_func()
 
-            for i, (key, observation) in enumerate(self._observation_container.items()):
+            # 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
 
+    @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)
+
+
+class BayesianModelFitter(BayesianModel, ABC):
     @abstractmethod
-    def fit(self, prior_distributions: HaikuDict[Distribution], **kwargs) -> FitResult: ...
+    def fit(self, **kwargs) -> FitResult: ...
 
 
-class BayesianFitter(ModelFitter):
+class NUTSFitter(BayesianModelFitter):
     """
     A class to fit a model to a given set of observation using a Bayesian approach. This class uses the NUTS sampler
     from numpyro to perform the inference on the model parameters.
@@ -204,21 +265,20 @@ class BayesianFitter(ModelFitter):
 
     def fit(
         self,
-        prior_distributions: HaikuDict[Distribution],
         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:
-            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.
             num_warmup: the number of warmup steps.
@@ -232,11 +292,8 @@ class BayesianFitter(ModelFitter):
             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
+        # bayesian_model = self.numpyro_model(prior_distributions)
 
         chain_kwargs = {
             "num_warmup": num_warmup,
@@ -244,50 +301,62 @@ 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 = filter_inference_data(inference_data, self._observation_container, self.background_model)
+        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
+        )
 
         return FitResult(
-            self.model,
-            self._observation_container,
+            self,
             inference_data,
             self.model.params,
             background_model=self.background_model,
         )
 
 
-class MinimizationFitter(ModelFitter):
+class MinimizationFitter(BayesianModelFitter):
     """
     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],
         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:
-            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.
@@ -296,37 +365,205 @@ class MinimizationFitter(ModelFitter):
             A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
         """
 
-        bayesian_model = self.numpyro_model(prior_distributions)
+        bayesian_model = self.numpyro_model
+        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)
+
+            with numpyro.handlers.seed(rng_seed=local_keys[0]):
+                starting_value = self.prior_distributions_func()
+
+            # 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,
+            inference_data,
+            self.model.params,
+            background_model=self.background_model,
+        )
+
+
+class NestedSamplingFitter(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.
+    Add Citation to jaxns
+    """
+
+    def fit(
+        self,
+        rng_key: int = 0,
+        num_samples: 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.
+
+        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(
+                num_parallel_workers=1,
+                verbose=verbose,
+                difficult_model=True,
+                max_samples=1e6,
+                parameter_estimation=True,
+                num_live_points=1_000,
+            ),
+            termination_kwargs=termination_kwargs if termination_kwargs else dict(),
+        )
+
+        ns.run(keys[0])
+
+        if plot_diagnostics:
+            ns.diagnostics()
+
+        posterior_samples = ns.get_samples(keys[1], num_samples=num_samples)
         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)(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,13 +571,15 @@ 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,
-            self._observation_container,
+            self,
             inference_data,
             self.model.params,
             background_model=self.background_model,

jaxspec/model/__init__.py

@@ -1 +0,0 @@
-