Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor of autodiff for MultiField #1070

Draft
wants to merge 9 commits into
base: master
Choose a base branch
from
Draft

Refactor of autodiff for MultiField #1070

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
bb0eafc
Started refactoring autodiff for multifield
JordiManyer Dec 21, 2024
f20a21c
Bugfixes
JordiManyer Dec 21, 2024
93eef54
Removed old code
JordiManyer Dec 21, 2024
eceacbe
Minor
JordiManyer Dec 21, 2024
bdcba1e
Recursive block structure
JordiManyer Dec 22, 2024
8fac19b
Ported skeleton AD to block interface
JordiManyer Dec 22, 2024
ca1393e
Deactivated SkeletonCellFieldPairTests
JordiManyer Dec 22, 2024
120499a
Changes to accommodate skeletons
JordiManyer Dec 23, 2024
6e649c5
Minor
JordiManyer Dec 23, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
82 changes: 33 additions & 49 deletions src/Arrays/Autodiff.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,19 +1,19 @@

function autodiff_array_gradient(a,i_to_x)
dummy_forwarddiff_tag = ()->()
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
dummy_tag = ()->()
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(),i_to_cfg,i_to_x)
i_to_ydual = a(i_to_xdual)
i_to_result = lazy_map(AutoDiffMap(ForwardDiff.gradient),i_to_ydual,i_to_x,i_to_cfg)
i_to_result = lazy_map(AutoDiffMap(),i_to_cfg,i_to_ydual)
i_to_result
end

function autodiff_array_jacobian(a,i_to_x)
dummy_forwarddiff_tag = ()->()
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
dummy_tag = ()->()
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(),i_to_cfg,i_to_x)
i_to_ydual = a(i_to_xdual)
i_to_result = lazy_map(AutoDiffMap(ForwardDiff.jacobian),i_to_ydual,i_to_x,i_to_cfg)
i_to_result = lazy_map(AutoDiffMap(),i_to_cfg,i_to_ydual)
i_to_result
end

Expand All @@ -23,22 +23,22 @@ function autodiff_array_hessian(a,i_to_x)
end

function autodiff_array_gradient(a,i_to_x,j_to_i)
dummy_forwarddiff_tag = ()->()
i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
dummy_tag = ()->()
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(),i_to_cfg,i_to_x)
j_to_ydual = a(i_to_xdual)
j_to_x = lazy_map(Reindex(i_to_x),j_to_i)
j_to_cfg = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x)
j_to_result = lazy_map(AutoDiffMap(ForwardDiff.gradient),j_to_ydual,j_to_x,j_to_cfg)
j_to_cfg = lazy_map(Reindex(i_to_cfg),j_to_i)
j_to_result = lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual)
j_to_result
end

function autodiff_array_jacobian(a,i_to_x,j_to_i)
dummy_forwarddiff_tag = ()->()
i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
dummy_tag = ()->()
i_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
i_to_xdual = lazy_map(DualizeMap(),i_to_cfg,i_to_x)
j_to_ydual = a(i_to_xdual)
j_to_x = lazy_map(Reindex(i_to_x),j_to_i)
j_to_cfg = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x)
j_to_result = lazy_map(AutoDiffMap(ForwardDiff.jacobian),j_to_ydual,j_to_x,j_to_cfg)
j_to_cfg = lazy_map(Reindex(i_to_cfg),j_to_i)
j_to_result = lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual)
j_to_result
end

Expand All @@ -54,7 +54,7 @@ struct ConfigMap{
f::F # ForwardDiff operation
tag::T # function for config tag name
end
# constructor with nothing as the tag, for backwards compatibility

ConfigMap(f) = ConfigMap(f,nothing)

# TODO Prescribing long chunk size can lead to slow compilation times!
Expand All @@ -72,52 +72,36 @@ function evaluate!(cfg,k::ConfigMap,x)
cfg
end

struct DualizeMap{
F <: Union{typeof(ForwardDiff.gradient),typeof(ForwardDiff.jacobian)},
T <: Union{<:Function,Nothing}} <: Map

f::F # ForwardDiff operation
tag::T # function for config tag name
end
# constructor with nothing as the tag, for backwards compatibility
DualizeMap(f) = DualizeMap(f,nothing)

function return_cache(k::DualizeMap,x)
return_cache(ConfigMap(k.f,k.tag),x)
end
struct DualizeMap <: Map end

function evaluate!(cfg,k::DualizeMap,x)
xdual = cfg.duals
ForwardDiff.seed!(xdual, x, cfg.seeds)
function evaluate!(cache,::DualizeMap,cfg,x)
xdual, seeds = cfg.duals, cfg.seeds
ForwardDiff.seed!(xdual, x, seeds)
xdual
end

struct AutoDiffMap{F} <: Map
f::F
end
struct AutoDiffMap <: Map end

function return_cache(k::AutoDiffMap,ydual,x,cfg::ForwardDiff.GradientConfig{T}) where T
function return_cache(::AutoDiffMap,cfg::ForwardDiff.GradientConfig,ydual)
ydual isa Real || throw(ForwardDiff.GRAD_ERROR)
result = similar(x, ForwardDiff.valtype(ydual))
result
result = similar(cfg.duals, ForwardDiff.valtype(ydual))
return result
end

function evaluate!(result,k::AutoDiffMap,ydual,x,cfg::ForwardDiff.GradientConfig{T}) where T
@notimplementedif ForwardDiff.chunksize(cfg) != length(x)
@notimplementedif length(result) != length(x)
function evaluate!(result,::AutoDiffMap,cfg::ForwardDiff.GradientConfig{T},ydual) where T
@check ForwardDiff.chunksize(cfg) == length(result)
result = ForwardDiff.extract_gradient!(T, result, ydual)
return result
end

function return_cache(k::AutoDiffMap,ydual,x,cfg::ForwardDiff.JacobianConfig{T,V,N}) where {T,V,N}
function return_cache(::AutoDiffMap,cfg::ForwardDiff.JacobianConfig{T,V,N},ydual) where {T,V,N}
ydual isa AbstractArray || throw(ForwardDiff.JACOBIAN_ERROR)
result = similar(ydual, ForwardDiff.valtype(eltype(ydual)), length(ydual), N)
result
return result
end

function evaluate!(result,k::AutoDiffMap,ydual,x,cfg::ForwardDiff.JacobianConfig{T,V,N}) where {T,V,N}
@notimplementedif ForwardDiff.chunksize(cfg) != length(x)
@notimplementedif size(result,2) != length(x)
function evaluate!(result,::AutoDiffMap,cfg::ForwardDiff.JacobianConfig{T,V,N},ydual) where {T,V,N}
@check ForwardDiff.chunksize(cfg) == size(result,2)
ForwardDiff.extract_jacobian!(T, result, ydual, N)
ForwardDiff.extract_value!(T, result, ydual)
return result
Expand Down
2 changes: 1 addition & 1 deletion src/CellData/SkeletonCellFieldPair.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ into it. Ideally, if we could parse and extract only the Skeleton integration
terms from the functional's julia function form, this fix is not required, but
this is not trivial to do. On the positive side, since the evaluations are all
lazy and not used, this doesn't put any noticeable memory or computational
overhead. Ofcourse, it is made sure that the such plus side pick doesn't happen
overhead. Of course, it is made sure that the such plus side pick doesn't happen
when the integration over the SkeletonTriangulation
=#
# If SkeletonCellFieldPair is evaluated we just pick the plus side parent
Expand Down
200 changes: 49 additions & 151 deletions src/FESpaces/FEAutodiff.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,31 +18,14 @@ function _gradient(f,uh,fuh::DomainContribution)
terms = DomainContribution()
for trian in get_domains(fuh)
g = _change_argument(gradient,f,trian,uh)
cell_u = get_cell_dof_values(uh)
cell_u = _get_cell_dof_values(uh,trian)
cell_id = _compute_cell_ids(uh,trian)
cell_grad = autodiff_array_gradient(g,cell_u,cell_id)
add_contribution!(terms,trian,cell_grad)
end
terms
end

function _compute_cell_ids(uh,ttrian)
strian = get_triangulation(uh)
if strian === ttrian
return collect(IdentityVector(Int32(num_cells(strian))))
end
@check is_change_possible(strian,ttrian)
D = num_cell_dims(strian)
sglue = get_glue(strian,Val(D))
tglue = get_glue(ttrian,Val(D))
@notimplementedif !isa(sglue,FaceToFaceGlue)
@notimplementedif !isa(tglue,FaceToFaceGlue)
scells = IdentityVector(Int32(num_cells(strian)))
mcells = extend(scells,sglue.mface_to_tface)
tcells = lazy_map(Reindex(mcells),tglue.tface_to_mface)
collect(tcells)
end

function jacobian(f::Function,uh::FEFunction)
fuh = f(uh)
_jacobian(f,uh,fuh)
Expand All @@ -61,7 +44,7 @@ function _jacobian(f,uh,fuh::DomainContribution)
terms = DomainContribution()
for trian in get_domains(fuh)
g = _change_argument(jacobian,f,trian,uh)
cell_u = get_cell_dof_values(uh)
cell_u = _get_cell_dof_values(uh,trian)
cell_id = _compute_cell_ids(uh,trian)
cell_grad = autodiff_array_jacobian(g,cell_u,cell_id)
add_contribution!(terms,trian,cell_grad)
Expand All @@ -87,15 +70,15 @@ function _hessian(f,uh,fuh::DomainContribution)
terms = DomainContribution()
for trian in get_domains(fuh)
g = _change_argument(hessian,f,trian,uh)
cell_u = get_cell_dof_values(uh)
cell_u = _get_cell_dof_values(uh,trian)
cell_id = _compute_cell_ids(uh,trian)
cell_grad = autodiff_array_hessian(g,cell_u,cell_id)
add_contribution!(terms,trian,cell_grad)
end
terms
end

function _change_argument(op,f,trian,uh::SingleFieldFEFunction)
function _change_argument(op,f,trian,uh)
U = get_fe_space(uh)
function g(cell_u)
cf = CellField(U,cell_u)
Expand All @@ -105,144 +88,59 @@ function _change_argument(op,f,trian,uh::SingleFieldFEFunction)
g
end

#= AD for DomainContribution involving SkeletonTriangulation

- Following are the constructs for performing gradient of DomainContribution
involving integrations over SkeletonTriangulation (Λ)
- The current approach followed to achieve the above is performing the Gridap
way of AD for plus and minus sides of the FEFunction occurring at Λ separately,
and combining the result. So as to Dualize only either plus side or minus
side of CellField/FEFunction we introduce the SkeletonCellFieldPair, which
stores two CellFields, one of which in the use case here is the dualized
version of the other.
- Currently, Jacobian and hence Hessian are not yet fully implemented and work-
in-progress.
- AD for integration over SkeletonTriangulation has not yet been implemented
for the MultiField case
_get_cell_dof_values(uh,trian) = get_cell_dof_values(uh)

=#
function _compute_cell_ids(uh,ttrian)
strian = get_triangulation(uh)
if strian === ttrian
return collect(IdentityVector(Int32(num_cells(strian))))
end
@check is_change_possible(strian,ttrian)
D = num_cell_dims(strian)
sglue = get_glue(strian,Val(D))
tglue = get_glue(ttrian,Val(D))
@notimplementedif !isa(sglue,FaceToFaceGlue)
@notimplementedif !isa(tglue,FaceToFaceGlue)
scells = IdentityVector(Int32(num_cells(strian)))
mcells = extend(scells,sglue.mface_to_tface)
tcells = lazy_map(Reindex(mcells),tglue.tface_to_mface)
collect(tcells)
end

function _change_argument(
op,f,
trian::SkeletonTriangulation,
uh::SingleFieldFEFunction)
# Skeleton AD

function _change_argument(op,f,trian::SkeletonTriangulation,uh)
U = get_fe_space(uh)
function g(cell_u)
uh_dual = CellField(U,cell_u)
scfp_plus = SkeletonCellFieldPair(uh_dual, uh)
scfp_minus = SkeletonCellFieldPair(uh, uh_dual)
cell_grad_plus = f(scfp_plus)
cell_grad_minus = f(scfp_minus)
get_contribution(cell_grad_plus,trian), get_contribution(cell_grad_minus,trian)
cf = SkeletonCellFieldPair(U,cell_u)
cell_grad = f(cf)
get_contribution(cell_grad,trian)
end
g
end

function _get_cell_dof_values(uh,trian::SkeletonTriangulation)
plus = _get_cell_dof_values(uh,trian.plus)
minus = _get_cell_dof_values(uh,trian.minus)
lazy_map(BlockMap(2,[1,2]),plus,minus)
end

function _compute_cell_ids(uh,ttrian::SkeletonTriangulation)
tcells_plus = _compute_cell_ids(uh,ttrian.plus)
tcells_minus = _compute_cell_ids(uh,ttrian.minus)
SkeletonPair(tcells_plus,tcells_minus)
end

## overloads for AD of SkeletonTriangulation DomainContribution ##

#= Notes regarding the placement of below AD functions for Skeleton Integration

- Earlier, the autodiff_array_### family of functions have been placed in the
`src/Arrays/Autodiff.jl`, but as below we are leveraging `BlockMap` for the
construction derivatives of Skeleton integration terms, this cannot be
placed in `Arrays/Autodiff.jl` due to circular dependency, this is due to the
fact that `BlockMap` belongs to Gridap.Fields which uses/imports functions
and constructs from Gridap.Arrays, so using `BlockMap` in Gridap.Arrays would
create a circular dependency.
- So as to have everything working we will need `use` some of the constructs
from Gridap.Arrays, `Gridap.Geometry.SkeletonPair` and
`Gridap.CellData.SkeletonCellFieldPair`
- This comes with an added advantage of being able use `SkeletonPair` in
`_compute_cell_ids` for `SkeletonTriangulation`
- we need to also think if all of AD can be moved into a separate module in
Gridap, where we can use or import required functionalities from other
modules without any circular dependency
=#
function autodiff_array_gradient(a, i_to_x, j_to_i::SkeletonPair)
dummy_forwarddiff_tag = ()->()
i_to_xdual = lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)

# dual output of both sides at once
j_to_ydual_plus, j_to_ydual_minus = a(i_to_xdual)

# Work for plus side
j_to_x_plus = lazy_map(Reindex(i_to_x),j_to_i.plus)
j_to_cfg_plus = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x_plus)
j_to_result_plus = lazy_map(AutoDiffMap(ForwardDiff.gradient),
j_to_ydual_plus,j_to_x_plus,j_to_cfg_plus)

# Work for minus side
j_to_x_minus = lazy_map(Reindex(i_to_x),j_to_i.minus)
j_to_cfg_minus = lazy_map(ConfigMap(ForwardDiff.gradient,dummy_forwarddiff_tag),j_to_x_minus)
j_to_result_minus = lazy_map(AutoDiffMap(ForwardDiff.gradient),
j_to_ydual_minus,j_to_x_minus,j_to_cfg_minus)

# Assemble on SkeletonTriangulation expects an array of interior of facets
# where each entry is a 2-block BlockVector with the first block being the
# contribution of the plus side and the second, the one of the minus side
lazy_map(BlockMap(2,[1,2]),j_to_result_plus,j_to_result_minus)
end

_length_1st_vector_vectorblock(a::VectorBlock{<:Vector}) = length(a.array[1])

function autodiff_array_jacobian(a,i_to_x,j_to_i::SkeletonPair)
dummy_forwarddiff_tag = ()->()
i_to_xdual = lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
j_to_ydual_plus, j_to_ydual_minus = a(i_to_xdual)

# Bilinear form when tested with test basis functions dv results in
# DomainContribution containing vector of 2-block `VectorBlock{Vector}`
# each block coming from plus side and minus side of dv.
# So we densify each of the `VectorBlock`` into plain vectors and construct
# back the jacobian contributions into a MatrixBlock

densify = DensifyInnerMostBlockLevelMap()
j_to_ydual_plus_dense = lazy_map(densify, j_to_ydual_plus)
j_to_ydual_minus_dense = lazy_map(densify, j_to_ydual_minus)

# Work for plus side
j_to_x_plus = lazy_map(Reindex(i_to_x),j_to_i.plus)
j_to_cfg_plus = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x_plus)
j_to_result_plus_dense = lazy_map(AutoDiffMap(ForwardDiff.jacobian),
j_to_ydual_plus_dense,j_to_x_plus,j_to_cfg_plus)

# Work for minus side
j_to_x_minus = lazy_map(Reindex(i_to_x),j_to_i.minus)
j_to_cfg_minus = lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),j_to_x_minus)
j_to_result_minus_dense = lazy_map(AutoDiffMap(ForwardDiff.jacobian),
j_to_ydual_minus_dense,j_to_x_minus,j_to_cfg_minus)

# j_to_result_plus_dense/j_to_result_minus_dense can be (and must be)
# laid out into 2x2 block matrices

lengths_1st_vector_vectorblocks_plus = lazy_map(_length_1st_vector_vectorblock,j_to_ydual_plus)
lengths_1st_vector_vectorblocks_minus = lazy_map(_length_1st_vector_vectorblock,j_to_ydual_minus)

J_11 = lazy_map((x,b)->view(x, 1:b,:),
j_to_result_plus_dense,lengths_1st_vector_vectorblocks_plus)

J_21 = lazy_map((x,b)->view(x, b+1:size(x,1),:),
j_to_result_plus_dense,lengths_1st_vector_vectorblocks_plus)

J_12 = lazy_map((x,b)->view(x, 1:b,:),
j_to_result_minus_dense,lengths_1st_vector_vectorblocks_minus)

J_22 = lazy_map((x,b)->view(x, b+1:size(x,1),:),
j_to_result_minus_dense,lengths_1st_vector_vectorblocks_minus)

# Assembly on SkeletonTriangulation expects an array of facets where each
# entry is a 2x2-block MatrixBlock with the blocks of the Jacobian matrix
bm_jacobian = BlockMap((2,2),[(1,1),(2,1),(1,2),(2,2)])
lazy_map(bm_jacobian, J_11, J_21, J_12, J_22)
end

function autodiff_array_hessian(a,i_to_x,i_to_j::SkeletonPair)
@notimplemented
plus = _compute_cell_ids(uh,ttrian.plus)
minus = _compute_cell_ids(uh,ttrian.minus)
lazy_map(BlockMap(2,[1,2]),plus,minus)
end

function CellData.SkeletonCellFieldPair(V::FESpace,cell_values)
plus = CellField(V,lazy_map(x->x.array[1],cell_values))
minus = CellField(V,lazy_map(x->x.array[2],cell_values))
CellData.SkeletonCellFieldPair(plus,minus)
end

function Arrays.lazy_map(r::Reindex, ids::LazyArray{<:Fill{BlockMap{1}}})
k = ids.maps.value
plus = lazy_map(r,ids.args[1])
# minus = lazy_map(r,ids.args[2])
# lazy_map(k,plus,minus)
plus
end
Loading
Loading