jaxspec 0.3.0__tar.gz → 0.3.2__tar.gz

{jaxspec-0.3.0 → jaxspec-0.3.2}/.gitignore

@@ -197,4 +197,6 @@ cython_debug/
 .Trashes
 ehthumbs.db
 Thumbs.db
+
+_old/
 # End of https://www.toptal.com/developers/gitignore/api/python,jupyternotebooks

{jaxspec-0.3.0 → jaxspec-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.3.0
+Version: 0.3.2
 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/
@@ -15,18 +15,18 @@ Requires-Dist: chainconsumer<2,>=1.1.2
 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.6,>=0.5.0
+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.19,>=0.17.0
-Requires-Dist: optimistix<0.0.11,>=0.0.10
+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.15
+Requires-Dist: scipy<1.16
 Requires-Dist: seaborn<0.14,>=0.13.1
 Requires-Dist: simpleeval<1.1.0,>=0.9.13
 Requires-Dist: sparse>0.15

{jaxspec-0.3.0 → jaxspec-0.3.2}/docs/examples/index.md

@@ -54,5 +54,11 @@
 
     Use `jaxspec` likelihood in other packages
 
+-   [__Cross-calibration between instruments__](calibration.ipynb)
+
+    ---
+
+    Add extra gains and shifts to account for the miscalibrations
+
 
 </div>

{jaxspec-0.3.0 → jaxspec-0.3.2}/docs/references/fitting.md

@@ -13,3 +13,11 @@
       show_root_full_path: false
       show_root_toc_entry: true
       heading_level: 3
+
+::: jaxspec.fit.VIFitter
+    options:
+      show_root_heading: true
+      show_root_full_path: false
+      show_root_toc_entry: true
+      heading_level: 3
+

jaxspec-0.3.2/docs/references/instrument.md

@@ -0,0 +1,4 @@
+::: jaxspec.model.instrument
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false

{jaxspec-0.3.0 → jaxspec-0.3.2}/mkdocs.yml

@@ -18,14 +18,13 @@ nav:
     - Add a background : examples/background.md
     - Good practices for MCMC: examples/work_with_arviz.ipynb
     - Interface with other frameworks : examples/external_samplers.ipynb
-  #- Models:
-  #    - models/index.md
-  #    - APEC: models/apec.md
+    - Cross calibration : examples/calibration.ipynb
   - API Reference:
     - Spectral model base: references/model.md
     - Additive models: references/additive.md
     - Multiplicative models: references/multiplicative.md
     - Background models: references/background.md
+    - Instrument models: references/instrument.md
     - Data containers: references/data.md
     - Fitting: references/fitting.md
     - Result containers: references/results.md

{jaxspec-0.3.0 → jaxspec-0.3.2}/pyproject.toml

@@ -1,17 +1,17 @@
 [project]
 name = "jaxspec"
-version = "0.3.0"
+version = "0.3.2"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = [{ name = "sdupourque", email = "sdupourque@irap.omp.eu" }]
 requires-python = ">=3.10,<3.13"
 readme = "README.md"
 license = "MIT"
 dependencies = [
-    "jax>=0.5.0,<0.6",
+    "jax>=0.5.0,<0.7",
     "numpy<3.0.0",
     "pandas>=2.2.0,<3",
     "astropy>=6.0.0,<8",
-    "numpyro>=0.17.0,<0.19",
+    "numpyro>=0.17.0,<0.20",
     "networkx~=3.1",
     "matplotlib>=3.8.0,<4",
     "arviz>=0.17.1,<0.23.0",
@@ -22,8 +22,8 @@ dependencies = [
     "tinygp>=0.3.0,<0.4",
     "seaborn>=0.13.1,<0.14",
     "sparse>0.15",
-    "optimistix>=0.0.10,<0.0.11",
-    "scipy<1.15",
+    "optimistix>=0.0.10,<0.0.12",
+    "scipy<1.16",
     "mendeleev>=0.15,<1.2",
     "jaxns>=2.6.7,<3",
     "pooch>=1.8.2,<2",
@@ -58,7 +58,7 @@ test = [
 ]
 dev = [
     "pre-commit>=3.5,<5.0",
-    "ruff>=0.2.1,<0.10.0",
+    "ruff>=0.2.1,<0.15.0",
     "jupyterlab>=4.0.7,<5",
     "notebook>=7.0.6,<8",
     "ipywidgets>=8.1.1,<9",

{jaxspec-0.3.0 → jaxspec-0.3.2}/src/jaxspec/analysis/results.py

@@ -122,7 +122,9 @@ class FitResult:
 
         samples_shape = (len(posterior.coords["chain"]), len(posterior.coords["draw"]))
 
-        total_shape = tuple(posterior.sizes[d] for d in posterior.coords)
+        total_shape = tuple(
+            posterior.sizes[d] for d in posterior.coords if not (("obs" in d) or ("bkg" in d))
+        )
 
         posterior = {key: posterior[key].data for key in posterior.data_vars}
 

{jaxspec-0.3.0 → jaxspec-0.3.2}/src/jaxspec/data/obsconf.py

@@ -163,10 +163,8 @@ class ObsConfiguration(xr.Dataset):
 
         if observation.folded_background is not None:
             folded_background = observation.folded_background.data[row_idx]
-            folded_background_unscaled = observation.folded_background_unscaled.data[row_idx]
         else:
             folded_background = np.zeros_like(folded_counts)
-            folded_background_unscaled = np.zeros_like(folded_counts)
 
         data_dict = {
             "transfer_matrix": (
@@ -208,14 +206,6 @@ class ObsConfiguration(xr.Dataset):
                     "unit": "photons",
                 },
             ),
-            "folded_background_unscaled": (
-                ["folded_channel"],
-                folded_background_unscaled,
-                {
-                    "description": "To be done",
-                    "unit": "photons",
-                },
-            ),
         }
 
         return cls(

{jaxspec-0.3.0 → jaxspec-0.3.2}/src/jaxspec/data/observation.py

@@ -46,16 +46,14 @@ class Observation(xr.Dataset):
         quality,
         exposure,
         background=None,
-        background_unscaled=None,
         backratio=1.0,
         attributes: dict | None = None,
     ):
         if attributes is None:
             attributes = {}
 
-        if background is None or background_unscaled is None:
+        if background is None:
             background = np.zeros_like(counts, dtype=np.int64)
-            background_unscaled = np.zeros_like(counts, dtype=np.int64)
 
         data_dict = {
             "counts": (
@@ -86,7 +84,9 @@ class Observation(xr.Dataset):
             ),
             "folded_backratio": (
                 ["folded_channel"],
-                np.asarray(np.ma.filled(grouping @ backratio), dtype=float),
+                np.asarray(
+                    np.ma.filled(grouping @ backratio) / grouping.sum(axis=1).todense(), dtype=float
+                ),
                 {"description": "Background scaling after grouping"},
             ),
             "background": (
@@ -94,11 +94,6 @@ class Observation(xr.Dataset):
                 np.asarray(background, dtype=np.int64),
                 {"description": "Background counts", "unit": "photons"},
             ),
-            "folded_background_unscaled": (
-                ["folded_channel"],
-                np.asarray(np.ma.filled(grouping @ background_unscaled), dtype=np.int64),
-                {"description": "Background counts", "unit": "photons"},
-            ),
             "folded_background": (
                 ["folded_channel"],
                 np.asarray(np.ma.filled(grouping @ background), dtype=np.float64),
@@ -147,8 +142,7 @@ class Observation(xr.Dataset):
             pha.quality,
             pha.exposure,
             backratio=backratio,
-            background=bkg.counts * backratio if bkg is not None else None,
-            background_unscaled=bkg.counts if bkg is not None else None,
+            background=bkg.counts if bkg is not None else None,
             attributes=metadata,
         )
 

{jaxspec-0.3.0 → jaxspec-0.3.2}/src/jaxspec/fit/_bayesian_model.py

@@ -10,9 +10,8 @@ import matplotlib.pyplot as plt
 import numpyro
 
 from flax import nnx
-from jax.experimental import mesh_utils
 from jax.random import PRNGKey
-from jax.sharding import PositionalSharding
+from jax.sharding import NamedSharding, PartitionSpec
 from numpyro.distributions import Poisson, TransformedDistribution
 from numpyro.infer import Predictive
 from numpyro.infer.inspect import get_model_relations
@@ -136,7 +135,7 @@ class BayesianModel(nnx.Module):
                 with numpyro.plate("obs_plate/~/" + name, len(observation.folded_counts)):
                     numpyro.sample(
                         "obs/~/" + name,
-                        Poisson(obs_countrate + bkg_countrate / observation.folded_backratio.data),
+                        Poisson(obs_countrate + bkg_countrate * observation.folded_backratio.data),
                         obs=observation.folded_counts.data if observed else None,
                     )
 
@@ -244,7 +243,7 @@ class BayesianModel(nnx.Module):
         return log_posterior_prob
 
     @cached_property
-    def _parameter_names(self) -> list[str]:
+    def parameter_names(self) -> list[str]:
         """
         A list of parameter names for the model.
         """
@@ -269,7 +268,7 @@ class BayesianModel(nnx.Module):
         """
         input_params = {}
 
-        for index, key in enumerate(self._parameter_names):
+        for index, key in enumerate(self.parameter_names):
             input_params[key] = theta[index]
 
         return input_params
@@ -279,9 +278,9 @@ class BayesianModel(nnx.Module):
         Convert a dictionary of parameters to an array of parameters.
         """
 
-        theta = jnp.zeros(len(self._parameter_names))
+        theta = jnp.zeros(len(self.parameter_names))
 
-        for index, key in enumerate(self._parameter_names):
+        for index, key in enumerate(self.parameter_names):
             theta = theta.at[index].set(dict_of_params[key])
 
         return theta
@@ -298,7 +297,7 @@ class BayesianModel(nnx.Module):
         @jax.jit
         def prior_sample(key):
             return Predictive(
-                self.numpyro_model, return_sites=self._parameter_names, num_samples=num_samples
+                self.numpyro_model, return_sites=self.parameter_names, num_samples=num_samples
             )(key, observed=False)
 
         return prior_sample(key)
@@ -324,7 +323,8 @@ class BayesianModel(nnx.Module):
         """
         key_prior, key_posterior = jax.random.split(key, 2)
         n_devices = len(jax.local_devices())
-        sharding = PositionalSharding(mesh_utils.create_device_mesh((n_devices,)))
+        mesh = jax.make_mesh((n_devices,), ("batch",))
+        sharding = NamedSharding(mesh, PartitionSpec("batch"))
 
         # Sample from prior and correct if the number of samples is not a multiple of the number of devices
         if num_samples % n_devices != 0:

{jaxspec-0.3.0 → jaxspec-0.3.2}/src/jaxspec/fit/_fitter.py

@@ -215,13 +215,29 @@ class VIFitter(BayesianModelFitter):
         self,
         rng_key: int = 0,
         num_steps: int = 10_000,
-        optimizer=numpyro.optim.Adam(step_size=0.0005),
-        loss=Trace_ELBO(),
+        optimizer: numpyro.optim._NumPyroOptim = numpyro.optim.Adam(step_size=0.0005),
+        loss: numpyro.infer.elbo.ELBO = Trace_ELBO(),
         num_samples: int = 1000,
-        guide=None,
+        guide: numpyro.infer.autoguide.AutoGuide | None = None,
         use_transformed_model: bool = True,
         plot_diagnostics: bool = False,
     ) -> FitResult:
+        """
+        Fit the model to the data using a variational inference approach from numpyro.
+
+        Parameters:
+            rng_key: the random key used to initialize the sampler.
+            num_steps: the number of steps for VI.
+            optimizer: the optimizer to use.
+            num_samples: the number of samples to draw.
+            loss: the loss function to use.
+            guide: the guide to use.
+            use_transformed_model: whether to use the transformed model to build the InferenceData.
+            plot_diagnostics: plot the loss during VI.
+
+        Returns:
+            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
+        """
         bayesian_model = (
             self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
         )

{jaxspec-0.3.0 → jaxspec-0.3.2}/src/jaxspec/model/instrument.py

@@ -8,13 +8,26 @@ from numpyro.distributions import Distribution
 
 
 class GainModel(ABC, nnx.Module):
+    """
+    Generic class for a gain model
+    """
+
     @abstractmethod
     def numpyro_model(self, observation_name: str):
         pass
 
 
 class ConstantGain(GainModel):
+    """
+    A constant gain model
+    """
+
     def __init__(self, prior_distribution: Distribution):
+        """
+        Parameters:
+            prior_distribution: the prior distribution for the gain value.
+        """
+
         self.prior_distribution = prior_distribution
 
     def numpyro_model(self, observation_name: str):
@@ -27,13 +40,25 @@ class ConstantGain(GainModel):
 
 
 class ShiftModel(ABC, nnx.Module):
+    """
+    Generic class for a shift model
+    """
+
     @abstractmethod
     def numpyro_model(self, observation_name: str):
         pass
 
 
 class ConstantShift(ShiftModel):
+    """
+    A constant shift model
+    """
+
     def __init__(self, prior_distribution: Distribution):
+        """
+        Parameters:
+            prior_distribution: the prior distribution for the shift value.
+        """
         self.prior_distribution = prior_distribution
 
     def numpyro_model(self, observation_name: str):
@@ -52,6 +77,15 @@ class InstrumentModel(nnx.Module):
         gain_model: GainModel | None = None,
         shift_model: ShiftModel | None = None,
     ):
+        """
+        Encapsulate an instrument model, build as a combination of a shift and gain model.
+
+        Parameters:
+            reference_observation_name : The observation to use as a reference
+            gain_model : The gain model
+            shift_model : The shift model
+        """
+
         self.reference = reference_observation_name
         self.gain_model = gain_model
         self.shift_model = shift_model

{jaxspec-0.3.0 → jaxspec-0.3.2}/tests/test_background.py

@@ -14,9 +14,9 @@ spectral_model_background = Powerlaw() + Blackbodyrad()
 
 prior_background = {
     "powerlaw_1_alpha": dist.Uniform(0, 5),
-    "powerlaw_1_norm": dist.LogUniform(1e-8, 1e-3),
+    "powerlaw_1_norm": dist.LogUniform(1e-7, 1e-3),
     "blackbodyrad_1_kT": dist.Uniform(0, 5),
-    "blackbodyrad_1_norm": dist.LogUniform(1e-6, 1e-1),
+    "blackbodyrad_1_norm": dist.LogUniform(1e-5, 1e-1),
 }
 
 
@@ -36,7 +36,7 @@ def test_background_model(obs_model_prior, bkg_model):
     obs_list, model, prior = obs_model_prior
     forward = MCMCFitter(model, prior, obs_list[0], background_model=bkg_model)
     result = forward.fit(
-        num_chains=4, num_warmup=100, num_samples=100, mcmc_kwargs={"progress_bar": False}
+        num_chains=4, num_warmup=1000, num_samples=1000, mcmc_kwargs={"progress_bar": False}
     )
     result.plot_ppc(title=f"Test {bkg_model.__class__.__name__}")
 

{jaxspec-0.3.0 → jaxspec-0.3.2}/tests/test_fakeit.py

@@ -63,10 +63,10 @@ def model():
 
 @pytest.fixture
 def sharded_parameters(unidimensional_parameters):
-    from jax.experimental import mesh_utils
-    from jax.sharding import PositionalSharding
+    from jax.sharding import NamedSharding, PartitionSpec
 
-    sharding = PositionalSharding(mesh_utils.create_device_mesh((4,)))
+    mesh = jax.make_mesh((4,), ("batch",))
+    sharding = NamedSharding(mesh, PartitionSpec("batch"))
 
     return jax.device_put(unidimensional_parameters, sharding)