def _run_benchmark(shared_dist_ctx, large_raw_file, nblocks, benchmark,
preserve_dim, min_size):
filename, shape, dtype = large_raw_file
chunking = tuple(dim // nb for dim, nb in zip(shape, nblocks))
print(chunking)
d_arr = _mk_dask_from_delayed(shape=shape,
dtype=dtype,
chunking=chunking,
filename=filename)
udf = NoOpUDF()
ds = shared_dist_ctx.load('dask',
d_arr,
sig_dims=2,
preserve_dimensions=preserve_dim,
min_size=min_size)
benchmark.pedantic(
shared_dist_ctx.run_udf,
kwargs=dict(
dataset=ds,
udf=udf,
),
warmup_rounds=0,
rounds=5,
iterations=1,
)
def test_real_correction_baseline(self, shared_dist_ctx, large_raw_file, benchmark):
filename, shape, dtype = large_raw_file
nav_dims = shape[:2]
sig_dims = shape[2:]
print("Nav dims: ", nav_dims)
print("Sig dims:", sig_dims)
udf = NoOpUDF()
ds = shared_dist_ctx.load(
'RAW',
path=str(filename),
scan_size=shape[:2],
dtype=dtype,
detector_size=shape[2:],
)
benchmark.pedantic(
shared_dist_ctx.run_udf,
kwargs=dict(
dataset=ds,
udf=udf,
),
warmup_rounds=0,
rounds=5,
iterations=1,
)
def test_real_correction(self, shared_dist_ctx, large_raw_file, benchmark,
gain, dark, num_excluded):
filename, shape, dtype = large_raw_file
nav_dims = shape[:2]
sig_dims = shape[2:]
if gain == 'use gain':
gain_map = (np.random.random(sig_dims) + 1).astype(np.float64)
elif gain == 'no gain':
gain_map = None
else:
raise ValueError
if dark == 'use dark':
dark_image = np.random.random(sig_dims).astype(np.float64)
elif dark == 'no dark':
dark_image = None
else:
raise ValueError
if num_excluded > 0:
excluded_coords = exclude_pixels(sig_dims=sig_dims,
num_excluded=num_excluded)
assert excluded_coords.shape[1] == num_excluded
exclude = sparse.COO(coords=excluded_coords,
shape=sig_dims,
data=True)
else:
exclude = None
print("Nav dims: ", nav_dims)
print("Sig dims:", sig_dims)
corrset = CorrectionSet(
dark=dark_image,
gain=gain_map,
excluded_pixels=exclude,
)
udf = NoOpUDF()
ds = shared_dist_ctx.load(
'RAW',
path=str(filename),
scan_size=shape[:2],
dtype=dtype,
detector_size=shape[2:],
)
benchmark.pedantic(
shared_dist_ctx.run_udf,
kwargs=dict(
dataset=ds,
udf=udf,
corrections=corrset,
),
warmup_rounds=0,
rounds=5,
iterations=1,
)
def test_default_ctx():
data = np.random.random((13, 17, 19, 23))
ctx = LiveContext()
aq = ctx.prepare_acquisition('memory', data=data)
udf1 = NoOpUDF()
ctx.run_udf(dataset=aq, udf=udf1)
for res in ctx.run_udf_iter(dataset=aq, udf=udf1):
pass
def test_multi_udf(lt_ctx, default_raw, progress, plots):
udfs = [NoOpUDF(), SumUDF(), SumSigUDF()]
combined_res = lt_ctx.run_udf(dataset=default_raw,
udf=udfs,
progress=progress,
plots=plots)
ref_res = (
lt_ctx.run_udf(dataset=default_raw, udf=udfs[0]),
lt_ctx.run_udf(dataset=default_raw, udf=udfs[1]),
lt_ctx.run_udf(dataset=default_raw, udf=udfs[2]),
)
assert isinstance(ref_res[0], dict)
for index, res in enumerate(ref_res):
for key in res.keys():
assert np.all(res[key].data == combined_res[index][key].data)
def test_plots(lt_ctx, default_raw, plots):
udfs = [NoOpUDF(), SumUDF(), SumSigUDF()]
if plots is None:
def test_channel(udf_result, damage):
result = udf_result['intensity'].data
if udf_result['intensity'].kind == 'nav':
res_damage = damage
else:
res_damage = True
return (result, res_damage)
real_plots = [
lt_ctx.plot_class(dataset=default_raw,
udf=udfs[1],
channel=test_channel,
title="Test title"),
lt_ctx.plot_class(dataset=default_raw,
udf=udfs[2],
channel=test_channel,
title="Test title")
]
mock_target = real_plots[0]
else:
real_plots = plots
mock_target = lt_ctx.plot_class
with mock.patch.object(mock_target, 'update') as plot_update:
with mock.patch.object(mock_target, 'display') as plot_display:
if plots is None:
real_plots[0].display()
real_plots[1].display()
combined_res = lt_ctx.run_udf(dataset=default_raw,
udf=udfs,
plots=real_plots)
plot_update.assert_called()
plot_display.assert_called()
ref_res = (
lt_ctx.run_udf(dataset=default_raw, udf=udfs[0]),
lt_ctx.run_udf(dataset=default_raw, udf=udfs[1]),
lt_ctx.run_udf(dataset=default_raw, udf=udfs[2]),
)
assert isinstance(ref_res[0], dict)
for index, res in enumerate(ref_res):
for key in res.keys():
assert np.all(res[key].data == combined_res[index][key].data)
def test_trigger(ltl_ctx):
data = np.random.random((13, 17, 19, 23))
triggered = np.array((False, ))
def trigger(acquisition):
triggered[:] = True
assert tuple(acquisition.shape.nav) == data.shape[:2]
aq = ltl_ctx.prepare_acquisition('memory', trigger=trigger, data=data)
udf1 = SignalMonitorUDF()
udf2 = NoOpUDF()
res = ltl_ctx.run_udf(dataset=aq, udf=[udf1, udf2])
assert np.all(res[0]['intensity'].data == data[-1, -1])
assert triggered
def test_udf_iter(lt_ctx, default_raw, progress, plots):
udfs = [NoOpUDF(), SumUDF(), SumSigUDF()]
for res in lt_ctx.run_udf_iter(dataset=default_raw,
udf=udfs,
progress=progress,
plots=plots):
ref_res = lt_ctx.run_udf(dataset=default_raw,
udf=copy.deepcopy(udfs),
roi=res.damage.data)
# first one is empty, skipping
# kind='sig'
for key in ref_res[1].keys():
ref_item = ref_res[1][key].data
res_item = res.buffers[1][key].data
assert np.all(ref_item == res_item)
# kind='nav'
for key in ref_res[2].keys():
ref_item = ref_res[2][key].data[res.damage.data]
res_item = res.buffers[2][key].data[res.damage.data]
assert np.all(ref_item == res_item)
async def test_udf_iter_async(lt_ctx, default_raw, progress, plots):
udfs = [NoOpUDF(), SumUDF(), SumSigUDF()]
async for res in lt_ctx.run_udf_iter(dataset=default_raw,
udf=udfs,
progress=progress,
plots=plots,
sync=False):
# Nested execution of UDFs doesn't work, so we take a copy
ref_res = lt_ctx.run_udf(dataset=default_raw,
udf=copy.deepcopy(udfs),
roi=res.damage.data)
# first one is empty, skipping
# kind='sig'
for key in ref_res[1].keys():
ref_item = ref_res[1][key].data
res_item = res.buffers[1][key].data
assert np.all(ref_item == res_item)
# kind='nav'
for key in ref_res[2].keys():
ref_item = ref_res[2][key].data[res.damage.data]
res_item = res.buffers[2][key].data[res.damage.data]
assert np.all(ref_item == res_item)
def test_plots_fail(lt_ctx, default_raw, plots):
udfs = [NoOpUDF()]
with pytest.raises(ValueError):
lt_ctx.run_udf(dataset=default_raw, udf=udfs, plots=plots)
def test_live_nochannels(default_raw):
udf = NoOpUDF()
with pytest.raises(ValueError):
Dummy2DPlot(dataset=default_raw, udf=udf)
def test_overhead(shared_dist_ctx, benchmark):
ds = shared_dist_ctx.load('memory',
data=np.zeros((1024, 2)),
sig_dims=1,
num_partitions=32)
benchmark(shared_dist_ctx.run_udf, dataset=ds, udf=NoOpUDF())
