Add Var.apply_landmask and Var.apply_oceanmask

Author: ph-kev

docs/src/api.md

@@ -68,6 +68,8 @@ Var.squared_error
 Var.global_mse
 Var.global_rmse
 Var.shift_to_start_of_previous_month
+Var.apply_landmask
+Var.apply_oceanmask
 ```
 
 ## Leaderboard

docs/src/howdoi.md

@@ -152,3 +152,15 @@ OrderedDict{String, Vector{Float64}} with 2 entries:
   "lon"      => [0.0, 1.0]
   "latitude" => [0.0, 1.0, 2.0]
 ```
+
+## How do I apply a land or sea mask to a `OutputVar`?
+
+You can use `apply_landmask` or `apply_oceanmask` to mask out the land or ocean,
+respectively, in a `OutputVar`. The result of `apply_landmask(var)` is data of `var`, where
+any coordinate corresponding to land is zero. Similarly, the result of `apply_oceanmask(var)` is
+data of `var`, where any coordinate corresponding to ocean is zero.
+
+```@julia masks
+var_no_land = ClimaAnalysis.apply_landmask(var)
+var_no_ocean = ClimaAnalysis.apply_oceanmask(var)
+```

masks/land_mask16.nc

@@ -53,7 +53,9 @@ export OutputVar,
     global_mse,
     global_rmse,
     set_units,
-    shift_to_start_of_previous_month
+    shift_to_start_of_previous_month,
+    apply_landmask,
+    apply_oceanmask
 
 """
     Representing an output variable
@@ -1447,6 +1449,78 @@ function shift_to_start_of_previous_month(var::OutputVar)
     return OutputVar(ret_attribs, ret_dims, ret_dim_attributes, ret_data)
 end
 
+"""
+    apply_landmask(var::OutputVar)
+
+Apply a land mask to `var` by zeroing any data whose coordinates are located on land.
+"""
+function apply_landmask(var::OutputVar)
+    _apply_lonlat_mask(var, LAND_MASK)
+end
+
+"""
+    apply_oceanmask(var::OutputVar)
+
+Apply an ocean mask to `var` by zeroing any data whose coordinates are in the ocean.
+"""
+function apply_oceanmask(var::OutputVar)
+    _apply_lonlat_mask(var, OCEAN_MASK)
+end
+
+"""
+    _apply_lonlat_mask(var, mask::AbstractString)
+
+Apply a mask using the NCDataset found at `mask`.
+
+The dimensions of the mask should only contain longitude and latitude and are in that order.
+"""
+function _apply_lonlat_mask(var, mask::AbstractString)
+    # Check if lon and lat exist
+    has_longitude(var) ||
+        error("Longitude does not exist as a dimension in var")
+    has_latitude(var) || error("Latitude does not exist as a dimension in var")
+
+    # Create OutputVar using mask
+    mask_var = OutputVar(mask)
+
+    # Resample so that the land mask match up with the grid of var
+    # Round because linear resampling is done and we want the mask to be only ones and zeros
+    land_mask =
+        [
+            mask_var(pt) for pt in Base.product(
+                var.dims[longitude_name(var)],
+                var.dims[latitude_name(var)],
+            )
+        ] .|> round
+
+    # Reshape data for broadcasting
+    lon_idx = var.dim2index[longitude_name(var)]
+    lat_idx = var.dim2index[latitude_name(var)]
+    lon_length = var.dims[longitude_name(mask_var)] |> length
+    lat_length = var.dims[latitude_name(mask_var)] |> length
+    if lon_idx > lat_idx
+        land_mask = transpose(land_mask)
+    end
+    size_to_reshape = (
+        if i == lon_idx
+            lon_length
+        elseif i == lat_idx
+            lat_length
+        else
+            1
+        end for i in 1:ndims(var.data)
+    )
+
+    # Mask data
+    land_mask = reshape(land_mask, size_to_reshape...)
+    masked_data = var.data .* land_mask
+
+    # Remake OutputVar with new data
+    ret_attribs = deepcopy(var.attributes)
+    ret_dims = deepcopy(var.dims)
+    ret_dim_attributes = deepcopy(var.dim_attributes)
+    return OutputVar(ret_attribs, ret_dims, ret_dim_attributes, masked_data)
+end
 
 """
     overload_binary_op(op)
@@ -1545,5 +1619,6 @@ end
 @overload_binary_op (/)
 
 include("outvar_dimensions.jl")
+include("constants.jl")
 
 end

masks/ocean_mask16.nc

@@ -0,0 +1,23 @@
+"""
+The variable LAND_MASK is a filepath to a NCDataset whose dimensions are latitude and
+longitude and data are ones and zeros. The zeros indicate land and ones indicate everything
+else.
+
+The land mask is generated from the ETOPO2022 dataset at 60 arc-second resolutions
+(https://www.ncei.noaa.gov/products/etopo-global-relief-model). The dimensions latitude and
+longitude are sampled every 16th element to thin out the dimensions. Any value for z greater
+than 0 is assigned 0.0 and all other values for z is assigned 1.0.
+"""
+LAND_MASK = joinpath(@__DIR__, "..", "masks", "land_mask16.nc")
+
+"""
+The variable OCEAN_MASK is a filepath to a NCDataset whose dimensions are latitude and
+longitude and data are 1s and 0s. The zeros indicate ocean and ones indicate everything
+else.
+
+The land mask is generated from the ETOPO2022 dataset at 60 arc-second resolutions
+(https://www.ncei.noaa.gov/products/etopo-global-relief-model). The dimensions latitude and
+longitude are sampled every 16th element to thin out the dimensions. Any value for z less
+than or equal to 0 is assigned 0.0 and all other values for z is assigned 1.0.
+"""
+OCEAN_MASK = joinpath(@__DIR__, "..", "masks", "ocean_mask16.nc")

src/Var.jl

@@ -136,4 +136,29 @@ using OrderedCollections
     )
     output_name = joinpath(tmp_dir, "plot_bias_kwargs.png")
     Makie.save(output_name, fig7)
+
+    # Test plots with apply_landmask and apply_oceanmask
+    fig8 = Makie.Figure()
+
+    lon = collect(range(-179.5, 179.5, 360))
+    lat = collect(range(-89.5, 89.5, 180))
+    data = ones(length(lon), length(lat))
+    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)
+    land_var = var |> ClimaAnalysis.apply_landmask
+    ClimaAnalysis.Visualize.heatmap2D_on_globe!(fig8, land_var)
+    output_name = joinpath(tmp_dir, "plot_apply_land_mask.png")
+    Makie.save(output_name, fig8)
+
+    fig9 = Makie.Figure()
+    ocean_var = var |> ClimaAnalysis.apply_oceanmask
+    ClimaAnalysis.Visualize.heatmap2D_on_globe!(fig9, ocean_var)
+    output_name = joinpath(tmp_dir, "plot_apply_ocean_mask.png")
+    Makie.save(output_name, fig9)
 end

src/constants.jl

@@ -1506,3 +1506,63 @@ end
         var_min,
     )
 end
+
+@testset "Land and ocean masks" begin
+    # Order of dimensions should not matter
+    lat = collect(range(-89.5, 89.5, 180))
+    lon = collect(range(-179.5, 179.5, 360))
+    data = ones(length(lat), length(lon))
+    dims = OrderedDict(["lat" => lat, "lon" => lon])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict([
+        "lat" => Dict("units" => "deg"),
+        "lon" => Dict("units" => "deg"),
+    ])
+    var_latlon = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+
+    lon = collect(range(-179.5, 179.5, 360))
+    lat = collect(range(-89.5, 89.5, 180))
+    data = ones(length(lon), length(lat))
+    dims = OrderedDict(["lon" => lon, "lat" => lat])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict([
+        "lon" => Dict("units" => "deg"),
+        "lat" => Dict("units" => "deg"),
+    ])
+    var_lonlat = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+
+    land_var_lonlat = ClimaAnalysis.apply_landmask(var_lonlat)
+    ocean_var_lonlat = ClimaAnalysis.apply_oceanmask(var_lonlat)
+    land_var_latlon = ClimaAnalysis.apply_landmask(var_latlon)
+    ocean_var_latlon = ClimaAnalysis.apply_oceanmask(var_latlon)
+    @test land_var_lonlat.data == land_var_latlon.data'
+    @test ocean_var_lonlat.data == ocean_var_latlon.data'
+
+    # Testing with another dimension
+    lat = collect(range(-89.5, 89.5, 180))
+    times = collect(range(0.0, 100, 2 * 180))
+    lon = collect(range(-179.5, 179.5, 360))
+    data = ones(length(lat), length(times), length(lon))
+    dims = OrderedDict(["lat" => lat, "time" => times, "lon" => lon])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict([
+        "lat" => Dict("units" => "deg"),
+        "time" => Dict("units" => "s"),
+        "lon" => Dict("units" => "deg"),
+    ])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    land_var = ClimaAnalysis.apply_landmask(var) |> ClimaAnalysis.average_time
+    ocean_var = ClimaAnalysis.apply_oceanmask(var) |> ClimaAnalysis.average_time
+    @test land_var.data |> transpose == land_var_lonlat.data
+    @test ocean_var.data |> transpose == ocean_var_lonlat.data
+
+    # Test error handling
+    times = collect(range(0.0, 100, 2 * 180))
+    data = ones(length(times))
+    dims = OrderedDict(["time" => times])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict(["time" => Dict("units" => "s")])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    @test_throws ErrorException ClimaAnalysis.apply_landmask(var)
+    @test_throws ErrorException ClimaAnalysis.apply_oceanmask(var)
+end