jaxspec 0.3.0__tar.gz → 0.3.1__tar.gz

{jaxspec-0.3.0 → jaxspec-0.3.1}/.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.1}/PKG-INFO

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

@@ -1,17 +1,17 @@
 [project]
 name = "jaxspec"
-version = "0.3.0"
+version = "0.3.1"
 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",
@@ -23,7 +23,7 @@ dependencies = [
     "seaborn>=0.13.1,<0.14",
     "sparse>0.15",
     "optimistix>=0.0.10,<0.0.11",
-    "scipy<1.15",
+    "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.1}/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.1}/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
@@ -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.1}/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)