jaxspec 0.2.2.dev0__py3-none-any.whl → 0.3.1__py3-none-any.whl

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/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)))
         )
@@ -115,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}
 
@@ -124,13 +133,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 +309,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 +318,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 +346,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 +460,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 +475,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 +503,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 +534,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 +560,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 +592,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 +650,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 +666,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/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/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/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/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}.")

jaxspec/experimental/interpolator.py

@@ -0,0 +1,74 @@
+import numpy as np
+
+from scipy.interpolate import RegularGridInterpolator
+
+
+class RegularGridInterpolatorWithGrad(RegularGridInterpolator):
+    """
+    A subclass of SciPy's RegularGridInterpolator that also returns the gradient
+    of each interpolated output with respect to input coordinates.
+
+    Supports:
+        - Linear interpolation
+        - Out-of-bounds handling (fill_value=0)
+        - Multi-dimensional output (e.g., RGB or vector fields)
+        - Batched or single-point evaluation
+
+    Returns:
+        - values: shape (..., output_dim)
+        - gradients: shape (..., input_dim, output_dim)
+    """
+
+    def __init__(self, points, values, **kwargs):
+        kwargs.setdefault("method", "linear")
+        kwargs.setdefault("bounds_error", False)
+        kwargs.setdefault("fill_value", 0.0)
+        self.points = [np.asarray(p) for p in points]
+        self.input_dim = len(self.points)
+
+        self.output_shape = values.shape[self.input_dim :]
+        values_reshaped = values.reshape(*[len(p) for p in self.points], -1)  # flatten output
+
+        super().__init__(self.points, values_reshaped, **kwargs)
+
+    def __call__(self, xi, return_gradient=True):
+        xi = np.atleast_2d(xi).astype(float)
+        n_points, n_dims = xi.shape
+        assert n_dims == self.input_dim, "Dim mismatch"
+
+        # Interpolate values
+        flat_vals = super().__call__(xi)  # shape: (n_points, output_dim)
+        values = flat_vals.reshape(n_points, *self.output_shape)
+
+        if not return_gradient:
+            return values[0] if values.shape[0] == 1 else values
+
+        gradients = np.zeros((n_points, self.input_dim, np.prod(self.output_shape)), dtype=float)
+
+        for d, grid in enumerate(self.points):
+            xq = xi[:, d]
+            idx_upper = np.searchsorted(grid, xq, side="right")
+            idx_lower = idx_upper - 1
+
+            idx_lower = np.clip(idx_lower, 0, len(grid) - 2)
+            idx_upper = np.clip(idx_upper, 1, len(grid) - 1)
+
+            xi_low = xi.copy()
+            xi_high = xi.copy()
+            xi_low[:, d] = grid[idx_lower]
+            xi_high[:, d] = grid[idx_upper]
+
+            f_low = super().__call__(xi_low)
+            f_high = super().__call__(xi_high)
+            delta = (grid[idx_upper] - grid[idx_lower])[:, np.newaxis]
+
+            grad = np.where(delta != 0, (f_high - f_low) / delta, 0.0)
+            gradients[:, d, :] = grad
+
+        # Reshape output properly
+        gradients = gradients.reshape(n_points, self.input_dim, *self.output_shape)
+
+        if values.shape[0] == 1:
+            return values[0], gradients[0]
+        else:
+            return values, gradients

jaxspec/experimental/interpolator_jax.py

@@ -0,0 +1,79 @@
+from itertools import product
+
+import jax.numpy as jnp
+
+from jax.scipy.interpolate import RegularGridInterpolator
+
+
+class RegularGridInterpolatorWithGrad(RegularGridInterpolator):
+    """
+    A subclass of SciPy's RegularGridInterpolator that also returns the gradient
+    of each interpolated output with respect to input coordinates.
+
+    Supports:
+        - Linear interpolation
+        - Out-of-bounds handling (fill_value=0)
+        - Multi-dimensional output (e.g., RGB or vector fields)
+        - Batched or single-point evaluation
+
+    Returns:
+        - values: shape (..., output_dim)
+        - gradients: shape (..., input_dim, output_dim)
+    """
+
+    def _ndim_coords_from_arrays(self, points):
+        """Convert a tuple of coordinate arrays to a (..., ndim)-shaped array."""
+        ndim = len(self.grid)
+
+        if isinstance(points, tuple) and len(points) == 1:
+            # handle argument tuple
+            points = points[0]
+        if isinstance(points, tuple):
+            p = jnp.broadcast_arrays(*points)
+            for p_other in p[1:]:
+                if p_other.shape != p[0].shape:
+                    raise ValueError("coordinate arrays do not have the same shape")
+            points = jnp.empty((*p[0].shape, len(points)), dtype=float)
+            for j, item in enumerate(p):
+                points = points.at[..., j].set(item)
+        else:
+            points = jnp.asarray(points)  # SciPy: asanyarray(points)
+            if points.ndim == 1:
+                if ndim is None:
+                    points = points.reshape(-1, 1)
+                else:
+                    points = points.reshape(-1, ndim)
+        return points
+
+    def __init__(self, points, values, **kwargs):
+        kwargs.setdefault("method", "linear")
+        kwargs.setdefault("bounds_error", False)
+        kwargs.setdefault("fill_value", 0.0)
+
+        super().__init__(points, values, **kwargs)
+
+    def value_and_grad(self, xi):
+        ndim = len(self.grid)
+        xi = self._ndim_coords_from_arrays(xi)
+        xi_shape = xi.shape
+        xi = xi.reshape(-1, xi_shape[-1])
+
+        indices, norm_distances, out_of_bounds = self._find_indices(xi.T)
+
+        vslice = (slice(None),) + (None,) * (self.values.ndim - len(indices))
+
+        # find relevant values
+        # each i and i+1 represents a edge
+        edges = product(*[[i, i + 1] for i in indices])
+        result = jnp.asarray(0.0)
+        for edge_indices in edges:
+            weight = jnp.asarray(1.0)
+            for ei, i, yi in zip(edge_indices, indices, norm_distances):
+                weight *= jnp.where(ei == 1, 1 - yi, yi)
+            result += self.values[edge_indices] * weight[vslice]
+
+        if not self.bounds_error and self.fill_value is not None:
+            bc_shp = result.shape[:1] + (1,) * (result.ndim - 1)
+            result = jnp.where(out_of_bounds.reshape(bc_shp), self.fill_value, result)
+
+        return result.reshape(xi_shape[:-1] + self.values.shape[ndim:])

jaxspec/experimental/intrument_models.py

@@ -0,0 +1,159 @@
+import jax.numpy as jnp
+import numpy as np
+import numpyro
+import numpyro.distributions as dist
+
+from scipy.interpolate import BSpline
+from tinygp import GaussianProcess, kernels
+
+from ..model.instrument import GainModel, ShiftModel
+
+
+def bspline_basis(n_basis: int, degree: int = 3, interval=(0.0, 1.0)):
+    """
+    Construct an open-uniform B-spline basis on a given interval.
+
+    Parameters
+    ----------
+    n_basis : int
+        Number of basis functions (X in the prompt). Must satisfy
+        n_basis >= degree + 1 for an open-uniform knot vector.
+    degree : int, optional
+        Polynomial degree of the splines (default 3 → cubic).
+    interval : tuple(float, float), optional
+        The (start, end) of the domain (default (0, 1)).
+
+    Returns:
+    -------
+    basis : list[BSpline]
+        List of `n_basis` BSpline objects forming a basis over the interval.
+    knots : ndarray
+        The full knot vector, including endpoint multiplicities.
+    """
+    a, b = interval
+    p = degree
+    if n_basis < p + 1:
+        raise ValueError(f"Need at least {p+1} basis functions (got {n_basis}).")
+
+    # How many *internal* knots (not counting the duplicated endpoints)?
+    n_internal = n_basis - p - 1  # open-uniform formula
+
+    # Equally spaced internal knots (could be user-supplied instead)
+    internal_knots = (
+        np.linspace(a, b, n_internal + 2)[1:-1]  # drop the two ends
+        if n_internal > 0
+        else np.empty(0)
+    )
+
+    # Open-uniform knot vector: endpoints repeated p+1 times
+    knots = np.concatenate((np.full(p + 1, a), internal_knots, np.full(p + 1, b)))
+
+    # Coefficient matrix: each column of I generates one basis spline
+    coeffs = np.eye(n_basis)
+
+    # Build BSpline objects
+    basis = [BSpline(knots, coeffs[i], p, extrapolate=False) for i in range(n_basis)]
+    return basis, knots
+
+
+class GaussianProcessGain(GainModel):
+    def __init__(self, e_min, e_max, n_nodes=30):
+        # self.prior_distribution = prior_distribution
+        self.e_min = e_min
+        self.e_max = e_max
+        self.n_nodes = n_nodes
+        self.kernel = kernels.Matern52
+
+    def numpyro_model(self, observation_name: str):
+        mean = numpyro.sample(f"ins/~/_{observation_name}_mean", dist.Normal(1.0, 0.3))
+
+        sigma = numpyro.sample(f"ins/~/_{observation_name}_sigma", dist.HalfNormal(3.0))
+        rho = numpyro.sample(f"ins/~/_{observation_name}_rho", dist.HalfNormal(10.0))
+
+        # Set up the kernel and GP objects
+        kernel = sigma**2 * self.kernel(rho)
+        nodes = jnp.linspace(0, 1, self.n_nodes)
+        gp = GaussianProcess(kernel, nodes, diag=1e-5 * jnp.ones_like(nodes), mean=mean)
+
+        gain_sample = numpyro.sample(f"ins/~/{observation_name}_gain", gp.numpyro_dist())
+
+        def gain(energy):
+            return jnp.interp(
+                energy.mean(axis=0),
+                nodes * (self.e_max - self.e_min) + self.e_min,
+                gain_sample,
+                left=1.0,
+                right=1.0,
+            )
+
+        return gain
+
+
+class BsplineGain(GainModel):
+    def __init__(self, e_min, e_max, n_nodes=6, grid_size=30):
+        self.e_min = e_min
+        self.e_max = e_max
+        self.n_nodes = n_nodes
+        self.egrid = jnp.linspace(e_min, e_max, grid_size)
+
+        basis, knots = bspline_basis(n_nodes, 3, (e_min, e_max))
+
+        self.gridded_basis = jnp.asarray([bi(self.egrid) for bi in basis])
+
+    def numpyro_model(self, observation_name: str):
+        coeff = numpyro.sample(
+            f"ins/~/_{observation_name}_coeff",
+            dist.Uniform(0 * jnp.ones(self.n_nodes), 2 * jnp.ones(self.n_nodes)),
+        )
+
+        def gain(energy):
+            gridded_gain = jnp.dot(coeff, self.gridded_basis)
+
+            return jnp.interp(energy.mean(axis=0), self.egrid, gridded_gain, left=1.0, right=1.0)
+
+        return gain
+
+
+class PolynomialGain(GainModel):
+    def __init__(self, prior_distribution):
+        self.prior_distribution = prior_distribution
+        distribution_shape = prior_distribution.shape()
+        self.degree = distribution_shape[0] if len(distribution_shape) > 0 else 0
+
+    def numpyro_model(self, observation_name: str):
+        polynomial_coefficient = numpyro.sample(
+            f"ins/~/gain_{observation_name}", self.prior_distribution
+        )
+
+        if self.degree == 0:
+
+            def gain(energy):
+                return polynomial_coefficient
+
+        else:
+
+            def gain(energy):
+                return jnp.polyval(polynomial_coefficient, energy.mean(axis=0))
+
+        return gain
+
+
+class PolynomialShift(ShiftModel):
+    def __init__(self, prior_distribution):
+        self.prior_distribution = prior_distribution
+        distribution_shape = prior_distribution.shape()
+        self.degree = distribution_shape[0] if len(distribution_shape) > 0 else 0
+
+    def numpyro_model(self, observation_name: str):
+        polynomial_coefficient = numpyro.sample(
+            f"ins/~/shift_{observation_name}", self.prior_distribution
+        )
+
+        if self.degree == 0:
+            # ensure that new_energy = energy + constant
+            polynomial_coefficient = jnp.asarray([1.0, polynomial_coefficient])
+
+        def shift(energy):
+            return jnp.polyval(polynomial_coefficient, energy)
+
+        return shift