Automatically generate a masking function using a OutputVar

NEWS.md

@@ -63,6 +63,38 @@ the units of a dimension.
 new_var = ClimaAnalysis.set_dim_units!(var, "lon", "degrees_east")
 ```
 
+### Automatically generate a mask from a OutputVar
+Masking function can automatically be generated from the function `make_lonlat_mask`. See
+the example below of generating a masking function that mask out any data that is `NaN` in
+`var`.
+
+```julia
+# Replace NaN with 0.0 and everything else with 1.0 for the mask
+mask_fn = ClimaAnalysis.make_lonlat_mask(var; set_to_val = isnan, true_val = 0.0, false_val = 1.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/api.md

@@ -71,6 +71,7 @@ Var.global_rmse
 Var.shift_to_start_of_previous_month
 Var.apply_landmask
 Var.apply_oceanmask
+Var.make_lonlat_mask
 Base.replace(var::OutputVar, old_new::Pair...)
 ```
 

docs/src/var.md

@@ -262,18 +262,45 @@ julia> units(se_var)
 "K^2"
 ```
 
-### Masking
-Bias and squared error can be computed only over the land or ocean through the `mask` parameter.
-As of now, the mask parameter takes in `apply_oceanmask` or `apply_oceanmask`. See the
-example below of this usage.
+## Masking
+Bias and squared error can be computed only over the land or ocean through the `mask`
+parameter. As of now, the mask parameter takes in `apply_oceanmask` or `apply_oceanmask`.
+See the example below of this usage.
 
 ```julia
 # Do not consider the ocean when computing the bias
-ClimaAnalysis.bias(sim_var, obs_var; mask = apply_oceanmask)
-ClimaAnalysis.global_bias(sim_var, obs_var; mask = apply_oceanmask)
+ClimaAnalysis.bias(sim_var, obs_var, mask = apply_oceanmask)
+ClimaAnalysis.global_bias(sim_var, obs_var, mask = apply_oceanmask)
 
 # Do not consider the land when computing the squared error
-ClimaAnalysis.squared_error(sim_var, obs_var; mask = apply_landmask)
-ClimaAnalysis.global_mse(sim_var, obs_var; mask = apply_landmask)
-ClimaAnalysis.global_rmse(sim_var, obs_var; mask = apply_landmask)
+ClimaAnalysis.squared_error(sim_var, obs_var, mask = apply_landmask)
+ClimaAnalysis.global_mse(sim_var, obs_var, mask = apply_landmask)
+ClimaAnalysis.global_rmse(sim_var, obs_var, mask = apply_landmask)
+```
+
+In other cases, you may want to generate a masking function using a `OutputVar`. For
+instance, you are comparing against observational data over land and you can't use an ocean
+mask since not all of the observational data is defined over the land. The function
+`make_lonlat_mask` allows you to generate a masking function. If
+the data is already zeros and ones, then you can use `make_lonlat_mask(var)`. Otherwise, you
+can specify `set_to_val` which takes in an element of `var.data` and return a boolean. If
+`set_to_val` returns `true`, then the value will be `true_val` in the mask and if `set_to_val`
+returns `false`, then the value will be `false_val` in the mask. See the example below of this
+usage.
+
+```julia
+# Any points that are NaNs should be zero in the mask
+mask_fn = ClimaAnalysis.make_lonlat_mask(
+    var;
+    set_to_val = isnan,
+    true_val = 0.0, # default is NaN
+    false_val = 1.0,
+)
+
+# Apply mask to another OutputVar
+another_masked_var = mask_fn(another_var)
+
+# Compute squared error and global MSE with custom masking function
+ClimaAnalysis.squared_error(sim_var, obs_var, mask = mask_fn)
+ClimaAnalysis.global_mse(sim_var, obs_var, mask = mask_fn)
 ```

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