Adds progress callback to corner plot and updates example to use scipy

Author: StephenNneji

ratapi/examples/normal_reflectivity/custom_XY_DSPC.py

@@ -1,8 +1,7 @@
 """A custom XY model for a supported DSPC bilayer."""
 
-import math
-
 import numpy as np
+from scipy.special import erf
 
 
 def custom_XY_DSPC(params, bulk_in, bulk_out, contrast):
@@ -114,7 +113,7 @@ def custom_XY_DSPC(params, bulk_in, bulk_out, contrast):
     totSLD = sldSilicon + sldOxide + sldHeadL + sldTails + sldHeadR + sldWat
 
     # Make the SLD array for output
-    SLD = [[a, b] for (a, b) in zip(z, totSLD)]
+    SLD = np.column_stack((z, totSLD))
 
     return SLD, subRough
 
@@ -132,8 +131,8 @@ def layer(z, prevLaySurf, thickness, height, Sigma_L, Sigma_R):
     a = (z - left) / ((2**0.5) * Sigma_L)
     b = (z - right) / ((2**0.5) * Sigma_R)
 
-    erf_a = np.array([math.erf(value) for value in a])
-    erf_b = np.array([math.erf(value) for value in b])
+    erf_a = erf(a)
+    erf_b = erf(b)
 
     VF = np.array((height / 2) * (erf_a - erf_b))
 

ratapi/utils/plotting.py

@@ -652,6 +652,7 @@ def plot_corner(
     return_fig: bool = False,
     hist_kwargs: Union[dict, None] = None,
     hist2d_kwargs: Union[dict, None] = None,
+    progress_callback: Union[Callable[[int, int], None], None] = None,
 ):
     """Create a corner plot from a Bayesian analysis.
 
@@ -697,29 +698,32 @@ def plot_corner(
         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))
+        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)
+        axes = fig.subplots(num_params, num_params, subplot_kw={"visible": False})
 
     # i is row, j is column
-    for i, row_param in enumerate(params):
-        for j, col_param in enumerate(params):
-            current_axes: Axes = axes[i][j]
+    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
                 )
-            elif i < j:  # upper triangle: no plot
-                current_axes.set_visible(False)
+
             # remove label if on inside of corner plot
             if j != 0:
                 current_axes.get_yaxis().set_visible(False)
@@ -732,6 +736,9 @@ def plot_corner(
             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)
@@ -1153,7 +1160,7 @@ def plot_chain(
 
     """
     chain = results.chain
-    nsimulations, nplots = chain.shape
+    nsimulations, _ = chain.shape
     # skip is to evenly distribute points plotted
     # all points will be plotted if maxpoints < nsimulations
     skip = max(floor(nsimulations / maxpoints), 1)