From cff55ca037e0ae93552d078db3f8a1b83664d1bb Mon Sep 17 00:00:00 2001
From: Kevin Phan <98072684+ph-kev@users.noreply.github.com>
Date: Fri, 16 Aug 2024 15:19:52 -0700
Subject: [PATCH] Add MakieExt._constrained_cmap

The function constrained_cmap is in ClimaCoupler and should be moved to
ClimaAnalysis.
---
 ext/MakieExt.jl          | 72 ++++++++++++++++++++++++++++++++++++++++
 test/test_GeoMakieExt.jl | 51 +++++++++++++++++++++++++++-
 2 files changed, 122 insertions(+), 1 deletion(-)

diff --git a/ext/MakieExt.jl b/ext/MakieExt.jl
index 5b54e49..23feb2c 100644
--- a/ext/MakieExt.jl
+++ b/ext/MakieExt.jl
@@ -551,4 +551,76 @@ function Visualize.plot!(
     )
 end
 
+"""
+    _to_unitrange(x::Number, lo::Number, hi::Number)
+
+Linearly transform x ∈ [lo, hi] to [0, 1]. 
+"""
+_to_unitrange(x::Number, lo::Number, hi::Number) = (x - lo) / (hi - lo)
+
+"""
+    _constrained_cmap(cols::Vector, lo, hi; mid = 0, categorical = false, rev = false)
+    _constrained_cmap(cols::Makie.ColorScheme, lo, hi; mid = 0, categorical = false, rev = false)
+
+Constrain a colormap to a given range.
+
+Given a colormap implicitly defined in `± maximum(abs, (lo, hi))`, constrain it to the range
+[lo, hi]. This is useful to ensure that a colormap which is desired to diverge
+symmetrically around zero maps the same color intensity to the same magnitude.
+
+# Arguments
+- `cols`: a vector of colors, or a ColorScheme
+- `lo`: lower bound of the range
+- `hi`: upper bound of the range
+
+# Keyword Arguments
+- `mid`: midpoint of the range  # TODO: test `mid` better
+- `categorical`: flag for whether returned colormap should be categorical or continuous
+- `rev`: flag for whether to reverse the colormap before constraining cmap
+
+# Returns
+- `cmap::Makie.ColorGradient`: a colormap
+"""
+function _constrained_cmap(
+    cols::Vector,
+    lo,
+    hi;
+    mid = 0,
+    categorical = false,
+    rev = false,
+)
+    _constrained_cmap(Makie.ColorScheme(cols), lo, hi; mid, categorical, rev)
+end
+
+function _constrained_cmap(
+    cols::Makie.ColorScheme,
+    lo,
+    hi;
+    mid = 0,
+    categorical = false,
+    rev = false,
+)
+    # Reverse colorscheme if requested, don't reverse below in `cgrad`
+    rev && (cols = reverse(cols))
+    absmax = maximum(abs, (lo, hi) .- mid)
+
+    # Map lo, hi ∈ [-absmax, absmax] onto [0,1] to sample their corresponding colors
+    lo_m, hi_m = _to_unitrange.((lo, hi) .- mid, -absmax, absmax)
+
+    # Values on [0,1] where each color in cols is defined
+    colsvals = range(0, 1; length = length(cols))
+
+    # Filter colsvals, keep only values in [lo_m, hi_m] + the endpoints lo_m and hi_m
+    filter_colsvals =
+        filter(x -> lo_m <= x <= hi_m, unique([lo_m; colsvals; hi_m]))
+
+    # Select colors in filtered range; interpolate new low and hi colors
+    newcols = Makie.get(cols, filter_colsvals)
+
+    # Values on [0,1] where the new colors are defined
+    new_colsvals = _to_unitrange.(filter_colsvals, lo_m, hi_m)
+    cmap = Makie.cgrad(newcols, new_colsvals; categorical, rev = false)
+    return cmap
+end
+
 end
diff --git a/test/test_GeoMakieExt.jl b/test/test_GeoMakieExt.jl
index efb8e88..ca37624 100644
--- a/test/test_GeoMakieExt.jl
+++ b/test/test_GeoMakieExt.jl
@@ -35,9 +35,58 @@ using OrderedCollections
 
     fig2 = Makie.Figure()
 
-    ClimaAnalysis.Visualize.contour2D_on_globe!(fig2, var2D)
+    ClimaAnalysis.Visualize.contour2D_on_globe!(
+        fig2,
+        var2D,
+        more_kwargs = Dict(
+            :plot =>
+                ClimaAnalysis.Utils.kwargs(colormap = Makie.colorschemes[:vik]),
+        ),
+    )
 
     output_name = joinpath(tmp_dir, "test_contours2D_globe.png")
     Makie.save(output_name, fig2)
 
+    # Test cmap 
+    MakieExt = Base.get_extension(ClimaAnalysis, :MakieExt)
+    test_cmap = MakieExt._constrained_cmap(
+        Makie.colorschemes[:vik],
+        0.0,
+        15000.0 + (5000.0 / 3.0),
+        mid = 5000.0,
+        categorical = true,
+    )
+
+    fig3 = Makie.Figure()
+
+    ClimaAnalysis.Visualize.contour2D_on_globe!(
+        fig3,
+        var2D,
+        more_kwargs = Dict(
+            :plot => ClimaAnalysis.Utils.kwargs(colormap = test_cmap),
+        ),
+    )
+
+    output_name = joinpath(tmp_dir, "test_contours2D_globe_with_test_cmap.png")
+    Makie.save(output_name, fig3)
+
+    test_cmap = MakieExt._constrained_cmap(
+        range(Makie.colorant"red", stop = Makie.colorant"green", length = 15),
+        0.0,
+        15000.0 + (5000.0 / 3.0),
+    )
+
+    fig4 = Makie.Figure()
+
+    ClimaAnalysis.Visualize.contour2D_on_globe!(
+        fig4,
+        var2D,
+        more_kwargs = Dict(
+            :plot => ClimaAnalysis.Utils.kwargs(colormap = test_cmap),
+        ),
+    )
+
+    output_name = joinpath(tmp_dir, "test_contours2D_globe_with_test_cmap2.png")
+    Makie.save(output_name, fig4)
+
 end