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

jaxspec/analysis/results.py

@@ -168,7 +168,8 @@ class FitResult:
         e_max: float,
         unit: Unit = u.photon / u.cm**2 / u.s,
         register: bool = False,
-        n_points: int = 100,
+        n_points: int = 5,
+        n_grid: int = 1_000,
     ) -> ArrayLike:
         """
         Compute the unfolded photon flux in a given energy band. The flux is then added to
@@ -179,23 +180,26 @@ 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.
+            n_points: The number of points per bin to use for computing the unfolded spectrum.
+            n_grid: The number of grid points to use for computing the unfolded spectrum.
 
         !!! warning
             Computation of the folded flux is not implemented yet. Feel free to open an
             [issue](https://github.com/renecotyfanboy/jaxspec/issues) in the GitHub repository.
         """
 
+        energy_grid = np.linspace(e_min, e_max, n_grid)
+
         @jax.jit
         @jnp.vectorize
         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=n_points
-            )[0]
+                parameters_pytree, energy_grid[:-1], energy_grid[1:], n_points=n_points
+            )
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
-        flux = vectorized_flux(*flat_tree)
+        flux = vectorized_flux(*flat_tree).sum(axis=-1)  # Sum over all bins
         conversion_factor = float((u.photon / u.cm**2 / u.s).to(unit))
         value = np.asarray(flux * conversion_factor)
 
@@ -220,7 +224,8 @@ class FitResult:
         e_max: float,
         unit: Unit = u.erg / u.cm**2 / u.s,
         register: bool = False,
-        n_points: int = 100,
+        n_points: int = 5,
+        n_grid: int = 1_000,
     ) -> ArrayLike:
         """
         Compute the unfolded energy flux in a given energy band. The flux is then added to
@@ -231,26 +236,29 @@ 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.
+            n_points: The number of points per bin to use for computing the unfolded spectrum.
+            n_grid: The number of grid points to use for computing the unfolded spectrum.
 
         !!! warning
             Computation of the folded flux is not implemented yet. Feel free to open an
             [issue](https://github.com/renecotyfanboy/jaxspec/issues) in the GitHub repository.
         """
 
+        energy_grid = np.linspace(e_min, e_max, n_grid)
+
         @jax.jit
         @jnp.vectorize
         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=n_points
-            )[0]
+                parameters_pytree, energy_grid[:-1], energy_grid[1:], n_points=n_points
+            )
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
-        flux = vectorized_flux(*flat_tree)
+        flux = vectorized_flux(*flat_tree).sum(axis=-1)  # Sum over all bins
         conversion_factor = float((u.keV / u.cm**2 / u.s).to(unit))
         value = np.asarray(flux * conversion_factor)
-
+        # TODO : ADD TESTS WITH BACKGROUND
         if register:
             self.inference_data.posterior[f"mod/~/energy_flux_{e_min:.1f}_{e_max:.1f}"] = (
                 xr.DataArray(
@@ -276,7 +284,8 @@ class FitResult:
         cosmology: Cosmology = Planck18,
         unit: Unit = u.erg / u.s,
         register: bool = False,
-        n_points=100,
+        n_points: int = 5,
+        n_grid: int = 1_000,
     ) -> ArrayLike:
         """
         Compute the luminosity of the source specifying its redshift. The luminosity is then added to
@@ -290,8 +299,12 @@ class FitResult:
             cosmology: Chosen cosmology.
             unit: The unit of the luminosity.
             register: Whether to register the flux with the other posterior parameters.
+            n_points: The number of points per bin to use for computing the unfolded spectrum.
+            n_grid: The number of grid points to use for computing the unfolded spectrum.
         """
 
+        energy_grid = np.linspace(e_min, e_max, n_grid)
+
         if not observer_frame:
             raise NotImplementedError()
 
@@ -312,13 +325,13 @@ class FitResult:
             parameters_pytree = jax.tree.unflatten(pytree_def, pars)
             return self.model.energy_flux(
                 parameters_pytree,
-                jnp.asarray([e_min]) * (1 + redshift),
-                jnp.asarray([e_max]) * (1 + redshift),
+                energy_grid[:-1] * (1 + redshift),
+                energy_grid[1:] * (1 + redshift),
                 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)
+        flux = vectorized_flux(*flat_tree).sum(axis=-1) * (u.keV / u.cm**2 / u.s)
         value = np.asarray(
             (flux * (4 * np.pi * cosmology.luminosity_distance(redshift) ** 2)).to(unit)
         )

jaxspec/data/__init__.py

@@ -6,5 +6,7 @@ from .observation import Observation
 
 u.add_enabled_aliases({"counts": u.count})
 u.add_enabled_aliases({"channel": u.dimensionless_unscaled})
+u.add_enabled_aliases({"ADU": u.dimensionless_unscaled})  # Appears in SIXTE outputs
+
 # Arbitrary units are found in .rsp files , let's hope it is compatible with what we would expect as the rmf x arf
 # u.add_enabled_aliases({"au": u.dimensionless_unscaled})

jaxspec/fit/_fitter.py

@@ -15,6 +15,7 @@ from numpyro.infer import AIES, ESS, MCMC, NUTS, SVI, Predictive, Trace_ELBO
 from numpyro.infer.autoguide import AutoMultivariateNormal
 
 from ..analysis.results import FitResult
+from ..model.background import SubtractedBackground
 from ._bayesian_model import BayesianModel
 
 
@@ -58,7 +59,11 @@ class BayesianModelFitter(BayesianModel, ABC):
             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:
+
+            # Subtracted background is not fitted so there is no likelihood
+            if self.background_model is not None and not isinstance(
+                self.background_model, SubtractedBackground
+            ):
                 log_likelihood["bkg/~/all"] = concatenate(
                     [ll for k, ll in log_likelihood.items() if "bkg" in k], axis=1
                 )

jaxspec/model/additive.py

@@ -34,8 +34,10 @@ class Powerlaw(AdditiveComponent):
         self.alpha = nnx.Param(1.7)
         self.norm = nnx.Param(1e-4)
 
-    def continuum(self, energy):
-        return self.norm * energy ** (-self.alpha)
+    def integrated_continuum(self, e_low, e_high):
+        return (
+            self.norm / (1 - self.alpha) * (e_high ** (1 - self.alpha) - e_low ** (1 - self.alpha))
+        )
 
 
 class Additiveconstant(AdditiveComponent):
@@ -166,28 +168,22 @@ class Gauss(AdditiveComponent):
         self.sigma = nnx.Param(1e-2)
         self.norm = nnx.Param(1.0)
 
-    def continuum(self, energy):
-        return (
-            self.norm
-            * jsp.stats.norm.pdf(energy, loc=jnp.asarray(self.El), scale=jnp.asarray(self.sigma))
-            / (1 - jsp.special.erf(-self.El / (self.sigma * jnp.sqrt(2))))
+    def integrated_continuum(self, e_low, e_high):
+        upper = jsp.stats.norm.cdf(
+            e_high,
+            loc=jnp.asarray(self.El),
+            scale=jnp.asarray(self.sigma),
         )
 
-    """
-    def integrated_continuum(self, e_low, e_high):
-        return self.norm * (
-            jsp.stats.norm.cdf(
-                e_high,
-                loc=jnp.asarray(self.El),
-                scale=jnp.asarray(self.sigma),
-            )
-            - jsp.stats.norm.cdf(
-                e_low,
-                loc=jnp.asarray(self.El),
-                scale=jnp.asarray(self.sigma),
-            ) #/  (1 - jsp.special.erf(- self.El / (self.sigma * jnp.sqrt(2))))
+        lower = jsp.stats.norm.cdf(
+            e_low,
+            loc=jnp.asarray(self.El),
+            scale=jnp.asarray(self.sigma),
         )
-    """
+
+        factor = 2 / (1 - jsp.special.erf(-self.El / (self.sigma * jnp.sqrt(2))))
+
+        return self.norm * (upper - lower) * factor
 
 
 class Cutoffpl(AdditiveComponent):

jaxspec/model/multiplicative.py

@@ -49,7 +49,9 @@ class Expfac(MultiplicativeComponent):
         self.E_c = nnx.Param(1.0)
 
     def factor(self, energy):
-        return jnp.where(energy >= self.E_c, 1.0 + self.A * jnp.exp(-self.f * energy), 1.0)
+        return jnp.where(
+            energy >= self.E_c, 1.0 + self.A * jnp.exp(-self.f * energy), 1.0
+        )
 
 
 class Tbabs(MultiplicativeComponent):
@@ -91,6 +93,49 @@ class Tbabs(MultiplicativeComponent):
         return jnp.exp(-self.nh * sigma)
 
 
+class zTbabs(MultiplicativeComponent):
+    r"""
+    The redshifted Tuebingen-Boulder ISM absorption model. See `Tbabs` for more details.
+    From Xspec manual:
+    This model assumes that 20% of the hydrogen is molecular
+    and that there is NO MATERIAL IN GRAINS.
+
+    $$
+        \mathcal{M}(E) = \exp^{-N_{\text{H}}\sigma(E)}
+    $$
+
+    !!! abstract "Parameters"
+        * $N_{\text{H}}$ (`nh`) $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$ : Equivalent hydrogen column density
+        * $z$ (`z`) $\left[\text{dimensionless}\right]$ : Redshift
+
+
+    !!! note
+        Abundances and cross-sections $\sigma$ can be found in Wilms et al. (2000).
+
+    """
+
+    def __init__(self):
+        table = Table.read(table_manager.fetch("xsect_tbabs_wilm.fits"))
+        self._energy = np.asarray(table["ENERGY"], dtype=np.float64)
+        self._sigma = np.asarray(table["SIGMA"], dtype=np.float64)
+        self.nh = nnx.Param(1.0)
+        self.z = nnx.Param(1.0)
+
+    def factor(self, energy):
+        z = jnp.asarray(self.z)
+        sigma = jnp.exp(
+            jnp.interp(
+                jnp.log(energy) + jnp.log1p(z),
+                jnp.log(self._energy),
+                jnp.log(self._sigma),
+                left="extrapolate",
+                right="extrapolate",
+            )
+        )
+
+        return jnp.exp(-self.nh * sigma)
+
+
 class Phabs(MultiplicativeComponent):
     r"""
     A photoelectric absorption model.
@@ -215,7 +260,9 @@ class Zedge(MultiplicativeComponent):
 
     def factor(self, energy):
         return jnp.where(
-            energy <= self.Ec, 1.0, jnp.exp(-self.D * (energy * (1 + self.z) / self.Ec) ** 3)
+            energy <= self.Ec,
+            1.0,
+            jnp.exp(-self.D * (energy * (1 + self.z) / self.Ec) ** 3),
         )
 
 
@@ -246,7 +293,11 @@ class Tbpcf(MultiplicativeComponent):
     def factor(self, energy):
         sigma = jnp.exp(
             jnp.interp(
-                energy, self._energy, jnp.log(self._sigma), left="extrapolate", right="extrapolate"
+                energy,
+                self._energy,
+                jnp.log(self._sigma),
+                left="extrapolate",
+                right="extrapolate",
             )
         )
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.3.3
+Version: 0.3.4
 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/
@@ -8,7 +8,7 @@ Author-email: sdupourque <sdupourque@irap.omp.eu>
 License-Expression: MIT
 License-File: LICENSE.md
 Requires-Python: <3.13,>=3.10
-Requires-Dist: arviz<0.23.0,>=0.17.1
+Requires-Dist: arviz<0.24.0,>=0.17.1
 Requires-Dist: astropy<8,>=6.0.0
 Requires-Dist: catppuccin<3,>=2.3.4
 Requires-Dist: chainconsumer<2,>=1.1.2
@@ -20,7 +20,7 @@ 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: numpyro<0.21,>=0.17.0
 Requires-Dist: pandas<3,>=2.2.0
 Requires-Dist: pooch<2,>=1.8.2
 Requires-Dist: scipy<1.16

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

@@ -2,8 +2,8 @@ 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=WoYpCR485TdMorTLZhBXXrPnw8qidD5LkOGmHHu-hZE,27923
-jaxspec/data/__init__.py,sha256=aantcYKC9kZFvaE-V2SIwSuLhIld17Kjrd9CIUu___Y,415
+jaxspec/analysis/results.py,sha256=nZ7JORgA6YYei8hRHmGvhUmnSpd11FfryP-E7UVVT9s,28650
+jaxspec/data/__init__.py,sha256=9fZRyB3eXdEi_ZsTcHT64xH3kW3jQFD0XS5eIAD0RDo,501
 jaxspec/data/instrument.py,sha256=weiPcEll1jZM6lqhxpF1aPIRwvaP6bygSB8jLBABXto,4815
 jaxspec/data/obsconf.py,sha256=bkYuD6mJgj8QmRaDVhcnXwUukVdo20xllzaI57prHag,10056
 jaxspec/data/observation.py,sha256=7FHJm1jHEEFyrqxg3COsGmfdh5dg-5XnfKCp1yb5fNY,7411
@@ -18,15 +18,15 @@ 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=KjOqKVoB9vd2Lg-8-j-3pS78ZmhA2_ICHLkxVEJHiYc,10160
+jaxspec/fit/_fitter.py,sha256=lkLFvwqE6NGpBk4i8gdTHGD8QQahCp9OgJS96j7N6FA,10383
 jaxspec/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/model/_graph_util.py,sha256=hPvHYmAxb7P3nyIecaZ7RqWOjwcZ1WvUByt_yNANiaY,4552
 jaxspec/model/abc.py,sha256=vvHM4teepc8VLqbpAtqf1b55oF00R_Lo_6nrBO5KmmQ,14793
-jaxspec/model/additive.py,sha256=rEONSy7b7lwfXIhuPqtI4y2Yhq55EqrlEtEckEe6TA0,20538
+jaxspec/model/additive.py,sha256=MGHqJ4Ai2kPkMrZOLom4vTNyeg2oczpg4A9Rv5JdU34,20366
 jaxspec/model/background.py,sha256=VLSrU0YCW9GSHCtaEdcth-sp74aPyEVSizIMFkTpM7M,7759
 jaxspec/model/instrument.py,sha256=1zLZgHmBZs8RLKTMT3Wu4bCx6JnxBUjhRIpYG2rLaZM,2947
 jaxspec/model/list.py,sha256=uC9rLEEeph10q6shat86WLACVuTSx73RGMl8Ij0jqQY,875
-jaxspec/model/multiplicative.py,sha256=odaOlF0K1KjzUDstPtWAk95ScHoZ7_XveOew6l3tbeU,8337
+jaxspec/model/multiplicative.py,sha256=L4zyCrYU364Cwb2bDJc9ydSc_EMzNQ21zNWIM8EbLKE,9793
 jaxspec/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/scripts/debug.py,sha256=qhyDtX4G5UdChmTLCM-5Wti4XZU-sU5S-wDb6TZjrvM,292
 jaxspec/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -35,8 +35,8 @@ jaxspec/util/integrate.py,sha256=7GwBSagmDzsF3P53tPs-oakeq0zHEwmZZS2zQlXngbE,463
 jaxspec/util/misc.py,sha256=O3qorCL1Y2X1BS2jdd36C1eDHK9QDXTSOr9kj3uqcJo,654
 jaxspec/util/online_storage.py,sha256=wwpowxmDgAqKzeUwmGUIxttA4VKUoR270Ew-F_0DrkE,2493
 jaxspec/util/typing.py,sha256=ZQM_l68qyYnIBZPz_1mKvwPMx64jvVBD8Uj6bx9sHv0,140
-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.4.dist-info/METADATA,sha256=VzNrPKy8_5ReZk5REcwADXU-m8QpiICd-Gx9ow_1X2w,4045
+jaxspec-0.3.4.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+jaxspec-0.3.4.dist-info/entry_points.txt,sha256=4ffU5AImfcEBxgWTqopQll2YffpFldOswXRh16pd0Dc,72
+jaxspec-0.3.4.dist-info/licenses/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
+jaxspec-0.3.4.dist-info/RECORD,,