jaxspec 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

jaxspec/analysis/results.py

@@ -77,8 +77,9 @@ class FitResult:
         r"""
         Convergence of the chain as computed by the $\hat{R}$ statistic.
         """
+        rhat = az.rhat(self.inference_data)
 
-        return all(az.rhat(self.inference_data) < 1.01)
+        return bool((rhat.to_array() < 1.01).all())
 
     def _ppc_folded_branches(self, obs_id):
         obs = self.obsconfs[obs_id]
@@ -167,6 +168,7 @@ class FitResult:
         e_max: float,
         unit: Unit = u.photon / u.cm**2 / u.s,
         register: bool = False,
+        n_points: int = 100,
     ) -> ArrayLike:
         """
         Compute the unfolded photon flux in a given energy band. The flux is then added to
@@ -177,6 +179,7 @@ class FitResult:
             e_max: The upper bound of the energy band in observer frame.
             unit: The unit of the photon flux.
             register: Whether to register the flux with the other posterior parameters.
+            n_points: The number of points to use for computing the unfolded spectrum.
 
         !!! warning
             Computation of the folded flux is not implemented yet. Feel free to open an
@@ -188,18 +191,25 @@ class FitResult:
         def vectorized_flux(*pars):
             parameters_pytree = jax.tree.unflatten(pytree_def, pars)
             return self.model.photon_flux(
-                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=100
+                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=n_points
             )[0]
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
         flux = vectorized_flux(*flat_tree)
-        conversion_factor = (u.photon / u.cm**2 / u.s).to(unit)
-        value = flux * conversion_factor
+        conversion_factor = float((u.photon / u.cm**2 / u.s).to(unit))
+        value = np.asarray(flux * conversion_factor)
 
         if register:
-            self.inference_data.posterior[f"photon_flux_{e_min:.1f}_{e_max:.1f}"] = (
-                list(self.inference_data.posterior.coords),
-                value,
+            self.inference_data.posterior[f"mod/~/photon_flux_{e_min:.1f}_{e_max:.1f}"] = (
+                xr.DataArray(
+                    value,
+                    dims=self.inference_data.posterior.dims,
+                    coords={
+                        "chain": self.inference_data.posterior.chain,
+                        "draw": self.inference_data.posterior.draw,
+                    },
+                    name=f"mod/~/photon_flux_{e_min:.1f}_{e_max:.1f}",
+                )
             )
 
         return value
@@ -210,6 +220,7 @@ class FitResult:
         e_max: float,
         unit: Unit = u.erg / u.cm**2 / u.s,
         register: bool = False,
+        n_points: int = 100,
     ) -> ArrayLike:
         """
         Compute the unfolded energy flux in a given energy band. The flux is then added to
@@ -220,6 +231,7 @@ class FitResult:
             e_max: The upper bound of the energy band in observer frame.
             unit: The unit of the energy flux.
             register: Whether to register the flux with the other posterior parameters.
+            n_points: The number of points to use for computing the unfolded spectrum.
 
         !!! warning
             Computation of the folded flux is not implemented yet. Feel free to open an
@@ -231,18 +243,25 @@ class FitResult:
         def vectorized_flux(*pars):
             parameters_pytree = jax.tree.unflatten(pytree_def, pars)
             return self.model.energy_flux(
-                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=100
+                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=n_points
             )[0]
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
         flux = vectorized_flux(*flat_tree)
-        conversion_factor = (u.keV / u.cm**2 / u.s).to(unit)
-        value = flux * conversion_factor
+        conversion_factor = float((u.keV / u.cm**2 / u.s).to(unit))
+        value = np.asarray(flux * conversion_factor)
 
         if register:
-            self.inference_data.posterior[f"energy_flux_{e_min:.1f}_{e_max:.1f}"] = (
-                list(self.inference_data.posterior.coords),
-                value,
+            self.inference_data.posterior[f"mod/~/energy_flux_{e_min:.1f}_{e_max:.1f}"] = (
+                xr.DataArray(
+                    value,
+                    dims=self.inference_data.posterior.dims,
+                    coords={
+                        "chain": self.inference_data.posterior.chain,
+                        "draw": self.inference_data.posterior.draw,
+                    },
+                    name=f"mod/~/energy_flux_{e_min:.1f}_{e_max:.1f}",
+                )
             )
 
         return value
@@ -257,6 +276,7 @@ class FitResult:
         cosmology: Cosmology = Planck18,
         unit: Unit = u.erg / u.s,
         register: bool = False,
+        n_points=100,
     ) -> ArrayLike:
         """
         Compute the luminosity of the source specifying its redshift. The luminosity is then added to
@@ -294,17 +314,26 @@ class FitResult:
                 parameters_pytree,
                 jnp.asarray([e_min]) * (1 + redshift),
                 jnp.asarray([e_max]) * (1 + redshift),
-                n_points=100,
+                n_points=n_points,
             )[0]
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
         flux = vectorized_flux(*flat_tree) * (u.keV / u.cm**2 / u.s)
-        value = (flux * (4 * np.pi * cosmology.luminosity_distance(redshift) ** 2)).to(unit)
+        value = np.asarray(
+            (flux * (4 * np.pi * cosmology.luminosity_distance(redshift) ** 2)).to(unit)
+        )
 
         if register:
-            self.inference_data.posterior[f"luminosity_{e_min:.1f}_{e_max:.1f}"] = (
-                list(self.inference_data.posterior.coords),
-                value,
+            self.inference_data.posterior[f"mod/~/luminosity_{e_min:.1f}_{e_max:.1f}"] = (
+                xr.DataArray(
+                    value,
+                    dims=self.inference_data.posterior.dims,
+                    coords={
+                        "chain": self.inference_data.posterior.chain,
+                        "draw": self.inference_data.posterior.draw,
+                    },
+                    name=f"mod/~/luminosity_{e_min:.1f}_{e_max:.1f}",
+                )
             )
 
         return value
@@ -315,10 +344,13 @@ class FitResult:
 
         Parameters:
             name: The name of the chain.
+            parameter_kind: The kind of parameters to keep.
         """
 
         keys_to_drop = [
-            key for key in self.inference_data.posterior.keys() if not key.startswith("mod")
+            key
+            for key in self.inference_data.posterior.keys()
+            if not key.startswith(parameter_kind)
         ]
 
         reduced_id = az.extract(
@@ -403,6 +435,7 @@ class FitResult:
         title: str | None = None,
         figsize: tuple[float, float] = (6, 6),
         x_lims: tuple[float, float] | None = None,
+        rescale_background: bool = False,
     ) -> list[plt.Figure]:
         r"""
         Plot the posterior predictive distribution of the model. It also features a residual plot, defined using the
@@ -423,6 +456,7 @@ class FitResult:
             title: The title of the plot.
             figsize: The size of the figure.
             x_lims: The limits of the x-axis.
+            rescale_background: Whether to rescale the background model to the data with backscal ratio.
 
         Returns:
             A list of matplotlib figures for each observation in the model.
@@ -573,10 +607,14 @@ class FitResult:
                         )
                     )
 
+                    rescale_background_factor = (
+                        obsconf.folded_backratio.data if rescale_background else 1.0
+                    )
+
                     model_bkg_plot = _plot_binned_samples_with_error(
                         ax[0],
                         xbins.value,
-                        y_samples_bkg.value,
+                        y_samples_bkg.value * rescale_background_factor,
                         color=BACKGROUND_COLOR,
                         alpha_envelope=alpha_envelope,
                         n_sigmas=n_sigmas,
@@ -585,9 +623,9 @@ class FitResult:
                     true_bkg_plot = _plot_poisson_data_with_error(
                         ax[0],
                         xbins.value,
-                        y_observed_bkg.value,
-                        y_observed_bkg_low.value,
-                        y_observed_bkg_high.value,
+                        y_observed_bkg.value * rescale_background_factor,
+                        y_observed_bkg_low.value * rescale_background_factor,
+                        y_observed_bkg_high.value * rescale_background_factor,
                         color=BACKGROUND_DATA_COLOR,
                         alpha=0.7,
                     )

jaxspec/data/instrument.py

@@ -1,4 +1,5 @@
 import os
+import sparse
 
 import numpy as np
 import xarray as xr
@@ -92,7 +93,7 @@ class Instrument(xr.Dataset):
 
         else:
             specresp = rmf.matrix.sum(axis=0)
-            rmf.sparse_matrix /= specresp
+            rmf.sparse_matrix = sparse.COO( rmf.matrix / specresp )
 
         return cls.from_matrix(
             rmf.sparse_matrix, specresp, rmf.energ_lo, rmf.energ_hi, rmf.e_min, rmf.e_max

jaxspec/data/util.py

@@ -152,11 +152,11 @@ def forward_model_with_multiple_inputs(
         transfer_matrix = BCOO.from_scipy_sparse(
             obs_configuration.transfer_matrix.data.to_scipy_sparse().tocsr()
         )
+        expected_counts = transfer_matrix @ flux_func(parameters).T
 
     else:
         transfer_matrix = np.asarray(obs_configuration.transfer_matrix.data.todense())
-
-    expected_counts = jnp.matvec(transfer_matrix, flux_func(parameters))
+        expected_counts = jnp.matvec(transfer_matrix, flux_func(parameters))
 
     # The result is clipped at 1e-6 to avoid 0 round-off and diverging likelihoods
     return jnp.clip(expected_counts, a_min=1e-6)

jaxspec/fit/_fitter.py

@@ -9,6 +9,7 @@ import matplotlib.pyplot as plt
 import numpyro
 
 from jax import random
+from jax.numpy import concatenate
 from jax.random import PRNGKey
 from numpyro.infer import AIES, ESS, MCMC, NUTS, SVI, Predictive, Trace_ELBO
 from numpyro.infer.autoguide import AutoMultivariateNormal
@@ -52,9 +53,18 @@ class BayesianModelFitter(BayesianModel, ABC):
         )
 
         log_likelihood = numpyro.infer.log_likelihood(numpyro_model, posterior_samples)
+        if len(log_likelihood.keys()) > 1:
+            log_likelihood["full"] = concatenate([ll for _, ll in log_likelihood.items()], axis=1)
+            log_likelihood["obs/~/all"] = concatenate(
+                [ll for k, ll in log_likelihood.items() if "obs" in k], axis=1
+            )
+            if self.background_model is not None:
+                log_likelihood["bkg/~/all"] = concatenate(
+                    [ll for k, ll in log_likelihood.items() if "bkg" in k], axis=1
+                )
 
         seeded_model = numpyro.handlers.substitute(
-            numpyro.handlers.seed(numpyro_model, keys[3]),
+            numpyro.handlers.seed(numpyro_model, keys[2]),
             substitute_fn=numpyro.infer.init_to_sample,
         )
 
@@ -108,12 +118,10 @@ class BayesianModelFitter(BayesianModel, ABC):
         predictive_parameters = []
 
         for key, value in self._observation_container.items():
+            predictive_parameters.append(f"obs/~/{key}")
             if self.background_model is not None:
-                predictive_parameters.append(f"obs/~/{key}")
                 predictive_parameters.append(f"bkg/~/{key}")
             #                predictive_parameters.append(f"ins/~/{key}")
-            else:
-                predictive_parameters.append(f"obs/~/{key}")
         #                predictive_parameters.append(f"ins/~/{key}")
 
         inference_data.posterior_predictive = inference_data.posterior_predictive[
@@ -247,7 +255,7 @@ class VIFitter(BayesianModelFitter):
 
         svi = SVI(bayesian_model, guide, optimizer, loss=loss)
 
-        keys = random.split(random.PRNGKey(rng_key), 3)
+        keys = random.split(random.PRNGKey(rng_key), 2)
         svi_result = svi.run(keys[0], num_steps)
         params = svi_result.params
 

jaxspec/model/abc.py

@@ -372,9 +372,10 @@ class AdditiveComponent(ModelComponent):
         continuum = self.continuum(energy)
         integrated_continuum = self.integrated_continuum(e_low, e_high)
 
-        return jsp.integrate.trapezoid(
-            continuum * energy**2, jnp.log(energy), axis=-1
-        ) + integrated_continuum * (e_high - e_low)
+        return (
+            jsp.integrate.trapezoid(continuum * energy**2, jnp.log(energy), axis=-1)
+            + integrated_continuum * (e_high + e_low) / 2.0
+        )
 
     @partial(jax.jit, static_argnums=0, static_argnames="n_points")
     def photon_flux(self, params, e_low, e_high, n_points=2):

jaxspec/util/online_storage.py

@@ -25,4 +25,5 @@ table_manager = pooch.create(
         "example_data/NGC7793_ULX4/MOS2.arf": "sha256:a126ff5a95a5f4bb93ed846944cf411d6e1c448626cb73d347e33324663d8b3f",
         "example_data/NGC7793_ULX4/PNbackground_spectrum.fits": "sha256:55e017e0c19b324245fef049dff2a7a2e49b9a391667ca9c4f667c4f683b1f49",
     },
+    retry_if_failed=10,
 )

{jaxspec-0.3.2.dist-info → jaxspec-0.3.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.3.2
+Version: 0.3.3
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
 Project-URL: Homepage, https://github.com/renecotyfanboy/jaxspec
 Project-URL: Documentation, https://jaxspec.readthedocs.io/en/latest/
@@ -16,19 +16,15 @@ Requires-Dist: cmasher<2,>=1.6.3
 Requires-Dist: flax>0.10.5
 Requires-Dist: interpax<0.4,>=0.3.5
 Requires-Dist: jax<0.7,>=0.5.0
-Requires-Dist: jaxns<3,>=2.6.7
-Requires-Dist: jaxopt<0.9,>=0.8.3
 Requires-Dist: matplotlib<4,>=3.8.0
 Requires-Dist: mendeleev<1.2,>=0.15
 Requires-Dist: networkx~=3.1
 Requires-Dist: numpy<3.0.0
 Requires-Dist: numpyro<0.20,>=0.17.0
-Requires-Dist: optimistix<0.0.12,>=0.0.10
 Requires-Dist: pandas<3,>=2.2.0
 Requires-Dist: pooch<2,>=1.8.2
 Requires-Dist: scipy<1.16
-Requires-Dist: seaborn<0.14,>=0.13.1
-Requires-Dist: simpleeval<1.1.0,>=0.9.13
+Requires-Dist: seaborn>=0.13.2
 Requires-Dist: sparse>0.15
 Requires-Dist: tinygp<0.4,>=0.3.0
 Requires-Dist: watermark<3,>=2.4.3

{jaxspec-0.3.2.dist-info → jaxspec-0.3.3.dist-info}/RECORD

@@ -2,13 +2,13 @@ jaxspec/__init__.py,sha256=Sbn02lX6Y-zNXk17N8dec22c5jeypiS0LkHmGfz7lWA,126
 jaxspec/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/analysis/_plot.py,sha256=0xEz-e_xk7XvU6PUfbNwxaWg1-SxAF2XAqhkxWEhIFs,6239
 jaxspec/analysis/compare.py,sha256=g2UFhmR9Zt-7cz5gQFOB6lXuklXB3yTyUvjTypOzoSY,725
-jaxspec/analysis/results.py,sha256=tIBWmLoX43EY2BXt50ec8A-DqQ98PMd3m-FqTRT4iRE,26073
+jaxspec/analysis/results.py,sha256=WoYpCR485TdMorTLZhBXXrPnw8qidD5LkOGmHHu-hZE,27923
 jaxspec/data/__init__.py,sha256=aantcYKC9kZFvaE-V2SIwSuLhIld17Kjrd9CIUu___Y,415
-jaxspec/data/instrument.py,sha256=RDiG_LkucvnF2XE_ghTFME6d_2YirgQUcEY0gEle6dk,4775
+jaxspec/data/instrument.py,sha256=weiPcEll1jZM6lqhxpF1aPIRwvaP6bygSB8jLBABXto,4815
 jaxspec/data/obsconf.py,sha256=bkYuD6mJgj8QmRaDVhcnXwUukVdo20xllzaI57prHag,10056
 jaxspec/data/observation.py,sha256=7FHJm1jHEEFyrqxg3COsGmfdh5dg-5XnfKCp1yb5fNY,7411
 jaxspec/data/ogip.py,sha256=eMmBuROW4eMRxRHkPPyGHf933e0IcREqB8WMQFMS2lY,9810
-jaxspec/data/util.py,sha256=4_f6ByGjUEZXTwrB37dCyYaTB1pjF10Z0ho7-4GrQuc,9761
+jaxspec/data/util.py,sha256=2JWoHsKJqGXUn74zPeoAqdU86x2n8NyfZGvpqC21ZaY,9832
 jaxspec/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/experimental/interpolator.py,sha256=mJRdCB4B71le3dQL_S_E6Wkqpb6QLT7Wdzlok-rU6Ok,2652
 jaxspec/experimental/interpolator_jax.py,sha256=13lflsjbImDRZTObSRDtZnujrXBvEP367Rn20eByONs,2967
@@ -18,10 +18,10 @@ jaxspec/experimental/tabulated.py,sha256=H0llUiso2KGH4xUzTUSVPy-6I8D3wm707lU_Z1P
 jaxspec/fit/__init__.py,sha256=OaS0-Hkb3Hd-AkE2o-KWfoWMX0NSCPY-_FP2znHf9l0,153
 jaxspec/fit/_bayesian_model.py,sha256=7c2Twgz06QV1S9DdctdVk5YT1v7P-ln100bWXAvv7Go,15179
 jaxspec/fit/_build_model.py,sha256=pNZVuVfwOq3Pg23opH7xRv28DsSkQZpvy2Z-1hQSfNs,3219
-jaxspec/fit/_fitter.py,sha256=92gd1P7CNIqusGb64x_DpBcb0KcoGyfvSDiEnRbfqP0,9709
+jaxspec/fit/_fitter.py,sha256=KjOqKVoB9vd2Lg-8-j-3pS78ZmhA2_ICHLkxVEJHiYc,10160
 jaxspec/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/model/_graph_util.py,sha256=hPvHYmAxb7P3nyIecaZ7RqWOjwcZ1WvUByt_yNANiaY,4552
-jaxspec/model/abc.py,sha256=RGrqDrXVNjCy7GYBZL-l1PZ3Lpr37SsMIw7L9_B8WJ4,14773
+jaxspec/model/abc.py,sha256=vvHM4teepc8VLqbpAtqf1b55oF00R_Lo_6nrBO5KmmQ,14793
 jaxspec/model/additive.py,sha256=rEONSy7b7lwfXIhuPqtI4y2Yhq55EqrlEtEckEe6TA0,20538
 jaxspec/model/background.py,sha256=VLSrU0YCW9GSHCtaEdcth-sp74aPyEVSizIMFkTpM7M,7759
 jaxspec/model/instrument.py,sha256=1zLZgHmBZs8RLKTMT3Wu4bCx6JnxBUjhRIpYG2rLaZM,2947
@@ -33,10 +33,10 @@ jaxspec/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/util/abundance.py,sha256=fsC313taIlGzQsZNwbYsJupDWm7ZbqzGhY66Ku394Mw,8546
 jaxspec/util/integrate.py,sha256=7GwBSagmDzsF3P53tPs-oakeq0zHEwmZZS2zQlXngbE,4634
 jaxspec/util/misc.py,sha256=O3qorCL1Y2X1BS2jdd36C1eDHK9QDXTSOr9kj3uqcJo,654
-jaxspec/util/online_storage.py,sha256=vm56RfcbFKpkRVfr0bXO7J9aQxuBq-I_oEgA26YIhCo,2469
+jaxspec/util/online_storage.py,sha256=wwpowxmDgAqKzeUwmGUIxttA4VKUoR270Ew-F_0DrkE,2493
 jaxspec/util/typing.py,sha256=ZQM_l68qyYnIBZPz_1mKvwPMx64jvVBD8Uj6bx9sHv0,140
-jaxspec-0.3.2.dist-info/METADATA,sha256=10PjN7QwhbU8BoZc9f1Lga2n1u1_j4p8Lk2Syy6cJC8,4199
-jaxspec-0.3.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-jaxspec-0.3.2.dist-info/entry_points.txt,sha256=4ffU5AImfcEBxgWTqopQll2YffpFldOswXRh16pd0Dc,72
-jaxspec-0.3.2.dist-info/licenses/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
-jaxspec-0.3.2.dist-info/RECORD,,
+jaxspec-0.3.3.dist-info/METADATA,sha256=RUeQs3cLTKOL_7lJHR2Nmjg-E0YMaHafSIN2YssSkx4,4045
+jaxspec-0.3.3.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+jaxspec-0.3.3.dist-info/entry_points.txt,sha256=4ffU5AImfcEBxgWTqopQll2YffpFldOswXRh16pd0Dc,72
+jaxspec-0.3.3.dist-info/licenses/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
+jaxspec-0.3.3.dist-info/RECORD,,

{jaxspec-0.3.2.dist-info → jaxspec-0.3.3.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: hatchling 1.27.0
+Generator: hatchling 1.28.0
 Root-Is-Purelib: true
 Tag: py3-none-any