diff --git a/src/rtichoke/discrimination/gains.py b/src/rtichoke/discrimination/gains.py index 51d506f4..817b88ee 100644 --- a/src/rtichoke/discrimination/gains.py +++ b/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) diff --git a/src/rtichoke/processing/time_reference_lines.py b/src/rtichoke/processing/time_reference_lines.py index c7fd8d1b..ccecb7c6 100644 --- a/src/rtichoke/processing/time_reference_lines.py +++ b/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, diff --git a/tests/test_time_curve_colors.py b/tests/test_time_curve_colors.py new file mode 100644 index 00000000..f0bbd3a9 --- /dev/null +++ b/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"])