Merge pull request #19 from bcube-project/dev_jacobian_physicalfunc

Gradient of a vector-valued PhysicalFunction

src/algebra/gradient.jl

@@ -103,12 +103,19 @@ function Gradient(
     MapOver(grad_shape_functionsNA(fs, n, ctype, cnodes, ξ))
 end
 
-function Gradient(cellFunction::AbstractCellFunction{<:PhysicalDomain}, cPoint::CellPoint)
+function Gradient(
+    cellFunction::AbstractCellFunction{<:PhysicalDomain, S},
+    cPoint::CellPoint,
+) where {S}
     f = get_function(cellFunction)
     x = change_domain(cPoint, PhysicalDomain()) # transparent if already in PhysicalDomain
-    return ForwardDiff.gradient(f, get_coord(x))
+    return _gradient_or_jacobian(Val(S), f, get_coord(x))
 end
 
+# dispatch on codomain size :
+_gradient_or_jacobian(::Val{1}, f, x) = ForwardDiff.gradient(f, x)
+_gradient_or_jacobian(::Val{S}, f, x) where {S} = ForwardDiff.jacobian(f, x)
+
 function grad_shape_functionsNA(fs::AbstractFunctionSpace, n::Val{1}, ctype, cnodes, ξ)
     grad = grad_shape_functions(fs, n, ctype, cnodes, ξ)
     _reshape_gradient_shape_function_impl(grad, n)

src/cellfunction/cellfunction.jl

@@ -1,19 +1,22 @@
 """
-    AbstractCellFunction{DS}
+    AbstractCellFunction{DS,S}
 
 Abstract type to represent a function defined in specific domain `DS`,
-which could be a `ReferenceDomain` or a `PhysicalDomain`.
+which could be a `ReferenceDomain` or a `PhysicalDomain`. `S` is the size
+of the codomain (i.e `S=length(f(x))`) where `f` is the `AbstractCellFunction`.
 
 # Subtypes should implement :
 * `get_function(f::AbstractCellFunction)`
 """
-abstract type AbstractCellFunction{DS} <: AbstractLazy end
+abstract type AbstractCellFunction{DS, S} <: AbstractLazy end
 
 """
     DomainStyle(f::AbstractCellFunction)
 """
 DomainStyle(f::AbstractCellFunction{DS}) where {DS} = DS()
 
+get_size(::AbstractCellFunction{DS, S}) where {DS, S} = S
+
 """
     get_function(f::AbstractCellFunction)
 """
@@ -70,7 +73,7 @@ LazyOperators.materialize(f::AbstractCellFunction, x::CellInfo) = f
 LazyOperators.materialize(f::AbstractCellFunction, x::CellPoint) = f(x)
 
 """
-    abstract type AbstractShapeFunction{DS,S,FS} <: AbstractCellFunction{DS} end
+    abstract type AbstractShapeFunction{DS,S,FS} <: AbstractCellFunction{DS,S} end
 
 Abstract type to represent a shape function defined in specific domain `DS`,
 which could be a `ReferenceDomain` or a `PhysicalDomain`. `S` is the size
@@ -82,9 +85,8 @@ Subtypes should implement `AbstractCellFunction`:
 
 and its own specitic interface: [empty]
 """
-abstract type AbstractShapeFunction{DS, S, FS} <: AbstractCellFunction{DS} end
+abstract type AbstractShapeFunction{DS, S, FS} <: AbstractCellFunction{DS, S} end
 
-get_size(::AbstractShapeFunction{DS, S}) where {DS, S} = S
 get_function_space(::AbstractShapeFunction{DS, S, FS}) where {DS, S, FS} = FS()
 
 """
@@ -199,37 +201,69 @@ end
 LazyOperators.pretty_name_style(a::CellShapeFunctions) = Dict(:color => :light_green)
 
 """
-    CellFunction{DS,F<:Function} <: AbstractCellFunction{DS}
+    CellFunction{DS,S,F<:Function} <: AbstractCellFunction{DS,S}
 
 Subtype of [`AbstractCellFunction`](@ref) used to wrap a function
 defined on a domain of style` `DS` in the cell
 """
-struct CellFunction{DS, F <: Function} <: AbstractCellFunction{DS}
+struct CellFunction{DS, S, F <: Function} <: AbstractCellFunction{DS, S}
     f::F
 end
 
 """
-    CellFunction(f::Function, domainstyle::DomainStyle)
+    CellFunction(f::Function, domainstyle::DomainStyle, ::Val{S}) where {S}
+
+`S` is the codomain size of `f`.
 """
-function CellFunction(f::Function, ds::DomainStyle)
-    CellFunction{typeof(ds), typeof(f)}(f)
+function CellFunction(f::Function, ds::DomainStyle, ::Val{S}) where {S}
+    CellFunction{typeof(ds), S, typeof(f)}(f)
 end
 
 get_function(f::CellFunction) = f.f
 
 """
-    PhysicalFunction(f::Function)
+    PhysicalFunction(f::Function, [size::Union{Integer, Tuple{Integer, Vararg{Integer}}} = 1])
+    PhysicalFunction(f::Function, ::Val{size}) where {size}
 
 Return a [`CellFunction`](@ref) defined on a `PhysicalDomain`.
+`size` is the size of the codomain of `f`.
+
+## Note:
+Using a `Val` to prescribe the size a of `PhysicalFunction` is
+recommended to improve type-stability and performance.
 """
-PhysicalFunction(f::Function) = CellFunction(f, PhysicalDomain())
+function PhysicalFunction(
+    f::Function,
+    size::Union{Integer, Tuple{Integer, Vararg{Integer}}} = 1,
+)
+    CellFunction(f, PhysicalDomain(), Val(size))
+end
+
+function PhysicalFunction(f::Function, s::Val{size}) where {size}
+    CellFunction(f, PhysicalDomain(), s)
+end
 
 """
-    ReferenceFunction(f::Function)
+    ReferenceFunction(f::Function, [size::Union{Integer, Tuple{Integer, Vararg{Integer}}} = 1])
+    ReferenceFunction(f::Function, ::Val{size}) where {size}
 
 Return a [`CellFunction`](@ref) defined on a `ReferenceDomain`.
+`size` is the size of the codomain of `f`.
+
+## Note:
+Using a `Val` to prescribe the size a of `ReferenceFunction` is
+recommended to improve type-stability and performance.
 """
-ReferenceFunction(f::Function) = CellFunction(f, ReferenceDomain())
+function ReferenceFunction(
+    f::Function,
+    size::Union{Integer, Tuple{Integer, Vararg{Integer}}} = 1,
+)
+    CellFunction(f, ReferenceDomain(), Val(size))
+end
+
+function ReferenceFunction(f::Function, s::Val{size}) where {size}
+    CellFunction(f, ReferenceDomain(), s)
+end
 
 function LazyOperators.materialize(::AbstractCellFunction, ::FaceInfo)
     error("Cannot integrate a `CellFunction` on a face : please select a `Side`")

src/fespace/fefunction.jl

@@ -28,7 +28,7 @@ function get_dof_values(f::AbstractFEFunction, icell, n::Val{N}) where {N}
     error("`get_dof_values` is not defined for type $(typeof(f))")
 end
 
-function Base.getindex(f::AbstractFEFunction, i::CellInfo)
+function Base.getindex(f::AbstractFEFunction{S}, i::CellInfo) where {S}
     feSpace = get_fespace(f)
     fSpace = get_function_space(feSpace)
     domainStyle = DomainStyle(fSpace)
@@ -38,7 +38,7 @@ function Base.getindex(f::AbstractFEFunction, i::CellInfo)
     ncomps = get_ncomponents(feSpace)
     q₀ = get_dof_values(f, cellindex(i), Val(ndofs))
     fcell = _interpolate(Val(ncomps), q₀, λ)
-    CellFunction(fcell, domainStyle)
+    CellFunction(fcell, domainStyle, Val(S))
 end
 
 # scalar case:

test/Project.toml

@@ -4,4 +4,5 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
-SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+SHA = "ea8e919c-243c-51af-8825-aaa63cd721ce"
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

test/interpolation/test_cellfunction.jl

@@ -57,10 +57,10 @@ end
         b = Bcube.CellPoint((x1_phys, x2_phys), cellinfo, Bcube.PhysicalDomain())
         _f = x -> x[1] + x[2]
 
-        λref = Bcube.CellFunction(_f, Bcube.ReferenceDomain())
+        λref = Bcube.CellFunction(_f, Bcube.ReferenceDomain(), Val(1))
         @test λref(a) == _f.((x1_ref, x2_ref))
 
-        λphys = Bcube.CellFunction(_f, Bcube.PhysicalDomain())
+        λphys = Bcube.CellFunction(_f, Bcube.PhysicalDomain(), Val(1))
         @test λphys(a) == _f.((x1_phys, x2_phys))
     end
 end

test/operator/test_algebra.jl

@@ -56,6 +56,20 @@
         scale!(mesh, 2.0)
         dΩ = Measure(CellDomain(mesh), 1)
         @test Bcube.compute(∫(∇(PhysicalFunction(x -> x[1])) ⋅ [1, 1])dΩ)[1] == 16.0
+
+        # Gradient of a vector PhysicalFunction
+        mesh = one_cell_mesh(:quad)
+        translate!(mesh, SA[π, -3.14]) # the translation vector can be anything
+        scale!(mesh, 2.0)
+        sizeU = spacedim(mesh)
+        U = TrialFESpace(FunctionSpace(:Lagrange, 1), mesh; sizeU)
+        V = TestFESpace(U)
+        _f = x -> SA[2 * x[1]^2, x[1] * x[2]]
+        f = PhysicalFunction(_f, sizeU)
+        ∇f = PhysicalFunction(x -> ForwardDiff.jacobian(_f, x), (sizeU, spacedim(mesh)))
+        dΩ = Measure(CellDomain(mesh), 3)
+        l(v) = ∫(tr(∇(f) - ∇f) ⋅ v)dΩ
+        @test assemble_linear(l, V) == [0.0, 0.0, 0.0, 0.0]
     end
 
     @testset "algebra" begin

test/runtests.jl

@@ -3,6 +3,7 @@ using Test
 using StaticArrays
 using LinearAlgebra
 using DelimitedFiles
+using ForwardDiff
 using SHA: sha1
 
 # from :