From f5d7fed40516d313cb2677eb233876d08b4996b7 Mon Sep 17 00:00:00 2001
From: Kevin Phan <98072684+ph-kev@users.noreply.github.com>
Date: Thu, 21 Nov 2024 17:39:16 -0800
Subject: [PATCH] Add Var.slice_intp

---
 NEWS.md          |  9 ++++++
 docs/src/api.md  |  1 +
 src/Var.jl       | 73 ++++++++++++++++++++++++++++++++++++++++++++++++
 test/test_Var.jl | 19 +++++++++++++
 4 files changed, 102 insertions(+)

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