jaxspec 0.3.1__tar.gz → 0.3.2__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.3.1
+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/
@@ -23,7 +23,7 @@ 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.11,>=0.0.10
+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

{jaxspec-0.3.1 → 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.1 → 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.1 → 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.1 → jaxspec-0.3.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxspec"
-version = "0.3.1"
+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"
@@ -22,7 +22,7 @@ dependencies = [
     "tinygp>=0.3.0,<0.4",
     "seaborn>=0.13.1,<0.14",
     "sparse>0.15",
-    "optimistix>=0.0.10,<0.0.11",
+    "optimistix>=0.0.10,<0.0.12",
     "scipy<1.16",
     "mendeleev>=0.15,<1.2",
     "jaxns>=2.6.7,<3",

{jaxspec-0.3.1 → 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.1 → 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.1 → jaxspec-0.3.2}/src/jaxspec/fit/_bayesian_model.py

@@ -135,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,
                     )
 

{jaxspec-0.3.1 → 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.1 → 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.1 → 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__}")