jaxspec 0.3.0__py3-none-any.whl → 0.3.1__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/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
@@ -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.dist-info → jaxspec-0.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.3.0
+Version: 0.3.1
 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: numpyro<0.20,>=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: 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.1.dist-info}/RECORD

@@ -2,7 +2,7 @@ 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
@@ -16,7 +16,7 @@ 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=jSCzAzoAhsmUX7mKUikbUR9A1ZNIaY6rdPOxq6OZSU0,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
@@ -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.1.dist-info/METADATA,sha256=8i1cuzZY4iwIjWEhIPXBiC-Z8Y4Vv27omIMwTKnoPwo,4199
+jaxspec-0.3.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+jaxspec-0.3.1.dist-info/entry_points.txt,sha256=4ffU5AImfcEBxgWTqopQll2YffpFldOswXRh16pd0Dc,72
+jaxspec-0.3.1.dist-info/licenses/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
+jaxspec-0.3.1.dist-info/RECORD,,