Coverage for dragonfly_sim/utils/petsc_utils.py: 100%

1import numpy as np 

2 

3 

4def set_petsc_vec_array(petsc_vec, np_array): 

5 """Set the locally-owned entries of a PETSc Vec from a numpy array. 

6 

7 Works for both sequential and distributed (MPI) Vecs. For a 

8 distributed Vec, ``np_array`` must have exactly as many elements as 

9 the local ownership range of ``petsc_vec`` on this rank, and values 

10 are placed at the correct *global* indices — not always starting at 0. 

11 

12 Parameters 

13 ---------- 

14 petsc_vec : PETSc.Vec 

15 Target Vec. May be sequential or MPI-distributed. 

16 np_array : np.ndarray, 1-D 

17 Values to write. Length must equal the local ownership size. 

18 """ 

19 # getOwnershipRange() returns the half-open global index interval 

20 # [rstart, rend) owned by this rank. For a sequential Vec this is 

21 # always [0, n), so the call is safe in both serial and parallel. 

22 rstart, rend = petsc_vec.getOwnershipRange() 

23 global_indices = np.arange(rstart, rend, dtype=np.int32) 

24 petsc_vec.setValues(global_indices, np_array) 

25 petsc_vec.assemble()