Add support for interpolating to pressure

Work in progress

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

src/ClimaAnalysis.jl

@@ -10,4 +10,6 @@ import .SimDir
 include("Visualize.jl")
 import .Visualize
 
+include("Atmos.jl")
+
 end # module ClimaAnalysis

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