def test_multi_partition_wins():
# FIXME: the constellatin partition+tile or partition+frame or similar
# can be optimized by using a fitting tiling scheme for tile/frame processing
# and accumulating the whole partition into a buffer, then running process_partition
# after the tile loop.
other_unlimited = UDFUnlimitedDepth()
other_best_fit = TestUDFBestFit()
other_deep = UDFWithLargeDepth()
udf_partition = TestUDFPartition()
udfs = [
udf_partition,
other_unlimited,
other_best_fit,
other_deep,
]
data = _mk_random(size=(32, 1860, 2048))
dataset = MemoryDataSet(
data=data,
num_partitions=1,
sig_dims=2,
base_shape=(1, 930, 16),
force_need_decode=False,
)
neg = Negotiator()
p = next(dataset.get_partitions())
scheme = neg.get_scheme(udfs=udfs, partition=p, read_dtype=np.float32, roi=None)
print(scheme._debug)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (32, 1860, 2048)
def test_multi_no_by_frame_small_size_wins():
other_unlimited = UDFUnlimitedDepth()
other_best_fit = TestUDFBestFit()
other_deep = UDFWithLargeDepth()
udfs = [
other_unlimited,
other_best_fit,
other_deep,
]
data = _mk_random(size=(32, 1860, 2048))
dataset = MemoryDataSet(
data=data,
num_partitions=1,
sig_dims=2,
base_shape=(1, 930, 16),
force_need_decode=False,
)
neg = Negotiator()
p = next(dataset.get_partitions())
scheme = neg.get_scheme(udfs=udfs, partition=p, read_dtype=np.float32, roi=None)
print(scheme._debug)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (17, 930, 16)
def test_hdf5_tileshape_negotation(lt_ctx, tmpdir_factory):
# try to hit the third case in _get_subslices:
datadir = tmpdir_factory.mktemp('data')
filename = os.path.join(datadir, 'tileshape-neg-test.h5')
data = _mk_random((4, 100, 256, 256), dtype=np.uint16)
with h5py.File(filename, "w") as f:
f.create_dataset("data", data=data, chunks=(2, 32, 32, 32))
ds = lt_ctx.load("hdf5", path=filename)
udfs = [UDFWithLargeDepth()]
p = next(ds.get_partitions())
neg = Negotiator()
tiling_scheme = neg.get_scheme(
udfs=udfs,
partition=p,
read_dtype=np.float32,
roi=None,
corrections=None,
)
assert len(tiling_scheme) > 1
next(
p.get_tiles(tiling_scheme=tiling_scheme,
roi=None,
dest_dtype=np.float32))
def test_hdf5_result_dtype(lt_ctx, tmpdir_factory, in_dtype, read_dtype,
use_roi):
datadir = tmpdir_factory.mktemp('data')
filename = os.path.join(datadir, 'result-dtype-checks.h5')
data = _mk_random((2, 2, 4, 4), dtype=in_dtype)
with h5py.File(filename, "w") as f:
f.create_dataset("data", data=data)
ds = lt_ctx.load("hdf5", path=filename)
if use_roi:
roi = np.zeros(ds.shape.nav, dtype=bool).reshape((-1, ))
roi[0] = 1
else:
roi = None
udfs = [SumSigUDF()] # need to have at least one UDF
p = next(ds.get_partitions())
neg = Negotiator()
tiling_scheme = neg.get_scheme(
udfs=udfs,
partition=p,
read_dtype=read_dtype,
roi=roi,
corrections=None,
)
tile = next(
p.get_tiles(tiling_scheme=tiling_scheme,
roi=roi,
dest_dtype=read_dtype))
assert tile.dtype == np.dtype(read_dtype)
def test_get_scheme_partition(default_raw):
neg = Negotiator()
p = next(default_raw.get_partitions())
udf = TestUDFPartition()
scheme = neg.get_scheme(udfs=[udf], partition=p, read_dtype=np.float32, roi=None)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (128, 128, 128)
def test_correction_size_overflow():
data = _mk_random(size=(32, 1860, 2048))
dataset = MemoryDataSet(
data=data,
num_partitions=1,
sig_dims=2,
base_shape=(1, 930, 16),
force_need_decode=True,
)
neg = Negotiator()
p = next(dataset.get_partitions())
udf = UDFWithLargeDepth()
excluded_coords = np.array([
(930, ),
(16, )
])
excluded_pixels = sparse.COO(coords=excluded_coords, shape=dataset.shape.sig, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
scheme = neg.get_scheme(
udfs=[udf], partition=p, read_dtype=np.float32, roi=None,
corrections=corr,
)
print(scheme._debug)
assert scheme._debug["need_decode"]
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (4, 1860, 32)
def prime_numba_cache(ds):
dtypes = (np.float32, None)
for dtype in dtypes:
roi = np.zeros(ds.shape.nav, dtype=bool).reshape((-1,))
roi[0] = 1
from libertem.udf.sum import SumUDF
from libertem.corrections.corrset import CorrectionSet
from libertem.io.dataset.base import Negotiator
udfs = [SumUDF()] # need to have at least one UDF
p = next(ds.get_partitions())
neg = Negotiator()
for corr_dtype in (np.float32, None):
if corr_dtype is not None:
corrections = CorrectionSet(dark=np.zeros(ds.shape.sig, dtype=corr_dtype))
else:
corrections = None
p.set_corrections(corrections)
tiling_scheme = neg.get_scheme(
udfs=udfs,
partition=p,
read_dtype=dtype,
roi=roi,
corrections=corrections,
)
next(p.get_tiles(tiling_scheme=tiling_scheme, roi=roi))
def _init_udfs(self, numpy_udfs, cupy_udfs, partition, roi, corrections,
device_class):
dtype = self._get_dtype(partition.dtype, corrections)
meta = UDFMeta(
partition_shape=partition.slice.adjust_for_roi(roi).shape,
dataset_shape=partition.meta.shape,
roi=roi,
dataset_dtype=partition.dtype,
input_dtype=dtype,
tiling_scheme=None,
corrections=corrections,
device_class=device_class,
)
for udf in numpy_udfs:
if device_class == 'cuda':
udf.set_backend('cuda')
else:
udf.set_backend('numpy')
if device_class == 'cpu':
assert not cupy_udfs
for udf in cupy_udfs:
udf.set_backend('cupy')
udfs = numpy_udfs + cupy_udfs
for udf in udfs:
udf.set_meta(meta)
udf.init_result_buffers()
udf.allocate_for_part(partition, roi)
udf.init_task_data()
if hasattr(udf, 'preprocess'):
udf.clear_views()
udf.preprocess()
neg = Negotiator()
# FIXME take compute backend into consideration as well
# Other boundary conditions when moving input data to device
tiling_scheme = neg.get_scheme(
udfs=udfs,
partition=partition,
read_dtype=dtype,
roi=roi,
corrections=corrections,
)
# print(tiling_scheme)
# FIXME: don't fully re-create?
meta = UDFMeta(
partition_shape=partition.slice.adjust_for_roi(roi).shape,
dataset_shape=partition.meta.shape,
roi=roi,
dataset_dtype=partition.dtype,
input_dtype=dtype,
tiling_scheme=tiling_scheme,
corrections=corrections,
device_class=device_class,
)
for udf in udfs:
udf.set_meta(meta)
return (meta, tiling_scheme, dtype)
def test_get_scheme_tile(default_raw):
neg = Negotiator()
p = next(default_raw.get_partitions())
udf = TestUDFBestFit()
scheme = neg.get_scheme(udfs=[udf], partition=p, read_dtype=np.float32, roi=None)
assert scheme.shape.sig.dims == 2
print(neg._get_udf_size_pref(udf))
print(scheme._debug)
print(p.shape)
assert tuple(scheme.shape) == (64, 32, 128)
def test_get_scheme_frame(default_raw):
neg = Negotiator()
p = next(default_raw.get_partitions())
udf = TilingUDFFrame()
scheme = neg.get_scheme(udfs=[udf],
dataset=default_raw,
approx_partition_shape=p.shape,
read_dtype=np.float32,
roi=None)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (16, 128, 128)
def test_get_scheme_upper_size_1():
"""
Test that will hit the 2**20 default size
"""
data = _mk_random(size=(1024, 144, 144))
dataset = MemoryDataSet(data=data, num_partitions=1, sig_dims=2)
neg = Negotiator()
p = next(dataset.get_partitions())
udf = TestUDFBestFit()
scheme = neg.get_scheme(udf=udf,
partition=p,
read_dtype=np.float32,
roi=None)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (65, 28, 144)
def test_scale_factors():
neg = Negotiator()
# scenario: k2is.
assert neg._get_scale_factors(
shape=(930, 16),
containing_shape=(1860, 2048),
size=1024,
) == [1, 1]
assert neg._get_scale_factors(
shape=(930, 16),
containing_shape=(1860, 2048),
size=930 * 16 * 16
) == [2, 8]
assert neg._get_scale_factors(
shape=(930, 16),
containing_shape=(1860, 2048),
size=930 * 16 * 128
) == [2, 64]
# slightly above, but not enough to fit another block: we err on the small side
assert neg._get_scale_factors(
shape=(930, 16),
containing_shape=(1860, 2048),
size=930 * 16 * 129
) == [2, 64]
# larger size than we can accomodate with our containing shape:
assert neg._get_scale_factors(
shape=(930, 16),
containing_shape=(1860, 2048),
size=1860 * 2048 * 2
) == [2, 128]
def test_depth_max_size_max():
data = _mk_random(size=(32, 1860, 2048))
dataset = MemoryDataSet(
data=data,
num_partitions=1,
sig_dims=2,
base_shape=(1, 930, 16),
force_need_decode=False,
)
neg = Negotiator()
p = next(dataset.get_partitions())
udf = UDFUnlimitedDepth()
scheme = neg.get_scheme(udfs=[udf], partition=p, read_dtype=np.float32, roi=None)
print(scheme._debug)
assert not scheme._debug["need_decode"]
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (32, 1860, 2048)
def test_get_scheme_upper_size_2():
"""
Test that will hit the 2**20 default size
"""
data = _mk_random(size=(2048, 264, 264))
dataset = MemoryDataSet(
data=data,
num_partitions=1,
sig_dims=2,
base_shape=(1, 8, 264),
)
neg = Negotiator()
p = next(dataset.get_partitions())
udf = TilingUDFBestFit()
scheme = neg.get_scheme(udfs=[udf],
partition=p,
read_dtype=np.float32,
roi=None)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (124, 8, 264)
def test_limited_depth():
data = _mk_random(size=(32, 1860, 2048))
dataset = MemoryDataSet(
data=data,
num_partitions=1,
sig_dims=2,
base_shape=(1, 930, 16),
force_need_decode=True,
)
neg = Negotiator()
p = next(dataset.get_partitions())
udf = UDFWithLargeDepth()
scheme = neg.get_scheme(udfs=[udf],
approx_partition_shape=p.shape,
dataset=dataset,
read_dtype=np.float32,
roi=None)
print(scheme._debug)
assert scheme._debug["need_decode"]
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (17, 930, 16)
def prime_numba_cache(ds):
dtypes = (np.float32, None)
for dtype in dtypes:
roi = np.zeros(ds.shape.nav, dtype=bool).reshape((-1, ))
roi[max(-ds._meta.sync_offset, 0)] = True
from libertem.udf.sum import SumUDF
from libertem.udf.raw import PickUDF
from libertem.io.corrections.corrset import CorrectionSet
from libertem.io.dataset.base import Negotiator
# need to have at least one UDF; here we run for both sum and pick
# to reduce the initial latency when switching to pick mode
udfs = [SumUDF(), PickUDF()]
neg = Negotiator()
for udf in udfs:
for corr_dtype in (np.float32, None):
if corr_dtype is not None:
corrections = CorrectionSet(
dark=np.zeros(ds.shape.sig, dtype=corr_dtype))
else:
corrections = None
found_first_tile = False
for p in ds.get_partitions():
if found_first_tile:
break
p.set_corrections(corrections)
tiling_scheme = neg.get_scheme(
udfs=[udf],
dataset=ds,
approx_partition_shape=p.shape,
read_dtype=dtype,
roi=roi,
corrections=corrections,
)
for t in p.get_tiles(tiling_scheme=tiling_scheme, roi=roi):
found_first_tile = True
break
def test_get_scheme_upper_size_roi():
"""
Confirm that a small ROI will not be split
up unnecessarily.
"""
data = _mk_random(size=(1024, 144, 144))
dataset = MemoryDataSet(data=data, num_partitions=1, sig_dims=2)
roi = np.zeros(dataset.shape.nav, dtype=bool)
# All in a single partition here
roi[0] = True
roi[512] = True
roi[-1] = True
neg = Negotiator()
p = next(dataset.get_partitions())
udf = TilingUDFBestFit()
scheme = neg.get_scheme(udfs=[udf],
partition=p,
read_dtype=np.float32,
roi=roi)
assert scheme.shape.sig.dims == 2
assert tuple(scheme.shape) == (3, 144, 144)
def run_for_partition(self, partition: Partition, roi):
with set_num_threads(1):
dtype = self._get_dtype(partition.dtype)
meta = UDFMeta(
partition_shape=partition.slice.adjust_for_roi(roi).shape,
dataset_shape=partition.meta.shape,
roi=roi,
dataset_dtype=partition.dtype,
input_dtype=dtype,
tiling_scheme=None,
)
udfs = self._udfs
for udf in udfs:
udf.set_meta(meta)
udf.init_result_buffers()
udf.allocate_for_part(partition, roi)
udf.init_task_data()
if hasattr(udf, 'preprocess'):
udf.clear_views()
udf.preprocess()
neg = Negotiator()
tiling_scheme = neg.get_scheme(
udfs=udfs,
partition=partition,
read_dtype=dtype,
roi=roi,
)
# FIXME: don't fully re-create?
meta = UDFMeta(
partition_shape=partition.slice.adjust_for_roi(roi).shape,
dataset_shape=partition.meta.shape,
roi=roi,
dataset_dtype=partition.dtype,
input_dtype=dtype,
tiling_scheme=tiling_scheme,
)
for udf in udfs:
udf.set_meta(meta)
# print("UDF TilingScheme: %r" % tiling_scheme.shape)
tiles = partition.get_tiles(tiling_scheme=tiling_scheme,
roi=roi,
dest_dtype=dtype)
for tile in tiles:
for udf in udfs:
method = udf.get_method()
if method == 'tile':
udf.set_contiguous_views_for_tile(partition, tile)
udf.set_slice(tile.tile_slice)
udf.process_tile(tile)
elif method == 'frame':
tile_slice = tile.tile_slice
for frame_idx, frame in enumerate(tile):
frame_slice = Slice(
origin=(tile_slice.origin[0] + frame_idx, ) +
tile_slice.origin[1:],
shape=Shape(
(1, ) + tuple(tile_slice.shape)[1:],
sig_dims=tile_slice.shape.sig.dims),
)
udf.set_slice(frame_slice)
udf.set_views_for_frame(partition, tile, frame_idx)
udf.process_frame(frame)
elif method == 'partition':
udf.set_views_for_tile(partition, tile)
udf.set_slice(partition.slice)
udf.process_partition(tile)
for udf in udfs:
udf.flush()
if hasattr(udf, 'postprocess'):
udf.clear_views()
udf.postprocess()
udf.cleanup()
udf.clear_views()
if self._debug:
try:
cloudpickle.loads(cloudpickle.dumps(partition))
except TypeError:
raise TypeError("could not pickle partition")
try:
cloudpickle.loads(
cloudpickle.dumps([u.results for u in udfs]))
except TypeError:
raise TypeError("could not pickle results")
return tuple(udf.results for udf in udfs)
