Add jacobi polynomial bases

NEWS.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Jacobi polynomial bases. Since PR [#896](https://github.com/gridap/Gridap.jl/pull/896).
+
 ### Fixed
 
 - Fixed the method `get_normal_vector` for `AdaptedTriangulation`. The method `get_facet_normal`

src/Polynomials/JacobiPolynomialBases.jl

@@ -0,0 +1,359 @@
+struct JacobiPolynomial <: Field end
+
+struct JacobiPolynomialBasis{D,T} <: AbstractVector{JacobiPolynomial}
+  orders::NTuple{D,Int}
+  terms::Vector{CartesianIndex{D}}
+  function JacobiPolynomialBasis{D}(
+    ::Type{T}, orders::NTuple{D,Int}, terms::Vector{CartesianIndex{D}}) where {D,T}
+    new{D,T}(orders,terms)
+  end
+end
+
+@inline Base.size(a::JacobiPolynomialBasis{D,T}) where {D,T} = (length(a.terms)*num_components(T),)
+@inline Base.getindex(a::JacobiPolynomialBasis,i::Integer) = JacobiPolynomial()
+@inline Base.IndexStyle(::JacobiPolynomialBasis) = IndexLinear()
+
+function JacobiPolynomialBasis{D}(
+  ::Type{T}, orders::NTuple{D,Int}, filter::Function=_q_filter) where {D,T}
+
+  terms = _define_terms(filter, orders)
+  JacobiPolynomialBasis{D}(T,orders,terms)
+end
+
+function JacobiPolynomialBasis{D}(
+  ::Type{T}, order::Int, filter::Function=_q_filter) where {D,T}
+
+  orders = tfill(order,Val{D}())
+  JacobiPolynomialBasis{D}(T,orders,filter)
+end
+
+# API
+
+function get_exponents(b::JacobiPolynomialBasis)
+  indexbase = 1
+  [Tuple(t) .- indexbase for t in b.terms]
+end
+
+function get_order(b::JacobiPolynomialBasis)
+  maximum(b.orders)
+end
+
+function get_orders(b::JacobiPolynomialBasis)
+  b.orders
+end
+
+return_type(::JacobiPolynomialBasis{D,T}) where {D,T} = T
+
+# Field implementation
+
+function return_cache(f::JacobiPolynomialBasis{D,T},x::AbstractVector{<:Point}) where {D,T}
+  @assert D == length(eltype(x)) "Incorrect number of point components"
+  np = length(x)
+  ndof = length(f.terms)*num_components(T)
+  n = 1 + _maximum(f.orders)
+  r = CachedArray(zeros(T,(np,ndof)))
+  v = CachedArray(zeros(T,(ndof,)))
+  c = CachedArray(zeros(eltype(T),(D,n)))
+  (r, v, c)
+end
+
+function evaluate!(cache,f::JacobiPolynomialBasis{D,T},x::AbstractVector{<:Point}) where {D,T}
+  r, v, c = cache
+  np = length(x)
+  ndof = length(f.terms)*num_components(T)
+  n = 1 + _maximum(f.orders)
+  setsize!(r,(np,ndof))
+  setsize!(v,(ndof,))
+  setsize!(c,(D,n))
+  for i in 1:np
+    @inbounds xi = x[i]
+    _evaluate_nd_jp!(v,xi,f.orders,f.terms,c)
+    for j in 1:ndof
+      @inbounds r[i,j] = v[j]
+    end
+  end
+  r.array
+end
+
+function return_cache(
+  fg::FieldGradientArray{1,JacobiPolynomialBasis{D,V}},
+  x::AbstractVector{<:Point}) where {D,V}
+
+  f = fg.fa
+  @assert D == length(eltype(x)) "Incorrect number of point components"
+  np = length(x)
+  ndof = length(f.terms)*num_components(V)
+  xi = testitem(x)
+  T = gradient_type(V,xi)
+  n = 1 + _maximum(f.orders)
+  r = CachedArray(zeros(T,(np,ndof)))
+  v = CachedArray(zeros(T,(ndof,)))
+  c = CachedArray(zeros(eltype(T),(D,n)))
+  g = CachedArray(zeros(eltype(T),(D,n)))
+  (r, v, c, g)
+end
+
+function evaluate!(
+  cache,
+  fg::FieldGradientArray{1,JacobiPolynomialBasis{D,T}},
+  x::AbstractVector{<:Point}) where {D,T}
+
+  f = fg.fa
+  r, v, c, g = cache
+  np = length(x)
+  ndof = length(f.terms) * num_components(T)
+  n = 1 + _maximum(f.orders)
+  setsize!(r,(np,ndof))
+  setsize!(v,(ndof,))
+  setsize!(c,(D,n))
+  setsize!(g,(D,n))
+  for i in 1:np
+    @inbounds xi = x[i]
+    _gradient_nd_jp!(v,xi,f.orders,f.terms,c,g,T)
+    for j in 1:ndof
+      @inbounds r[i,j] = v[j]
+    end
+  end
+  r.array
+end
+
+function return_cache(
+  fg::FieldGradientArray{2,JacobiPolynomialBasis{D,V}},
+  x::AbstractVector{<:Point}) where {D,V}
+
+  f = fg.fa
+  @assert D == length(eltype(x)) "Incorrect number of point components"
+  np = length(x)
+  ndof = length(f.terms)*num_components(V)
+  xi = testitem(x)
+  T = gradient_type(gradient_type(V,xi),xi)
+  n = 1 + _maximum(f.orders)
+  r = CachedArray(zeros(T,(np,ndof)))
+  v = CachedArray(zeros(T,(ndof,)))
+  c = CachedArray(zeros(eltype(T),(D,n)))
+  g = CachedArray(zeros(eltype(T),(D,n)))
+  h = CachedArray(zeros(eltype(T),(D,n)))
+  (r, v, c, g, h)
+end
+
+function evaluate!(
+  cache,
+  fg::FieldGradientArray{2,JacobiPolynomialBasis{D,T}},
+  x::AbstractVector{<:Point}) where {D,T}
+
+  f = fg.fa
+  r, v, c, g, h = cache
+  np = length(x)
+  ndof = length(f.terms) * num_components(T)
+  n = 1 + _maximum(f.orders)
+  setsize!(r,(np,ndof))
+  setsize!(v,(ndof,))
+  setsize!(c,(D,n))
+  setsize!(g,(D,n))
+  setsize!(h,(D,n))
+  for i in 1:np
+    @inbounds xi = x[i]
+    _hessian_nd_jp!(v,xi,f.orders,f.terms,c,g,h,T)
+    for j in 1:ndof
+      @inbounds r[i,j] = v[j]
+    end
+  end
+  r.array
+end
+
+# Optimizing evaluation at a single point
+
+function return_cache(f::JacobiPolynomialBasis{D,T},x::Point) where {D,T}
+  ndof = length(f.terms)*num_components(T)
+  r = CachedArray(zeros(T,(ndof,)))
+  xs = [x]
+  cf = return_cache(f,xs)
+  r, cf, xs
+end
+
+function evaluate!(cache,f::JacobiPolynomialBasis{D,T},x::Point) where {D,T}
+  r, cf, xs = cache
+  xs[1] = x
+  v = evaluate!(cf,f,xs)
+  ndof = size(v,2)
+  setsize!(r,(ndof,))
+  a = r.array
+  copyto!(a,v)
+  a
+end
+
+function return_cache(
+  f::FieldGradientArray{N,JacobiPolynomialBasis{D,V}}, x::Point) where {N,D,V}
+  xs = [x]
+  cf = return_cache(f,xs)
+  v = evaluate!(cf,f,xs)
+  r = CachedArray(zeros(eltype(v),(size(v,2),)))
+  r, cf, xs
+end
+
+function evaluate!(
+  cache, f::FieldGradientArray{N,JacobiPolynomialBasis{D,V}}, x::Point) where {N,D,V}
+  r, cf, xs = cache
+  xs[1] = x
+  v = evaluate!(cf,f,xs)
+  ndof = size(v,2)
+  setsize!(r,(ndof,))
+  a = r.array
+  copyto!(a,v)
+  a
+end
+
+# Helpers
+
+function _evaluate_1d_jp!(v::AbstractMatrix{T},x,order,d) where T
+  n = order + 1
+  z = one(T)
+  @inbounds v[d,1] = z
+  if n > 1
+    ξ = ( 2*x[d] - 1 )
+    for i in 2:n
+      @inbounds v[d,i] = sqrt(2*i-1)*jacobi(ξ,i-1,0,0)
+    end
+  end
+end
+
+function _gradient_1d_jp!(v::AbstractMatrix{T},x,order,d) where T
+  n = order + 1
+  z = zero(T)
+  @inbounds v[d,1] = z
+  if n > 1
+    ξ = ( 2*x[d] - 1 )
+    for i in 2:n
+      @inbounds v[d,i] = sqrt(2*i-1)*i*jacobi(ξ,i-2,1,1)
+    end
+  end
+end
+
+function _hessian_1d_jp!(v::AbstractMatrix{T},x,order,d) where T
+  n = order + 1
+  z = zero(T)
+  @inbounds v[d,1] = z
+  if n > 1
+    @inbounds v[d,2] = z
+    ξ = ( 2*x[d] - 1 )
+    for i in 3:n
+      @inbounds v[d,i] = sqrt(2*i-1)*(i*(i+1)/2)*jacobi(ξ,i-3,2,2)
+    end
+  end
+end
+
+function _evaluate_nd_jp!(
+  v::AbstractVector{V},
+  x,
+  orders,
+  terms::AbstractVector{CartesianIndex{D}},
+  c::AbstractMatrix{T}) where {V,T,D}
+
+  dim = D
+  for d in 1:dim
+    _evaluate_1d_jp!(c,x,orders[d],d)
+  end
+
+  o = one(T)
+  k = 1
+
+  for ci in terms
+
+    s = o
+    for d in 1:dim
+      @inbounds s *= c[d,ci[d]]
+    end
+
+    k = _set_value!(v,s,k)
+
+  end
+
+end
+
+function _gradient_nd_jp!(
+  v::AbstractVector{G},
+  x,
+  orders,
+  terms::AbstractVector{CartesianIndex{D}},
+  c::AbstractMatrix{T},
+  g::AbstractMatrix{T},
+  ::Type{V}) where {G,T,D,V}
+
+  dim = D
+  for d in 1:dim
+    _evaluate_1d_jp!(c,x,orders[d],d)
+    _gradient_1d_jp!(g,x,orders[d],d)
+  end
+
+  z = zero(Mutable(VectorValue{D,T}))
+  o = one(T)
+  k = 1
+
+  for ci in terms
+
+    s = z
+    for i in eachindex(s)
+      @inbounds s[i] = o
+    end
+    for q in 1:dim
+      for d in 1:dim
+        if d != q
+          @inbounds s[q] *= c[d,ci[d]]
+        else
+          @inbounds s[q] *= g[d,ci[d]]
+        end
+      end
+    end
+
+    k = _set_gradient!(v,s,k,V)
+
+  end
+
+end
+
+function _hessian_nd_jp!(
+  v::AbstractVector{G},
+  x,
+  orders,
+  terms::AbstractVector{CartesianIndex{D}},
+  c::AbstractMatrix{T},
+  g::AbstractMatrix{T},
+  h::AbstractMatrix{T},
+  ::Type{V}) where {G,T,D,V}
+
+  dim = D
+  for d in 1:dim
+    _evaluate_1d_jp!(c,x,orders[d],d)
+    _gradient_1d_jp!(g,x,orders[d],d)
+    _hessian_1d_jp!(h,x,orders[d],d)
+  end
+
+  z = zero(Mutable(TensorValue{D,D,T}))
+  o = one(T)
+  k = 1
+
+  for ci in terms
+
+    s = z
+    for i in eachindex(s)
+      @inbounds s[i] = o
+    end
+    for r in 1:dim
+      for q in 1:dim
+        for d in 1:dim
+          if d != q && d != r
+            @inbounds s[r,q] *= c[d,ci[d]]
+          elseif d == q && d ==r
+            @inbounds s[r,q] *= h[d,ci[d]]
+          else
+            @inbounds s[r,q] *= g[d,ci[d]]
+          end
+        end
+      end
+    end
+
+    k = _set_gradient!(v,s,k,V)
+
+  end
+
+end

src/Polynomials/Polynomials.jl

@@ -25,6 +25,7 @@ export QGradMonomialBasis
 export QCurlGradMonomialBasis
 export PCurlGradMonomialBasis
 export ModalC0Basis
+export JacobiPolynomialBasis
 export get_exponents
 
 export get_order
@@ -41,4 +42,6 @@ include("PCurlGradMonomialBases.jl")
 
 include("ModalC0Bases.jl")
 
+include("JacobiPolynomialBases.jl")
+
 end # module