diff --git a/Project.toml b/Project.toml index 56acb76..4c8729b 100644 --- a/Project.toml +++ b/Project.toml @@ -4,17 +4,26 @@ version = "0.1.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +Atomix = "a9b6321e-bd34-4604-b9c9-b65b8de01458" EllipsisNotation = "da5c29d0-fa7d-589e-88eb-ea29b0a81949" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01" +[weakdeps] +LinearMaps = "7a12625a-238d-50fd-b39a-03d52299707e" + +[extensions] +NDInterpolationsLinearMapsExt = "LinearMaps" + [compat] Adapt = "4.3.0" Aqua = "0.8" +Atomix = "1.1.1" DataInterpolations = "8" EllipsisNotation = "1.8.0" ForwardDiff = "0" KernelAbstractions = "0.9.34" +LinearMaps = "3.11.4" Random = "1" RecipesBase = "1.3.4" SafeTestsets = "0.1" diff --git a/ext/NDInterpolationsLinearMapsExt.jl b/ext/NDInterpolationsLinearMapsExt.jl new file mode 100644 index 0000000..ac4ea95 --- /dev/null +++ b/ext/NDInterpolationsLinearMapsExt.jl @@ -0,0 +1,69 @@ +module NDInterpolationsLinearMapsExt +using NDInterpolations +using NDInterpolations: validate_derivative_orders, get_output_size +using LinearMaps + +# A linear map interp.u -> grid evaluation +function LinearMaps.LinearMap( + interp::NDInterpolation{N_in}, ::Val{:grid}; + derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in) +) where {N_in} + validate_derivative_orders(derivative_orders, interp) + + T = NDInterpolations.output_type(interp) + + grid_size = map(itp_dim -> length(itp_dim.t_eval), interp.interp_dims) + size_out = (grid_size..., get_output_size(interp)...) + N_input = length(interp.u) + N_output = prod(size_out) + + function map!(out_flat, u_flat) + u_reshaped = reshape(u_flat, size(interp.u)) + out_reshaped = reshape(out_flat, size_out) + + interp_ = NDInterpolation(u_reshaped, interp.interp_dims, interp.cache) + eval_grid!(out_reshaped, interp_; derivative_orders) + return out_flat + end + + function map_adjoint!(u_flat, out_flat) + u_reshaped = reshape(u_flat, size(interp.u)) + out_reshaped = reshape(out_flat, size_out) + + interp_ = NDInterpolation(u_reshaped, interp.interp_dims, interp.cache) + eval_grid!(out_reshaped, interp_; derivative_orders, adjoint = true) + return u_flat + end + + return FunctionMap{T}(map!, map_adjoint!, N_output, N_input) +end + +# A linear map interp.u -> unstructured evaluation +function LinearMaps.LinearMap( + interp::NDInterpolation{N_in}, ::Val{:unstructured}; + derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in) +) where {N_in} + validate_derivative_orders(derivative_orders, interp) + + T = NDInterpolations.output_type(interp) + + N_input = length(interp.u) + size_out = (length(first(interp.interp_dims).t_eval), get_output_size(interp)...) + N_output = prod(size_out) + + function map!(out_flat, u_flat) + u_reshaped = reshape(u_flat, size(interp.u)) + out_reshaped = reshape(out_flat, size_out) + + interp_ = NDInterpolation(u_reshaped, interp.interp_dims, interp.cache) + eval_unstructured!(out_reshaped, interp_; derivative_orders) + return out_flat + end + + # function map_adjoint!() + # end + + return FunctionMap{T}(map!, N_output, N_input) +end + +end # module NDInterpolationsLinearMapsExt \ No newline at end of file diff --git a/src/NDInterpolations.jl b/src/NDInterpolations.jl index a2eb17b..dbc36eb 100644 --- a/src/NDInterpolations.jl +++ b/src/NDInterpolations.jl @@ -1,5 +1,6 @@ module NDInterpolations -using KernelAbstractions # Keep as dependency or make extension? +import Atomix +using KernelAbstractions using Adapt: @adapt_structure using EllipsisNotation using RecipesBase @@ -56,6 +57,7 @@ include("interpolation_dimensions.jl") include("spline_utils.jl") include("interpolation_utils.jl") include("interpolation_methods.jl") +include("interpolation_methods_adjoint.jl") include("interpolation_parallel.jl") include("plot_rec.jl") diff --git a/src/interpolation_methods_adjoint.jl b/src/interpolation_methods_adjoint.jl new file mode 100644 index 0000000..f2e69db --- /dev/null +++ b/src/interpolation_methods_adjoint.jl @@ -0,0 +1,39 @@ +function _interpolate_adjoint!( + A::NDInterpolation{N_in, N_out, ID}, + out, + t::Tuple{Vararg{Number, N_in}}, + idx::NTuple{N_in, <:Integer}, + derivative_orders::NTuple{N_in, <:Integer}, + multi_point_index +)::Nothing where {N_in, N_out, ID <: LinearInterpolationDimension} + any(>(1), derivative_orders) && return out + + tᵢ = ntuple(i -> A.interp_dims[i].t[idx[i]], N_in) + tᵢ₊₁ = ntuple(i -> A.interp_dims[i].t[idx[i] + 1], N_in) + + # Size of the (hyper)rectangle `t` is in + t_vol = one(eltype(tᵢ)) + for (t₁, t₂) in zip(tᵢ, tᵢ₊₁) + t_vol *= t₂ - t₁ + end + + # Loop over the corners of the (hyper)rectangle `t` is in + for I in Iterators.product(ntuple(i -> (false, true), N_in)...) + c = eltype(out)(inv(t_vol)) + for (t_, right_point, d, t₁, t₂) in zip(t, I, derivative_orders, tᵢ, tᵢ₊₁) + c *= if right_point + iszero(d) ? t_ - t₁ : one(t_) + else + iszero(d) ? t₂ - t_ : -one(t_) + end + end + J = (ntuple(i -> idx[i] + I[i], N_in)..., ..) + + if iszero(N_out) + Atomix.@atomic A.u[J...] += c * out + else + Atomix.@atomic A.u[J...] += c * out + end + end + return nothing +end \ No newline at end of file diff --git a/src/interpolation_parallel.jl b/src/interpolation_parallel.jl index cce0a5d..2b4ad8f 100644 --- a/src/interpolation_parallel.jl +++ b/src/interpolation_parallel.jl @@ -85,34 +85,37 @@ in place. function eval_grid!( out::AbstractArray, interp::NDInterpolation{N_in}; - derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in) + derivative_orders::NTuple{N_in, <:Integer} = ntuple(_ -> 0, N_in), + adjoint::Bool = false ) where {N_in} validate_derivative_orders(derivative_orders, interp; multi_point = true) backend = get_backend(out) @assert all(i -> size(out, i) == length(interp.interp_dims[i].t_eval), N_in) "For the first N_in dimensions of out the length must match the t_eval of the corresponding interpolation dimension." @assert size(out)[(N_in + 1):end]==get_output_size(interp) "The size of the last N_out dimensions of out must be the same as the output size of the interpolation." - eval_kernel(backend)( - out, - interp, - derivative_orders, - true, - ndrange = size(out)[1:N_in] - ) + if adjoint + eval_kernel_adjoint(backend)( + interp, + out, + derivative_orders, + true, + ndrange = size(out)[1:N_in] + ) + else + eval_kernel(backend)( + out, + interp, + derivative_orders, + true, + ndrange = size(out)[1:N_in] + ) + end synchronize(backend) return out end -@kernel function eval_kernel( - out, - @Const(A), - derivative_orders, - eval_grid -) - N_in = length(A.interp_dims) - N_out = ndims(A.u) - N_in - - k = @index(Global, NTuple) - +function get_eval_params( + A::NDInterpolation{N_in, N_out}, eval_grid::Bool, k::NTuple{N_in, Int} +) where {N_in, N_out} if eval_grid t_eval = ntuple(i -> A.interp_dims[i].t_eval[k[i]], N_in) idx_eval = ntuple(i -> A.interp_dims[i].idx_eval[k[i]], N_in) @@ -121,6 +124,18 @@ end idx_eval = ntuple(i -> A.interp_dims[i].idx_eval[only(k)], N_in) end + return N_out, t_eval, idx_eval +end + +@kernel function eval_kernel( + out, + @Const(A), + derivative_orders, + eval_grid +) + k = @index(Global, NTuple) + N_out, t_eval, idx_eval = get_eval_params(A, eval_grid, k) + if iszero(N_out) out[k...] = _interpolate!( make_out(A, t_eval), A, t_eval, idx_eval, derivative_orders, k) @@ -130,3 +145,21 @@ end A, t_eval, idx_eval, derivative_orders, k) end end + +@kernel function eval_kernel_adjoint( + A, + @Const(out), + derivative_orders, + eval_grid +) + k = @index(Global, NTuple) + N_out, t_eval, idx_eval = get_eval_params(A, eval_grid, k) + + if iszero(N_out) + _interpolate_adjoint!( + A, out[k...], t_eval, idx_eval, derivative_orders, k) + else + _interpolate_adjoint!( + A, view(out, k..., ..), t_eval, idx_eval, derivative_orders, k) + end +end \ No newline at end of file diff --git a/src/interpolation_utils.jl b/src/interpolation_utils.jl index 45db3aa..450bed6 100644 --- a/src/interpolation_utils.jl +++ b/src/interpolation_utils.jl @@ -87,19 +87,29 @@ function make_zero!!(v::T) where {T <: AbstractArray} v end +function output_type(interp::NDInterpolation) + promote_type(eltype(interp.u), map(itp_dim -> eltype(itp_dim.t), interp.interp_dims)...) +end + +function output_type( + interp::NDInterpolation{N_in}, + t::NTuple{N_in, >:Number} +) where {N_in} + promote_type(output_type(interp), map(typeof, t)...) +end + function make_out( interp::NDInterpolation{N_in, 0}, t::NTuple{N_in, >:Number} ) where {N_in} - zero(promote_type(eltype(interp.u), map(typeof, t)...)) + zero(output_type(interp, t)) end function make_out( interp::NDInterpolation{N_in}, t::NTuple{N_in, >:Number} ) where {N_in} - similar( - interp.u, promote_type(eltype(interp.u), map(eltype, t)...), get_output_size(interp)) + similar(interp.u, ouput_type(interp, t)) end get_left(::AbstractInterpolationDimension) = false