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