Improve PETSc integration

feectools/linalg/topetsc.py

@@ -332,6 +332,48 @@ def get_npts_per_block(V : VectorSpace) -> list:
     return npts_local_per_block
 
 
+def _own_process_index(cart):
+    """ Flat process index (as used by get_npts_per_block) of the *current* process
+    within its Cartesian process grid, or 0 if running serially.
+    """
+    if not cart.comm:
+        return 0
+
+    nprocs = cart.nprocs
+    coords = cart.coords
+    if len(nprocs) == 1:
+        return coords[0]
+    elif len(nprocs) == 2:
+        return coords[1] + coords[0] * nprocs[1]
+    elif len(nprocs) == 3:
+        return coords[2] + coords[1] * nprocs[2] + coords[0] * nprocs[1] * nprocs[2]
+    else:
+        raise NotImplementedError("Cannot handle more than 3 dimensions.")
+
+
+def _own_process_index_shift_and_shape(V, b, cart):
+    """ For block `b`, the PETSc global-index offset of the *current* process' data,
+    and this process' local shape for that block.
+
+    Since a call to `vec_topetsc`/`mat_topetsc` only ever needs the PETSc index of data
+    *local to the calling process*, the owning process of every such index is always the
+    calling process itself -- there is no need to search for it index-by-index (as
+    `psydac_to_petsc_global` does), which is what makes this vectorizable.
+    """
+    npts_local_per_block_per_process = xp.array(get_npts_per_block(V)) # indexed [block, process, dim]
+    local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) # indexed [block, process]
+
+    proc_index = _own_process_index(cart)
+
+    index_shift = int(
+        xp.sum(local_sizes_per_block_per_process[:, :proc_index])
+        + xp.sum(local_sizes_per_block_per_process[:b, proc_index])
+    )
+    own_shape = npts_local_per_block_per_process[b, proc_index]
+
+    return index_shift, own_shape
+
+
 def vec_topetsc(vec):
     """ Convert vector from Psydac format to a PETSc.Vec object.
 
@@ -370,7 +412,7 @@ def vec_topetsc(vec):
     # Sum over the blocks to get the total local size
     localsize = xp.sum(xp.prod(npts_local, axis=1))
 
-    gvec  = PETSc.Vec().create(comm=carts[0].global_comm)    
+    gvec  = PETSc.Vec().create(comm=carts[0].global_comm)
 
     # Set global and local size:
     gvec.setSizes(size=(localsize, globalsize))
@@ -383,45 +425,43 @@ def vec_topetsc(vec):
 
     vec_block = vec
 
-    for b in range(n_blocks): 
+    for b in range(n_blocks):
         if isinstance(vec, BlockVector):
             vec_block = vec.blocks[b]
-
-        s = carts[b].starts
+
         ghost_size = [pi*mi for pi,mi in zip(carts[b].pads, carts[b].shifts)]
+        index_shift, own_shape = _own_process_index_shift_and_shape(vec.space, b, carts[b])
+
+        local_slices = tuple(slice(gs, gs + n) for gs, n in zip(ghost_size, npts_local[b]))
+        local_data = vec_block._data[local_slices]
 
         if ndims[b] == 1:
-            for i1 in range(npts_local[b][0]):
-                value = vec_block._data[i1 + ghost_size[0]]
-                if value != 0:
-                    i1_n = s[0] + i1
-                    i_g = psydac_to_petsc_global(vec.space, (b,), (i1_n,))
-                    petsc_indices.append(i_g)
-                    petsc_data.append(value)        
+            (nz1,) = xp.nonzero(local_data)
+            i_g = index_shift + nz1
+            petsc_indices.append(i_g)
+            petsc_data.append(local_data[nz1])
 
         elif ndims[b] == 2:
-            for i1 in range(npts_local[b][0]):
-                for i2 in range(npts_local[b][1]):
-                    value = vec_block._data[i1 + ghost_size[0], i2 + ghost_size[1]]
-                    if value != 0:
-                        i1_n = s[0] + i1
-                        i2_n = s[1] + i2                    
-                        i_g = psydac_to_petsc_global(vec.space, (b,), (i1_n, i2_n))
-                        petsc_indices.append(i_g)
-                        petsc_data.append(value)
+            nz1, nz2 = xp.nonzero(local_data)
+            i_g = index_shift + nz2 + nz1 * own_shape[1]
+            petsc_indices.append(i_g)
+            petsc_data.append(local_data[nz1, nz2])
 
         elif ndims[b] == 3:
-            for i1 in xp.arange(npts_local[b][0]):             
-                for i2 in xp.arange(npts_local[b][1]):
-                    for i3 in xp.arange(npts_local[b][2]):
-                        value = vec_block._data[i1 + ghost_size[0], i2 + ghost_size[1], i3 + ghost_size[2]]
-                        if value != 0:
-                            i1_n = s[0] + i1
-                            i2_n = s[1] + i2
-                            i3_n = s[2] + i3    
-                            i_g = psydac_to_petsc_global(vec.space, (b,), (i1_n, i2_n, i3_n))                    
-                            petsc_indices.append(i_g)
-                            petsc_data.append(value)        
+            nz1, nz2, nz3 = xp.nonzero(local_data)
+            i_g = index_shift + nz3 + nz2 * own_shape[2] + nz1 * own_shape[1] * own_shape[2]
+            petsc_indices.append(i_g)
+            petsc_data.append(local_data[nz1, nz2, nz3])
+
+        else:
+            raise NotImplementedError("Cannot handle more than 3 dimensions.")
+
+    petsc_indices = xp.concatenate(petsc_indices) if petsc_indices else xp.array([], dtype=PETSc.IntType)
+    petsc_indices = petsc_indices.astype(PETSc.IntType)
+    petsc_data = xp.concatenate(petsc_data) if petsc_data else xp.array([])
+    # if PETSc was built with a real scalar type but the source data is complex,
+    # drop the imaginary part (silently, matching the previous per-scalar behavior)
+    petsc_data = petsc_data.astype(PETSc.ScalarType)
 
     # Set the values. The values are stored in a cache memory.
     gvec.setValues(petsc_indices, petsc_data, addv=PETSc.InsertMode.ADD_VALUES) #The addition mode the values is necessary when periodic BC

feectools/linalg/utilities.py

@@ -6,7 +6,7 @@
 from feectools.linalg.basic   import Vector
 from feectools.linalg.stencil import StencilVector, StencilVectorSpace
 from feectools.linalg.block   import BlockVector, BlockVectorSpace
-from feectools.linalg.topetsc import petsc_local_to_psydac, get_npts_per_block
+from feectools.linalg.topetsc import get_npts_local
 
 __all__ = (
     'array_to_psydac',
@@ -99,34 +99,36 @@ def petsc_to_psydac(x, Xh, out=None):
     if isinstance(Xh, BlockVectorSpace):
         if any([isinstance(Xh.spaces[b], BlockVectorSpace) for b in range(len(Xh.spaces))]):
             raise NotImplementedError('Block of blocks not implemented.')
-        
+
         if out is not None:
             assert isinstance(out, BlockVector)
             assert out.space is Xh
             u = out
         else:
             u = BlockVector(Xh)
 
-        comm       = x.comm
         dtype      = Xh._dtype
         localsize, globalsize = x.getSizes()
         assert globalsize == u.shape[0], 'Sizes of global vectors do not match'
 
-        # Find shift for process k:
-        # ..get number of points for each block, each process and each dimension:
-        npts_local_per_block_per_process = xp.array(get_npts_per_block(Xh)) #indexed [b,k,d] for block b and process k and dimension d
-        # ..get local sizes for each block and each process:
-        local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) #indexed [b,k] for block b and process k
-        # ..sum the sizes over all the blocks and the previous processes:
-        index_shift = 0 + xp.sum(local_sizes_per_block_per_process[:,:comm.Get_rank()], dtype=int) #global variable
-
-        for local_petsc_index in range(localsize):
-            block_index, psydac_index = petsc_local_to_psydac(Xh, local_petsc_index)
-            # Get value of local PETSc vector passing the global PETSc index
-            value = x.getValue(local_petsc_index + index_shift) 
-            if value != 0:
-                u[block_index[0]]._data[psydac_index] = value if dtype is complex else value.real # PETSc always handles dtype specified in the installation configuration
-
+        # local PETSc data (this process only), ordered block-by-block (see vec_topetsc)
+        values = x.getArray(readonly=True)
+
+        # this process' local shape, per block
+        npts_local_per_block = get_npts_local(Xh) #indexed [b][d]
+
+        offset = 0
+        for bb in range(Xh.n_blocks):
+            npts_local = npts_local_per_block[bb]
+            n = int(xp.prod(npts_local))
+            block_values = values[offset:offset + n]
+            offset += n
+
+            block_indices = _local_petsc_indices_to_data_indices(
+                xp.arange(n), npts_local, Xh.spaces[bb].pads, Xh.spaces[bb].shifts,
+            )
+            u[bb]._data[block_indices] = block_values if dtype is complex else block_values.real
+
     elif isinstance(Xh, StencilVectorSpace):
 
         if out is not None:
@@ -136,25 +138,16 @@ def petsc_to_psydac(x, Xh, out=None):
         else:
             u = StencilVector(Xh)
 
-        comm       = x.comm
         dtype      = Xh.dtype
         localsize, globalsize = x.getSizes()
         assert globalsize == u.shape[0], 'Sizes of global vectors do not match'
 
-        # Find shift for process k:
-        # ..get number of points for each process and each dimension:
-        npts_local_per_block_per_process = xp.array(get_npts_per_block(Xh))[0] #indexed [k,d] for process k and dimension d
-        # ..get local sizes for each process:
-        local_sizes_per_block_per_process = xp.prod(npts_local_per_block_per_process, axis=-1) #indexed [k] for process k
-        # ..sum the sizes over all the previous processes:
-        index_shift = 0 + xp.sum(local_sizes_per_block_per_process[:comm.Get_rank()], dtype=int) #global variable
-
-        for local_petsc_index in range(localsize):
-            block_index, psydac_index = petsc_local_to_psydac(Xh, local_petsc_index) 
-            # Get value of local PETSc vector passing the global PETSc index
-            value = x.getValue(local_petsc_index + index_shift)
-            if value != 0:
-                u._data[psydac_index] = value if dtype is complex else value.real # PETSc always handles dtype specified in the installation configuration            
+        # local PETSc data (this process only)
+        values = x.getArray(readonly=True)
+        npts_local = get_npts_local(Xh)[0]
+
+        indices = _local_petsc_indices_to_data_indices(xp.arange(localsize), npts_local, Xh.pads, Xh.shifts)
+        u._data[indices] = values if dtype is complex else values.real
 
     else:
         raise ValueError('Xh must be a StencilVectorSpace or a BlockVectorSpace')
@@ -163,6 +156,34 @@ def petsc_to_psydac(x, Xh, out=None):
 
     return u
 
+
+def _local_petsc_indices_to_data_indices(local_petsc_indices, npts_local, pads, shifts):
+    """ Vectorized equivalent of calling `petsc_local_to_psydac` for every index in
+    `local_petsc_indices` (all belonging to the same block), returning a tuple of
+    integer arrays directly usable to index a StencilVector's `._data` array.
+    """
+    ndim = len(npts_local)
+
+    if ndim == 1:
+        i0 = local_petsc_indices + pads[0] * shifts[0]
+        return (i0,)
+
+    elif ndim == 2:
+        i0 = local_petsc_indices // npts_local[1] + pads[0] * shifts[0]
+        i1 = local_petsc_indices % npts_local[1] + pads[1] * shifts[1]
+        return (i0, i1)
+
+    elif ndim == 3:
+        n1n2 = npts_local[1] * npts_local[2]
+        i0 = local_petsc_indices // n1n2 + pads[0] * shifts[0]
+        rem = local_petsc_indices % n1n2
+        i1 = rem // npts_local[2] + pads[1] * shifts[1]
+        i2 = rem % npts_local[2] + pads[2] * shifts[2]
+        return (i0, i1, i2)
+
+    else:
+        raise NotImplementedError("Cannot handle more than 3 dimensions.")
+
 #==============================================================================
 def _sym_ortho(a, b):
     """

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.1.9"
+version         = "0.1.11"
 description     = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
 readme          = "README.md"
 requires-python = ">= 3.10"