def loadlocalmaps(infiles, windower, *args, **kwargs):
'''
Invoke findmaps(*args, **kwargs) to identify a map from file names for
WaveformMap serializations to desired keys for that file, then load
each file and extract the subset of the contained WaveformMap
corresponding to those keys.
If window is not None, it should be a callable that will be applied to
each Waveform before it is added to the map.
'''
wmap = WaveformMap()
for f, pairs in findmaps(infiles, *args, **kwargs).items():
# Define a filter to only load locally assigned keys
for key, wave in WaveformMap.generate(f):
if key not in pairs: continue
if windower: wave = windower(wave)
wmap[key] = wave
return wmap
def procmessages(sendreqs, recvreqs, recvbufs):
'''
Enter a loop to process incoming messages and close out pending sends,
yielding (t, r) pairs and Waveform records as they are received.
The arguments sendreqs and recvreqs are, respectively, lists send and
receive requests as prepared by postmessages. The argument recvbufs is
a map from source ranks to lists of BytesIO buffers that will be
populated with the incoming messages associated with recvreqs.
No action is taken when send requests are ready, except to wait for
their completion.
'''
# Track the number of receive requests to differentiate sends and receives
nrecvs = len(recvreqs)
# Lump all requests together for processing
requests = recvreqs + sendreqs
# Begin processing messages
status = MPI.Status()
while True:
# Wait until a message can be processed
idx = MPI.Request.Waitany(requests, status)
if idx == MPI.UNDEFINED: break
# Figure out the rank, tag and size of this message
tag = status.tag
if 0 <= idx < nrecvs:
# Parse the incoming WaveformMap stream
bstr = recvbufs[status.source][tag]
bstr.seek(0)
# Yield the keys and waveforms in turn
yield from WaveformMap.generate(bstr)
# Free buffer by closing the stream
bstr.close()
elif idx < 0 or idx >= len(requests):
raise ValueError(f'Unexpected MPI request index {idx}')
def findmaps(infiles, start=0, stride=1):
'''
Parse all of the WaveformMap instances encapsulated in infiles (a list
of files or globs) to identify for each file a set of all keys (t, r)
in that file.
If (start, stride) is other than (0, 1), the keys in each file will be
pared to only sorted(keys)[start::stride].
A map from file names to the (optionally strided) set of pairs for each
file is returned.
'''
pairmaps = {}
for f in infiles:
# Build the key generator
keys = WaveformMap.generate(f, keys_only=True)
if (start, stride) != (0, 1):
# Sort for striding, if desired
keys = sorted(keys)[start::stride]
# Convert to set and store if nonempty
keys = set(keys)
if keys: pairmaps[f] = keys
return pairmaps