jaxspec 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

jaxspec/model/instrument.py

@@ -0,0 +1,81 @@
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+
+import numpyro
+
+from flax import nnx
+from numpyro.distributions import Distribution
+
+
+class GainModel(ABC, nnx.Module):
+    @abstractmethod
+    def numpyro_model(self, observation_name: str):
+        pass
+
+
+class ConstantGain(GainModel):
+    def __init__(self, prior_distribution: Distribution):
+        self.prior_distribution = prior_distribution
+
+    def numpyro_model(self, observation_name: str):
+        factor = numpyro.sample(f"ins/~/gain_{observation_name}", self.prior_distribution)
+
+        def gain(energy):
+            return factor
+
+        return gain
+
+
+class ShiftModel(ABC, nnx.Module):
+    @abstractmethod
+    def numpyro_model(self, observation_name: str):
+        pass
+
+
+class ConstantShift(ShiftModel):
+    def __init__(self, prior_distribution: Distribution):
+        self.prior_distribution = prior_distribution
+
+    def numpyro_model(self, observation_name: str):
+        shift_offset = numpyro.sample(f"ins/~/shift_{observation_name}", self.prior_distribution)
+
+        def shift(energy):
+            return energy + shift_offset
+
+        return shift
+
+
+class InstrumentModel(nnx.Module):
+    def __init__(
+        self,
+        reference_observation_name: str,
+        gain_model: GainModel | None = None,
+        shift_model: ShiftModel | None = None,
+    ):
+        self.reference = reference_observation_name
+        self.gain_model = gain_model
+        self.shift_model = shift_model
+
+    def get_gain_and_shift_model(
+        self, observation_name: str
+    ) -> tuple[Callable | None, Callable | None]:
+        """
+        Return the gain and shift models for the given observation. It should be called within a numpyro model.
+        """
+
+        if observation_name == self.reference:
+            return None, None
+
+        else:
+            gain = (
+                self.gain_model.numpyro_model(observation_name)
+                if self.gain_model is not None
+                else None
+            )
+            shift = (
+                self.shift_model.numpyro_model(observation_name)
+                if self.shift_model is not None
+                else None
+            )
+
+            return gain, shift

jaxspec/model/list.py

@@ -1,4 +1,5 @@
 import inspect
+
 from .abc import ModelComponent
 from .additive import *  # noqa: F403
 from .multiplicative import *  # noqa: F403
@@ -8,7 +9,9 @@ def all_models(cls: ModelComponent) -> list[ModelComponent]:
     """
     Return a list of all the subclasses of a given ModelComponent class
     """
-    subclasses = list(set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in all_models(c)]))
+    subclasses = list(
+        set(cls.__subclasses__()).union([s for c in cls.__subclasses__() for s in all_models(c)])
+    )
 
     return [s for s in subclasses if not inspect.isabstract(s)]
 

jaxspec/model/multiplicative.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import flax.nnx as nnx
 import jax.numpy as jnp
+import jax.scipy.special as jsp
 import numpy as np
 
 from astropy.table import Table
@@ -72,12 +73,20 @@ class Tbabs(MultiplicativeComponent):
 
     def __init__(self):
         table = Table.read(table_manager.fetch("xsect_tbabs_wilm.fits"))
-        self.energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
-        self.sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
+        self._energy = np.asarray(table["ENERGY"], dtype=np.float64)
+        self._sigma = np.asarray(table["SIGMA"], dtype=np.float64)
         self.nh = nnx.Param(1.0)
 
     def factor(self, energy):
-        sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
+        sigma = jnp.exp(
+            jnp.interp(
+                jnp.log(energy),
+                jnp.log(self._energy),
+                jnp.log(self._sigma),
+                left="extrapolate",
+                right="extrapolate",
+            )
+        )
 
         return jnp.exp(-self.nh * sigma)
 
@@ -89,7 +98,6 @@ class Phabs(MultiplicativeComponent):
     !!! abstract "Parameters"
         * $N_{\text{H}}$ (`nh`) $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$ : Equivalent hydrogen column density
 
-
     """
 
     def __init__(self):
@@ -148,13 +156,20 @@ class Gabs(MultiplicativeComponent):
         self.sigma = nnx.Param(1.0)
         self.E0 = nnx.Param(1.0)
 
-    def factor(self, energy):
-        return jnp.exp(
-            -self.tau
-            / (jnp.sqrt(2 * jnp.pi) * self.sigma)
-            * jnp.exp(-0.5 * ((energy - self.E0) / self.sigma) ** 2)
+    def g(self, E):
+        return (
+            2
+            / (
+                self.sigma
+                * jnp.sqrt(2 * jnp.pi)
+                * (1 - jsp.erf(-self.E0 / (self.sigma * jnp.sqrt(2))))
+            )
+            * jnp.exp(-((E - self.E0) ** 2) / (2 * self.sigma**2))
         )
 
+    def factor(self, energy):
+        return jnp.exp(-self.tau * self.g(energy))
+
 
 class Highecut(MultiplicativeComponent):
     r"""
@@ -223,13 +238,18 @@ class Tbpcf(MultiplicativeComponent):
 
     def __init__(self):
         table = Table.read(table_manager.fetch("xsect_tbabs_wilm.fits"))
-        self.energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
-        self.sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
+        self._energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
+        self._sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
         self.nh = nnx.Param(1.0)
         self.f = nnx.Param(0.2)
 
     def factor(self, energy):
-        sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
+        sigma = jnp.exp(
+            jnp.interp(
+                energy, self._energy, jnp.log(self._sigma), left="extrapolate", right="extrapolate"
+            )
+        )
+
         return self.f * jnp.exp(-self.nh * sigma) + (1.0 - self.f)
 
 

jaxspec/util/integrate.py

@@ -9,11 +9,13 @@ It mainly relies on tanh-sinh (or double exponential) quadrature to perform the
     * [Tanh-sinh quadrature](https://en.wikipedia.org/wiki/Tanh-sinh_quadrature) from Wikipedia
 """
 
+from collections.abc import Callable
+
 import jax
 import jax.numpy as jnp
+
 from jax import Array
 from jax.scipy.integrate import trapezoid
-from typing import Callable
 
 
 def interval_weights(a: float, b: float, n: int) -> tuple[Array, Array, Array]:
@@ -93,9 +95,15 @@ def integrate_interval(integrand: Callable, n: int = 51) -> Callable:
         primal_out = f(a, b, *args)
 
         # Partial derivatives along other parameters
-        jac = trapezoid(jnp.nan_to_num(jnp.asarray(jax.jacfwd(lambda args: integrand(x, *args))(args)) * dx), x=t, axis=-1)
-
-        tangent_out = -integrand(a, *args) * a_dot + integrand(b, *args) * b_dot + jac @ jnp.asarray(args_dot)
+        jac = trapezoid(
+            jnp.nan_to_num(jnp.asarray(jax.jacfwd(lambda args: integrand(x, *args))(args)) * dx),
+            x=t,
+            axis=-1,
+        )
+
+        tangent_out = (
+            -integrand(a, *args) * a_dot + integrand(b, *args) * b_dot + jac @ jnp.asarray(args_dot)
+        )
         return primal_out, tangent_out
 
     return f
@@ -142,7 +150,11 @@ def integrate_positive(integrand: Callable, n: int = 51) -> Callable:
         primal_out = f(*args)
 
         # Partial derivatives along other parameters
-        jac = trapezoid(jnp.nan_to_num(jnp.asarray(jax.jacfwd(lambda args: integrand(x, *args))(args)) * dx), x=t, axis=-1)
+        jac = trapezoid(
+            jnp.nan_to_num(jnp.asarray(jax.jacfwd(lambda args: integrand(x, *args))(args)) * dx),
+            x=t,
+            axis=-1,
+        )
 
         tangent_out = jac @ jnp.asarray(args_dot)
         return primal_out, tangent_out

{jaxspec-0.2.2.dist-info → jaxspec-0.3.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.2.2
+Version: 0.3.0
 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,28 +8,28 @@ Author-email: sdupourque <sdupourque@irap.omp.eu>
 License-Expression: MIT
 License-File: LICENSE.md
 Requires-Python: <3.13,>=3.10
-Requires-Dist: arviz<0.21.0,>=0.17.1
-Requires-Dist: astropy<7,>=6.0.0
+Requires-Dist: arviz<0.23.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
 Requires-Dist: cmasher<2,>=1.6.3
-Requires-Dist: flax<0.11,>=0.10.3
+Requires-Dist: flax>0.10.5
 Requires-Dist: interpax<0.4,>=0.3.5
 Requires-Dist: jax<0.6,>=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<0.20,>=0.15
+Requires-Dist: mendeleev<1.2,>=0.15
 Requires-Dist: networkx~=3.1
-Requires-Dist: numpy<2.0.0
-Requires-Dist: numpyro<0.18,>=0.17.0
+Requires-Dist: numpy<3.0.0
+Requires-Dist: numpyro<0.19,>=0.17.0
 Requires-Dist: optimistix<0.0.11,>=0.0.10
 Requires-Dist: pandas<3,>=2.2.0
 Requires-Dist: pooch<2,>=1.8.2
 Requires-Dist: scipy<1.15
 Requires-Dist: seaborn<0.14,>=0.13.1
 Requires-Dist: simpleeval<1.1.0,>=0.9.13
-Requires-Dist: sparse<0.16,>=0.15.4
+Requires-Dist: sparse>0.15
 Requires-Dist: tinygp<0.4,>=0.3.0
 Requires-Dist: watermark<3,>=2.4.3
 Description-Content-Type: text/markdown
@@ -44,7 +44,7 @@ Description-Content-Type: text/markdown
 
 
 [![PyPI - Version](https://img.shields.io/pypi/v/jaxspec?style=for-the-badge&logo=pypi&color=rgb(37%2C%20150%2C%20190))](https://pypi.org/project/jaxspec/)
-[![Python package](https://img.shields.io/pypi/pyversions/jaxspec?style=for-the-badge)](https://pypi.org/project/jaxspec/)
+![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Frenecotyfanboy%2Fjaxspec%2Frefs%2Fheads%2Fmain%2Fpyproject.toml&style=for-the-badge)
 [![Read the Docs](https://img.shields.io/readthedocs/jaxspec?style=for-the-badge)](https://jaxspec.readthedocs.io/en/latest/)
 [![Codecov](https://img.shields.io/codecov/c/github/renecotyfanboy/jaxspec?style=for-the-badge)](https://app.codecov.io/gh/renecotyfanboy/jaxspec)
 [![Slack](https://img.shields.io/badge/Slack-4A154B?style=for-the-badge&logo=slack&logoColor=white)](https://join.slack.com/t/jaxspec/shared_invite/zt-2cuxkdl2f-t0EEAKP~HBEHKvIUZJL2sg)

jaxspec-0.3.0.dist-info/RECORD

@@ -0,0 +1,42 @@
+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=_qwDSsThI7FOAR6nMaJltGWlKO5Sz2wc1EQ73Y0Ghho,26013
+jaxspec/data/__init__.py,sha256=aantcYKC9kZFvaE-V2SIwSuLhIld17Kjrd9CIUu___Y,415
+jaxspec/data/instrument.py,sha256=RDiG_LkucvnF2XE_ghTFME6d_2YirgQUcEY0gEle6dk,4775
+jaxspec/data/obsconf.py,sha256=G0RwNshvbDQzw_ba8Y8NdI-cRsgEj-OlSNdeYCANqVM,10484
+jaxspec/data/observation.py,sha256=OHXfs7ApC8JAiG6h1teEmXu3iaNyhrqI5BCwl-qFvoM,7820
+jaxspec/data/ogip.py,sha256=eMmBuROW4eMRxRHkPPyGHf933e0IcREqB8WMQFMS2lY,9810
+jaxspec/data/util.py,sha256=4_f6ByGjUEZXTwrB37dCyYaTB1pjF10Z0ho7-4GrQuc,9761
+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
+jaxspec/experimental/intrument_models.py,sha256=vuRw7xypPI9YV-Hv8chVNP4ti24dCGjbtDBNbtoVBZQ,5283
+jaxspec/experimental/nested_sampler.py,sha256=8jCAXQAe2mD5YSNSF0jia_rFWES_MzwRM3FrQQS_x7w,2807
+jaxspec/experimental/tabulated.py,sha256=H0llUiso2KGH4xUzTUSVPy-6I8D3wm707lU_Z1P5uq4,9429
+jaxspec/fit/__init__.py,sha256=OaS0-Hkb3Hd-AkE2o-KWfoWMX0NSCPY-_FP2znHf9l0,153
+jaxspec/fit/_bayesian_model.py,sha256=BeYukXr86Y1kEmSyiv-6QC4M2rM78Kx_MgGecu4ML98,15179
+jaxspec/fit/_build_model.py,sha256=pNZVuVfwOq3Pg23opH7xRv28DsSkQZpvy2Z-1hQSfNs,3219
+jaxspec/fit/_fitter.py,sha256=doBTJqTP5CN1OJhZHVlS3oMVOzPJyH4YqOnGevIIU68,8893
+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/additive.py,sha256=rEONSy7b7lwfXIhuPqtI4y2Yhq55EqrlEtEckEe6TA0,20538
+jaxspec/model/background.py,sha256=VLSrU0YCW9GSHCtaEdcth-sp74aPyEVSizIMFkTpM7M,7759
+jaxspec/model/instrument.py,sha256=FSGgWCQjUfSVOzJ3YfrbQQ2abPlZ4P4ndtLL9Axcl-g,2217
+jaxspec/model/list.py,sha256=uC9rLEEeph10q6shat86WLACVuTSx73RGMl8Ij0jqQY,875
+jaxspec/model/multiplicative.py,sha256=odaOlF0K1KjzUDstPtWAk95ScHoZ7_XveOew6l3tbeU,8337
+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
+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/typing.py,sha256=ZQM_l68qyYnIBZPz_1mKvwPMx64jvVBD8Uj6bx9sHv0,140
+jaxspec-0.3.0.dist-info/METADATA,sha256=92shp3kcwQIbKTSVSD7SU68InowsGVZXST0uJYvRwnQ,4199
+jaxspec-0.3.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+jaxspec-0.3.0.dist-info/entry_points.txt,sha256=4ffU5AImfcEBxgWTqopQll2YffpFldOswXRh16pd0Dc,72
+jaxspec-0.3.0.dist-info/licenses/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
+jaxspec-0.3.0.dist-info/RECORD,,

jaxspec-0.2.2.dist-info/RECORD

@@ -1,34 +0,0 @@
-jaxspec/__init__.py,sha256=Sbn02lX6Y-zNXk17N8dec22c5jeypiS0LkHmGfz7lWA,126
-jaxspec/fit.py,sha256=gO0H8CR76WP7lrG2X-Z3HLMndOSkCulzrp2jll1b4ow,24337
-jaxspec/_fit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxspec/_fit/_build_model.py,sha256=NgLobtysyhCoiRguAazYGRIpG75iW-pG0Ss_hQGuCT4,2019
-jaxspec/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxspec/analysis/_plot.py,sha256=1VbWentCLHR8XzM3jZtGyppbQRFWnChNztrwpZ2reA4,6185
-jaxspec/analysis/compare.py,sha256=g2UFhmR9Zt-7cz5gQFOB6lXuklXB3yTyUvjTypOzoSY,725
-jaxspec/analysis/results.py,sha256=9ybhWMtwXD_f9reTt_UcSUKQaY4Raxm0Fi7QBv7fqow,25448
-jaxspec/data/__init__.py,sha256=aantcYKC9kZFvaE-V2SIwSuLhIld17Kjrd9CIUu___Y,415
-jaxspec/data/instrument.py,sha256=RDiG_LkucvnF2XE_ghTFME6d_2YirgQUcEY0gEle6dk,4775
-jaxspec/data/obsconf.py,sha256=r0deDgzpfKJZFMeiKQYUnkqQEmoDYTzl3eQFjb5xYjo,10263
-jaxspec/data/observation.py,sha256=dfb-hJFpeSPMS52oNmQl39CB9GD_MK1gG9vHlQXhhwU,7741
-jaxspec/data/ogip.py,sha256=57lfUgvMhM4j5tdfnOmwLI0ehx09pG7SvJx6VNIJFHE,9543
-jaxspec/data/util.py,sha256=FnIqrnwm29WsFV3FayeLXpjJDS10FcZ613B3SzGQCTk,9331
-jaxspec/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxspec/model/_graph_util.py,sha256=hPvHYmAxb7P3nyIecaZ7RqWOjwcZ1WvUByt_yNANiaY,4552
-jaxspec/model/abc.py,sha256=sHaFKOGfpD1ocXKlzEMdsmGLluikmBSmX8FAu7duFMQ,15489
-jaxspec/model/additive.py,sha256=2MXttFw7eZ6dZ2yWyNpMGDDFFE0kNo6yAP-G-YtL2Qk,20276
-jaxspec/model/background.py,sha256=INRuuHIYgfu2E1a8Ddqpr_gF1wnhPtOpciWCteZfVVg,7687
-jaxspec/model/list.py,sha256=0RPAoscVz_zM1CWdx_Gd5wfrQWV5Nv4Kd4bSXu2ayUA,860
-jaxspec/model/multiplicative.py,sha256=5LHBAXCV9EQO1BGm1HzAI56dqjVbafBdV8Eq7_qXdbM,7890
-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
-jaxspec/util/abundance.py,sha256=fsC313taIlGzQsZNwbYsJupDWm7ZbqzGhY66Ku394Mw,8546
-jaxspec/util/integrate.py,sha256=_Ax_knpC7d4et2-QFkOUzVtNeQLX1-cwLvm-FRBxYcw,4505
-jaxspec/util/misc.py,sha256=O3qorCL1Y2X1BS2jdd36C1eDHK9QDXTSOr9kj3uqcJo,654
-jaxspec/util/online_storage.py,sha256=vm56RfcbFKpkRVfr0bXO7J9aQxuBq-I_oEgA26YIhCo,2469
-jaxspec/util/typing.py,sha256=ZQM_l68qyYnIBZPz_1mKvwPMx64jvVBD8Uj6bx9sHv0,140
-jaxspec-0.2.2.dist-info/METADATA,sha256=mKdmXpp8G1qug5j2-LEIrQQDbbjvVmleZkY1Pnq0O0c,4111
-jaxspec-0.2.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-jaxspec-0.2.2.dist-info/entry_points.txt,sha256=4ffU5AImfcEBxgWTqopQll2YffpFldOswXRh16pd0Dc,72
-jaxspec-0.2.2.dist-info/licenses/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
-jaxspec-0.2.2.dist-info/RECORD,,