Solving ad multifield

src/Arrays/Autodiff.jl

@@ -106,4 +106,3 @@ function evaluate!(result,k::AutoDiffMap,ydual,x,cfg::ForwardDiff.JacobianConfig
   ForwardDiff.extract_value!(T, result, ydual)
   return result
 end
-

src/MultiField/MultiField.jl

@@ -20,6 +20,7 @@ using Gridap.FESpaces: SingleFieldFEBasis, TestBasis, TrialBasis
 using Gridap.CellData: CellFieldAt
 import Gridap.Fields: gradient, DIV, ∇∇
 
+using ForwardDiff
 using FillArrays
 using SparseArrays
 using LinearAlgebra

src/MultiField/MultiFieldFEAutodiff.jl

@@ -1,21 +1,96 @@
 
+# TO IMPROVE .... Perhaps via a Map?
+# Do we have already this function in Gridap?
+# What if there are no cells? I am assuming that there is at least one cell
+# What if the number of dofs per field per cell is different among cells?
+function _get_cell_dofs_field_offsets(uh::MultiFieldFEFunction)
+  U = get_fe_space(uh)
+  uh_dofs = get_cell_dof_values(uh)[1]
+  nfields = length(U.spaces)
+  dofs_field_offsets=Vector{Int}(undef,nfields+1)
+  dofs_field_offsets[1]=1
+  for i in 1:nfields
+    dofs_field_offsets[i+1]=dofs_field_offsets[i]+length(uh_dofs.array[i])
+  end
+  dofs_field_offsets
+end
+
+function FESpaces._gradient(f,uh::MultiFieldFEFunction,fuh::DomainContribution)
+  terms = DomainContribution()
+  U = get_fe_space(uh)
+  for trian in get_domains(fuh)
+    g = FESpaces._change_argument(gradient,f,trian,uh)
+    cell_u = lazy_map(DensifyInnerMostBlockLevelMap(),get_cell_dof_values(uh))
+    cell_id = FESpaces._compute_cell_ids(uh,trian)
+    cell_grad = autodiff_array_gradient(g,cell_u,cell_id)
+    monolithic_result=cell_grad
+    blocks = [] # TO-DO type unstable. How can I infer the type of its entries?
+    nfields = length(U.spaces)
+    cell_dofs_field_offsets=_get_cell_dofs_field_offsets(uh)
+    for i in 1:nfields
+      view_range=cell_dofs_field_offsets[i]:cell_dofs_field_offsets[i+1]-1
+      block=lazy_map(x->view(x,view_range),monolithic_result)
+      append!(blocks,[block])
+    end
+    cell_grad=lazy_map(BlockMap(nfields,collect(1:nfields)),blocks...)
+    add_contribution!(terms,trian,cell_grad)
+  end
+  terms
+end
+
+function FESpaces._jacobian(f,uh::MultiFieldFEFunction,fuh::DomainContribution)
+  terms = DomainContribution()
+  U = get_fe_space(uh)
+  for trian in get_domains(fuh)
+    g = FESpaces._change_argument(jacobian,f,trian,uh)
+    cell_u = lazy_map(DensifyInnerMostBlockLevelMap(),get_cell_dof_values(uh))
+    cell_id = FESpaces._compute_cell_ids(uh,trian)
+    cell_grad = autodiff_array_jacobian(g,cell_u,cell_id)
+    monolithic_result=cell_grad
+    blocks        = [] # TO-DO type unstable. How can I infer the type of its entries?
+    blocks_coords = Tuple{Int,Int}[]
+    nfields = length(U.spaces)
+    cell_dofs_field_offsets=_get_cell_dofs_field_offsets(uh)
+    for j=1:nfields
+      view_range_j=cell_dofs_field_offsets[j]:cell_dofs_field_offsets[j+1]-1
+      for i=1:nfields
+        view_range_i=cell_dofs_field_offsets[i]:cell_dofs_field_offsets[i+1]-1
+        # TO-DO: depending on the residual being differentiated, we may end with
+        #        blocks [i,j] full of zeros. I guess that it might desirable to early detect
+        #        these zero blocks and use a touch[i,j]==false block in ArrayBlock.
+        #        How can we detect that we have a zero block?
+        block=lazy_map(x->view(x,view_range_i,view_range_j),monolithic_result)
+        append!(blocks,[block])
+        append!(blocks_coords,[(i,j)])
+      end
+    end
+    cell_grad=lazy_map(BlockMap((nfields,nfields),blocks_coords),blocks...)
+    add_contribution!(terms,trian,cell_grad)
+  end
+  terms
+end
+
 function FESpaces._change_argument(
   op::typeof(jacobian),f,trian,uh::MultiFieldFEFunction)
 
   U = get_fe_space(uh)
   function g(cell_u)
     single_fields = GenericCellField[]
     nfields = length(U.spaces)
+    cell_dofs_field_offsets=_get_cell_dofs_field_offsets(uh)
     for i in 1:nfields
-      cell_values_field = lazy_map(a->a.array[i],cell_u)
+      view_range=cell_dofs_field_offsets[i]:cell_dofs_field_offsets[i+1]-1
+      cell_values_field = lazy_map(a->view(a,view_range),cell_u)
       cf = CellField(U.spaces[i],cell_values_field)
-      cell_data = lazy_map(BlockMap((1,nfields),i),get_data(cf))
-      uhi = GenericCellField(cell_data,get_triangulation(cf),DomainStyle(cf))
-      push!(single_fields,uhi)
+      push!(single_fields,cf)
     end
     xh = MultiFieldCellField(single_fields)
     cell_grad = f(xh)
-    get_contribution(cell_grad,trian)
+    cell_grad_cont_block=get_contribution(cell_grad,trian)
+    bs = [cell_dofs_field_offsets[i+1]-cell_dofs_field_offsets[i] for i=1:nfields]
+    lazy_map(DensifyInnerMostBlockLevelMap(),
+             Fill(bs,length(cell_grad_cont_block)),
+             cell_grad_cont_block)
   end
   g
 end
@@ -27,12 +102,12 @@ function FESpaces._change_argument(
   function g(cell_u)
     single_fields = GenericCellField[]
     nfields = length(U.spaces)
+    cell_dofs_field_offsets=_get_cell_dofs_field_offsets(uh)
     for i in 1:nfields
-      cell_values_field = lazy_map(a->a.array[i],cell_u)
+      view_range=cell_dofs_field_offsets[i]:cell_dofs_field_offsets[i+1]-1
+      cell_values_field = lazy_map(a->view(a,view_range),cell_u)
       cf = CellField(U.spaces[i],cell_values_field)
-      cell_data = lazy_map(BlockMap(nfields,i),get_data(cf))
-      uhi = GenericCellField(cell_data,get_triangulation(cf),DomainStyle(cf))
-      push!(single_fields,uhi)
+      push!(single_fields,cf)
     end
     xh = MultiFieldCellField(single_fields)
     cell_grad = f(xh)

test/MultiFieldTests/MultiFieldFEAutodiffTests.jl

@@ -39,7 +39,7 @@ xh = FEFunction(Y,rand(num_free_dofs(Y)))
 
 @test j(xh,dx,dy)[Ω][end][1,1] != nothing
 @test j(xh,dx,dy)[Ω][end][2,1] != nothing
-@test j(xh,dx,dy)[Ω][end][1,2] == nothing
+@test j(xh,dx,dy)[Ω][end][1,2] != nothing
 @test j(xh,dx,dy)[Ω][end][2,2] != nothing
 @test g(xh,dy)[Ω][end][1] != nothing
 @test g(xh,dy)[Ω][end][2] != nothing
@@ -52,11 +52,57 @@ dx = get_trial_fe_basis(Y)
 dy = get_fe_basis(Y)
 xh = FEFunction(Y,rand(num_free_dofs(Y)))
 
-@test_broken j(xh,dx,dy)[Ω][end][1,1] != nothing
-@test_broken j(xh,dx,dy)[Ω][end][2,1] != nothing
-@test_broken j(xh,dx,dy)[Ω][end][1,2] == nothing
-@test_broken j(xh,dx,dy)[Ω][end][2,2] != nothing
-@test_broken g(xh,dy)[Ω][end][1] != nothing
-@test_broken g(xh,dy)[Ω][end][2] != nothing
+@test j(xh,dx,dy)[Ω][end][1,1] != nothing
+@test j(xh,dx,dy)[Ω][end][2,1] != nothing
+@test j(xh,dx,dy)[Ω][end][1,2] != nothing
+@test j(xh,dx,dy)[Ω][end][2,2] != nothing
+@test g(xh,dy)[Ω][end][1] != nothing
+@test g(xh,dy)[Ω][end][2] != nothing
+
+eu(uh) = ∫( uh*uh )dΩ
+ep(ph) = ∫( ph*ph )dΩ
+eup((uh,ph)) = ∫( uh*uh + ph*ph )dΩ
+geu(xh,dy) = gradient(xh->eu(xh),xh)
+gep(xh,dy) = gradient(xh->ep(xh),xh)
+geup(xh,dy) = gradient(xh->eup(xh),xh)
+
+uh,ph=xh
+geu_uh=geu(uh,dy)
+gep_ph=gep(ph,dy)
+geup_uh_ph=geup(xh,dy)
+
+a=geu_uh[Ω]
+b=gep_ph[Ω]
+c=geup_uh_ph[Ω]
+
+@test all(a .== map(x->x.array[1],c))
+@test all(b .== map(x->x.array[2],c))
+
+
+ru(uh,v) = ∫( v*uh*uh )dΩ
+rp(ph,q) = ∫( q*ph*ph )dΩ
+rup((uh,ph),(v,q)) = ∫( v*uh*uh + q*ph*ph )dΩ
+ju(xh,dx,dy) = jacobian(xh->ru(xh,dy),xh)
+jp(xh,dx,dy) = jacobian(xh->rp(xh,dy),xh)
+jup(xh,dx,dy) = jacobian(xh->rup(xh,dy),xh)
+
+uh=FEFunction(V1,get_free_dof_values(xh[1]))
+ph=FEFunction(V2,get_free_dof_values(xh[2]))
+du = get_trial_fe_basis(V1)
+dv = get_fe_basis(V1)
+dp = get_trial_fe_basis(V2)
+dq = get_fe_basis(V2)
+
+
+ju_uh=ju(uh,du,dv)
+jp_ph=jp(ph,dp,dq)
+jup_uh_ph=jup(xh,dx,dy)
+
+a=ju_uh[Ω]
+b=jp_ph[Ω]
+c=jup_uh_ph[Ω]
+
+@test all(a .== map(x->x.array[1,1],c))
+@test all(b .== map(x->x.array[2,2],c))
 
 end # module

test/MultiFieldTests/MultiFieldFEOperatorsTests.jl

@@ -59,12 +59,9 @@ b = residual(op,xh)
 A = jacobian(op,xh)
 test_fe_operator(op,get_free_dof_values(xh),b)
 
-@test_broken begin
 op_auto = FEOperator(r,X,Y)
 A_auto = jacobian(op_auto,xh)
 @test A ≈ A_auto
-true
-end
 
 r_const((u,p),(v,q)) = -1.0 * (∫( v*1.0 )*dΩ)
 op_const = FEOperator(r_const,j,X,Y)