ft: add gridded interpolation to DimensionalData

Project.toml

@@ -11,6 +11,7 @@ DataAPI = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Extents = "411431e0-e8b7-467b-b5e0-f676ba4f2910"
 Interfaces = "85a1e053-f937-4924-92a5-1367d23b7b87"
+Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f"
 IteratorInterfaceExtensions = "82899510-4779-5014-852e-03e436cf321d"
@@ -22,6 +23,7 @@ RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 TableTraits = "3783bdb8-4a98-5b6b-af9a-565f29a5fe9c"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [weakdeps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -34,13 +36,16 @@ NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 PythonCall = "6099a3de-0909-46bc-b1f4-468b9a2dfc0d"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
+UnrolledUtilities = "0fe1646c-419e-43be-ac14-22321958931b"
 
 [extensions]
 DimensionalDataAbstractFFTsExt = "AbstractFFTs"
 DimensionalDataAlgebraOfGraphicsExt = "AlgebraOfGraphics"
 DimensionalDataCategoricalArraysExt = "CategoricalArrays"
 DimensionalDataChainRulesCoreExt = "ChainRulesCore"
 DimensionalDataDiskArraysExt = "DiskArrays"
+DimensionalDataInterpolationsExt = ["Interpolations", "UnrolledUtilities"]
 DimensionalDataMakieExt = "Makie"
 DimensionalDataNearestNeighborsExt = "NearestNeighbors"
 DimensionalDataPythonCallExt = "PythonCall"
@@ -73,6 +78,7 @@ GPUArrays = "10"
 ImageFiltering = "0.7"
 ImageTransformations = "0.10"
 Interfaces = "0.3"
+Interpolations = "0.16.2"
 IntervalSets = "0.5, 0.6, 0.7"
 InvertedIndices = "1"
 IteratorInterfaceExtensions = "1"
@@ -97,6 +103,7 @@ TableTraits = "1"
 Tables = "1"
 Test = "1"
 Unitful = "1"
+UnrolledUtilities = "0.1.9"
 julia = "1.9"
 
 [extras]

ext/DimensionalDataInterpolationsExt.jl

@@ -0,0 +1,191 @@
+module DimensionalDataInterpolationsExt
+
+using DimensionalData
+import DimensionalData as DD
+import DimensionalData: AbstractBasicDimArray
+# import Interpolations as Intp
+import Interpolations:AbstractInterpolation, Gridded, NoInterp, Linear, interpolate, extrapolate
+
+import UnrolledUtilities
+
+"""
+    DimGriddedInterpolation{T,N,A<:AbstractDimArray{T,N},I<:AbstractInterpolation,ICoords}
+
+A gridded interpolation object for `AbstractDimArray`s.
+
+# Fields
+- `array::A`: The array to interpolate.
+- `itp::I`: The interpolation object.
+- `itp_coords::ICoords`: The interpolation coordinates.
+
+Construct using `gridded_interpolate`.
+"""
+struct DimGriddedInterpolation{T,N,A<:AbstractDimArray{T,N},I<:AbstractInterpolation,ICoords}
+    array::A
+    itp::I
+    itp_coords::ICoords
+end
+
+"""
+    gridded_interpolate(A::AbstractDimArray, opts::NamedTuple; extrapolation_bc, [default_opt = NoInterp()])
+
+Interpolate `A` using the interpolation options `opts`.
+
+# Arguments
+- `A`: The array to interpolate.
+- `opts`: A named tuple of interpolation options for each dimension, see `Interpolations.jl` for available options.
+- `extrapolation_bc`: The extrapolation boundary conditions. One of:
+    `Throw()`, `Flat()`, `Line()`, `Free()`, `Reflect()`, `InPlace()`, `InPlaceQ()`, `Periodic()`
+- `default_opt`: The default interpolation option for each dimension.
+
+# Returns
+- `dgi::DimGriddedInterpolation`: The interpolated array.
+
+# Examples
+```julia-repl
+julia> using DimensionalData, Interpolations
+
+julia> A = DimArray(reshape(1.0:100.0, 10, 10), (X(1:10), Y(1:10)));
+
+julia> itp = DimensionalData.gridded_interpolate(A, (X = Gridded(Linear()), Y = Gridded(Linear())); extrapolation_bc = Line())
+10×10 DimGriddedInterpolation{Float64, 2} DimensionalData.NoName()
+  ↓ X Gridded(Linear()) [1, …, 10]
+  → Y Gridded(Linear()) [1, …, 10]
+Extrapolation: Line()
+
+julia> itp(X(5), Y(5))
+45.0
+
+julia> B = itp(X(5), Y([2, 4, 6]))
+┌ 3-element DimArray{Float64, 1} ┐
+├────────────────────────────────┴───────────────────────── dims ┐
+  ↓ Y Sampled{Int64} [2, …, 6] ForwardOrdered Irregular Points
+└────────────────────────────────────────────────────────────────┘
+ 2  15.0
+ 4  35.0
+ 6  55.0
+
+julia> refdims(B)  # Reference to "scalar" dimensions are retained
+(↓ X 5)
+```
+"""
+function DD.gridded_interpolate(A::AbstractDimArray, opts::NamedTuple; extrapolation_bc, default_opt = NoInterp())
+    # Set all dimensions to `default_opt` by default
+    default_opts = NamedTuple{DD.name(dims(A))}(ntuple(i -> default_opt, length(dims(A))))
+    opts = merge(default_opts, opts)
+    # Assert all opts are either Gridded or NoInterp
+    @assert all(Base.Fix2(isa, Union{Gridded, NoInterp}), opts) "All interpolation options must be either Gridded or NoInterp"
+
+    # `NoInterp` requires 1:length(dim) as nodes and coords.
+    nodes = map(dims(A)) do dim
+        opts[DD.name(dim)] isa Gridded ? lookup(dim) : Base.OneTo(length(dim))
+    end
+    itp_coords = map(dims(A)) do dim
+        opts[DD.name(dim)] isa Gridded ? dim : rebuild(dim, Base.OneTo(length(dim)))
+    end
+
+    itp = interpolate(nodes, A, values(opts))
+    extp = extrapolate(itp, extrapolation_bc)
+    DimGriddedInterpolation(A, extp, itp_coords)
+end
+
+function (dgi::DimGriddedInterpolation)(interp_dims...; da_kwargs...)
+    # Interpolate along the dimensions specified in `interp_dims`,
+    # use the original dimensions for the rest.
+    itp_dims = DD.setdims(dgi.itp_coords, interp_dims)  # insert dimensions in the correct order
+    result_arr = dgi.itp(val.(itp_dims)...)
+
+    # If `dgi` has any `NoInterp()`-dimensions, we need to fetch the values from the original dimensions.
+    nointerp_dims = DD.otherdims(dims(dgi.array), interp_dims)
+    dd_coords = DD.setdims(itp_dims, nointerp_dims)
+
+    # Filter out scalar dimensions (UnrolledUtilities needed for type-stability)
+    isvector(dim) = val(dim) isa AbstractVector
+    vector_coords, scalar_coords = UnrolledUtilities.unrolled_split(isvector, dd_coords)
+    isempty(vector_coords) && return result_arr  # Return scalar result if all dimensions were filtered
+
+    return DD.DimArray(result_arr, vector_coords; refdims = scalar_coords, da_kwargs...)
+end
+
+"""
+    interpolate!(dgi::DimGriddedInterpolation, A_to::AbstractDimArray)
+
+Interpolate using `dgi` and write the result into `A_to` in-place.
+
+The dimensions of `A_to` determine where to interpolate.
+
+# Arguments
+- `dgi`: The interpolation object.
+- `A_to`: The output array to write to. Its dimensions determine the interpolation coordinates.
+
+# Returns
+- `A_to`: Returns the modified array.
+
+# Notes:
+- (For now,) the order of the dimensions of `A_to` must match the underlying data of `dgi`.
+
+# Examples
+```julia-repl
+julia> using DimensionalData, Interpolations
+
+julia> A_from = DimArray(reshape(1.0:100.0, 10, 10), (X(1:10), Y(1:10)));
+
+julia> itp = DimensionalData.gridded_interpolate(A_from, (X = Gridded(Linear()), Y = Gridded(Linear())); extrapolation_bc = Line())
+10×10 DimGriddedInterpolation{Float64, 2} DimensionalData.NoName()
+  ↓ X Gridded(Linear()) [1, …, 10]
+  → Y Gridded(Linear()) [1, …, 10]
+Extrapolation: Line()
+
+julia> A_to = zeros((X(1:0.5:3), Y(1:0.1:2)));  # `DimArray` with desired dimensions
+
+julia> DimensionalData.interpolate!(itp, A_to)  # does in-place interpolation, zero allocations
+┌ 5×11 DimArray{Float64, 2} ┐
+├───────────────────────────┴──────────────────────────────── dims ┐
+  ↓ X Sampled{Float64} 1.0:0.5:3.0 ForwardOrdered Regular Points,
+  → Y Sampled{Float64} 1.0:0.1:2.0 ForwardOrdered Regular Points
+└──────────────────────────────────────────────────────────────────┘
+ ↓ →  1.0  1.1  1.2  1.3  1.4  1.5  1.6   1.7   1.8   1.9   2.0
+ 1.0  1.0  2.0  3.0  4.0  5.0  6.0  7.0   8.0   9.0  10.0  11.0
+ ⋮                        ⋮                           ⋮    
+ 3.0  3.0  4.0  5.0  6.0  7.0  8.0  9.0  10.0  11.0  12.0  13.0
+```
+"""
+function DD.interpolate!(dgi::DimGriddedInterpolation, A_to::AbstractDimArray)
+    # Interpolate along the dimensions specified in `A_to`,
+    # Check that `A_to` has the same dimensions as `dgi.itp_coords`, in the same order.
+    @assert DD.name(dims(A_to)) == DD.name(dgi.itp_coords) "`A_to` has different dimensions than `dgi.itp_coords`, or in a different order."
+    # Insert the interpolation dimensions in the correct order
+    itp_dims = DD.setdims(dgi.itp_coords, dims(A_to))
+    # Interpolate
+    itp_splat = splat(dgi.itp)
+    dd_pts = DimPoints(itp_dims)
+    @. A_to = itp_splat(dd_pts)
+    return A_to
+end
+
+function Base.show(io::IO, dgi::DimGriddedInterpolation)
+    A = dgi.array
+    dims_A = DD.dims(A)
+
+    # compact header: size and type
+    s_size = join(size(A), "×")
+    elty = eltype(A)
+    nd = ndims(A)
+    name = DD.name(A)
+    println(io, "$s_size DimGriddedInterpolation{$elty, $nd} $name")
+
+    # arrow glyphs to indicate axis directions (cycle if more dims)
+    arrows = ["↓", "→", "↗", "↖", "←", "↘"]
+
+    opts = dgi.itp.itp.it
+    for (i, d) in enumerate(dims_A)
+        arrow = arrows[mod1(i, length(arrows))]
+        nm = DD.name(d)
+        left, right = DD.extrema(d)
+        opt = opts[i]
+        println(io, "  $arrow $nm $opt [$left, …, $right]")
+    end
+    print(io, "Extrapolation: $(dgi.itp.et)")
+end
+
+end  # end module

src/DimensionalData.jl

@@ -111,6 +111,7 @@ include("tree/tree.jl")
 include("tree/show.jl")
 # Other
 include("tables.jl")
+include("interpolations.jl")
 # Combined (easier to work on these in one file)
 include("plotrecipes.jl")
 include("utils.jl")

src/interpolations.jl

@@ -0,0 +1,7 @@
+function gridded_interpolate(args...)
+    throw(ArgumentError("Load Interpolations.jl to use `gridded_interpolate`"))
+end
+
+function interpolate!(args...)
+    throw(ArgumentError("Load Interpolations.jl to use `interpolate!`"))
+end