relaxed.ops package#

relaxed.ops.cramer_rao_uncert(model: Any, pars: dict[str, jax.Array], data: Array, return_tree=True) → Array#

Approximate uncertainties on MLE parameters for a model with a logpdf method. Defined as the square root of the diagonal of the Fisher information matrix, valid via the Cramer-Rao lower bound.

Parameters:

  • model (Any) – The model to compute the Cramer-Rao uncertainty for. Needs to have a logpdf method (that returns list[float] for now).

  • pars (Array) – The (MLE) parameters of the model.

  • data (Array) – The data to compute the uncertainty for.

Returns:

Cramer-Rao uncertainty on the MLE parameters.

Return type:

Array

relaxed.ops.cut(data: Array, cut_val: float, slope: float = 1.0, keep: str = 'above') → Array#

Use a sigmoid function as an approximate cut. Same as a hard cut in the limit of infinite slope. Note: this function returns weights, not indices.

Parameters:

  • data (Array) – The data to cut.

  • cut_val (float) – The value to cut at.

  • slope (float) – The slope of the sigmoid function.

  • keep (str, optional) – Whether to keep the data above or below the cut. One of: - “above” (default) - “below”

Returns:

Weighted yields of the data after the cut.

Return type:

Array

relaxed.ops.fisher_info(model: Any, pars: dict[str, jax.Array], data: Array) → Array#

Fisher information matrix for a model with a logpdf method.

Parameters:

  • model (Any) – The model to compute the Fisher information matrix for. Needs to have a logpdf method (that returns list[float] for now).

  • pars (dict[str, Array]) – The (MLE) parameters of the model, as a dict of arrays/floats.

  • data (Array) – The data to compute the Fisher information matrix for.

Returns:

Fisher information matrix of shape (num_pars, num_pars). Order of columns is the same as the order of the parameters in pars. Parameters with multiple dimensions are flattened into their own columns.

Return type:

Array

relaxed.ops.hist(data: Array, bins: Array, bandwidth: float, density: bool = False, reflect_infinities: bool = False) → Array#

Differentiable histogram, defined via a binned kernel density estimate (bKDE).

Parameters:

  • data (Array) – 1D array of data to histogram.

  • bins (Array) – 1D array of bin edges.

  • bandwidth (float) – The bandwidth of the kernel. Bigger == lower gradient variance, but more bias.

  • density (bool) – Normalise the histogram to unit area.

  • reflect_infinities (bool) – If True, define bins at +/- infinity, and reflect their mass into the edge bins.

Returns:

1D array of bKDE counts.

Return type:

Array