Skip to content

Commit

Permalink
Merge pull request #169 from gridap/autodiff
Browse files Browse the repository at this point in the history 
AD for multifield and skeletons
  • Loading branch information
@JordiManyer
JordiManyer authored Mar 4, 2025
2 parents b5a5944 + 86d0d93 commit 3e4fafe
Show file tree
Hide file tree
Showing 7 changed files with 306 additions and 79 deletions.
6 changes: 6 additions & 0 deletions NEWS.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).

## [0.4.7] 2025-03-04

### Added

- Extended support for automatic differentiation to multi-field spaces and skeleton triangulations. Since PR[#169](https://github.com/gridap/GridapDistributed.jl/pull/169).

## [0.4.6] 2024-12-03

### Added
Expand Down
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "GridapDistributed"
uuid = "f9701e48-63b3-45aa-9a63-9bc6c271f355"
authors = ["S. Badia <santiago.badia@monash.edu>", "A. F. Martin <alberto.f.martin@anu.edu.au>", "F. Verdugo <f.verdugo.rojano@vu.nl>"]
version = "0.4.6"
version = "0.4.7"

[deps]
BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
Expand Down
227 changes: 187 additions & 40 deletions src/Autodiff.jl
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@

function Fields.gradient(f::Function,uh::DistributedCellField)
const DistributedADTypes = Union{DistributedCellField,DistributedMultiFieldCellField}

# Distributed counterpart of: src/FESpaces/FEAutodiff.jl

function Fields.gradient(f::Function,uh::DistributedADTypes)
fuh = f(uh)
FESpaces._gradient(f,uh,fuh)
end
Expand All @@ -9,15 +13,15 @@ function FESpaces._gradient(f,uh,fuh::DistributedDomainContribution)
local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh)))
for local_trians in local_domains
g = FESpaces._change_argument(gradient,f,local_trians,uh)
cell_u = map(get_cell_dof_values,local_views(uh))
cell_u = map(FESpaces.get_cell_dof_values,local_views(uh))
cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians)
cell_grad = distributed_autodiff_array_gradient(g,cell_u,cell_id)
map(add_contribution!,local_terms,local_trians,cell_grad)
end
DistributedDomainContribution(local_terms)
end

function Fields.jacobian(f::Function,uh::DistributedCellField)
function Fields.jacobian(f::Function,uh::DistributedADTypes)
fuh = f(uh)
FESpaces._jacobian(f,uh,fuh)
end
Expand All @@ -27,79 +31,222 @@ function FESpaces._jacobian(f,uh,fuh::DistributedDomainContribution)
local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh)))
for local_trians in local_domains
g = FESpaces._change_argument(jacobian,f,local_trians,uh)
cell_u = map(get_cell_dof_values,local_views(uh))
cell_u = map(FESpaces.get_cell_dof_values,local_views(uh))
cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians)
cell_grad = distributed_autodiff_array_jacobian(g,cell_u,cell_id)
map(add_contribution!,local_terms,local_trians,cell_grad)
end
DistributedDomainContribution(local_terms)
end

function FESpaces._change_argument(op,f,trian,uh::DistributedCellField)
spaces = map(get_fe_space,local_views(uh))
function g(cell_u)
cf = DistributedCellField(
map(CellField,spaces,cell_u),
get_triangulation(uh),
function Fields.hessian(f::Function,uh::DistributedADTypes)
fuh = f(uh)
FESpaces._hessian(f,uh,fuh)
end

function FESpaces._hessian(f,uh,fuh::DistributedDomainContribution)
local_terms = map(r -> DomainContribution(), get_parts(fuh))
local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh)))
for local_trians in local_domains
g = FESpaces._change_argument(hessian,f,local_trians,uh)
cell_u = map(FESpaces.get_cell_dof_values,local_views(uh))
cell_id = map(FESpaces._compute_cell_ids,local_views(uh),local_trians)
cell_grad = distributed_autodiff_array_hessian(g,cell_u,cell_id)
map(add_contribution!,local_terms,local_trians,cell_grad)
end
DistributedDomainContribution(local_terms)
end

function FESpaces._change_argument(op,f,local_trians,uh::DistributedADTypes)
function dist_cf(uh::DistributedCellField,cfs)
DistributedCellField(cfs,get_triangulation(uh))
end
function dist_cf(uh::DistributedMultiFieldCellField,cfs)
sf_cfs = map(DistributedCellField,
[tuple_of_arrays(map(cf -> Tuple(cf.single_fields),cfs))...],
map(get_triangulation,uh)
)
cell_grad = f(cf)
map(get_contribution,local_views(cell_grad),trian)
DistributedMultiFieldCellField(sf_cfs,cfs)
end

uhs = local_views(uh)
spaces = map(get_fe_space,uhs)
function g(cell_u)
cfs = map(CellField,spaces,cell_u)
cf = dist_cf(uh,cfs)
cg = f(cf)
map(get_contribution,local_views(cg),local_trians)
end
g
end

# Distributed counterpart of: src/Arrays/Autodiff.jl
# autodiff_array_xxx

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

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

function distributed_autodiff_array_hessian(a,i_to_x)
agrad = i_to_y -> distributed_autodiff_array_gradient(a,i_to_y)
distributed_autodiff_array_jacobian(agrad,i_to_x)
end

function distributed_autodiff_array_gradient(a,i_to_x,j_to_i)
dummy_forwarddiff_tag = ()->()
i_to_xdual = map(i_to_x) do i_to_x
lazy_map(DualizeMap(ForwardDiff.gradient,dummy_forwarddiff_tag),i_to_x)
dummy_tag = ()->()
i_to_cfg = map(i_to_x) do i_to_x
lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
end
i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
lazy_map(DualizeMap(),i_to_cfg,i_to_x)
end
j_to_ydual = a(i_to_xdual)
j_to_result = map(i_to_x,j_to_i,j_to_ydual) do i_to_x,j_to_i,j_to_ydual
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)
lazy_map(AutoDiffMap(ForwardDiff.gradient),j_to_ydual,j_to_x,j_to_cfg)
j_to_result = map(i_to_cfg,j_to_i,j_to_ydual) do i_to_cfg,j_to_i,j_to_ydual
j_to_cfg = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i)
lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual)
end
return j_to_result
end

function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i)
dummy_forwarddiff_tag = ()->()
i_to_xdual = map(i_to_x) do i_to_x
lazy_map(DualizeMap(ForwardDiff.jacobian,dummy_forwarddiff_tag),i_to_x)
dummy_tag = ()->()
i_to_cfg = map(i_to_x) do i_to_x
lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
end
i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
lazy_map(DualizeMap(),i_to_cfg,i_to_x)
end
j_to_ydual = a(i_to_xdual)
j_to_result = map(i_to_x,j_to_i,j_to_ydual) do i_to_x,j_to_i,j_to_ydual
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)
lazy_map(AutoDiffMap(ForwardDiff.jacobian),j_to_ydual,j_to_x,j_to_cfg)
j_to_result = map(i_to_cfg,j_to_i,j_to_ydual) do i_to_cfg,j_to_i,j_to_ydual
j_to_cfg = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i)
lazy_map(AutoDiffMap(),j_to_cfg,j_to_ydual)
end
j_to_result
return j_to_result
end

function distributed_autodiff_array_hessian(a,i_to_x,i_to_j)
agrad = i_to_y -> distributed_autodiff_array_gradient(a,i_to_y,i_to_j)
distributed_autodiff_array_jacobian(agrad,i_to_x,i_to_j)
end

# Skeleton AD

function FESpaces._change_argument(op,f,local_trians::AbstractArray{<:SkeletonTriangulation},uh::DistributedADTypes)
function dist_cf(uh::DistributedCellField,cfs)
DistributedCellField(cfs, get_triangulation(uh))
end
function dist_cf(uh::DistributedMultiFieldCellField,cfs)
sf_cfs = map(DistributedCellField,
[tuple_of_arrays(map(cf -> Tuple(cf.single_fields),cfs))...],
map(get_triangulation,uh)
)
DistributedMultiFieldCellField(sf_cfs,cfs)
end

uhs = local_views(uh)
spaces = map(get_fe_space,uhs)
function g(cell_u)
uhs_dual = map(CellField,spaces,cell_u)
cf_plus = dist_cf(uh,map(SkeletonCellFieldPair,uhs_dual,uhs))
cf_minus = dist_cf(uh,map(SkeletonCellFieldPair,uhs,uhs_dual))
cg_plus = f(cf_plus)
cg_minus = f(cf_minus)
plus = map(get_contribution,local_views(cg_plus),local_trians)
minus = map(get_contribution,local_views(cg_minus),local_trians)
plus, minus
end
g
end

function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair})
dummy_tag = ()->()
i_to_cfg = map(i_to_x) do i_to_x
lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x)
end
i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
lazy_map(DualizeMap(),i_to_cfg,i_to_x)
end

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

j_to_result = map(i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus) do i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus
# Work for plus side
j_to_cfg_plus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.plus)
j_to_result_plus = lazy_map(AutoDiffMap(),j_to_cfg_plus,j_to_ydual_plus)

# Work for minus side
j_to_cfg_minus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.minus)
j_to_result_minus = lazy_map(AutoDiffMap(),j_to_cfg_minus,j_to_ydual_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
is_single_field = eltype(eltype(j_to_result_plus)) <: Number
k = is_single_field ? BlockMap(2,[1,2]) : Fields.BlockBroadcasting(BlockMap(2,[1,2]))
lazy_map(k,j_to_result_plus,j_to_result_minus)
end

return j_to_result
end

function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair})
dummy_tag = ()->()
i_to_cfg = map(i_to_x) do i_to_x
lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x)
end
i_to_xdual = map(i_to_cfg,i_to_x) do i_to_cfg, i_to_x
lazy_map(DualizeMap(),i_to_cfg,i_to_x)
end

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

j_to_result = map(i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus) do i_to_cfg,j_to_i,j_to_ydual_plus,j_to_ydual_minus
# Work for plus side
j_to_cfg_plus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.plus)
j_to_result_plus = lazy_map(AutoDiffMap(),j_to_cfg_plus,j_to_ydual_plus)

# Work for minus side
j_to_cfg_minus = Arrays.autodiff_array_reindex(i_to_cfg,j_to_i.minus)
j_to_result_minus = lazy_map(AutoDiffMap(),j_to_cfg_minus,j_to_ydual_minus)

# Merge the columns into a 2x2 block matrix
# I = [[(CartesianIndex(i,),CartesianIndex(i,j)) for i in 1:2] for j in 1:2]
I = [
[(CartesianIndex(1,), CartesianIndex(1, 1)), (CartesianIndex(2,), CartesianIndex(2, 1))], # Plus -> First column
[(CartesianIndex(1,), CartesianIndex(1, 2)), (CartesianIndex(2,), CartesianIndex(2, 2))] # Minus -> Second column
]
is_single_field = eltype(eltype(j_to_result_plus)) <: AbstractArray
k = is_single_field ? Fields.MergeBlockMap((2,2),I) : Fields.BlockBroadcasting(Fields.MergeBlockMap((2,2),I))
lazy_map(k,j_to_result_plus,j_to_result_minus)
end

return j_to_result
end
Loading

2 comments on commit 3e4fafe

@JordiManyer
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register()

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/126259

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.4.7 -m "<description of version>" 3e4fafec239ba98000839f6affb8f0dcb69217c3
git push origin v0.4.7

Please sign in to comment.