diff --git a/NEWS.md b/NEWS.md index bf10ca4..f8a21d1 100644 --- a/NEWS.md +++ b/NEWS.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- Added tag support for AD to align with Gridap[#1181](https://github.com/gridap/Gridap.jl/pull/1181) and Gridap[#1297](https://github.com/gridap/Gridap.jl/pull/1297). Since PR[#204](https://github.com/gridap/GridapDistributed.jl/pull/204). + ## [0.4.15] - 2026-03-23 ### Fixed @@ -15,7 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.4.14] - 2026-03-20 -### Changed +### Changed - Added support for Gridap 0.20.0. Since PR[#200](https://github.com/gridap/GridapDistributed.jl/pull/200). diff --git a/src/Autodiff.jl b/src/Autodiff.jl index 8d3a55d..773cc6e 100644 --- a/src/Autodiff.jl +++ b/src/Autodiff.jl @@ -11,11 +11,12 @@ end function FESpaces._gradient(f,uh,fuh::DistributedDomainContribution) local_terms = map(r -> DomainContribution(), get_parts(fuh)) local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) + tag = x->Fields.gradient(f, uh) for local_trians in local_domains g = FESpaces._change_argument(gradient,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_gradient(g,cell_u,cell_id) + cell_grad = distributed_autodiff_array_gradient(g,cell_u,cell_id,tag) map(add_contribution!,local_terms,local_trians,cell_grad) end DistributedDomainContribution(local_terms) @@ -29,11 +30,12 @@ end function FESpaces._jacobian(f,uh,fuh::DistributedDomainContribution) local_terms = map(r -> DomainContribution(), get_parts(fuh)) local_domains = tuple_of_arrays(map(Tuple∘get_domains,local_views(fuh))) + tag = x->Fields.jacobian(f, uh) for local_trians in local_domains g = FESpaces._change_argument(jacobian,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_jacobian(g,cell_u,cell_id) + cell_grad = distributed_autodiff_array_jacobian(g,cell_u,cell_id,tag) map(add_contribution!,local_terms,local_trians,cell_grad) end DistributedDomainContribution(local_terms) @@ -112,9 +114,13 @@ end # autodiff_array_xxx function distributed_autodiff_array_gradient(a,i_to_x) - dummy_tag = ()->() + tag = x->ForwardDiff.gradient(a, x) + distributed_autodiff_array_gradient(a,i_to_x,tag) +end + +function distributed_autodiff_array_gradient(a,i_to_x,tag::Function) i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x) + lazy_map(ConfigMap(ForwardDiff.gradient,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) @@ -127,9 +133,13 @@ function distributed_autodiff_array_gradient(a,i_to_x) end function distributed_autodiff_array_jacobian(a,i_to_x) - dummy_tag = ()->() + tag = x->ForwardDiff.jacobian(a, x) + distributed_autodiff_array_jacobian(a,i_to_x,tag) +end + +function distributed_autodiff_array_jacobian(a,i_to_x,tag::Function) i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x) + lazy_map(ConfigMap(ForwardDiff.jacobian,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) @@ -147,9 +157,13 @@ function distributed_autodiff_array_hessian(a,i_to_x) end function distributed_autodiff_array_gradient(a,i_to_x,j_to_i) - dummy_tag = ()->() + tag = x->ForwardDiff.gradient(a, x) + distributed_autodiff_array_gradient(a,i_to_x,j_to_i,tag) +end + +function distributed_autodiff_array_gradient(a,i_to_x,j_to_i,tag::Function) i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x) + lazy_map(ConfigMap(ForwardDiff.gradient,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) @@ -163,9 +177,13 @@ function distributed_autodiff_array_gradient(a,i_to_x,j_to_i) end function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i) - dummy_tag = ()->() + tag = x->ForwardDiff.jacobian(a, x) + distributed_autodiff_array_jacobian(a,i_to_x,j_to_i,tag) +end + +function distributed_autodiff_array_jacobian(a,i_to_x,j_to_i,tag::Function) i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x) + lazy_map(ConfigMap(ForwardDiff.jacobian,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) @@ -213,9 +231,13 @@ function FESpaces._change_argument(op,f,local_trians::AbstractArray{<:SkeletonTr end function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) - dummy_tag = ()->() + tag = x->ForwardDiff.gradient(a, x) + distributed_autodiff_array_gradient(a, i_to_x, j_to_i, tag) +end + +function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}, tag::Function) i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.gradient,dummy_tag),i_to_x) + lazy_map(ConfigMap(ForwardDiff.gradient,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) @@ -245,9 +267,13 @@ function distributed_autodiff_array_gradient(a, i_to_x, j_to_i::AbstractArray{<: end function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}) - dummy_tag = ()->() + tag = x->ForwardDiff.jacobian(a, x) + distributed_autodiff_array_jacobian(a, i_to_x, j_to_i, tag) +end + +function distributed_autodiff_array_jacobian(a, i_to_x, j_to_i::AbstractArray{<:SkeletonPair}, tag::Function) i_to_cfg = map(i_to_x) do i_to_x - lazy_map(ConfigMap(ForwardDiff.jacobian,dummy_tag),i_to_x) + lazy_map(ConfigMap(ForwardDiff.jacobian,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) diff --git a/test/AutodiffTests.jl b/test/AutodiffTests.jl index f86a79d..95da0ac 100644 --- a/test/AutodiffTests.jl +++ b/test/AutodiffTests.jl @@ -43,6 +43,40 @@ function main_sf(distribute,parts) b_AD = assemble_vector(gradient(g,uh),U) @test b ≈ b_AD + # Tags + dv = get_fe_basis(V) + dp = get_trial_fe_basis(U) + ph = interpolate(x->rand(),U) + ener(uh,ph) = ∫( 0.5*∇(uh)⋅∇(uh)*ph )*dΩ + res(uh,ph) = ∫( ∇(uh)⋅∇(dv)*ph )*dΩ + jac(uh,ph) = ∫( ∇(uh)⋅∇(dv)*dp )*dΩ + + ∂2L∂u∂p_auto = assemble_matrix(jacobian(ph->gradient(uh->ener(uh,ph),uh),ph),U,U) + ∂2L∂u∂p = assemble_matrix(jac(uh,ph),U,U) + @test reduce(&,map(≈,partition(∂2L∂u∂p_auto),partition(∂2L∂u∂p))) + + Γ_reg = BoundaryTriangulation(model) + Λ_reg = SkeletonTriangulation(model) + dΓ_reg = Measure(Γ_reg, 2) + dΛ_reg = Measure(Λ_reg, 2) + + reffe_l2 = ReferenceFE(lagrangian, Float64, 1) + V_reg = TestFESpace(model, reffe_l2, conformity=:L2) + U_reg = TrialFESpace(V_reg) + dv_reg = get_fe_basis(V_reg) + dp_reg = get_trial_fe_basis(U_reg) + uh_reg = FEFunction(U_reg, rand(num_free_dofs(U_reg))) + ph_reg = FEFunction(U_reg, rand(num_free_dofs(U_reg))) + + ener_reg(u, p) = ∫(0.5*u*u*p)*dΩ + ∫(0.5*u*u*p)*dΓ_reg + ∫(0.5*mean(u)*mean(u)*mean(p))*dΛ_reg + nested_ad_contrib = jacobian(p -> gradient(u -> ener_reg(u, p), uh_reg), ph_reg) + mat_nested_ad = assemble_matrix(nested_ad_contrib, U_reg, V_reg) + analytic_contrib = ∫(uh_reg*dv_reg*dp_reg)*dΩ + ∫(uh_reg*dv_reg*dp_reg)*dΓ_reg + + ∫(mean(uh_reg)*mean(dv_reg)*mean(dp_reg))*dΛ_reg + mat_analytic = assemble_matrix(analytic_contrib, U_reg, V_reg) + + @test reduce(&,map(≈,partition(mat_nested_ad),partition(mat_analytic))) + # Skeleton AD # I would like to compare the results, but we cannot be using FD in parallel... Λ = SkeletonTriangulation(model)