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()