jaxspec 0.2.2.dev0__tar.gz → 0.3.0__tar.gz

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/.github/workflows/test-and-coverage.yml

@@ -42,7 +42,7 @@ jobs:
 
     - name: "Upload coverage to Codecov"
       if: steps.filter.outputs.src == 'true'
-      uses: codecov/codecov-action@v4
+      uses: codecov/codecov-action@v5
       with:
         token: ${{ secrets.CODECOV_TOKEN }}
         fail_ci_if_error: true

jaxspec-0.3.0/.python-version

@@ -0,0 +1 @@
+>=3.10, <3.13

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/.readthedocs.yaml

@@ -20,7 +20,7 @@ build:
       - uv pip install -r pyproject.toml
       # Using insiders versions of mkdocs-material & mkdocstrings
       - uv pip uninstall mkdocs-material # mkdocstrings mkdocstrings-python
-      - uv pip install git+https://$GH_TOKEN@github.com/squidfunk/mkdocs-material-insiders.git@9.5.36-insiders-4.53.13
+      # - uv pip install git+https://$GH_TOKEN@github.com/squidfunk/mkdocs-material-insiders.git@9.5.36-insiders-4.53.13
       - uv pip install mkdocstrings mkdocstrings-python
       - uv pip install mkdocs-autorefs
       - uv pip install mkdocs-jupyter # This is bugged, I enforced it manually, let's see if it works

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.2.2.dev0
+Version: 0.3.0
 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/
@@ -8,28 +8,28 @@ Author-email: sdupourque <sdupourque@irap.omp.eu>
 License-Expression: MIT
 License-File: LICENSE.md
 Requires-Python: <3.13,>=3.10
-Requires-Dist: arviz<0.21.0,>=0.17.1
-Requires-Dist: astropy<7,>=6.0.0
+Requires-Dist: arviz<0.23.0,>=0.17.1
+Requires-Dist: astropy<8,>=6.0.0
 Requires-Dist: catppuccin<3,>=2.3.4
 Requires-Dist: chainconsumer<2,>=1.1.2
 Requires-Dist: cmasher<2,>=1.6.3
-Requires-Dist: flax<0.11,>=0.10.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: jaxns<3,>=2.6.7
 Requires-Dist: jaxopt<0.9,>=0.8.3
 Requires-Dist: matplotlib<4,>=3.8.0
-Requires-Dist: mendeleev<0.20,>=0.15
+Requires-Dist: mendeleev<1.2,>=0.15
 Requires-Dist: networkx~=3.1
-Requires-Dist: numpy<2.0.0
-Requires-Dist: numpyro<0.18,>=0.17.0
+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: pandas<3,>=2.2.0
 Requires-Dist: pooch<2,>=1.8.2
 Requires-Dist: scipy<1.15
 Requires-Dist: seaborn<0.14,>=0.13.1
 Requires-Dist: simpleeval<1.1.0,>=0.9.13
-Requires-Dist: sparse<0.16,>=0.15.4
+Requires-Dist: sparse>0.15
 Requires-Dist: tinygp<0.4,>=0.3.0
 Requires-Dist: watermark<3,>=2.4.3
 Description-Content-Type: text/markdown
@@ -44,7 +44,7 @@ Description-Content-Type: text/markdown
 
 
 [![PyPI - Version](https://img.shields.io/pypi/v/jaxspec?style=for-the-badge&logo=pypi&color=rgb(37%2C%20150%2C%20190))](https://pypi.org/project/jaxspec/)
-[![Python package](https://img.shields.io/pypi/pyversions/jaxspec?style=for-the-badge)](https://pypi.org/project/jaxspec/)
+![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Frenecotyfanboy%2Fjaxspec%2Frefs%2Fheads%2Fmain%2Fpyproject.toml&style=for-the-badge)
 [![Read the Docs](https://img.shields.io/readthedocs/jaxspec?style=for-the-badge)](https://jaxspec.readthedocs.io/en/latest/)
 [![Codecov](https://img.shields.io/codecov/c/github/renecotyfanboy/jaxspec?style=for-the-badge)](https://app.codecov.io/gh/renecotyfanboy/jaxspec)
 [![Slack](https://img.shields.io/badge/Slack-4A154B?style=for-the-badge&logo=slack&logoColor=white)](https://join.slack.com/t/jaxspec/shared_invite/zt-2cuxkdl2f-t0EEAKP~HBEHKvIUZJL2sg)

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/README.md

@@ -8,7 +8,7 @@
 
 
 [![PyPI - Version](https://img.shields.io/pypi/v/jaxspec?style=for-the-badge&logo=pypi&color=rgb(37%2C%20150%2C%20190))](https://pypi.org/project/jaxspec/)
-[![Python package](https://img.shields.io/pypi/pyversions/jaxspec?style=for-the-badge)](https://pypi.org/project/jaxspec/)
+![Python Version from PEP 621 TOML](https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Frenecotyfanboy%2Fjaxspec%2Frefs%2Fheads%2Fmain%2Fpyproject.toml&style=for-the-badge)
 [![Read the Docs](https://img.shields.io/readthedocs/jaxspec?style=for-the-badge)](https://jaxspec.readthedocs.io/en/latest/)
 [![Codecov](https://img.shields.io/codecov/c/github/renecotyfanboy/jaxspec?style=for-the-badge)](https://app.codecov.io/gh/renecotyfanboy/jaxspec)
 [![Slack](https://img.shields.io/badge/Slack-4A154B?style=for-the-badge&logo=slack&logoColor=white)](https://join.slack.com/t/jaxspec/shared_invite/zt-2cuxkdl2f-t0EEAKP~HBEHKvIUZJL2sg)

jaxspec-0.3.0/codecov.yml

@@ -0,0 +1,2 @@
+ignore:
+  - "src/jaxspec/experimental"

jaxspec-0.3.0/docs/css/extra.css

@@ -0,0 +1,13 @@
+a.card-link center {
+    width: fit-content;
+    display: block;
+    margin: 0 auto;
+}
+
+a.card-link .card-title {
+    font-size: 1.5em;
+}
+
+a.card-link {
+    color: initial; // Make cards have white text until hovered.
+}

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/docs/faq/index.md

@@ -1,5 +1,10 @@
 # Frequently asked questions
 
+## How can I load multiple spectra to fit ?
+
+Simply pass a list or dictionnary of [`ObsConfiguration`][jaxspec.data.ObsConfiguration] objects when building your
+fitter object.
+
 ## Why should I use `jaxspec` over `xspec` or associated ?
 
 We have taken great care to make `jaxspec` as easy to use as possible. It can be installed with `pip install jaxspec`

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/docs/index.md

@@ -11,13 +11,21 @@ by combining components, and fit it to one or multiple observed spectra using Ba
 
 ## Getting started
 
-<div class="grid">
-
-  <a href="frontpage/installation/" class="card" style="font-size: 1.2em;">🛠️ Installation</a>
-  <a href="examples/fitting_example/" class="card" style="font-size: 1.2em;">🚀 Quickstart</a>
-  <a href="examples/" class="card" style="font-size: 1.2em;">📚 Examples</a>
-  <a href="contribute/" class="card" style="font-size: 1.2em;">🤝 Contribute</a>
 
+<div class="grid cards" markdown>
+
+-   <a class="card-link" href="frontpage/installation/" target="_blank" rel="noreferrer">
+        <span class="card-title center">🛠️ Installation</span>
+    </a>
+-   <a class="card-link" href="examples/fitting_example/" target="_blank" rel="noreferrer">
+    <span class="card-title center">🚀 Quickstart</span>
+    </a>
+-   <a class="card-link" href="examples/" target="_blank" rel="noreferrer">
+    <span class="card-title center">📚 Examples</span>
+    </a>
+-   <a class="card-link" href="contribute/" target="_blank" rel="noreferrer">
+    <span class="card-title center">🤝 Contribute</span>
+    </a>
 </div>
 
 ## How does it work?
@@ -33,7 +41,7 @@ by combining components, and fit it to one or multiple observed spectra using Ba
 Basically, the use of `JAX` as backend allows our models to be differentiable and computable on accelerators, and `numpyro`
 gives access to appropriate samplers such as the No U-Turn Sampler (NUTS) and Hamiltonian Monte Carlo (HMC).
 
-## Citation 
+## Citation
 
 If you use `jaxspec` in your research, please consider citing the following article
 

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

@@ -7,24 +7,9 @@
 
 ## Fitter classes
 
-::: jaxspec.fit.BayesianModelFitter
-    options:
-      show_root_heading: true
-      show_root_full_path: false
-      show_root_toc_entry: true
-      heading_level: 3
-
-
 ::: jaxspec.fit.MCMCFitter
     options:
       show_root_heading: true
       show_root_full_path: false
       show_root_toc_entry: true
       heading_level: 3
-
-::: jaxspec.fit.NSFitter
-    options:
-      show_root_heading: true
-      show_root_full_path: false
-      show_root_toc_entry: true
-      heading_level: 3

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/mkdocs.yml

@@ -87,7 +87,7 @@ theme:
 plugins:
   - search
   - autorefs
-  - typeset
+#  - typeset
   - mkdocs-jupyter:
       include_source: True
       ignore_h1_titles: True
@@ -172,6 +172,7 @@ extra_css:
   - css/material.css
   - css/mkdocstrings.css
   - css/xarray.css
+  - css/extra.css
 
 extra_javascript:
   - javascripts/mathjax.js

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxspec"
-version = "0.2.2dev"
+version = "0.3.0"
 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"
@@ -8,29 +8,29 @@ readme = "README.md"
 license = "MIT"
 dependencies = [
     "jax>=0.5.0,<0.6",
-    "numpy<2.0.0",
+    "numpy<3.0.0",
     "pandas>=2.2.0,<3",
-    "astropy>=6.0.0,<7",
-    "numpyro>=0.17.0,<0.18",
+    "astropy>=6.0.0,<8",
+    "numpyro>=0.17.0,<0.19",
     "networkx~=3.1",
     "matplotlib>=3.8.0,<4",
-    "arviz>=0.17.1,<0.21.0",
+    "arviz>=0.17.1,<0.23.0",
     "chainconsumer>=1.1.2,<2",
     "simpleeval>=0.9.13,<1.1.0",
     "cmasher>=1.6.3,<2",
     "jaxopt>=0.8.3,<0.9",
     "tinygp>=0.3.0,<0.4",
     "seaborn>=0.13.1,<0.14",
-    "sparse>=0.15.4,<0.16",
+    "sparse>0.15",
     "optimistix>=0.0.10,<0.0.11",
     "scipy<1.15",
-    "mendeleev>=0.15,<0.20",
+    "mendeleev>=0.15,<1.2",
     "jaxns>=2.6.7,<3",
     "pooch>=1.8.2,<2",
     "interpax>=0.3.5,<0.4",
     "watermark>=2.4.3,<3",
     "catppuccin>=2.3.4,<3",
-    "flax>=0.10.3,<0.11",
+    "flax>0.10.5",
 ]
 
 [project.urls]

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/src/jaxspec/analysis/_plot.py

@@ -59,8 +59,8 @@ def _plot_poisson_data_with_error(
         y,
         xerr=np.abs(x_bins - np.sqrt(x_bins[0] * x_bins[1])),
         yerr=[
-            y - y_low,
-            y_high - y,
+            np.maximum(y - y_low, 0),
+            np.maximum(y_high - y, 0),
         ],
         color=color,
         linestyle=linestyle,
@@ -149,13 +149,13 @@ def _compute_effective_area(
     mid_bins_arf = obsconf.in_energies.mean(axis=0) * u.keV
     mid_bins_arf = mid_bins_arf.to(x_unit, u.spectral())
     e_grid = np.linspace(*xbins, 10)
-    interpolated_arf = np.interp(e_grid, mid_bins_arf, obsconf.area)
+    interpolated_arf = np.interp(e_grid.value, mid_bins_arf.value, obsconf.area)
     integrated_arf = (
-        trapezoid(interpolated_arf, x=e_grid, axis=0)
+        trapezoid(interpolated_arf, x=e_grid.value, axis=0)
         / (
             np.abs(
                 xbins[1] - xbins[0]
-            )  # Must fold in abs because some units reverse the ordering of the bins
+            ).value  # Must fold in abs because some units reverse the ordering of the bins
         )
         * u.cm**2
     )

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

@@ -42,6 +42,11 @@ V = TypeVar("V")
 T = TypeVar("T")
 
 
+def auto_in_axes(pytree, axis=0):
+    """Return a pytree of 0/None depending on whether the leaf is batched."""
+    return jax.tree.map(lambda x: axis if (hasattr(x, "ndim") and x.ndim > 0) else None, pytree)
+
+
 class FitResult:
     """
     Container for the result of a fit using any ModelFitter class.
@@ -54,17 +59,17 @@ class FitResult:
         inference_data: az.InferenceData,
         background_model: BackgroundModel = None,
     ):
-        self.model = bayesian_fitter.model
+        self.model = bayesian_fitter.spectral_model
         self.bayesian_fitter = bayesian_fitter
         self.inference_data = inference_data
-        self.obsconfs = bayesian_fitter.observation_container
+        self.obsconfs = bayesian_fitter._observation_container
         self.background_model = background_model
 
         # Add the model used in fit to the metadata
         for group in self.inference_data.groups():
             group_name = group.split("/")[-1]
             metadata = getattr(self.inference_data, group_name).attrs
-            metadata["model"] = str(self.model)
+            # metadata["model"] = str(self.model)
             # TODO : Store metadata about observations used in the fitting process
 
     @property
@@ -78,6 +83,7 @@ class FitResult:
     def _ppc_folded_branches(self, obs_id):
         obs = self.obsconfs[obs_id]
 
+        # Slice the parameters corresponding to the current ObsID
         if len(next(iter(self.input_parameters.values())).shape) > 2:
             idx = list(self.obsconfs.keys()).index(obs_id)
             obs_parameters = jax.tree.map(lambda x: x[..., idx], self.input_parameters)
@@ -85,7 +91,7 @@ class FitResult:
         else:
             obs_parameters = self.input_parameters
 
-        if self.bayesian_fitter.sparse:
+        if self.bayesian_fitter.settings.get("sparse", False):
             transfer_matrix = BCOO.from_scipy_sparse(
                 obs.transfer_matrix.data.to_scipy_sparse().tocsr()
             )
@@ -98,6 +104,7 @@ class FitResult:
         flux_func = jax.jit(
             jax.vmap(jax.vmap(lambda p: self.model.photon_flux(p, *energies, split_branches=True)))
         )
+
         convolve_func = jax.jit(
             jax.vmap(jax.vmap(lambda flux: jnp.clip(transfer_matrix @ flux, a_min=1e-6)))
         )
@@ -124,13 +131,14 @@ class FitResult:
 
         for key, value in input_parameters.items():
             module, parameter = key.rsplit("_", 1)
+            key_to_search = f"mod/~/{module}_{parameter}"
 
-            if f"{module}_{parameter}" in posterior.keys():
+            if key_to_search in posterior.keys():
                 # We add as extra dimension as there might be different values per observation
-                if posterior[f"{module}_{parameter}"].shape == samples_shape:
-                    to_set = posterior[f"{module}_{parameter}"][..., None]
+                if posterior[key_to_search].shape == samples_shape:
+                    to_set = posterior[key_to_search][..., None]
                 else:
-                    to_set = posterior[f"{module}_{parameter}"]
+                    to_set = posterior[key_to_search]
 
                 input_parameters[f"{module}_{parameter}"] = to_set
 
@@ -299,7 +307,7 @@ class FitResult:
 
         return value
 
-    def to_chain(self, name: str) -> Chain:
+    def to_chain(self, name: str, parameter_kind="mod") -> Chain:
         """
         Return a ChainConsumer Chain object from the posterior distribution of the parameters_type.
 
@@ -308,9 +316,7 @@ class FitResult:
         """
 
         keys_to_drop = [
-            key
-            for key in self.inference_data.posterior.keys()
-            if (key.startswith("_") or key.startswith("bkg"))
+            key for key in self.inference_data.posterior.keys() if not key.startswith("mod")
         ]
 
         reduced_id = az.extract(
@@ -338,6 +344,8 @@ class FitResult:
 
         df = pd.concat(df_list, axis=1)
 
+        df = df.rename(columns=lambda x: x.split("/~/")[-1])
+
         return Chain(samples=df, name=name)
 
     @property
@@ -450,7 +458,7 @@ class FitResult:
                 legend_labels = []
 
                 count = az.extract(
-                    self.inference_data, var_names=f"obs_{obs_id}", group="posterior_predictive"
+                    self.inference_data, var_names=f"obs/~/{obs_id}", group="posterior_predictive"
                 ).values.T
 
                 xbins, exposure, integrated_arf = _compute_effective_area(obsconf, x_unit)
@@ -465,7 +473,9 @@ class FitResult:
                     case "photon_flux_density":
                         denominator = (xbins[1] - xbins[0]) * integrated_arf * exposure
 
-                y_samples = (count * u.ct / denominator).to(y_units)
+                y_samples = count * u.ct / denominator
+
+                y_samples = y_samples.to(y_units)
 
                 y_observed, y_observed_low, y_observed_high = _error_bars_for_observed_data(
                     obsconf.folded_counts.data, denominator, y_units
@@ -491,8 +501,8 @@ class FitResult:
                     alpha=0.7,
                 )
 
-                lowest_y = y_observed.min()
-                highest_y = y_observed.max()
+                lowest_y = np.nanmin(y_observed)
+                highest_y = np.nanmax(y_observed)
 
                 legend_plots.append((true_data_plot,))
                 legend_labels.append("Observed")
@@ -522,7 +532,10 @@ class FitResult:
                             count.reshape((count.shape[0] * count.shape[1], -1))
                             * u.ct
                             / denominator
-                        ).to(y_units)
+                        )
+
+                        y_samples = y_samples.to(y_units)
+
                         component_plot = _plot_binned_samples_with_error(
                             ax[0],
                             xbins.value,
@@ -545,7 +558,7 @@ class FitResult:
                         if self.background_model is None
                         else az.extract(
                             self.inference_data,
-                            var_names=f"bkg_{obs_id}",
+                            var_names=f"bkg/~/{obs_id}",
                             group="posterior_predictive",
                         ).values.T
                     )
@@ -577,18 +590,18 @@ class FitResult:
                         alpha=0.7,
                     )
 
-                    lowest_y = min(lowest_y, y_observed_bkg.min())
-                    highest_y = max(highest_y, y_observed_bkg.max())
+                    # lowest_y = np.nanmin(lowest_y.min, np.nanmin(y_observed_bkg.value).astype(float))
+                    # highest_y = np.nanmax(highest_y.value.astype(float), np.nanmax(y_observed_bkg.value).astype(float))
 
                     legend_plots.append((true_bkg_plot,))
                     legend_labels.append("Observed (bkg)")
                     legend_plots += model_bkg_plot
                     legend_labels.append("Model (bkg)")
 
-                max_residuals = np.max(np.abs(residual_samples))
+                max_residuals = min(3.5, np.nanmax(np.abs(residual_samples)))
 
                 ax[0].loglog()
-                ax[1].set_ylim(-max(3.5, max_residuals), +max(3.5, max_residuals))
+                ax[1].set_ylim(-np.nanmax([3.5, max_residuals]), +np.nanmax([3.5, max_residuals]))
                 ax[0].set_ylabel(f"Folded spectrum\n [{y_units:latex_inline}]")
                 ax[1].set_ylabel("Residuals \n" + r"[$\sigma$]")
 
@@ -635,9 +648,9 @@ class FitResult:
 
                 fig.align_ylabels()
                 plt.subplots_adjust(hspace=0.0)
+                fig.suptitle(f"Posterior predictive - {obs_id}" if title is None else title)
                 fig.tight_layout()
                 figure_list.append(fig)
-                fig.suptitle(f"Posterior predictive - {obs_id}" if title is None else title)
                 # fig.show()
 
         plt.tight_layout()
@@ -651,9 +664,9 @@ class FitResult:
         """
 
         consumer = ChainConsumer()
-        consumer.add_chain(self.to_chain(self.model.to_string()))
+        consumer.add_chain(self.to_chain("Model"))
 
-        return consumer.analysis.get_latex_table(caption="Results of the fit", label="tab:results")
+        return consumer.analysis.get_latex_table(caption="Fit result", label="tab:results")
 
     def plot_corner(
         self,

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

@@ -85,13 +85,20 @@ class ObsConfiguration(xr.Dataset):
 
         from .util import data_path_finder
 
-        arf_path_default, rmf_path_default, bkg_path_default = data_path_finder(pha_path)
+        arf_path_default, rmf_path_default, bkg_path_default = data_path_finder(
+            pha_path,
+            require_arf=(arf_path is None) and (arf_path != ""),
+            require_rmf=rmf_path is None,
+            require_bkg=bkg_path is None,
+        )
 
         arf_path = arf_path_default if arf_path is None else arf_path
         rmf_path = rmf_path_default if rmf_path is None else rmf_path
         bkg_path = bkg_path_default if bkg_path is None else bkg_path
 
-        instrument = Instrument.from_ogip_file(rmf_path, arf_path=arf_path)
+        instrument = Instrument.from_ogip_file(
+            rmf_path, arf_path=arf_path if arf_path != "" else None
+        )
         observation = Observation.from_pha_file(pha_path, bkg_path=bkg_path)
 
         return cls.from_instrument(
@@ -141,7 +148,6 @@ class ObsConfiguration(xr.Dataset):
         transfer_matrix = grouping @ (redistribution * area * exposure)
 
         # Exclude bins out of the considered energy range, and bins without contribution from the RMF
-
         row_idx = (e_min > low_energy) & (e_max < high_energy) & (grouping.sum(axis=1) > 0)
         col_idx = (e_min_unfolded > 0) & (redistribution.sum(axis=0) > 0)
 

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

@@ -164,7 +164,9 @@ class Observation(xr.Dataset):
         """
         from .util import data_path_finder
 
-        arf_path, rmf_path, bkg_path_default = data_path_finder(pha_path)
+        arf_path, rmf_path, bkg_path_default = data_path_finder(
+            pha_path, require_arf=False, require_rmf=False, require_bkg=False
+        )
         bkg_path = bkg_path_default if bkg_path is None else bkg_path
 
         pha = DataPHA.from_file(pha_path)

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/src/jaxspec/data/ogip.py

@@ -109,7 +109,7 @@ class DataPHA:
             raise ValueError("No QUALITY column found in the PHA file.")
 
         if "BACKSCAL" in header:
-            backscal = header["BACKSCAL"] * np.ones_like(data["CHANNEL"])
+            backscal = header["BACKSCAL"] * np.ones_like(data["CHANNEL"], dtype=float)
         elif "BACKSCAL" in data.colnames:
             backscal = data["BACKSCAL"]
         else:
@@ -138,7 +138,14 @@ class DataPHA:
             "flags": flags,
         }
 
-        return cls(data["CHANNEL"], data["COUNTS"], header["EXPOSURE"], **kwargs)
+        if "COUNTS" in data.colnames:
+            counts = data["COUNTS"]
+        elif "RATE" in data.colnames:
+            counts = data["RATE"] * header["EXPOSURE"]
+        else:
+            raise ValueError("No COUNTS or RATE column found in the PHA file.")
+
+        return cls(data["CHANNEL"], counts, header["EXPOSURE"], **kwargs)
 
 
 class DataARF:

{jaxspec-0.2.2.dev0 → jaxspec-0.3.0}/src/jaxspec/data/util.py

@@ -228,12 +228,17 @@ def fakeit_for_multiple_parameters(
     return fakeits[0] if len(fakeits) == 1 else fakeits
 
 
-def data_path_finder(pha_path: str) -> tuple[str | None, str | None, str | None]:
+def data_path_finder(
+    pha_path: str, require_arf: bool = True, require_rmf: bool = True, require_bkg: bool = False
+) -> tuple[str | None, str | None, str | None]:
     """
     Function which tries its best to find the ARF, RMF and BKG files associated with a given PHA file.
 
     Parameters:
         pha_path: The PHA file path.
+        require_arf: Whether to raise an error if the ARF file is not found.
+        require_rmf: Whether to raise an error if the RMF file is not found.
+        require_bkg: Whether to raise an error if the BKG file is not found.
 
     Returns:
         arf_path: The ARF file path.
@@ -241,23 +246,24 @@ def data_path_finder(pha_path: str) -> tuple[str | None, str | None, str | None]
         bkg_path: The BKG file path.
     """
 
-    def find_path(file_name: str, directory: str) -> str | None:
-        if file_name.lower() != "none" and file_name != "":
-            return find_file_or_compressed_in_dir(file_name, directory)
-        else:
-            return None
+    def find_path(file_name: str, directory: str, raise_err: bool = True) -> str | None:
+        if raise_err:
+            if file_name.lower() != "none" and file_name != "":
+                return find_file_or_compressed_in_dir(file_name, directory, raise_err)
+
+        return None
 
     header = fits.getheader(pha_path, "SPECTRUM")
     directory = str(Path(pha_path).parent)
 
-    arf_path = find_path(header.get("ANCRFILE", "none"), directory)
-    rmf_path = find_path(header.get("RESPFILE", "none"), directory)
-    bkg_path = find_path(header.get("BACKFILE", "none"), directory)
+    arf_path = find_path(header.get("ANCRFILE", "none"), directory, require_arf)
+    rmf_path = find_path(header.get("RESPFILE", "none"), directory, require_rmf)
+    bkg_path = find_path(header.get("BACKFILE", "none"), directory, require_bkg)
 
     return arf_path, rmf_path, bkg_path
 
 
-def find_file_or_compressed_in_dir(path: str | Path, directory: str | Path) -> str:
+def find_file_or_compressed_in_dir(path: str | Path, directory: str | Path, raise_err: bool) -> str:
     """
     Try to find a file or its .gz compressed version in a given directory and return
     the full path of the file.
@@ -275,5 +281,5 @@ def find_file_or_compressed_in_dir(path: str | Path, directory: str | Path) -> s
         if file.suffix == ".gz":
             return str(file)
 
-    else:
+    elif raise_err:
         raise FileNotFoundError(f"Can't find {path}(.gz) in {directory}.")