From 08e32654effaf4dacdfe77dab1182151d071e0d0 Mon Sep 17 00:00:00 2001
From: Ziyi Yin <ziyi.yin@gatech.edu>
Date: Thu, 13 Oct 2022 13:54:34 -0400
Subject: [PATCH] add function to assign different portions of a single
 colormap to different sub-colormaps

---
 examples/plot_example.jl |  7 +++++++
 src/SlimPlotting.jl      | 41 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+)

diff --git a/examples/plot_example.jl b/examples/plot_example.jl
index f557968..610c5b8 100644
--- a/examples/plot_example.jl
+++ b/examples/plot_example.jl
@@ -68,3 +68,10 @@ tight_layout()
 # Wiggle plot
 figure(figsize=(10, 5))
 wiggle_plot(shot[1:5:end, 1:10:end], xloc[1:10:end], 0:0.02:4.6; new_fig=false)
+
+# plot with multiple colorscales
+figure(figsize=(10, 5))
+intervals = [(-3,-2),(-1,1),(2,3)]
+cmap_types = [matplotlib.cm.YlGn,matplotlib.cm.Purples,matplotlib.cm.Reds]
+cmap = create_multi_color_map(intervals, cmap_types)
+plot_simage(randn(Float32, 100, 100), (1, 1); cmap=cmap, new_fig=false, interp="none", name="colorful white noise")
diff --git a/src/SlimPlotting.jl b/src/SlimPlotting.jl
index af172a4..842a6cb 100644
--- a/src/SlimPlotting.jl
+++ b/src/SlimPlotting.jl
@@ -17,6 +17,7 @@ end
 
 export plot_fslice, plot_velocity, plot_simage, plot_sdata, wiggle_plot
 export colorschemes
+export create_multi_color_map
 
 """
     _plot_with_units(image, spacing; perc=95, cmap=:cet_CET_L1, 
@@ -286,4 +287,44 @@ function wiggle_plot(data::Array{Td, 2}, xrec=nothing, time_axis=nothing;
     ax.set_ylabel("Time")
 end
 
+"""
+    create_multi_color_map(intervals, cmap_types)
+
+Create a single colormap that assigns different portions of it to different sub-colormaps.
+
+# Arguments
+  - `intervals::Vector{Tuple{T, T}}`: intervals of values where each interval will be assigned to a different colormap
+  - `cmap_types::Vector{ColorMap}`: colormaps for each interval
+  - `out_of_interval_color`: (Optional) the color assigned to values that live outside of all the intervals, default is black [0.,0.,0.,1.]
+
+# Example
+    create_multi_color_map([(1,2),(3,4),(5,6)], [matplotlib.cm.YlGn,matplotlib.cm.Purples,matplotlib.cm.Reds])
+"""
+function create_multi_color_map(intervals::Vector{Tuple{T, T}}, cmap_types::Vector{ColorMap}; out_of_interval_color=[0.,0.,0.,1.]) where T
+
+    @assert length(intervals) == length(cmap_types) "number of intervals and number of colormaps do not match"
+    
+    # ensure at least 100 points in each interval
+    colorrange = range(minimum(intervals)[1], stop=maximum(intervals)[2],
+                length=Int(round(1f2/minimum([intervals[i][2]-intervals[i][1] for i = 1:length(intervals)]))))
+             
+    # number of points in each colormap
+    colorlength = [Int.(round.((intervals[i][2]-intervals[i][1])./(colorrange[2]-colorrange[1]))) for i = 1:length(intervals)]
+    
+    # create each colormap
+    cmap_sections = [cmap_types[i](range(0f0, stop=1f0, length=colorlength[i])) for i = 1:length(cmap_types)]
+
+    # fill in out-of-interval by out_of_interval_color
+    between_interval_sections = [reshape(repeat(out_of_interval_color, inner=Int(round((intervals[i][1]-intervals[i-1][2])/(colorrange[2]-colorrange[1])))), :, 4) for i = 2:length(intervals)]
+    
+    # concatenate all the sections to get the entire colormap
+    total_sections = vcat(cmap_sections[1], vcat([vcat(between_interval_sections[i], cmap_sections[i+1]) for i = 1:length(cmap_sections)-1]...))
+
+    # construct the colormaps
+    cmaps = matplotlib.colors.ListedColormap(total_sections)
+    
+    return cmaps
+    
+end
+
 end # module