From 1ead937ef49beb744876df405ce537734ede9b70 Mon Sep 17 00:00:00 2001 From: Haakon Ludvig Langeland Ervik <45243236+haakon-e@users.noreply.github.com> Date: Sun, 2 Nov 2025 17:40:51 -0800 Subject: [PATCH] ft: add gridded interpolation to DimensionalData Add functions `gridded_interpolate` and `interpolate!` in a project extension to interpolate `DimArray`s using the `Interpolations.jl` package. - `gridded_interpolate` creates a `DimGriddedInterpolation` object, which can be used to interpolate the data in-place or out-of-place. - out-of-place interpolation is supported by calling the `DimGriddedInterpolation` object with the desired dimensions. - in-place interpolation is supported by calling `interpolate!` with the `DimGriddedInterpolation` object and an existing `DimArray`. --- Project.toml | 7 + ext/DimensionalDataInterpolationsExt.jl | 191 ++++++++++++++++++++++++ src/DimensionalData.jl | 1 + src/interpolations.jl | 7 + 4 files changed, 206 insertions(+) create mode 100644 ext/DimensionalDataInterpolationsExt.jl create mode 100644 src/interpolations.jl diff --git a/Project.toml b/Project.toml index 81e9b465a..ca63c1915 100644 --- a/Project.toml +++ b/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,6 +36,8 @@ 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" @@ -41,6 +45,7 @@ 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] diff --git a/ext/DimensionalDataInterpolationsExt.jl b/ext/DimensionalDataInterpolationsExt.jl new file mode 100644 index 000000000..ce5bfaf39 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/DimensionalData.jl b/src/DimensionalData.jl index 1f188c6cf..690efd2e4 100644 --- a/src/DimensionalData.jl +++ b/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") diff --git a/src/interpolations.jl b/src/interpolations.jl new file mode 100644 index 000000000..1bb0d85e2 --- /dev/null +++ b/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 \ No newline at end of file