diff --git a/NEWS.md b/NEWS.md index 91014a6b..eb2c5cee 100644 --- a/NEWS.md +++ b/NEWS.md @@ -50,6 +50,15 @@ new_var = ClimaAnalysis.convert_dim_units( ) ``` +### Slicing by interpolating +Slicing is possible by nearest value using `ClimaAnalysis.slice`. There is now support by +slicing using linear interpolation through `ClimaAnalysis.slice_intp`. See the example +below. + +```julia +ClimaAnalysis.slice_intp(var, lat = 30, lon = 20, time = 100) +``` + ## Bug fixes - `Atmos.to_pressure_coordinates` now works with Unitful units. - `Atmos.to_pressure_coordinates` now uses reasonable pressure values when `target_pressure` diff --git a/docs/src/api.md b/docs/src/api.md index 496c1e67..fe09e111 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -26,6 +26,7 @@ Var.long_name Var.units Var.has_units Var.slice +Var.slice_intp Var.average_lat Var.weighted_average_lat Var.average_lon diff --git a/src/Var.jl b/src/Var.jl index b6ef6d29..d17f0762 100644 --- a/src/Var.jl +++ b/src/Var.jl @@ -30,6 +30,7 @@ export OutputVar, average_time, is_z_1D, slice, + slice_intp, window, arecompatible, center_longitude!, @@ -931,6 +932,9 @@ end Return a new OutputVar by slicing across dimensions as defined by the keyword arguments. +Slicing is done by using the nearest value. See [`slice_intp`](@ref) for slicing by linear +interpolation. + Example =========== ```julia @@ -945,6 +949,75 @@ function slice(var; kwargs...) return sliced_var end +""" + _slice_intp_general(var::OutputVar, val, dim_name) + +Return a new OutputVar by interpolating to the given `val` for the given dimension and +slicing across it. +""" +function _slice_intp_general(var::OutputVar, val, dim_name) + haskey(var.dims, dim_name) || + error("Var does not have dimension $dim_name, found $(keys(var.dims))") + + reduced_var = _reduce_over(_slice_intp_over, dim_name, var, var, val) + + # Let's try adding this operation to the long_name, if possible (ie, if the correct + # attributes are available) + if haskey(var.attributes, "long_name") && + haskey(var.dim_attributes, dim_name) && + haskey(var.dim_attributes[dim_name], "units") + dim_units = var.dim_attributes[dim_name]["units"] + if (dim_name == "time" || dim_name == "t") && dim_units == "s" + # Dimension is time and units are seconds. Let's convert them to something nicer + pretty_timestr = seconds_to_prettystr(val) + reduced_var.attributes["long_name"] *= " $dim_name = $pretty_timestr" + else + reduced_var.attributes["long_name"] *= " $dim_name = $val $dim_units" + end + reduced_var.attributes["slice_$dim_name"] = "$val" + reduced_var.attributes["slice_$(dim_name)_units"] = dim_units + end + return reduced_var +end + +""" + _slice_intp_over(data, var::OutputVar, val; dims) + +Slice a `OutputVar` by using its interpolat at `val` for the dimension whose index is +`dims`. + +The argument `data` is included, but we do not use `data` in the function. This is because +`_reduce_over` requires the first argument to be `data`. +""" +function _slice_intp_over(data, var::OutputVar, val; dims) + dim_arrays = [values(var.dims)...] + dim_arrays[dims] = [val] + ret_data = [var(pt) for pt in Base.product(dim_arrays...)] + return ret_data +end + +""" + slice_intp(var::OutputVar, kwargs...) + +Return a new OutputVar by slicing across dimensions as defined by the keyword arguments. + +Slicing is done by linear interpolation. See [`slice`](@ref) for slicing by using the +nearest value. + +Example +=========== +```julia +slice_intp(var, lat = 30, lon = 20, time = 100) +``` +""" +function slice_intp(var::OutputVar; kwargs...) + sliced_var = var + for (dim_name, val) in kwargs + sliced_var = _slice_intp_general(sliced_var, val, String(dim_name)) + end + return sliced_var +end + """ window(var::OutputVar, dim_name; left = nothing, right = nothing) diff --git a/test/test_Var.jl b/test/test_Var.jl index fd453853..331289b2 100644 --- a/test/test_Var.jl +++ b/test/test_Var.jl @@ -1945,3 +1945,22 @@ end "rads", ) end + +@testset "Slice by interpolating" begin + lat = collect(range(-89.5, 89.5, 180)) + lon = collect(range(-42.0, 42.0, 360)) + data = reshape(collect(1:(360 * 180)), (180, 360)) + dims = OrderedDict(["lat" => lat, "lon" => lon]) + attribs = Dict("long_name" => "hi") + dim_attribs = OrderedDict(["lon" => Dict("units" => "deg")]) + var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data) + + # Check slice and slice_intp produces the same result when the value already exists + var_nearest = ClimaAnalysis.slice(var, lat = lat[1]) + var_intp = ClimaAnalysis.slice_intp(var, lat = lat[1]) + + @test var_nearest.dims == var_intp.dims + @test var_nearest.data == var_intp.data + @test var_nearest.attributes == var_intp.attributes + @test var_nearest.dim_attributes == var_nearest.dim_attributes +end