Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
# If any of the inputs were trimmed due to early splits,
# trim the others too.
# In very hairy cases this can take multiple passes.
# TODO: can we optimize this, or code it more elegantly?
max_passes_left = 10
while max_passes_left > 0:
this_chunk_end = min([x.end for x in inputs.values()]
+ [this_chunk_end])
if len(set([x.end for x in inputs.values()])) <= 1:
break
for d in self.depends_on:
inputs[d], back_to_buffer = \
inputs[d].split(
t=this_chunk_end,
allow_early_split=True)
self.input_buffer[d] = strax.Chunk.concatenate(
[back_to_buffer, self.input_buffer[d]])
max_passes_left -= 1
else:
raise RuntimeError(
f"{self} was unable to get time-consistent "
f"inputs after ten passess. Inputs: \n{inputs}\n"
f"Input buffer:\n{self.input_buffer}")
# Merge inputs of the same kind
inputs_merged = {
kind: strax.Chunk.merge([inputs[d] for d in deps_of_kind])
for kind, deps_of_kind in self.dependencies_by_kind().items()}
# Submit the computation
# print(f"{self} calling with {inputs_merged}")
if self.parallel and executor is not None:
def do_compute(self, chunk_i=None, **kwargs):
if not len(kwargs):
raise RuntimeError("OverlapWindowPlugin must have a dependency")
# Add cached inputs to compute arguments
for k, v in kwargs.items():
if len(self.cached_input):
kwargs[k] = strax.Chunk.concatenate(
[self.cached_input[k], v])
# Compute new results
result = super().do_compute(chunk_i=chunk_i, **kwargs)
# Throw away results we already sent out
_, result = result.split(t=self.sent_until,
allow_early_split=False)
# When does this batch of inputs end?
ends = [v.end for v in kwargs.values()]
if not len(set(ends)) == 1:
raise RuntimeError(
f"OverlapWindowPlugin got incongruent inputs: {kwargs}")
end = ends[0]
def save_from(self, source: typing.Generator, rechunk=True, executor=None):
"""Iterate over source and save the results under key
along with metadata
"""
pending = []
exhausted = False
chunk_i = 0
try:
while not exhausted:
chunk = None
try:
if rechunk and self.allow_rechunk:
while chunk is None or chunk.data.nbytes < 1e8:
chunk = strax.Chunk.concatenate(
[chunk, next(source)])
else:
chunk = next(source)
except StopIteration:
exhausted = True
if chunk is None:
break
new_f = self.save(chunk=chunk,
chunk_i=chunk_i, executor=executor)
pending = [f for f in pending if not f.done()]
if new_f is not None:
pending += [new_f]
chunk_i += 1
def _squash_buffer(self):
"""Squash the buffer into a single chunk using np.concatenate"""
if len(self.buffer) > 1:
self.buffer = [strax.Chunk.concatenate(self.buffer)]
# Sanity check on buffer_items
if not len(self.buffer):
assert self.buffer_items == 0
else:
assert self.buffer_items == len(self.buffer[0])
def _fetch_chunk(self, d, iters, check_end_not_before=None):
"""Add a chunk of the datatype d to the input buffer.
Return True if this succeeded, False if the source is exhausted.
:param d: data type to fetch
:param iters: iterators that produce data
:param check_end_not_before: Raise a runtimeError if the source
is exhausted, but the input buffer ends before this time.
"""
try:
# print(f"Fetching {d} in {self}, hope to see {hope_to_see}")
self.input_buffer[d] = strax.Chunk.concatenate(
[self.input_buffer[d], next(iters[d])])
# print(f"Fetched {d} in {self}, "
# f"now have {self.input_buffer[d]}")
return True
except StopIteration:
# print(f"Got StopIteration while fetching for {d} in {self}")
if (check_end_not_before is not None
and self.input_buffer[d].end < check_end_not_before):
raise RuntimeError(
f"Tried to get data until {check_end_not_before}, but {d} "
f"ended prematurely at {self.input_buffer[d].end}")
return False