Add option to mask ocean/land

docs/make.jl

@@ -29,6 +29,7 @@ makedocs(;
     pages = [
         "Home" => "index.md",
         "OutputVars" => "var.md",
+        "Visualizing OutputVars" => "visualize.md",
         "APIs" => "api.md",
         "How do I?" => "howdoi.md",
     ],

docs/src/api.md

@@ -93,6 +93,8 @@ Visualize.plot!
 ## GeoMakie
 
 ```@docs
+Visualize.oceanmask
+Visualize.landmask
 Visualize.contour2D_on_globe!
 Visualize.heatmap2D_on_globe!
 ```

docs/src/var.md

@@ -29,7 +29,7 @@ When possible, `ClimaAnalysis` uses
 [Unitful](https://painterqubits.github.io/Unitful.jl/stable) to handle units.
 This enables automatic unit conversion for `OutputVar`s.
 
-Consider the following example: 
+Consider the following example:
 ```julia
 import ClimaAnalysis
 values = 0:100.0 |> collect

docs/src/visualize.md

@@ -0,0 +1,47 @@
+# Visualizing `OutputVar`s
+
+This page is under construction, in the meantime, consult [`Visualize`](@ref).
+
+### Masking part of the output in `GeoMakie`
+
+When performing ocean or land simulations, it is often convenient to hide the
+other component (e.g., hide the ocean and focus on the continents). For
+`GeoMakie` plots, there is a direct way to accomplish this. In this section, we
+discuss this feature.
+
+The main `GeoMakie` plots are [`Visualize.contour2D_on_globe!`](@ref) and
+[`Visualize.heatmap2D_on_globe!`](@ref). Both these functions take a `mask` argument. By
+default, `mask` is `nothing`, meaning that the entire output is displayed on the
+globe. Alternatively, `mask` can be a collection of polygons that can be plotted
+with [`Makie.poly`](https://docs.makie.org/v0.21/reference/plots/poly).
+`ClimaAnalysis` comes with the most important ones [`Visualize.oceanmask`](@ref) and
+[`Visualize.landmask`](@ref), to hide ocean and continents respectively.
+
+For example, suppose `var` it the variable we want to plot with a ocean mask
+```julia
+import ClimaAnalysis.Visualize: contour2D_on_globe!, oceanmask
+import ClimaAnalysis.Utils: kwargs as ca_kwargs
+import GeoMakie
+import CairoMakie
+
+fig = CairoMakie.Figure()
+
+contour2D_on_globe!(fig,
+                    var,
+                    mask = oceanmask(),
+                    more_kwargs = Dict(:mask => ca_kwargs(color = :blue)),
+                   )
+
+CairoMakie.save("myfigure.pdf", fig)
+```
+
+In this example, we plotted `var` on the globe and overplotted a blue ocean.
+`ca_kwargs` (`Utils.kwargs`) is a convenience function to pass keyword arguments
+more easily.
+
+!!! note Masking does not affect the colorbar. If you have values defined
+    beneath the map, they can still affect the colorbar.
+
+The output might look something like:
+
+![oceanmask](./assets/oceanmask.png)

ext/ClimaAnalysisGeoMakieExt.jl

@@ -7,17 +7,42 @@ import ClimaAnalysis: Visualize
 
 MakiePlace = Union{Makie.Figure, Makie.GridLayout}
 
+"""
+    oceanmask()
+
+Return a collection of polygons to mask out the ocean.
+
+Plot with `Makie.poly`.
+"""
+function Visualize.oceanmask()
+    elevation = 0
+    return GeoMakie.NaturalEarth.bathymetry(elevation)
+end
+
+"""
+    landmask()
+
+Return a collection of polygons to mask out the continents.
+
+Plot with `Makie.poly`.
+"""
+function Visualize.landmask()
+    return GeoMakie.land()
+end
+
 function _geomakie_plot_on_globe!(
     place::MakiePlace,
     var::ClimaAnalysis.OutputVar;
     p_loc = (1, 1),
     plot_coastline = true,
     plot_colorbar = true,
+    mask = nothing,
     more_kwargs = Dict(
         :plot => Dict(),
         :cb => Dict(),
         :axis => Dict(),
         :coast => Dict(:color => :black),
+        :mask => Dict(),
     ),
     plot_fn = Makie.surface!,
 )
@@ -47,6 +72,9 @@ function _geomakie_plot_on_globe!(
     plot_kwargs = get(more_kwargs, :plot, Dict())
     cb_kwargs = get(more_kwargs, :cb, Dict())
     coast_kwargs = get(more_kwargs, :coast, Dict(:color => :black))
+    mask_kwargs = get(more_kwargs, :mask, Dict(:color => :white))
+
+    plot_mask = !isnothing(mask)
 
     var.attributes["long_name"] =
         ClimaAnalysis.Utils.warp_string(var.attributes["long_name"])
@@ -56,6 +84,7 @@ function _geomakie_plot_on_globe!(
     GeoMakie.GeoAxis(place[p_loc...]; title, axis_kwargs...)
 
     plot = plot_fn(lon, lat, var.data; plot_kwargs...)
+    plot_mask && Makie.poly!(mask; mask_kwargs...)
     plot_coastline && Makie.lines!(GeoMakie.coastlines(); coast_kwargs...)
 
     if plot_colorbar
@@ -75,12 +104,14 @@ end
                         p_loc = (1,1),
                         plot_coastline = true,
                         plot_colorbar = true,
+                        mask = nothing,
                         more_kwargs)
     heatmap2D_on_globe!(grid_layout::Makie.GridLayout,
                         var::ClimaAnalysis.OutputVar;
                         p_loc = (1,1),
                         plot_coastline = true,
                         plot_colorbar = true,
+                        mask = nothing,
                         more_kwargs)
 
 
@@ -95,6 +126,13 @@ This function assumes that the following attributes are available:
 
 The dimensions have to be longitude and latitude.
 
+`mask` has to be an object that can be plotted by `Makie.poly`. Typically, an ocean or land
+mask. `ClimaAnalysis` comes with predefined masks, check out [`oceanmask`](@ref) and
+[`landmask`](@ref).
+
+!!! note Masking does not affect the colorbar. If you have values defined beneath the map,
+    they can still affect the colorbar.
+
 Additional arguments to the plotting and axis functions
 =======================================================
 
@@ -103,6 +141,7 @@ Additional arguments to the plotting and axis functions
 - the plotting function (`:plot`)
 - the colorbar (`:cb`)
 - the coastline (`:coast`)
+- the mask (`:mask`)
 
 The coastline is plotted from `GeoMakie.coastline` using the `lines!` plotting function.
 
@@ -115,11 +154,13 @@ function Visualize.heatmap2D_on_globe!(
     p_loc = (1, 1),
     plot_coastline = true,
     plot_colorbar = true,
+    mask = nothing,
     more_kwargs = Dict(
         :plot => Dict(),
         :cb => Dict(),
         :axis => Dict(),
         :coast => Dict(:color => :black),
+        :mask => Dict(),
     ),
 )
     return _geomakie_plot_on_globe!(
@@ -128,6 +169,7 @@ function Visualize.heatmap2D_on_globe!(
         p_loc,
         plot_coastline,
         plot_colorbar,
+        mask,
         more_kwargs,
         plot_fn = Makie.surface!,
     )
@@ -140,13 +182,15 @@ end
                         plot_coastline = true,
                         plot_colorbar = true,
                         plot_contours = true,
+                        mask = nothing,
                         more_kwargs)
     contours2D_on_globe!(grid_layout::Makie.GridLayout,
                         var::ClimaAnalysis.OutputVar;
                         p_loc = (1,1),
                         plot_coastline = true,
                         plot_colorbar = true,
                         plot_contours = true,
+                        mask = nothing,
                         more_kwargs)
 
 
@@ -161,6 +205,13 @@ This function assumes that the following attributes are available:
 
 The dimensions have to be longitude and latitude.
 
+`mask` has to be an object that can be plotted by `Makie.poly`. Typically, an ocean or land
+mask. `ClimaAnalysis` comes with predefined masks, check out [`oceanmask`](@ref) and
+[`landmask`](@ref).
+
+!!! note Masking does not affect the colorbar. If you have values defined beneath the map,
+    they can still affect the colorbar.
+
 Additional arguments to the plotting and axis functions
 =======================================================
 
@@ -169,6 +220,7 @@ Additional arguments to the plotting and axis functions
 - the plotting function (`:plot`)
 - the colorbar (`:cb`)
 - the coastline (`:coast`)
+- the mask (`:mask`)
 
 The coastline is plotted from `GeoMakie.coastline` using the `lines!` plotting function.
 
@@ -181,11 +233,13 @@ function Visualize.contour2D_on_globe!(
     p_loc = (1, 1),
     plot_coastline = true,
     plot_colorbar = true,
+    mask = nothing,
     more_kwargs = Dict(
         :plot => Dict(),
         :cb => Dict(),
         :axis => Dict(),
         :coast => Dict(:color => :black),
+        :mask => Dict(),
     ),
 )
     _geomakie_plot_on_globe!(
@@ -194,6 +248,7 @@ function Visualize.contour2D_on_globe!(
         p_loc,
         plot_coastline,
         plot_colorbar,
+        mask,
         more_kwargs,
         plot_fn = Makie.contourf!,
     )

src/Visualize.jl

@@ -4,6 +4,10 @@ export plot!
 
 function _constrained_cmap end
 
+function oceanmask end
+
+function landmask end
+
 function heatmap2D! end
 
 function sliced_heatmap! end

test/test_GeoMakieExt.jl

@@ -88,4 +88,16 @@ using OrderedCollections
     output_name = joinpath(tmp_dir, "test_contours2D_globe_with_test_cmap2.png")
     Makie.save(output_name, fig4)
 
+    # Test with oceanmask
+    fig5 = Makie.Figure()
+
+    ClimaAnalysis.Visualize.heatmap2D_on_globe!(
+        fig5,
+        var2D,
+        mask = ClimaAnalysis.Visualize.oceanmask(),
+        more_kwargs = Dict(:mask => ClimaAnalysis.Utils.kwargs(color = :blue)),
+    )
+
+    output_name = joinpath(tmp_dir, "test_contours2D_globe_with_oceanmask.png")
+    Makie.save(output_name, fig5)
 end