jaxspec 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

jaxspec/fit.py

@@ -10,11 +10,12 @@ import arviz as az
 import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
-import numpy as np
 import numpyro
 
 from jax import random
+from jax.experimental import mesh_utils
 from jax.random import PRNGKey
+from jax.sharding import PositionalSharding
 from numpyro.contrib.nested_sampling import NestedSampler
 from numpyro.distributions import Poisson, TransformedDistribution
 from numpyro.infer import AIES, ESS, MCMC, NUTS, Predictive
@@ -23,12 +24,16 @@ from numpyro.infer.reparam import TransformReparam
 from numpyro.infer.util import log_density
 
 from ._fit._build_model import build_prior, forward_model
-from .analysis._plot import _plot_poisson_data_with_error
+from .analysis._plot import (
+    _error_bars_for_observed_data,
+    _plot_binned_samples_with_error,
+    _plot_poisson_data_with_error,
+)
 from .analysis.results import FitResult
 from .data import ObsConfiguration
 from .model.abc import SpectralModel
 from .model.background import BackgroundModel
-from .util.typing import PriorDictModel, PriorDictType
+from .util.typing import PriorDictType
 
 
 class BayesianModel:
@@ -63,10 +68,12 @@ class BayesianModel:
 
         if not callable(prior_distributions):
             # Validate the entry with pydantic
-            prior = PriorDictModel.from_dict(prior_distributions).nested_dict
+            # prior = PriorDictModel.from_dict(prior_distributions).
 
             def prior_distributions_func():
-                return build_prior(prior, expand_shape=(len(self.observation_container),))
+                return build_prior(
+                    prior_distributions, expand_shape=(len(self.observation_container),)
+                )
 
         else:
             prior_distributions_func = prior_distributions
@@ -74,6 +81,22 @@ class BayesianModel:
         self.prior_distributions_func = prior_distributions_func
         self.init_params = self.prior_samples()
 
+        # Check the priors are suited for the observations
+        split_parameters = [
+            (param, shape[-1])
+            for param, shape in jax.tree.map(lambda x: x.shape, self.init_params).items()
+            if (len(shape) > 1)
+            and not param.startswith("_")
+            and not param.startswith("bkg")  # hardcoded for subtracted background
+        ]
+
+        for parameter, proposed_number_of_obs in split_parameters:
+            if proposed_number_of_obs != len(self.observation_container):
+                raise ValueError(
+                    f"Invalid splitting in the prior distribution. "
+                    f"Expected {len(self.observation_container)} but got {proposed_number_of_obs} for {parameter}"
+                )
+
     @cached_property
     def observation_container(self) -> dict[str, ObsConfiguration]:
         """
@@ -137,7 +160,9 @@ class BayesianModel:
                 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,
                     )
 
@@ -289,41 +314,48 @@ class BayesianModel:
         Check if the prior distribution include the observed data.
         """
         key_prior, key_posterior = jax.random.split(key, 2)
+        n_devices = len(jax.local_devices())
+        sharding = PositionalSharding(mesh_utils.create_device_mesh((n_devices,)))
+
+        # 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:
+            num_samples = num_samples + n_devices - (num_samples % n_devices)
+
         prior_params = self.prior_samples(key=key_prior, num_samples=num_samples)
-        posterior_observations = self.mock_observations(prior_params, key=key_posterior)
+
+        # Split the parameters on every device
+        sharded_parameters = jax.device_put(prior_params, sharding)
+        posterior_observations = self.mock_observations(sharded_parameters, key=key_posterior)
 
         for key, value in self.observation_container.items():
-            fig, axs = plt.subplots(
+            fig, ax = plt.subplots(
                 nrows=2, ncols=1, sharex=True, figsize=(5, 6), height_ratios=[3, 1]
             )
 
-            _plot_poisson_data_with_error(
-                axs[0],
+            legend_plots = []
+            legend_labels = []
+
+            y_observed, y_observed_low, y_observed_high = _error_bars_for_observed_data(
+                value.folded_counts.values, 1.0, "ct"
+            )
+
+            true_data_plot = _plot_poisson_data_with_error(
+                ax[0],
                 value.out_energies,
-                value.folded_counts.values,
-                percentiles=(16, 84),
+                y_observed.value,
+                y_observed_low.value,
+                y_observed_high.value,
+                alpha=0.7,
             )
 
-            for i, (envelop_percentiles, color, alpha) in enumerate(
-                zip(
-                    [(16, 86), (2.5, 97.5), (0.15, 99.85)],
-                    ["#03045e", "#0077b6", "#00b4d8"],
-                    [0.5, 0.4, 0.3],
-                )
-            ):
-                lower, upper = np.percentile(
-                    posterior_observations["obs_" + key], envelop_percentiles, axis=0
-                )
+            prior_plot = _plot_binned_samples_with_error(
+                ax[0], value.out_energies, posterior_observations["obs_" + key], n_sigmas=3
+            )
 
-                axs[0].stairs(
-                    upper,
-                    edges=[*list(value.out_energies[0]), value.out_energies[1][-1]],
-                    baseline=lower,
-                    alpha=alpha,
-                    fill=True,
-                    color=color,
-                    label=rf"${1+i}\sigma$",
-                )
+            legend_plots.append((true_data_plot,))
+            legend_labels.append("Observed")
+            legend_plots += prior_plot
+            legend_labels.append("Prior Predictive")
 
             # rank = np.vstack((posterior_observations["obs_" + key], value.folded_counts.values)).argsort(axis=0)[-1] / (num_samples) * 100
             counts = posterior_observations["obs_" + key]
@@ -336,22 +368,22 @@ class BayesianModel:
 
             rank = (less_than_obs + 0.5 * equal_to_obs) / num_samples * 100
 
-            axs[1].stairs(rank, edges=[*list(value.out_energies[0]), value.out_energies[1][-1]])
+            ax[1].stairs(rank, edges=[*list(value.out_energies[0]), value.out_energies[1][-1]])
 
-            axs[1].plot(
+            ax[1].plot(
                 (value.out_energies.min(), value.out_energies.max()),
                 (50, 50),
                 color="black",
                 linestyle="--",
             )
 
-            axs[1].set_xlabel("Energy (keV)")
-            axs[0].set_ylabel("Counts")
-            axs[1].set_ylabel("Rank (%)")
-            axs[1].set_ylim(0, 100)
-            axs[0].set_xlim(value.out_energies.min(), value.out_energies.max())
-            axs[0].loglog()
-            axs[0].legend(loc="upper right")
+            ax[1].set_xlabel("Energy (keV)")
+            ax[0].set_ylabel("Counts")
+            ax[1].set_ylabel("Rank (%)")
+            ax[1].set_ylim(0, 100)
+            ax[0].set_xlim(value.out_energies.min(), value.out_energies.max())
+            ax[0].loglog()
+            ax[0].legend(legend_plots, legend_labels)
             plt.suptitle(f"Prior Predictive coverage for {key}")
             plt.tight_layout()
             plt.show()
@@ -544,7 +576,6 @@ class MCMCFitter(BayesianModelFitter):
         return FitResult(
             self,
             inference_data,
-            self.model.params,
             background_model=self.background_model,
         )
 
@@ -590,11 +621,13 @@ class NSFitter(BayesianModelFitter):
         ns = NestedSampler(
             bayesian_model,
             constructor_kwargs=dict(
-                num_parallel_workers=1,
                 verbose=verbose,
                 difficult_model=True,
-                max_samples=1e6,
+                max_samples=1e5,
                 parameter_estimation=True,
+                gradient_guided=True,
+                devices=jax.devices(),
+                # init_efficiency_threshold=0.01,
                 num_live_points=num_live_points,
             ),
             termination_kwargs=termination_kwargs if termination_kwargs else dict(),
@@ -613,6 +646,5 @@ class NSFitter(BayesianModelFitter):
         return FitResult(
             self,
             inference_data,
-            self.model.params,
             background_model=self.background_model,
         )

jaxspec/model/_graph_util.py

@@ -0,0 +1,151 @@
+"""Helper functions to deal with the graph logic within model building"""
+
+import re
+
+from collections.abc import Callable
+from uuid import uuid4
+
+import networkx as nx
+
+
+def get_component_names(graph: nx.DiGraph):
+    """
+    Get the set of component names from the nodes of a graph.
+
+    Parameters:
+        graph: The graph to get the component names from.
+    """
+    return set(
+        data["name"] for _, data in graph.nodes(data=True) if "component" in data.get("type")
+    )
+
+
+def increment_name(name: str, used_names: set):
+    """
+    Increment the suffix number in a name if it is formated as 'name_1'.
+
+    Parameters:
+        name: The name to increment.
+        used_names: The set of names that are already used.
+    """
+    # Use regex to extract base name and suffix number
+    match = re.match(r"(.*?)(?:_(\d+))?$", name)
+    base_name = match.group(1)
+    suffix = match.group(2)
+    if suffix:
+        number = int(suffix)
+    else:
+        number = 1  # Start from 1 if there is no suffix
+
+    new_name = name
+    while new_name in used_names:
+        number += 1
+        new_name = f"{base_name}_{number}"
+
+    return new_name
+
+
+def compose_with_rename(graph_1: nx.DiGraph, graph_2: nx.DiGraph):
+    """
+    Compose two graphs by updating the 'name' attributes of nodes in graph_2,
+    and return the graph joined on the 'out' node.
+
+    Parameters:
+        graph_1: The first graph to compose.
+        graph_2: The second graph to compose.
+    """
+
+    # Initialize the set of used names with names from graph_1
+    used_names = get_component_names(graph_1)
+
+    # Update the 'name' attributes in graph_2 to make them unique
+    for node, data in graph_2.nodes(data=True):
+        if "component" in data.get("type"):
+            original_name = data["name"]
+            new_name = original_name
+
+            if new_name in used_names:
+                new_name = increment_name(original_name, used_names)
+                data["name"] = new_name
+                used_names.add(new_name)
+
+            else:
+                used_names.add(new_name)
+
+    # Now you can safely compose the graphs
+    composed_graph = nx.compose(graph_1, graph_2)
+
+    return composed_graph
+
+
+def compose(
+    graph_1: nx.DiGraph,
+    graph_2: nx.DiGraph,
+    operation: str = "",
+    operation_func: Callable = lambda x, y: None,
+):
+    """
+    Compose two graphs by joining the 'out' node of graph_1 and graph_2, and turning
+    it to an 'operation' node with the relevant operator and add a new 'out' node.
+
+    Parameters:
+        graph_1: The first graph to compose.
+        graph_2: The second graph to compose.
+        operation: The string describing the operation to perform.
+        operation_func: The callable that performs the operation.
+    """
+
+    combined_graph = compose_with_rename(graph_1, graph_2)
+    node_id = str(uuid4())
+    graph = nx.relabel_nodes(combined_graph, {"out": node_id})
+    nx.set_node_attributes(graph, {node_id: f"{operation}_operation"}, "type")
+    nx.set_node_attributes(graph, {node_id: operation_func}, "operator")
+
+    # Now add the output node and link it to the operation node
+    graph.add_node("out", type="out")
+    graph.add_edge(node_id, "out")
+
+    # Compute the new depth of each node
+    longest_path = nx.dag_longest_path_length(graph)
+
+    for node in graph.nodes:
+        nx.set_node_attributes(
+            graph,
+            {node: longest_path - nx.shortest_path_length(graph, node, "out")},
+            "depth",
+        )
+
+    return graph
+
+
+def export_to_mermaid(graph, file=None):
+    mermaid_code = "graph LR\n"  # LR = left to right
+
+    # Add nodes
+    for node, attributes in graph.nodes(data=True):
+        if attributes["type"] == "out":
+            mermaid_code += f'    {node}("Output")\n'
+
+        else:
+            operation_type, node_type = attributes["type"].split("_")
+
+            if node_type == "component":
+                name, number = attributes["name"].split("_")
+                mermaid_code += f'    {node}("{name.capitalize()} ({number})")\n'
+
+            elif node_type == "operation":
+                if operation_type == "add":
+                    mermaid_code += f"    {node}{{**+**}}\n"
+
+                elif operation_type == "mul":
+                    mermaid_code += f"    {node}{{**x**}}\n"
+
+    # Draw connexion between nodes
+    for source, target in graph.edges():
+        mermaid_code += f"    {source} --> {target}\n"
+
+    if file is None:
+        return mermaid_code
+    else:
+        with open(file, "w") as f:
+            f.write(mermaid_code)