non_conforming_octrees_wip.jl

using P4est_wrapper
using GridapP4est
using Gridap
using PartitionedArrays
using GridapDistributed
using MPI
using Gridap.FESpaces

# Generate a local numbering of vertices that includes hanging vertices 
# Generate a local numbering of faces out of the one generated by vertices (automatic? to confirm)

# Establish the correspondence among local numbering of vertices and p4est global numbering 
# Establish the correspondence among local numbering of faces and p4est global numbering 

# Generate a global numbering of (regular,hanging) vertices?
# Generate a global numbering of (regular,hanging) faces?


## Better to use a C-enum. But I did not use it in order to keep the Julia
## version of this C example as simple as possible
const nothing_flag = Cint(0)
const refine_flag = Cint(1)

## Refine those cells with even identifier    (0,2,4,6,8,...)
## Leave untouched cells with odd identifier  (1,3,5,7,9,...)
function allocate_and_set_refinement_and_coarsening_flags(forest_ptr::Ptr{p4est_t})
  forest = forest_ptr[]
  tree = p4est_tree_array_index(forest.trees, 0)[]
  return [i != 1 ? nothing_flag : refine_flag for i = 1:tree.quadrants.elem_count]
end

## Global variable which is updated across calls to init_fn_callback_2d
current_quadrant_index = Cint(0)
## Global variable which is updated across calls to refine_replace_callback_2d
num_calls = Cint(0)

# This C callback function is called once per quadtree quadrant. Here we are assuming
# that p4est->user_pointer has been set prior to the first call to this call
# back function to an array of ints with as many entries as forest quadrants. This call back function
# initializes the quadrant->p.user_data void * pointer of all quadrants such that it
# points to the corresponding entry in the global array mentioned in the previous sentence.
function init_fn_callback_2d(forest_ptr::Ptr{p4est_t},
  which_tree::p4est_topidx_t,
  quadrant_ptr::Ptr{p4est_quadrant_t})
  @assert which_tree == 0
  # Extract a reference to the first (and uniquely allowed) tree
  forest = forest_ptr[]
  tree = p4est_tree_array_index(forest.trees, 0)[]
  quadrant = quadrant_ptr[]
  q = P4est_wrapper.p4est_quadrant_array_index(tree.quadrants, current_quadrant_index)
  @assert p4est_quadrant_compare(q, quadrant_ptr) == 0
  user_data = unsafe_wrap(Array, Ptr{Cint}(forest.user_pointer), current_quadrant_index + 1)[current_quadrant_index+1]
  unsafe_store!(Ptr{Cint}(quadrant.p.user_data), user_data, 1)
  global current_quadrant_index = (current_quadrant_index + 1) % (tree.quadrants.elem_count)
  return nothing
end

const init_fn_callback_2d_c = @cfunction(init_fn_callback_2d, Cvoid, (Ptr{p4est_t}, p4est_topidx_t, Ptr{p4est_quadrant_t}))


function refine_callback_2d(::Ptr{p4est_t},
  which_tree::p4est_topidx_t,
  quadrant_ptr::Ptr{p4est_quadrant_t})
  @assert which_tree == 0
  quadrant = quadrant_ptr[]
  return Cint(unsafe_wrap(Array, Ptr{Cint}(quadrant.p.user_data), 1)[] == refine_flag)
end

const refine_callback_2d_c = @cfunction(refine_callback_2d, Cint, (Ptr{p4est_t}, p4est_topidx_t, Ptr{p4est_quadrant_t}))

MPI.Init()
parts = get_part_ids(MPIBackend(), 1)
# run(parts, (1, 1))
# MPI.Finalize()

# This is for debuging
Dc=2
domain = (0, 1, 0, 1)
subdomains = (1, 1)
coarse_model = CartesianDiscreteModel(domain, subdomains)
model = OctreeDistributedDiscreteModel(parts, coarse_model, 1)

# TO-DO: refine and coarsening flags should be an input argument, instead of being hard-coded
# function adapt_non_conforming_work_in_progress(model::OctreeDistributedDiscreteModel{Dc,Dp}) where {Dc,Dp}
# Copy and refine input p4est
ptr_new_pXest = GridapP4est.pXest_copy(Val{Dc}, model.ptr_pXest)

user_data = allocate_and_set_refinement_and_coarsening_flags(ptr_new_pXest)
p4est_reset_data(ptr_new_pXest, Cint(sizeof(Cint)), init_fn_callback_2d_c, pointer(user_data))
p4est_refine_ext(ptr_new_pXest, 0, -1, refine_callback_2d_c, C_NULL, C_NULL)
p4est_partition(ptr_new_pXest, 1, C_NULL)

p4est_vtk_write_file(ptr_new_pXest, C_NULL, string("adapted_forest"))

ptr_pXest_ghost = GridapP4est.setup_pXest_ghost(Val{Dc}, ptr_new_pXest)
ptr_pXest_lnodes = GridapP4est.p4est_lnodes_new(ptr_new_pXest, ptr_pXest_ghost, -2)

dmodel,non_conforming_glue=setup_non_conforming_distributed_discrete_model(Val{2},
                                                       parts,
                                                       coarse_model,
                                                       model.ptr_pXest_connectivity,
                                                       ptr_new_pXest,
                                                       ptr_pXest_ghost,
                                                       ptr_pXest_lnodes)

# FE Spaces
order=1
reffe = ReferenceFE(lagrangian,Float64,order)
V = TestFESpace(dmodel,reffe,dirichlet_tags="boundary")
U = TrialFESpace(V)   

function _h_refined_reffe(reffe::Tuple{<:Lagrangian,Any,Any})
  (reffe[1],(reffe[2][1],2*reffe[2][2]),reffe[3])
end

basis, reffe_args,reffe_kwargs = reffe
cell_reffe = ReferenceFE(QUAD,basis,reffe_args...;reffe_kwargs...)
reffe_cell = cell_reffe

h_refined_reffe = _h_refined_reffe(reffe)
basis, reffe_args,reffe_kwargs = h_refined_reffe
cell_reffe_h_refined = ReferenceFE(QUAD,basis,reffe_args...;reffe_kwargs...)
reffe_cell_h_refined = cell_reffe_h_refined

dof_basis_h_refined = Gridap.CellData.get_dof_basis(reffe_cell_h_refined)

coarse_shape_funs=Gridap.ReferenceFEs.get_shapefuns(reffe_cell)
ref_constraints=evaluate(dof_basis_h_refined,coarse_shape_funs)

rr = Gridap.Adaptivity.RedRefinementRule(QUAD)
face_subface_ldof_to_cell_ldof = Gridap.Adaptivity.coarse_nodes_above_fine_nodes(rr,(2*order-1,2*order-1),order)

function generate_constraints(dmodel, 
                              V,
                              reffe,
                              non_conforming_glue,
                              ref_constraints, 
                              face_subface_ldof_to_cell_ldof)
  gridap_cells_vertices,
  num_regular_vertices, num_hanging_vertices,
  hanging_vertices_owner_cell_and_lface,
  gridap_cells_faces,
  num_regular_faces, num_hanging_faces,
  hanging_faces_owner_cell_and_lface = non_conforming_glue

  sDOF_to_dof, sDOF_to_dofs, sDOF_to_coeffs = map_parts(gridap_cells_vertices,
            num_regular_vertices,
            num_hanging_vertices,
            hanging_vertices_owner_cell_and_lface,
            gridap_cells_faces,
            num_regular_faces,
            num_hanging_faces,
            hanging_faces_owner_cell_and_lface, model.dmodel.models, V.spaces) do gridap_cells_vertices,
                                                   num_regular_vertices, num_hanging_vertices,
                                                   hanging_vertices_owner_cell_and_lface,
                                                   gridap_cells_faces,
                                                   num_regular_faces, num_hanging_faces,
                                                   hanging_faces_owner_cell_and_lface,
                                                   model, V
   
      # Locate for each hanging vertex a cell to which it belongs 
      # and local position within that cell 
      hanging_vertices_to_cell    = Vector{Int}(undef, num_hanging_vertices) 
      hanging_vertices_to_lvertex = Vector{Int}(undef, num_hanging_vertices)
      for cell=1:length(gridap_cells_vertices)
        s=gridap_cells_vertices.ptrs[cell]
        e=gridap_cells_vertices.ptrs[cell+1]
        l=e-s
        for j=1:l
          vid=gridap_cells_vertices.data[s+j-1]
          if vid>num_regular_vertices 
            vid_hanging=vid-num_regular_vertices
            hanging_vertices_to_cell[vid_hanging]=cell
            hanging_vertices_to_lvertex[vid_hanging]=j
          end 
        end 
      end

      basis, reffe_args,reffe_kwargs = reffe
      cell_reffe = ReferenceFE(QUAD,basis,reffe_args...;reffe_kwargs...)
      reffe_cell = cell_reffe

      cell_dof_ids  = get_cell_dof_ids(V)
      face_own_dofs = Gridap.ReferenceFEs.get_face_own_dofs(reffe_cell)
      face_dofs     = Gridap.ReferenceFEs.get_face_dofs(reffe_cell)


      ptrs=Vector{Int}(undef, num_hanging_vertices+1)
      ptrs[1]=1
      for vid_hanging=1:num_hanging_vertices
        (_,ocell_lface)=hanging_vertices_owner_cell_and_lface[vid_hanging]
        ptrs[vid_hanging+1] = ptrs[vid_hanging] + length(face_dofs[ocell_lface])
      end 
      data_owner_face_dofs=Vector{Int}(undef, ptrs[num_hanging_vertices+1]-1)
      data_owner_face_ldofs=Vector{Int}(undef, ptrs[num_hanging_vertices+1]-1)
      for vid_hanging=1:num_hanging_vertices
        (ocell,ocell_lface)=hanging_vertices_owner_cell_and_lface[vid_hanging]
        s=ptrs[vid_hanging]
        e=ptrs[vid_hanging+1]-1
        for (j,ldof) in enumerate(face_dofs[ocell_lface])
          data_owner_face_dofs[s+j-1]=cell_dof_ids[ocell][ldof]
          data_owner_face_ldofs[s+j-1]=ldof
        end
      end 
      hanging_vertices_owner_face_dofs  = Gridap.Arrays.Table(data_owner_face_dofs , ptrs)
      hanging_vertices_owner_face_ldofs = Gridap.Arrays.Table(data_owner_face_ldofs, ptrs)

      sDOF_to_dof    = Int[]
      sDOF_to_dofs   = Vector{Int}[]
      sDOF_to_coeffs = Vector{Float64}[]

      basis, reffe_args,reffe_kwargs = reffe
      face_reffe = ReferenceFE(SEGMENT,basis,reffe_args...;reffe_kwargs...)
      Dc = 2
      hanging_lvertex_within_first_subface = 2^(Dc-1)
      subface_own_dofs = Gridap.ReferenceFEs.get_face_own_dofs(face_reffe)
      subface_own_dofs = subface_own_dofs[hanging_lvertex_within_first_subface]
      num_faces = 2^Dc

      for vid_hanging=1:num_hanging_vertices
         # Go over the dofs of hanging_vertices_to_cell[vid_hanging]
         # on hanging_vertices_to_lvertex lvertx
         cell=hanging_vertices_to_cell[vid_hanging]
         lvertex=hanging_vertices_to_lvertex[vid_hanging]
         (_,ocell_lface)=hanging_vertices_owner_cell_and_lface[vid_hanging]
         ocell_lface = ocell_lface - num_faces
         for ((ldof,dof),ldof_subface) in zip(enumerate(face_own_dofs[lvertex]),subface_own_dofs)
           push!(sDOF_to_dof,cell_dof_ids[cell][dof])
           push!(sDOF_to_dofs,hanging_vertices_owner_face_dofs[vid_hanging])
           coeffs=Vector{Float64}(undef,length(hanging_vertices_owner_face_ldofs[vid_hanging]))
           for (i,ldof_coarse) in enumerate(hanging_vertices_owner_face_ldofs[vid_hanging])
            coeffs[i]=ref_constraints[face_subface_ldof_to_cell_ldof[ocell_lface][1][ldof_subface],ldof_coarse]
           end
           push!(sDOF_to_coeffs,coeffs) 
         end 
      end
      # TO-DO: the tables can be generated more efficiently
      sDOF_to_dof, Gridap.Arrays.Table(sDOF_to_dofs), Gridap.Arrays.Table(sDOF_to_coeffs)
    end 
end

sDOF_to_dof, sDOF_to_dofs,sDOF_to_coeffs=
    generate_constraints(model.dmodel, V, reffe, 
                                  non_conforming_glue, ref_constraints, face_subface_ldof_to_cell_ldof)
println(sDOF_to_dof)
println(sDOF_to_dofs)
println(sDOF_to_coeffs)

# Define manufactured functions
u(x) = x[1]
f(x) = 0.0

map_parts(dmodel.models,V.spaces,U.spaces,sDOF_to_dof,sDOF_to_dofs,sDOF_to_coeffs) do model,V,U,sDOF_to_dof,sDOF_to_dofs,sDOF_to_coeffs
  Vc = FESpaceWithLinearConstraints(
    sDOF_to_dof,
    sDOF_to_dofs,
    sDOF_to_coeffs,
    V)
  Uc = TrialFESpace(Vc,u)

  # Define integration mesh and quadrature
  degree = 2
  Ω = Triangulation(model)
  dΩ = Measure(Ω,degree)

  a(u,v) = ∫( ∇(v)⊙∇(u) )*dΩ
  b(v) = ∫(v*f)*dΩ

  # op = AffineFEOperator(a,b,U,V0)
  op = AffineFEOperator(a,b,Uc,Vc)
  uh = solve(op)

  # Define exact solution and error
  e = u - uh

  # Compute errors
  el2 = sqrt(sum( ∫( e*e )*dΩ ))
  eh1 = sqrt(sum( ∫( e*e + ∇(e)⋅∇(e) )*dΩ ))

  tol=1e-8
  @assert el2 < tol
  @assert eh1 < tol
end