jaxspec 0.1.0__tar.gz → 0.1.1__tar.gz

{jaxspec-0.1.0 → jaxspec-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxspec
-Version: 0.1.0
+Version: 0.1.1
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
 Home-page: https://github.com/renecotyfanboy/jaxspec
 License: MIT
@@ -24,7 +24,7 @@ Requires-Dist: jaxns (>=2.5.1,<3.0.0)
 Requires-Dist: jaxopt (>=0.8.1,<0.9.0)
 Requires-Dist: matplotlib (>=3.8.0,<4.0.0)
 Requires-Dist: mendeleev (>=0.15,<0.18)
-Requires-Dist: mkdocstrings (>=0.24,<0.26)
+Requires-Dist: mkdocstrings (>=0.24,<0.27)
 Requires-Dist: networkx (>=3.1,<4.0)
 Requires-Dist: numpy (<2.0.0)
 Requires-Dist: numpyro (>=0.15.2,<0.16.0)

{jaxspec-0.1.0 → jaxspec-0.1.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jaxspec"
-version = "0.1.0"
+version = "0.1.1"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = ["sdupourque <sdupourque@irap.omp.eu>"]
 license = "MIT"
@@ -28,7 +28,7 @@ gpjax = "^0.8.0"
 jaxopt = "^0.8.1"
 tinygp = "^0.3.0"
 seaborn = "^0.13.1"
-mkdocstrings = ">=0.24,<0.26"
+mkdocstrings = ">=0.24,<0.27"
 sparse = "^0.15.1"
 optimistix = "^0.0.7"
 scipy = "<1.15"
@@ -43,8 +43,8 @@ watermark = "^2.4.3"
 [tool.poetry.group.docs.dependencies]
 mkdocs = "^1.5.3"
 mkdocs-material = "^9.4.6"
-mkdocstrings = {extras = ["python"], version = ">=0.24,<0.26"}
-mkdocs-jupyter = "^0.24.6"
+mkdocstrings = {extras = ["python"], version = ">=0.24,<0.27"}
+mkdocs-jupyter = ">=0.24.6,<0.26.0"
 
 
 [tool.poetry.group.test.dependencies]
@@ -59,7 +59,7 @@ testbook = "^0.4.2"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = "^3.5.0"
-ruff = ">=0.2.1,<0.6.0"
+ruff = ">=0.2.1,<0.7.0"
 jupyterlab = "^4.0.7"
 notebook = "^7.0.6"
 ipywidgets = "^8.1.1"
@@ -118,4 +118,4 @@ requires = ["poetry-core"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.poetry.scripts]
-jaxspec-debug-info = "jaxspec.scripts.debug:debug_info"
+jaxspec-debug-info = "jaxspec.scripts.debug:debug_info"

{jaxspec-0.1.0 → jaxspec-0.1.1}/src/jaxspec/fit.py

@@ -1,4 +1,5 @@
 import operator
+import warnings
 
 from abc import ABC, abstractmethod
 from collections.abc import Callable
@@ -20,7 +21,7 @@ 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 import AIES, ESS, 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
@@ -181,7 +182,7 @@ class BayesianModel:
 
         if not callable(prior_distributions):
             # Validate the entry with pydantic
-            prior = PriorDictModel(nested_dict=prior_distributions).nested_dict
+            prior = PriorDictModel.from_dict(prior_distributions).nested_dict
 
             def prior_distributions_func():
                 return build_prior(prior, expand_shape=(len(self.observation_container),))
@@ -293,6 +294,9 @@ class BayesianModel:
         that can be fetched with the [`parameter_names`][jaxspec.fit.BayesianModel.parameter_names].
         """
 
+        # This is required as numpyro.infer.util.log_densities does not check parameter validity by itself
+        numpyro.enable_validation()
+
         @jax.jit
         def log_posterior_prob(constrained_params):
             log_posterior_prob, _ = log_density(
@@ -350,6 +354,69 @@ class BayesianModel:
 
 
 class BayesianModelFitter(BayesianModel, ABC):
+    def build_inference_data(
+        self,
+        posterior_samples,
+        num_chains: int = 1,
+        num_predictive_samples: int = 1000,
+        key: PRNGKey = PRNGKey(0),
+        use_transformed_model: bool = False,
+    ) -> az.InferenceData:
+        """
+        Build an InferenceData object from the posterior samples.
+        """
+
+        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()
+        }
+
+        return az.from_dict(
+            posterior_samples,
+            prior=prior,
+            posterior_predictive=posterior_predictive,
+            log_likelihood=log_likelihood,
+            observed_data=observations,
+        )
+
     @abstractmethod
     def fit(self, **kwargs) -> FitResult: ...
 
@@ -411,18 +478,89 @@ class NUTSFitter(BayesianModelFitter):
 
         mcmc.run(keys[0])
 
-        posterior_predictive = Predictive(bayesian_model, mcmc.get_samples())(
-            keys[1], observed=False
-        )
+        posterior = mcmc.get_samples()
 
-        prior = Predictive(bayesian_model, num_samples=num_samples)(keys[2], observed=False)
+        inference_data = filter_inference_data(
+            self.build_inference_data(posterior, num_chains=num_chains),
+            self.observation_container,
+            self.background_model,
+        )
 
-        inference_data = az.from_numpyro(
-            mcmc, prior=prior, posterior_predictive=posterior_predictive
+        return FitResult(
+            self,
+            inference_data,
+            self.model.params,
+            background_model=self.background_model,
         )
 
+
+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",
+        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.
+            max_tree_depth: the recursion depth of NUTS sampler.
+            target_accept_prob: the target acceptance probability for the NUTS sampler.
+            dense_mass: whether to use a dense mass for the NUTS sampler.
+            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
+        # bayesian_model = self.numpyro_model(prior_distributions)
+
+        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 = filter_inference_data(
-            inference_data, self.observation_container, self.background_model
+            self.build_inference_data(posterior, num_chains=num_chains),
+            self.observation_container,
+            self.background_model,
         )
 
         return FitResult(

{jaxspec-0.1.0 → jaxspec-0.1.1}/src/jaxspec/model/abc.py

@@ -1,12 +1,15 @@
 from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from uuid import uuid4
+
 import haiku as hk
 import jax
 import jax.numpy as jnp
 import networkx as nx
+
 from haiku._src import base
-from uuid import uuid4
 from jax.scipy.integrate import trapezoid
-from abc import ABC
 from simpleeval import simple_eval
 
 
@@ -215,16 +218,18 @@ class SpectralModel:
                 continuum[node_id] = runtime_modules[node_id].continuum(energies)
 
             elif node and node["type"] == "operation":
-                component_1 = list(self.graph.in_edges(node_id))[0][0]
+                component_1 = list(self.graph.in_edges(node_id))[0][0]  # noqa: RUF015
                 component_2 = list(self.graph.in_edges(node_id))[1][0]
-                continuum[node_id] = node["function"](continuum[component_1], continuum[component_2])
+                continuum[node_id] = node["function"](
+                    continuum[component_1], continuum[component_2]
+                )
 
         if n_points == 2:
-            flux_1D = continuum[list(self.graph.in_edges("out"))[0][0]]
+            flux_1D = continuum[list(self.graph.in_edges("out"))[0][0]]  # noqa: RUF015
             flux = jnp.stack((flux_1D[:-1], flux_1D[1:]))
 
         else:
-            flux = continuum[list(self.graph.in_edges("out"))[0][0]]
+            flux = continuum[list(self.graph.in_edges("out"))[0][0]]  # noqa: RUF015
 
         if energy_flux:
             continuum_flux = trapezoid(
@@ -234,7 +239,9 @@ class SpectralModel:
             )
 
         else:
-            continuum_flux = trapezoid(flux * energies_to_integrate, x=jnp.log(energies_to_integrate), axis=0)
+            continuum_flux = trapezoid(
+                flux * energies_to_integrate, x=jnp.log(energies_to_integrate), axis=0
+            )
 
         # Iterate from the root nodes to the output node and
         # compute the fine structure contribution for each component
@@ -249,14 +256,18 @@ class SpectralModel:
             path = nx.shortest_path(self.graph, source=root_node_id, target="out")
             nodes_id_in_path = [node_id for node_id in path]
 
-            flux_from_component, mean_energy = runtime_modules[root_node_id].emission_lines(e_low, e_high)
+            flux_from_component, mean_energy = runtime_modules[root_node_id].emission_lines(
+                e_low, e_high
+            )
 
             multiplicative_nodes = []
 
             # Search all multiplicative components connected to this node
             # and apply them at mean energy
             for node_id in nodes_id_in_path[::-1]:
-                multiplicative_nodes.extend([node_id for node_id in self.find_multiplicative_components(node_id)])
+                multiplicative_nodes.extend(
+                    [node_id for node_id in self.find_multiplicative_components(node_id)]
+                )
 
             for mul_node in multiplicative_nodes:
                 flux_from_component *= runtime_modules[mul_node].continuum(mean_energy)
@@ -309,7 +320,10 @@ class SpectralModel:
         if component.type == "additive":
 
             def lam_func(e):
-                return component(**kwargs).continuum(e) + component(**kwargs).emission_lines(e, e + 1)[0]
+                return (
+                    component(**kwargs).continuum(e)
+                    + component(**kwargs).emission_lines(e, e + 1)[0]
+                )
 
         elif component.type == "multiplicative":
 
@@ -342,7 +356,9 @@ class SpectralModel:
 
         return cls(graph, labels)
 
-    def compose(self, other: SpectralModel, operation=None, function=None, name=None) -> SpectralModel:
+    def compose(
+        self, other: SpectralModel, operation=None, function=None, name=None
+    ) -> SpectralModel:
         """
         This function operate a composition between the operation graph of two models
         1) It fuses the two graphs using which joins at the 'out' nodes
@@ -524,3 +540,10 @@ class AdditiveComponent(ModelComponent, ABC):
 
         return jnp.trapz(self(x) * dx, x=t)
     '''
+
+
+class MultiplicativeComponent(ModelComponent, ABC):
+    type = "multiplicative"
+
+    @abstractmethod
+    def continuum(self, energy): ...

{jaxspec-0.1.0 → jaxspec-0.1.1}/src/jaxspec/model/additive.py

@@ -38,7 +38,7 @@ class Powerlaw(AdditiveComponent):
         return norm * energy ** (-alpha)
 
 
-class AdditiveConstant(AdditiveComponent):
+class Additiveconstant(AdditiveComponent):
     r"""
     A constant model
 

{jaxspec-0.1.0 → jaxspec-0.1.1}/src/jaxspec/model/multiplicative.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
-
 import haiku as hk
 import jax.numpy as jnp
 import numpy as np
@@ -10,14 +8,7 @@ from astropy.table import Table
 from haiku.initializers import Constant as HaikuConstant
 
 from ..util.online_storage import table_manager
-from .abc import ModelComponent
-
-
-class MultiplicativeComponent(ModelComponent, ABC):
-    type = "multiplicative"
-
-    @abstractmethod
-    def continuum(self, energy): ...
+from .abc import MultiplicativeComponent
 
 
 class Expfac(MultiplicativeComponent):
@@ -226,6 +217,7 @@ class Tbpcf(MultiplicativeComponent):
 
         return f * jnp.exp(-nh * sigma) + (1 - f)
 
+
 class FDcut(MultiplicativeComponent):
     r"""
     A Fermi-Dirac cutoff model.
@@ -243,4 +235,4 @@ class FDcut(MultiplicativeComponent):
         cutoff = hk.get_parameter("E_c", [], init=HaikuConstant(1))
         folding = hk.get_parameter("E_f", [], init=HaikuConstant(1))
 
-        return (1 + jnp.exp((energy - cutoff)/folding)) ** -1
+        return (1 + jnp.exp((energy - cutoff) / folding)) ** -1

{jaxspec-0.1.0 → jaxspec-0.1.1}/src/jaxspec/util/typing.py

@@ -9,6 +9,13 @@ from pydantic import BaseModel, field_validator
 PriorDictType = dict[str, dict[str, dist.Distribution | ArrayLike]]
 
 
+def is_flat_dict(input_data: dict[str, Any]) -> bool:
+    """
+    Check if the input data is a flat dictionary with string keys and non-dictionary values.
+    """
+    return all(isinstance(k, str) and not isinstance(v, dict) for k, v in input_data.items())
+
+
 class PriorDictModel(BaseModel):
     """
     Pydantic model for a nested dictionary of NumPyro distributions or JAX arrays.
@@ -21,6 +28,23 @@ class PriorDictModel(BaseModel):
     class Config:  # noqa D106
         arbitrary_types_allowed = True
 
+    @classmethod
+    def from_dict(cls, input_prior: dict[str, Any]):
+        if is_flat_dict(input_prior):
+            nested_dict = {}
+
+            for key, obj in input_prior.items():
+                component, component_number, *parameter = key.split("_")
+
+                sub_dict = nested_dict.get(f"{component}_{component_number}", {})
+                sub_dict["_".join(parameter)] = obj
+
+                nested_dict[f"{component}_{component_number}"] = sub_dict
+
+            return cls(nested_dict=nested_dict)
+
+        return cls(nested_dict=input_prior)
+
     @field_validator("nested_dict", mode="before")
     def check_and_cast_nested_dict(cls, value: dict[str, Any]):
         if not isinstance(value, dict):
@@ -35,9 +59,10 @@ class PriorDictModel(BaseModel):
                     try:
                         # Attempt to cast to JAX array
                         value[key][inner_key] = jnp.array(obj, dtype=float)
+
                     except Exception as e:
                         raise ValueError(
-                            f'The value for key "{inner_key}" in inner dictionary must '
-                            f"be a NumPyro distribution or castable to JAX array. Error: {e}"
+                            f'The value for key "{inner_key}" in {key} be a NumPyro '
+                            f"distribution or castable to JAX array. Error: {e}"
                         )
         return value