gwkokab.analysis.utils.marginals¶

Classes¶

PlotStyle

A named tuple representing the style for plotting marginal densities, including

Functions¶

calculate_dist_layouts(→ list[tuple[int, Ellipsis]])

Calculate the layout of the distribution's support based on its shaped values.

calculate_marginals_over_axes(→ list[jaxtyping.Array])

Calculate marginal densities for each axis of a joint density array.

compute_batched_marginals(model_meta_cls, ...[, ...])

Compute marginal densities.

generate_marginal_probs(model_meta_cls, ...[, ...])

Generate marginal probability densities.

plot_marginal_with_intervals(ax, filename, parameter, ...)

Plot marginal densities with confidence intervals for a specified parameter.

plot_with_intervals(ax, xx, yy, style)

Help function to plot densities with confidence intervals.

remove_comoving_volume_factor(→ jaxtyping.Array)

Remove the comoving volume factor from a marginal density over redshift.

save_results_to_hdf5(constants, variables_index, ...)

Save the computed marginal densities to an HDF5 file.

Module Contents¶

class gwkokab.analysis.utils.marginals.PlotStyle¶

Bases: NamedTuple

A named tuple representing the style for plotting marginal densities, including the color, label, and additional keyword arguments for line plots and fill-between plots.

color: str¶
fill_between_kwargs: dict¶
label: str¶
line_plot_kwargs: dict¶
gwkokab.analysis.utils.marginals.calculate_dist_layouts(shaped_values: list[tuple[int, int] | int]) → list[tuple[int, Ellipsis]]¶

Calculate the layout of the distribution’s support based on its shaped values.

Parameters:

shaped_values (list[tuple[int, int] | int]) – A list of shaped values for the distribution. Each element can be either a tuple representing a multi-dimensional support (with the second element indicating the number of dimensions), or an integer representing a one-dimensional support.

Returns:

A list of tuples representing the layout of the distribution’s support. Each tuple contains the indices of the dimensions that correspond to a particular component of the distribution.

Return type:

list[tuple[int, …]]

gwkokab.analysis.utils.marginals.calculate_marginals_over_axes(probs: jaxtyping.Array, domains: list[jaxtyping.Array], normalize: list[bool] | None = None) → list[jaxtyping.Array]¶

Calculate marginal densities for each axis of a joint density array.

This function iteratively integrates out all dimensions except the specified axis to derive the marginal densities for each axis.

Parameters:

  • probs (Array) – An array representing the joint density over multiple dimensions.

  • domains (list[Array]) – A list of arrays representing the domain values for each dimension of the joint density. The length of this list should match the number of dimensions in probs.

  • normalize (list[bool] | None, optional) – A list of booleans indicating whether to normalize the marginal densities for each dimension. If None, all dimensions will be normalized, by default None

Returns:

A list of arrays representing the marginal densities for each dimension.

Return type:

list[Array]

gwkokab.analysis.utils.marginals.compute_batched_marginals(model_meta_cls: type, samples_batch: jaxtyping.Array, constants: dict, variables_index: dict, domains: list[jaxtyping.Array], normalize: list[bool], batch_size: int | None = None)¶

Compute marginal densities.

Parameters:

  • model_meta_cls (type) – A class representing the meta-information of the model, which includes a method for constructing the model given specific parameters.

  • samples_batch (Array) – A batch of samples, where each row corresponds to a single sample and each column corresponds to a specific parameter of the model.

  • constants (dict) – A dictionary of constant values required for constructing the model. The keys should match the parameter names expected by the model’s constructor.

  • variables_index (dict) – A dictionary mapping parameter names to their corresponding column indices in the samples_batch array. This mapping is used to extract the relevant parameter values from the samples when constructing the model.

  • domains (list[Array]) – A list of arrays representing the domain values for each parameter of the model. The length of this list should match the number of parameters in the model.

  • normalize (list[bool], optional) – A list of booleans indicating whether to normalize the marginal densities for each parameter.

  • batch_size (int | None, optional) – The size of the batch to process at a time, by default None

gwkokab.analysis.utils.marginals.generate_marginal_probs(model_meta_cls: type, input_file_path: str | pathlib.Path, output_file_path: str | pathlib.Path, domain_cfg: dict[str, tuple[float, float, int]], max_samples: int | None = None, batch_size: int | None = None)¶

Generate marginal probability densities.

Parameters:

  • model_meta_cls (type) – A class representing the meta-information of the model, which includes a method for constructing the model given specific parameters.

  • input_file_path (str | Path) – Path to hdf5 file containing the input data.

  • output_file_path (str | Path) – The path to the HDF5 file where the computed marginal densities will be saved.

  • domain_cfg (dict[str, tuple[float, float, int]]) – A dictionary mapping parameter names to their corresponding domain specifications. Each value in the dictionary should be a tuple containing the start, stop, and number of points for the domain of the parameter.

  • filename (str) – The path to the HDF5 file where the results will be saved.

  • max_samples (int | None, optional) – The maximum number of samples to use for computing marginal densities, by default None

  • batch_size (int | None, optional) – The batch size for computing marginal densities, by default None

Raises:

FileNotFoundError – If the required samples file cannot be found in the specified base directory.

gwkokab.analysis.utils.marginals.plot_marginal_with_intervals(ax: matplotlib.pyplot.Axes, filename: str, parameter: str, style: PlotStyle, component_idxs: list[int], scale: float | Callable = 1.0, weights: list[float | Callable] | None = None, normalize: bool = False)¶

Plot marginal densities with confidence intervals for a specified parameter.

Parameters:

  • ax (plt.Axes) – The Matplotlib Axes object on which to plot the marginal densities and confidence intervals.

  • filename (str) – The path to the HDF5 file containing the marginal density data.

  • parameter (str) – The name of the parameter for which to plot the marginal densities. This should correspond to a dataset in the HDF5 file under the “probs/component_{i}” groups.

  • style (PlotStyle) – A list of PlotStyle objects specifying the plotting style for each component’s marginal density. If an element is None, the corresponding component will be skipped in the plot.

  • component_idxs (list[int]) – A list of indices specifying which components to plot.

  • scale (float | Callable, optional) – A scaling factor for the marginal densities. If a callable is provided, it will be evaluated with the parameters from the HDF5 file, by default 1.0

  • weights (list[float | Callable] | None, optional) – The weights for each component’s marginal density. If None, equal weights are assumed, by default None

  • normalize (bool, optional) – Whether to normalize the marginal densities, by default False

gwkokab.analysis.utils.marginals.plot_with_intervals(ax: matplotlib.pyplot.Axes, xx: numpy.ndarray, yy: numpy.ndarray, style: PlotStyle)¶

Help function to plot densities with confidence intervals.

Parameters:

  • ax (plt.Axes) – Matplotlib Axes object on which to plot the densities and confidence intervals.

  • xx (np.ndarray) – X-axis values of shape (N,) representing the domain of the density.

  • yy (np.ndarray) – Y-axis values of shape (N, M) representing the densities for M samples.

  • style (PlotStyle) – A PlotStyle object specifying the plotting style for the density and confidence intervals.

gwkokab.analysis.utils.marginals.remove_comoving_volume_factor(marginal_density: jaxtyping.Array, redshift_domain: jaxtyping.Array) → jaxtyping.Array¶

Remove the comoving volume factor from a marginal density over redshift.

This function takes a marginal density that includes the comoving volume factor and time dilation factor and divides it by the comoving volume element to obtain the underlying density without the volume factor.

Parameters:

  • marginal_density (Array) – The marginal density over redshift that includes the comoving volume factor.

  • redshift_domain (Array) – The array of redshift values corresponding to the marginal density. This is used to compute the comoving volume element.

Returns:

The marginal density with the comoving volume and time dilation factors removed, representing the underlying density over redshift.

Return type:

Array

gwkokab.analysis.utils.marginals.save_results_to_hdf5(constants: dict, variables_index: dict[str, int], samples: jaxtyping.Array, batched_results: list[list[list[jaxtyping.Array]]], parameters: list[str], domain_cfg: dict[str, tuple[float, float, int]], filepath: str | pathlib.Path)¶

Save the computed marginal densities to an HDF5 file.

Parameters:

  • constants (dict) – A dictionary of constant values required for constructing the model. The keys should match the parameter names expected by the model’s constructor.

  • variables_index (dict[str, int]) – A dictionary mapping parameter names to their corresponding column indices in the samples array. This mapping is used to extract the relevant parameter values from the samples when constructing the model.

  • samples (Array) – An array of samples used for computing the marginal densities. Each row corresponds to a single sample, and each column corresponds to a specific parameter of the model.

  • batched_results (list[list[Array]]) – A nested list containing the computed marginal densities for each component of the model. The outer list corresponds to the components of the model, the inner list corresponds to the parameters of the model, and each leaf is a 2D array of shape (num_samples, domain_size).

  • parameters (list[str]) – A list of parameter names for the model.

  • domain_cfg (dict[str, tuple[float, float, int]]) – A dictionary mapping parameter names to their corresponding domain specifications.

  • filepath (str | Path) – The path to the HDF5 file where the results will be saved.