add histogram to residuals plot and center zero

Author: plakrisenko

petab/visualize/plot_residuals.py

@@ -13,17 +13,16 @@
 from ..calculate import calculate_residuals
 from ..core import get_simulation_df
 from ..problem import Problem
-from ..C import NOISE_DISTRIBUTION, NORMAL, LIN, OBSERVABLE_TRANSFORMATION, \
-    OBSERVABLE_ID
+from ..C import *
 
 __all__ = ['plot_residuals_vs_simulation']
 
 
 def plot_residuals_vs_simulation(
         petab_problem: Problem,
         simulations_df: Union[str, Path, pd.DataFrame],
-        size: Tuple = (10, 7),
-        ax: Optional[plt.Axes] = None
+        size: Optional[Tuple] = (10, 7),
+        axes: Optional[Tuple[plt.Axes, plt.Axes]] = None
 ) -> matplotlib.axes.Axes:
     """
     Plot residuals versus simulation values for measurements with normal noise
@@ -38,7 +37,7 @@ def plot_residuals_vs_simulation(
         output data file.
     size:
         Figure size.
-    ax:
+    axes:
         Axis object.
 
     Returns
@@ -66,9 +65,11 @@ def plot_residuals_vs_simulation(
         raise ValueError("Residuals plot is only applicable for normal "
                          "additive noise assumption")
 
-    if ax is None:
-        fig, ax = plt.subplots(figsize=size)
+    if axes is None:
+        fig, axes = plt.subplots(1, 2, sharey=True, figsize=size,
+                                 width_ratios=[2, 1])
         fig.set_layout_engine("tight")
+        fig.suptitle("Residuals")
 
     residual_df = calculate_residuals(
         measurement_dfs=petab_problem.measurement_df,
@@ -82,21 +83,30 @@ def plot_residuals_vs_simulation(
         simulations_df[OBSERVABLE_ID].isin(observable_ids)]
 
     # compare to standard normal distribution
-    ks_result = stats.kstest(normal_residuals['residual'], stats.norm.cdf)
-
-    ax.hlines(y=0, xmin=min(simulations_normal['simulation']),
-              xmax=max(simulations_normal['simulation']), ls='--',
-              color='gray')
-    ax.scatter(simulations_normal['simulation'],
-               normal_residuals['residual'])
-    ax.text(0.3, 0.85,
-            f'Kolmogorov-Smirnov test results:\n'
-            f'statistic: {ks_result[0]:.2f}\n'
-            f'pvalue: {ks_result[1]:.2e} ', transform=ax.transAxes)
-
-    ax.set_title("Residuals")
-    ax.set_xlabel('simulated values')
-    ax.set_ylabel('residuals')
-
-    plt.tight_layout()
-    return ax
+    ks_result = stats.kstest(normal_residuals[RESIDUAL], stats.norm.cdf)
+
+    # plot the residuals plot
+    axes[0].hlines(y=0, xmin=min(simulations_normal[SIMULATION]),
+                   xmax=max(simulations_normal[SIMULATION]), ls='--',
+                   color='gray')
+    axes[0].scatter(simulations_normal[SIMULATION],
+                    normal_residuals[RESIDUAL])
+    axes[0].text(0.3, 0.85,
+                 f'Kolmogorov-Smirnov test results:\n'
+                 f'statistic: {ks_result[0]:.2f}\n'
+                 f'pvalue: {ks_result[1]:.2e} ', transform=axes[0].transAxes)
+    axes[0].set_xlabel('simulated values')
+    axes[0].set_ylabel('residuals')
+
+    # plot histogram
+    axes[1].hist(normal_residuals[RESIDUAL], density=True,
+                 orientation='horizontal')
+    axes[1].set_xlabel('distribution')
+
+    ymin, ymax = axes[0].get_ylim()
+    ylim = max(abs(ymin), abs(ymax))
+    axes[0].set_ylim(-ylim, ylim)
+    axes[1].tick_params(left=False, labelleft=False, right=True,
+                        labelright=True)
+
+    return axes