def test_tileshape_adjustment_10():
sig_shape = (122, 455)
tile_shape = (8, 1)
base_shape = (2, 1)
excluded_coords = np.array([
(121, ),
(454, )
])
excluded_pixels = sparse.COO(coords=excluded_coords, shape=sig_shape, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
adjusted = corr.adjust_tileshape(
tile_shape=tile_shape, sig_shape=sig_shape, base_shape=base_shape
)
assert adjusted == (8, 3)
_validate(excluded_coords=excluded_coords, adjusted=adjusted, sig_shape=sig_shape)
def test_tileshape_adjustment_7():
sig_shape = (123, 456)
tile_shape = (14, 42)
base_shape = (7, 1)
excluded_coords = np.array([
(14, ),
(42, )
])
excluded_pixels = sparse.COO(coords=excluded_coords, shape=sig_shape, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
adjusted = corr.adjust_tileshape(
tile_shape=tile_shape, sig_shape=sig_shape, base_shape=base_shape
)
assert adjusted == (21, 41)
_validate(excluded_coords=excluded_coords, adjusted=adjusted, sig_shape=sig_shape)
def test_tileshape_adjustment_6_3():
sig_shape = (123, 456)
tile_shape = (1, 1)
base_shape = (1, 1)
excluded_coords = np.array([
range(123),
range(0, 246, 2)
])
excluded_pixels = sparse.COO(coords=excluded_coords, shape=sig_shape, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
adjusted = corr.adjust_tileshape(
tile_shape=tile_shape, sig_shape=sig_shape, base_shape=base_shape
)
assert adjusted == (123, 256)
_validate(excluded_coords=excluded_coords, adjusted=adjusted, sig_shape=sig_shape)
def test_tileshape_adjustment_6_1():
sig_shape = (123, 456)
tile_shape = (122, 1)
base_shape = (1, 1)
excluded_coords = np.array([
range(123),
np.zeros(123, dtype=int)
])
excluded_pixels = sparse.COO(coords=excluded_coords, shape=sig_shape, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
adjusted = corr.adjust_tileshape(
tile_shape=tile_shape, sig_shape=sig_shape, base_shape=base_shape
)
print(adjusted)
assert adjusted == (123, 2)
_validate(excluded_coords=excluded_coords, adjusted=adjusted, sig_shape=sig_shape)
def test_tileshape_adjustment_8():
sig_shape = (1014, 1024)
tile_shape = (1, 1)
base_shape = (1, 1)
# These magic numbers are "worst case" to produce collisions
# 2*3*4*5*6*7
excluded_coords = np.array([
(720, 210, 306),
(120, 210, 210)
])
excluded_pixels = sparse.COO(coords=excluded_coords, shape=sig_shape, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
adjusted = corr.adjust_tileshape(
tile_shape=tile_shape, sig_shape=sig_shape, base_shape=base_shape
)
print(adjusted)
assert adjusted != (1014, 1024)
_validate(excluded_coords=excluded_coords, adjusted=adjusted, sig_shape=sig_shape)
def test_tileshape_adjustment_fuzz():
for n in range(10):
sig_shape = (np.random.randint(1, 2**12), np.random.randint(1, 2**12))
print("Sig shape", sig_shape)
tile_shape = (1, 1)
base_shape = (1, 1)
size = max(1, max(sig_shape) // 10)
excluded_coords = np.vstack([
np.random.randint(0, sig_shape[0], size=size),
np.random.randint(0, sig_shape[1], size=size),
])
print("excluded_coords", excluded_coords.shape, excluded_coords)
excluded_pixels = sparse.COO(coords=excluded_coords, shape=sig_shape, data=True)
corr = CorrectionSet(excluded_pixels=excluded_pixels)
adjusted = corr.adjust_tileshape(
tile_shape=tile_shape, sig_shape=sig_shape, base_shape=base_shape
)
print(adjusted)
_validate(excluded_coords=excluded_coords, adjusted=adjusted, sig_shape=sig_shape)
def get_scheme(self,
udfs,
partition,
read_dtype: np.dtype,
roi: np.ndarray,
corrections: CorrectionSet = None):
"""
Generate a :class:`TilingScheme` instance that is
compatible with both the given `udf` and the
:class:~`libertem.io.dataset.base.DataSet`.
Parameters
----------
udfs : List[UDF]
The concrete UDF to optimize the tiling scheme for.
Depending on the method (tile, frame, partition)
and preferred total input size and depth.
partition : Partition
The `TilingScheme` is created specifically
for the given `Partition`, so it can adjust
even in the face of different partition sizes/shapes.
read_dtype
The dtype in which the data will be fed into the UDF
roi : np.ndarray
Region of interest
corrections : CorrectionSet
Correction set to consider in negotiation
"""
itemsize = np.dtype(read_dtype).itemsize
# FIXME: itemsize != native_dtype.itemsize! use partition.meta.raw_dtype.itemsize?
# try not to waste page faults:
# FIXME: let the UDF define upper bound for signal size (lower bound, too?)
# (signal buffers should fit into the L2 cache)
min_sig_size = 4 * 4096 // itemsize
# This already takes corrections into account through a different pathway
need_decode = partition.need_decode(roi=roi, read_dtype=read_dtype)
if need_decode:
io_max_size = 1 * 2**20
else:
io_max_size = itemsize * np.prod(partition.shape, dtype=np.int64)
depths = [self._get_min_depth(udf, partition) for udf in udfs]
depth = max(depths) # take the largest min-depth
base_shape = self._get_base_shape(udfs, partition)
sizes = [
self._get_size(
io_max_size,
udf,
itemsize,
partition,
base_shape,
) for udf in udfs
]
if any(udf.get_method() == "partition" for udf in udfs):
size = max(sizes) # by partition wants to be big, ...
else:
size = min(sizes)
size_px = size // itemsize
if corrections is not None and corrections.have_corrections():
# The correction has to make sure that there are no excluded pixels
# at tile boundaries
base_shape = corrections.adjust_tileshape(
tile_shape=base_shape,
sig_shape=tuple(partition.shape.sig),
base_shape=base_shape,
)
# first, scale `base_shape` up to contain at least `min_sig_size` items:
min_factors = self._get_scale_factors(
base_shape,
containing_shape=partition.shape.sig,
size=min_sig_size,
)
min_base_shape = self._scale_base_shape(base_shape, min_factors)
# considering the min size, calculate the max depth:
max_depth = size_px // np.prod(min_base_shape, dtype=np.int64)
if depth > max_depth:
depth = max_depth
full_base_shape = (1, ) + tuple(base_shape)
min_factors = (depth, ) + tuple(min_factors)
factors = self._get_scale_factors(
full_base_shape,
containing_shape=partition.shape,
size=size_px,
min_factors=min_factors,
)
tileshape = self._scale_base_shape(full_base_shape, factors)
# the partition has a "veto" on the tileshape:
tileshape = partition.adjust_tileshape(tileshape)
self.validate(tileshape, partition, size, itemsize, full_base_shape)
return TilingScheme.make_for_shape(
tileshape=Shape(tileshape, sig_dims=partition.shape.sig.dims),
dataset_shape=partition.meta.shape,
debug={
"min_factors": min_factors,
"factors": factors,
"tileshape": tileshape,
"size": size,
"size_px": size_px,
"full_base_shape": full_base_shape,
"need_decode": need_decode,
"depth": depth,
})