Add Var._make_interpolant and boundary conditions

ph-kev · ph-kev · commit e42b99615ee1 · 2024-10-01T18:14:02.000Z
This commit refactors the code making interpolats for OutputVars into the
function `Var._make_interpolant`. A periodic boundary condition is added
for the longitude dimension and a flat boundary condition is added for
the latitude dimension when the array is equispaced and spans the entire
range. Lastly, if the element type of any dimension array is DateTime,
then no interpolant will be created.
diff --git a/NEWS.md b/NEWS.md
@@ -14,6 +14,16 @@ the units or units are missing.
 new_var = ClimaAnalysis.set_units(var, "kg m s^-1")
 ```
 
+### Extrapolating `OutputVar` on longitude and latitude
+Extrapolation is now possible for the longitude and latitude dimensions. If the dimension
+arrays are equispaced and span the entire range, then a periodic boundary condition is added
+for the longitude dimension and a flat boundary condition is added for the latitude
+dimension.
+
+## Bug fixes
+
+- Interpolation is not possible with dates. When dates are detected in any dimension, an
+  interpolat will not be made.
 
 v0.5.9
 ------
diff --git a/docs/src/var.md b/docs/src/var.md
@@ -63,6 +63,23 @@ new_var = ClimaAnalysis.set_units(var, "kg m s^-1")
 !!! warning "Override existing units"
     If units already exist, this will override the units for data in `var`.
 
+## Interpolations and extrapolations
+
+Interpolating a `OutputVar` onto coordinates can be done by doing the following:
+```julia
+var((0.0, 0.0)) # var is a two-dimensional OutputVar
+```
+
+A multilinear interpolation is used to determine the value at the coordinate (0, 0).
+!!! warning "Interpolate on dates"
+    If any of the dimensions contains `Dates.DateTime` elements, interpolation is not
+    possible. `Interpolations.jl` does not support interpolating on dates.
+
+Extrapolating is supported only on the longitude and latitude dimensions. For the longitude
+and latitude dimensions, a periodic boundary condition and a flat boundary condition are
+added, respectively, when the dimension array is equispaced and spans the entire range. For
+all other cases, extrapolating beyond the domain of the dimension will throw an error.
+
 ## Integration
 
 `OutputVar`s can be integrated with respect to longitude, latitude, or both using
diff --git a/src/Var.jl b/src/Var.jl
@@ -16,7 +16,8 @@ import ..Utils:
     split_by_season,
     time_to_date,
     date_to_time,
-    _data_at_dim_vals
+    _data_at_dim_vals,
+    _isequispaced
 
 export OutputVar,
     read_var,
@@ -80,37 +81,91 @@ struct OutputVar{T <: AbstractArray, A <: AbstractArray, B, C, ITP}
     interpolant::ITP
 end
 
-function OutputVar(attribs, dims, dim_attribs, data)
-    index2dim = keys(dims) |> collect
-    dim2index =
-        Dict([dim_name => index for (index, dim_name) in enumerate(keys(dims))])
+"""
+    _make_interpolant(dims, data)
+
+Make a linear interpolant from `dims`, a dictionary mapping dimension name to array and
+`data`, an array containing data. Used in constructing a `OutputVar`.
+
+If any element of the arrays in `dims` is a Dates.DateTime, then no interpolant is returned.
+Interpolations.jl does not support interpolating on dates. If the longitudes span the entire
+range and are equispaced, then a periodic boundary condition is added for the longitude
+dimension. If the latitudes span the entire range and are equispaced, then a flat boundary
+condition is added for the latitude dimension. In all other cases, an error is thrown when
+extrapolating outside of `dim_array`.
+"""
+function _make_interpolant(dims, data)
+    # If any element is DateTime, then return nothing for the interpolant because
+    # Interpolations.jl do not support DateTimes
+    for dim_array in values(dims)
+        eltype(dim_array) <: Dates.DateTime && return nothing
+    end
 
     # We can only create interpolants when we have 1D dimensions
-    if isempty(index2dim) ||
-       any(d -> ndims(d) != 1 || length(d) == 1, values(dims))
-        itp = nothing
-    else
-        # Dimensions are all 1D, check that they are compatible with data
-        size_data = size(data)
-        for (dim_index, (dim_name, dim_array)) in enumerate(dims)
-            dim_length = length(dim_array)
-            data_length = size_data[dim_index]
-            if dim_length != data_length
-                error(
-                    "Dimension $dim_name has inconsistent size with provided data ($dim_length != $data_length)",
-                )
-            end
+    if isempty(dims) || any(d -> ndims(d) != 1 || length(d) == 1, values(dims))
+        return nothing
+    end
+
+    # Dimensions are all 1D, check that they are compatible with data
+    size_data = size(data)
+    for (dim_index, (dim_name, dim_array)) in enumerate(dims)
+        dim_length = length(dim_array)
+        data_length = size_data[dim_index]
+        if dim_length != data_length
+            error(
+                "Dimension $dim_name has inconsistent size with provided data ($dim_length != $data_length)",
+            )
         end
+    end
 
-        dims_tuple = tuple(values(dims)...)
+    # Find boundary conditions for extrapolation
+    extp_bound_conds = (
+        _find_extp_bound_cond(dim_name, dim_array) for
+        (dim_name, dim_array) in dims
+    )
 
-        # TODO: Make this lazy: we should compute the spline the first time we use
-        # it, not when we create the object
-        itp = Intp.extrapolate(
-            Intp.interpolate(dims_tuple, data, Intp.Gridded(Intp.Linear())),
-            Intp.Throw(),
-        )
-    end
+    dims_tuple = tuple(values(dims)...)
+    extp_bound_conds_tuple = tuple(extp_bound_conds...)
+    return Intp.extrapolate(
+        Intp.interpolate(dims_tuple, data, Intp.Gridded(Intp.Linear())),
+        extp_bound_conds_tuple,
+    )
+end
+
+"""
+    _find_extp_bound_cond(dim_name, dim_array)
+
+Find the appropriate boundary condition for the `dim_name` dimension.
+"""
+function _find_extp_bound_cond(dim_name, dim_array)
+    min_of_dim, max_of_dim = extrema(dim_array)
+    dim_size = max_of_dim - min_of_dim
+    dsize = dim_array[begin + 1] - dim_array[begin]
+
+    # If the dimension array span the entire range and is equispaced, then add the
+    # appropriate boundary condition
+    # We do not handle the cases when the array is not equispaced
+    (
+        conventional_dim_name(dim_name) == "longitude" &&
+        _isequispaced(dim_array) &&
+        isapprox(dim_size + dsize, 360.0)
+    ) && return Intp.Periodic()
+    (
+        conventional_dim_name(dim_name) == "latitude" &&
+        _isequispaced(dim_array) &&
+        isapprox(dim_size + dsize, 180.0)
+    ) && return Intp.Flat()
+    return Intp.Throw()
+end
+
+function OutputVar(attribs, dims, dim_attribs, data)
+    index2dim = keys(dims) |> collect
+    dim2index =
+        Dict([dim_name => index for (index, dim_name) in enumerate(keys(dims))])
+
+    # TODO: Make this lazy: we should compute the spline the first time we use
+    # it, not when we create the object
+    itp = _make_interpolant(dims, data)
 
     function _maybe_process_key_value(k, v)
         k != "units" && return k => v
@@ -710,7 +765,14 @@ end
 Interpolate variable `x` onto the given `target_coord` coordinate using
 multilinear interpolation.
 
-Extrapolation is now allowed and will throw a `BoundsError`.
+Extrapolation is now allowed and will throw a `BoundsError` in most cases.
+
+If any element of the arrays of the dimensions is a Dates.DateTime, then interpolation is
+not possible. Interpolations.jl do not support making interpolations for dates. If the
+longitudes span the entire range and are equispaced, then a periodic boundary condition is
+added for the longitude dimension. If the latitudes span the entire range and are
+equispaced, then a flat boundary condition is added for the latitude dimension. In all other
+cases, an error is thrown when extrapolating outside of the array of the dimension.
 
 Example
 =======
diff --git a/test/test_Var.jl b/test/test_Var.jl
@@ -6,6 +6,7 @@ import NaNStatistics: nanmean
 import NCDatasets: NCDataset
 import OrderedCollections: OrderedDict
 import Unitful: @u_str
+import Dates
 
 @testset "General" begin
     # Add test for short constructor
@@ -85,6 +86,61 @@ import Unitful: @u_str
     )
 end
 
+@testset "Interpolant boundary conditions" begin
+    # Check boundary condtions for lon (equispaced and span), lat (equispaced and span), and
+    # time
+    lon = 0.5:1.0:359.5 |> collect
+    lat = -89.5:1.0:89.5 |> collect
+    time = 1.0:100 |> collect
+    data = ones(length(lon), length(lat), length(time))
+    dims = OrderedDict(["lon" => lon, "lat" => lat, "time" => time])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict([
+        "long" => Dict("units" => "test_units1"),
+        "lat" => Dict("units" => "test_units2"),
+        "time" => Dict("units" => "test_units3"),
+    ])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    @test var.interpolant.et == (Intp.Periodic(), Intp.Flat(), Intp.Throw())
+
+    # Not equispaced for lon and lat
+    lon = 0.5:1.0:359.5 |> collect |> x -> push!(x, 42.0) |> sort
+    lat = -89.5:1.0:89.5 |> collect |> x -> push!(x, 42.0) |> sort
+    time = 1.0:100 |> collect
+    data = ones(length(lon), length(lat), length(time))
+    dims = OrderedDict(["lon" => lon, "lat" => lat, "time" => time])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    @test var.interpolant.et == (Intp.Throw(), Intp.Throw(), Intp.Throw())
+
+    # Does not span entire range for and lat
+    lon = 0.5:1.0:350.5 |> collect
+    lat = -89.5:1.0:80.5 |> collect
+    time = 1.0:100 |> collect
+    data = ones(length(lon), length(lat), length(time))
+    dims = OrderedDict(["lon" => lon, "lat" => lat, "time" => time])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    @test var.interpolant.et == (Intp.Throw(), Intp.Throw(), Intp.Throw())
+
+    # Dates for the time dimension
+    lon = 0.5:1.0:359.5 |> collect
+    lat = -89.5:1.0:89.5 |> collect
+    time = [
+        Dates.DateTime(2020, 3, 1, 1, 1),
+        Dates.DateTime(2020, 3, 1, 1, 2),
+        Dates.DateTime(2020, 3, 1, 1, 3),
+    ]
+    data = ones(length(lon), length(lat), length(time))
+    dims = OrderedDict(["lon" => lon, "lat" => lat, "time" => time])
+    attribs = Dict("long_name" => "hi")
+    dim_attribs = OrderedDict([
+        "long" => Dict("units" => "test_units1"),
+        "lat" => Dict("units" => "test_units2"),
+        "time" => Dict("units" => "test_units3"),
+    ])
+    var = ClimaAnalysis.OutputVar(attribs, dims, dim_attribs, data)
+    @test isnothing(var.interpolant)
+end
+
 @testset "empty" begin
     dims = OrderedDict{String, Vector{Float64}}()
     data = Float64[]
@@ -423,7 +479,7 @@ end
 
 @testset "Interpolation" begin
     # 1D interpolation with linear data, should yield correct results
-    long = -180.0:180.0 |> collect
+    long = -175.0:175.0 |> collect
     data = copy(long)
 
     longvar = ClimaAnalysis.OutputVar(Dict("long" => long), data)
