Skip to content

Commit

Permalink
Second implementation of differentiable triangulations
Browse files Browse the repository at this point in the history
  • Loading branch information
@JordiManyer
JordiManyer committed Sep 5, 2024
1 parent 8c746ca commit 317b754
Show file tree
Hide file tree
Showing 3 changed files with 329 additions and 105 deletions.
40 changes: 20 additions & 20 deletions scripts/Embedded/MWEs/jordi_embedded.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import Gridap.Geometry: get_node_coordinates, collect1d
include("../differentiable_trians.jl")

order = 1
model = CartesianDiscreteModel((0,1,0,1),(8,8))
model = CartesianDiscreteModel((0,1,0,1),(20,20))
Ω = Triangulation(model)
dΩ = Measure(Ω,2*order)

Expand All @@ -18,40 +18,40 @@ V = TestFESpace(model,reffe_scalar)
U = TrialFESpace(V)
V_φ = TestFESpace(model,reffe_scalar)

φh = interpolate(x->max(abs(x[1]-0.5),abs(x[2]-0.5))-0.25,V_φ)
# φh = interpolate(x->max(abs(x[1]-0.5),abs(x[2]-0.5))-0.11,V_φ)
#φh = interpolate(x->sqrt((x[1]-0.5)^2+(x[2]-0.5)^2)-0.302,V_φ)
#φh = interpolate(x->max(abs(x[1]-0.5),abs(x[2]-0.5))-0.25,V_φ)
φh = interpolate(x->max(abs(x[1]-0.5),abs(x[2]-0.5))-0.11,V_φ)
#φh = interpolate(x->sqrt((x[1]-0.5)^2+(x[2]-0.5)^2)-0.303,V_φ)
x_φ = get_free_dof_values(φh)
any(iszero,x_φ)
any(x -> x < 0,x_φ)
any(x -> x > 0,x_φ)

Ωin = DifferentiableTriangulation(φh) do φh
geo = DiscreteGeometry(φh,model)
cutgeo = cut(model,geo)
return Triangulation(cutgeo,PHYSICAL_IN)
end
geo = DiscreteGeometry(φh,model)
cutgeo = cut(model,geo)

Ωin = Triangulation(cutgeo,PHYSICAL_IN).a
Ωout = Triangulation(cutgeo,PHYSICAL_OUT).a

Ωout = DifferentiableTriangulation(φh) do φh
geo = DiscreteGeometry(φh,model)
cutgeo = cut(model,geo)
return Triangulation(cutgeo,PHYSICAL_OUT)
end
diff_Ωin = DifferentiableTriangulation(Ωin)
diff_Ωout = DifferentiableTriangulation(Ωout)

dΩin = Measure(Ωin,5*order)
dΩin = Measure(diff_Ωin,5*order)
j_in(φ) = ∫(1)dΩin
dj_in = gradient(j_in,φh)
dj_vec_in = assemble_vector(dj_in,V_φ)
norm(dj_vec_in)

dj_in_a = dj_in[Ωin].a

dΩout = Measure(Ωout,5*order)
dΩout = Measure(diff_Ωout,5*order)
j_out(φ) = ∫(1)dΩout
dj_out = gradient(j_out,φh)
dj_vec_out = -assemble_vector(dj_out,V_φ)
norm(dj_vec_out)

geo = DiscreteGeometry(φh,model)
cutgeo = cut(model,geo)
Γ = EmbeddedBoundary(cutgeo)
dΓ = Measure(Γ,2*order)
dj_expected(q) = ∫(-q)dΓ
dj_exp_vec = assemble_vector(dj_expected,V_φ)
norm(dj_exp_vec)

norm(dj_vec_in-dj_exp_vec)
norm(dj_vec_out-dj_exp_vec)
Expand Down
244 changes: 159 additions & 85 deletions scripts/Embedded/differentiable_trians.jl
View file
Original file line number Diff line number Diff line change
@@ -1,91 +1,162 @@

using Gridap
using Gridap.CellData, Gridap.Geometry, Gridap.Helpers, Gridap.Arrays
using Gridap.Fields, Gridap.ReferenceFEs

function CellData.get_contribution(a::DomainContribution,trian::Geometry.AppendedTriangulation)
if haskey(a.dict,trian)
return a.dict[trian]
elseif haskey(a.dict,trian.a)
return a.dict[trian.a]
elseif haskey(a.dict,trian.b)
return a.dict[trian.b]
else
@unreachable """\n
There is not contribution associated with the given mesh in this DomainContribution object.
"""
end
end
using Gridap.Geometry: num_nodes

function FESpaces._compute_cell_ids(uh,ttrian::AppendedTriangulation)
ids_a = FESpaces._compute_cell_ids(uh,ttrian.a)
ids_b = FESpaces._compute_cell_ids(uh,ttrian.b)
lazy_append(ids_a,ids_b)
end
using FillArrays

using Gridap.Geometry
using Gridap.Geometry: num_nodes
using GridapEmbedded
using GridapEmbedded.LevelSetCutters, GridapEmbedded.Interfaces

function GridapEmbedded.LevelSetCutters._get_value_at_coords(φh::CellField,model::DiscreteModel{Dc,Dp}) where {Dc,Dp}
@assert DomainStyle(φh) == ReferenceDomain()
# Cell-to-node map for the original model
c2n_map = collect1d(get_cell_node_ids(model))
# GridapEmbedded

# Cell-wise node coordinates (in ReferenceDomain coordinates)
cell_reffe = get_cell_reffe(model)
cell_node_coords = lazy_map(get_node_coordinates,cell_reffe)
function compute_cell_maps(cell_coords,cell_reffes)
cell_shapefuns = lazy_map(get_shapefuns,cell_reffes)
default_cell_map = lazy_map(linear_combination,cell_coords,cell_shapefuns)
default_cell_grad = lazy_map(∇,default_cell_map)
cell_poly = lazy_map(get_polytope,cell_reffes)
cell_q0 = lazy_map(p->zero(first(get_vertex_coordinates(p))),cell_poly)
origins = lazy_map(evaluate,default_cell_map,cell_q0)
gradients = lazy_map(evaluate,default_cell_grad,cell_q0)
cell_map = lazy_map(Gridap.Fields.affine_map,gradients,origins)
return cell_map
end

weights = fill(0.0,num_nodes(model))
for cell in eachindex(c2n_map)
for node in c2n_map[cell]
weights[node] += 1.0
end
end
for node in 1:num_nodes(model)
weights[node] = 1.0 / weights[node]
end
function get_edge_list(poly::Polytope)
ltcell_to_lpoints, simplex = simplexify(poly)
simplex_edges = get_faces(simplex,1,0)
ltcell_to_edges = map(pts -> map(e -> pts[e], simplex_edges), ltcell_to_lpoints)
return unique(sort,vcat(ltcell_to_edges...))
end

# Get cell data
φh_data = CellData.get_data(φh)
T = return_type(testitem(CellData.get_data(φh)),testitem(testitem(cell_node_coords)))
values =zeros(T,num_nodes(model))
cell_node_coords_cache = array_cache(cell_node_coords)
# Loop over cells
for cell in eachindex(c2n_map)
field = φh_data[cell]
node_coords = getindex!(cell_node_coords_cache,cell_node_coords,cell)
for (iN,node) in enumerate(c2n_map[cell])
val = field(node_coords[iN])
values[node] += val * weights[node]
function belongs_to_edge(p,edge,bgpts)
p1,p2 = bgpts[edge]
return iszero(cross(p-p1,p2-p1))
end

function compute_coords(
coords ::Vector{<:Point{Dp,Tp}},
rcoords::Vector{<:Point{Dp,Tp}},
bg_coords::Vector{<:Point{Dp,Tp}},
bg_rcoords::Vector{<:Point{Dp,Tp}},
values::Vector{Tv},
edges
) where {Dp,Tp,Tv}
T = Point{Dp,Tv}
new_coords = Vector{T}(undef,length(coords))
new_rcoords = Vector{T}(undef,length(rcoords))
for (i,(q,p)) in enumerate(zip(coords,rcoords))
if p ∈ bg_rcoords
new_coords[i] = q
new_rcoords[i] = p
continue
end
e = findfirst(edge -> belongs_to_edge(p,edge,bg_rcoords), edges)
n1, n2 = edges[e]
q1, q2 = bg_coords[n1], bg_coords[n2]
p1, p2 = bg_rcoords[n1], bg_rcoords[n2]
v1, v2 = values[n1], values[n2]
w1, w2 = abs(v1), abs(v2)
λ = w1/(w1+w2)
new_coords[i] = q1 + λ*(q2-q1)
new_rcoords[i] = p1 + λ*(p2-p1)
end
return new_coords, new_rcoords
end

for node in 1:num_nodes(model)
if iszero(values[node])
values[node] -= eps(T)
end
function compute_coords(
cell_to_coords::AbstractVector{<:AbstractVector{<:Point{Dp,Tp}}},
cell_to_rcoords::AbstractVector{<:AbstractVector{<:Point{Dp,Tp}}},
cell_to_bgcoords::AbstractVector{<:AbstractVector{<:Point{Dp,Tp}}},
cell_to_bg_rcoords::AbstractVector{<:AbstractVector{<:Point{Dp,Tp}}},
cell_to_values::AbstractVector{<:AbstractVector{Tv}},
cell_to_edges
) where {Dp,Tp,Tv}
T = Point{Dp,Tv}
results = lazy_map(compute_coords,cell_to_coords,cell_to_rcoords,cell_to_bgcoords,cell_to_bg_rcoords,cell_to_values,cell_to_edges)
cache = array_cache(results)
ncells = length(cell_to_coords)
new_cell_to_coords = Vector{Vector{T}}(undef,ncells)
new_cell_to_rcoords = Vector{Vector{T}}(undef,ncells)
for cell in 1:ncells
new_coords, new_rcoords = getindex!(cache,results,cell)
new_cell_to_coords[cell] = new_coords
new_cell_to_rcoords[cell] = new_rcoords
end
return new_cell_to_coords, new_cell_to_rcoords
end

function compute_dual_values(
trian::GridapEmbedded.Interfaces.SubCellTriangulation,
φh::CellField
)
bgmodel = get_background_model(trian)
subcells = trian.subcells

cell_to_bgcell = subcells.cell_to_bgcell
cell_to_points = subcells.cell_to_points
point_to_rcoords = subcells.point_to_rcoords
point_to_coords = subcells.point_to_coords

ctypes = get_cell_type(bgmodel)
bgcell_to_polys = expand_cell_data(get_polytopes(bgmodel),ctypes)
bgcell_to_coords = get_cell_coordinates(bgmodel)
bgcell_to_rcoords = lazy_map(get_vertex_coordinates,bgcell_to_polys)
bgcell_to_edges = lazy_map(get_edge_list,bgcell_to_polys)
bgcell_to_values = lazy_map(evaluate,CellData.get_data(φh),bgcell_to_rcoords)

cell_to_reffes = Fill(LagrangianRefFE(Float64,TRI,1),length(cell_to_points))
cell_to_bgcoords = lazy_map(Reindex(bgcell_to_coords),cell_to_bgcell)
cell_to_bgrcoords = lazy_map(Reindex(bgcell_to_rcoords),cell_to_bgcell)
cell_to_values = lazy_map(Reindex(bgcell_to_values),cell_to_bgcell)
cell_to_edges = lazy_map(Reindex(bgcell_to_edges),cell_to_bgcell)

old_cell_to_rcoords = lazy_map(Broadcasting(Reindex(point_to_rcoords)),cell_to_points)
old_cell_to_coords = lazy_map(Broadcasting(Reindex(point_to_coords)),cell_to_points)

new_cell_to_coords, new_cell_to_rcoords = compute_coords(
old_cell_to_coords,old_cell_to_rcoords,cell_to_bgcoords,cell_to_bgrcoords,cell_to_values,cell_to_edges
)
#@assert all(new_cell_to_rcoords .== old_cell_to_rcoords)
#@assert all(new_cell_to_coords .== old_cell_to_coords)

new_cmaps = compute_cell_maps(new_cell_to_coords,cell_to_reffes)
new_ref_cmaps = compute_cell_maps(new_cell_to_rcoords,cell_to_reffes)
return new_cell_to_coords, new_cell_to_rcoords, new_cmaps, new_ref_cmaps
end

function GridapEmbedded.LevelSetCutters._get_value_at_coords(
φh::FEFunction,model::DiscreteModel{Dc,Dp}
) where {Dc,Dp}
values = get_free_dof_values(φh)
return values
end

# Autodiff

function FESpaces._compute_cell_ids(uh,ttrian::AppendedTriangulation)
ids_a = FESpaces._compute_cell_ids(uh,ttrian.a)
ids_b = FESpaces._compute_cell_ids(uh,ttrian.b)
lazy_append(ids_a,ids_b)
end

# DifferentiableTriangulation

mutable struct DifferentiableTriangulation{Dc,Dp} <: Triangulation{Dc,Dp}
recipe :: Function
state :: Triangulation{Dc,Dp}
children :: IdDict{UInt64,Measure}
trian :: Triangulation{Dc,Dp}
state
end

function DifferentiableTriangulation(f::Function,φh)
DifferentiableTriangulation(f,f(φh),IdDict{UInt64,Measure}())
function DifferentiableTriangulation(trian::Triangulation)
DifferentiableTriangulation(trian,nothing)
end

(t::DifferentiableTriangulation)(φh) = update_trian!(t,φh)

function update_trian!(trian::DifferentiableTriangulation,φh)
trian.state = trian.recipe(φh)
for child in values(trian.children)
update_measure!(child,trian.state)
end
trian.state = compute_dual_values(trian.trian,φh)
return trian
end

Expand All @@ -103,39 +174,42 @@ function FESpaces._change_argument(
end

function FESpaces._compute_cell_ids(uh,ttrian::DifferentiableTriangulation)
FESpaces._compute_cell_ids(uh,ttrian.state)
FESpaces._compute_cell_ids(uh,ttrian.trian)
end

function Geometry.get_background_model(t::DifferentiableTriangulation)
get_background_model(t.state)
get_background_model(t.trian)
end

Geometry.get_glue(ttrian::DifferentiableTriangulation,val::Val{d}) where d = get_glue(ttrian.state,val)

# Mutable measure

mutable struct MutableMeasure <: Measure
state :: Measure
trian :: Triangulation
params
function Geometry.get_grid(t::DifferentiableTriangulation)
get_grid(t.trian)
end

function Measure(trian::DifferentiableTriangulation,args...;kwargs...)
state = Measure(trian.state,args...;kwargs...)
meas = MutableMeasure(state, trian,(args,kwargs))
push!(trian.children, objectid(meas) => meas)
return meas
function CellData.get_cell_points(ttrian::DifferentiableTriangulation)
if isnothing(ttrian.state)
return get_cell_points(ttrian.trian)
end
cell_to_coords, cell_to_rcoords, _, _ = ttrian.state
return CellPoint(cell_to_rcoords, cell_to_coords, ttrian, ReferenceDomain())
end

function update_measure!(meas::MutableMeasure,trian::Triangulation)
args, kwargs = meas.params
meas.state = Measure(trian,args...;kwargs...)
return meas
function Geometry.get_cell_map(ttrian::DifferentiableTriangulation)
if isnothing(ttrian.state)
return get_cell_map(ttrian.trian)
end
_, _, cmaps, _ = ttrian.state
return cmaps
end

function CellData.integrate(f,b::MutableMeasure)
c = integrate(f,b.state.quad)
cont = DomainContribution()
add_contribution!(cont,b.trian,c)
cont
function Geometry.get_glue(ttrian::DifferentiableTriangulation{Dc},val::Val{d}) where {Dc,d}
if isnothing(ttrian.state)
get_glue(ttrian.trian,val)
end
if d != Dc
return nothing
end
_, _, _, ref_cmaps = ttrian.state
tface_to_mface = ttrian.trian.subcells.cell_to_bgcell
tface_to_mface_map = ref_cmaps
FaceToFaceGlue(tface_to_mface,tface_to_mface_map,nothing)
end
Loading

0 comments on commit 317b754

Please sign in to comment.