add n_row and n_col argument where applicable

bayesflow/diagnostics.py

@@ -98,6 +98,10 @@ def plot_recovery(
  A flag for adding R^2 between true and estimates to the plot
  color : str, optional, default: '#8f2727'
  The color for the true vs. estimated scatter points and error bars
+ n_row : int, optional, default: None
+ The number of rows for the subplots. Dynamically determined if None.
+ n_col : int, optional, default: None
+ The number of columns for the subplots. Dynamically determined if None.
  xlabel : str, optional, default: 'Ground truth'
  The label on the x-axis of the plot
  ylabel : str, optional, default: 'Estimated'
@@ -232,7 +236,7 @@ def plot_z_score_contraction(
  tick_fontsize=12,
  color="#8f2727",
  n_col=None,
- n_row=None,
+ n_row=None
 ):
  """Implements a graphical check for global model sensitivity by plotting the posterior
  z-score over the posterior contraction for each set of posterior samples in ``post_samples``
@@ -279,6 +283,10 @@ def plot_z_score_contraction(
  The font size of the axis ticklabels
  color : str, optional, default: '#8f2727'
  The color for the true vs. estimated scatter points and error bars
+ n_row : int, optional, default: None
+ The number of rows for the subplots. Dynamically determined if None.
+ n_col : int, optional, default: None
+ The number of columns for the subplots. Dynamically determined if None.
 
  Returns
  -------
@@ -379,6 +387,8 @@ def plot_sbc_ecdf(
  tick_fontsize=12,
  rank_ecdf_color="#a34f4f",
  fill_color="grey",
+ n_row=None,
+ n_col=None,
  **kwargs,
 ):
  """Creates the empirical CDFs for each marginal rank distribution and plots it against
@@ -419,6 +429,10 @@ def plot_sbc_ecdf(
  The color to use for the rank ECDFs
  fill_color : str, optional, default: 'grey'
  The color of the fill arguments.
+ n_row : int, optional, default: None
+ The number of rows for the subplots. Dynamically determined if None.
+ n_col : int, optional, default: None
+ The number of columns for the subplots. Dynamically determined if None.
  **kwargs : dict, optional, default: {}
  Keyword arguments can be passed to control the behavior of ECDF simultaneous band computation
  through the ``ecdf_bands_kwargs`` dictionary. See `simultaneous_ecdf_bands` for keyword arguments
@@ -447,9 +461,14 @@ def plot_sbc_ecdf(
  n_row, n_col = 1, 1
  f, ax = plt.subplots(1, 1, figsize=fig_size)
  else:
- # Determine n_subplots dynamically
- n_row = int(np.ceil(n_params / 6))
- n_col = int(np.ceil(n_params / n_row))
+ # Determine number of rows and columns for subplots based on inputs
+ if n_row is None and n_col is None:
+ n_row = int(np.ceil(n_params / 6))
+ n_col = int(np.ceil(n_params / n_row))
+ elif n_row is None and n_col is not None:
+ n_row = int(np.ceil(n_params / n_col))
+ elif n_row is not None and n_col is None:
+ n_col = int(np.ceil(n_params / n_row))
 
  # Determine fig_size dynamically, if None
  if fig_size is None:
@@ -543,6 +562,8 @@ def plot_sbc_histograms(
  title_fontsize=18,
  tick_fontsize=12,
  hist_color="#a34f4f",
+ n_row=None,
+ n_col=None
 ):
  """Creates and plots publication-ready histograms of rank statistics for simulation-based calibration
  (SBC) checks according to [1].
@@ -576,6 +597,10 @@ def plot_sbc_histograms(
  The font size of the axis ticklabels
  hist_color : str, optional, default '#a34f4f'
  The color to use for the histogram body
+ n_row : int, optional, default: None
+ The number of rows for the subplots. Dynamically determined if None.
+ n_col : int, optional, default: None
+ The number of columns for the subplots. Dynamically determined if None.
 
  Returns
  -------
@@ -615,9 +640,14 @@ def plot_sbc_histograms(
  if param_names is None:
  param_names = [f"$\\theta_{{{i}}}$" for i in range(1, n_params + 1)]
 
- # Determine n_subplots dynamically
- n_row = int(np.ceil(n_params / 6))
- n_col = int(np.ceil(n_params / n_row))
+ # Determine number of rows and columns for subplots based on inputs
+ if n_row is None and n_col is None:
+ n_row = int(np.ceil(n_params / 6))
+ n_col = int(np.ceil(n_params / n_row))
+ elif n_row is None and n_col is not None:
+ n_row = int(np.ceil(n_params / n_col))
+ elif n_row is not None and n_col is None:
+ n_col = int(np.ceil(n_params / n_row))
 
  # Initialize figure
  if fig_size is None:
@@ -1026,6 +1056,8 @@ def plot_calibration_curves(
  epsilon=0.02,
  fig_size=None,
  color="#8f2727",
+ n_row=None,
+ n_col=None
 ):
  """Plots the calibration curves, the ECEs and the marginal histograms of predicted posterior model probabilities
  for a model comparison problem. The marginal histograms inform about the fraction of predictions in each bin.
@@ -1055,6 +1087,10 @@ def plot_calibration_curves(
  The figure size passed to the ``matplotlib`` constructor. Inferred if ``None``
  color : str, optional, default: '#8f2727'
  The color of the calibration curves
+ n_row : int, optional, default: None
+ The number of rows for the subplots. Dynamically determined if None.
+ n_col : int, optional, default: None
+ The number of columns for the subplots. Dynamically determined if None.
 
  Returns
  -------
@@ -1065,9 +1101,15 @@ def plot_calibration_curves(
  if model_names is None:
  model_names = [rf"$M_{{{m}}}$" for m in range(1, num_models + 1)]
 
- # Determine n_subplots dynamically
- n_row = int(np.ceil(num_models / 6))
- n_col = int(np.ceil(num_models / n_row))
+ # Determine number of rows and columns for subplots based on inputs
+ if n_row is None and n_col is None:
+ n_row = int(np.ceil(num_models / 6))
+ n_col = int(np.ceil(num_models / n_row))
+ elif n_row is None and n_col is not None:
+ n_row = int(np.ceil(num_models / n_col))
+ elif n_row is not None and n_col is None:
+ n_col = int(np.ceil(num_models / n_row))
+
 
  # Compute calibration
  cal_errs, probs_true, probs_pred = expected_calibration_error(true_models, pred_models, num_bins)
@@ -1233,6 +1275,7 @@ def plot_confusion_matrix(
  ax.set_title("Confusion Matrix", fontsize=title_fontsize)
  return fig
 
+
 def plot_mmd_hypothesis_test(
  mmd_null,
  mmd_observed=None,