Aquila-Consortium/csiborgtools
0
0
Fork 0
mirror of https://github.com/Richard-Sti/csiborgtools.git synced 2024-12-22 23:28:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

add nparts_to_start_ind func

Browse source
This commit is contained in:
rstiskalek 2022-10-10 18:08:49 +01:00
parent b119d21acf
commit 075536b839
1 changed files with 18 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

21
galomatch/io/readsim.py
View file

@ -119,6 +119,23 @@ def read_sp(dtype, partfile):
raise TypeError("Unexpected dtype `{}`.".format(dtype)) raise TypeError("Unexpected dtype `{}`.".format(dtype))
def nparts_to_start_ind(nparts):
"""
Convert `nparts` array to starting indices in a pre-allocated array for looping over the CPU number.
Parameters
----------
nparts : 1-dimensional array
Number of parts assosiated with each CPU.
Returns
-------
start_ind : 1-dimensional array
The starting indices calculated as a cumulative sum starting at 0.
"""
return numpy.hstack([[0], numpy.cumsum(nparts[:-1])])
def read_particle(pars_extract, n, simpath, verbose=True): def read_particle(pars_extract, n, simpath, verbose=True):
""" """
Read particle files of a simulation at a given snapshot and return Read particle files of a simulation at a given snapshot and return
@ -163,9 +180,7 @@ def read_particle(pars_extract, n, simpath, verbose=True):
"formats": [forder[fnames.index(p)][1] for p in pars_extract]} "formats": [forder[fnames.index(p)][1] for p in pars_extract]}
# Allocate the output structured array # Allocate the output structured array
out = numpy.full(npart_tot, numpy.nan, dtype) out = numpy.full(npart_tot, numpy.nan, dtype)
# Loop indices start_ind = nparts_to_start_ind((nparts))
start_ind = numpy.zeros(ncpu, dtype=int)
start_ind[1:] = numpy.cumsum(nparts)[:-1]
iters = tqdm(range(ncpu)) if verbose else range(ncpu) iters = tqdm(range(ncpu)) if verbose else range(ncpu)
for cpu in iters: for cpu in iters:
i = start_ind[cpu] i = start_ind[cpu]