Merge pull request #142 from renecotyfanboy/v0.0.5

V0.0.5

.devcontainer.json

@@ -0,0 +1,4 @@
+{
+  "name": "jaxspec-xspec-container",
+  "dockerFile": "Dockerfile"
+}

.pre-commit-config.yaml

@@ -2,7 +2,7 @@ fail_fast: true
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.5.0
+    rev: v4.6.0
     hooks:
      # - id: check-added-large-files
      #   args: ["--maxkb=5000"]
@@ -18,13 +18,13 @@ repos:
         exclude: \.(fits|pha|arf|rmf)$  # Exclude .fits, .pha, .arf, .rmf files
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.1.3
+    rev: v0.3.7
     hooks:
       - id: ruff-format
         exclude: \.(fits|pha|arf|rmf)$  # Exclude .fits, .pha, .arf, .rmf files
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.1.3
+    rev: v0.3.7
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

.readthedocs.yaml

@@ -23,6 +23,7 @@ build:
       # https://python-poetry.org/docs/managing-dependencies/#dependency-groups
       - poetry install --with docs
       # Using insiders versions of mkdocs-material & mkdocstrings
+      - pip uninstall mkdocs-material mkdocstrings mkdocstrings-python -y
       - pip install git+https://${GH_TOKEN}@github.com/squidfunk/mkdocs-material-insiders.git
       - pip install git+https://${GH_TOKEN}@github.com/pawamoy-insiders/mkdocstrings-python.git
       - pip install mkdocs-jupyter # This is bugged, I enforced it manually, let's see if it works

README.md

@@ -16,7 +16,7 @@ Documentation : https://jaxspec.readthedocs.io/en/latest/
 
 ## Installation
 
-We recommend the users to start from a fresh Python 3.10 [conda environment](https://conda.io/projects/conda/en/latest/user-guide/install/index.html). 
+We recommend the users to start from a fresh Python 3.10 [conda environment](https://conda.io/projects/conda/en/latest/user-guide/install/index.html).
 
 ```
 conda create -n jaxspec python=3.10

docs/examples/background.md

@@ -12,7 +12,7 @@ approach is equivalent to subtract the background to the observed spectrum when
 ``` python
 from jaxspec.model.background import SubtractedBackground
 
-forward = BayesianModel(model, obs, background_model=SubtractedBackground())
+forward = BayesianFitter(model, obs, background_model=SubtractedBackground())
 result = forward.fit(prior, num_chains=4, num_warmup=1000, num_samples=1000, mcmc_kwargs={"progress_bar": True})
 
 result.plot_ppc()
@@ -28,7 +28,7 @@ it is to consider each background bin as a Poisson realisation of a counting pro
 ``` python
 from jaxspec.model.background import BackgroundWithError
 
-forward = BayesianModel(model, obs, background_model=BackgroundWithError())
+forward = BayesianFitter(model, obs, background_model=BackgroundWithError())
 result = forward.fit(prior, num_chains=4, num_warmup=1000, num_samples=1000, mcmc_kwargs={"progress_bar": True})
 
 result.plot_ppc()
@@ -43,7 +43,7 @@ nodes will drive the flexibility of the Gaussian process, and it should always b
 ``` python
 from jaxspec.model.background import GaussianProcessBackground
 
-forward = BayesianModel(model, obs, background_model=GaussianProcessBackground(e_min=0.3, e_max=8, n_nodes=20))
+forward = BayesianFitter(model, obs, background_model=GaussianProcessBackground(e_min=0.3, e_max=8, n_nodes=20))
 result = forward.fit(prior, num_chains=4, num_warmup=1000, num_samples=1000, mcmc_kwargs={"progress_bar": True})
 
 result.plot_ppc()

docs/examples/build_model.md

@@ -5,7 +5,7 @@ your favorite spectral fitting library. The following example shows how to build
 models using additive and multiplicative components.
 
 ```python
-from jaxspec.model.additive import Powerlaw
+from jaxspec.model._additive import Powerlaw
 from jaxspec.model.multiplicative import Tbabs
 
 model_simple = Tbabs() * Powerlaw()
@@ -30,7 +30,7 @@ build arbitrary complex models, in the same fashion as you would do in other
 spectral fitting libraries.
 
 ```python
-from jaxspec.model.additive import Blackbody, Powerlaw
+from jaxspec.model._additive import Blackbody, Powerlaw
 from jaxspec.model.multiplicative import Tbabs, Phabs
 
 model_complex = Tbabs() * (Powerlaw() + Phabs() * Blackbody()) + Blackbody()

docs/examples/fitting_example.md

@@ -49,7 +49,7 @@ obs = ObsConfiguration.from_pha_file('obs_1.pha', low_energy=0.3, high_energy=12
 
 ``` python
 import numpyro.distributions as dist
-from jaxspec.fit import BayesianModel
+from jaxspec.fit import BayesianFitter
 
 prior = {
     "powerlaw_1": {
@@ -64,7 +64,7 @@ prior = {
         {"N_H": dist.Uniform(0, 1)}
 }
 
-forward = BayesianModel(model, obs)
+forward = BayesianFitter(model, obs)
 result = forward.fit(prior, num_chains=4, num_warmup=5000, num_samples=5000)
 ```
 

docs/examples/result_analysis.ipynb

@@ -72,7 +72,7 @@
     "numpyro.set_host_device_count(4)\n",
     "\n",
     "import numpyro.distributions as dist\n",
-    "from jaxspec.fit import BayesianModel\n",
+    "from jaxspec.fit import BayesianFitter\n",
     "from jaxspec.data import ObsConfiguration\n",
     "from jaxspec.model.background import ConjugateBackground\n",
     "from jaxspec.model.additive import Diskbb, Powerlaw\n",
@@ -86,7 +86,7 @@
     "    high_energy=7.5\n",
     ")\n",
     "\n",
-    "forward = BayesianModel(model, obs, background_model=ConjugateBackground())\n",
+    "forward = BayesianFitter(model, obs, background_model=ConjugateBackground())\n",
     "\n",
     "prior = {\n",
     "    \"diskbb_1\": {\n",

docs/models/apec.md

@@ -0,0 +1,6 @@
+# Additive components
+
+::: jaxspec.model._additive.apec
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false

docs/models/index.md

@@ -0,0 +1,3 @@
+# Welcome to JAXspec's documentation!
+
+Some models are highlighted in this section

docs/references/abundance.md

@@ -0,0 +1,4 @@
+::: jaxspec.util.abundance
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false

docs/references/fitting.md

@@ -1,11 +1,18 @@
-::: jaxspec.fit.BayesianModelAbstract
+::: jaxspec.fit.BayesianFitter
     options:
       show_root_heading: true
       show_root_full_path: false
       show_root_toc_entry: true
       heading_level: 3
 
-::: jaxspec.fit.BayesianModel
+::: jaxspec.fit.MinimizationFitter
+    options:
+      show_root_heading: true
+      show_root_full_path: false
+      show_root_toc_entry: true
+      heading_level: 3
+
+::: jaxspec.fit.ModelFitter
     options:
       show_root_heading: true
       show_root_full_path: false

docs/references/results.md

@@ -1,12 +1,13 @@
-::: jaxspec.analysis.results.ChainResult
+::: jaxspec.analysis.results.FitResult
     options:
       members:
+        - converged
+        - samples_haiku
+        - samples_flat
         - photon_flux
         - energy_flux
         - luminosity
         - plot_ppc
         - plot_corner
         - table
-        - chain
-      show_root_heading: true
-      show_root_toc_entry: false
+      show_root_heading: true

mkdocs.yml

@@ -16,6 +16,9 @@ nav:
     - Add a background : examples/background.md
     - Work with arviz : examples/result_analysis.ipynb
     - Build a custom model : examples/build_custom_model.ipynb
+  - Models:
+      - models/index.md
+      - APEC: models/apec.md
   - API Reference:
     - Spectral model base: references/model.md
     - Additive models: references/additive.md
@@ -24,6 +27,7 @@ nav:
     - Data containers: references/data.md
     - Fitting: references/fitting.md
     - Result containers: references/results.md
+    - Abundance: references/abundance.md
     - Integration: references/integrate.md
   - FAQ:
       - faq/index.md
@@ -89,7 +93,9 @@ plugins:
             members_order: alphabetical
             docstring_section_style: spacy
             docstring_style: google
+            show_symbol_type_heading: true
             show_symbol_type_toc: true
+            modernize_annotations: true
 #  - group:
 #      enabled: !ENV INSIDERS
 #      plugins:

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jaxspec"
-version = "0.0.4"
+version = "0.0.5"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = ["sdupourque <sdupourque@irap.omp.eu>"]
 license = "MIT"
@@ -28,6 +28,8 @@ seaborn = "^0.13.1"
 mkdocstrings = "^0.24.0"
 sparse = "^0.15.1"
 scipy = "<1.13"
+mendeleev = "^0.15.0"
+pyzmq = "<26"
 
 
 [tool.poetry.group.docs.dependencies]

src/jaxspec/__init__.py

@@ -1,5 +1,5 @@
 """
-This is the root of jaxspec's module. TODO: clear this up.
+This is the root of jaxspec's module.
 """
 
 import importlib.metadata

src/jaxspec/analysis/compare.py

@@ -1,9 +1,9 @@
 from typing import Dict
-from .results import ChainResult
+from .results import FitResult
 from chainconsumer import ChainConsumer
 
 
-def plot_corner_comparison(obs_dict: Dict[str, ChainResult], **kwargs):
+def plot_corner_comparison(obs_dict: Dict[str, FitResult], **kwargs):
     """
     Plot the correlation plot of parameters from different fitted observations. Observations are passed in as a
     dictionary. Each observation is named according to its key. It shall be used to compare the same model independently
@@ -16,6 +16,6 @@ def plot_corner_comparison(obs_dict: Dict[str, ChainResult], **kwargs):
     c = ChainConsumer()
 
     for name, obs in obs_dict.items():
-        c.add_chain(obs.chain(name))
+        c.add_chain(obs.to_chain(name))
 
     return c.plotter.plot(**kwargs)