Merge pull request #29 from bcube-project/grad_lazyop

Gradient for operators

src/algebra/gradient.jl

@@ -21,10 +21,30 @@ function LazyOperators.materialize(
     Gradient(f, cPoint)
 end
 
+function LazyOperators.materialize(
+    lOp::Gradient{O, <:Tuple{AbstractLazy}},
+    cPoint::CellPoint,
+) where {O}
+    f, = get_args(lOp)
+    Gradient(f, cPoint)
+end
+
 """
 Materialization of a `Gradient` on a `CellPoint`. Only valid for a function and a `CellPoint` defined on
 the reference domain.
 
+# Implementation
+The user writes mathematical expressions in the PhysicalDomain. So the gradient always represents the derivation
+with respect to the physical domain spatial coordinates, even if evaluated on a point expressed in the ReferenceDomain.
+
+The current Gradient implementation consists in applying ForwardDiff on the given operator 'u', on a point in the
+ReferenceDomain. That is to say, we compute ForwarDiff.derivative(ξ -> u ∘ F, ξ) (where F is the identity if u is defined
+is the ReferenceDomain). This gives ∇(u ∘ F)(ξ) = t∇(F)(ξ) * ∇(u)(F(x)).
+However, we only want ∇(u)(F(x)) : that's why a multiplication by the transpose of the inverse mapping jacobian is needed.
+
+An alternative approach would be to apply ForwardDiff in the PhysicalDomain : ForwarDiff.derivative(x -> u ∘ F^-1, x). The
+problem is that the inverse mapping F^-1 is not always defined.
+
 # Maths notes
 We use the following convention for any function f:R^n->R^p : ∇f is a tensor/matrix equal to ∂fi/∂xj. However when f
 is a scalar function, i.e f:R^n -> R, then ∇f is a column vector (∂f/∂x1, ∂f/∂x2, ...).
@@ -77,19 +97,31 @@ TODO:
 * improve formulae with a reshape
 * Specialize for a ShapeFunction to use the hardcoded version instead of ForwardDiff
 """
-function Gradient(
-    cellFunction::AbstractCellFunction{<:ReferenceDomain},
-    cPoint::CellPoint{ReferenceDomain},
-)
-    cnodes = get_cellnodes(cPoint)
-    ctype = get_celltype(cPoint)
-    ξ = get_coord(cPoint)
-    f = get_function(cellFunction)
-    m = mapping_jacobian_inv(cnodes, ctype, ξ)
-    _Gradient(cellFunction(cPoint), f, ξ, m)
+function Gradient(op::AbstractLazy, cPoint::CellPoint{ReferenceDomain})
+    m = mapping_jacobian_inv(get_cellnodes(cPoint), get_celltype(cPoint), get_coord(cPoint))
+    f(ξ) = op(CellPoint(ξ, get_cellinfo(cPoint), ReferenceDomain()))
+    valS = _size_codomain(f, get_coord(cPoint))
+    return _gradient(valS, f, get_coord(cPoint), m)
+
+    # ForwarDiff applied in the PhysicalDomain : not viable because
+    # the inverse mapping is not always known.
+    # cPoint_phys = change_domain(cPoint, PhysicalDomain())
+    # f(ξ) = op(CellPoint(ξ, get_cellinfo(cPoint_phys), PhysicalDomain()))
+    # fx = f(get_coord(cPoint_phys))
+    # return _gradient_or_jacobian(Val(length(fx)), f, get_coord(cPoint_phys))
+end
+
+_size_codomain(f, x) = Val(length(f(x)))
+_size_codomain(f::AbstractCellFunction, x) = Val(get_size(f))
+
+_gradient(::Val{1}, f, ξ::AbstractArray, m) = transpose(m) * ForwardDiff.gradient(f, ξ)
+_gradient(::Val{S}, f, ξ::AbstractArray, m) where {S} = ForwardDiff.jacobian(f, ξ) * m
+
+function Gradient(op::AbstractLazy, cPoint::CellPoint{PhysicalDomain})
+    f(x) = op(CellPoint(x, cPoint.cellinfo, PhysicalDomain()))
+    valS = _size_codomain(f, get_coord(cPoint))
+    return _gradient_or_jacobian(valS, f, get_coord(cPoint))
 end
-_Gradient(::Real, f, ξ::AbstractArray, m) = transpose(m) * ForwardDiff.gradient(f, ξ)
-_Gradient(::AbstractArray, f, ξ::AbstractArray, m) = ForwardDiff.jacobian(f, ξ) * m
 
 function Gradient(
     cellFunction::AbstractCellShapeFunctions{<:ReferenceDomain},
@@ -103,15 +135,6 @@ function Gradient(
     MapOver(grad_shape_functionsNA(fs, n, ctype, cnodes, ξ))
 end
 
-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 _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)

src/mapping/mapping.jl

@@ -53,6 +53,9 @@ end
 
 Jacobian matrix of the inverse mapping : ``\\dfrac{\\partial F_i^{-1}}{\\partial x_j}``
 
+Contrary to `mapping_jacobian_inv`, this function is not always defined because the
+inverse mapping, F^-1, is not always defined.
+
 # Implementation
 Default version using LinearAlgebra to inverse the matrix, but can be specified for each shape (if it exists).
 """

src/mesh/mesh_generator.jl

@@ -624,14 +624,14 @@ function transform!(mesh::AbstractMesh, fun)
 end
 
 """
-    translate(mesh::AbstractMesh, t::Vector{Float64})
+    translate(mesh::AbstractMesh, t::AbstractVector)
 
 Translate the input mesh with vector `t`.
 
 Usefull for debugging.
 """
-translate(mesh::AbstractMesh, t::AbstractVector{Float64}) = transform(mesh, x -> x + t)
-translate!(mesh::AbstractMesh, t::AbstractVector{Float64}) = transform!(mesh, x -> x + t)
+translate(mesh::AbstractMesh, t::AbstractVector) = transform(mesh, x -> x + t)
+translate!(mesh::AbstractMesh, t::AbstractVector) = transform!(mesh, x -> x + t)
 
 """
     _duplicate_mesh(mesh::AbstractMesh)

test/mapping/test_mapping.jl

@@ -120,6 +120,17 @@
         @test isapprox_arrays(mapping_inv(n, ct, x3), SA[1.0, 1.0]; rtol = tol)
         @test isapprox_arrays(mapping_inv(n, ct, x4), SA[-1.0, 1.0]; rtol = tol)
 
+        θ = π / 5
+        s = 3
+        t = SA[-1, 2]
+        R(θ) = SA[cos(θ) -sin(θ); sin(θ) cos(θ)]
+        mesh = one_cell_mesh(:quad)
+        mesh = transform(mesh, x -> R(θ) * (s .* x .+ t)) # scale, translate and rotate
+        c = CellInfo(mesh, 1)
+        cnodes = nodes(c)
+        ctype = Bcube.celltype(c)
+        @test all(mapping_jacobian_inv(cnodes, ctype, SA[0.0, 0.0]) .≈ R(-θ) ./ s)
+
         # Quad order 2
         xmin, ymin = -rand(2)
         xmax, ymax = rand(2)

test/operator/test_algebra.jl

@@ -69,7 +69,110 @@
         ∇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]
+        _a = assemble_linear(l, V)
+        @test all(isapprox.(_a, [0.0, 0.0, 0.0, 0.0]; atol = 100 * eps()))
+
+        @testset "AbstractLazy" begin
+            mesh = one_cell_mesh(:quad)
+            scale!(mesh, 3.0)
+            translate!(mesh, [4.0, 0.0])
+            U_sca = TrialFESpace(FunctionSpace(:Lagrange, 1), mesh)
+            U_vec = TrialFESpace(FunctionSpace(:Lagrange, 1), mesh; size = 2)
+            V_sca = TestFESpace(U_sca)
+            V_vec = TestFESpace(U_vec)
+            u_sca = FEFunction(U_sca)
+            u_vec = FEFunction(U_vec)
+            dΩ = Measure(CellDomain(mesh), 2)
+            projection_l2!(u_sca, PhysicalFunction(x -> 3 * x[1] - 4x[2]), dΩ)
+            projection_l2!(
+                u_vec,
+                PhysicalFunction(x -> SA[2x[1] + 5x[2], 4x[1] - 3x[2]]),
+                dΩ,
+            )
+
+            l1(v) = ∫((∇(π * u_sca) ⋅ ∇(2 * u_sca)) ⋅ v)dΩ
+            l2(v) = ∫((∇(u_sca) ⋅ ∇(u_sca)) ⋅ v)dΩ
+
+            a1_sca = assemble_linear(l1, V_sca)
+            a2_sca = assemble_linear(l2, V_sca)
+            @test all(a1_sca .≈ (2π .* a2_sca))
+
+            V_vec = TestFESpace(U_vec)
+            l1_vec(v) = ∫((∇(π * u_vec) * u_vec) ⋅ v)dΩ
+            l2_vec(v) = ∫((∇(u_vec) * u_vec) ⋅ v)dΩ
+            a1_vec = assemble_linear(l1_vec, V_vec)
+            a2_vec = assemble_linear(l2_vec, V_vec)
+            @test all(a1_vec .≈ (π .* a2_vec))
+
+            # Testing without assemble_*linear
+            θ = π / 5
+            s = 3
+            t = SA[-1, 2]
+            R = SA[cos(θ) -sin(θ); sin(θ) cos(θ)]
+            mesh = one_cell_mesh(:quad)
+
+            transform!(mesh, x -> R * (s .* x .+ t)) # scale, translate and rotate
+
+            # Select a cell and get its info
+            c = CellInfo(mesh, 1)
+            cnodes = nodes(c)
+            ctype = Bcube.celltype(c)
+            F = mapping(cnodes, ctype)
+            # tJinv = transpose(R ./ s) # if we want the analytic one...
+            tJinv(ξ) = transpose(mapping_jacobian_inv(cnodes, ctype, ξ))
+
+            # Test 1
+            u1 = PhysicalFunction(x -> x[1])
+            u2 = PhysicalFunction(x -> x[2])
+            u_a = u1 * u1 + 2 * u1 * u2 + u2 * u2 * u2
+            u_b = PhysicalFunction(x -> x[1]^2 + 2 * x[1] * x[2] + x[2]^3)
+            ∇u_ana = x -> SA[2 * (x[1] + x[2]); 2 * x[1] + 3 * x[2]^2]
+
+            ξ = Bcube.CellPoint(SA[0.5, -0.1], c, Bcube.ReferenceDomain())
+            x = Bcube.change_domain(ξ, Bcube.PhysicalDomain())
+            ∇u = ∇u_ana(Bcube.get_coord(x))
+            ∇u_a_ref = Bcube.materialize(∇(u_a), ξ)
+            ∇u_b_ref = Bcube.materialize(∇(u_b), ξ)
+            ∇u_a_phy = Bcube.materialize(∇(u_a), x)
+            ∇u_b_phy = Bcube.materialize(∇(u_b), x)
+            @test all(∇u_a_ref .≈ ∇u)
+            @test all(∇u_b_ref .≈ ∇u)
+            @test all(∇u_a_phy .≈ ∇u)
+            @test all(∇u_b_phy .≈ ∇u)
+
+            # Test 2
+            u1 = ReferenceFunction(ξ -> ξ[1])
+            u2 = ReferenceFunction(ξ -> ξ[2])
+            u_a = u1 * u1 + 2 * u1 * u2 + u2 * u2 * u2
+            u_b = ReferenceFunction(ξ -> ξ[1]^2 + 2 * ξ[1] * ξ[2] + ξ[2]^3)
+            ∇u_ana = ξ -> SA[2 * (ξ[1] + ξ[2]); 2 * ξ[1] + 3 * ξ[2]^2]
+
+            x = Bcube.CellPoint(SA[0.5, -0.1], c, Bcube.PhysicalDomain())
+            ξ = Bcube.change_domain(x, Bcube.ReferenceDomain()) # not always possible, but ok of for quad
+            ∇u = ∇u_ana(Bcube.get_coord(ξ))
+            _tJinv = tJinv(Bcube.get_coord(ξ))
+            ∇u_a_ref = Bcube.materialize(∇(u_a), ξ)
+            ∇u_b_ref = Bcube.materialize(∇(u_b), ξ)
+            ∇u_a_phy = Bcube.materialize(∇(u_a), x)
+            ∇u_b_phy = Bcube.materialize(∇(u_b), x)
+            @test all(∇u_a_ref .≈ _tJinv * ∇u)
+            @test all(∇u_b_ref .≈ _tJinv * ∇u)
+            @test all(∇u_a_phy .≈ _tJinv * ∇u)
+            @test all(∇u_b_phy .≈ _tJinv * ∇u)
+
+            # Test 3
+            u_phy = PhysicalFunction(x -> x[1]^2 + 2 * x[1] * x[2] + x[2]^3)
+            u_ref = ReferenceFunction(ξ -> ξ[1]^2 + 2 * ξ[1] * ξ[2] + ξ[2]^3)
+            ∇u_ana = t -> SA[2 * (t[1] + t[2]); 2 * t[1] + 3 * t[2]^2]
+
+            ξ = Bcube.CellPoint(SA[0.5, -0.1], c, Bcube.ReferenceDomain())
+            x = Bcube.change_domain(ξ, Bcube.PhysicalDomain())
+            ∇u_ref = ∇u_ana(Bcube.get_coord(ξ))
+            ∇u_phy = ∇u_ana(Bcube.get_coord(x))
+            _tJinv = tJinv(Bcube.get_coord(ξ))
+            @test all(Bcube.materialize(∇(u_phy + u_ref), ξ) .≈ ∇u_phy + _tJinv * ∇u_ref)
+            @test all(Bcube.materialize(∇(u_phy + u_ref), x) .≈ ∇u_phy + _tJinv * ∇u_ref)
+        end
     end
 
     @testset "algebra" begin