Merge pull request #173 from Joroks/nghosts

Allow uniform_partition to have arbitrarily many ghost layers

Author: fverdugo

src/p_range.jl

@@ -548,12 +548,14 @@ interface.
 - `ranks`: Array containing the distribution of ranks.
 -  `np::NTuple{N}`: Number of parts per direction.
 -  `n::NTuple{N}`: Number of global indices per direction.
--  `ghost::NTuple{N}=ntuple(i->false,N)`: Use or not ghost indices per direction.
+-  `ghost::NTuple{N}=ntuple(i->false,N)`: Number of ghost indices per direction.
 -  `periodic::NTuple{N}=ntuple(i->false,N)`: Use or not periodic boundaries per direction.
 
 For convenience, one can also provide scalar inputs instead tuples
 to create 1D block partitions. In this case, the argument `np` can be omitted
-and it will be computed as `np=length(ranks)`.
+and it will be computed as `np=length(ranks)`. At the moment, it's only possible
+to use this syntax for zero (with `ghost=false`) or one (with `ghost=true`) layer(s)
+of ghost indices. If you wish to have more ghost indices, use tuples instead.
 
 # Examples
 
@@ -622,42 +624,35 @@ function block_with_constant_size(rank,np,n,ghost,periodic=map(i->false,ghost))
     p = CartesianIndices(np)[rank]
     own_ranges = map(local_range,Tuple(p),np,n)
     local_ranges = map(local_range,Tuple(p),np,n,ghost,periodic)
-    owners = map(Tuple(p),own_ranges,local_ranges) do p,or,lr
+    owners = map(Tuple(p), np, n, local_ranges) do p, np, n, lr
         myowners = zeros(Int32,length(lr))
-        for i in 1:length(lr)
-            if lr[i] in or
+        i = 1
+        for p in Iterators.cycle(1:np)
+            plr = local_range(p, np, n)
+            while mod(lr[i]-1, n)+1 in plr
                 myowners[i] = p
+                (i += 1) > length(myowners) && return myowners
             end
         end
-        if myowners[1] == 0
-            myowners[1] = p-1
-        end
-        if myowners[end] == 0
-            myowners[end] = p+1
-        end
-        myowners
-    end
-    n_ghost = 0
-    cis = CartesianIndices(map(length,local_ranges))
-    predicate(p,i,owners) = owners[i] == p
-    for ci in cis
-        flags = map(predicate,Tuple(p),Tuple(ci),owners)
-        if !all(flags)
-            n_ghost += 1
-        end
     end
+    n_local = prod(map(length, local_ranges))
+    n_own = prod(map(length, own_ranges))
+    n_ghost = n_local - n_own
+
     ghost_to_global = zeros(Int,n_ghost)
     ghost_to_owner = zeros(Int32,n_ghost)
-    n_local = prod(map(length,local_ranges))
     perm = zeros(Int32,n_local)
     i_ghost = 0
     i_own = 0
-    n_own = prod(map(length,own_ranges))
+
+    cis = CartesianIndices(map(length,local_ranges))
     lis = CircularArray(LinearIndices(n))
     local_cis = CartesianIndices(local_ranges)
-    owner_lis = CircularArray(LinearIndices(np))
+    owner_lis = LinearIndices(np)
     for (i,ci) in enumerate(cis)
-        flags = map(predicate,Tuple(p),Tuple(ci),owners)
+        flags = map(Tuple(ci), own_ranges, local_ranges) do i, or, lr
+            i in (or .- first(lr) .+ 1)
+        end
         if !all(flags)
             i_ghost += 1
             ghost_to_global[i_ghost] = lis[local_cis[i]]
@@ -809,39 +804,32 @@ function renumber_partition(partition_in;renumber_local_indices=true)
 end
 
 function local_range(p,np,n,ghost=false,periodic=false)
-    l = n ÷ np
+    l, rem = divrem(n, np)
     offset = l * (p-1)
-    rem = n % np
     if rem >= (np-p+1)
-        l = l + 1
-        offset = offset + p - (np-rem) - 1
-    end
-    start = 1+offset
-    stop = l+offset
-    if ghost && np != 1
-        if periodic || p!=1
-            start -= 1
-        end
-        if periodic || p!=np
-            stop += 1
-        end
-    end
-    start:stop
-end
-
-function boundary_owner(p,np,n,ghost=false,periodic=false)
-    start = p
-    stop = p
-    if ghost && np != 1
-        if periodic || p!=1
-            start -= 1
-        end
-        if periodic || p!=np
-            stop += 1
-        end
+        l += 1
+        offset += p - (np-rem) - 1
     end
-    (start,p,stop)
-end
+    start = 1+offset-ghost
+    stop = l+offset+ghost
+
+    periodic && return start:stop
+    return max(1, start):min(n,stop)
+end
+
+## unused
+# function boundary_owner(p,np,n,ghost=false,periodic=false)
+#     start = p
+#     stop = p
+
+#     if periodic || p!=1
+#         start -= ghost
+#     end
+#     if periodic || p!=np
+#         stop += ghost
+#     end
+#     (start,p,stop)
+# end
 
 struct VectorFromDict{Tk,Tv} <: AbstractVector{Tv}
     dict::Dict{Tk,Tv}

test/p_range_tests.jl

@@ -51,6 +51,22 @@ function p_range_tests(distribute)
    periodic = (true,true)
    uniform_partition(rank,np,n,ghost,periodic)
 
+   # uniform Cartesian partition with two layers of ghosts
+   # in the selected directions
+   # pzeros fails, if the partition is not consistent
+   np = (2,2)
+   n = (10,10)
+   ghost = (2,2)
+   uniform_partition(rank,np,n,ghost) |> pzeros
+
+   # uniform Cartesian partition with two layers of ghosts
+   # in the selected directions
+   # pzeros fails, if the partition is not consistent
+   np = (2,2)
+   n = (10,10)
+   periodic = (true,true)
+   uniform_partition(rank,np,n,ghost,periodic) |> pzeros
+
    # Custom linear partition with no ghost
    n_own = map(rank) do rank
        mod(rank,3) + 2