revert

Author: StephenNneji

ratapi/utils/plotting.py

@@ -906,6 +906,262 @@ def _y_update_offset_text_position(axis, _bboxes, bboxes2):
     axis.offsetText.set_position((x - x_offset, y))
 
 
+@assert_bayesian("Corner")
+def plot_corner(
+    results: ratapi.outputs.BayesResults,
+    params: list[int | str] | None = None,
+    smooth: bool = True,
+    block: bool = False,
+    fig: matplotlib.figure.Figure | None = None,
+    return_fig: bool = False,
+    hist_kwargs: dict | None = None,
+    hist2d_kwargs: dict | None = None,
+    progress_callback: Callable[[int, int], None] | None = None,
+):
+    """Create a corner plot from a Bayesian analysis.
+
+    Parameters
+    ----------
+    results : BayesResults
+        The results from a Bayesian calculation.
+    params : list[int or str], default None
+        The indices or names of a subset of parameters if required.
+        If None, uses all indices.
+    smooth : bool, default True
+        Whether to apply Gaussian smoothing to the corner plot.
+    block : bool, default False
+        Whether Python should block until the plot is closed.
+    fig : matplotlib.figure.Figure, optional
+        The figure object to use for plot.
+    return_fig: bool, default False
+        If True, return the figure as an object instead of showing it.
+    hist_kwargs : dict
+        Extra keyword arguments to pass to the 1d histograms.
+        Default is {'density': True, 'bins': 25}
+    hist2d_kwargs : dict
+        Extra keyword arguments to pass to the 2d histograms.
+        Default is {'density': True, 'bins': 25}
+    progress_callback: Union[Callable[[int, int], None], None]
+        Callback function for providing progress during plot creation
+        First argument is current completed sub plot and second is total number of sub plots
+
+    Returns
+    -------
+    Figure or None
+        If `return_fig` is True, return the figure - otherwise, return nothing.
+
+    """
+    fitname_to_index = partial(name_to_index, names=results.fitNames)
+
+    if params is None:
+        params = range(0, len(results.fitNames))
+    else:
+        params = list(map(fitname_to_index, params))
+
+    # defaults are applied inside each function - just pass blank dicts for now
+    if hist_kwargs is None:
+        hist_kwargs = {}
+    if hist2d_kwargs is None:
+        hist2d_kwargs = {}
+
+    num_params = len(params)
+    total_count = num_params + (num_params**2 - num_params) // 2
+
+    if fig is None:
+        fig, axes = plt.subplots(num_params, num_params, figsize=(11, 10), subplot_kw={"visible": False})
+    else:
+        fig.clf()
+        axes = fig.subplots(num_params, num_params, subplot_kw={"visible": False})
+
+    # i is row, j is column
+    current_count = 0
+    for i in range(num_params):
+        for j in range(i + 1):
+            row_param = params[i]
+            col_param = params[j]
+            current_axes: Axes = axes if isinstance(axes, matplotlib.axes.Axes) else axes[i][j]
+            current_axes.tick_params(which="both", labelsize="medium")
+            current_axes.xaxis.offsetText.set_fontsize("small")
+            current_axes.yaxis.offsetText.set_fontsize("small")
+            current_axes.set_visible(True)
+            if i == j:  # diagonal: histograms
+                plot_one_hist(results, param=row_param, smooth=smooth, axes=current_axes, **hist_kwargs)
+            elif i > j:  # lower triangle: 2d histograms
+                plot_contour(
+                    results, x_param=col_param, y_param=row_param, smooth=smooth, axes=current_axes, **hist2d_kwargs
+                )
+
+            # remove label if on inside of corner plot
+            if j != 0:
+                current_axes.get_yaxis().set_visible(False)
+            if i != len(params) - 1:
+                current_axes.get_xaxis().set_visible(False)
+            # make labels invisible as titles cover that
+            current_axes.yaxis._update_offset_text_position = types.MethodType(
+                _y_update_offset_text_position, current_axes.yaxis
+            )
+            current_axes.yaxis.offset_text_position = "center"
+            current_axes.set_ylabel("")
+            current_axes.set_xlabel("")
+            if progress_callback is not None:
+                current_count += 1
+                progress_callback(current_count, total_count)
+    if return_fig:
+        return fig
+    plt.show(block=block)
+
+
+@assert_bayesian("Histogram")
+def plot_one_hist(
+    results: ratapi.outputs.BayesResults,
+    param: int | str,
+    smooth: bool = True,
+    sigma: float | None = None,
+    estimated_density: Literal["normal", "lognor", "kernel", None] = None,
+    axes: Axes | None = None,
+    block: bool = False,
+    return_fig: bool = False,
+    **hist_settings,
+):
+    """Plot the marginalised posterior for a parameter of a Bayesian analysis.
+
+    Parameters
+    ----------
+    results : BayesResults
+        The results from a Bayesian calculation.
+    param : Union[int, str]
+        Either the index or name of a parameter.
+    block : bool, default False
+        Whether Python should block until the plot is closed.
+    smooth : bool, default True
+        Whether to apply Gaussian smoothing to the histogram.
+        Defaults to True.
+    sigma: float or None, default None
+        If given, is used as the sigma-parameter for the Gaussian smoothing.
+        If None, the default (1/3rd of parameter chain standard deviation) is used.
+    estimated_density : 'normal', 'lognor', 'kernel' or None, default None
+        If None (default), ignore. Else, add an estimated density
+        of the given form on top of the histogram by the following estimations:
+        'normal': normal Gaussian.
+        'lognor': Log-normal probability density.
+        'kernel': kernel density estimation.
+    axes: Axes or None, default None
+        If provided, plot on the given Axes object.
+    block : bool, default False
+        Whether Python should block until the plot is closed.
+    return_fig: bool, default False
+        If True, return the figure as an object instead of showing it.
+    **hist_settings :
+        Settings passed to `np.histogram`. By default, the settings
+        passed are `bins = 25` and `density = True`.
+
+    Returns
+    -------
+    Figure or None
+        If `return_fig` is True, return the figure - otherwise, return nothing.
+
+    """
+    chain = results.chain
+    param = name_to_index(param, results.fitNames)
+
+    if axes is None:
+        fig, axes = plt.subplots(1, 1)
+    else:
+        fig = None
+
+    # apply default settings if not set by user
+    default_settings = {"bins": 25, "density": True}
+    hist_settings = {**default_settings, **hist_settings}
+
+    parameter_chain = chain[:, param]
+    counts, bins = np.histogram(parameter_chain, **hist_settings)
+    mean_y = np.mean(parameter_chain)
+    sd_y = np.std(parameter_chain)
+
+    if smooth:
+        if sigma is None:
+            sigma = sd_y / 2
+        counts = gaussian_filter1d(counts, sigma)
+    axes.hist(
+        bins[:-1],
+        bins,
+        weights=counts,
+        edgecolor="black",
+        linewidth=1.2,
+        color="white",
+    )
+
+    axes.set_title(results.fitNames[param], loc="left", fontsize="medium")
+
+    if estimated_density:
+        dx = bins[1] - bins[0]
+        if estimated_density == "normal":
+            t = np.linspace(mean_y - 3.5 * sd_y, mean_y + 3.5 * sd_y)
+            axes.plot(t, norm.pdf(t, loc=mean_y, scale=sd_y**2))
+        elif estimated_density == "lognor":
+            t = np.linspace(bins[0] - 0.5 * dx, bins[-1] + 2 * dx)
+            axes.plot(t, lognorm.pdf(t, np.mean(np.log(parameter_chain)), np.std(np.log(parameter_chain))))
+        elif estimated_density == "kernel":
+            t = np.linspace(bins[0] - 2 * dx, bins[-1] + 2 * dx, 200)
+            kde = gaussian_kde(parameter_chain)
+            axes.plot(t, kde.evaluate(t))
+        else:
+            raise ValueError(
+                f"{estimated_density} is not a supported estimated density function."
+                " Supported functions are 'normal' 'lognor' or 'kernel'."
+            )
+
+    # adding the estimated density extends the figure range - reset it to histogram range
+    x_range = hist_settings.get("range", (parameter_chain.min(), parameter_chain.max()))
+    axes.set_xlim(x_range)
+
+    if fig is not None:
+        if return_fig:
+            return fig
+        plt.show(block=block)
+
+
+def _y_update_offset_text_position(axis, _bboxes, bboxes2):
+    """Update the position of the Y axis offset text using the provided bounding boxes.
+
+    Adapted from https://github.com/matplotlib/matplotlib/issues/4476#issuecomment-105627334.
+
+    Parameters
+    ----------
+    axis : matplotlib.axis.YAxis
+        Y axis to update.
+    _bboxes : List
+        list of bounding boxes
+    bboxes2 : List
+        list of bounding boxes
+    """
+    x, y = axis.offsetText.get_position()
+
+    if axis.offset_text_position == "left":
+        # y in axes coords, x in display coords
+        axis.offsetText.set_transform(
+            mtransforms.blended_transform_factory(axis.axes.transAxes, mtransforms.IdentityTransform())
+        )
+
+        top = axis.axes.bbox.ymax
+        y = top + axis.OFFSETTEXTPAD * axis.figure.dpi / 72.0
+
+    else:
+        # x & y in display coords
+        axis.offsetText.set_transform(mtransforms.IdentityTransform())
+
+        # Northwest of upper-right corner of right-hand extent of tick labels
+        if bboxes2:
+            bbox = mtransforms.Bbox.union(bboxes2)
+        else:
+            bbox = axis.axes.bbox
+        center = bbox.ymin + (bbox.ymax - bbox.ymin) / 2
+        x = bbox.xmin - axis.OFFSETTEXTPAD * axis.figure.dpi / 72.0
+        y = center
+    x_offset = 110
+    axis.offsetText.set_position((x - x_offset, y))
+
+
 @assert_bayesian("Contour")
 def plot_contour(
     results: ratapi.outputs.BayesResults,
@@ -982,7 +1238,10 @@ def plot_contour(
 
 
 def panel_plot_helper(
-    plot_func: Callable, indices: list[int], fig: matplotlib.figure.Figure | None = None
+    plot_func: Callable,
+    indices: list[int],
+    fig: matplotlib.figure.Figure | None = None,
+    progress_callback: Callable[[int, int], None] | None = None,
 ) -> matplotlib.figure.Figure:
     """Generate a panel-based plot from a single plot function.
 
@@ -994,6 +1253,9 @@ def panel_plot_helper(
         The list of indices to pass into ``plot_func``.
     fig : matplotlib.figure.Figure, optional
         The figure object to use for plot.
+    progress_callback: Union[Callable[[int, int], None], None]
+        Callback function for providing progress during plot creation
+        First argument is current completed sub plot and second is total number of sub plots
 
     Returns
     -------
@@ -1005,21 +1267,19 @@ def panel_plot_helper(
     nrows, ncols = ceil(sqrt(nplots)), round(sqrt(nplots))
 
     if fig is None:
-        fig = plt.subplots(nrows, ncols, figsize=(11, 10))[0]
+        fig = plt.subplots(nrows, ncols, figsize=(11, 10), subplot_kw={"visible": False})[0]
     else:
         fig.clf()
-        fig.subplots(nrows, ncols)
+        fig.subplots(nrows, ncols, subplot_kw={"visible": False})
     axs = fig.get_axes()
-
-    for plot_num, index in enumerate(indices):
-        axs[plot_num].tick_params(which="both", labelsize="medium")
-        axs[plot_num].xaxis.offsetText.set_fontsize("small")
-        axs[plot_num].yaxis.offsetText.set_fontsize("small")
-        plot_func(axs[plot_num], index)
-
-    # blank unused plots
-    for i in range(nplots, len(axs)):
-        axs[i].set_visible(False)
+    for index, plot_num in enumerate(indices):
+        axs[index].tick_params(which="both", labelsize="medium")
+        axs[index].xaxis.offsetText.set_fontsize("small")
+        axs[index].yaxis.offsetText.set_fontsize("small")
+        axs[index].set_visible(True)
+        plot_func(axs[index], plot_num)
+        if progress_callback is not None:
+            progress_callback(index, nplots)
 
     fig.tight_layout()
     return fig
@@ -1036,6 +1296,7 @@ def plot_hists(
     block: bool = False,
     fig: matplotlib.figure.Figure | None = None,
     return_fig: bool = False,
+    progress_callback: Callable[[int, int], None] | None = None,
     **hist_settings,
 ):
     """Plot marginalised posteriors for several parameters from a Bayesian analysis.
@@ -1072,6 +1333,9 @@ def plot_hists(
         The figure object to use for plot.
     return_fig: bool, default False
         If True, return the figure as an object instead of showing it.
+    progress_callback: Union[Callable[[int, int], None], None]
+        Callback function for providing progress during plot creation
+        First argument is current completed sub plot and second is total number of sub plots
     hist_settings :
         Settings passed to `np.histogram`. By default, the settings
         passed are `bins = 25` and `density = True`.
@@ -1130,6 +1394,7 @@ def validate_dens_type(dens_type: str | None, param: str):
         ),
         params,
         fig,
+        progress_callback,
     )
     if return_fig:
         return fig
@@ -1144,6 +1409,7 @@ def plot_chain(
     block: bool = False,
     fig: matplotlib.figure.Figure | None = None,
     return_fig: bool = False,
+    progress_callback: Callable[[int, int], None] | None = None,
 ):
     """Plot the MCMC chain for each parameter of a Bayesian analysis.
 
@@ -1162,6 +1428,9 @@ def plot_chain(
         The figure object to use for plot.
     return_fig: bool, default False
         If True, return the figure as an object instead of showing it.
+    progress_callback: Union[Callable[[int, int], None], None]
+        Callback function for providing progress during plot creation
+        First argument is current completed sub plot and second is total number of sub plots
 
     Returns
     -------
@@ -1187,7 +1456,7 @@ def plot_one_chain(axes: Axes, i: int):
         axes.plot(range(0, nsimulations, skip), chain[:, i][0:nsimulations:skip])
         axes.set_title(results.fitNames[i], fontsize="small")
 
-    fig = panel_plot_helper(plot_one_chain, params, fig=fig)
+    fig = panel_plot_helper(plot_one_chain, params, fig, progress_callback)
     if return_fig:
         return fig
     plt.show(block=block)