Multi field autodiff

src/Arrays/BlockArraysCoo.jl

@@ -108,9 +108,11 @@ function Base.similar(a::BlockArrayCoo{S,N} where S,::Type{T}, axs::NTuple{N,<:B
 end
 
 function _similar_block_array_coo(a,::Type{T},axs) where T
-  S = eltype(a)
-  @notimplementedif S != T
-  A = eltype(a.blocks)
+  if length(a.blocks) != 0
+    A = typeof(similar(first(a.blocks),T))
+  else
+    A = typeof(similar(first(a.zero_blocks),T))
+  end
   blocks = A[]
   for p in 1:length(a.blocks)
     I = a.blockids[p]
@@ -231,6 +233,22 @@ function Base.getindex(a::BlockVectorCoo,i::Integer,j::Integer...)
   ai
 end
 
+function Base.setindex!(a::BlockArrayCoo{T,N} where T,v,i::Vararg{Integer,N}) where N
+  s = map(findblockindex,a.axes,i)
+  s = map(findblockindex,a.axes,(i,))
+  I = Block(map(i->i.I[1],s)...)
+  α = CartesianIndex(map(BlockArrays.blockindex,s))
+  a[I][α] = v
+end
+
+function Base.setindex!(a::BlockVectorCoo,v,i::Integer,j::Integer...)
+  @assert all( j .== 1)
+  s = map(findblockindex,a.axes,(i,))
+  I = Block(map(i->i.I[1],s)...)
+  α = CartesianIndex(map(BlockArrays.blockindex,s))
+  a[I][α] = v
+end
+
 function Base.:+(a::BlockArrayCoo{Ta,N},b::BlockArrayCoo{Tb,N}) where {Ta,Tb,N}
   @assert axes(a) == axes(b)
   @assert blockaxes(a) == blockaxes(b)

src/MultiField/MultiFieldFESpaces.jl

@@ -91,7 +91,13 @@ function FESpaces.EvaluationFunction(fe::MultiFieldFESpace, free_values)
 end
 
 function CellData.CellField(fe::MultiFieldFESpace,cell_values)
-  @notimplemented
+  single_fields = CellField[]
+  for field in 1:num_fields(fe)
+    cell_values_field = apply(a->a[Block(field)],cell_values)
+    cf = CellField(fe.spaces[field],cell_values_field)
+    push!(single_fields,cf)
+  end
+  MultiFieldCellField(single_fields)
 end
 
 function CellData.CellField(fe::MultiFieldFESpace,cell_values::VectorOfBlockArrayCoo)

test/ArraysTests/AutodiffTests.jl

@@ -1,6 +1,7 @@
 module AutodiffTests
 
 using Gridap.Arrays
+using BlockArrays
 
 function user_cell_energy(cell_u)
   function f(u)
@@ -15,12 +16,14 @@ end
 
 function user_cell_residual(cell_u)
   function f(u)
-    r = copy(u)
-    r .= zero(eltype(u))
-    for (i,ui) in enumerate(u)
-      r[i] = 2*ui
-    end
+    r = 2*u
     r
+    #r = copy(u)
+    #r .= zero(eltype(u))
+    #for (i,ui) in enumerate(u)
+    #  r[i] = 2*ui
+    #end
+    #r
   end
   apply(f,cell_u)
 end
@@ -39,6 +42,24 @@ end
 
 L = 10
 l = 24 # Do not use a number <13 (too easy for ForwardDiff)
+
+blocksids = [(1,),(2,)]
+axs = (blockedrange([4,4]),)
+cell_u = [ BlockArrayCoo([rand(4),rand(4)],blocksids,axs) for i in 1:L ]
+
+cell_e = user_cell_energy(cell_u)
+cell_r = user_cell_residual(cell_u)
+cell_j = user_cell_jacobian(cell_u)
+cell_h = cell_j
+
+cell_r_auto = autodiff_array_gradient(user_cell_energy,cell_u)
+cell_j_auto = autodiff_array_jacobian(user_cell_residual,cell_u)
+cell_h_auto = autodiff_array_hessian(user_cell_energy,cell_u)
+
+test_array(cell_r_auto,cell_r)
+test_array(cell_j_auto,cell_j)
+test_array(cell_h_auto,cell_h)
+
 cell_u = [ rand(l) for i in 1:L ]
 
 cell_e = user_cell_energy(cell_u)
@@ -84,4 +105,10 @@ test_array(cell_r_Γ_auto,cell_r_Γ)
 test_array(cell_j_Γ_auto,cell_j_Γ)
 test_array(cell_h_Γ_auto,cell_h_Γ)
 
+
+
+
+#cell_u = [ rand(l) for i in 1:L ]
+
+
 end # module

test/MultiFieldTests/MultiFieldFEOperatorsTests.jl

@@ -67,5 +67,39 @@ b = residual(op,xh)
 A = jacobian(op,xh)
 test_fe_operator(op,get_free_values(xh),b)
 
+function r(x,y)
+  u,p = x
+  v,q = y
+  v*(u*u) + v*p*u - q*p - v*4 + q
+end
+
+function j(x,dx,y)
+  u,p = x
+  du,dp = dx
+  v,q = y
+  2*v*u*du + v*dp*u + v*p*du - q*dp
+end
+
+t_Ω = FETerm(r,j,trian,quad)
+t_Ω_auto = FETerm(r,trian,quad)
+
+x = FEFunction(X,rand(num_free_dofs(X)))
+dx = get_cell_basis(X)
+y = get_cell_basis(Y)
+
+cell_r = get_cell_residual(t_Ω,x,y)
+cell_j = get_cell_jacobian(t_Ω,x,dx,y)
+
+cell_r_auto = get_cell_residual(t_Ω_auto,x,y)
+cell_j_auto = get_cell_jacobian(t_Ω_auto,x,dx,y)
+test_array(cell_r_auto,cell_r,≈)
+test_array(cell_j_auto,cell_j,≈)
+
+op = FEOperator(X,Y,t_Ω_auto)
+xh = zero(X)
+b = residual(op,xh)
+A = jacobian(op,xh)
+test_fe_operator(op,get_free_values(xh),b)
+
 
 end # module