FIX correct functional as labels in plot_reliability_diagram (#125)

lorentzenchr · Dec 5, 2023 · 37d43bc · 37d43bc
1 parent 78ace74
commit 37d43bc
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 21 deletions.
diff --git a/docs/examples/quantile_regression.ipynb b/docs/examples/quantile_regression.ipynb
diff --git a/src/model_diagnostics/calibration/plots.py b/src/model_diagnostics/calibration/plots.py
@@ -205,13 +205,26 @@ def iso_statistic(y_obs, y_pred, weights=None, x_values=None):
         )
         ax.plot(iso.X_thresholds_, y_plot, label=label)
 
+    xlabel_mapping = {
+        "mean": "E(Y|X)",
+        "median": "median(Y|X)",
+        "expectile": f"{level}-expectile(Y|X)",
+        "quantile": f"{level}-quantile(Y|X)",
+    }
+    ylabel_mapping = {
+        "mean": "E(Y|prediction)",
+        "median": "median(Y|prediction)",
+        "expectile": f"{level}-expectile(Y|prediction)",
+        "quantile": f"{level}-quantile(Y|prediction)",
+    }
+    xlabel = "prediction for " + xlabel_mapping[functional]
     if diagram_type == "reliability":
-        ylabel = "estimated E(Y|prediction)"
+        ylabel = "estimated " + ylabel_mapping[functional]
         title = "Reliability Diagram"
     else:
-        ylabel = "prediction - estimated E(Y|prediction)"
+        ylabel = "prediction - estimated " + ylabel_mapping[functional]
         title = "Bias Reliability Diagram"
-    ax.set(xlabel="prediction for E(Y|X)", ylabel=ylabel)
+    ax.set(xlabel=xlabel, ylabel=ylabel)
 
     if n_pred >= 2:
         ax.set_title(title)

diff --git a/src/model_diagnostics/calibration/tests/test_plots.py b/src/model_diagnostics/calibration/tests/test_plots.py
@@ -79,14 +79,30 @@ def test_plot_reliability_diagram(diagram_type, functional, n_bootstrap, weights
         diagram_type=diagram_type,
     )
 
+    xlabel_mapping = {
+        "mean": "E(Y|X)",
+        "median": "median(Y|X)",
+        "expectile": "0.8-expectile(Y|X)",
+        "quantile": "0.8-quantile(Y|X)",
+    }
+    ylabel_mapping = {
+        "mean": "E(Y|prediction)",
+        "median": "median(Y|prediction)",
+        "expectile": "0.8-expectile(Y|prediction)",
+        "quantile": "0.8-quantile(Y|prediction)",
+    }
+
     if ax is not None:
         assert ax is plt_ax
-    assert plt_ax.get_xlabel() == "prediction for E(Y|X)"
+    assert plt_ax.get_xlabel() == "prediction for " + xlabel_mapping[functional]
     if diagram_type == "reliability":
-        assert plt_ax.get_ylabel() == "estimated E(Y|prediction)"
+        assert plt_ax.get_ylabel() == "estimated " + ylabel_mapping[functional]
         assert plt_ax.get_title() == "Reliability Diagram"
     else:
-        assert plt_ax.get_ylabel() == "prediction - estimated E(Y|prediction)"
+        assert (
+            plt_ax.get_ylabel()
+            == "prediction - estimated " + ylabel_mapping[functional]
+        )
         assert plt_ax.get_title() == "Bias Reliability Diagram"
 
     plt_ax = plot_reliability_diagram(