Merge pull request #12 from kylebeggs/data-centers

add ability to evaluate at only subset of the input data

src/linalg/stencil.jl

@@ -13,9 +13,9 @@ function _build_weightmx(
     ℒrbf = ℒ(basis)
 
     # allocate arrays to build sparse matrix
-    I = zeros(Int, k * length(data))
+    I = zeros(Int, k * length(adjl))
     J = reduce(vcat, adjl)
-    V = zeros(TD, k * length(data))
+    V = zeros(TD, k * length(adjl))
 
     n_threads = Threads.nthreads()
 
@@ -27,15 +27,52 @@ function _build_weightmx(
     d = Vector{Vector{eltype(data)}}(undef, n_threads)
 
     # build stencil for each data point and store in global weight matrix
-    Threads.@threads for i in eachindex(data)
+    Threads.@threads for i in eachindex(adjl)
         range[Threads.threadid()] = ((i - 1) * k + 1):(i * k)
         @turbo I[range[Threads.threadid()]] .= i
         d[Threads.threadid()] = data[adjl[i]]
         V[range[Threads.threadid()]] = @views _build_stencil!(
             A, b, Threads.threadid(), ℒrbf, ℒmon, d, basis, mon, k
         )
     end
-    return sparse(I, J, V)
+
+    return sparse(I, J, V, length(adjl), length(data))
+end
+
+function _build_weight_vec(
+    ℒ, data::AbstractVector{D}, adjl::Vector{Vector{T}}, basis::B
+) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    TD = eltype(first(data))
+    dim = length(first(data)) # dimension of data
+    nmon = binomial(dim + basis.poly_deg, basis.poly_deg)
+    k = length(first(adjl))  # number of data in influence/support domain
+    sizes = (k, nmon)
+
+    # build monomial basis and operator
+    mon = MonomialBasis(dim, basis.poly_deg)
+    ℒmon = ℒ(mon)
+    ℒrbf = ℒ(basis)
+
+    # allocate arrays to build sparse matrix
+    V = [zeros(TD, k) for _ in 1:length(adjl)]
+
+    n_threads = Threads.nthreads()
+
+    # create work arrays
+    n = sum(sizes)
+    A = Symmetric[Symmetric(zeros(TD, n, n), :U) for _ in 1:n_threads]
+    b = Vector[zeros(TD, n) for _ in 1:n_threads]
+    d = Vector{Vector{eltype(data)}}(undef, n_threads)
+
+    # build stencil for each data point and store in global weight matrix
+    Threads.@threads for i in eachindex(adjl)
+        d[Threads.threadid()] = data[adjl[i]]
+        V[i] = @views _build_stencil!(
+            A, b, Threads.threadid(), ℒrbf, ℒmon, d, basis, mon, k
+        )
+    end
+
+    return V
 end
 
 function _build_stencil!(

src/operators/gradient.jl

@@ -9,28 +9,67 @@ end
 
 # convienience constructors
 function gradient(
-    data::AbstractVector{D}, basis::B=PHS(3; poly_deg=2); k::T=autoselect_k(data, basis)
+    data::AbstractVector{D},
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+    sparse=true,
+) where {D<:AbstractArray,B<:AbstractRadialBasis,T<:Int}
+    f = ntuple(length(first(data))) do dim
+        return let dim = dim
+            x -> ∂(x, 1, dim)
+        end
+    end
+    ℒ = Gradient(f)
+    return RadialBasisOperator(ℒ, data, basis; k=k, sparse=sparse)
+end
+
+function gradient(
+    data::AbstractVector{D},
+    centers::AbstractVector{D},
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+    sparse=true,
 ) where {D<:AbstractArray,B<:AbstractRadialBasis,T<:Int}
     f = ntuple(length(first(data))) do dim
         return let dim = dim
             x -> ∂(x, 1, dim)
         end
     end
     ℒ = Gradient(f)
-    return RadialBasisOperator(ℒ, data, basis; k=k)
+    return RadialBasisOperator(ℒ, data, centers, basis; k=k, sparse=sparse)
 end
 
 function RadialBasisOperator(
     ℒ::Gradient,
     data::AbstractVector{D},
     basis::B=PHS(3; poly_deg=2);
     k::T=autoselect_k(data, basis),
+    sparse=true,
 ) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    TD = eltype(D)
     adjl = find_neighbors(data, k)
-    N = length(data)
-    weights = ntuple(_ -> spzeros(N, N), length(ℒ.ℒ))
+    N = length(adjl)
+    weights = ntuple(_ -> _allocate_weights(TD, N, N, k; sparse=sparse), length(ℒ.ℒ))
     return RadialBasisOperator(ℒ, weights, data, adjl, basis)
 end
 
+function RadialBasisOperator(
+    ℒ::Gradient,
+    data::AbstractVector{D},
+    centers::AbstractVector{D},
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+    sparse=true,
+) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    TD = eltype(D)
+    adjl = find_neighbors(data, centers, k)
+    Na = length(adjl)
+    Nd = length(data)
+    weights = ntuple(_ -> _allocate_weights(TD, Na, Nd, k; sparse=sparse), length(ℒ.ℒ))
+    return RadialBasisOperator(ℒ, weights, data, adjl, basis)
+end
+
+Base.size(op::RadialBasisOperator{<:Gradient}) = size(first(op.weights))
+
 # pretty printing
 print_op(op::Gradient) = "Gradient (∇)"

src/operators/laplacian.jl

@@ -15,5 +15,15 @@ function laplacian(
     return RadialBasisOperator(ℒ, data, basis; k=k)
 end
 
+function laplacian(
+    data::AbstractVector{D},
+    centers::AbstractVector{D},
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    ℒ = Laplacian(∇²)
+    return RadialBasisOperator(ℒ, data, centers, basis; k=k)
+end
+
 # pretty printing
 print_op(op::Laplacian) = "Laplacian (∇² or Δ)"

src/operators/operators.jl

@@ -1,14 +1,13 @@
-abstract type AbstractRadialBasisOperator end
 abstract type AbstractOperator end
 abstract type ScalarValuedOperator <: AbstractOperator end
 abstract type VectorValuedOperator <: AbstractOperator end
 
 """
-    struct RadialBasisOperator <: AbstractRadialBasisOperator
+    struct RadialBasisOperator
 
 Operator of data using a radial basis with potential monomial augmentation.
 """
-struct RadialBasisOperator{L,W,D,A,B<:AbstractRadialBasis} <: AbstractRadialBasisOperator
+struct RadialBasisOperator{L,W,D,A,B<:AbstractRadialBasis}
     ℒ::L
     weights::W
     data::D
@@ -24,18 +23,95 @@ end
 
 # convienience constructors
 function RadialBasisOperator(
-    ℒ, data::AbstractVector{D}, basis::B=PHS(3; poly_deg=2); k::T=autoselect_k(data, basis)
+    ℒ,
+    data::AbstractVector{D},
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+    sparse=true,
 ) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    TD = eltype(D)
     adjl = find_neighbors(data, k)
-    N = length(data)
-    weights = spzeros(N, N)
+    Na = length(adjl)
+    Nd = length(data)
+    weights = _allocate_weights(TD, Na, Nd, k; sparse=sparse)
     return RadialBasisOperator(ℒ, weights, data, adjl, basis)
 end
 
+function RadialBasisOperator(
+    ℒ,
+    data::AbstractVector{D},
+    centers::AbstractVector{D},
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+    sparse=true,
+) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    TD = eltype(D)
+    adjl = find_neighbors(data, centers, k)
+    Na = length(adjl)
+    Nd = length(data)
+    weights = _allocate_weights(TD, Na, Nd, k; sparse=sparse)
+    return RadialBasisOperator(ℒ, weights, data, adjl, basis)
+end
+
+# extend Base methods
+Base.length(op::RadialBasisOperator) = length(op.adjl)
+Base.size(op::RadialBasisOperator) = size(op.weights)
+function Base.size(op::RadialBasisOperator{<:VectorValuedOperator})
+    return ntuple(i -> size(op.weights[i]), embeddim(op))
+end
+Base.getindex(op::RadialBasisOperator, i) = nonzeros(op.weights[i, :])
+function Base.getindex(op::RadialBasisOperator{VectorValuedOperator}, i)
+    return ntuple(j -> nonzeros(op.weights[j][i, :]), embeddim(op))
+end
+
+# convienience methods
+embeddim(op::RadialBasisOperator) = length(first(op.data))
+
+# caching
+invalidate_cache(op::RadialBasisOperator) = op.valid_cache[] = false
+validate_cache(op::RadialBasisOperator) = op.valid_cache[] = true
+is_cache_valid(op::RadialBasisOperator) = op.valid_cache[]
+
+# for operator types such as Partial, Gradient, Laplacian, etc...
+(op::AbstractOperator)(x) = op.ℒ(x)
+
+# dispatches for evaluation
+_eval_op(op::RadialBasisOperator, x::AbstractVector) = _eval_op(op.weights, x)
+
+function _eval_op(op::RadialBasisOperator{<:VectorValuedOperator}, x::AbstractVector)
+    return ntuple(i -> _eval_op(op.weights[i], x), embeddim(op))
+end
+
+function _eval_op(
+    op::RadialBasisOperator{<:VectorValuedOperator,W}, x::AbstractVector
+) where {W<:Tuple}
+    if first(op.weights) isa Vector{<:Vector}
+        return ntuple(i -> _eval_op(op.weights[i], x, op.adjl), embeddim(op))
+    else
+        return ntuple(i -> _eval_op(op.weights[i], x), embeddim(op))
+    end
+end
+
+function _eval_op(
+    op::RadialBasisOperator{L,W}, x::AbstractVector
+) where {L,W<:Vector{<:Vector}}
+    return _eval_op(op.weights, x, op.adjl)
+end
+
+_eval_op(w::AbstractMatrix, x::AbstractVector) = w * x
+
+function _eval_op(w::AbstractVector{<:AbstractVector{T}}, x::AbstractVector, adjl) where {T}
+    y = zeros(T, length(w))
+    Threads.@threads for i in eachindex(adjl)
+        @views y[i] = w[i] ⋅ x[adjl[i]]
+    end
+    return y
+end
+
 # evaluate
 function (op::RadialBasisOperator)(x)
     !is_cache_valid(op) && update_weights!(op)
-    return op.weights * x
+    return _eval_op(op, x)
 end
 
 function LinearAlgebra.mul!(
@@ -51,11 +127,6 @@ function LinearAlgebra.mul!(
     return mul!(y, op.weights, x, α, β)
 end
 
-# evaluate
-function (op::RadialBasisOperator{<:VectorValuedOperator})(x::AbstractVector)
-    !is_cache_valid(op) && update_weights!(op)
-    return map(w -> w * x, op.weights)
-end
 function LinearAlgebra.:⋅(
     op::RadialBasisOperator{<:VectorValuedOperator}, x::AbstractVector
 )
@@ -76,26 +147,50 @@ function LinearAlgebra.mul!(
 end
 
 # update weights
+function _build_weights(ℒ, op::RadialBasisOperator)
+    if op.weights isa Vector{<:Vector}
+        return _build_weight_vec(ℒ, op.data, op.adjl, op.basis)
+    else
+        return _build_weightmx(ℒ, op.data, op.adjl, op.basis)
+    end
+end
+
+function _build_weights(ℒ, op::RadialBasisOperator{<:VectorValuedOperator})
+    if first(op.weights) isa Vector{<:Vector}
+        return _build_weight_vec(ℒ, op.data, op.adjl, op.basis)
+    else
+        return _build_weightmx(ℒ, op.data, op.adjl, op.basis)
+    end
+end
+
 function update_weights!(op::RadialBasisOperator)
-    op.weights .= _build_weightmx(op.ℒ, op.data, op.adjl, op.basis)
+    op.weights .= _build_weights(op.ℒ, op)
     validate_cache(op)
     return nothing
 end
 
 function update_weights!(op::RadialBasisOperator{<:VectorValuedOperator})
+    return _update_weights!(op, op.weights)
+end
+
+function _update_weights!(op, weights::NTuple{N,AbstractMatrix}) where {N}
     for (i, ℒ) in enumerate(op.ℒ.ℒ)
-        op.weights[i] .= _build_weightmx(ℒ, op.data, op.adjl, op.basis)
+        weights[i] .= _build_weights(ℒ, op)
     end
     validate_cache(op)
     return nothing
 end
 
-Base.getindex(op::O, i) where {O<:AbstractRadialBasisOperator} = nonzeros(op.weights[i, :])
-invalidate_cache(op::RadialBasisOperator) = op.valid_cache[] = false
-validate_cache(op::RadialBasisOperator) = op.valid_cache[] = true
-is_cache_valid(op::RadialBasisOperator) = op.valid_cache[]
-
-(op::AbstractOperator)(x) = op.ℒ(x)
+function _update_weights!(op, weights::NTuple{N,AbstractVector}) where {N}
+    for (i, ℒ) in enumerate(op.ℒ.ℒ)
+        w = _build_weights(ℒ, op)
+        for j in eachindex(weights[i])
+            weights[i][j] .= w[j]
+        end
+    end
+    validate_cache(op)
+    return nothing
+end
 
 # pretty printing
 function Base.show(io::IO, op::RadialBasisOperator)

src/operators/partial.jl

@@ -24,5 +24,20 @@ function partial(
     return RadialBasisOperator(ℒ, data, basis; k=k)
 end
 
+function partial(
+    data::AbstractVector{D},
+    centers::AbstractVector{D},
+    order::T,
+    dim::T,
+    basis::B=PHS(3; poly_deg=2);
+    k::T=autoselect_k(data, basis),
+) where {D<:AbstractArray,T<:Int,B<:AbstractRadialBasis}
+    f = let o = order, dim = dim
+        x -> ∂(x, o, dim)
+    end
+    ℒ = Partial(f, order, dim)
+    return RadialBasisOperator(ℒ, data, centers, basis; k=k)
+end
+
 # pretty printing
 print_op(op::Partial) = "∂ⁿ/∂xᵢ (n = $(op.order), i = $(op.dim))"

src/utils.jl

@@ -1,9 +1,15 @@
-function find_neighbors(data::Vector, k::Int)
+function find_neighbors(data::AbstractVector, k::Int)
     tree = KDTree(data)
     adjl, _ = knn(tree, data, k, true)
     return adjl
 end
 
+function find_neighbors(data::AbstractVector, centers::AbstractVector, k::Int)
+    tree = KDTree(data)
+    adjl, _ = knn(tree, centers, k, true)
+    return adjl
+end
+
 function get_num_params(f)
     return length.(getfield.(getfield.(methods(f), :sig), :parameters)) .- 1
 end
@@ -51,3 +57,8 @@ function check_poly_deg(poly_deg)
     end
     return nothing
 end
+
+_allocate_weights(m, n, k; sparse=true) = _allocate_weights(Float64, m, n, k; sparse=sparse)
+function _allocate_weights(T, m, n, k; sparse=true)
+    return sparse ? spzeros(T, m, n) : [zeros(T, k) for _ in 1:m]
+end