Skip to content

Parrays 0.5 #178

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.

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

Draft
wants to merge 4 commits into
base: parrays-0.5
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
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
33 changes: 29 additions & 4 deletions src/Assembly.jl
Original file line number Diff line number Diff line change
Expand Up @@ -225,17 +225,42 @@ function Algebra.create_from_nz(a::PVectorAllocation{<:Assembled,<:AbstractVecto
@assert false
end

function rhs_callback(values,rows)
# Based on the old _rhs_callback pattern:
# The issue is that values come from the FE space structure but need to be
# organized according to the assembly matrix row structure.
# We need to create the vector with proper index alignment.

if isa(rows, PRange)
assembly_indices = partition(rows)
else
assembly_indices = rows
end

# Create the vector using the assembly indices structure
# This ensures proper alignment with matrix rows
return PVector(values, assembly_indices)
end

function Algebra.create_from_nz(a::PVectorAllocation{<:Union{LocallyAssembled,SubAssembled}})
rows = remove_ghost(unpermute(axes(a,1)))
test_ids = partition(axes(a,1))
rows = map(remove_ghost,map(unpermute,test_ids))
values = local_views(a)
b = rhs_callback(values,rows)

# Use the same pattern as joint assembly: create FE space vector then repartition
b_fespace = PVector(values, test_ids)
b = locally_repartition(b_fespace, rows)
return b
end

function Algebra.create_from_nz(a::PVectorAllocation{<:Assembled})
rows = collect_touched_ids(a)
new_indices = collect_touched_ids(a)
values = map(ai -> ai.values, local_views(a))
b = rhs_callback(values,rows)

# Use the same pattern as joint assembly: create FE space vector then repartition
b_fespace = PVector(values, partition(axes(a,1)))
b = locally_repartition(b_fespace, new_indices)

t2 = assemble!(b)
if t2 !== nothing
wait(t2)
Expand Down
73 changes: 72 additions & 1 deletion src/FESpaces.jl
Original file line number Diff line number Diff line change
Expand Up @@ -194,7 +194,76 @@ end
const DistributedSingleFieldFEFunction{A,B,T} = DistributedCellField{A,B,DistributedFEFunctionData{T}}

function FESpaces.get_free_dof_values(uh::DistributedSingleFieldFEFunction)
uh.metadata.free_values
free_values = uh.metadata.free_values

# Check if we have periodic boundary conditions that require index structure conversion
# This is needed when the FE space uses PermutedLocalIndices but the linear system
# requires LocalIndicesWithVariableBlockSize for compatibility
if _needs_periodic_bc_fix(free_values)
return _convert_to_compatible_index_structure(free_values)
end

return free_values
end

"""
Check if the vector needs periodic BC index structure conversion
"""
function _needs_periodic_bc_fix(v::PVector)
# Check if any partition uses PermutedLocalIndices which typically indicates periodic BCs
indices = partition(axes(v, 1))
checks = map(indices) do idx
isa(idx, PartitionedArrays.PermutedLocalIndices)
end
# Use reduction to avoid scalar indexing on DebugArray
return reduce(|, checks, init=false)
end

"""
Convert a vector with PermutedLocalIndices to one with LocalIndicesWithVariableBlockSize
This is needed for periodic BC compatibility with linear system operators
"""
function _convert_to_compatible_index_structure(v::PVector)
# Get the current index partition
old_indices = partition(axes(v, 1))

# Create new indices using LocalIndicesWithVariableBlockSize
new_indices = map(old_indices) do idx
if isa(idx, PartitionedArrays.PermutedLocalIndices)
# Extract the underlying LocalIndicesWithVariableBlockSize
return idx.indices
else
return idx
end
end

# Create a new PRange with the unpermuted indices
new_prange = PRange(new_indices)

# Create a new vector with the compatible index structure
new_vector = similar(v, new_prange)

# Copy values handling the index mapping
map(partition(new_vector), partition(v), new_indices, old_indices) do new_part, old_part, new_idx, old_idx
if isa(old_idx, PartitionedArrays.PermutedLocalIndices)
# Handle permuted indices by mapping through global indices
old_loc_to_glo = local_to_global(old_idx)
new_glo_to_loc = global_to_local(new_idx)

for old_lid in 1:local_length(old_idx)
gid = old_loc_to_glo[old_lid]
new_lid = new_glo_to_loc[gid]
if new_lid > 0
new_part[new_lid] = old_part[old_lid]
end
end
else
# Direct copy if no permutation
copy!(new_part, old_part)
end
end

return new_vector
end

# Single field related
Expand Down Expand Up @@ -834,3 +903,5 @@ function FESpaces.collect_cell_matrix_and_vector(
collect_cell_matrix_and_vector(u,v,m,l,f)
end
end


12 changes: 6 additions & 6 deletions src/PArraysExtras.jl
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,8 @@ function unpermute(indices::AbstractLocalIndices)
@notimplemented
end

unpermute(prange::PRange) = PRange(map(unpermute, partition(prange)))
PArrays.remove_ghost(prange::PRange) = PRange(map(PArrays.remove_ghost, partition(prange)))
unpermute(indices::PermutedLocalIndices) = unpermute(indices.indices)
unpermute(indices::PartitionedArrays.LocalIndicesInBlockPartition) = indices
unpermute(indices::OwnAndGhostIndices) = indices
Expand All @@ -64,6 +66,10 @@ function unpermute(indices::LocalIndices)
rank = part_id(indices)
own = OwnIndices(nglobal,rank,own_to_global(indices))
ghost = GhostIndices(nglobal,ghost_to_global(indices),ghost_to_owner(indices))

# For LocalIndices, we can use the original global_to_owner mapping directly
# This should work properly with PartitionedArrays since LocalIndices
# has a complete global_to_owner mapping
OwnAndGhostIndices(own,ghost,global_to_owner(indices))
end

Expand Down Expand Up @@ -161,12 +167,6 @@ end
# end

function PArrays.remove_ghost(indices::PermutedLocalIndices)
perm = indices.perm
own_to_local = indices.own_to_local
display(own_to_local)
println(issorted(own_to_local))
display(perm)
println("-------------------------------")
remove_ghost(indices.indices)
end

Expand Down
33 changes: 32 additions & 1 deletion test/PeriodicBCsTests.jl
Original file line number Diff line number Diff line change
Expand Up @@ -29,9 +29,40 @@ function main(distribute,parts)
l(v) = ∫( v*f )dΩ
op = AffineFEOperator(a,l,U,V)

# Test 1: Basic solve(op) functionality
uh = solve(op)
eh = u - uh
@test sqrt(sum( ∫(abs2(eh))dΩ )) < 0.00122
error_solve = sqrt(sum( ∫(abs2(eh))dΩ ))
@test error_solve < 0.00122

# Test 2: LinearFESolver consistency (fix verification)
# This tests the specific fix for periodic BCs where LinearFESolver
# uses compatible vector index structures
solver = LinearFESolver(BackslashSolver())
uh_solver = solve(solver, op)
eh_solver = u - uh_solver
error_solver = sqrt(sum( ∫(abs2(eh_solver))dΩ ))

# Test 3: Direct matrix solve for comparison
A = get_matrix(op)
b = get_vector(op)
x_direct = A \ b
uh_direct = FEFunction(U, x_direct)
eh_direct = u - uh_direct
error_direct = sqrt(sum( ∫(abs2(eh_direct))dΩ ))

# Test 4: Verify all methods give consistent results (within numerical precision)
# This is the key test for the periodic BC fix - all ratios should be ≈ 1.0
ratio_solve_direct = error_solve / error_direct
ratio_solver_direct = error_solver / error_direct

@test abs(ratio_solve_direct - 1.0) < 1e-10 # solve(op) should match direct solve
@test abs(ratio_solver_direct - 1.0) < 1e-10 # LinearFESolver should match direct solve
@test abs(error_solve - error_solver) < 1e-14 # solve(op) and LinearFESolver should be identical

# Test 5: Error magnitudes should be reasonable
@test error_solver < 0.00122
@test error_direct < 0.00122

end

Expand Down