-
Notifications
You must be signed in to change notification settings - Fork 100
/
Copy pathMultiFieldFEAutodiff.jl
142 lines (133 loc) · 5.17 KB
/
MultiFieldFEAutodiff.jl
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
# 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
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)
push!(single_fields,cf)
end
xh = MultiFieldCellField(single_fields)
cell_grad = f(xh)
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
function FESpaces._change_argument(
op::typeof(gradient),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
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)
push!(single_fields,cf)
end
xh = MultiFieldCellField(single_fields)
cell_grad = f(xh)
get_contribution(cell_grad,trian)
end
g
end
function Algebra.hessian(f::Function,uh::MultiFieldFEFunction)
@notimplemented
end
#function FESpaces._change_argument(
# op::typeof(hessian),f,trian,uh::MultiFieldFEFunction)
#
# U = get_fe_space(uh)
# function g(cell_u)
# single_fields = GenericCellField[]
# nfields = length(U.spaces)
# for i in 1:nfields
# cell_values_field = lazy_map(a->a.array[i],cell_u)
# cf = CellField(U.spaces[i],cell_values_field)
# cell_data = lazy_map(BlockMap((nfields,nfields),i),get_data(cf))
# uhi = GenericCellField(cell_data,get_triangulation(cf),DomainStyle(cf))
# push!(single_fields,uhi)
# end
# xh = MultiFieldCellField(single_fields)
# cell_grad = f(xh)
# get_contribution(cell_grad,trian)
# end
# g
#end