devinterp.vis_utils package

Module contents

class devinterp.vis_utils.EpsilonBetaAnalyzer

Bases: object

A class for analyzing and visualizing the local learning coefficient (LLC) across different epsilon and beta values.

This class provides methods to configure, run, and visualize sweeps of the local learning coefficient over epsilon and beta.

configure_sweep(llc_estimator: Callable, llc_estimator_kwargs: dict, min_epsilon: float | None = 1e-06, max_epsilon: float | None = 0.01, epsilon_samples: float = 8, epsilon_range: List[float] | None = None, min_beta: float | None = None, max_beta: float | None = None, beta_samples: float = 8, beta_range: List[float] | None = None, dataloader: DataLoader | None = None) → None

Configure the sampling parameters for the LLC analysis.

Parameters:

• llc_estimator (Callable) –

  Callable function to estimate the local learning coefficient. Must have the signature:

  def estimator(epsilon: float, beta: float, **kwargs) -> Dict[str, Any]
  

  where kwargs are the arguments to estimate_learning_coeff_with_summary.

  The return value must be a dict with a "llc/trace" key corresponding to a numpy array of shape (num_chains, num_draws) Additional keys can represent other values of interest (e.g. acceptance rates, true LLC, etc.)

• llc_estimator_kwargs (Dict[str, Any]) – Keyword arguments for the llc_estimator function.

• min_epsilon (Optional[float]) – Minimum value for epsilon range (if epsilon_range not provided).

• max_epsilon (Optional[float]) – Maximum value for epsilon range (if epsilon_range not provided).

• epsilon_samples (int) – Number of samples in epsilon range (if epsilon_range not provided).

• epsilon_range (Optional[List[float]]) – Explicit range of epsilon values to use (overrides min/max_epsilon).

• min_beta (Optional[float]) – Minimum value for beta range (if beta_range not provided).

• max_beta (Optional[float]) – Maximum value for beta range (if beta_range not provided).

• beta_samples (int) – Number of samples in beta range (if beta_range not provided).

• beta_range (Optional[List[float]]) – Explicit range of beta values to use (overrides min/max_beta).

• dataloader (Optional[torch.utils.data.DataLoader]) – Optional dataloader for calculating optimal beta.

plot(true_lambda: float | int | str | Container | None = None, num_last_steps_to_average: int = 50, color: str | None = None, slider: str | None = None, slider_plane: str | None = False, **kwargs) → Figure

Plot the results of the LLC sweep.

Parameters:

• true_lambda (float | int | str | Container) – True value of lambda for comparison (optional). Can be a scalar, a list, or a string column name of sweep_df. Will plot a horizontal plane at the true_lambda value.

• num_last_steps_to_average (int) – Number of last steps to average for final LLC value.

• color (Optional[str]) – Column name to use for coloring the scatter points.

• slider (Optional[str]) – Column name to use for creating a slider in the plot.

• slider_plane (bool) – If True, adds a plane for each slider value.

• kwargs (Any) – Additional keyword arguments to pass to the plotting function f.e. range_color=[0, 0.15].

Returns:

A plotly Figure object containing the LLC sweep visualization.

Return type:

go.Figure

Raises:

AssertionError – If no data is available to plot.

save_fig(path: str) → None

Save the current figure to a file.

Parameters:

path (str) – Path to save the figure to.

Raises:

AssertionError – If no figure is available to save.

sweep(add_to_existing=False) → None

Perform the LLC sweep using the configured parameters.

This method runs the LLC estimator for each combination of epsilon and beta values and stores the results in self.sweep_df.

Parameters:

add_to_existing (bool) – If True, adds new sweep results to existing ones. If False, replaces existing results. Useful for sweeping over multiple models or datasets.

Raises:

AssertionError – If sweep configuration is not set.