Merge pull request #334 from uriahf/fix/color-values-propagation

Honor custom colors in time-dependent curves

Author: uriahf

src/rtichoke/discrimination/gains.py

@@ -13,6 +13,7 @@
     _create_reference_lines_data,
     _check_if_multiple_populations_are_being_validated_times,
 )
+from rtichoke.processing.time_reference_lines import _apply_color_values_times
 from rtichoke.performance_data.performance_data_times import prepare_performance_data_times
 import numpy as np
 import polars as pl
@@ -251,6 +252,7 @@ def create_gains_curve_times(
         color_value=color_values,
         curve="gains",
     )
+    curve_list = _apply_color_values_times(curve_list, color_values)
     curve_list = _replace_gains_reference_data_times(curve_list, performance_data)
 
     return _create_plotly_curve_times(curve_list)

src/rtichoke/processing/time_reference_lines.py

@@ -27,6 +27,33 @@ def _get_reference_aj_estimates_times(performance_data: pl.DataFrame) -> pl.Data
     )
 
 
+def _apply_color_values_times(curve_list: dict, color_values) -> dict:
+    """Apply custom colors while retaining R's single-model black styling."""
+    if color_values is None or not curve_list["multiple_reference_groups"]:
+        return curve_list
+
+    reference_groups = curve_list["reference_group_keys"]
+    if len(color_values) < len(reference_groups):
+        raise ValueError(
+            "color_values must contain at least one color per reference group"
+        )
+
+    colors_dictionary = curve_list["colors_dictionary"]
+    for index, reference_group in enumerate(reference_groups):
+        color = color_values[index]
+        for key in (
+            reference_group,
+            f"random_guess_{reference_group}",
+            f"perfect_model_{reference_group}",
+            f"treat_none_{reference_group}",
+            f"treat_all_{reference_group}",
+        ):
+            if key in colors_dictionary:
+                colors_dictionary[key] = color
+
+    return curve_list
+
+
 def _replace_reference_data_times(
     curve_list: dict,
     performance_data: pl.DataFrame,
@@ -88,6 +115,7 @@ def _create_rtichoke_plotly_curve_times_reference_safe(
         min_p_threshold=min_p_threshold,
         max_p_threshold=max_p_threshold,
     )
+    curve_list = _apply_color_values_times(curve_list, color_values)
     return _create_plotly_curve_times(
         _replace_reference_data_times(
             curve_list,

tests/test_time_curve_colors.py

@@ -0,0 +1,66 @@
+import pytest
+
+from rtichoke.processing.time_reference_lines import _apply_color_values_times
+
+
+def _curve_list(reference_groups, multiple_reference_groups=True):
+    colors_dictionary = {
+        "random_guess": "#BEBEBE",
+        "perfect_model": "#BEBEBE",
+        "treat_none": "#BEBEBE",
+        "treat_all": "#BEBEBE",
+    }
+    for group in reference_groups:
+        for key in (
+            group,
+            f"random_guess_{group}",
+            f"perfect_model_{group}",
+            f"treat_none_{group}",
+            f"treat_all_{group}",
+        ):
+            colors_dictionary[key] = "#000000"
+
+    return {
+        "reference_group_keys": reference_groups,
+        "multiple_reference_groups": multiple_reference_groups,
+        "colors_dictionary": colors_dictionary,
+    }
+
+
+def test_time_curve_custom_colors_propagate_to_groups_and_references():
+    curve_list = _curve_list(["model_a", "model_b"])
+
+    result = _apply_color_values_times(curve_list, ["#111111", "#222222"])
+
+    for group, expected_color in (
+        ("model_a", "#111111"),
+        ("model_b", "#222222"),
+    ):
+        for key in (
+            group,
+            f"random_guess_{group}",
+            f"perfect_model_{group}",
+            f"treat_none_{group}",
+            f"treat_all_{group}",
+        ):
+            assert result["colors_dictionary"][key] == expected_color
+
+    assert result["colors_dictionary"]["random_guess"] == "#BEBEBE"
+    assert result["colors_dictionary"]["perfect_model"] == "#BEBEBE"
+    assert result["colors_dictionary"]["treat_none"] == "#BEBEBE"
+    assert result["colors_dictionary"]["treat_all"] == "#BEBEBE"
+
+
+def test_time_curve_single_model_keeps_existing_black_style():
+    curve_list = _curve_list(["model"], multiple_reference_groups=False)
+
+    result = _apply_color_values_times(curve_list, ["#123456"])
+
+    assert result["colors_dictionary"]["model"] == "#000000"
+
+
+def test_time_curve_custom_colors_require_one_per_reference_group():
+    curve_list = _curve_list(["model_a", "model_b"])
+
+    with pytest.raises(ValueError, match="one color per reference group"):
+        _apply_color_values_times(curve_list, ["#111111"])