update example notebook, add annotate_textcolors_threshold parameter

PCMDI · Nov 8, 2023 · dca315b · dca315b
1 parent 3b4959d
commit dca315b
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 58 deletions.
diff --git a/pcmdi_metrics/graphics/portrait_plot/portrait_plot_example.ipynb b/pcmdi_metrics/graphics/portrait_plot/portrait_plot_example.ipynb
diff --git a/pcmdi_metrics/graphics/portrait_plot/portrait_plot_lib.py b/pcmdi_metrics/graphics/portrait_plot/portrait_plot_lib.py
@@ -17,6 +17,7 @@ def portrait_plot(
     annotate=False,
     annotate_data=None,
     annotate_textcolors=("black", "white"),
+    annotate_textcolors_threshold=(-2, 2),
     annotate_fontsize=15,
     annotate_format="{x:.2f}",
     figsize=(12, 10),
@@ -63,6 +64,7 @@ def portrait_plot(
                   but work only for heatmap style map (i.e., no triangles)
     - `annotate_data`: 2d numpy array, default=None. If None, the image's data is used.  Optional.
     - `annotate_textcolors`: Tuple. A pair of colors for annotation text. Default is ("black", "white")
+    - `annotate_textcolors_threshold`: Tuple or float. Value in data units according to which the colors from textcolors are applied. Default=(-2, 2)
     - `annotate_fontsize`: number (int/float), default=15. Font size for annotation
     - `annotate_format`: format for annotate value, default="{x:.2f}"
     - `figsize`: tuple of two numbers (width, height), default=(12, 10), figure size in inches
@@ -171,14 +173,14 @@ def portrait_plot(
                     sys.exit("Error: annotate_data has different size than data")
             else:
                 annotate_data = data
-            annotate_heatmap(
+            ax = annotate_heatmap(
                 im,
                 ax=ax,
                 data=data,
                 annotate_data=annotate_data,
                 valfmt=annotate_format,
                 textcolors=annotate_textcolors,
-                threshold=(2, -2),
+                threshold=annotate_textcolors_threshold,
                 fontsize=annotate_fontsize,
             )
 
@@ -503,6 +505,8 @@ def annotate_heatmap(
                 kw.update(color=textcolors[int(data[i, j] > threshold)])
             text = ax.text(j + 0.5, i + 0.5, valfmt(annotate_data[i, j], None), **kw)
             texts.append(text)
+
+    return ax
 
 
 # ======================================================================