Merge pull request #113 from voetberg/36-coverage

36 coverage

Author: bnord

src/deepdiagnostics/metrics/coverage_fraction.py

@@ -51,34 +51,35 @@ def calculate(self) -> tuple[Sequence, Sequence]:
         Returns:
             tuple[Sequence, Sequence]: A tuple of the samples tested (M samples, Samples per inference, N parameters) and the coverage over those samples. 
         """
+
         all_samples = np.empty(
             (self.number_simulations, self.samples_per_inference, np.shape(self.thetas)[1])
         )
-        count_array = []
         iterator = range(self.number_simulations)
         if self.use_progress_bar:
             iterator = tqdm(
                 iterator,
                 desc="Sampling from the posterior for each observation",
                 unit=" observation",
             )
+        n_theta_samples = self.thetas.shape[0]
+        count_array = np.zeros((self.number_simulations, len(self.percentiles), self.thetas.shape[1]))
+
         for sample_index in iterator:
             context_sample = self.context[self.data.rng.integers(0, len(self.context))]
             samples = self._run_model_inference(self.samples_per_inference, context_sample)
 
             all_samples[sample_index] = samples
 
-            count_vector = []
             # step through the percentile list
-            for cov in self.percentiles:
+            for index, cov in enumerate(self.percentiles):
                 percentile_lower = 50.0 - cov / 2
                 percentile_upper = 50.0 + cov / 2
 
                 # find the percentile for the posterior for this observation
                 # this is n_params dimensional
                 # the units are in parameter space
                 confidence_lower = np.percentile(samples, percentile_lower, axis=0)
-                
                 confidence_upper = np.percentile(samples, percentile_upper, axis=0)
 
 
@@ -87,22 +88,25 @@ def calculate(self) -> tuple[Sequence, Sequence]:
                 # upper and lower confidence intervals
                 # checks separately for each side of the 50th percentile
 
-                count = np.logical_and(
-                    confidence_upper - self.thetas[sample_index, :].numpy() > 0,
-                    self.thetas[sample_index, :].numpy() - confidence_lower > 0,
+                c = np.logical_and(
+                    confidence_upper - self.thetas.numpy() > 0,
+                    self.thetas.numpy() - confidence_lower > 0,
                 )
-                count_vector.append(count)
+                count_array[sample_index, index] = np.sum(c.astype(int), axis=0)/n_theta_samples
 
             # each time the above is > 0, adds a count
-            count_array.append(count_vector)
+            #count_array[sample_index] = count_vector
+
+        coverage_mean = np.mean(count_array, axis=0)
+        coverage_std = np.std(count_array, axis=0)
 
-        count_sum_array = np.sum(count_array, axis=0)
-        frac_lens_within_vol = np.array(count_sum_array)
-        coverage = frac_lens_within_vol / len(self.context)
+        self.output = {
+            "coverage": coverage_mean,
+            "coverage_std": coverage_std,
 
-        self.output = coverage
+        }
 
-        return all_samples, coverage
+        return all_samples, (coverage_mean, coverage_std)
 
     def __call__(self, **kwds: Any) -> Any:
         self.calculate()

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,53 +49,130 @@ 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"]:
+        y_label="Coverage fraction within posterior volume",
+        residual_y_label="Coverage Fraction Residual",
+        title=""
+    ) -> tuple["fig", "ax"]:
         """
-        Args:
+        Plot the coverage fraction and residuals if specified.
+
+        Args: 
+            data_display (Union[DataDisplay, str]): DataDisplay object or path to h5 file containing the data. If str, it will be loaded and requires the fields "coverage_fractions", "coverage_percentiles", and optionally "coverage_std".
             figure_alpha (float, optional): Opacity of parameter lines. Defaults to 1.0.
             line_width (int, optional): Width of parameter lines. Defaults to 3.
-            legend_loc (str, optional): Location of the legend, str based on `matplotlib <https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.legend.html>`_. Defaults to "lower right".
+            legend_loc (str, optional): Location of the legend. Defaults to matplotlib specified. 
+            include_coverage_std (bool, optional): Whether to include the standard deviation shading for coverage fractions . Defaults to False.
+            include_coverage_residual (bool, optional): Whether to include the residual plot (coverage fraction - diagonal). Creates an additional subplot under the original plot. Defaults to False.
+            include_coverage_residual_std (bool, optional): Whether to include the standard deviation shading for residuals. Defaults to False.
+            include_ideal_range (bool, optional): Whether to include the ideal range shading (0.1/0.2 around the diagonal). Defaults to True.
             reference_line_label (str, optional): Label name for the diagonal ideal line. Defaults to "Reference Line".
             reference_line_style (str, optional): Line style for the reference line. Defaults to "k--".
             x_label (str, optional): y label. Defaults to "Confidence Interval of the Posterior Volume".
             y_label (str, optional): y label. Defaults to "Fraction of Lenses within Posterior Volume".
+            residual_y_label (str, optional): y label for the residual plot. Defaults to "Coverage Fraction Residual".
             title (str, optional): plot title. Defaults to "NPE".
+
         """
-    
+
         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)
+            subplots[1].set_xlabel(x_label)
+
+        else:
+            fig, ax = plt.subplots(1, 1, figsize=self.figure_size)
+            ax.set_xlabel(x_label)
 
         # 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 +182,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,15 +200,36 @@ 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)
         ax.set_title(title)
 

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