From 70a3654bc2b0cc724d8ecaebd1c613682d9d63c8 Mon Sep 17 00:00:00 2001 From: Gabriele Bozzola Date: Sun, 2 Jun 2024 11:38:46 -0700 Subject: [PATCH] Add support for interpolating to pressure Work in progress --- src/Atmos.jl | 96 ++++++++++++++++++++++++++++++++++++++++++++ src/ClimaAnalysis.jl | 2 + test/Atmos.jl | 93 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 191 insertions(+) create mode 100644 src/Atmos.jl create mode 100644 test/Atmos.jl diff --git a/src/Atmos.jl b/src/Atmos.jl new file mode 100644 index 00000000..b5afd7e6 --- /dev/null +++ b/src/Atmos.jl @@ -0,0 +1,96 @@ +""" +The `Atmos` module contains functions that are primarily useful when working +with atmospheric simulations. +""" +module Atmos + +import Interpolations as Intp +import ..OutputVar +import ..short_name + +""" + + TODO +""" +function _resample_column!(dest, var1d, origin_pressure, target_pressure) + # Interpolations.jl require increasing knots, but pressure is decreasing, so + # we have to reverse it + var1d_of_P = Intp.extrapolate( + Intp.interpolate( + (reverse(origin_pressure),), + reverse(var1d), + Intp.Gridded(Intp.Linear()), + ), + Intp.Linear(), + ) + + dest .= var1d_of_P.(target_pressure) + return nothing +end + +""" + +TODO +""" +function to_pressure_coordinates(var::OutputVar, pressure::OutputVar) + # TODO: Add error checking: + # - `var` and `pressure` have to be compatible + + ALTITUDE_NAMES = Set(["z", "z_physical", "z_reference"]) + + # Pick the correct longitude name and check that we have an altitude variable + z_name = "" + for possible_z_name in ALTITUDE_NAMES + haskey(var.dims, possible_z_name) && (z_name = possible_z_name; break) + end + z_name != "" || error("var does not have altitude among its dimensions") + z_index = var.dim2index[z_name] + pressure_name = short_name(pressure) + + # First, we construct the target pressure grid. For this, we take the + # extrema of pressure and divide the interval linearly with the same number + # of points we originally had in z + + # TODO: Pick this more sensibly + # TODO: Make it go from max to min? (This is not supported by Interpolations.jl...) + target_pressure = range( + minimum(pressure.data), + maximum(pressure.data), + length = length(var.dims[z_name]), + ) + + # Then, we prepare the output variable + ret_attributes = copy(var.attributes) + TypeOfDims = typeof(var.dims) + ret_dims = TypeOfDims( + k != z_name ? k => v : pressure_name => target_pressure for + (k, v) in var.dims + ) + TypeOfDimAttributes = typeof(var.dim_attributes) + ret_dim_attributes = TypeOfDims( + k != z_name ? k => v : pressure_name => pressure.attributes for + (k, v) in var.dim_attributes + ) + + ret_data = zeros(size(var.data)...) + + # We have to loop over all the possible columns + num_dims = ndims(var.data) + ranges = [1:size(var.data)[i] for i in 1:num_dims if i != z_index] + + # Iterate over the Cartesian product of these ranges + for idx in CartesianIndices(Tuple(ranges)) + indices = ntuple(i -> (i == z_index ? Colon() : idx[i]), num_dims) + + _resample_column!( + view(ret_data, indices...), + view(var.data, indices...), + view(pressure.data, indices...), + target_pressure, + ) + end + + return OutputVar(ret_attributes, ret_dims, ret_dim_attributes, ret_data) +end + +end diff --git a/src/ClimaAnalysis.jl b/src/ClimaAnalysis.jl index 0f185c64..77666bba 100644 --- a/src/ClimaAnalysis.jl +++ b/src/ClimaAnalysis.jl @@ -10,4 +10,6 @@ import .SimDir include("Visualize.jl") import .Visualize +include("Atmos.jl") + end # module ClimaAnalysis diff --git a/test/Atmos.jl b/test/Atmos.jl new file mode 100644 index 00000000..41236eac --- /dev/null +++ b/test/Atmos.jl @@ -0,0 +1,93 @@ +using Test +import ClimaAnalysis +import OrderedCollections: OrderedDict + +@testset "To pressure coordinates" begin + + # Let start with testing a single column + + z_alt = 0:100.0 |> collect + data = copy(z_alt) + + zvar = ClimaAnalysis.OutputVar(Dict("z" => z_alt), data) + + # Fake pressure, linearly decreasing, so that we can check precisely + pressure = 300.0:-2.0:100.0 |> collect + pdata = copy(pressure) + + attribs = Dict("short_name" => "pfull") + dim_attribs = Dict{String, Any}() + pressure_var = + ClimaAnalysis.OutputVar(attribs, Dict("z" => z_alt), dim_attribs, pdata) + + pressure_in_pressure_coordinates = + ClimaAnalysis.Atmos.to_pressure_coordinates(pressure_var, pressure_var) + + @test collect(keys(pressure_in_pressure_coordinates.dims)) == ["pfull"] + # reverse because we go from min to max for pressure (to have it increasing + # for Interpolations.jl) + @test pressure_in_pressure_coordinates.dims["pfull"] == reverse(pdata) + @test pressure_in_pressure_coordinates.data == reverse(pdata) + + # Fake var, again linear. When everything is linear we should obtain the + # input variable back + myvardata = 500.0:10.0:1500.0 |> collect + mydata = copy(myvardata) + + attribs = Dict("short_name" => "myvar") + dim_attribs = Dict{String, Any}() + myvar = ClimaAnalysis.OutputVar( + attribs, + Dict("z" => z_alt), + dim_attribs, + mydata, + ) + + myvar_in_pressure_coordinates = + ClimaAnalysis.Atmos.to_pressure_coordinates(myvar, pressure_var) + + @test collect(keys(myvar_in_pressure_coordinates.dims)) == ["pfull"] + @test myvar_in_pressure_coordinates.dims["pfull"] == reverse(pdata) + @test myvar_in_pressure_coordinates.data == reverse(mydata) + + exp_pressure = exp.(1:-0.01:0 |> collect) + exp_pdata = copy(exp_pressure) + + attribs = Dict("short_name" => "pfull") + dim_attribs = Dict{String, Any}() + exp_pressure_var = ClimaAnalysis.OutputVar( + attribs, + Dict("z" => z_alt), + dim_attribs, + exp_pdata, + ) + + myvar_in_exp_pressure_coordinates = + ClimaAnalysis.Atmos.to_pressure_coordinates(myvar, exp_pressure_var) + + # Linear range from min to max + @test myvar_in_exp_pressure_coordinates.dims["pfull"] == collect(range(1, exp(1), length = length(exp_pdata))) + # TODO: Add tests here + + # 3D test + long = 0.0:180.0 |> collect + lat = 0.0:90.0 |> collect + zzz = 0.0:10.0 |> collect + data = reshape(1.0:(91 * 181 * 11), (181, 91, 11)) + dims = OrderedDict(["lon" => long, "lat" => lat, "z" => zzz]) + dim_attribs = Dict{String, Any}() + var3D = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data) + pdata3D = data.^-1 + + attribs = Dict("short_name" => "pfull") + dim_attribs = Dict{String, Any}() + pressure3D = ClimaAnalysis.OutputVar( + attribs, + dims, + dim_attribs, + pdata3D, + ) + + ClimaAnalysis.Atmos.to_pressure_coordinates(var3D, pressure3D) + +end