Update plot styling, add residuals

Author: voetberg

src/deepdiagnostics/plots/coverage_fraction.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Optional, Union
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib.axes import Axes as ax
@@ -49,25 +49,82 @@ def __init__(
     def plot_name(self):
         return "coverage_fraction.png"
 
-    def _data_setup(self):
-        _, coverage = coverage_fraction_metric(
-            self.model, self.data, self.run_id, out_dir=None
+    def _data_setup(self, percentile_step_size:float=1) -> DataDisplay:
+        _, (coverage_mean, coverage_std) = coverage_fraction_metric(
+            self.model, self.data, self.run_id, out_dir=None, percentiles=np.arange(0, 100, percentile_step_size), use_progress_bar=self.use_progress_bar
         ).calculate()
         return DataDisplay(
-            coverage_fractions=coverage
+            coverage_fractions=coverage_mean, 
+            coverage_percentiles=np.arange(0, 100, percentile_step_size),
+            coverage_std=coverage_std
         )
 
+    def _plot_residual(self, data_display, ax, figure_alpha, line_width, reference_line_style, include_coverage_residual_std, include_ideal_range):
+        color_cycler = iter(plt.cycler("color", self.parameter_colors))
+        line_style_cycler = iter(plt.cycler("line_style", self.line_cycle))
+        percentile_array = data_display.coverage_percentiles / 100.0
+
+        ax.plot([0,1], [0, 0], reference_line_style, lw=line_width, zorder=1000)
+
+        for i in range(self.n_parameters):
+            color = next(color_cycler)["color"]
+            line_style = next(line_style_cycler)["line_style"]
+
+            residual = data_display.coverage_fractions[:, i] - np.linspace(0, 1, len(data_display.coverage_fractions[:,i]))
+
+            ax.plot(
+                percentile_array,
+                residual,
+                alpha=figure_alpha,
+                lw=line_width*.8,
+                linestyle=line_style,
+                color=color,
+                label=self.parameter_names[i],
+            )
+            if include_coverage_residual_std:
+
+                ax.fill_between(
+                    percentile_array,
+                    residual - data_display.coverage_std[:, i],
+                    residual + data_display.coverage_std[:, i],
+                    color=color,
+                    alpha=0.2,
+                )
+
+        if include_ideal_range:
+
+            ax.fill_between(
+                [0, 1],
+                [-0.2]*2,
+                [0.2]*2,
+                color="gray",
+                alpha=0.1,
+            )
+            ax.fill_between(
+                [0, 1],
+                [-0.1]*2,
+                [0.1]*2,
+                color="gray",
+                alpha=0.2,
+            )
+
     def plot(
         self,
         data_display: Union[DataDisplay, str],
         figure_alpha=1.0,
         line_width=3,
-        legend_loc="lower right",
+        legend_loc:Optional[str]=None,
+        include_coverage_std:bool = False, 
+        include_coverage_residual:bool = False,
+        include_coverage_residual_std:bool = False,
+        include_ideal_range: bool=True,
         reference_line_label="Reference Line",
         reference_line_style="k--",
         x_label="Confidence Interval of the Posterior Volume",
         y_label="Fraction of Lenses within Posterior Volume",
-        title="NPE") -> tuple["fig", "ax"]:
+        residual_y_label="Coverage Fraction Residual",
+        title="NPE"
+    ) -> tuple["fig", "ax"]:
         """
         Args:
             figure_alpha (float, optional): Opacity of parameter lines. Defaults to 1.0.
@@ -83,19 +140,28 @@ def plot(
         if not isinstance(data_display, DataDisplay):
             data_display = DataDisplay().from_h5(data_display, self.plot_name)
 
-        n_steps = data_display.coverage_fractions.shape[0]
-        percentile_array = np.linspace(0, 1, n_steps)
+
+        percentile_array = data_display.coverage_percentiles / 100.0
         color_cycler = iter(plt.cycler("color", self.parameter_colors))
         line_style_cycler = iter(plt.cycler("line_style", self.line_cycle))
 
         # Plotting
-        fig, ax = plt.subplots(1, 1, figsize=self.figure_size)
+        if include_coverage_residual:
+            fig, subplots = plt.subplots(2, 1, figsize=(self.figure_size[0], self.figure_size[1]*1.2), height_ratios=[3, 1], sharex=True)
+            ax = subplots[0]
+
+            self._plot_residual(
+                data_display, subplots[1], figure_alpha, line_width, reference_line_style, include_coverage_residual_std, include_ideal_range
+            )
+            subplots[1].set_ylabel(residual_y_label)
+
+        else:
+            fig, ax = plt.subplots(1, 1, figsize=self.figure_size)
 
         # Iterate over the number of parameters in the model
         for i in range(self.n_parameters):
             color = next(color_cycler)["color"]
             line_style = next(line_style_cycler)["line_style"]
-
             ax.plot(
                 percentile_array,
                 data_display.coverage_fractions[:, i],
@@ -105,6 +171,14 @@ def plot(
                 color=color,
                 label=self.parameter_names[i],
             )
+            if include_coverage_std:
+                ax.fill_between(
+                    percentile_array,
+                    data_display.coverage_fractions[:, i] - data_display.coverage_std[:, i],
+                    data_display.coverage_fractions[:, i] + data_display.coverage_std[:, i],
+                    color=color,
+                    alpha=0.2,
+                )
 
         ax.plot(
             [0, 0.5, 1],
@@ -115,13 +189,35 @@ def plot(
             label=reference_line_label,
         )
 
+        if include_ideal_range:
+            def add_clearance(ax, clearance=0.1, clearance_alpha=0.2):
+                x_values = np.linspace(0, 1, 100)  # More points for smoother curves
+                y_lower = np.maximum(0, x_values - clearance)  # Lower bound with clearance
+                y_upper = np.minimum(1, x_values + clearance)  # Upper bound with clearance
+                
+                # Fill the area between the bounds
+                ax.fill_between(
+                x_values,
+                y_lower,
+                y_upper,
+                color="gray",
+                alpha=clearance_alpha,
+                )
+                
+            add_clearance(ax, clearance=0.2, clearance_alpha=0.2)
+            add_clearance(ax, clearance=0.1, clearance_alpha=0.1)
+            
+
         ax.set_xlim([-0.05, 1.05])
         ax.set_ylim([-0.05, 1.05])
 
-        ax.text(0.03, 0.93, "Under-confident", horizontalalignment="left")
-        ax.text(0.3, 0.05, "Overconfident", horizontalalignment="left")
+        # ax.text(-0.03, 0.93, "Under-confident", horizontalalignment="left")
+        # ax.text(0.3, 0.05, "Overconfident", horizontalalignment="left")
 
-        ax.legend(loc=legend_loc)
+        if legend_loc is not None:
+            ax.legend(loc=legend_loc)
+        else:
+            ax.legend()
 
         ax.set_xlabel(x_label)
         ax.set_ylabel(y_label)

tests/conftest.py

@@ -82,8 +82,8 @@ def setUp(result_output):
     sim_paths = f"{simulator_config_path.strip('/')}/simulators.json"
     os.remove(sim_paths)
 
-    # out_dir = get_item("common", "out_dir", raise_exception=True)
-    # shutil.rmtree(out_dir)
+    out_dir = get_item("common", "out_dir", raise_exception=True)
+    shutil.rmtree(out_dir)
 
 @pytest.fixture
 def model_path():

tests/test_plots.py

@@ -20,7 +20,7 @@ def plot_config(config_factory):
     metrics_settings = {
         "use_progress_bar": False,
         "samples_per_inference": 10,
-        "percentiles": [95],
+        "percentiles": [95, 75, 50],
     }
     config = config_factory(metrics_settings=metrics_settings)
     return config