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

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/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/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/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/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/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.dist-info → jaxspec-0.3.2.dist-info}/METADATA

@@ -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.dist-info → jaxspec-0.3.2.dist-info}/RECORD

@@ -2,11 +2,11 @@ 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/analysis/results.py,sha256=tIBWmLoX43EY2BXt50ec8A-DqQ98PMd3m-FqTRT4iRE,26073
 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/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/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -16,15 +16,15 @@ jaxspec/experimental/intrument_models.py,sha256=vuRw7xypPI9YV-Hv8chVNP4ti24dCGjb
 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/_bayesian_model.py,sha256=7c2Twgz06QV1S9DdctdVk5YT1v7P-ln100bWXAvv7Go,15179
 jaxspec/fit/_build_model.py,sha256=pNZVuVfwOq3Pg23opH7xRv28DsSkQZpvy2Z-1hQSfNs,3219
-jaxspec/fit/_fitter.py,sha256=doBTJqTP5CN1OJhZHVlS3oMVOzPJyH4YqOnGevIIU68,8893
+jaxspec/fit/_fitter.py,sha256=92gd1P7CNIqusGb64x_DpBcb0KcoGyfvSDiEnRbfqP0,9709
 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/instrument.py,sha256=1zLZgHmBZs8RLKTMT3Wu4bCx6JnxBUjhRIpYG2rLaZM,2947
 jaxspec/model/list.py,sha256=uC9rLEEeph10q6shat86WLACVuTSx73RGMl8Ij0jqQY,875
 jaxspec/model/multiplicative.py,sha256=odaOlF0K1KjzUDstPtWAk95ScHoZ7_XveOew6l3tbeU,8337
 jaxspec/scripts/__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=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.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,,