Repository: epic-open-source/seismometer

Issue:
Cohort selection for generating analytics tables
## Is your feature request related to a problem? Please describe

Unable to filter to specific cohorts for generating analytics tables via `ExploreAnalyticsTable()` interactive visualization and `table.analytics_table.binary_analytics_table()` function. 

## Describe the solution you'd like

Add a widget/parameter for cohort selection similar to `ExploreModelEvaluation()` and `plot_model_evaluation({'cohort column': ('cohort value', ...)}, ...)`.

Hints from issue comments:
@seanrmeyer: closed as duplicate

Solution patch:
diff --git a/changelog/131.feature.rst b/changelog/131.feature.rst
new file mode 100644
index 00000000..179fe193
--- /dev/null
+++ b/changelog/131.feature.rst
@@ -0,0 +1 @@
+Addresses #131 by adding cohort selection to analytics table options.
\ No newline at end of file
diff --git a/changelog/134.bugfix.rst b/changelog/134.bugfix.rst
new file mode 100644
index 00000000..605cf8f7
--- /dev/null
+++ b/changelog/134.bugfix.rst
@@ -0,0 +1 @@
+Fix an issue with reading parquet files without Pandas metadata (e.g. those created by Polars)
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index ebef71d1..50f04efa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,3 +10,7 @@ directory = "changelog"
 issue_format = "`#{issue} <https://github.com/epic-open-source/seismometer/issues/{issue}>`__"
 title_format = "{version}"
 underlines = ["-","~"]
+
+[tool.black]
+line-length = 119
+target-version = ["py310"]
\ No newline at end of file
diff --git a/src/seismometer/data/binary_performance.py b/src/seismometer/data/binary_performance.py
index f94a4edb..24a4e29b 100644
--- a/src/seismometer/data/binary_performance.py
+++ b/src/seismometer/data/binary_performance.py
@@ -1,10 +1,11 @@
 import itertools
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import numpy as np
 import pandas as pd
 
 from seismometer.data import pandas_helpers as pdh
+from seismometer.data.filter import FilterRule
 from seismometer.seismogram import Seismogram
 
 from . import BinaryClassifierMetricGenerator
@@ -32,7 +33,7 @@ def calculate_stats(
     metric_values: List[str],
     metrics_to_display: Optional[List[str]] = None,
     decimals: int = 3,
-):
+) -> dict:
     """
     Calculates overall performance statistics such as AUROC and threshold-specific statistics
     such as sensitivity and specificity for the provided list of metric values.
@@ -124,10 +125,12 @@ def generate_analytics_data(
     metric_values: List[float],
     *,
     top_level: str = "Score",
+    cohort_dict: Optional[dict[str, tuple[Any]]] = None,
     per_context: bool = False,
     metrics_to_display: Optional[List[str]] = None,
     decimals: int = 3,
-):
+    censor_threshold: int = 10,
+) -> Optional[pd.DataFrame]:
     """
     Generates a DataFrame containing calculated statistics for each combination of scores and targets.
 
@@ -143,6 +146,8 @@ def generate_analytics_data(
         A list of metric values for which corresponding statistics are calculated.
     top_level : str, optional
         The primary grouping category in the performance table, by default "Score".
+    cohort_dict : Optional[dict[str, tuple[Any]]], optional
+        dictionary of cohort columns and values used to subselect a population for evaluation, by default None.
     per_context : bool
         If scores should be grouped by context, by default False.
     metrics_to_display : Optional[List[str]], optional
@@ -150,10 +155,12 @@ def generate_analytics_data(
         in GENERATED_COLUMNS.
     decimals : int, optional
         The number of decimal places for rounding numerical results, by default 3.
+    censor_threshold : int, optional
+        Minimum rows required to generate analytics data, by default 10.
 
     Returns
     -------
-    pd.DataFrame
+    Optional[pd.DataFrame]
         A DataFrame containing the calculated statistics for each combination of scores and targets.
     """
     rows_list = []
@@ -164,19 +171,23 @@ def generate_analytics_data(
     )
     second_level = "Target" if top_level == "Score" else "Score"
     sg = Seismogram()
+    cohort_dict = cohort_dict or sg.available_cohort_groups
+    cohort_filter = FilterRule.from_cohort_dictionary(cohort_dict)
+    cohort_filter.MIN_ROWS = censor_threshold
+    data = cohort_filter.filter(sg.dataframe)
+    if len(data) <= censor_threshold:
+        return None
     for first, second in product:
         current_row = {top_level: first, second_level: second}
         (score, target) = (first, second) if top_level == "Score" else (second, first)
         if per_context:
             data = pdh.event_score(
-                sg.dataframe,
+                data,
                 sg.entity_keys,
                 score=score,
                 ref_event=sg.predict_time,
                 aggregation_method=sg.event_aggregation_method(target),
             )
-        else:
-            data = sg.dataframe
         current_row.update(
             calculate_stats(
                 data[[target, score]],
@@ -190,5 +201,5 @@ def generate_analytics_data(
         )
         rows_list.append(current_row)
     # Create a DataFrame from the rows data
-    data = pd.DataFrame(rows_list)
-    return data
+    analytics_data = pd.DataFrame(rows_list)
+    return analytics_data
diff --git a/src/seismometer/data/loader/prediction.py b/src/seismometer/data/loader/prediction.py
index e661aa11..5c5b8165 100644
--- a/src/seismometer/data/loader/prediction.py
+++ b/src/seismometer/data/loader/prediction.py
@@ -38,11 +38,9 @@ def parquet_loader(config: ConfigProvider) -> pd.DataFrame:
 
         if config.target in present_columns:
             desired_columns.add(config.target)
-
         actual_columns = desired_columns & present_columns
         _log_column_mismatch(actual_columns, desired_columns, present_columns)
-
-        dataframe = pd.read_parquet(config.prediction_path, columns=actual_columns)
+        dataframe = pd.read_parquet(config.prediction_path, columns=list(actual_columns))
 
     dataframe = _rename_targets(config, dataframe)
 
diff --git a/src/seismometer/table/analytics_table.py b/src/seismometer/table/analytics_table.py
index ec6bdf0e..3f68f142 100644
--- a/src/seismometer/table/analytics_table.py
+++ b/src/seismometer/table/analytics_table.py
@@ -1,20 +1,22 @@
 from enum import Enum
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import pandas as pd
 
 # import ipywidgets as widgets
 import traitlets
 from great_tables import GT, loc, style
-from ipywidgets import HTML, Box, Dropdown, FloatRangeSlider, GridBox, Layout
+from ipywidgets import HTML, Dropdown, GridBox, Layout, VBox
 from pandas.api.types import is_integer_dtype, is_numeric_dtype
 
 from seismometer.controls.explore import ExplorationWidget, _combine_scores_checkbox
-from seismometer.controls.selection import MultiselectDropdownWidget
-from seismometer.controls.styles import BOX_GRID_LAYOUT, WIDE_LABEL_STYLE
+from seismometer.controls.selection import MultiselectDropdownWidget, MultiSelectionListWidget
+from seismometer.controls.styles import BOX_GRID_LAYOUT, html_title
+from seismometer.controls.thresholds import MonotonicProbabilitySliderListWidget
 from seismometer.data import pandas_helpers as pdh
 from seismometer.data.binary_performance import GENERATED_COLUMNS, generate_analytics_data
 from seismometer.data.performance import MONOTONIC_METRICS, OVERALL_PERFORMANCE, STATNAMES, THRESHOLD
+from seismometer.html import template
 from seismometer.seismogram import Seismogram
 
 from .analytics_table_config import COLORING_CONFIG_DEFAULT, AnalyticsTableConfig
@@ -52,39 +54,46 @@ def __init__(
         metrics_to_display: Optional[List[str]] = None,
         title: str = "Model Performance Statistics",
         top_level: str = "Score",
+        cohort_dict: Optional[dict[str, tuple]] = None,
         table_config: Optional[AnalyticsTableConfig] = AnalyticsTableConfig(),
         statistics_data: Optional[pd.DataFrame] = None,
         per_context: bool = False,
+        censor_threshold: int = 10,
     ):
         """
         Initializes the AnalyticsTable object with the necessary data and parameters.
 
         Parameters
         ----------
-        score_columns: Optional[List[str]]
+        score_columns: Optional[List[str]], optional
             A list of column names corresponding to model prediction scores, by default None.
-        target_columns: Optional[List[str]]
+        target_columns: Optional[List[str]], optional
             A list of column names corresponding to (binary) targets, by default None.
-        metric: str
-            Performance metrics will be presented for the provided values of this metric.
-        metric_values: List[float]
+        metric: str, optional
+            Performance metrics will be presented for the provided values of this metric, by default "Threshold".
+        metric_values: List[float], optional
             Values for the specified metric to derive detailed performance statistics, by default [0.1, 0.2].
-        metrics_to_display: Optional[List[str]]
+        metrics_to_display: Optional[List[str]], optional
             List of metrics to include in the table, by default None. The default behavior is to show all columns
             in GENERATED_COLUMNS.
-        title: str
+        title: str, optional
             The title for the performance statistics table, by default "Model Performance Statistics".
-        top_level: str
+        top_level: str, optional
             The primary grouping category in the performance table, by default 'Score'.
-        decimals: int
+        cohort_dict : Optional[dict[str,tuple]], optional
+            dictionary of cohort columns and values used to subselect a population for evaluation, by default None.
+        decimals: int, optional
             The number of decimal places for rounding numerical results, by default 3.
-        table_config: Optional[AnalyticsTableConfig]
+        table_config: Optional[AnalyticsTableConfig], optional
             Configuration for the analytics table, including formatting and display options.
-        statistics_data: Optional[pd.DataFrame]
+        statistics_data: Optional[pd.DataFrame], optional
             Additional performance metrics statistics, will be joined with the statistics data generated
             by the code, by default None.
-        per_context : bool
+        per_context : bool, optional
             If scores should be grouped by context, by default False.
+        censor_threshold : int, optional
+            Minimum number of rows required in the cohort data to enable the generation of an analytics table,
+            by default 10.
 
         Raises
         ------
@@ -115,12 +124,14 @@ def __init__(
         self.metrics_to_display = metrics_to_display or GENERATED_COLUMNS
         self.title = title
         self.top_level = top_level
+        self.cohort_dict = cohort_dict
         self.columns_show_percentages = table_config.columns_show_percentages
 
         self.percentages_decimals = table_config.percentages_decimals
 
         self._initializing = False
         self.per_context = per_context
+        self.censor_threshold = censor_threshold
 
     def _validate_df_statistics_data(self):
         if not self._initializing:  # Skip validation during initial setup
@@ -214,7 +225,7 @@ def top_level(self, value):
                 f"{[member.value for member in TopLevel]}"
             ) from e
 
-    def generate_initial_table(self, data):
+    def generate_initial_table(self, data) -> GT:
         """
         Generates the initial table with formatted headers, stubs, and numeric/percentage formatting.
 
@@ -249,7 +260,7 @@ def generate_initial_table(self, data):
         )
         return gt
 
-    def group_columns_by_metric_value(self, gt, columns, value):
+    def group_columns_by_metric_value(self, gt, columns, value) -> GT:
         """
         Groups columns by a specified metric value and adds borders.
 
@@ -273,7 +284,7 @@ def group_columns_by_metric_value(self, gt, columns, value):
         )
         return self.add_borders(gt, columns[0], columns[-1])
 
-    def add_borders(self, gt, left_column, right_column):
+    def add_borders(self, gt, left_column, right_column) -> GT:
         """
         Adds borders to the left of left_column and right of right_column in the table.
 
@@ -300,23 +311,26 @@ def add_borders(self, gt, left_column, right_column):
         )
         return gt
 
-    def analytics_table(self):
+    def analytics_table(self) -> HTML:
         """
         Generates an analytics table based on calculated performance statistics.
 
         Returns
         -------
-        GT
-            A `GT` (from great_tables package) object representing the formatted analytics table.
+        HTML
+            An HTML object representing the formatted analytics table.
         """
         data = self._generate_table_data()
+        if data is None:
+            return template.render_censored_plot_message(self.censor_threshold)
 
         gt = self.generate_initial_table(data)
 
         # Group columns of the form value_*** together
         for value in self.metric_values:
             columns = [column for column in data.columns if column.startswith(f"{value}_")]
-            gt = self.group_columns_by_metric_value(gt, columns, value)
+            if columns:
+                gt = self.group_columns_by_metric_value(gt, columns, value)
 
         gt = (
             gt.opt_horizontal_padding(scale=3)
@@ -337,7 +351,7 @@ def analytics_table(self):
 
         return HTML(gt.as_raw_html(), layout=Layout(max_height="800px"))
 
-    def _generate_table_data(self):
+    def _generate_table_data(self) -> Optional[pd.DataFrame]:
         """
         Generates a DataFrame containing calculated statistics for each combination of scores and targets.
 
@@ -353,10 +367,14 @@ def _generate_table_data(self):
             self.metric,
             self.metric_values,
             top_level=self.top_level,
+            cohort_dict=self.cohort_dict,
             per_context=self.per_context,
             metrics_to_display=self.metrics_to_display,
             decimals=self.decimals,
+            censor_threshold=self.censor_threshold,
         )
+        if data is None:
+            return None
 
         # Update target columns name to drop possible _Value at the end
         data["Target"] = data["Target"].apply(pdh.event_name)
@@ -392,6 +410,7 @@ def binary_analytics_table(
     metric_values: list[float],
     metrics_to_display: list[str],
     group_by: str,
+    cohort_dict: dict[str, tuple[Any]],
     *,
     title: str = None,
     per_context: bool = False,
@@ -413,6 +432,8 @@ def binary_analytics_table(
         List of performance metrics to include in the table.
     group_by : str
         The primary grouping category in the performance table. It could be "Score" or "Target".
+    cohort_dict : dict[str, tuple[Any]]
+        dictionary of cohort columns and values used to subselect a population for evaluation.
     title : str, optional
         The title for the performance statistics table, by default None.
     per_context : bool, optional
@@ -421,8 +442,9 @@ def binary_analytics_table(
     Returns
     -------
     HTML
-        The HTML table for the fairness evaluation.
+        The HTML table representing the corresponding analytics table.
     """
+    sg = Seismogram()
     table_config = AnalyticsTableConfig(**COLORING_CONFIG_DEFAULT)
     performance_metrics = AnalyticsTable(
         score_columns=score_cols,
@@ -432,8 +454,10 @@ def binary_analytics_table(
         metrics_to_display=metrics_to_display,
         title=title or "Model Performance Statistics",
         top_level=group_by,
+        cohort_dict=cohort_dict,
         table_config=table_config,
         per_context=per_context,
+        censor_threshold=sg.censor_threshold,
     )
     return performance_metrics.analytics_table()
 
@@ -457,6 +481,7 @@ def __init__(self, title: Optional[str] = None):
                 metric="Threshold",
                 metric_values=None,
                 metrics_to_display=None,
+                cohort_dict=sg.available_cohort_groups,
                 title=title,
             ),
             plot_function=binary_analytics_table,
@@ -466,18 +491,19 @@ def __init__(self, title: Optional[str] = None):
     def generate_plot_args(self) -> tuple[tuple, dict]:
         """Generates the plot arguments for the analytics table."""
         args = (
-            tuple(map(pdh.event_value, self.option_widget.target_cols)),  # Updated to use target_columns
-            self.option_widget.score_cols,  # Updated to use score_columns
+            tuple(map(pdh.event_value, self.option_widget.target_cols)),  # Updated to use target_cols
+            self.option_widget.score_cols,  # Updated to use score_cols
             self.option_widget.metric,  # Updated to use metric
             self.option_widget.metric_values,  # Updated to use metric_values
             list(self.option_widget.metrics_to_display),  # Updated to use metrics_to_display
             self.option_widget.group_by,  # Updated to use group_by
+            self.option_widget.cohort_dict,
         )
         kwargs = {"title": self.title, "per_context": self.option_widget.group_scores}
         return args, kwargs
 
 
-class AnalyticsTableOptionsWidget(Box, traitlets.HasTraits):
+class AnalyticsTableOptionsWidget(VBox, traitlets.HasTraits):
     value = traitlets.Dict(help="The selected values for the analytics table options.")
 
     def __init__(
@@ -489,6 +515,7 @@ def __init__(
         model_options_widget=None,
         metric_values=None,
         metrics_to_display: Optional[tuple[str]] = None,
+        cohort_dict: Optional[dict[str, tuple[Any]]] = None,
         title: str = None,
     ):
         """
@@ -508,6 +535,8 @@ def __init__(
             Metric values for which performance metric will be computed, by default None.
         metrics_to_display : Optional[tuple[str]]
             Metrics to include in the analytics table, by default None.
+        cohort_dict : Optional[dict[str, tuple[Any]]]
+            dictionary of cohort columns and values used to subselect a population for evaluation, by default None.
         title : Optional[str]
             Title of the widget, by default None.
         """
@@ -532,43 +561,45 @@ def __init__(
             options=[THRESHOLD] + MONOTONIC_METRICS,
             value=metric,
             description="Metric",
-            style=WIDE_LABEL_STYLE,
+            style={"description_width": "min-content"},
+            layout=Layout(width="calc(max(max-content, var(--jp-widgets-inline-width-short)))", min_width="200px"),
         )
         self._metrics_to_display = MultiselectDropdownWidget(
             options=[THRESHOLD] + STATNAMES + OVERALL_PERFORMANCE,
             value=metrics_to_display or GENERATED_COLUMNS,
             title="Performance Metrics to Display",
         )
-        self._metric_values_slider = FloatRangeSlider(
-            min=0.01,
-            max=1.00,
-            step=0.01,
-            value=metric_values or [0.2, 0.8],
-            description="Metric Values",
-            style=WIDE_LABEL_STYLE,
+        metric_values = metric_values or (0.8, 0.2)
+        self._metric_values = MonotonicProbabilitySliderListWidget(
+            names=("Metric Value 1", "Metric Value 2"), value=tuple(metric_values), ascending=False
         )
+        for slider in self._metric_values.sliders.values():
+            slider.style.description_width = "min-content"
         self._group_by = Dropdown(
             options=["Score", "Target"],
             value="Score",
             description="Group By",
-            style=WIDE_LABEL_STYLE,
+            style={"description_width": "min-content"},
+            layout=Layout(width="calc(max(max-content, var(--jp-widgets-inline-width-short)))", min_width="200px"),
         )
+        self._cohort_dict = MultiSelectionListWidget(cohort_dict or sg.available_cohort_groups, title="Cohort Filter")
         self.per_context_checkbox = _combine_scores_checkbox(per_context=False)
 
         self._target_cols.observe(self._on_value_changed, names="value")
         self._score_cols.observe(self._on_value_changed, names="value")
         self._metric.observe(self._on_value_changed, names="value")
-        self._metric_values_slider.observe(self._on_value_changed, names="value")
+        self._metric_values.observe(self._on_value_changed, "value")
         self._metrics_to_display.observe(self._on_value_changed, names="value")
         self._group_by.observe(self._on_value_changed, names="value")
         self.per_context_checkbox.observe(self._on_value_changed, names="value")
+        self._cohort_dict.observe(self._on_value_changed, names="value")
 
         v_children = [
             self._target_cols,
             self._score_cols,
             self._metrics_to_display,
             self._metric,
-            self._metric_values_slider,
+            self._metric_values,
             self._group_by,
             self.per_context_checkbox,
         ]
@@ -577,14 +608,25 @@ def __init__(
             self.model_options_widget.observe(self._on_value_changed, names="value")
 
         grid_layout = Layout(
-            width="80%", grid_template_columns="repeat(3, 1fr)", justify_items="flex-end", grid_gap="10px"
+            width="100%", grid_template_columns="repeat(3, 1fr)", justify_items="flex-start", grid_gap="10px"
         )  # Three columns
 
         # Create a GridBox with the specified layout
         grid_box = GridBox(children=v_children, layout=grid_layout)
 
+        # Add title
+        grid_with_title = VBox(
+            children=[
+                html_title("Analytics Table Options"),
+                grid_box,
+            ],
+            layout=Layout(
+                align_items="flex-start",
+            ),
+        )
+
         super().__init__(
-            children=[grid_box],
+            children=[self._cohort_dict, grid_with_title],
             layout=BOX_GRID_LAYOUT,
         )
 
@@ -601,9 +643,10 @@ def disabled(self, value):
         self._target_cols.disabled = value
         self._score_cols.disabled = value
         self._metric.disabled = value
-        self._metric_values_slider.disabled = value
+        self._metric_values.disabled = value
         self._metrics_to_display.disabled = value
         self._group_by.disabled = value
+        self._cohort_dict.disabled = value
         self.per_context_checkbox.disabled = value
         if self.model_options_widget:
             self.model_options_widget.disabled = value
@@ -613,9 +656,10 @@ def _on_value_changed(self, change=None):
             "target_cols": self._target_cols.value,
             "score_cols": self._score_cols.value,
             "metric": self._metric.value,
-            "metric_values": self._metric_values_slider.value,
+            "metric_values": self._metric_values.value,
             "metrics_to_display": self._metrics_to_display.value,
             "group_by": self._group_by.value,
+            "cohort_dict": self._cohort_dict.value,
             "group_scores": self.per_context_checkbox.value,
         }
         if self.model_options_widget:
@@ -636,7 +680,7 @@ def metric(self):
 
     @property
     def metric_values(self):
-        return self._metric_values_slider.value
+        return tuple(self._metric_values.value.values())
 
     @property
     def metrics_to_display(self):
@@ -646,6 +690,10 @@ def metrics_to_display(self):
     def group_by(self):
         return self._group_by.value
 
+    @property
+    def cohort_dict(self):
+        return self._cohort_dict.value
+
     @property
     def group_scores(self):
         return self.per_context_checkbox.value
diff --git a/src/seismometer/table/fairness.py b/src/seismometer/table/fairness.py
index 0295fe14..03be3c97 100644
--- a/src/seismometer/table/fairness.py
+++ b/src/seismometer/table/fairness.py
@@ -388,13 +388,15 @@ def __init__(
         Parameters
         ----------
         metric_names : tuple[str]
-            Metrics that can be evaluated for fairness
+            Metrics that can be evaluated for fairness.
         cohort_dict : dict[str, tuple[Any]]
             Dictionary of cohort groups.
         fairness_ratio : float, optional
-            Allowed difference by cohort, by default 0.2
-        model_options_widget : _type_, optional
-            Additional model options if needed, will appear before fairness options, by default None
+            Allowed difference by cohort, by default 0.2.
+        model_options_widget : Optional[widget], optional
+            Additional model options if needed, will appear before fairness options, by default None.
+        default_metrics : Optional[tuple[str]], optional
+            Default list of metrics to select initially for fairness evaluation, by default None.
         """
         self.model_options_widget = model_options_widget
         default_metrics = default_metrics or metric_names


Test patch:
diff --git a/tests/controls/test_explore.py b/tests/controls/test_explore.py
index 45cba4b8..518f9073 100644
--- a/tests/controls/test_explore.py
+++ b/tests/controls/test_explore.py
@@ -848,6 +848,7 @@ def test_init(self, mock_seismo, mock_plot_function):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1_Value", "T2_Value"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
         mock_plot_function.__name__ = "plot_function"
         mock_plot_function.__module__ = "test_explore"
@@ -859,15 +860,16 @@ def test_init(self, mock_seismo, mock_plot_function):
         widget.update_plot()
         assert (
             widget.current_plot_code
-            == "test_explore.plot_function(('T1_Value', 'T2_Value'), ('S1', 'S2'), 'Threshold', (0.2, 0.8), "
+            == "test_explore.plot_function(('T1_Value', 'T2_Value'), ('S1', 'S2'), 'Threshold', (0.8, 0.2), "
             + "['Positives', 'Prevalence', 'AUROC', 'AUPRC', 'Accuracy', 'PPV', 'Sensitivity', 'Specificity', "
-            + "'Flag Rate', 'Threshold'], 'Score', title='Unit Test Title', per_context=False)"
+            + "'Flag Rate', 'Threshold'], 'Score', {}, "
+            + "title='Unit Test Title', per_context=False)"
         )
         mock_plot_function.assert_called_once_with(
             ("T1_Value", "T2_Value"),
             ("S1", "S2"),
             "Threshold",
-            (0.2, 0.8),
+            (0.8, 0.2),
             [
                 "Positives",
                 "Prevalence",
@@ -881,6 +883,7 @@ def test_init(self, mock_seismo, mock_plot_function):
                 "Threshold",
             ],
             "Score",
+            {},
             title="Unit Test Title",
             per_context=False,
         )
@@ -890,6 +893,7 @@ def test_generate_plot_args(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1_Value", "T2_Value"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
         widget = ExploreBinaryModelAnalytics(title="Unit Test Title")
 
@@ -898,7 +902,7 @@ def test_generate_plot_args(self, mock_seismo):
             ("T1_Value", "T2_Value"),
             ("S1", "S2"),
             "Threshold",
-            (0.2, 0.8),
+            (0.8, 0.2),
             [
                 "Positives",
                 "Prevalence",
@@ -912,6 +916,7 @@ def test_generate_plot_args(self, mock_seismo):
                 "Threshold",
             ],
             "Score",
+            {},
         )
         assert kwargs == {"title": "Unit Test Title", "per_context": False}
 
@@ -922,62 +927,88 @@ def test_init(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1", "T2"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
         widget = AnalyticsTableOptionsWidget(
             target_cols=("T1", "T2"),
             score_cols=("S1", "S2"),
             metric="Threshold",
-            metric_values=[0.2, 0.8],
+            metric_values=[0.8, 0.2],
             metrics_to_display=("Accuracy", "PPV"),
+            cohort_dict=fake_seismo.available_cohort_groups,
             title="Unit Test Title",
         )
 
         assert widget._target_cols.value == ("T1", "T2")
         assert widget._score_cols.value == ("S1", "S2")
         assert widget._metric.value == "Threshold"
-        assert widget._metric_values_slider.value == (0.2, 0.8)
+        assert widget._metric_values.value == {"Metric Value 1": 0.8, "Metric Value 2": 0.2}
         assert widget._metrics_to_display.value == ("Accuracy", "PPV")
         assert widget._group_by.value == "Score"
         assert widget.per_context_checkbox.value is False
+        assert widget._cohort_dict.value == {}
 
     @patch.object(seismogram, "Seismogram", return_value=Mock())
     def test_disabled_property(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1", "T2"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
-        widget = AnalyticsTableOptionsWidget(target_cols=("T1", "T2"), score_cols=("S1", "S2"), metric="Threshold")
+        widget = AnalyticsTableOptionsWidget(
+            target_cols=("T1", "T2"),
+            score_cols=("S1", "S2"),
+            metric="Threshold",
+            cohort_dict=fake_seismo.available_cohort_groups,
+        )
         widget.disabled = True
         assert widget._target_cols.disabled is True
         assert widget._score_cols.disabled is True
         assert widget._metric.disabled is True
-        assert widget._metric_values_slider.disabled is True
+        assert widget._metric_values.disabled is True
         assert widget._metrics_to_display.disabled is True
         assert widget._group_by.disabled is True
         assert widget.per_context_checkbox.disabled is True
+        assert widget._cohort_dict.disabled is True
 
     @patch.object(seismogram, "Seismogram", return_value=Mock())
     def test_on_value_changed(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1", "T2"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
-        widget = AnalyticsTableOptionsWidget(target_cols=("T1", "T2"), score_cols=("S1", "S2"), metric="Threshold")
+        widget = AnalyticsTableOptionsWidget(
+            target_cols=("T1", "T2"),
+            score_cols=("S1", "S2"),
+            metric="Threshold",
+            cohort_dict=fake_seismo.available_cohort_groups,
+        )
         widget._target_cols.value = ("T1",)
         widget._score_cols.value = ("S1",)
         widget._metric.value = "Sensitivity"
-        widget._metric_values_slider.value = [0.1, 0.9]
+        widget._metric_values.value = {"Metric Value 1": 0.9, "Metric Value 2": 0.1}
         widget._metrics_to_display.value = ("Accuracy",)
         widget._group_by.value = "Target"
+        widget._cohort_dict.value = {
+            "C1": [
+                "C1.1",
+            ]
+        }
         widget.per_context_checkbox.value = True
 
         expected_value = {
             "target_cols": ("T1",),
             "score_cols": ("S1",),
             "metric": "Sensitivity",
-            "metric_values": (0.1, 0.9),
+            "metric_values": {"Metric Value 1": 0.9, "Metric Value 2": 0.1},
             "metrics_to_display": ("Accuracy",),
             "group_by": "Target",
+            "cohort_dict": {
+                "C1": [
+                    "C1.1",
+                ]
+            },
             "group_scores": True,
         }
         assert widget.value == expected_value
@@ -987,6 +1018,7 @@ def test_model_options_widget(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1", "T2"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
         model_options_widget = ipywidgets.Dropdown(
             options=["Val1", "Val2"],
@@ -997,6 +1029,7 @@ def test_model_options_widget(self, mock_seismo):
             target_cols=("T1", "T2"),
             score_cols=("S1", "S2"),
             metric="Threshold",
+            cohort_dict=fake_seismo.available_cohort_groups,
             model_options_widget=model_options_widget,
         )
         assert widget.model_options_widget == model_options_widget
@@ -1006,8 +1039,14 @@ def test_group_scores(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1", "T2"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
-        widget = AnalyticsTableOptionsWidget(target_cols=("T1", "T2"), score_cols=("S1", "S2"), metric="Threshold")
+        widget = AnalyticsTableOptionsWidget(
+            target_cols=("T1", "T2"),
+            score_cols=("S1", "S2"),
+            metric="Threshold",
+            cohort_dict=fake_seismo.available_cohort_groups,
+        )
         widget.per_context_checkbox.value = True
         assert widget.group_scores is True
 
@@ -1019,6 +1058,7 @@ def test_edge_cases(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = []
         fake_seismo.output_list = []
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
         widget = AnalyticsTableOptionsWidget(
             target_cols=(),
@@ -1026,13 +1066,14 @@ def test_edge_cases(self, mock_seismo):
             metric="Threshold",
             metric_values=[],
             metrics_to_display=(),
+            cohort_dict={},
             title="Unit Test Title",
         )
 
         assert widget._target_cols.value == ()
         assert widget._score_cols.value == ()
         assert widget._metric.value == "Threshold"
-        assert widget._metric_values_slider.value == (0.2, 0.8)
+        assert widget._metric_values.value == {"Metric Value 1": 0.8, "Metric Value 2": 0.2}
         assert widget._metrics_to_display.value == (
             "Positives",
             "Prevalence",
@@ -1053,13 +1094,15 @@ def test_state_changes(self, mock_seismo):
         fake_seismo = mock_seismo()
         fake_seismo.get_binary_targets.return_value = ["T1", "T2"]
         fake_seismo.output_list = ["S1", "S2"]
+        fake_seismo.available_cohort_groups = {"C1": ["C1.1", "C1.2"], "C2": ["C2.1", "C2.2"]}
 
         widget = AnalyticsTableOptionsWidget(
             target_cols=("T1", "T2"),
             score_cols=("S1", "S2"),
             metric="Threshold",
-            metric_values=[0.2, 0.8],
+            metric_values=[0.8, 0.2],
             metrics_to_display=("Accuracy", "PPV"),
+            cohort_dict=fake_seismo.available_cohort_groups,
             title="Unit Test Title",
         )
 
@@ -1067,18 +1110,28 @@ def test_state_changes(self, mock_seismo):
         widget._target_cols.value = ("T2",)
         widget._score_cols.value = ("S2",)
         widget._metric.value = "Sensitivity"
-        widget._metric_values_slider.value = [0.1, 0.9]
+        widget._metric_values.value = {"Metric Value 1": 0.9, "Metric Value 2": 0.1}
         widget._metrics_to_display.value = ("PPV",)
         widget._group_by.value = "Target"
+        widget._cohort_dict.value = {
+            "C1": [
+                "C1.2",
+            ]
+        }
         widget.per_context_checkbox.value = True
 
         # Verify the state changes
         assert widget._target_cols.value == ("T2",)
         assert widget._score_cols.value == ("S2",)
         assert widget._metric.value == "Sensitivity"
-        assert widget._metric_values_slider.value == (0.1, 0.9)
+        assert widget._metric_values.value == {"Metric Value 1": 0.9, "Metric Value 2": 0.1}
         assert widget._metrics_to_display.value == ("PPV",)
         assert widget._group_by.value == "Target"
+        assert widget._cohort_dict.value == {
+            "C1": [
+                "C1.2",
+            ]
+        }
         assert widget.per_context_checkbox.value is True
 
 
diff --git a/tests/data/loaders/test_load_events.py b/tests/data/loaders/test_load_events.py
index d2305ce6..1c0d7320 100644
--- a/tests/data/loaders/test_load_events.py
+++ b/tests/data/loaders/test_load_events.py
@@ -5,6 +5,8 @@
 import numpy as np
 import pandas as pd
 import pandas.testing as pdt
+import pyarrow as pa
+import pyarrow.parquet as pq
 import pytest
 
 import seismometer.data.loader.event as undertest
@@ -60,7 +62,7 @@ def event_frame(rename=False):
 
 # endregion
 # region File-type setup functions
-def parquet_setup():
+def pandas_parquet_setup():
     # Expects current directory to have been modified by the testcase,
     # such as using tmp_as_current
     file = Path("events.parquet")
@@ -71,8 +73,23 @@ def parquet_setup():
     return fake_config(file)
 
 
+def non_pandas_parquet_setup():
+    file = Path("predictions.parquet")
+
+    data = event_frame()
+    # remove the pandas metadata in the parquet file to simulate the data
+    # coming from a different parquet source (e.g. polars)
+    table = pa.Table.from_pandas(data).replace_schema_metadata({})
+
+    pq.write_table(table, file)
+    return fake_config(file)
+
+
 # region Tests
-@pytest.mark.parametrize("setup_fn,load_fn", [[parquet_setup, undertest.parquet_loader]])
+@pytest.mark.parametrize(
+    "setup_fn,load_fn",
+    [[pandas_parquet_setup, undertest.parquet_loader], [non_pandas_parquet_setup, undertest.parquet_loader]],
+)
 @pytest.mark.usefixtures("tmp_as_current")
 class TestPredictionLoad:
     def test_nofile_warns_and_returns_empty(self, setup_fn, load_fn, caplog):
diff --git a/tests/data/loaders/test_load_predictions.py b/tests/data/loaders/test_load_predictions.py
index 7a317588..82a617bf 100644
--- a/tests/data/loaders/test_load_predictions.py
+++ b/tests/data/loaders/test_load_predictions.py
@@ -5,6 +5,8 @@
 import numpy as np
 import pandas as pd
 import pandas.testing as pdt
+import pyarrow as pa
+import pyarrow.parquet as pq
 import pytest
 from conftest import tmp_as_current  # noqa
 
@@ -61,7 +63,7 @@ def pred_frame():
 
 # endregion
 # region File-type setup functions
-def parquet_setup():
+def pandas_parquet_setup():
     file = Path("predictions.parquet")
 
     data = pred_frame()
@@ -70,11 +72,29 @@ def parquet_setup():
     return fake_config(file)
 
 
+def non_pandas_parquet_setup():
+    file = Path("predictions.parquet")
+
+    data = pred_frame()
+    # remove the pandas metadata in the parquet file to simulate the data
+    # coming from a different parquet source (e.g. polars)
+    table = pa.Table.from_pandas(data).replace_schema_metadata({})
+
+    pq.write_table(table, file)
+    return fake_config(file)
+
+
 # endregion
 
 
 # region Tests
-@pytest.mark.parametrize("setup_fn,load_fn", [[parquet_setup, undertest.parquet_loader]])
+@pytest.mark.parametrize(
+    "setup_fn,load_fn",
+    [
+        [pandas_parquet_setup, undertest.parquet_loader],
+        [non_pandas_parquet_setup, undertest.parquet_loader],
+    ],
+)
 @pytest.mark.usefixtures("tmp_as_current")
 class TestPredictionLoad:
     def test_load_all_columns(self, setup_fn, load_fn):
@@ -107,7 +127,10 @@ def test_load_selected_columns(self, extra_columns, setup_fn, load_fn):
 
     @pytest.mark.parametrize(
         "extra_columns",
-        [pytest.param([], id="all features"), pytest.param(["column1", "column2", "column3"], id="full features")],
+        [
+            pytest.param([], id="all features"),
+            pytest.param(["column1", "column2", "column3"], id="full features"),
+        ],
     )
     def test_target_inclusion_is_renamed(self, extra_columns, setup_fn, load_fn):
         config = setup_fn()
@@ -123,7 +146,10 @@ def test_target_inclusion_is_renamed(self, extra_columns, setup_fn, load_fn):
         "desired_columns",
         [
             pytest.param(["column1", "column2", "not_in_file"], id="one unseen columns"),
-            pytest.param(["column1", "column2", "not_in_file", "another_unseen"], id="multiple unseen columns"),
+            pytest.param(
+                ["column1", "column2", "not_in_file", "another_unseen"],
+                id="multiple unseen columns",
+            ),
             pytest.param(["not_in_file"], id="only missing feature"),
             pytest.param(["column1", "not_in_file"], id="one present, one missing feature"),
         ],
@@ -136,7 +162,14 @@ def test_target_inclusion_is_renamed(self, extra_columns, setup_fn, load_fn):
         ],
     )
     def test_column_mismatch_logs_warning(
-        self, log_level, debug_present, warning_present, desired_columns, setup_fn, load_fn, caplog
+        self,
+        log_level,
+        debug_present,
+        warning_present,
+        desired_columns,
+        setup_fn,
+        load_fn,
+        caplog,
     ):
         config = setup_fn()
         config.features = desired_columns
@@ -183,7 +216,11 @@ def test_assumed_types_convert_scores(self, target_values, expected_values):
         percentage_like = [1, 50, 99]
         proba_like = [0.01, 0.50, 0.99]
         dataframe = pd.DataFrame(
-            {"target1": target_values, "nottarget_big": percentage_like, "nottarget_small": proba_like}
+            {
+                "target1": target_values,
+                "nottarget_big": percentage_like,
+                "nottarget_small": proba_like,
+            }
         )
 
         expected = dataframe.copy()
@@ -219,8 +256,16 @@ class TestDictionaryTypes:
         "defined_types,expected_dtypes",
         [
             pytest.param(
-                [{"name": "downcast", "dtype": "float16"}, {"name": "hardint", "dtype": "int"}],
-                {"downcast": "float16", "assumedTime": "datetime64[ns]", "keep": "int64", "hardint": "int"},
+                [
+                    {"name": "downcast", "dtype": "float16"},
+                    {"name": "hardint", "dtype": "int"},
+                ],
+                {
+                    "downcast": "float16",
+                    "assumedTime": "datetime64[ns]",
+                    "keep": "int64",
+                    "hardint": "int",
+                },
                 id="basic types",
             )
         ],
@@ -258,7 +303,13 @@ def test_value_error_raises_configuration_error(
             DictionaryItem(name="bad_col2", dtype="datetime64[ns]"),
         ]
 
-        dataframe = pd.DataFrame({"bad_col1": ["a", "b", "c"], "bad_col2": ["a", "b", "c"], "keep_col": [4, 5, 6]})
+        dataframe = pd.DataFrame(
+            {
+                "bad_col1": ["a", "b", "c"],
+                "bad_col2": ["a", "b", "c"],
+                "keep_col": [4, 5, 6],
+            }
+        )
 
         with pytest.raises(undertest.ConfigurationError, match="Could not convert") as cerr:
             undertest.dictionary_types(config, dataframe)
diff --git a/tests/data/test_binary_performance.py b/tests/data/test_binary_performance.py
index 3b97148d..9610f2f2 100644
--- a/tests/data/test_binary_performance.py
+++ b/tests/data/test_binary_performance.py
@@ -1,9 +1,76 @@
+from unittest.mock import Mock
+
 import numpy as np
 import pandas as pd
 import pytest
 from sklearn.metrics import auc, precision_recall_curve, roc_auc_score
 
-from seismometer.data.binary_performance import calculate_stats
+from seismometer.configuration import ConfigProvider
+from seismometer.configuration.model import Cohort, Event
+from seismometer.data.binary_performance import calculate_stats, generate_analytics_data
+from seismometer.data.loader import SeismogramLoader
+from seismometer.seismogram import Seismogram
+
+
+def get_test_config(tmp_path):
+    mock_config = Mock(autospec=ConfigProvider)
+    mock_config.output_dir.return_value
+    mock_config.events = {
+        "event1": Event(source="event1", display_name="event1", window_hr=1),
+        "event2": Event(source="event2", display_name="event2", window_hr=2, aggregation_method="min"),
+        "event3": Event(source="event3", display_name="event3", window_hr=1),
+    }
+    mock_config.target = "event1"
+    mock_config.entity_keys = ["entity"]
+    mock_config.predict_time = "time"
+    mock_config.cohorts = [Cohort(source=name) for name in ["cohort1"]]
+    mock_config.features = ["one"]
+    mock_config.config_dir = tmp_path / "config"
+    mock_config.censor_min_count = 0
+    mock_config.targets = ["event1", "event2", "event3"]
+    mock_config.output_list = ["prediction", "score1", "score2"]
+
+    return mock_config
+
+
+def get_test_loader(config):
+    mock_loader = Mock(autospec=SeismogramLoader)
+    mock_loader.config = config
+
+    return mock_loader
+
+
+def get_test_data():
+    return pd.DataFrame(
+        {
+            "entity": ["A", "A", "B", "C"],
+            "prediction": [1, 2, 3, 4],
+            "time": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
+            "event1_Value": [0, 1, 0, 1],
+            "event1_Time": ["2022-01-01", "2022-01-02", "2022-01-03", "2021-12-31"],
+            "event2_Value": [0, 1, 0, 1],
+            "event2_Time": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
+            "event3_Value": [0, 2, 5, 1],
+            "event3_Time": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
+            "cohort1": ["A", "A", "A", "B"],
+            "score1": [0.1, 0.4, 0.35, 0.8],
+            "score2": [0.2, 0.5, 0.3, 0.7],
+            "target1": [0, 1, 0, 1],
+            "target2": [1, 0, 1, 0],
+            "target3": [1, 1, 1, 0],
+        }
+    )
+
+
+@pytest.fixture
+def fake_seismo(tmp_path):
+    config = get_test_config(tmp_path)
+    loader = get_test_loader(config)
+    sg = Seismogram(config, loader)
+    sg.dataframe = get_test_data()
+    yield sg
+
+    Seismogram.kill()
 
 
 class TestCalculateStats:
@@ -123,3 +190,40 @@ def test_computed_threshold_edge_cases_all_ones(self, metric, expected_threshold
         stats = calculate_stats(df, "target", "score", metric, metric_values, metrics_to_display=["Threshold"])
         computed_thresholds = [stats[f"{val}_Threshold"] for val in metric_values]
         assert np.array_equal(computed_thresholds, expected_thresholds)
+
+
+class TestGenerateAnalyticsData:
+    def test_censor_threshold_below(self, fake_seismo):
+        # Seismogram().dataframe has fewer rows than the censor_threshold
+        result = generate_analytics_data(
+            score_columns=["score1"],
+            target_columns=["target1"],
+            metric="Sensitivity",
+            metric_values=[0.5, 0.7],
+            censor_threshold=5,
+        )
+        # Expected None when data rows are below censor_threshold
+        assert result is None
+
+    def test_censor_threshold_above(self, fake_seismo):
+        # Seismogram().dataframe has more rows than the censor_threshold
+        result = generate_analytics_data(
+            score_columns=["score1"],
+            target_columns=["target1"],
+            metric="Sensitivity",
+            metric_values=[0.5, 0.7],
+            censor_threshold=3,
+        )
+        assert result is not None
+        assert not result.empty
+
+    def test_censor_threshold_exact(self, fake_seismo):
+        # Seismogram().dataframe has exactly the same number of rows as the censor_threshold
+        result = generate_analytics_data(
+            score_columns=["score1"],
+            target_columns=["target1"],
+            metric="Sensitivity",
+            metric_values=[0.5, 0.7],
+            censor_threshold=4,
+        )
+        assert result is None
diff --git a/tests/table/test_analytics_table.py b/tests/table/test_analytics_table.py
index c5970d64..9ca1316a 100644
--- a/tests/table/test_analytics_table.py
+++ b/tests/table/test_analytics_table.py
@@ -22,12 +22,12 @@ def get_test_config(tmp_path):
     mock_config.target = "event1"
     mock_config.entity_keys = ["entity"]
     mock_config.predict_time = "time"
-    mock_config.cohorts = [Cohort(source=name) for name in ["cohort1", "cohort2"]]
+    mock_config.cohorts = [Cohort(source=name) for name in ["cohort1"]]
     mock_config.features = ["one"]
     mock_config.config_dir = tmp_path / "config"
     mock_config.censor_min_count = 0
     mock_config.targets = ["event1", "event2", "event3"]
-    mock_config.output_list = ["prediction"]
+    mock_config.output_list = ["prediction", "score1", "score2"]
 
     return mock_config
 
@@ -51,8 +51,12 @@ def get_test_data():
             "event2_Time": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
             "event3_Value": [0, 2, 5, 1],
             "event3_Time": ["2022-01-01", "2022-01-02", "2022-01-03", "2022-01-04"],
-            "cohort1": [1, 0, 1, 0],
-            "cohort2": [0, 1, 0, 1],
+            "cohort1": ["A", "A", "A", "B"],
+            "score1": [0.1, 0.4, 0.35, 0.8],
+            "score2": [0.2, 0.5, 0.3, 0.7],
+            "target1": [0, 1, 0, 1],
+            "target2": [1, 0, 1, 0],
+            "target3": [1, 1, 1, 0],
         }
     )
 
@@ -63,6 +67,7 @@ def fake_seismo(tmp_path):
     loader = get_test_loader(config)
     sg = Seismogram(config, loader)
     sg.dataframe = get_test_data()
+    sg.available_cohort_groups = {"cohort1": ["A", "B"]}
     yield sg
 
     Seismogram.kill()
@@ -91,6 +96,8 @@ def test_initialization(self, fake_seismo):
             metric_values=metric_values,
             statistics_data=statistics_data,
             table_config=table_config,
+            censor_threshold=10,
+            cohort_dict={"cohort1": ("A", "B")},
         )
 
         assert table.score_columns == scores
@@ -103,6 +110,8 @@ def test_initialization(self, fake_seismo):
         assert table.columns_show_percentages == ["Prevalence"]
         assert table.percentages_decimals == 0
         assert table.statistics_data.equals(statistics_data)
+        assert table.censor_threshold == 10
+        assert table.cohort_dict == {"cohort1": ("A", "B")}
 
     def test_invalid_metric(self, fake_seismo):
         with pytest.raises(
@@ -137,16 +146,7 @@ def test_empty_df_and_statistics_data(self, fake_seismo):
             AnalyticsTable(metric="Sensitivity")
 
     def test_non_empty_df_empty_scores(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-            }
-        )
         sg = Seismogram()
-        sg.dataframe = df
         sg.output_list = None
         with pytest.raises(
             ValueError,
@@ -164,14 +164,6 @@ def test_non_empty_df_empty_targets(self, fake_seismo):
             AnalyticsTable(score_columns=["score1", "score2"], metric="Sensitivity")
 
     def test_analytics_table_with_statistics_data(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-            }
-        )
         statistics_data = pd.DataFrame(
             {
                 "Score": ["score1", "score2", "score1", "score2"],
@@ -180,8 +172,6 @@ def test_analytics_table_with_statistics_data(self, fake_seismo):
                 "extra-stats-2": [0.6, 0.7, 0.55, 0.75],
             }
         )
-        sg = Seismogram()
-        sg.dataframe = df
         scores = ["score1", "score2"]
         targets = ["target1", "target2"]
         metric_values = [0.7, 0.8]
@@ -191,11 +181,13 @@ def test_analytics_table_with_statistics_data(self, fake_seismo):
             metric="Sensitivity",
             metric_values=metric_values,
             statistics_data=statistics_data,
+            censor_threshold=1,
+            cohort_dict={"cohort1": ("A", "B")},
         )
         gt = table.analytics_table()
         assert gt is not None
 
-    def test_analytics_table_with_statistics_data_only(self, fake_seismo):
+    def test_analytics_table_with_only_global_stats(self, fake_seismo):
         statistics_data = pd.DataFrame(
             {"Score": ["prediction"], "Target": ["event1_Value"], "Stat1": [0.75], "Stat2": [0.6]}
         )
@@ -204,81 +196,53 @@ def test_analytics_table_with_statistics_data_only(self, fake_seismo):
         table = AnalyticsTable(
             metric="Sensitivity",
             metric_values=metric_values,
+            metrics_to_display=["Positives"],
             statistics_data=statistics_data,
+            censor_threshold=1,
+            cohort_dict={"cohort1": ("A", "B")},
         )
         gt = table.analytics_table()
         assert gt is not None
 
     def test_generate_initial_table(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-                "target3": [1, 1, 1, 0],
-            }
-        )
         scores = ["score1", "score2"]
         targets = ["target1", "target2", "target3"]
-        sg = Seismogram()
-        sg.dataframe = df
-        table = AnalyticsTable(score_columns=scores, target_columns=targets, metric="Sensitivity")
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=1,
+            cohort_dict={"cohort1": ("A", "B")},
+        )
         data = table._generate_table_data()
         gt = table.generate_initial_table(data)
         assert gt is not None
 
     def test_generate_initial_table_group_by_target(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-                "target3": [1, 1, 1, 0],
-            }
-        )
         scores = ["score1", "score2"]
         targets = ["target1", "target2", "target3"]
-        sg = Seismogram()
-        sg.dataframe = df
-        table = AnalyticsTable(score_columns=scores, target_columns=targets, top_level="Target", metric="Sensitivity")
-        data = table._generate_table_data()
-        gt = table.generate_initial_table(data)
-        assert gt is not None
-
-    def test_generate_color_bar(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-            }
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            top_level="Target",
+            metric="Sensitivity",
+            censor_threshold=1,
+            cohort_dict={"cohort1": ("A", "B")},
         )
-        scores = ["score1", "score2"]
-        targets = ["target1", "target2"]
-        sg = Seismogram()
-        sg.dataframe = df
-        table = AnalyticsTable(score_columns=scores, target_columns=targets, metric="Sensitivity")
         data = table._generate_table_data()
         gt = table.generate_initial_table(data)
         assert gt is not None
 
     def test_group_columns_by_metric_value(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-            }
-        )
         scores = ["score1", "score2"]
         targets = ["target1", "target2"]
-        sg = Seismogram()
-        sg.dataframe = df
-        table = AnalyticsTable(score_columns=scores, target_columns=targets, metric="Sensitivity")
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=1,
+            cohort_dict={"cohort1": ("A", "B")},
+        )
         data = table._generate_table_data()
         gt = table.generate_initial_table(data)
         for value in table.metric_values:
@@ -287,18 +251,72 @@ def test_group_columns_by_metric_value(self, fake_seismo):
         assert gt is not None
 
     def test_analytics_table(self, fake_seismo):
-        df = pd.DataFrame(
-            {
-                "score1": [0.1, 0.4, 0.35, 0.8],
-                "score2": [0.2, 0.5, 0.3, 0.7],
-                "target1": [0, 1, 0, 1],
-                "target2": [1, 0, 1, 0],
-            }
-        )
         scores = ["score1", "score2"]
         targets = ["target1", "target2"]
-        sg = Seismogram()
-        sg.dataframe = df
-        table = AnalyticsTable(score_columns=scores, target_columns=targets, metric="Sensitivity")
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=1,
+            cohort_dict={"cohort1": ("A", "B")},
+        )
         gt = table.analytics_table()
         assert gt is not None
+
+    def test_censor_threshold(self, fake_seismo):
+        scores = ["score1", "score2"]
+        targets = ["target1", "target2"]
+
+        # Case where data is above the censor threshold
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=2,
+            cohort_dict={"cohort1": ("A", "B")},
+        )
+        data = table._generate_table_data()
+        assert data is not None  # Data should be generated as it meets the threshold
+
+        # Case where data is below the censor threshold
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=10,
+            cohort_dict={"cohort1": ("A", "B")},
+        )
+        data = table._generate_table_data()
+        assert data is None  # Data should be None as it does not meet the threshold
+
+    def test_cohort_dict_filtering(self, fake_seismo):
+        scores = ["score1", "score2"]
+        targets = ["target1", "target2"]
+
+        # Case where cohort_dict filters the data correctly
+        cohort_dict = {"cohort1": ("A",)}
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=1,
+            cohort_dict=cohort_dict,
+        )
+        data = table._generate_table_data()
+
+        # Check if the data is filtered correctly
+        assert data is not None
+
+        # Case where cohort_dict filters the data to an empty set
+        cohort_dict = {"cohort1": ("Z",)}  # "Z" is not in the cohort1 column
+        table = AnalyticsTable(
+            score_columns=scores,
+            target_columns=targets,
+            metric="Sensitivity",
+            censor_threshold=1,
+            cohort_dict=cohort_dict,
+        )
+        data = table._generate_table_data()
+
+        # Check if the data is None due to filtering to an empty set
+        assert data is None