From 11ebcd418dee98f6521d08a543040751e0bf1f86 Mon Sep 17 00:00:00 2001 From: Max Date: Fri, 7 Aug 2026 09:13:50 +0200 Subject: [PATCH 1/2] Improve vec_topetsc and mat_topetsc --- feectools/linalg/topetsc.py | 102 +++++++++++++++++++++++----------- feectools/linalg/utilities.py | 87 ++++++++++++++++++----------- 2 files changed, 125 insertions(+), 64 deletions(-) diff --git a/feectools/linalg/topetsc.py b/feectools/linalg/topetsc.py index 47e5d0e77..2e8ad8b20 100644 --- a/feectools/linalg/topetsc.py +++ b/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 diff --git a/feectools/linalg/utilities.py b/feectools/linalg/utilities.py index eebb061fd..e08d97649 100644 --- a/feectools/linalg/utilities.py +++ b/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,7 +99,7 @@ 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 @@ -107,26 +107,28 @@ def petsc_to_psydac(x, Xh, out=None): 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): """ From c10542411b89884d9a2167d0f71f49d00f3006e1 Mon Sep 17 00:00:00 2001 From: Max Date: Sun, 9 Aug 2026 09:13:24 +0200 Subject: [PATCH 2/2] Update version number to 0.1.11 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 8e22c4993..ac462afe6 100644 --- a/pyproject.toml +++ 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"