def test_map_overlap_multiarray():
# Same ndim, same numblocks, same chunks
x = da.arange(10, chunks=5)
y = da.arange(10, chunks=5)
z = da.map_overlap(lambda x, y: x + y, x, y, depth=1)
assert_eq(z, 2 * np.arange(10))
# Same ndim, same numblocks, different chunks
x = da.arange(10, chunks=(2, 3, 5))
y = da.arange(10, chunks=(5, 3, 2))
z = da.map_overlap(lambda x, y: x + y, x, y, depth=1)
assert z.chunks == ((2, 3, 3, 2), )
assert_eq(z, 2 * np.arange(10))
# Same ndim, different numblocks, different chunks
x = da.arange(10, chunks=(10, ))
y = da.arange(10, chunks=(4, 4, 2))
z = da.map_overlap(lambda x, y: x + y, x, y, depth=1)
assert z.chunks == ((4, 4, 2), )
assert_eq(z, 2 * np.arange(10))
# Different ndim, different numblocks, different chunks
x = da.arange(10, chunks=(10, ))
y = da.arange(10).reshape(1, 10).rechunk((1, (4, 4, 2)))
z = da.map_overlap(lambda x, y: x + y, x, y, depth=1)
assert z.chunks == ((1, ), (4, 4, 2))
assert z.shape == (1, 10)
assert_eq(z, 2 * np.arange(10)[np.newaxis])
def test_map_overlap_multiarray_uneven_numblocks_exception():
x = da.arange(10, chunks=(10,))
y = da.arange(10, chunks=(5, 5))
with pytest.raises(ValueError):
# Fail with chunk alignment explicitly disabled
da.map_overlap(
lambda x, y: x + y, x, y, align_arrays=False, boundary="none"
).compute()
def test_map_overlap_assumes_shape_matches_first_array_if_trim_is_false():
# https://github.com/dask/dask/issues/6681
x1 = da.ones((10, ), chunks=(5, 5))
x2 = x1.rechunk(10)
def oversum(x):
return x[2:-2]
z1 = da.map_overlap(oversum, x1, depth=2, trim=False)
assert z1.shape == (10, )
z2 = da.map_overlap(oversum, x2, depth=2, trim=False)
assert z2.shape == (10, )
def test_map_overlap_multiarray_variadic():
# Test overlapping row slices from 3D arrays
xs = [
# Dim 0 will unify to chunks of size 4 for all:
da.ones((12, 1, 1), chunks=((12,), 1, 1)),
da.ones((12, 8, 1), chunks=((8, 4), 8, 1)),
da.ones((12, 8, 4), chunks=((4, 8), 8, 4)),
]
def func(*args):
return np.array([sum(x.size for x in args)])
x = da.map_overlap(
func,
*xs,
chunks=(1,),
depth=1,
trim=False,
drop_axis=[1, 2],
boundary="reflect",
)
# Each func call should get 4 rows from each array padded by 1 in each dimension
size_per_slice = sum(np.pad(x[:4], 1, mode="constant").size for x in xs)
assert x.shape == (3,)
assert all(x.compute() == size_per_slice)
def run(depth):
return da.map_overlap(lambda x, y: x.sum() + y.sum(),
x,
y,
depth=depth,
chunks=(0, ),
trim=False).compute()
def stitch_blocks(blocks, blocksize, overlap):
"""
"""
# blocks may be a vector fields
# conditional is too general, but works for now
if blocks.ndim != len(blocksize):
blocksize = list(blocksize) + [3,]
overlap = tuple(overlap) + (0,)
# weight block edges
weighted_blocks = da.map_blocks(
weight_block, blocks,
blocksize=blocksize[:3],
overlap=overlap[:3],
dtype=np.float32,
)
# stitch overlap regions, return
return da.map_overlap(
merge_overlaps, weighted_blocks,
overlap=overlap[:3],
depth=overlap,
boundary=0.,
trim=False,
dtype=np.float32,
chunks=blocksize,
)
def _process_dask(raster, xs, ys):
if max_distance >= max_possible_distance:
# consider all targets in the whole raster
# the data array is computed at once,
# make sure your data fit your memory
height, width = raster.shape
raster.data = raster.data.rechunk({0: height, 1: width})
xs = xs.rechunk({0: height, 1: width})
ys = ys.rechunk({0: height, 1: width})
pad_y = pad_x = 0
else:
cellsize_x, cellsize_y = get_dataarray_resolution(raster)
# calculate padding for each chunk
pad_y = int(max_distance / cellsize_y + 0.5)
pad_x = int(max_distance / cellsize_x + 0.5)
out = da.map_overlap(
_process_numpy,
raster.data, xs, ys,
depth=(pad_y, pad_x),
boundary=np.nan,
meta=np.array(()),
)
return out
def stitch_fields(fields, blocksize):
"""
"""
# weight block edges
weighted_fields = da.map_blocks(
weight_block,
fields,
blocksize=blocksize,
dtype=np.float32,
)
# remove block index dimensions
sh = fields.shape[:3]
list_of_blocks = [[[[weighted_fields[i, j, k]] for k in range(sh[2])]
for j in range(sh[1])] for i in range(sh[0])]
aug_fields = da.block(list_of_blocks)
# merge overlap regions
overlaps = tuple([int(round(x / 8)) for x in blocksize] + [
0,
])
return da.map_overlap(
merge_overlaps,
aug_fields,
blocksize=blocksize,
depth=overlaps,
boundary=0.,
trim=False,
dtype=np.float32,
chunks=blocksize + [
3,
],
)
def prune(X,
pos,
window,
step,
threshold,
align_chunks=True,
short_circuit=False,
windows_per_chunk=10,
numba=False,
compute=True):
assert step < window
assert window % step == 0
assert X.ndim == 2
# Gather information defining how the array should be chunked
chunk_size = windows_per_chunk * window
chunk_info = get_chunk_info(pos, chunk_size)
# Create list of individual chunk sizes, which will almost always be uneven
# due to breaks at contig boundaries
chunk_lens = tuple(cs for ci in chunk_info.chunks for cs in ci.chunk_size)
# Rechunk the provided array to match assumptions made by this method
if not align_chunks and X.chunks[0] != chunk_lens:
raise ValueError(f'Expected chunks {chunk_lens}, found {X.chunks[0]}')
if align_chunks:
X = X.rechunk(chunks=(chunk_lens, X.chunks[1]))
# Determine how many rows should be shared in each block calculation
overlap_depth = window - step
fn = jit(_prune, nopython=True, nogil=True) if numba else _prune
R = da.map_overlap(
X,
fn,
window=window,
step=step,
overlap_depth=overlap_depth,
depth=(overlap_depth, 0), # No overlap in axis 1
threshold=threshold,
boundary=-1,
short_circuit=short_circuit,
# Use tuples because numbda dict results in this when pickled for serialized tasks:
# TypeError: can't pickle _nrt_python._MemInfo objects"
contig_boundary=tuple(chunk_info.get_contig_chunk_boundary().items()),
chunk_offset=tuple(chunk_info.get_chunk_offset().items()),
short_circuit_step_rate=2,
# TODO: What are the consequences of providing chunk row size that is <= result here?
chunks=([v for v in X.chunks[0]], 2),
dtype=np.float64,
trim=False,
return_ld_matrix=False)
# Result contains rows indices to keep in first column (block index in second)
if compute:
R = R.compute()
Xp = X[np.unique(R[:, 0]).astype(int)]
return Xp, (X, R, chunk_info, overlap_depth)
return R, (X, chunk_info, overlap_depth)
def test_map_overlap_multiarray_defaults():
# Check that by default, chunk alignment and arrays of varying dimensionality
# are supported by with no effect on result shape
# (i.e. defaults are pass-through to map_blocks)
x = da.ones((10, ), chunks=10)
y = da.ones((1, 10), chunks=5)
z = da.map_overlap(lambda x, y: x + y, x, y)
# func should be called twice and get (5,) and (1, 5) arrays of ones each time
assert_eq(z.shape, (1, 10))
assert_eq(z.sum(), 20.0)
def test_map_overlap_deprecated_signature():
def func(x):
return np.array(x.sum())
x = da.ones(3)
# Old positional signature: func, depth, boundary, trim
with pytest.warns(FutureWarning):
y = da.map_overlap(x, func, 0, "reflect", True)
assert y.compute() == 3
assert y.shape == (3, )
with pytest.warns(FutureWarning):
y = da.map_overlap(x, func, 1, "reflect", True)
assert y.compute() == 5
assert y.shape == (3, )
with pytest.warns(FutureWarning):
y = da.map_overlap(x, func, 1, "reflect", False)
assert y.compute() == 5
assert y.shape == (3, )
def dog_ransac_affine_distributed(
fixed, moving,
fixed_vox, moving_vox,
cc_radius,
nspots,
match_threshold,
align_threshold,
blocksize=[256,]*3,
cluster_extra=["-P multifish",]
):
"""
"""
# get number of blocks required
block_grid = np.ceil(np.array(fixed.shape) / blocksize).astype(int)
nblocks = np.prod(block_grid)
# distributed computations done in cluster context
# TODO: generalize w.r.t. workstations and cluster managers
with distributed.distributedState() as ds:
# set up the cluster
ds.initializeLSFCluster(job_extra=cluster_extra)
ds.initializeClient()
ds.scaleCluster(njobs=nblocks)
# wrap images as dask arrays
fixed_da = da.from_array(fixed, chunks=blocksize)
moving_da = da.from_array(moving, chunks=blocksize)
# wrap affine function
def my_dog_ransac_affine(x, y):
affine = dog_ransac_affine(
x, y, fixed_vox, moving_vox,
cc_radius, nspots, match_threshold, align_threshold,
)
return affine.reshape((1,1,1,3,4))
# affine align all chunks
affines = da.map_overlap(
my_dog_ransac_affine, fixed_da, moving_da,
depth=tuple([int(round(x/8)) for x in blocksize]),
boundary='reflect',
trim=False,
align_arrays=False,
dtype=np.float64,
new_axis=[3, 4],
chunks=[1, 1, 1, 3, 4],
).compute()
return interpolate_affines(affines)
def lr_deconvolution(image,psf,iterations=50):
"""
Tiled Lucy-Richardson deconvolution using DECON_LIBRARY
:param image: ndarray
raw data
:param psf: ndarray
theoretical PSF
:param iterations: int
number of iterations to run
:return deconvolved: ndarray
deconvolved image
"""
# create dask array
scan_chunk_size = 512
if image.shape[0]<scan_chunk_size:
dask_raw = da.from_array(image,chunks=(image.shape[0],image.shape[1],image.shape[2]))
overlap_depth = (0,2*psf.shape[1],2*psf.shape[1])
else:
dask_raw = da.from_array(image,chunks=(scan_chunk_size,image.shape[1],image.shape[2]))
overlap_depth = 2*psf.shape[0]
del image
gc.collect()
if DECON_LIBRARY=='dexp':
# define dask dexp partial function for GPU LR deconvolution
lr_dask = partial(dexp_lr_decon,psf=psf,num_iterations=iterations,padding=2*psf.shape[0],internal_dtype=np.float16)
else:
lr_dask = partial(mv_lr_decon,psf=psf,num_iterations=iterations)
# create dask plan for overlapped blocks
dask_decon = da.map_overlap(lr_dask,dask_raw,depth=overlap_depth,boundary=None,trim=True,meta=np.array((), dtype=np.uint16))
# perform LR deconvolution in blocks
if DECON_LIBRARY=='dexp':
with CupyBackend(enable_cutensor=True,enable_cub=True,enable_fft_planning=True):
with ProgressBar():
decon_data = dask_decon.compute(scheduler='single-threaded')
else:
with ProgressBar():
decon_data = dask_decon.compute(scheduler='single-threaded')
# clean up memory
cp.clear_memo()
del dask_decon
gc.collect()
return decon_data.astype(np.uint16)
def life_step_dask(x):
def access_roll(x):
#da.roll extracts an array where all values are shifted left, right, up or down
l_roll = da.roll(x, 1, axis=0)
r_roll = da.roll(x, -1, axis=0)
return l_roll + r_roll + da.roll(x,1,axis=1) + da.roll(x,-1,axis=1) + \
da.roll(l_roll, 1, axis=1) + da.roll(l_roll, -1, axis=1) + \
da.roll(r_roll, 1, axis=1) + da.roll(r_roll, -1, axis=1)
c_grid = da.map_overlap(x, access_roll, depth=1, boundary="periodic")
nx = x - ((x == 1) & ((c_grid < 2) | (c_grid > 3))).astype(np.int)
nx = nx + ((x == 0) & (c_grid == 3)).astype(np.int)
return nx
def test_map_overlap_multiarray_block_broadcast():
def func(x, y):
# Return result with expected padding
z = x.size + y.size
return np.ones((3, 3)) * z
# Chunks in trailing dimension will be unified to two chunks of size 6
# and block broadcast will allow chunks from x to repeat
x = da.ones((12,), chunks=12) # numblocks = (1,) -> (2, 2) after broadcast
y = da.ones((16, 12), chunks=(8, 6)) # numblocks = (2, 2)
z = da.map_overlap(func, x, y, chunks=(3, 3), depth=1, trim=True)
assert_eq(z, z)
assert z.shape == (2, 2)
# func call will receive (8,) and (10, 8) arrays for each of 4 blocks
assert_eq(z.sum(), 4 * (10 * 8 + 8))
def test_map_overlap_trim_using_drop_axis_and_different_depths(drop_axis):
x = da.random.standard_normal((5, 10, 8), chunks=(2, 5, 4))
def _mean(x):
return x.mean(axis=drop_axis)
expected = _mean(x)
# unique boundary and depth value per axis
boundary = (0, "reflect", "nearest")
depth = (1, 3, 2)
# to match expected result, dropped axes must have depth 0
_drop_axis = (drop_axis,) if np.isscalar(drop_axis) else drop_axis
depth = tuple(0 if i in _drop_axis else d for i, d in enumerate(depth))
y = da.map_overlap(
_mean, x, depth=depth, boundary=boundary, drop_axis=drop_axis, dtype=float
).compute()
assert_array_almost_equal(expected, y)
def get_poss_bergs_fr_raster(onedem, usedask):
# trans=onedem.attrs['transform']
flipax = []
# if trans[0] < 0:
# flipax.append(1)
# if trans[4] < 0:
# flipax.append(0)
if pd.Series(onedem.x).is_monotonic_decreasing:
flipax.append(1)
if pd.Series(onedem.y).is_monotonic_increasing:
flipax.append(0)
fjord = onedem.attrs['fjord']
min_area = fjord_props.get_min_berg_area(fjord)
res = onedem.attrs['res'][
0] #Note: the pixel area will be inaccurate if the resolution is not the same in x and y
if usedask == True:
# Daskify the iceberg segmentation process. Note that dask-image has some functionality to operate
# directly on dask arrays (e.g. dask_image.ndfilters.sobel), which would need to be put into utils.raster.py
# https://dask-image.readthedocs.io/en/latest/dask_image.ndfilters.html
# However, as of yet there doesn't appear to be a way to easily implement the watershed segmentation, other than in chunks
# print(onedem)
# see else statement with non-dask version for descriptions of what each step is doing
def seg_wrapper(tiles):
return raster_ops.labeled_from_segmentation(tiles, [3, 10],
resolution=res,
min_area=min_area,
flipax=[])
def filter_wrapper(tiles, elevs):
return raster_ops.border_filtering(tiles, elevs, flipax=[])
elev_copy = onedem.elevation.data # should return a dask array
for ax in flipax:
elev_copy = da.flip(elev_copy, axis=ax)
# print(type(elev_copy))
elev_overlap = da.overlap.overlap(elev_copy,
depth=10,
boundary='nearest')
seglabeled_overlap = da.map_overlap(
seg_wrapper, elev_overlap,
trim=False) # including depth=10 here will ADD another overlap
print("Got labeled raster of potential icebergs for an image")
labeled_overlap = da.map_overlap(filter_wrapper,
seglabeled_overlap,
elev_overlap,
trim=False,
dtype='int32')
labeled_arr = da.overlap.trim_overlap(labeled_overlap, depth=10)
# re-flip the labeled_arr so it matches the orientation of the original elev data that's within the xarray
for ax in flipax:
labeled_arr = da.flip(labeled_arr, axis=ax)
# import matplotlib.pyplot as plt
# print(plt.imshow(labeled_arr))
try:
del elev_copy
del elev_overlap
del seglabeled_overlap
del labeled_overlap
print("deleted the intermediate steps")
except NameError:
pass
# print(da.min(labeled_arr).compute())
# print(da.max(labeled_arr).compute())
print("about to get the list of possible bergs")
print(
'Please note the transform computation is very application specific (negative y coordinates) and may need generalizing'
)
print(
"this transform computation is particularly sensitive to axis order (y,x) because it is accessed by index number"
)
poss_bergs_list = []
'''
# I think that by using concatenate=True, it might not actually be using dask for the computation
def get_bergs(labeled_blocks):
# Note: features.shapes returns a generator. However, if we try to iterate through it with a for loop, the StopIteration exception
# is not passed up into the for loop and execution hangs when it hits the end of the for loop without completing the function
block_bergs = list(poly[0]['coordinates'][0] for poly in rasterio.features.shapes(
labeled_blocks.astype('int32'), transform=onedem.attrs['transform']))[:-1]
poss_bergs_list.append(block_bergs)
da.blockwise(get_bergs, '', labeled_arr, 'ij',
meta=pd.DataFrame({'c':[]}), concatenate=True).compute()
# print(poss_bergs_list[0])
# print(type(poss_bergs_list))
poss_bergs_gdf = gpd.GeoDataFrame({'geometry':[Polygon(poly) for poly in poss_bergs_list[0]]})
# another approach could be to try and coerce the output from map_blocks into an array, but I suspect you'd still have the geospatial issue
# https://github.com/dask/dask/issues/3590#issuecomment-464609620
'''
# URL: https://stackoverflow.com/questions/66232232/produce-vector-output-from-a-dask-array/66245347?noredirect=1#comment117119583_66245347
@dask.delayed
def get_bergs(labeled_blocks, pointer, chunk0, chunk1):
print("running the dask delayed function")
def getpx(chunkid, chunksz):
amin = chunkid[0] * chunksz[0][0]
amax = amin + chunksz[0][0]
bmin = chunkid[1] * chunksz[1][0]
bmax = bmin + chunksz[1][0]
return (amin, amax, bmin, bmax)
# order of all inputs (and outputs) should be y, x when axis order is used
chunksz = (onedem.chunks['y'], onedem.chunks['x'])
# rasterio_trans = rasterio.transform.guard_transform(onedem.attrs["transform"])
# print(rasterio_trans)
ymini, ymaxi, xmini, xmaxi = getpx((chunk0, chunk1), chunksz)
# print(chunk0, chunk1)
# print(xmini)
# print(xmaxi)
# print(ymini)
# print(ymaxi)
# use rasterio Windows and rioxarray to construct transform
# https://rasterio.readthedocs.io/en/latest/topics/windowed-rw.html#window-transforms
chwindow = rasterio.windows.Window(xmini, ymini, xmaxi - xmini,
ymaxi - ymini)
trans = onedem.rio.isel_window(chwindow).rio.transform(recalc=True)
# print(trans)
return list(
poly[0]['coordinates'][0] for poly in rasterio.features.shapes(
labeled_blocks.astype('int32'), transform=trans))[:-1]
for __, obj in enumerate(labeled_arr.to_delayed()):
for bl in obj:
piece = dask.delayed(get_bergs)(bl, *bl.key)
poss_bergs_list.append(piece)
del piece
poss_bergs_list = dask.compute(*poss_bergs_list)
# tried working with this instead of the for loops above
# poss_bergs_list = dask.compute([get_bergs(bl, *bl.key) for bl in obj for __, obj in enumerate(labeled_arr.to_delayed())])[0]
# print(poss_bergs_list)
# unnest the list of polygons returned by using dask to polygonize
concat_list = [
item for sublist in poss_bergs_list for item in sublist
if len(item) != 0
]
# print(concat_list)
poss_bergs_gdf = gpd.GeoDataFrame(
{'geometry': [Polygon(poly) for poly in concat_list]})
# convert to a geodataframe, combine geometries (in case any bergs were on chunk borders), and generate new polygon list
print(poss_bergs_gdf)
# print(poss_bergs_gdf.geometry.plot())
poss_berg_combined = gpd.overlay(poss_bergs_gdf,
poss_bergs_gdf,
how='union')
# print(poss_berg_combined)
# print(poss_berg_combined.geometry.plot())
poss_bergs = [berg for berg in poss_berg_combined.geometry]
# print(poss_bergs)
print(len(poss_bergs))
try:
del labeled_arr
del poss_bergs_list
del concat_list
del poss_berg_combined
except NameError:
pass
else:
print("NOT USING DASK")
# create copy of elevation values so original dataset values are not impacted by image manipulations
# and positive/negative coordinate systems can be ignored (note flipax=[] below)
# something wonky is happening and when I ran this code on Pangeo I needed to NOT flip the elevation values here and then switch the bounding box y value order
# Not entirely sure what's going on, but need to be aware of this!!
# print("Note: check for proper orientation of results depending on compute environment. Pangeo results were upside down.")
elev_copy = np.copy(np.flip(onedem.elevation.values, axis=flipax))
# flipax=[]
# generate a labeled array of potential iceberg features, excluding those that are too large or small
seglabeled_arr = raster_ops.labeled_from_segmentation(
elev_copy, [3, 10], resolution=res, min_area=min_area, flipax=[])
print("Got labeled raster of potential icebergs for an image")
# remove features whose borders are >50% no data values (i.e. the "iceberg" edge is really a DEM edge)
labeled_arr = raster_ops.border_filtering(seglabeled_arr,
elev_copy,
flipax=[]).astype(
seglabeled_arr.dtype)
# apparently rasterio can't handle int64 inputs, which is what border_filtering returns
# import matplotlib.pyplot as plt
# print(plt.imshow(labeled_arr))
# create iceberg polygons
# somehow a < 1 pixel berg made it into this list... I'm doing a secondary filtering by area in the iceberg filter step for now
poss_bergs = list(
poly[0]['coordinates'][0] for poly in rasterio.features.shapes(
labeled_arr, transform=onedem.attrs['transform']))[:-1]
try:
del elev_copy
del seglabeled_arr
del labeled_arr
except NameError:
pass
return poss_bergs
def prepare_piecewise_ransac_affine(
fix,
mov,
fix_spacing,
mov_spacing,
min_radius,
max_radius,
match_threshold,
blocksize,
**kwargs,
):
"""
"""
# get number of blocks required
block_grid = np.ceil(np.array(fix.shape) / blocksize).astype(int)
nblocks = np.prod(block_grid)
overlap = [int(round(x / 8)) for x in blocksize]
# wrap images as dask arrays
fix_da = da.from_array(fix, chunks=blocksize)
mov_da = da.from_array(mov, chunks=blocksize)
# wrap affine function
def wrapped_ransac_affine(x, y, block_info=None):
# compute affine
affine = ransac_affine(
x,
y,
fix_spacing,
mov_spacing,
min_radius,
max_radius,
match_threshold,
**kwargs,
)
# adjust for block origin
idx = np.array(block_info[0]['chunk-location'])
origin = (idx * blocksize - overlap) * fix_spacing
tl, tr = np.eye(4), np.eye(4)
tl[:3, -1], tr[:3, -1] = origin, -origin
affine = np.matmul(tl, np.matmul(affine, tr))
# return with block index axes
return affine.reshape((1, 1, 1, 4, 4))
# affine align all chunks
return da.map_overlap(
wrapped_ransac_affine,
fix_da,
mov_da,
depth=tuple(overlap),
boundary='reflect',
trim=False,
align_arrays=False,
dtype=np.float64,
new_axis=[3, 4],
chunks=[1, 1, 1, 4, 4],
)
def diff_overlap(a):
return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
def diff_overlap(
a: Annotated[np.ndarray, "X:center"]
) -> Annotated[np.ndarray, "X:left"]:
return map_overlap(diff_center_to_left, a, depth=1, boundary="periodic")
def distributed_piecewise_alignment_pipeline(
fix,
mov,
fix_spacing,
mov_spacing,
nblocks,
overlap=0.5,
fix_mask=None,
mov_mask=None,
steps=['rigid', 'affine'],
random_kwargs={},
rigid_kwargs={},
affine_kwargs={},
deform_kwargs={},
cluster=None,
cluster_kwargs={},
**kwargs,
):
"""
Piecewise affine alignment of moving to fixed image.
Overlapping blocks are given to `affine_align` in parallel
on distributed hardware. Can include random initialization,
rigid alignment, and affine alignment.
Parameters
----------
fix : ndarray
the fixed image
mov : ndarray
the moving image; `fix.shape` must equal `mov.shape`
I.e. typically piecewise affine alignment is done after
a global affine alignment wherein the moving image has
been resampled onto the fixed image voxel grid.
fix_spacing : 1d array
The spacing in physical units (e.g. mm or um) between voxels
of the fixed image.
Length must equal `fix.ndim`
mov_spacing : 1d array
The spacing in physical units (e.g. mm or um) between voxels
of the moving image.
Length must equal `mov.ndim`
nblocks : iterable
The number of blocks to use along each axis.
Length should be equal to `fix.ndim`
overlap : float in range [0, 1] (default: 0.5)
Block overlap size as a percentage of block size
fix_mask : binary ndarray (default: None)
A mask limiting metric evaluation region of the fixed image
mov_mask : binary ndarray (default: None)
A mask limiting metric evaluation region of the moving image
Due to the distribution aspect, if a mov_mask is provided
you must also provide a fix_mask. A reasonable choice if
no fix_mask exists is an array of all ones.
steps : list of type string (default: ['rigid', 'affine'])
Flags to indicate which steps to run. An empty list will guarantee
all affines are the identity. Any of the following may be in the list:
'random': run `random_affine_search` first
'rigid': run `affine_align` with rigid=True
'affine': run `affine_align` with rigid=False
If all steps are present they are run in the order given above.
Steps share parameters given to kwargs. Parameters for individual
steps override general settings with `random_kwargs`, `rigid_kwargs`,
and `affine_kwargs`. If `random` is in the list, `random_kwargs`
must be defined.
random_kwargs : dict (default: {})
Keyword arguments to pass to `random_affine_search`. This is only
necessary if 'random' is in `steps`. If so, the following keys must
be given:
'max_translation'
'max_rotation'
'max_scale'
'max_shear'
'random_iterations'
However any argument to `random_affine_search` may be defined. See
documentation for `random_affine_search` for descriptions of these
parameters. If 'random' and 'rigid' are both in `steps` then
'max_scale' and 'max_shear' must both be 0.
rigid_kwargs : dict (default: {})
If 'rigid' is in `steps`, these keyword arguments are passed
to `affine_align` during the rigid=True step. They override
any common general kwargs.
affine_kwargs : dict (default: {})
If 'affine' is in `steps`, these keyword arguments are passed
to `affine_align` during the rigid=False (affine) step. They
override any common general kwargs.
cluster_kwargs : dict (default: {})
Arguments passed to ClusterWrap.cluster
If working with an LSF cluster, this will be
ClusterWrap.janelia_lsf_cluster. If on a workstation
this will be ClusterWrap.local_cluster.
This is how distribution parameters are specified.
kwargs : any additional arguments
Passed to calls `random_affine_search` and `affine_align` calls
Returns
-------
affines : nd array
Affine matrix for each block. Shape is (X, Y, ..., 4, 4)
for X blocks along first axis and so on.
field : nd array
Local affines stitched together into a displacement field
Shape is `fix.shape` + (3,) as the last dimension contains
the displacement vector.
"""
# wait for at least one worker to be fully instantiated
while ((cluster.client.status == "running")
and (len(cluster.client.scheduler_info()["workers"]) < 1)):
time.sleep(1.0)
# compute block size and overlaps
blocksize = np.array(fix.shape).astype(np.float32) / nblocks
blocksize = np.ceil(blocksize).astype(np.int16)
overlaps = tuple(np.round(blocksize * overlap).astype(np.int16))
# pad the ends to fill in the last blocks
# blocks must all be exact for stitch to work correctly
pads = [(0, y - x % y) if x % y > 0 else (0, 0)
for x, y in zip(fix.shape, blocksize)]
fix_p = np.pad(fix, pads)
mov_p = np.pad(mov, pads)
# pad masks if necessary
if fix_mask is not None:
fm_p = np.pad(fix_mask, pads)
if mov_mask is not None:
mm_p = np.pad(mov_mask, pads)
# CONSTRUCT DASK ARRAY VERSION OF OBJECTS
# fix
fix_future = cluster.client.scatter(fix_p)
fix_da = da.from_delayed(fix_future, shape=fix_p.shape,
dtype=fix_p.dtype).rechunk(tuple(blocksize))
# mov
mov_future = cluster.client.scatter(mov_p)
mov_da = da.from_delayed(mov_future, shape=mov_p.shape,
dtype=mov_p.dtype).rechunk(tuple(blocksize))
# fix mask
if fix_mask is not None:
fm_future = cluster.client.scatter(fm_p)
fm_da = da.from_delayed(fm_future, shape=fm_p.shape,
dtype=fm_p.dtype).rechunk(tuple(blocksize))
else:
fm_da = da.from_array([
None,
], chunks=(1, ))
# mov mask
if mov_mask is not None:
mm_future = cluster.client.scatter(mm_p)
mm_da = da.from_delayed(mm_future, shape=mm_p.shape,
dtype=mm_p.dtype).rechunk(tuple(blocksize))
else:
mm_da = da.from_array([
None,
], chunks=(1, ))
# establish all keyword arguments
random_kwargs = {**kwargs, **random_kwargs}
rigid_kwargs = {**kwargs, **rigid_kwargs}
affine_kwargs = {**kwargs, **affine_kwargs}
deform_kwargs = {**kwargs, **deform_kwargs}
# closure for alignment pipeline
def align_single_block(fix, mov, fm, mm, block_info=None):
# check for masks
if fm.shape != fix.shape: fm = None
if mm.shape != mov.shape: mm = None
# run alignment pipeline
transform = alignment_pipeline(
fix,
mov,
fix_spacing,
mov_spacing,
steps,
fix_mask=fm,
mov_mask=mm,
random_kwargs=random_kwargs,
rigid_kwargs=rigid_kwargs,
affine_kwargs=affine_kwargs,
deform_kwargs=deform_kwargs,
)
# convert to single vector field
if isinstance(transform, tuple):
affine, deform = transform[0], transform[1][1]
transform = compose_affine_and_displacement_vector_field(
affine,
deform,
fix_spacing,
)
else:
transform = ut.matrix_to_displacement_field(
fix,
transform,
fix_spacing,
)
# get block index and block grid size
block_grid = block_info[0]['num-chunks']
block_index = block_info[0]['chunk-location']
# create weights array
core, pad_ones, pad_linear = [], [], []
for i in range(3):
# get core shape and pad sizes
o = max(0, 2 * overlaps[i] - 1)
c = blocksize[i] - o + 1
p_ones, p_linear = [0, 0], [o, o]
if block_index[i] == 0:
p_ones[0], p_linear[0] = o // 2, 0
if block_index[i] == block_grid[i] - 1:
p_ones[1], p_linear[1] = o // 2, 0
core.append(c)
pad_ones.append(tuple(p_ones))
pad_linear.append(tuple(p_linear))
# create weights core
weights = np.ones(core, dtype=np.float32)
# extend weights
weights = np.pad(
weights,
pad_ones,
mode='constant',
constant_values=1,
)
weights = np.pad(
weights,
pad_linear,
mode='linear_ramp',
end_values=0,
)
# apply weights
transform = transform * weights[..., None]
# package and return result
result = np.empty((1, 1, 1), dtype=object)
result[0, 0, 0] = transform
return result
# END CLOSURE
# determine overlaps list
overlaps_list = [overlaps, overlaps]
if fix_mask is not None:
overlaps_list.append(overlaps)
else:
overlaps_list.append(0)
if mov_mask is not None:
overlaps_list.append(overlaps)
else:
overlaps_list.append(0)
# align all chunks
fields = da.map_overlap(
align_single_block,
fix_da,
mov_da,
fm_da,
mm_da,
depth=overlaps_list,
dtype=object,
boundary='none',
trim=False,
align_arrays=False,
chunks=[1, 1, 1],
).compute()
# reconstruct single transform
transform = np.zeros(fix_p.shape + (3, ), dtype=np.float32)
# adds fields to transform
block_grid = fields.shape[:3]
for ijk in range(np.prod(block_grid)):
i, j, k = np.unravel_index(ijk, block_grid)
x, y, z = (max(0, x * y - z)
for x, y, z in zip((i, j, k), blocksize, overlaps))
xr, yr, zr = fields[i, j, k].shape[:3]
transform[x:x + xr, y:y + yr, z:z + zr] += fields[i, j, k][...]
# crop back to original shape and return
x, y, z = fix.shape
return transform[:x, :y, :z]
def tiled_deformable_align(
fixed,
moving,
fixed_spacing,
moving_spacing,
blocksize,
transpose=[False] * 2,
global_affine=None,
local_affines=None,
write_path=None,
lazy=True,
deform_kwargs={},
# cluster_kwargs={},
):
"""
"""
# get number of blocks required
block_grid = np.ceil(np.array(fixed.shape) / blocksize)
nblocks = np.prod(block_grid)
# get true field shape
original_shape = fixed.shape
if transpose[0]:
original_shape = original_shape[::-1]
# get affine position field
affine_pf = None
if global_affine is not None or local_affines is not None:
if local_affines is None:
local_affines = np.empty(
block_grid + (3, 4),
dtype=np.float32,
)
local_affines[..., :, :] = np.eye(4)[:3, :]
affine_pf = transform.local_affines_to_position_field(
original_shape,
fixed_spacing,
blocksize,
local_affines,
global_affine=global_affine,
lazy=True,
#cluster_kwargs=cluster_kwargs,
)
# distributed computations done in cluster context
#with ClusterWrap.cluster(**cluster_kwargs) as cluster:
# if write_path is not None or not lazy:
# cluster.scale_cluster(nblocks + WORKER_BUFFER)
# wrap images as dask arrays
fixed_da = da.from_array(fixed)
moving_da = da.from_array(moving)
# in case xyz convention is flipped for input file
if transpose[0]:
fixed_da = fixed_da.transpose(2, 1, 0)
if transpose[1]:
moving_da = moving_da.transpose(2, 1, 0)
# pad the ends to fill in the last blocks
pads = []
for x, y in zip(original_shape, blocksize):
pads += [(0, y - x % y) if x % y > 0 else (0, 0)]
fixed_da = da.pad(fixed_da, pads)
moving_da = da.pad(moving_da, pads)
# chunk to blocksize
fixed_da = fixed_da.rechunk(tuple(blocksize))
moving_da = moving_da.rechunk(tuple(blocksize))
# wrap deformable function
def wrapped_deformable_align(x, y):
warp = deformable_align(
x,
y,
fixed_spacing,
moving_spacing,
**deform_kwargs,
)
return warp.reshape((1, 1, 1) + warp.shape)
# deform all chunks
overlaps = tuple([int(round(x / 8)) for x in blocksize])
out_blocks = [x + 2 * y for x, y in zip(blocksize, overlaps)]
out_blocks = [1, 1, 1] + out_blocks + [
3,
]
warps = da.map_overlap(
wrapped_deformable_align,
fixed_da,
moving_da,
depth=overlaps,
boundary=0,
trim=False,
align_arrays=False,
dtype=np.float32,
new_axis=[
3,
4,
5,
6,
],
chunks=out_blocks,
)
# stitch neighboring displacement fields
warps = stitch.stitch_fields(warps, blocksize)
# crop any pads
warps = warps[:original_shape[0], :original_shape[1], :original_shape[2]]
# TODO refactor transform.compose_position_fields
# replace this approximation
# compose with affine position field
if affine_pf is not None:
final_field = affine_pf + warps
else:
final_field = warps + transform.position_grid_dask(
original_shape,
blocksize,
)
# if user wants to write to disk
if write_path is not None:
compressor = Blosc(cname='zstd', clevel=9, shuffle=Blosc.BITSHUFFLE)
final_field_disk = zarr.open(
write_path,
'w',
shape=final_field.shape,
chunks=tuple(blocksize + [
3,
]),
dtype=final_field.dtype,
compressor=compressor,
)
da.to_zarr(final_field, final_field_disk)
# if user wants to compute and return full field
if not lazy:
return final_field.compute()
# if user wants to return compute graph w/o executing
if lazy:
return final_field
def prepare_piecewise_deformable_align(
fix,
mov,
fix_spacing,
mov_spacing,
blocksize,
transpose=[False] * 2,
global_affine=None,
local_affines=None,
**kwargs,
):
"""
"""
# get number of blocks required
block_grid = np.ceil(np.array(fix.shape) / blocksize)
# get true field shape
original_shape = fix.shape
if transpose[0]:
original_shape = original_shape[::-1]
# compose global/local affines
total_affines = None
if local_affines is not None and global_affine is not None:
total_affines = transform.compose_affines(
global_affine,
local_affines,
)
elif global_affine is not None:
total_affines = np.empty(tuple(block_grid) + (4, 4))
total_affines[..., :, :] = global_affine
elif local_affines is not None:
total_affines = np.copy(local_affines)
# get affine position field
overlap = tuple([int(round(x / 8)) for x in blocksize])
affine_pf = None
if total_affines is not None:
affine_pf = ds.local_affine.local_affines_to_field(
original_shape,
fix_spacing,
total_affines,
blocksize,
overlap,
displacement=False,
)
# wrap images as dask arrays
fix_da = da.from_array(fix)
mov_da = da.from_array(mov)
# in case xyz convention is flipped for input file
if transpose[0]:
fix_da = fix_da.transpose(2, 1, 0)
if transpose[1]:
mov_da = mov_da.transpose(2, 1, 0)
# pad the ends to fill in the last blocks
pads = []
for x, y in zip(original_shape, blocksize):
pads += [(0, y - x % y) if x % y > 0 else (0, 0)]
fix_da = da.pad(fix_da, pads)
mov_da = da.pad(mov_da, pads)
# chunk to blocksize
fix_da = fix_da.rechunk(tuple(blocksize))
mov_da = mov_da.rechunk(tuple(blocksize))
# wrap deformable function
def wrapped_deformable_align(x, y):
return deformable_align(
x,
y,
fix_spacing,
mov_spacing,
**kwargs,
)
# deform all chunks
out_blocks = [x + 2 * y for x, y in zip(blocksize, overlap)] + [
3,
]
warps = da.map_overlap(
wrapped_deformable_align,
fix_da,
mov_da,
depth=overlap,
boundary=0,
trim=False,
align_arrays=False,
dtype=np.float32,
new_axis=[
3,
],
chunks=out_blocks,
)
# stitch neighboring displacement fields
warps = ds.stitch.stitch_blocks(warps, blocksize, overlap)
# crop any pads
warps = warps[:original_shape[0], :original_shape[1], :original_shape[2]]
# compose with affine position field
# TODO refactor transform.compose_position_fields
# replace this approximation
if affine_pf is not None:
final_field = affine_pf + warps
else:
final_field = warps + ds.local_affine.position_grid(
original_shape,
blocksize,
)
return final_field
def spatial_chunked_regrid(
src_cube,
tgt_cube,
scheme,
min_src_chunk_size=2,
max_src_chunk_size=dask.config.get("array.chunk-size"),
min_tgt_chunk_size=None,
max_tgt_chunk_size=dask.config.get("array.chunk-size"),
tol=1e-16,
):
"""Spatially chunked regridding using dask.
Only the y-coordinate is chunked. This is done because the x-coordinate may be
circular (global), which may require additional logic to be implemented.
Args:
src_cube (iris.cube.Cube): Cube to be regridded onto the coordinate system
defined by the target cube.
tgt_cube (iris.cube.Cube): Target cube. This is solely required to specify the
target coordinate system and may contain dummy data.
scheme: The type of regridding to use to regrid the source cube onto the
target grid, e.g. `iris.analysis.Linear`, `iris.analysis.Nearest`, and
`iris.analysis.AreaWeighted`.
min_src_chunk_size (None, int): Minimum source cube chunk size along the
y-dimension, specified in terms of the number of elements per chunk along
this axis. Note that some regridders, e.g. `iris.analysis.Linear()`
require at least a chunk size of 2 here.
max_src_chunk_size (None, int, str): The maximum size of chunks along the
source cube's y-dimension. Can be given in bytes, e.g. '10MB' or '100KB'.
If None is given, the chunks will be as large as possible.
min_tgt_chunk_size (None, int): Analogous to `min_src_chunk_size` for the
target cube.
max_tgt_chunk_size (None, int, str): Analogous to `max_src_chunk_size` for the
target cube.
Raises:
TypeError: If `src_cube` does not have lazy data.
ValueError: If the source cube is not 2D.
ValueError: If the source or target cube do not define x and y coordinates.
ValueError: If source and target cubes do not define their x and y coordinates
along the same dimensions.
ValueError: If any of the x, y coordinates are not monotonic.
ValueError: If the given maximum chunk sizes are smaller than required for the
regridding of a single data chunk.
"""
if not src_cube.has_lazy_data():
raise TypeError("Source cube needs to have lazy data.")
if src_cube.core_data().ndim != 2:
raise ValueError("Source cube data needs to be 2D.")
coord_err = "{name} cube needs to define x and y coordinates."
try:
src_x_coord, src_y_coord = get_xy_dim_coords(src_cube)
except Exception as exc:
raise ValueError(coord_err.format("Source")) from exc
try:
tgt_x_coord, tgt_y_coord = get_xy_dim_coords(tgt_cube)
except Exception as exc:
raise ValueError(coord_err.format("Target")) from exc
y_dim = src_y_dim = src_cube.coord_dims(src_y_coord)[0]
x_dim = src_x_dim = src_cube.coord_dims(src_x_coord)[0]
tgt_y_dim = tgt_cube.coord_dims(tgt_y_coord)[0]
tgt_x_dim = tgt_cube.coord_dims(tgt_x_coord)[0]
if (src_y_dim, src_x_dim) != (tgt_y_dim, tgt_x_dim):
raise ValueError("Coordinates are not aligned.")
monotonic_err_msg = "{:}-coordinate needs to be monotonic."
src_x_coord_monotonic, src_x_coord_direction = iris.util.monotonic(
src_x_coord.points, return_direction=True)
if not src_x_coord_monotonic:
raise ValueError(monotonic_err_msg.format("Source x"))
if src_x_coord_direction < 0:
# Coordinate is monotonically decreasing, so we need to invert it.
flip_slice = [slice(None)] * src_cube.ndim
flip_slice[src_x_dim] = slice(None, None, -1)
src_cube = src_cube[tuple(flip_slice)]
src_x_coord, src_y_coord = get_xy_dim_coords(src_cube)
src_x_coord.bounds = src_x_coord.bounds[:, ::-1]
src_y_coord_monotonic, src_y_coord_direction = iris.util.monotonic(
src_y_coord.points, return_direction=True)
if not src_y_coord_monotonic:
raise ValueError(monotonic_err_msg.format("Source y"))
if src_y_coord_direction < 0:
# Coordinate is monotonically decreasing, so we need to invert it.
flip_slice = [slice(None)] * src_cube.ndim
flip_slice[src_y_dim] = slice(None, None, -1)
src_cube = src_cube[tuple(flip_slice)]
src_x_coord, src_y_coord = get_xy_dim_coords(src_cube)
src_y_coord.bounds = src_y_coord.bounds[:, ::-1]
tgt_x_coord_monotonic, tgt_x_coord_direction = iris.util.monotonic(
tgt_x_coord.points, return_direction=True)
if not tgt_x_coord_monotonic:
raise ValueError(monotonic_err_msg.format("Target x"))
if tgt_x_coord_direction < 0:
# Coordinate is monotonically decreasing, so we need to invert it.
flip_slice = [slice(None)] * tgt_cube.ndim
flip_slice[tgt_x_dim] = slice(None, None, -1)
tgt_cube = tgt_cube[tuple(flip_slice)]
tgt_x_coord, tgt_y_coord = get_xy_dim_coords(tgt_cube)
tgt_x_coord.bounds = tgt_x_coord.bounds[:, ::-1]
tgt_y_coord_monotonic, tgt_y_coord_direction = iris.util.monotonic(
tgt_y_coord.points, return_direction=True)
if not tgt_y_coord_monotonic:
raise ValueError(monotonic_err_msg.format("Target y"))
if tgt_y_coord_direction < 0:
# Coordinate is monotonically decreasing, so we need to invert it.
flip_slice = [slice(None)] * tgt_cube.ndim
flip_slice[tgt_y_dim] = slice(None, None, -1)
tgt_cube = tgt_cube[tuple(flip_slice)]
tgt_x_coord, tgt_y_coord = get_xy_dim_coords(tgt_cube)
tgt_y_coord.bounds = tgt_y_coord.bounds[:, ::-1]
max_src_chunk_size = convert_chunk_size(
max_src_chunk_size,
# The number of elements along the non-chunked dimension.
factor=src_cube.shape[x_dim],
dtype=src_cube.dtype,
masked=isinstance(src_cube.core_data()._meta, np.ma.MaskedArray),
)
max_tgt_chunk_size = convert_chunk_size(
max_tgt_chunk_size,
# The number of elements along the non-chunked dimension.
factor=tgt_cube.shape[x_dim],
# NOTE: Is this true?
dtype=src_cube.dtype,
# Set masked to True here since we will add a mask later in all cases.
masked=True,
)
max_chunk_msg = (
"Maximum {:} chunk size was smaller than the minimum required for a single "
"chunk.")
if max_src_chunk_size == 0:
raise ValueError(max_chunk_msg.format("source"))
if max_tgt_chunk_size == 0:
raise ValueError(max_chunk_msg.format("target"))
# Calculate all possible chunks along the y dimension.
overlap_indices, valid_src_y_slice, valid_tgt_y_slice = get_overlapping(
src_y_coord.contiguous_bounds(),
tgt_y_coord.contiguous_bounds(),
tol=tol,
)
# Some regridding methods care about cells with points outside of overlapping
# bounds, like `iris.analysis.Linear()`. Include an additional source cell on
# either end if possible to account for this.
# NOTE: These additions may be superfluous, but would require re-writing
# `get_overlapping()` to determine.
if valid_src_y_slice.start > 0:
overlap_indices[valid_tgt_y_slice][0].insert(
0, overlap_indices[valid_tgt_y_slice][0][0] - 1)
if valid_src_y_slice.stop < src_y_coord.shape[0]:
overlap_indices[valid_tgt_y_slice][-1].append(
overlap_indices[valid_tgt_y_slice][-1][-1] + 1)
valid_src_y_slice = slice(
max(0, valid_src_y_slice.start - 1),
min(src_y_coord.shape[0], valid_src_y_slice.stop + 1),
)
cell_numbers, overlap_y = get_cell_numbers(overlap_indices)
cell_mapping = get_valid_cell_mapping(cell_numbers[valid_tgt_y_slice])
tgt_y_slices = []
src_y_chunks = []
tgt_y_chunks = []
for tgt_cells, src_cells in cell_mapping.items():
tgt_y_slices.append(slice(tgt_cells[0], tgt_cells[-1] + 1))
src_y_chunks.append(len(src_cells))
tgt_y_chunks.append(len(tgt_cells))
# XXX: This override is sometimes needed due to floating point errors, e.g. for
# test_regrid case
# 100_200-50_120--90_90--90_90--180_180--180_180-AreaWeighted_1-Block-20KB
# where tgt_cube.coord('latitude')[24].bounds[0][1] is 3.55e-15 instead of 0,
# causing src_cube.coord('latitude')[50] (with bounds [0, 1.8] to be required to
# match the masking behaviour in Iris regrid even though this is not expected to
# influence the final result to the small overlap. Another solution is to decrease
# the `tol` parameter of `get_overlapping()`, in this case below 3.55e-16
# (e.g. 1e-16).
# overlap_y = True
src_y_chunks, tgt_y_slices = calculate_blocks(
src_y_chunks,
tgt_y_chunks,
tgt_y_slices,
min_src_chunk_size,
max_src_chunk_size,
min_tgt_chunk_size,
max_tgt_chunk_size,
)
valid_src_slice = [slice(None)] * src_cube.ndim
valid_src_slice[src_y_dim] = valid_src_y_slice
valid_src_slice = tuple(valid_src_slice)
# Re-chunk the data and coordinate along the y-dimension.
block_src_data = src_cube.core_data()[valid_src_slice].rechunk((
src_y_chunks,
-1,
))
# 2D arrays are created here to enable consistent map_overlap behaviour.
block_src_y_pnts = (da.from_array(
src_y_coord.points[valid_src_y_slice]).rechunk(
(src_y_chunks, )).reshape(-1, 1))
block_src_low_y_bnds = (da.from_array(
src_y_coord.bounds[valid_src_y_slice, 0]).rechunk(
(src_y_chunks, )).reshape(-1, 1))
block_src_upp_y_bnds = (da.from_array(
src_y_coord.bounds[valid_src_y_slice, 1]).rechunk(
(src_y_chunks, )).reshape(-1, 1))
chunks_spec = [None] * 2
chunks_spec[x_dim] = tgt_x_coord.shape[0]
chunks_spec[y_dim] = tuple(slice_len(s) for s in tgt_y_slices)
chunks_spec = tuple(chunks_spec)
output = da.map_overlap(
regrid_chunk,
block_src_data,
block_src_y_pnts,
block_src_low_y_bnds,
block_src_upp_y_bnds,
# Store metadata for the y-coordinate for which points and bounds will be
# filled in during the course of blocked regridding.
src_y_coord_metadata=src_y_coord.metadata,
src_x_coord=src_x_coord,
src_cube_metadata=src_cube.metadata,
tgt_y_coord=tgt_y_coord[valid_tgt_y_slice],
tgt_x_coord=tgt_x_coord,
tgt_y_slices=tgt_y_slices,
tgt_cube_metadata=tgt_cube.metadata,
y_dim=y_dim,
x_dim=x_dim,
scheme=scheme,
depth={
# The y-coordinate may need to be overlapped.
y_dim: 1 if overlap_y else 0,
# The x-coordinate will not be overlapped as it is never chunked.
x_dim: 0,
},
boundary="none",
trim=False,
dtype=np.float64,
chunks=chunks_spec,
meta=np.array([], dtype=np.float64),
)
if not isinstance(output._meta, np.ma.MaskedArray):
# XXX: Ideally this should not be needed, but the mask appears to vanish in
# some cases.
output = da.ma.masked_array(output, mask=False)
if slice_len(valid_tgt_y_slice) < tgt_y_coord.shape[0]:
# Embed the output data calculated above into the final target cube by padding
# with masked data as necessary.
seq = []
if valid_tgt_y_slice.start > 0:
# Pad at the start.
start_pad_shape = [tgt_x_coord.shape[0]] * 2
start_pad_shape[y_dim] = valid_tgt_y_slice.start
seq.append(
da.ma.masked_array(
da.zeros(start_pad_shape),
mask=True,
))
seq.append(output)
if valid_tgt_y_slice.stop < tgt_y_coord.shape[0]:
# Pad at the end.
end_pad_shape = [tgt_x_coord.shape[0]] * 2
end_pad_shape[
y_dim] = tgt_y_coord.shape[0] - valid_tgt_y_slice.stop
seq.append(da.ma.masked_array(
da.zeros(end_pad_shape),
mask=True,
))
output = da.concatenate(seq, axis=y_dim)
return tgt_cube.copy(data=output)
def deformable_align_distributed(
fixed, moving,
fixed_vox, moving_vox,
write_path,
cc_radius,
gradient_smoothing,
field_smoothing,
iterations,
shrink_factors,
smooth_sigmas,
step,
blocksize=[256,]*3,
cluster_extra=["-P multifish"],
transpose=False,
):
"""
"""
# distributed computations done in cluster context
with distributed.distributedState() as ds:
# get number of blocks required
block_grid = np.ceil(np.array(fixed.shape) / blocksize)
nblocks = np.prod(block_grid)
# set up the cluster
ds.initializeLSFCluster(
job_extra=cluster_extra,
cores=4, memory="64GB", ncpus=4, threads_per_worker=8, mem=64000,
)
ds.initializeClient()
ds.scaleCluster(njobs=nblocks)
# wrap images as dask arrays
fixed_da = da.from_array(fixed)
moving_da = da.from_array(moving)
# in case xyz convention is flipped for input file
if transpose:
fixed_da = fixed_da.transpose(2,1,0)
# pad the ends to fill in the last blocks
orig_sh = fixed_da.shape
pads = [(0, y - x % y) if x % y != 0 else (0, 0) for x, y in zip(orig_sh, blocksize)]
fixed_da = da.pad(fixed_da, pads)
moving_da = da.pad(moving_da, pads)
fixed_da = fixed_da.rechunk(tuple(blocksize))
moving_da = moving_da.rechunk(tuple(blocksize))
# wrap deformable function to simplify passing parameters
def my_deformable_align(x, y):
return deformable_align(
x, y, fixed_vox, moving_vox,
cc_radius, gradient_smoothing, field_smoothing,
iterations, shrink_factors, smooth_sigmas, step,
)
# deform all chunks
overlaps = tuple([int(round(x/8)) for x in blocksize])
out_blocks = [1,1,1] + [x + 2*y for x, y in zip(blocksize, overlaps)] + [3,]
warps = da.map_overlap(
my_deformable_align, fixed_da, moving_da,
depth=overlaps,
boundary='reflect',
trim=False,
align_arrays=False,
dtype=np.float32,
new_axis=[3,4,5,6,],
chunks=out_blocks,
)
# stitch neighboring displacement fields
warps = stitch.stitch_fields(warps, blocksize)
# crop any pads
warps = warps[:orig_sh[0], :orig_sh[1], :orig_sh[2]]
# convert to position field
warps = warps + transform.position_grid_dask(orig_sh, blocksize)
# write result to zarr file
compressor = Blosc(cname='zstd', clevel=9, shuffle=Blosc.BITSHUFFLE)
warps_disk = zarr.open(write_path, 'w',
shape=warps.shape, chunks=tuple(blocksize + [3,]),
dtype=warps.dtype, compressor=compressor,
)
da.to_zarr(warps, warps_disk)
# return reference to zarr data store
return warps_disk
def apply_position_field(
mov,
mov_spacing,
fix,
fix_spacing,
transform,
output,
blocksize=[
256,
] * 3,
order=1,
transform_spacing=None,
transpose=[
False,
] * 3,
depth=(32, 32, 32),
):
"""
"""
with distributed.distributedState() as ds:
# get number of jobs needed
block_grid = np.ceil(np.array(mov.shape) / blocksize).astype(int)
nblocks = np.prod(block_grid)
# set up the cluster
ds.initializeLSFCluster(job_extra=["-P multifish"])
ds.initializeClient()
ds.scaleCluster(njobs=nblocks)
# determine mov/fix relative chunking
m_blocksize = blocksize * fix_spacing / mov_spacing
m_blocksize = list(np.round(m_blocksize).astype(np.int16))
m_depth = depth * fix_spacing / mov_spacing
m_depth = tuple(np.round(m_depth).astype(np.int16))
# determine trans/fix relative chunking
if transform_spacing is not None:
t_blocksize = blocksize * fix_spacing / transform_spacing
t_blocksize = list(np.round(t_blocksize).astype(np.int16))
t_depth = depth * fix_spacing / transform_spacing
t_depth = tuple(np.round(t_depth).astype(np.int16))
else:
t_blocksize = blocksize
t_depth = depth
# wrap objects as dask arrays
fix_da = da.from_array(fix)
if transpose[0]:
fix_da = fix_da.transpose(2, 1, 0)
mov_da = da.from_array(mov)
if transpose[1]:
mov_da = mov_da.transpose(2, 1, 0)
block_grid = block_grid[::-1]
transform_da = da.from_array(transform)
if transpose[2]:
transform_da = transform_da.transpose(2, 1, 0, 3)
transform_da = transform_da[..., ::-1]
# chunk dask arrays
fix_da = da.reshape(fix_da, fix_da.shape + (1, )).rechunk(
tuple(blocksize + [
1,
]))
mov_da = da.reshape(mov_da, mov_da.shape + (1, )).rechunk(
tuple(m_blocksize + [
1,
]))
transform_da = transform_da.rechunk(tuple(t_blocksize + [
3,
]))
# put transform in voxel units
transform_da = transform_da / mov_spacing
# map the interpolate function with overlaps
# TODO: depth should be computed automatically from transform maximum?
d = [depth + (0, ), m_depth + (0, ), t_depth + (0, )]
aligned = da.map_overlap(
interpolate_image_dask,
fix_da,
mov_da,
transform_da,
blocksize=m_blocksize,
margin=m_depth,
depth=d,
boundary=0,
dtype=np.uint16,
align_arrays=False,
)
# remove degenerate dimension
aligned = da.reshape(aligned, aligned.shape[:-1])
# write in parallel as 3D array to zarr file
compressor = Blosc(cname='zstd', clevel=9, shuffle=Blosc.BITSHUFFLE)
aligned_disk = zarr.open(
output,
'w',
shape=aligned.shape,
chunks=aligned.chunksize,
dtype=aligned.dtype,
compressor=compressor,
)
da.to_zarr(aligned, aligned_disk)
# return pointer to zarr file
return aligned_disk
def tiled_ransac_affine(
fixed,
moving,
fixed_vox,
moving_vox,
min_radius,
max_radius,
match_threshold,
blocksize,
# cluster_kwargs={},
**kwargs,
):
"""
"""
# get number of blocks required
block_grid = np.ceil(np.array(fixed.shape) / blocksize).astype(int)
nblocks = np.prod(block_grid)
overlap = [int(round(x / 8)) for x in blocksize]
# distributed computations done in cluster context
#with ClusterWrap.cluster(**cluster_kwargs) as cluster:
# cluster.scale_cluster(nblocks + WORKER_BUFFER)
# wrap images as dask arrays
fixed_da = da.from_array(fixed, chunks=blocksize)
moving_da = da.from_array(moving, chunks=blocksize)
# wrap affine function
def wrapped_ransac_affine(x, y):
affine = ransac_affine(
x,
y,
fixed_vox,
moving_vox,
min_radius,
max_radius,
match_threshold,
**kwargs,
)
return affine.reshape((1, 1, 1, 3, 4))
# affine align all chunks
affines = da.map_overlap(
wrapped_ransac_affine,
fixed_da,
moving_da,
depth=tuple(overlap),
boundary='reflect',
trim=False,
align_arrays=False,
dtype=np.float64,
new_axis=[3, 4],
chunks=[1, 1, 1, 3, 4],
).compute()
# improve on identity affines
affines = interpolate_affines(affines)
# adjust affines for block origins
for i in range(nblocks):
x, y, z = np.unravel_index(i, block_grid)
origin = (np.array(blocksize) * [x, y, z] - overlap) * fixed_vox
tl, tr, a = np.eye(4), np.eye(4), np.eye(4)
tl[:3, -1], tr[:3, -1] = origin, -origin
a[:3, :] = affines[x, y, z]
affines[x, y, z] = np.matmul(tl, np.matmul(a, tr))[:3, :]
return affines
