def test_to_event_model_new_api_e_stop(RE, hw):
source = Stream()
t = FromEventStream("event", ("data", "motor"), source, principle=True)
assert t.principle
n = simple_to_event_stream_new_api(
{t: {"data_keys": {"ct": {"units": "arb", "precision": 2}}}}
)
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
def f(*x):
if x[0] == "stop":
return
source.emit(x)
RE.subscribe(f)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
rs = d[0]["uid"]
assert tt
assert set(p) == {"start", "event", "descriptor"}
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
assert d[1]["data_keys"]["ct"]["units"] == "arb"
ll = len(d)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert d[ll]["run_start"] == rs
assert set(p) == {"start", "stop", "event", "descriptor"}
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_to_event_model_new_api_multi(RE, hw):
source = Stream()
stop = FromEventStream("stop", (), source)
t = FromEventStream(
"event", ("data", "motor"), source, principle=True, stream_name="hi"
)
assert t.principle
tt = t.zip(stop)
n = simple_to_event_stream_new_api(
{
t: {"data_keys": {"ct": {"units": "arb", "precision": 2}}},
tt: {
"name": "final",
"data_keys": {"ct": {"units": "arb", "precision": 2}},
},
},
hello="world",
)
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert tt
assert set(p) == {"start", "stop", "event", "descriptor"}
assert d[0]["hello"] == "world"
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
assert d[1]["data_keys"]["ct"]["units"] == "arb"
assert d[-3]["name"] == "final"
assert d[-1]["run_start"]
def test_last_cache(RE, hw):
source = Stream()
t = FromEventStream("event", ("data", "motor"), source, principle=True)
assert t.principle
n = ToEventStream(
t, ("ct",), data_key_md={"ct": {"units": "arb"}}
).LastCache()
tt = t.sink_to_list()
names = n.pluck(0).sink_to_list()
docs = n.pluck(1).sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert len(docs) == 10 + 3 + 2
assert names[-3] == "descriptor"
assert names[-2] == "event"
assert tt
assert set(names) == {"start", "stop", "event", "descriptor"}
assert docs[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
assert docs[1]["data_keys"]["ct"]["units"] == "arb"
assert docs[-1]["run_start"]
def test_to_event_model_new_api_multi_parent(RE, hw):
source = Stream()
t = FromEventStream("event", ("data", "motor"), source, principle=True)
t2 = FromEventStream("event", ("data", "motor"), source, principle=True)
assert t.principle
n = simple_to_event_stream_new_api(
{
t.zip(t2).pluck(0): {
"data_keys": {"ct": {"units": "arb", "precision": 2}}
}
}
)
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert tt
assert set(p) == {"start", "stop", "event", "descriptor"}
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
assert d[1]["data_keys"]["ct"]["units"] == "arb"
assert d[-1]["run_start"]
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_no_parent_nodes():
# build the graph
g1 = FromEventStream(
"event", ("data", "det_image"), stream_name="g1", asynchronous=True
)
g11 = FromEventStream(
"event", ("data", "det_image"), stream_name="g11", asynchronous=True
)
g2 = g1.zip(g11).starmap(op.mul, stream_name="mul")
g2.SimpleToEventStream(("img2",))
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 test_from_event_model_all(RE, hw):
source = Stream()
t = FromEventStream("event", (), source, principle=True)
L = t.sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert len(L) == 10
for i, ll in enumerate(L):
assert i == ll["data"]["motor"]
def test_double_buffer_to_event_model():
source = Stream(asynchronous=True)
t = FromEventStream("event", ("data", "det_image"), source, principle=True)
assert t.principle
ts = t
a = ts.map(slow_inc)
aa = ts.map(slow_inc)
n = a.zip(aa).SimpleToEventStream(("ct", ))
b = n
b.sink(print)
L = b.sink_to_list()
futures_L = []
tt = t.sink_to_list()
p = b.pluck(0).sink_to_list()
d = b.pluck(1).sink_to_list()
t0 = time.time()
for gg in y(10):
futures_L.append(gg)
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
assert t1 - t0 > .5 * 10
assert tt
assert p == ["start", "descriptor"] + ["event"] * 10 + ["stop"]
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
t0 = time.time()
for gg in y(10):
futures_L.append(gg)
yield source.emit(gg)
while len(L) < len(futures_L):
yield gen.sleep(.01)
print(len(L), len(futures_L))
t1 = time.time()
# check that this was faster than running in series
assert t1 - t0 > .5 * 10
assert tt
assert p == (["start", "descriptor"] + ["event"] * 10 + ["stop"]) * 2
assert d[14]["hints"] == {"analyzer": {"fields": ["ct"]}}
for i, j in zip([0, 1, 12], [13, 14, 25]):
assert p[i] == p[j]
assert d[i] != d[j]
def test_to_event_model_new_api_no_principle(RE, hw):
source = Stream()
stop = FromEventStream("stop", (), source)
t = FromEventStream("event", ("data", "motor"), source, stream_name="hi")
tt = t.zip(stop)
simple_to_event_stream_new_api(
{
t: {"data_keys": {"ct": {"units": "arb", "precision": 2}}},
tt: {
"name": "final",
"data_keys": {"ct": {"units": "arb", "precision": 2}},
},
},
hello="world",
)
def test_walk_up():
raw = Stream()
a_translation = FromEventStream("start", ("time",), raw, principle=True)
b_translation = FromEventStream("event", ("data", "pe1_image"), raw)
d = b_translation.zip_latest(a_translation)
dd = d.map(op.truediv)
e = ToEventStream(dd, ("data",))
g = nx.DiGraph()
walk_to_translation(e, g)
att = []
for node, attrs in g.nodes.items():
att.append(attrs["stream"])
s = {a_translation, b_translation, d, dd, e}
assert s == set(att)
assert {_hash_or_uid(k) for k in s} == set(g.nodes)
def test_align_interrupted(RE, hw):
a = Stream()
b = FromEventStream("event", ("data", "img"), a, principle=True).map(
op.add, 1
)
b.sink(print)
c = ToEventStream(b, ("out",))
z = move_to_first(a.AlignEventStreams(c))
sl = z.sink_to_list()
L = []
RE.subscribe(lambda *x: L.append(x))
RE(count([hw.img]))
for nd in L:
name, doc = nd
# cause an exception
if name == "event":
doc["data"]["img"] = "hi"
try:
a.emit((name, doc))
except TypeError:
pass
assert {"start", "stop"} == set(list(zip(*sl))[0])
# check that buffers are not cleared, yet
sl.clear()
# If there are elements in the buffer they need to be cleared when all
# start docs come in.
for nd in L:
name, doc = nd
# cause an exception
if name == "event":
doc["data"]["img"] = 1
a.emit((name, doc))
if name == "start":
# now buffers should be clear
assert not any(
[b for n, tb in z.true_buffers.items() for u, b in tb.items()]
)
assert {"start", "descriptor", "event", "stop"} == set(list(zip(*sl))[0])
# now buffers should be clear (as all docs were emitted)
assert not any(
[b for n, tb in z.true_buffers.items() for u, b in tb.items()]
)
def test_parent_nodes():
# build the graph
g1 = FromEventStream(
"event",
("data", "det_image"),
principle=True,
stream_name="g1",
asynchronous=True,
)
g11 = FromEventStream(
"event", ("data", "det_image"), stream_name="g11", asynchronous=True
)
g2 = g1.zip(g11).starmap(op.mul, stream_name="mul")
g = g2.SimpleToEventStream(("img2",))
l1 = g.sink_to_list()
# g.sink(print)
assert len(g.translation_nodes) == 2
print("start experiment")
# run the experiment
l0 = []
for yy in y(5):
l0.append(yy)
g11.update(yy)
g1.update(yy)
print(g11.start_uid)
assert len(l1[0][1]["parent_node_map"]) == 2
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 test_from_event_model_stream_name2():
def data():
suid = str(uuid.uuid4())
duid = str(uuid.uuid4())
yield "start", {"hi": "world", "uid": suid}
yield "descriptor", {
"name": "hi",
"data_keys": {"ct"},
"uid": duid,
"run_start": suid,
}
for i in range(10):
yield "event", {
"uid": str(uuid.uuid4()),
"data": {"ct": i},
"descriptor": duid,
}
duid = str(uuid.uuid4())
yield "descriptor", {
"name": "not hi",
"data_keys": {"ct"},
"uid": duid,
"run_start": suid,
}
for i in range(100, 110):
yield "event", {
"uid": str(uuid.uuid4()),
"data": {"ct": i},
"descriptor": duid,
}
yield "stop", {"uid": str(uuid.uuid4()), "run_start": suid}
g = data()
source = Stream()
t = FromEventStream(
"event", ("data", "ct"), source, event_stream_name="not hi"
)
L = t.sink_to_list()
for gg in g:
source.emit(gg)
assert len(L) == 10
for i, ll in enumerate(L):
assert i + 100 == ll
def test_to_event_model_stream_syntax(RE, hw):
source = Stream()
t = FromEventStream("event", ("data", "motor"), source, principle=True)
assert t.principle
n = t.simple_to_event_stream(("ct",), data_key_md={"ct": {"units": "arb"}})
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert tt
assert set(p) == {"start", "stop", "event", "descriptor"}
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
assert d[1]["data_keys"]["ct"]["units"] == "arb"
assert d[-1]["run_start"]
def test_to_event_model_new_api_clobber(RE, hw):
source = Stream()
t = FromEventStream("event", ("data", "motor"), source, principle=True)
assert t.principle
n = simple_to_event_stream_new_api(
{t: {"data_keys": {"ct": {"units": "arb", "dtype": "array"}}}}
)
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert tt
assert set(p) == {"start", "stop", "event", "descriptor"}
assert d[1]["data_keys"]["ct"]["dtype"] == "array"
assert d[-1]["run_start"]
def test_to_event_model_new_api_no_data_keys(RE, hw):
source = Stream()
t = FromEventStream("event", ("data",), source, principle=True)
assert t.principle
n = simple_to_event_stream_new_api({t: {}})
tt = t.sink_to_list()
p = n.pluck(0).sink_to_list()
d = n.pluck(1).sink_to_list()
RE.subscribe(unstar(source.emit))
RE.subscribe(print)
RE(scan([hw.motor], hw.motor, 0, 9, 10))
assert tt
assert set(p) == {"start", "stop", "event", "descriptor"}
assert d[1]["hints"] == {
"analyzer": {"fields": ["motor", "motor_setpoint"]}
}
assert d[1]["data_keys"]["motor"]
assert d[-1]["run_start"]
def test_slow_to_event_model_parallel_dask(c, s, a, b):
source = Stream(asynchronous=True)
t = FromEventStream("event", ("data", "det_image"), source, principle=True)
assert t.principle
ts = t.scatter(backend="dask")
# futures_L = t.sink_to_list()
a = ts.map(slow_inc)
n = a.SimpleToEventStream(("ct", ))
b = n.buffer(100).gather()
b.sink(print)
L = b.sink_to_list()
tt = t.sink_to_list()
p = b.pluck(0).sink_to_list()
d = b.pluck(1).sink_to_list()
t0 = time.time()
futures_L = []
for gg in y(10):
futures_L.append(gg)
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 "uid" in d[0]
assert d[1]["hints"] == {"analyzer": {"fields": ["ct"]}}
assert d[1]["data_keys"]["ct"]["dtype"] == "number"
assert d[2]["data"]["ct"] == 2
def test_replay_export_test():
def y():
suid = str(uuid.uuid4())
yield ("start", {"uid": suid, "time": time.time()})
duid = str(uuid.uuid4())
yield (
"descriptor",
{
"uid": duid,
"run_start": suid,
"name": "primary",
"data_keys": {"det_image": {"dtype": "int", "units": "arb"}},
"time": time.time(),
},
)
for i in range(5):
yield (
"event",
{
"uid": str(uuid.uuid4()),
"data": {"det_image": i},
"timestamps": {"det_image": time.time()},
"seq_num": i + 1,
"time": time.time(),
"descriptor": duid,
},
)
yield (
"stop",
{"uid": str(uuid.uuid4()), "time": time.time(), "run_start": suid},
)
print("build graph")
g1 = FromEventStream(
"event", ("data", "det_image"), principle=True, stream_name="g1"
)
g11 = FromEventStream("event", ("data", "det_image"), stream_name="g11")
g11_1 = g1.zip(g11)
g2 = g11_1.starmap(op.mul).map(np.log)
g = g2.SimpleToEventStream(("img2",))
from pprint import pprint
g.sink(pprint)
L = g.sink_to_list()
print("run experiment")
for yy in y():
print(yy[0])
g11.update(yy)
g1.update(yy)
assert L[-1][1]["run_start"]
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()