diff --git a/schedview/data/sample_opsim.db b/schedview/data/sample_opsim.db index b8e79f3c..5fa78f84 100644 Binary files a/schedview/data/sample_opsim.db and b/schedview/data/sample_opsim.db differ diff --git a/schedview/data/sample_rewards.h5 b/schedview/data/sample_rewards.h5 index 90f9d6e3..2bf73ec5 100644 Binary files a/schedview/data/sample_rewards.h5 and b/schedview/data/sample_rewards.h5 differ diff --git a/schedview/data/sample_scheduler.pickle.xz b/schedview/data/sample_scheduler.pickle.xz index bae805e8..512e080a 100644 Binary files a/schedview/data/sample_scheduler.pickle.xz and b/schedview/data/sample_scheduler.pickle.xz differ diff --git a/schedview/plot/visits.py b/schedview/plot/visits.py index 10085776..7669fbc0 100644 --- a/schedview/plot/visits.py +++ b/schedview/plot/visits.py @@ -4,6 +4,7 @@ import bokeh.io import bokeh.layouts import bokeh.models +import bokeh.palettes import bokeh.plotting import bokeh.transform import hvplot @@ -17,6 +18,89 @@ from .colors import make_band_cmap + +def _get_observation_reason_mapping( + source: bokeh.models.ColumnarDataSource, + observation_reason_threshold: int, +) -> tuple[tuple[str, ...] | None, str | None, list[str] | None]: + """Get color palette and column mapping for observation_reason values. + + Parameters + ---------- + source : ColumnarDataSource + Data source for the plot (already contains observation_reason column). + observation_reason_threshold : int + Maximum number of distinct observation_reason values to assign + individual colors. If there are more distinct values than this + threshold, the most common ones get individual colors and the + remainder are grouped into an "other" category. + + Returns + ------- + color_palette : tuple of str or None + Color palette for the factors, or None if no observation_reason data. + color_column_name : str or None + Column name to use for the color mapping, or None if no + observation_reason data. + color_factors : list of str or None + List of observation_reason values to map, or None if + no observation_reason data. + """ + # Get observation_reason values from source data + observation_reasons = source.data.get("observation_reason", np.array([])) + + if len(observation_reasons) == 0: + return None, None, None + + # Get unique observation_reason values and their counts + unique, counts = np.unique(observation_reasons, return_counts=True) + value_counts = pd.Series(counts, index=unique) + + # Get sorted list of observation_reasons by frequency + sorted_reasons = list(value_counts.index) + + color_column_name: str = "observation_reason" + color_factors = sorted_reasons + + # If there are too many reasons for a reasonable number of categories, + # collapse reasons with common prefixes into common categories + # until the number of categories becomes manageable. + match_prefixes = [p.split("_")[0] for p in color_factors] + for match_prefix in match_prefixes: + if len(color_factors) > observation_reason_threshold: + unmatched_reasons = [r for r in color_factors if not r.startswith(match_prefix + "_")] + + source.data["obs_reason_col"] = np.where( + np.isin(source.data.get(color_column_name, np.array([])), unmatched_reasons), + source.data.get(color_column_name, np.array([])), + match_prefix, + ) + color_column_name = "obs_reason_col" + color_factors = unmatched_reasons + [match_prefix] + + # If there are still too many categories, collapse the least common + # into an "other" category + if len(color_factors) > observation_reason_threshold: + common_reasons = color_factors[:observation_reason_threshold] + # Create a new column with "other" for rare values + source.data[color_column_name] = np.where( + np.isin(source.data.get(color_column_name, np.array([])), common_reasons), + source.data.get(color_column_name, np.array([])), + "other", + ) + color_column_name = "obs_reason_col" + color_factors = common_reasons + ["other"] + + # Create the color palette + num_colors_needed = len(color_factors) + if num_colors_needed <= 8: + color_palette = tuple(bokeh.palettes.Colorblind8[:num_colors_needed]) + else: + color_palette = tuple(bokeh.palettes.TolRainbow[num_colors_needed]) + + return color_palette, color_column_name, color_factors + + TIMELINE_TOOLTIPS = [ ( "Start time", @@ -134,6 +218,8 @@ def plot_visit_param_vs_time( user_figure_kwargs: dict | None = None, default_sim: str = "All", sim_alpha: float = 0.2, + color_by_observation_reason: bool = False, + observation_reason_threshold: int = 10, **kwargs, ) -> bokeh.models.ui.ui_element.UIElement: """Plot a column in the visit table vs. time. @@ -164,6 +250,15 @@ def plot_visit_param_vs_time( Label of default sim to show (defaults to ``All``) sim_alpha: float The alpha of points for visible simulations. + color_by_observation_reason: `bool` + If True, color points by observation_reason instead of band. + Defaults to False. + observation_reason_threshold: `int` + Maximum number of distinct observation_reason values to assign + individual colors. If there are more distinct values than + this threshold, the most common ones will get individual + colors and the remainder will be grouped into an "other" + category. Defaults to 10. **kwargs Additional keyword arguments to be passed to `bokeh.plotting.figure.scatter`. @@ -227,7 +322,18 @@ def plot_visit_param_vs_time( scatter_kwargs = {"fill_alpha": 0.0, "line_alpha": "sim_alpha", "name": "timeline"} - if color_transform is None: + if color_by_observation_reason: + color_palette, color_column_name, color_factors = _get_observation_reason_mapping( + source, observation_reason_threshold + ) + if color_palette is not None and color_column_name is not None and color_factors is not None: + scatter_kwargs["color"] = bokeh.transform.factor_cmap( + color_column_name, color_palette, color_factors + ) + else: + # No observation_reason data, fall back to band colors + scatter_kwargs["color"] = make_band_cmap(band_column(visits)) + elif color_transform is None: scatter_kwargs["color"] = make_band_cmap(band_column(visits)) else: scatter_kwargs["color"] = bokeh.transform.factor_cmap( diff --git a/tests/test_plot_visits.py b/tests/test_plot_visits.py index fe12371c..0ee8a640 100644 --- a/tests/test_plot_visits.py +++ b/tests/test_plot_visits.py @@ -26,3 +26,71 @@ def test_plot_visit_param_vs_time(self): def test_create_visit_table(self): plot = schedview.plot.create_visit_table(self.visits, show=False) self.assertIsInstance(plot, bokeh.models.ui.ui_element.UIElement) + + def test_plot_visit_param_vs_time_color_by_observation_reason(self): + plot = schedview.plot.plot_visit_param_vs_time( + self.visits, "seeingFwhmEff", color_by_observation_reason=True + ) + self.assertIsInstance(plot, bokeh.models.ui.ui_element.UIElement) + + # The plot contains figures with renderers + # Get the first figure from the layout + ui_children = plot.children if hasattr(plot, "children") else [] + figures = [child for child in ui_children if hasattr(child, "renderers")] + self.assertGreater(len(figures), 0) + + # Check that the scatter glyph uses a factor_cmap for + # observation_reason + scatter_renderer = figures[0].renderers[0] + fill_color = scatter_renderer.glyph.fill_color + # fill_color is a Field object with .transform containing + # the factor_cmap. + transform = fill_color.transform + # Check that transform is a factor_cmap by checking it has + # expected attributes. + self.assertTrue(hasattr(transform, "factors")) + self.assertTrue(hasattr(transform, "palette")) + + # Check that observation_reason values are in the factors + # The factors should contain the unique observation_reason values + unique_reasons = set(self.visits["observation_reason"].unique()) + for reason in unique_reasons: + self.assertIn(reason, transform.factors) + + # Check that the number of factors matches the number of unique + # observation_reason values in the data. + self.assertEqual(len(transform.factors), len(unique_reasons)) + + def test_plot_visit_param_vs_time_color_by_observation_reason_with_threshold(self): + # Add many observation_reason values to test the threshold behavior + visits = self.visits.copy() + # Create 15 unique observation_reason values (more than default + # threshold of 10). + # Assign unique names to the first 15 visits. + for i in range(15): + visits.loc[visits.index[i], "observation_reason"] = f"reason_{i}" + # Keep the rest as the original values. + + # With 15 unique values (more than threshold of 10), we should + # collapse names with common prefixes into single categories. + plot = schedview.plot.plot_visit_param_vs_time( + visits, "seeingFwhmEff", color_by_observation_reason=True + ) + self.assertIsInstance(plot, bokeh.models.ui.ui_element.UIElement) + + # Get the figures from the layout + ui_children = plot.children if hasattr(plot, "children") else [] + figures = [child for child in ui_children if hasattr(child, "renderers")] + self.assertGreater(len(figures), 0) + + # Check that the scatter glyph uses a factor_cmap + scatter_renderer = figures[0].renderers[0] + fill_color = scatter_renderer.glyph.fill_color + transform = fill_color.transform + + # Check that transform has the expected attributes + self.assertTrue(hasattr(transform, "factors")) + self.assertTrue(hasattr(transform, "palette")) + + # Check that there are 11 factors (10 common reasons + "other") + self.assertLess(len(transform.factors), 11)