Add weighted_average_lat

And add option to add weights to  `average_lat`

NEWS.md

@@ -4,6 +4,7 @@ ClimaAnalysis.jl Release Notes
 v0.5.6
 ------
 - Fix finding variables with name like `clwup_1m_40s_inst.nc` (composed period).
+- Add support for weighted averages in `average_lat`.
 
 v0.5.5
 ------

docs/src/api.md

@@ -25,6 +25,7 @@ Var.long_name
 Var.units
 Var.slice
 Var.average_lat
+Var.weighted_average_lat
 Var.average_lon
 Var.average_x
 Var.average_y

src/Var.jl

@@ -10,6 +10,7 @@ import ..Utils: nearest_index, seconds_to_prettystr, squeeze
 export OutputVar,
     read_var,
     average_lat,
+    weighted_average_lat,
     average_lon,
     average_x,
     average_y,
@@ -203,11 +204,32 @@ function _reduce_over(
 end
 
 """
-    average_lat(var::OutputVar)
+    average_lat(var::OutputVar; weighted = false)
 
 Return a new OutputVar where the values on the latitudes are averaged arithmetically.
-"""
-function average_lat(var)
+
+When `weighted` is `true`, weight the average over `cos(lat)`.
+"""
+function average_lat(var; weighted = false)
+    if weighted
+        var = copy(var)
+        lats = latitudes(var)
+        maximum(lats) >= 0.5π ||
+            @warn "Detected latitudes are small. If units are radians, results will be wrong"
+
+        weights_1d = cosd.(lats)
+        lat_index = var.dim2index[latitude_name(var)]
+        weights = ones(size(var.data))
+        for index in CartesianIndices(weights)
+            index_tuple = ntuple(
+                d -> d == lat_index ? Colon() : index[d], ndims(weights)
+            )
+
+            weights[index_tuple...] .= weights_1d
+        end
+        var.data .*= weights
+    end
+
     reduced_var = _reduce_over(mean, latitude_name(var), var)
 
     if haskey(var.attributes, "long_name")
@@ -216,6 +238,14 @@ function average_lat(var)
     return reduced_var
 end
 
+"""
+    weighted_average_lat(var::OutputVar)
+
+Return a new OutputVar where the values on the latitudes are averaged arithmetically
+with weights of `cos(lat)`.
+"""
+weighted_average_lat(var) = average_lat(var; weighted = true)
+
 """
     average_lon(var::OutputVar)
 

test/test_Var.jl

@@ -49,7 +49,7 @@ end
     lat = 0.0:90.0 |> collect
     time = 0.0:10.0 |> collect
 
-    data1 = reshape(1.0:(91 * 181 * 10), (10, 181, 91))
+    data1 = reshape(1.0:(91 * 181 * 11), (11, 181, 91))
 
     dims = OrderedDict(["time" => time, "lon" => long, "lat" => lat])
     dim_attributes = OrderedDict([
@@ -68,7 +68,7 @@ end
 
     var2 = ClimaAnalysis.OutputVar(attribs, dims, dim_attributes2, data1)
 
-    data3 = 5.0 .+ reshape(1.0:(91 * 181 * 10), (10, 181, 91))
+    data3 = 5.0 .+ reshape(1.0:(91 * 181 * 11), (11, 181, 91))
     attribs3 = Dict("long_name" => "bob", "short_name" => "bula")
     var3 = ClimaAnalysis.OutputVar(attribs3, dims, dim_attributes, data3)
 
@@ -100,7 +100,7 @@ end
     lat = 0.0:90.0 |> collect
     time = 0.0:10.0 |> collect
 
-    data = reshape(1.0:(91 * 181 * 10), (10, 181, 91))
+    data = reshape(1.0:(91 * 181 * 11), (11, 181, 91))
 
     dims = OrderedDict(["time" => time, "lon" => long, "lat" => lat])
     dim_attributes = OrderedDict([
@@ -126,6 +126,34 @@ end
           OrderedDict(["lon" => Dict("b" => 2), "time" => Dict()])
     @test lat_avg.data == dropdims(mean(data, dims = 3), dims = 3)
 
+    wei_lat_avg = ClimaAnalysis.weighted_average_lat(var)
+    @test wei_lat_avg.dims == OrderedDict(["lon" => long, "time" => time])
+    @test wei_lat_avg.dim_attributes ==
+          OrderedDict(["lon" => Dict("b" => 2), "time" => Dict()])
+    weights = ones(size(data))
+    for i in eachindex(time)
+        for j in eachindex(long)
+            for k in eachindex(lat)
+                weights[i, j, k] = cosd(lat[k])
+            end
+        end
+    end
+    expected_avg = dropdims(mean(data .* weights, dims = 3), dims = 3)
+    @test wei_lat_avg.data == expected_avg
+
+    wrong_dims = OrderedDict(["lat" => [0.0, 0.1]])
+    wrong_dim_attributes = OrderedDict(["lat" => Dict("a" => 1)])
+    wrong_var = ClimaAnalysis.OutputVar(
+        Dict{String, Any}(),
+        wrong_dims,
+        wrong_dim_attributes,
+        [0.0, 0.1],
+    )
+    @test_logs (
+        :warn,
+        "Detected latitudes are small. If units are radians, results will be wrong",
+    )
+
     lat_lon_avg = ClimaAnalysis.average_lon(lat_avg)
     @test lat_lon_avg.dims == OrderedDict(["time" => time])
     @test lat_lon_avg.dim_attributes == OrderedDict(["time" => Dict()])
@@ -147,7 +175,7 @@ end
     y = 0.0:90.0 |> collect
     time = 0.0:10.0 |> collect
 
-    data = reshape(1.0:(91 * 181 * 10), (10, 181, 91))
+    data = reshape(1.0:(91 * 181 * 11), (11, 181, 91))
 
     # Identical test pattern to sphere setup, with `dims` modified.
     dims = OrderedDict(["time" => time, "x" => x, "y" => y])
@@ -273,7 +301,7 @@ end
     lat = 0.0:90.0 |> collect
     time = 0.0:10.0 |> collect
 
-    data = reshape(1.0:(91 * 181 * 10), (10, 181, 91))
+    data = reshape(1.0:(91 * 181 * 11), (11, 181, 91))
 
     dims = OrderedDict(["time" => time, "lon" => long, "lat" => lat])
     attribs = Dict("short_name" => "bob", "long_name" => "hi")