Add mask function as a parameter to plot_bias

This commit adds a new parameter to `Visualize.plot_bias_on_globe!`
which is `apply_mask_to_var`. This parameter is a masking function used
when computing the bias.

NEWS.md

@@ -74,6 +74,27 @@ mask_fn = ClimaAnalysis.make_lonlat_mask(var; set_to_val = isnan, true_val = 0.0
 another_masked_var = mask_fn(another_var)
 ```
 
+### Using masking function when plotting
+Masking functions can now be passed for the `mask` keyword for plotting functions. See the
+example below of plotting with a masking function.
+
+```julia
+import ClimaAnalysis
+import ClimaAnalysis.Visualize: plot_bias_on_globe!, oceanmask
+import GeoMakie
+import CairoMakie
+
+mask_var = ClimaAnalysis.OutputVar("ocean_mask.nc")
+mask_fn = ClimaAnalysis.make_lonlat_mask(mask_var; set_to_val = isnan)
+
+obs_var = ClimaAnalysis.OutputVar("ta_1d_average.nc")
+sim_var = ClimaAnalysis.get(ClimaAnalysis.simdir("simulation_output"), "ta")
+
+fig = CairoMakie.Figure()
+plot_bias_on_globe!(fig, var, mask = mask_fn)
+CairoMakie.save("myfigure.pdf", fig)
+```
+
 ## Bug fixes
 - Masking now affects the colorbar.
 - `Var.shift_to_start_of_previous_month` now checks for duplicate dates and throws an error

docs/src/visualize.md

@@ -93,3 +93,36 @@ CairoMakie.save("myfigure.pdf", fig)
 The output produces something like:
 
 ![biasplot_oceanmask](./assets/bias_plot_oceanmask.png)
+
+We can also plot the bias using a custom mask generated from `make_lonlat_mask`.
+
+!!! note "Passing a masking function for `mask`"
+    ClimaAnalysis do not support mask keyword arguments for masking functions. If you want
+    the values of the mask to not show in a plot, then pass `true_val = NaN` as a keyword
+    argument to `make_lonlat_mask`. The color of `NaN` is controlled by the keyword
+    `nan_color` which can be passed for the plotting function (`:plot`).
+
+    Note that if the backend is CairoMakie, then the keyword `nan_color` does nothing. See
+    this [issue](https://github.com/MakieOrg/Makie.jl/issues/4524).
+
+```julia
+import ClimaAnalysis
+import ClimaAnalysis.Visualize: plot_bias_on_globe!, oceanmask
+import GeoMakie
+import CairoMakie
+
+mask_var = ClimaAnalysis.OutputVar("ocean_mask.nc")
+mask_fn = ClimaAnalysis.make_lonlat_mask(mask_var; set_to_val = isnan)
+
+obs_var = ClimaAnalysis.OutputVar("ta_1d_average.nc")
+sim_var = ClimaAnalysis.get(ClimaAnalysis.simdir("simulation_output"), "ta")
+
+fig = CairoMakie.Figure()
+plot_bias_on_globe!(fig, var, mask = mask_fn)
+CairoMakie.save("myfigure.pdf", fig)
+```
+
+The output produces something like:
+
+![bias_with_custom_mask_plot](./assets/plot_bias_with_custom_mask.png)
+

ext/ClimaAnalysisGeoMakieExt.jl

@@ -48,7 +48,7 @@ function _geomakie_plot_on_globe!(
 )
     length(var.dims) == 2 || error("Can only plot 2D variables")
 
-    apply_mask = _find_mask_to_apply(mask)
+    viz_mask, apply_mask = _find_mask_to_apply(mask)
     !isnothing(apply_mask) && (var = apply_mask(var))
 
     lon_name = ""
@@ -77,7 +77,7 @@ function _geomakie_plot_on_globe!(
     coast_kwargs = get(more_kwargs, :coast, Dict(:color => :black))
     mask_kwargs = get(more_kwargs, :mask, Dict(:color => :white))
 
-    plot_mask = !isnothing(mask)
+    plot_mask = !isnothing(viz_mask)
 
     var.attributes["long_name"] =
         ClimaAnalysis.Utils.warp_string(var.attributes["long_name"])
@@ -87,7 +87,7 @@ function _geomakie_plot_on_globe!(
     ax = GeoMakie.GeoAxis(place[p_loc...]; title, axis_kwargs...)
 
     plot = plot_fn(ax, lon, lat, var.data; plot_kwargs...)
-    plot_mask && Makie.poly!(ax, mask; mask_kwargs...)
+    plot_mask && Makie.poly!(ax, viz_mask; mask_kwargs...)
     plot_coastline && Makie.lines!(ax, GeoMakie.coastlines(); coast_kwargs...)
 
     if plot_colorbar
@@ -129,9 +129,19 @@ 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 [`Visualize.oceanmask`](@ref) and
-[`Visualize.landmask`](@ref).
+`mask` has to be an object that can be plotted by `Makie.poly` or a masking function.
+`ClimaAnalysis` comes with predefined masks, check out [`Visualize.oceanmask`](@ref) and
+[`Visualize.landmask`](@ref). Also, the corresponding mask is applied to the `OutputVar`s.
+For instance, using `Visualize.landmask` means `ClimaAnalysis.apply_landmask` is applied to
+the `OutputVar`s when computing the bias. One can also pass in
+`ClimaAnalysis.apply_landmask`, `ClimaAnalysis.apply_oceanmask`, or a custom masking
+function ([`ClimaAnalysis.Var.make_lonlat_mask`](@ref)).
+
+!!! note "Passing a masking function for `mask`"
+    ClimaAnalysis do not support mask keyword arguments for masking functions. If you want
+    the values of the mask to not show, then pass `true_val = NaN` as a keyword argument
+    to `make_lonlat_mask`. The color of `NaN` is controlled by the keyword `nan_color` which
+    can be passed for the plotting function (`:plot`).
 
 Additional arguments to the plotting and axis functions
 =======================================================
@@ -208,9 +218,19 @@ 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 [`Visualize.oceanmask`](@ref) and
-[`Visualize.landmask`](@ref).
+`mask` has to be an object that can be plotted by `Makie.poly` or a masking function.
+`ClimaAnalysis` comes with predefined masks, check out [`Visualize.oceanmask`](@ref) and
+[`Visualize.landmask`](@ref). Also, the corresponding mask is applied to the `OutputVar`s.
+For instance, using `Visualize.landmask` means `ClimaAnalysis.apply_landmask` is applied to
+the `OutputVar`s when computing the bias. One can also pass in
+`ClimaAnalysis.apply_landmask`, `ClimaAnalysis.apply_oceanmask`, or a custom masking
+function ([`ClimaAnalysis.Var.make_lonlat_mask`](@ref)).
+
+!!! note "Passing a masking function for `mask`"
+    ClimaAnalysis do not support mask keyword arguments for masking functions. If you want
+    the values of the mask to not show, then pass `true_val = NaN` as a keyword argument
+    to `make_lonlat_mask`. The color of `NaN` is controlled by the keyword `nan_color` which
+    can be passed for the plotting function (`:plot`).
 
 Additional arguments to the plotting and axis functions
 =======================================================
@@ -285,11 +305,20 @@ based on the values of `cmap_extrema`.
 
 The dimensions have to be longitude and latitude.
 
-`mask` has to be an object that can be plotted by `Makie.poly`. `ClimaAnalysis` comes with
-predefined masks, check out [`Visualize.oceanmask`](@ref) and [`Visualize.landmask`](@ref).
-Also, the corresponding mask is applied to the `OutputVar`s. For instance, using
-`Visualize.landmask` means `ClimaAnalysis.apply_landmask` is applied to the `OutputVar`s
-when computing the bias.
+`mask` has to be an object that can be plotted by `Makie.poly` or a masking function.
+`ClimaAnalysis` comes with predefined masks, check out [`Visualize.oceanmask`](@ref) and
+[`Visualize.landmask`](@ref). Also, the corresponding mask is applied to the `OutputVar`s.
+For instance, using `Visualize.landmask` means `ClimaAnalysis.apply_landmask` is applied to
+the `OutputVar`s when computing the bias. One can also pass in
+`ClimaAnalysis.apply_landmask`, `ClimaAnalysis.apply_oceanmask`, or a custom masking
+function ([`ClimaAnalysis.Var.make_lonlat_mask`](@ref)). The masking function is used for
+computing the bias.
+
+!!! note "Passing a masking function for `mask`"
+    ClimaAnalysis do not support mask keyword arguments for masking functions. If you want
+    the values of the mask to not show, then pass `true_val = NaN` as a keyword argument
+    to `make_lonlat_mask`. The color of `NaN` is controlled by the keyword `nan_color` which
+    can be passed for the plotting function (`:plot`).
 
 Additional arguments to the plotting and axis functions
 =======================================================
@@ -323,7 +352,7 @@ function Visualize.plot_bias_on_globe!(
         :mask => Dict(),
     ),
 )
-    apply_mask = _find_mask_to_apply(mask)
+    _, apply_mask = _find_mask_to_apply(mask)
 
     bias_var = ClimaAnalysis.bias(sim, obs, mask = apply_mask)
     global_bias = round(bias_var.attributes["global_bias"], sigdigits = 3)
@@ -380,22 +409,28 @@ function Visualize.plot_bias_on_globe!(
 end
 
 """
-    Return the appropriate mask to apply to a `OutputVar` from the plotting `mask`.
+    _find_mask_to_apply(mask)
 
-If `mask` is `Visualize.landmask()`, return `ClimaAnalysis.apply_landmask`. If `mask` is
-`Visualize.oceanmask()`, return `ClimaAnalysis.apply_oceanmask`. In all other cases, return
-nothing.
+Return the appropriate mask for visualizing and applying the appropriate masking function to
+the `OutputVar` from `mask`. The return type is a Tuple where the first element is the mask
+used for visualizing (e.g. Visualize.landmask, Visualize.oceanmask or a vector of polygons) and
+the second element is a mask function (e.g. Var.apply_landmask, Var.apply_oceanmask, or a
+custom masking function).
 """
 function _find_mask_to_apply(mask)
+    # The first element being returned is a mask used for plotting (e.g. Visualize.landmask)
+    # The second element being returned is a masking function (e.g. Var.apply_landmask)
     if isnothing(mask)
-        return nothing
+        return nothing, nothing
+    elseif mask isa Function
+        return nothing, mask
     elseif mask == Visualize.landmask()
-        return ClimaAnalysis.apply_landmask
+        return mask, ClimaAnalysis.apply_landmask
     elseif mask == Visualize.oceanmask()
-        return ClimaAnalysis.apply_oceanmask
+        return mask, ClimaAnalysis.apply_oceanmask
     else
         @warn "Mask not recognized, overplotting it. The colorbar will not be correct"
-        return nothing
+        return mask, nothing
     end
 end
 

test/test_GeoMakieExt.jl

@@ -237,4 +237,45 @@ using OrderedCollections
     output_name = joinpath(tmp_dir, "plot_custom_mask.png")
     Makie.save(output_name, fig13)
 
+    # Make bias with mask generating function
+    fig14 = Makie.Figure()
+
+    lon = collect(range(-179.5, 179.5, 360))
+    lat = collect(range(-89.5, 89.5, 180))
+    data = collect(reshape(-32400:32399, (360, 180))) ./ (32399.0 / 5.0)
+    dims = OrderedDict(["lon" => lon, "lat" => lat])
+    attribs =
+        Dict("long_name" => "idk", "short_name" => "short", "units" => "kg")
+    dim_attribs = OrderedDict([
+        "lon" => Dict("units" => "deg"),
+        "lat" => Dict("units" => "deg"),
+    ])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+
+    data_zero = zeros(length(lon), length(lat))
+    var_zero = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data_zero)
+
+    var_mask = ClimaAnalysis.OutputVar(ncpath, "gpp")
+    var_mask.attributes["short_name"] = "gpp"
+    var_mask = ClimaAnalysis.replace(var_mask, missing => NaN)
+    var_mask = ClimaAnalysis.slice(
+        var_mask,
+        time = ClimaAnalysis.times(var_mask) |> first,
+    )
+    mask_fn = ClimaAnalysis.make_lonlat_mask(var_mask; set_to_val = isnan)
+
+    ClimaAnalysis.Visualize.plot_bias_on_globe!(
+        fig14,
+        var,
+        var_zero,
+        mask = mask_fn,
+        cmap_extrema = (-5.0, 5.0),
+        # The keyword `nan_color` do not work right now for CairoMakie.
+        # See https://github.com/MakieOrg/Makie.jl/issues/4524
+        more_kwargs = Dict(
+            :plot => ClimaAnalysis.Utils.kwargs(nan_color = :red),
+        ),
+    )
+    output_name = joinpath(tmp_dir, "plot_bias_with_custom_mask.png")
+    Makie.save(output_name, fig14)
 end