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

jaxspec/model/abc.py

@@ -1,37 +1,86 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
+import operator
+
+from abc import ABC
+from functools import partial
 from uuid import uuid4
 
-import haiku as hk
+import flax.nnx as nnx
 import jax
 import jax.numpy as jnp
+import jax.scipy as jsp
 import networkx as nx
-import rich
 
-from haiku._src import base
-from jax.scipy.integrate import trapezoid
-from rich.table import Table
 from simpleeval import simple_eval
 
+from jaxspec.util.typing import PriorDictType
+
+from ._graph_util import compose, export_to_mermaid
+
+
+def set_parameters(params: PriorDictType, state: nnx.State) -> nnx.State:
+    """
+    Set the parameters of a spectral model using `nnx'` routines.
+
+    Parameters:
+        params: Dictionary of parameters to set.
+        model: Spectral model.
+
+    Returns:
+        A spectral model with the newly set parameters.
+    """
+
+    state_dict = state.to_pure_dict()  # haiku-like 2 level dictionary
+
+    for key, value in params.items():
+        # Split the key to extract the module name and parameter name
+        module_name, param_name = key.rsplit("_", 1)
+        state_dict["modules"][module_name][param_name] = value
+
+    state.replace_by_pure_dict(state_dict)
+
+    return state
 
-class SpectralModel:
+
+class Composable(ABC):
     """
-    This class is supposed to handle the composition of models through basic
-    operations, and allows tracking of the operation graph and individual parameters.
+    Defines the set of operations between model components and spectral models
     """
 
-    raw_graph: nx.DiGraph
-    graph: nx.DiGraph
-    labels: dict[str, str]
-    n_parameters: int
+    def sanitize_inputs(self, other):
+        if isinstance(self, ModelComponent):
+            model_1 = SpectralModel.from_component(self)
+        else:
+            model_1 = self
+
+        if isinstance(other, ModelComponent):
+            model_2 = SpectralModel.from_component(other)
+        else:
+            model_2 = other
 
-    def __init__(self, internal_graph, labels):
-        self.raw_graph = internal_graph
-        self.labels = labels
-        self.graph = self.build_namespace()
+        return model_1, model_2
 
-        self.n_parameters = hk.data_structures.tree_size(self.params)
+    def __add__(self, other):
+        model_1, model_2 = self.sanitize_inputs(other)
+        return model_1.compose(model_2, operation="add", operation_func=operator.add)
+
+    def __mul__(self, other):
+        model_1, model_2 = self.sanitize_inputs(other)
+        return model_1.compose(model_2, operation="mul", operation_func=operator.mul)
+
+
+class SpectralModel(nnx.Module, Composable):
+    # graph: nx.DiGraph = eqx.field(static=True)
+    # modules: dict
+
+    def __init__(self, graph: nx.DiGraph):
+        self.graph = graph
+        self.modules = {}
+
+        for node, data in self.graph.nodes(data=True):
+            if "component" in data["type"]:
+                self.modules[data["name"]] = data["component"]  # (**data['kwargs'])
 
     @classmethod
     def from_string(cls, string: str) -> SpectralModel:
@@ -67,10 +116,10 @@ class SpectralModel:
             >>> model.to_string()
             "Tbabs()*(Powerlaw() + Blackbody())"
         """
-
         return str(self)
 
-    def __str__(self) -> SpectralModel:
+    """
+    def __str__(self) -> str:
         def build_expression(node_id):
             node = self.graph.nodes[node_id]
             if node["type"] == "component":
@@ -92,485 +141,297 @@ class SpectralModel:
                 predecessors = list(self.graph.predecessors(node_id))
                 return build_expression(predecessors[0])
 
-        return build_expression("out")[1:-1]
-
-    @property
-    def transformed_func_photon(self):
-        def func_to_transform(e_low, e_high, n_points=2):
-            return self.flux(e_low, e_high, n_points=n_points, energy_flux=False)
-
-        return hk.without_apply_rng(hk.transform(func_to_transform))
-
-    @property
-    def transformed_func_energy(self):
-        def func_to_transform(e_low, e_high, n_points=2):
-            return self.flux(e_low, e_high, n_points=n_points, energy_flux=True)
-
-        return hk.without_apply_rng(hk.transform(func_to_transform))
-
-    @property
-    def params(self):
-        return self.transformed_func_photon.init(None, jnp.ones(10), jnp.ones(10))
-
-    def __rich_repr__(self):
-        table = Table(title=str(self))
-
-        table.add_column("Component", justify="right", style="bold", no_wrap=True)
-        table.add_column("Parameter")
-
-        params = self.params
-
-        for component in params.keys():
-            once = True
-
-            for parameters in params[component].keys():
-                table.add_row(component if once else "", parameters)
-                once = False
-
-        return table
-
-    def __repr_html_(self):
-        return self.__rich_repr__()
-
-    def __repr__(self):
-        if not base.frame_stack:
-            rich.print(self.__rich_repr__())
-        return ""
-
-    def photon_flux(self, params, e_low, e_high, n_points=2):
-        r"""
-        Compute the expected counts between $E_\min$ and $E_\max$ by integrating the model.
-
-        $$ \Phi_{\text{photon}}\left(E_\min, ~E_\max\right) =
-        \int _{E_\min}^{E_\max}\text{d}E ~ \mathcal{M}\left( E \right)
-        \quad \left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$$
-
-        Parameters:
-            params : The parameters of the model.
-            e_low : The lower bound of the energy bins.
-            e_high : The upper bound of the energy bins.
-            n_points : The number of points used to integrate the model in each bin.
-
-        !!! info
-            This method is internally used in the inference process and should not be used directly. See
-            [`photon_flux`][jaxspec.analysis.results.FitResult.photon_flux] to compute
-            the photon flux associated with a set of fitted parameters in a
-            [`FitResult`][jaxspec.analysis.results.FitResult]
-            instead.
-        """
-
-        params = jax.tree_map(lambda x: jnp.asarray(x), params)
-        e_low = jnp.asarray(e_low)
-        e_high = jnp.asarray(e_high)
-
-        return self.transformed_func_photon.apply(params, e_low, e_high, n_points=n_points)
-
-    def energy_flux(self, params, e_low, e_high, n_points=2):
-        r"""
-        Compute the expected energy flux between $E_\min$ and $E_\max$ by integrating the model.
-
-        $$ \Phi_{\text{energy}}\left(E_\min, ~E_\max\right) =
-        \int _{E_\min}^{E_\max}\text{d}E ~ E ~ \mathcal{M}\left( E \right)
-        \quad \left[\frac{\text{keV}}{\text{cm}^2\text{s}}\right]$$
-
-        Parameters:
-            params : The parameters of the model.
-            e_low : The lower bound of the energy bins.
-            e_high : The upper bound of the energy bins.
-            n_points : The number of points used to integrate the model in each bin.
-
-        !!! info
-            This method is internally used in the inference process and should not be used directly. See
-            [`energy_flux`](/references/results/#jaxspec.analysis.results.FitResult.energy_flux) to compute
-            the energy flux associated with a set of fitted parameters in a
-            [`FitResult`](/references/results/#jaxspec.analysis.results.FitResult)
-            instead.
-        """
-
-        params = jax.tree_map(lambda x: jnp.asarray(x), params)
-        e_low = jnp.asarray(e_low)
-        e_high = jnp.asarray(e_high)
-
-        return self.transformed_func_energy.apply(params, e_low, e_high, n_points=n_points)
-
-    def build_namespace(self):
-        """
-        This method build a namespace for the model components, to avoid name collision
-        """
-
-        name_space = []
-        new_graph = self.raw_graph.copy()
-
-        for node_id in nx.dag.topological_sort(new_graph):
-            node = new_graph.nodes[node_id]
-
-            if node and node["type"] == "component":
-                name_space.append(node["name"])
-                n = name_space.count(node["name"])
-                nx.set_node_attributes(new_graph, {node_id: name_space[-1] + f"_{n}"}, "name")
-
-        return new_graph
+        return "This must be changed"  # build_expression("out")[1:-1]
+    """
 
-    def flux(self, e_low, e_high, energy_flux=False, n_points=2):
+    def compose(self, other, operation=None, operation_func=None):
         """
-        This method return the expected counts between e_low and e_high by integrating the model.
-        It contains most of the "usine à gaz" which makes jaxspec works.
-        It evaluates the graph of operations and returns the result.
-        It should be transformed using haiku.
+        This function operate a composition between the operation graph of two models
+        1) It fuses the two graphs using which joins at the 'out' nodes and change components name to unique identifiers
+        2) It relabels the 'out' node with a unique identifier and labels it with the operation
+        3) It links the operation to a new 'out' node
         """
 
-        # TODO : enable interpolation and integration with more than 2 points for the continuum
-
-        if n_points == 2:
-            energies = jnp.hstack((e_low, e_high[-1]))
-            energies_to_integrate = jnp.stack((e_low, e_high))
-
-        else:
-            energies_to_integrate = jnp.linspace(e_low, e_high, n_points)
-            energies = energies_to_integrate
-
-        fine_structures_flux = jnp.zeros_like(e_low)
-        runtime_modules = {}
-        continuum = {}
-
-        # Iterate through the graph in topological order and
-        # compute the continuum contribution for each component
-
-        for node_id in nx.dag.topological_sort(self.graph):
-            node = self.graph.nodes[node_id]
-
-            # Instantiate the haiku modules
-            if node and node["type"] == "component":
-                runtime_modules[node_id] = node["component"](name=node["name"], **node["kwargs"])
-                continuum[node_id] = runtime_modules[node_id].continuum(energies)
-
-            elif node and node["type"] == "operation":
-                component_1 = list(self.graph.in_edges(node_id))[0][0]  # noqa: RUF015
-                component_2 = list(self.graph.in_edges(node_id))[1][0]
-                continuum[node_id] = node["function"](
-                    continuum[component_1], continuum[component_2]
-                )
-
-        if n_points == 2:
-            flux_1D = continuum[list(self.graph.in_edges("out"))[0][0]]  # noqa: RUF015
-            flux = jnp.stack((flux_1D[:-1], flux_1D[1:]))
-
-        else:
-            flux = continuum[list(self.graph.in_edges("out"))[0][0]]  # noqa: RUF015
-
-        if energy_flux:
-            continuum_flux = trapezoid(
-                flux * energies_to_integrate**2,
-                x=jnp.log(energies_to_integrate),
-                axis=0,
-            )
-
-        else:
-            continuum_flux = trapezoid(
-                flux * energies_to_integrate, x=jnp.log(energies_to_integrate), axis=0
-            )
-
-        # Iterate from the root nodes to the output node and
-        # compute the fine structure contribution for each component
-
-        root_nodes = [
-            node_id
-            for node_id, in_degree in self.graph.in_degree(self.graph.nodes)
-            if in_degree == 0 and self.graph.nodes[node_id].get("component_type") == "additive"
-        ]
-
-        for root_node_id in root_nodes:
-            path = nx.shortest_path(self.graph, source=root_node_id, target="out")
-            nodes_id_in_path = [node_id for node_id in path]
-
-            flux_from_component, mean_energy = runtime_modules[root_node_id].emission_lines(
-                e_low, e_high
-            )
-
-            multiplicative_nodes = []
+        composed_graph = compose(
+            self.graph, other.graph, operation=operation, operation_func=operation_func
+        )
 
-            # Search all multiplicative components connected to this node
-            # and apply them at mean energy
-            for node_id in nodes_id_in_path[::-1]:
-                multiplicative_nodes.extend(
-                    [node_id for node_id in self.find_multiplicative_components(node_id)]
-                )
+        return SpectralModel(composed_graph)
 
-            for mul_node in multiplicative_nodes:
-                flux_from_component *= runtime_modules[mul_node].continuum(mean_energy)
+    @classmethod
+    def from_component(cls, component):
+        node_id = str(uuid4())
+        graph = nx.DiGraph()
 
-            if energy_flux:
-                fine_structures_flux += trapezoid(
-                    flux_from_component * energies_to_integrate,
-                    x=jnp.log(energies_to_integrate),
-                    axis=0,
-                )
+        node_properties = {
+            "type": f"{component.type}_component",
+            "name": f"{component.__class__.__name__}_1".lower(),
+            "component": component,
+            "depth": 0,
+        }
 
-            else:
-                fine_structures_flux += flux_from_component
+        graph.add_node(node_id, **node_properties)
+        graph.add_node("out", type="out", depth=1)
+        graph.add_edge(node_id, "out")
 
-        return continuum_flux + fine_structures_flux
+        return cls(graph)
 
-    def find_multiplicative_components(self, node_id):
+    def _find_multiplicative_components(self, node_id):
         """
         Recursively finds all the multiplicative components connected to the node with the given ID.
         """
         node = self.graph.nodes[node_id]
         multiplicative_nodes = []
 
-        if node.get("operation_type") == "mul":
+        if node.get("type") == "mul_operation":
             # Recursively find all the multiplicative components using the predecessors
             predecessors = self.graph.pred[node_id]
             for node_id in predecessors:
-                if self.graph.nodes[node_id].get("component_type") == "multiplicative":
+                if "multiplicative_component" == self.graph.nodes[node_id].get("type"):
                     multiplicative_nodes.append(node_id)
-                elif self.graph.nodes[node_id].get("operation_type") == "mul":
-                    multiplicative_nodes.extend(self.find_multiplicative_components(node_id))
+                elif "mul_operation" == self.graph.nodes[node_id].get("type"):
+                    multiplicative_nodes.extend(self._find_multiplicative_components(node_id))
 
         return multiplicative_nodes
 
-    def __call__(self, pars, e_low, e_high, **kwargs):
-        return self.photon_flux(pars, e_low, e_high, **kwargs)
+    @property
+    def root_nodes(self) -> list[str]:
+        return [
+            node_id
+            for node_id, in_degree in self.graph.in_degree(self.graph.nodes)
+            if in_degree == 0 and ("additive" in self.graph.nodes[node_id].get("type"))
+        ]
 
-    @classmethod
-    def from_component(cls, component, **kwargs) -> SpectralModel:
-        """
-        Build a model from a single component
-        """
+    @property
+    def branches(self) -> list[str]:
+        branches = []
 
-        graph = nx.DiGraph()
+        for root_node_id in self.root_nodes:
+            root_node_name = self.graph.nodes[root_node_id].get("name")
+            path = nx.shortest_path(self.graph, source=root_node_id, target="out")
+            multiplicative_components = []
 
-        # Add the component node
-        # Random static node id to keep it trackable in the graph
-        node_id = str(uuid4())
+            # Search all multiplicative components connected to this node
+            # and apply them at mean energy
+            for node_id in path[::-1]:
+                multiplicative_components.extend(
+                    [node_id for node_id in self._find_multiplicative_components(node_id)]
+                )
 
-        if component.type == "additive":
+            branch = ""
 
-            def lam_func(e):
-                return (
-                    component(**kwargs).continuum(e)
-                    + component(**kwargs).emission_lines(e, e + 1)[0]
-                )
+            for multiplicative_node_id in multiplicative_components:
+                multiplicative_node_name = self.graph.nodes[multiplicative_node_id].get("name")
+                branch += f"{multiplicative_node_name}*"
 
-        elif component.type == "multiplicative":
+            branch += f"{root_node_name}"
+            branches.append(branch)
 
-            def lam_func(e):
-                return component().continuum(e)
+        return branches
 
-        else:
+    def turbo_flux(self, e_low, e_high, energy_flux=False, n_points=2, return_branches=False):
+        continuum = {}
 
-            def lam_func(e):
-                return print("Some components are not working at this stage")
+        ## Evaluate the expected contribution for each component
+        for node_id in nx.dag.topological_sort(self.graph):
+            node = self.graph.nodes[node_id]
 
-        node_properties = {
-            "type": "component",
-            "component_type": component.type,
-            "name": component.__name__.lower(),
-            "component": component,
-            # "params": hk.transform(lam_func).init(None, jnp.ones(1)),
-            "fine_structure": False,
-            "kwargs": kwargs,
-            "depth": 0,
-        }
+            if node["type"] == "additive_component":
+                node_name = node["name"]
+                runtime_modules = self.modules[node_name]
 
-        graph.add_node(node_id, **node_properties)
+                if not energy_flux:
+                    continuum[node_name] = runtime_modules._photon_flux(
+                        e_low, e_high, n_points=n_points
+                    )
 
-        # Add the output node
-        labels = {node_id: component.__name__.lower(), "out": "out"}
+                else:
+                    continuum[node_name] = runtime_modules._energy_flux(
+                        e_low, e_high, n_points=n_points
+                    )
 
-        graph.add_node("out", type="out", depth=1)
-        graph.add_edge(node_id, "out")
+            elif node["type"] == "multiplicative_component":
+                node_name = node["name"]
+                runtime_modules = self.modules[node_name]
+                continuum[node_name] = runtime_modules.factor((e_low + e_high) / 2)
 
-        return cls(graph, labels)
+            else:
+                pass
 
-    def compose(
-        self, other: SpectralModel, operation=None, function=None, name=None
-    ) -> SpectralModel:
-        """
-        This function operate a composition between the operation graph of two models
-        1) It fuses the two graphs using which joins at the 'out' nodes
-        2) It relabels the 'out' node with a unique identifier and labels it with the operation
-        3) It links the operation to a new 'out' node
-        """
+        ## Propagate the absorption for each branch
+        root_nodes = self.root_nodes
 
-        # Compose the two graphs with their output as common node
-        # and add the operation node by overwriting the 'out' node
-        node_id = str(uuid4())
-        graph = nx.relabel_nodes(nx.compose(self.raw_graph, other.raw_graph), {"out": node_id})
-        nx.set_node_attributes(graph, {node_id: "operation"}, "type")
-        nx.set_node_attributes(graph, {node_id: operation}, "operation_type")
-        nx.set_node_attributes(graph, {node_id: function}, "function")
-        nx.set_node_attributes(graph, {node_id: name}, "operation_label")
-
-        # Merge label dictionaries
-        labels = self.labels | other.labels
-        labels[node_id] = operation
-
-        # Now add the output node and link it to the operation node
-        graph.add_node("out", type="out")
-        graph.add_edge(node_id, "out")
+        branches = {}
 
-        # Compute the new depth of each node
-        longest_path = nx.dag_longest_path_length(graph)
+        for root_node_id in root_nodes:
+            root_node_name = self.graph.nodes[root_node_id].get("name")
+            root_continuum = continuum[root_node_name]
 
-        for node in graph.nodes:
-            nx.set_node_attributes(
-                graph,
-                {node: longest_path - nx.shortest_path_length(graph, node, "out")},
-                "depth",
-            )
+            path = nx.shortest_path(self.graph, source=root_node_id, target="out")
+            multiplicative_components = []
 
-        return SpectralModel(graph, labels)
+            # Search all multiplicative components connected to this node
+            # and apply them at mean energy
+            for node_id in path[::-1]:
+                multiplicative_components.extend(
+                    [node_id for node_id in self._find_multiplicative_components(node_id)]
+                )
 
-    def __add__(self, other: SpectralModel) -> SpectralModel:
-        return self.compose(other, operation="add", function=lambda x, y: x + y, name="+")
+            branch = ""
 
-    def __mul__(self, other: SpectralModel) -> SpectralModel:
-        return self.compose(other, operation="mul", function=lambda x, y: x * y, name=r"*")
+            for multiplicative_node_id in multiplicative_components:
+                multiplicative_node_name = self.graph.nodes[multiplicative_node_id].get("name")
+                root_continuum *= continuum[multiplicative_node_name]
+                branch += f"{multiplicative_node_name}*"
 
-    def export_to_mermaid(self, file=None):
-        mermaid_code = "graph LR\n"  # LR = left to right
+            branch += f"{root_node_name}"
+            branches[branch] = root_continuum
 
-        # Add nodes
-        for node, attributes in self.graph.nodes(data=True):
-            if attributes["type"] == "component":
-                name, number = attributes["name"].split("_")
-                mermaid_code += f'    {node}("{name.capitalize()} ({number})")\n'
+        if return_branches:
+            return branches
 
-            if attributes["type"] == "operation":
-                if attributes["operation_type"] == "add":
-                    mermaid_code += f"    {node}{{+}}\n"
+        return sum(branches.values())
 
-                if attributes["operation_type"] == "mul":
-                    mermaid_code += f"    {node}{{x}}\n"
+    def to_mermaid(self, file: str | None = None):
+        """
+        This method returns the mermaid representation of the model.
 
-            if attributes["type"] == "out":
-                mermaid_code += f'    {node}("Output")\n'
+        Parameters:
+            file : The file to write the mermaid representation to.
 
-        # Draw connexion between nodes
-        for source, target in self.graph.edges():
-            mermaid_code += f"    {source} --> {target}\n"
+        Returns:
+            A string containing the mermaid representation of the model.
+        """
+        return export_to_mermaid(self.graph, file)
 
-        if file is None:
-            return mermaid_code
-        else:
-            with open(file, "w") as f:
-                f.write(mermaid_code)
+    @partial(jax.jit, static_argnums=0, static_argnames=("n_points", "split_branches"))
+    def photon_flux(self, params, e_low, e_high, n_points=2, split_branches=False):
+        r"""
+        Compute the expected counts between $E_\min$ and $E_\max$ by integrating the model.
 
-    def plot(self, figsize=(8, 8)):
-        import matplotlib.pyplot as plt
+        $$ \Phi_{\text{photon}}\left(E_\min, ~E_\max\right) =
+        \int _{E_\min}^{E_\max}\text{d}E ~ \mathcal{M}\left( E \right)
+        \quad \left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$$
 
-        plt.figure(figsize=figsize)
+        Parameters:
+            params : The parameters of the model.
+            e_low : The lower bound of the energy bins.
+            e_high : The upper bound of the energy bins.
+            n_points : The number of points used to integrate the model in each bin.
 
-        pos = nx.multipartite_layout(self.graph, subset_key="depth", scale=1)
+        !!! info
+            This method is internally used in the inference process and should not be used directly. See
+            [`photon_flux`][jaxspec.analysis.results.FitResult.photon_flux] to compute
+            the photon flux associated with a set of fitted parameters in a
+            [`FitResult`][jaxspec.analysis.results.FitResult]
+            instead.
+        """
 
-        nodes_out = [x for x, y in self.graph.nodes(data=True) if y["type"] == "out"]
-        nx.draw_networkx_nodes(self.graph, pos, nodelist=nodes_out, node_color="tab:green")
-        nx.draw_networkx_edges(self.graph, pos, width=1.0)
+        graphdef, state = nnx.split(self)
+        state = set_parameters(params, state)
 
-        nx.draw_networkx_labels(
-            self.graph,
-            pos,
-            labels=nx.get_node_attributes(self.graph, "name"),
-            font_size=12,
-            font_color="black",
-            bbox={"fc": "tab:red", "boxstyle": "round", "pad": 0.3},
-        )
-        nx.draw_networkx_labels(
-            self.graph,
-            pos,
-            labels=nx.get_node_attributes(self.graph, "operation_label"),
-            font_size=12,
-            font_color="black",
-            bbox={"fc": "tab:blue", "boxstyle": "circle", "pad": 0.3},
-        )
+        return nnx.call((graphdef, state)).turbo_flux(
+            e_low, e_high, n_points=n_points, return_branches=split_branches
+        )[0]
 
-        plt.axis("equal")
-        plt.axis("off")
-        plt.tight_layout()
-        plt.show()
+    @partial(jax.jit, static_argnums=0, static_argnames="n_points")
+    def energy_flux(self, params, e_low, e_high, n_points=2):
+        r"""
+        Compute the expected energy flux between $E_\min$ and $E_\max$ by integrating the model.
 
+        $$ \Phi_{\text{energy}}\left(E_\min, ~E_\max\right) =
+        \int _{E_\min}^{E_\max}\text{d}E ~ E ~ \mathcal{M}\left( E \right)
+        \quad \left[\frac{\text{keV}}{\text{cm}^2\text{s}}\right]$$
 
-class ComponentMetaClass(type(hk.Module)):
-    """
-    This metaclass enable the construction of model from components with a simple
-    syntax while style enabling the components to be used as haiku modules.
-    """
+        Parameters:
+            params : The parameters of the model.
+            e_low : The lower bound of the energy bins.
+            e_high : The upper bound of the energy bins.
+            n_points : The number of points used to integrate the model in each bin.
 
-    def __call__(self, **kwargs) -> SpectralModel:
-        """
-        This method enable to use model components as haiku modules when folded in a haiku transform
-        function and also to instantiate them as SpectralModel when out of a haiku transform
+        !!! info
+            This method is internally used in the inference process and should not be used directly. See
+            [`energy_flux`](/references/results/#jaxspec.analysis.results.FitResult.energy_flux) to compute
+            the energy flux associated with a set of fitted parameters in a
+            [`FitResult`](/references/results/#jaxspec.analysis.results.FitResult)
+            instead.
         """
 
-        if not base.frame_stack:
-            return SpectralModel.from_component(self, **kwargs)
+        graphdef, state = nnx.split(self)
+        state = set_parameters(params, state)
 
-        else:
-            return super().__call__(**kwargs)
+        return nnx.call((graphdef, state)).turbo_flux(
+            e_low, e_high, n_points=n_points, energy_flux=True
+        )[0]
 
 
-class ModelComponent(hk.Module, ABC, metaclass=ComponentMetaClass):
+class ModelComponent(nnx.Module, Composable, ABC):
     """
     Abstract class for model components
     """
 
-    type: str
+    ...
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
 
-
-class AdditiveComponent(ModelComponent, ABC):
+class AdditiveComponent(ModelComponent):
     type = "additive"
 
     def continuum(self, energy):
-        """
-        Method for computing the continuum associated to the model.
-        By default, this is set to 0, which means that the model has no continuum.
-        This should be overloaded by the user if the model has a continuum.
-        """
+        r"""
+        Compute the continuum of the component.
 
+        Parameters:
+            energy : The energy at which to compute the continuum.
+        """
         return jnp.zeros_like(energy)
 
-    def emission_lines(self, e_min, e_max) -> (jax.Array, jax.Array):
-        """
-        Method for computing the fine structure of an additive model between two energies.
-        By default, this is set to 0, which means that the model has no emission lines.
-        This should be overloaded by the user if the model has a fine structure.
+    def integrated_continuum(self, e_low, e_high):
+        r"""
+        Compute the integrated continuum between $E_\min$ and $E_\max$.
+
+        Parameters:
+            e_low: Lower bound of the energy bin.
+            e_high: Upper bound of the energy bin.
         """
+        return jnp.zeros_like((e_low + e_high) / 2)
 
-        return jnp.zeros_like(e_min), (e_min + e_max) / 2
+    def _photon_flux(self, e_low, e_high, n_points=2):
+        energy = jnp.linspace(e_low, e_high, n_points, axis=-1)
+        continuum = self.continuum(energy)
+        integrated_continuum = self.integrated_continuum(e_low, e_high)
 
-    '''
-    def integral(self, e_min, e_max):
-        r"""
-        Method for integrating an additive model between two energies. It relies on
-        double exponential quadrature for finite intervals to compute an approximation
-        of the integral of a model.
+        return jsp.integrate.trapezoid(continuum, energy, axis=-1) + integrated_continuum
 
-        references
-        ----------
-        * $Takahasi and Mori (1974) <https://ems.press/journals/prims/articles/2686>$_
-        * $Mori and Sugihara (2001) <https://doi.org/10.1016/S0377-0427(00)00501-X>$_
-        * $Tanh-sinh quadrature <https://en.wikipedia.org/wiki/Tanh-sinh_quadrature>$_ from Wikipedia
+    def _energy_flux(self, e_low, e_high, n_points=2):
+        energy = jnp.linspace(e_low, e_high, n_points, axis=-1)
+        continuum = self.continuum(energy)
+        integrated_continuum = self.integrated_continuum(e_low, e_high)
 
-        """
+        return jsp.integrate.trapezoid(
+            continuum * energy**2, jnp.log(energy), axis=-1
+        ) + integrated_continuum * (e_high - e_low)
 
-        t = jnp.linspace(-4, 4, 71) # The number of points used is hardcoded and this is not ideal
-        # Quadrature nodes as defined in reference
-        phi = jnp.tanh(jnp.pi / 2 * jnp.sinh(t))
-        dphi = jnp.pi / 2 * jnp.cosh(t) * (1 / jnp.cosh(jnp.pi / 2 * jnp.sinh(t)) ** 2)
-        # Change of variable to turn the integral from E_min to E_max into an integral from -1 to 1
-        x = (e_max - e_min) / 2 * phi + (e_max + e_min) / 2
-        dx = (e_max - e_min) / 2 * dphi
+    @partial(jax.jit, static_argnums=0, static_argnames="n_points")
+    def photon_flux(self, params, e_low, e_high, n_points=2):
+        return SpectralModel.from_component(self).photon_flux(
+            params, e_low, e_high, n_points=n_points
+        )
 
-        return jnp.trapz(self(x) * dx, x=t)
-    '''
+    @partial(jax.jit, static_argnums=0, static_argnames="n_points")
+    def energy_flux(self, params, e_low, e_high, n_points=2):
+        return SpectralModel.from_component(self).energy_flux(
+            params, e_low, e_high, n_points=n_points
+        )
 
 
-class MultiplicativeComponent(ModelComponent, ABC):
+class MultiplicativeComponent(ModelComponent):
     type = "multiplicative"
 
-    @abstractmethod
-    def continuum(self, energy): ...
+    def factor(self, energy):
+        """
+        Absorption factor applied for a given energy
+
+        Parameters:
+            energy : The energy at which to compute the factor.
+        """
+        return jnp.ones_like(energy)