diff --git a/NEWS.md b/NEWS.md index a8a7e650..5138ede2 100644 --- a/NEWS.md +++ b/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 ------ diff --git a/docs/src/api.md b/docs/src/api.md index 226ec993..11158bf8 100644 --- a/docs/src/api.md +++ b/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 diff --git a/src/Var.jl b/src/Var.jl index c6966b34..413eea09 100644 --- a/src/Var.jl +++ b/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) diff --git a/test/test_Var.jl b/test/test_Var.jl index 24c36b7a..98a51b08 100644 --- a/test/test_Var.jl +++ b/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")