def test_execution_order():
L = []
for i in range(5):
s = Stream()
b = s.pluck(1)
a = s.pluck(0)
l = a.combine_latest(b, emit_on=a).sink_to_list()
z = [(1, "red"), (2, "blue"), (3, "green")]
for zz in z:
s.emit(zz)
L.append((l, ))
for ll in L:
assert ll == L[0]
L2 = []
for i in range(5):
s = Stream()
a = s.pluck(0)
b = s.pluck(1)
l = a.combine_latest(b, emit_on=a).sink_to_list()
z = [(1, "red"), (2, "blue"), (3, "green")]
for zz in z:
s.emit(zz)
L2.append((l, ))
for ll, ll2 in zip(L, L2):
assert ll2 == L2[0]
assert ll != ll2
def test_filter_str():
def iseven(x):
return x % 2 == 0
source = Stream()
s = source.filter(iseven)
assert str(s) == "<filter: iseven>"
def test_map_str():
def add(x=0, y=0):
return x + y
source = Stream()
s = source.map(add, y=10)
assert str(s) == "<map: add>"
def test_create_file():
source1 = Stream(stream_name="source1")
source2 = Stream(stream_name="source2")
n1 = source1.zip(source2)
n2 = n1.map(add).scan(mul).map(lambda x: x + 1)
n2.sink(source1.emit)
with tmpfile(extension="png") as fn:
visualize(n1, filename=fn)
assert os.path.exists(fn)
with tmpfile(extension="svg") as fn:
n1.visualize(filename=fn, rankdir="LR")
assert os.path.exists(fn)
with tmpfile(extension="dot") as fn:
n1.visualize(filename=fn, rankdir="LR")
with open(fn) as f:
text = f.read()
for word in [
"rankdir",
"source1",
"source2",
"zip",
"map",
"add",
"shape=box",
"shape=ellipse",
]:
assert word in text
def test_buffer(c, s, a, b):
source = Stream(asynchronous=True)
L = (
source.scatter()
.map(slowinc, delay=0.5)
.buffer(5)
.gather()
.sink_to_list()
)
start = time.time()
for i in range(5):
yield source.emit(i)
end = time.time()
assert end - start < 0.5
for i in range(5, 10):
yield source.emit(i)
end2 = time.time()
assert end2 - start > (0.5 / 3)
while len(L) < 10:
yield gen.sleep(0.01)
assert time.time() - start < 5
assert L == list(map(inc, range(10)))
assert source.loop == c.loop
def test_slow_to_event_model():
"""This doesn't use threads so it should be slower due to sleep"""
source = Stream(asynchronous=True)
t = FromEventStream("event", ("data", "det_image"), source, principle=True)
assert t.principle
a = t.map(slow_inc)
L = a.sink_to_list()
futures_L = a.sink_to_list()
n = a.SimpleToEventStream(("ct", ))
n.sink(print)
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
t0 = time.time()
for gg in y(10):
yield source.emit(gg)
while len(L) < len(futures_L):
yield gen.sleep(.01)
t1 = time.time()
# check that this was faster than running in series
td = t1 - t0
ted = .5 * 10
assert td > ted
assert tt
assert p == ["start", "descriptor"] + ["event"] * 10 + ["stop"]
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
def test_amorphous_pipeline():
pdf = Stream()
ns = amorphsivity_pipeline(pdf)
L = ns["amorphsivity"].sink_to_list()
a = np.ones(10)
pdf.emit(a)
assert L[0] == np.sum(a[6:])
def run_exp(delay): # pragma: no cover
time.sleep(delay)
print("running exp")
p = Publisher(proxy[0], prefix=b"raw")
RE.subscribe(p)
# Tiny fake pipeline
pp = Publisher(proxy[0], prefix=b"an")
raw_source = Stream()
SimpleFromEventStream(
"event",
("data", "img"),
raw_source.starmap(Retrieve({"NPY_SEQ": NumpySeqHandler})),
principle=True,
).map(lambda x: x * 2).SimpleToEventStream(
("img2",), analysis_stage="pdf"
).starsink(
pp
)
RE.subscribe(lambda *x: raw_source.emit(x))
RE(bp.count([hw.img], md=dict(analysis_stage="raw")))
print("finished exp")
p.close()
def test_same_hdr_many_times(hw, RE):
source = Stream()
fes1 = FromEventStream("start", ("number",), source, principle=True)
fes2 = FromEventStream("event", ("data", "motor"), source, principle=True)
out1 = fes1.map(op.add, 1)
out2 = fes2.combine_latest(out1, emit_on=0).starmap(op.mul)
a = ToEventStream(out1, ("out1",))
b = ToEventStream(out2, ("out2",))
la = a.sink_to_list()
lb = b.sink_to_list()
L = []
RE.subscribe(lambda *x: L.append(x))
RE(count([hw.motor], md={"number": 5}))
for i in range(1, 3):
for ll in L:
source.emit(ll)
for lst in [la, lb]:
o1 = [z[0] for z in lst]
o2 = ["start", "descriptor", "event", "stop"] * i
assert o1 == o2
def test_separate_thread_without_time(loop, thread):
assert thread.is_alive()
source = Stream(loop=loop)
L = source.map(inc).sink_to_list()
for i in range(10):
source.emit(i)
assert L[-1] == i + 1
def test_pluck(backend):
source = Stream(asynchronous=True)
L = source.scatter(backend=backend).pluck(0).gather().sink_to_list()
for i in range(5):
yield source.emit((i, i))
assert L == list(range(5))
def test_sync_in_event_loop():
a = Stream()
assert not a.asynchronous
L = a.timed_window(0.01).sink_to_list()
sleep(0.05)
assert L
assert a.loop
assert a.loop is not IOLoop.current()
def test_orderedweakset_index():
source = Stream()
b1 = source.map(inc)
b2 = source.map(double)
c = b1.scan(add)
assert source.downstreams.index(b1) == 0
def test_filter():
source = Stream()
L = source.filter(lambda x: x % 2 == 0).sink_to_list()
for i in range(10):
source.emit(i)
assert L == [0, 2, 4, 6, 8]
def test_zip_same():
a = Stream()
b = a.zip(a)
L = b.sink_to_list()
a.emit(1)
a.emit(2)
assert L == [(1, 1), (2, 2)]
def test_partition():
source = Stream()
L = source.partition(2).sink_to_list()
for i in range(10):
source.emit(i)
assert L == [(0, 1), (2, 3), (4, 5), (6, 7), (8, 9)]
def test_remove():
source = Stream()
L = source.remove(lambda x: x % 2 == 0).sink_to_list()
for i in range(10):
source.emit(i)
assert L == [1, 3, 5, 7, 9]
def test_filter_none():
source = Stream()
L = source.filter(None).sink_to_list()
for i in range(10):
source.emit(i % 3)
assert L == [1, 2, 1, 2, 1, 2]
def test_map_errors_log():
a = Stream(asynchronous=True)
b = a.delay(0.001).map(lambda x: 1 / x)
with captured_logger("rapidz") as logger:
a._emit(0)
yield gen.sleep(0.1)
out = logger.getvalue()
assert "ZeroDivisionError" in out
def test_rate_limit():
source = Stream(asynchronous=True)
L = source.rate_limit(0.05).sink_to_list()
start = time()
for i in range(5):
yield source.emit(i)
stop = time()
assert stop - start > 0.2
assert len(L) == 5
def test_zip_latest_first():
a = Stream()
b = Stream()
c = a.zip_latest(b).starmap(operator.sub)
d = a.zip_latest(b, first=True).starmap(operator.add)
L = c.union(d).sink_to_list()
a.emit(1)
b.emit(1)
assert L == [2, 0]
def test_filter_args_kwargs():
def f(x, y, z=False):
print(y)
print(z)
return y and z
source = Stream()
L = source.filter(f, True, z=True).sink_to_list()
source.emit(1)
assert L[0] is 1
def test_map():
def add(x=0, y=0):
return x + y
source = Stream()
L = source.map(add, y=10).sink_to_list()
source.emit(1)
assert L[0] == 11
def test_accumulate_errors_log():
a = Stream(asynchronous=True)
b = a.delay(0.001).accumulate(lambda x, y: x / y)
with captured_logger("rapidz") as logger:
a._emit(1)
a._emit(0)
yield gen.sleep(0.1)
out = logger.getvalue()
assert "ZeroDivisionError" in out
def test_starmap():
def add(x=0, y=0):
return x + y
source = Stream()
L = source.starmap(add).sink_to_list()
source.emit((1, 10))
assert L[0] == 11
def __init__(self, upstream, root, writer, resource_kwargs=None, **kwargs):
Stream.__init__(self, upstream, **kwargs)
if writer is None:
writer = {}
self.writer = writer
self.root = root
self.resource_kwargs = resource_kwargs
self.init_writers = {}
self.descriptors = {}
self.not_issued_descriptors = set()
def test_align_stream_syntax():
a = Stream()
b = Stream()
z = a.align_event_streams(b)
sl = z.sink_to_list()
# TODO: use real run engine here
for n, d, dd in zip(
["start", "descriptor", "event", "stop"],
[
{"a": "hi", "b": {"hi": "world"}, "uid": "hi", "time": 123},
{"bla": "foo", "uid": "abc"},
{"data": "now", "descriptor": "abc"},
{"stop": "doc"},
],
[
{"a": "hi2", "b": {"hi2": "world"}},
{"bla": "foo", "uid": "123"},
{"data": "now", "descriptor": "123"},
{"stop": "doc"},
],
):
a.emit((n, d))
b.emit((n, dd))
assert len(sl) == 4
assert sl[0][1].get("b") == {"hi": "world", "hi2": "world"}
assert "original_start_time" in sl[0][1]
def test_scan(c, s, a, b):
source = Stream(asynchronous=True)
futures = scatter(source).map(inc).scan(add)
futures_L = futures.sink_to_list()
L = futures.gather().sink_to_list()
for i in range(5):
yield source.emit(i)
assert L == [1, 3, 6, 10, 15]
assert all(isinstance(f, Future) for f in futures_L)
def test_starmap(c, s, a, b):
def add(x, y, z=0):
return x + y + z
source = Stream(asynchronous=True)
L = source.scatter().starmap(add, z=10).gather().sink_to_list()
for i in range(5):
yield source.emit((i, i))
assert L == [10, 12, 14, 16, 18]
def test_sync_2(loop):
with cluster() as (s, [a, b]):
with Client(s["address"], loop=loop): # flake8: noqa
source = Stream()
L = source.scatter().map(inc).gather().sink_to_list()
for i in range(10):
source.emit(i)
assert len(L) == i + 1
assert L == list(map(inc, range(10)))
def test_filter_map(c, s, a, b):
source = Stream(asynchronous=True)
futures = scatter(source).filter(lambda x: x % 2 == 0).map(inc)
futures_L = futures.sink_to_list()
L = futures.gather().sink_to_list()
for i in range(5):
yield source.emit(i)
assert L == [1, 3, 5]
assert all(isinstance(f, Future) for f in futures_L)
def full_field_tomo(source: Stream, qoi_name, rotation, **kwargs):
theta = SimpleFromEventStream(
"event", ("data", rotation), upstream=source
).map(np.deg2rad)
qoi = SimpleFromEventStream(
"event", ("data", qoi_name), upstream=source, principle=True
)
center = SimpleFromEventStream(
"start", ("tomo", "center"), upstream=source
)
source.starsink(StartStopCallback())
return locals()
def pencil_tomo(source: Stream, qoi_name, translation, rotation, stack=None,
**kwargs):
"""Extract data from a raw stream for pencil beam tomography
Parameters
----------
source : Stream
The stream of raw event model data
qoi_name : str
The name of the QOI for this reconstruction
kwargs
Returns
-------
dict :
The namespace
"""
start = SimpleFromEventStream('start', (), upstream=source)
if stack:
stack_position = SimpleFromEventStream("event", ("data", stack),
upstream=source)
x = SimpleFromEventStream("event", ("data", translation), upstream=source)
th = SimpleFromEventStream("event", ("data", rotation), upstream=source)
# Extract the index for the translation and rotation so we can
# extract the dimensions and extents
# TODO: turn into proper function
translation_position = SimpleFromEventStream(
"start", ("motors",), upstream=source
).map(lambda x: x.index(translation))
rotation_position = SimpleFromEventStream(
"start", ("motors",), upstream=source
).map(lambda x: x.index(rotation))
dims = SimpleFromEventStream("start", ("shape",), upstream=source)
th_dim = dims.zip(rotation_position).starmap(op.getitem)
x_dim = dims.zip(translation_position).starmap(op.getitem)
extents = SimpleFromEventStream("start", ("extents",), upstream=source)
th_extents = extents.zip(rotation_position).starmap(op.getitem)
x_extents = extents.zip(translation_position).starmap(op.getitem)
qoi = SimpleFromEventStream(
"event", ("data", qoi_name), upstream=source, principle=True
)
center = SimpleFromEventStream(
"start", ("tomo", "center"), upstream=source
)
source.starsink(StartStopCallback())
return locals()
def run_server(
data_dir,
outbound_proxy_address=glbl_dict["outbound_proxy_address"],
prefix=b"an",
):
"""Start up the databroker server for analyzed data.
Parameters
----------
data_dir : str
The directory to save the array data into.
outbound_proxy_address : str, optional
The address and port of the zmq proxy. Defaults to
``glbl_dict["outbound_proxy_address"]``
prefix : bytes or list of bytes, optional
The Publisher channels to listen to. Defaults to
``b"an"``
"""
d = RemoteDispatcher(outbound_proxy_address, prefix=prefix)
an_broker = glbl_dict["an_db"]
an_source = Stream()
an_source.Store(data_dir, NpyWriter).starsink(an_broker.insert)
rr = RunRouter(
[
lambda x: (lambda *nd: an_source.emit(nd))
if x.get("analysis_stage", None) == "pdf"
else None,
lambda x: (lambda *nd: an_source.emit(nd))
if x.get("analysis_stage", None) == "integration"
else None,
]
)
d.subscribe(rr)
print("Starting DB Server")
d.start()
yield from trigger_and_read(list(detectors) + list(motors))
install_kicker()
p = Publisher(glbl_dict["inbound_proxy_address"])
hw = hw()
import numpy as np
rand_img = SynSignal(
func=lambda: np.array(np.random.random((10, 10))),
name="img",
labels={"detectors"},
)
RE = RunEngine()
# build the pipeline
raw_source = Stream()
raw_output = SimpleFromEventStream(
"event", ("data", "det_a"), raw_source, principle=True
)
raw_output2 = SimpleFromEventStream("event", ("data", "noisy_det"), raw_source)
raw_output3 = SimpleFromEventStream("event", ("data", "img"), raw_source)
pipeline = (
raw_output.union(raw_output2, raw_output3.map(np.sum))
.map(lambda x: x ** 2)
.accumulate(lambda x, y: x + y)
)
res = SimpleToEventStream(pipeline, ("result",))
merge = AlignEventStreams(raw_source.starmap(StripDepVar()), res)
def run_server(
prefix=None,
outbound_proxy_address=glbl_dict["outbound_proxy_address"],
inbound_proxy_address=glbl_dict["inbound_proxy_address"],
_publisher=None,
**kwargs
):
"""Start up the QOI server
Parameters
----------
prefix : bytes or list of bytes, optional
The Publisher channels to listen to. Defaults to
``[b"an", b"raw"]``
outbound_proxy_address : str, optional
The address and port of the zmq proxy. Defaults to
``glbl_dict["outbound_proxy_address"]``
inbound_proxy_address : str, optional
The inbound ip address for the ZMQ server. Defaults to the value
from the global dict
"""
if prefix is None:
prefix = [b"an", b"raw"]
d = RemoteDispatcher(outbound_proxy_address, prefix=prefix)
install_qt_kicker(loop=d.loop)
if _publisher is None:
an_with_ind_pub = Publisher(inbound_proxy_address, prefix=b"qoi")
else:
an_with_ind_pub = _publisher
raw_source = Stream()
# create amorphous pipeline
amorphous_ns = link(
*[amorphsivity_fem, amorphsivity_pipeline, amorphsivity_tem],
source=Stream(),
**kwargs
)
# Combine the data outputs with the raw independent data
amorphous_ns.update(
to_event_stream_with_ind(
move_to_first(raw_source.starmap(StripDepVar())),
*[
node
for node in amorphous_ns.values()
if isinstance(node, SimpleToEventStream)
],
publisher=an_with_ind_pub
)
)
rr = RunRouter(
[
lambda x: lambda *y: raw_source.emit(y)
if x["analysis_stage"] == "raw"
else None,
lambda x: lambda *y: amorphous_ns["source"].emit(y)
if x["analysis_stage"] == "pdf"
else None,
]
)
d.subscribe(rr)
print("Starting QOI Server")
d.start()
def run_server(
folder,
outbound_proxy_address=glbl_dict["outbound_proxy_address"],
prefix=None,
handlers=None,
):
"""Start up the portable databroker server
Parameters
----------
folder : str
The location where to save the portable databrokers
outbound_proxy_address : str, optional
The address and port of the zmq proxy. Defaults to
``glbl_dict["outbound_proxy_address"]``
prefix : bytes or list of bytes, optional
The Publisher channels to listen to. Defaults to
``[b"an", b"raw"]``
handlers : dict
The map between handler specs and handler classes, defaults to
the map used by the experimental databroker if possible
"""
# TODO: convert to bytestrings if needed
# TODO: maybe separate this into different processes?
# TODO: support multiple locations for folders
if prefix is None:
prefix = [b"an", b"raw"]
d = RemoteDispatcher(outbound_proxy_address, prefix=prefix)
portable_folder = folder
portable_configs = {}
for folder_name in ["an", "raw"]:
fn = os.path.join(portable_folder, folder_name)
os.makedirs(fn, exist_ok=True)
# if the path doesn't exist then make the databrokers
with open(
os.path.join(portable_folder, f"{folder_name}.yml"), "w"
) as f:
f.write(portable_template.format(folder_name))
print(portable_template.format(folder_name))
print(fn)
portable_configs[folder_name] = yaml.load(
io.StringIO(portable_template.format(fn))
)
os.makedirs(os.path.join(fn, "data"), exist_ok=True)
# TODO: add more files here, eg. a databroker readme/tutorial
with open(os.path.join(portable_folder, "db_load.py"), "w") as f:
f.write(load_script)
an_broker = Broker.from_config(portable_configs["an"])
an_source = Stream()
zed = an_source.Store(
os.path.join(
portable_configs["an"]["metadatastore"]["config"]["directory"],
"data",
),
NpyWriter,
)
zed.starsink(an_broker.insert)
raw_broker = Broker.from_config(portable_configs["raw"])
if handlers is None:
handlers = raw_broker.reg.handler_reg
raw_source = Stream()
raw_source.starmap(
ExportCallback(
os.path.join(
portable_configs["raw"]["metadatastore"]["config"][
"directory"
],
"data",
),
handler_reg=handlers,
)
).starsink(raw_broker.insert)
rr = RunRouter(
[
lambda x: (lambda *nd: raw_source.emit(nd))
if x.get("analysis_stage", "") == "raw"
else None
]
+ [
lambda x: (lambda *nd: an_source.emit(nd))
if x.get("analysis_stage", None) == "pdf"
else None,
lambda x: (lambda *nd: an_source.emit(nd))
if x.get("analysis_stage", None) == "integration"
else None,
]
)
d.subscribe(rr)
print("Starting Portable DB Server")
d.start()