def process_volumes(parameters): '''process original volumes and save into nifti format''' import os import copy import numpy as np from scipy import interpolate from types import SimpleNamespace from voluseg._tools.load_volume import load_volume from voluseg._tools.save_volume import save_volume from voluseg._tools.plane_name import plane_name from voluseg._tools.constants import ori, ali, nii, hdf from voluseg._tools.evenly_parallelize import evenly_parallelize p = SimpleNamespace(**parameters) volume_nameRDD = evenly_parallelize( p.volume_names0 if p.planes_packed else p.volume_names) for color_i in range(p.n_colors): fullname_volmean = os.path.join(p.dir_output, 'volume%d' % (color_i)) if os.path.isfile(fullname_volmean + hdf): continue dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) os.makedirs(dir_volume, exist_ok=True) def initial_processing(tuple_name_volume): name_volume = tuple_name_volume[1] # try: # load input volumes fullname_input = os.path.join(p.dir_input, name_volume) volume = load_volume(fullname_input + p.ext) # get number of planes if p.planes_packed: name_volume0 = copy.deepcopy(name_volume) volume0 = copy.deepcopy(volume) lp = len(volume0) else: lp = 1 for pi in range(lp): if p.planes_packed: name_volume = plane_name(name_volume0, pi) volume = volume0[pi] fullname_volume = os.path.join(dir_volume, name_volume) # skip processing if volume exists if load_volume(fullname_volume+ori+nii) is not None \ or load_volume(fullname_volume+ali+hdf): continue if volume.ndim == 2: volume = volume[None, :, :] volume = volume.transpose(2, 1, 0) # get dimensions lx, ly, lz = volume.shape # split two-color volumes into two halves if p.n_colors == 2: # ensure that two-frames have even number of y-dim voxels assert (ly % 2 == 0) ly //= 2 if color_i == 0: volume = volume[:, :ly, :] elif color_i == 1: volume = volume[:, ly:, :] # downsample in the x-y if specified if p.ds > 1: if (lx % p.ds) or (ly % p.ds): lx -= (lx % p.ds) ly -= (ly % p.ds) volume = volume[:lx, :ly, :] # make grid for computing downsampled values sx_ds = np.arange(0.5, lx, p.ds) sy_ds = np.arange(0.5, ly, p.ds) xy_grid_ds = np.dstack( np.meshgrid(sx_ds, sy_ds, indexing='ij')) # get downsampled volume volume_ds = np.zeros((len(sx_ds), len(sy_ds), lz)) for zi in np.arange(lz): interpolation_fx = interpolate.RegularGridInterpolator( (np.arange(lx), np.arange(ly)), volume[:, :, zi], method='linear') volume_ds[:, :, zi] = interpolation_fx(xy_grid_ds) volume = volume_ds # pad planes as necessary if p.registration and p.planes_pad: volume = np.lib.pad( volume, ((0, 0), (0, 0), (p.planes_pad, p.planes_pad)), 'constant', constant_values=(np.percentile(volume, 1), )) # save volume in output directory if p.registration: save_volume(fullname_volume + ori + nii, volume, p.affine_mat) else: volume = volume.T save_volume(fullname_volume + ali + hdf, volume) volume_nameRDD.foreach(initial_processing)
def mask_volumes(parameters): '''create intensity mask from the average registered volume''' import os import h5py import pyspark import numpy as np from scipy import stats from sklearn import mixture from skimage import morphology from types import SimpleNamespace from scipy.ndimage.filters import median_filter from voluseg._tools.ball import ball from voluseg._tools.constants import ali, hdf from voluseg._tools.load_volume import load_volume from voluseg._tools.clean_signal import clean_signal from voluseg._tools.evenly_parallelize import evenly_parallelize # set up spark from pyspark.sql.session import SparkSession spark = SparkSession.builder.getOrCreate() sc = spark.sparkContext # set up matplotlib import warnings import matplotlib with warnings.catch_warnings(): warnings.simplefilter('ignore') matplotlib.use('Agg') import matplotlib.pyplot as plt p = SimpleNamespace(**parameters) # compute mean timeseries and ranked dff fullname_timemean = os.path.join(p.dir_output, 'mean_timeseries') volume_nameRDD = evenly_parallelize(p.volume_names) if not os.path.isfile(fullname_timemean + hdf): dff_rank = np.zeros(p.lt) mean_timeseries_raw = np.zeros((p.n_colors, p.lt)) mean_timeseries = np.zeros((p.n_colors, p.lt)) mean_baseline = np.zeros((p.n_colors, p.lt)) for color_i in range(p.n_colors): dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) def mean_volume(tuple_name_volume): name_volume = tuple_name_volume[1] fullname_volume = os.path.join(dir_volume, name_volume) return np.mean(load_volume(fullname_volume + ali + hdf), dtype='float64') mean_timeseries_raw[color_i] = volume_nameRDD.map( mean_volume).collect() time, base = clean_signal(parameters, mean_timeseries_raw[color_i]) mean_timeseries[color_i], mean_baseline[color_i] = time, base dff_rank += stats.rankdata((time - base) / time) # get high delta-f/f timepoints if not p.nt: timepoints = np.range(p.lt) else: timepoints = np.sort(np.argsort(dff_rank)[::-1][:p.nt]) with h5py.File(fullname_timemean + hdf, 'w') as file_handle: file_handle['mean_timeseries_raw'] = mean_timeseries_raw file_handle['mean_timeseries'] = mean_timeseries file_handle['mean_baseline'] = mean_baseline file_handle['timepoints'] = timepoints # load timepoints with h5py.File(fullname_timemean + hdf, 'r') as file_handle: timepoints = file_handle['timepoints'][()] for color_i in range(p.n_colors): fullname_volmean = os.path.join(p.dir_output, 'volume%d' % (color_i)) if os.path.isfile(fullname_volmean + hdf): continue dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) dir_plot = os.path.join(p.dir_output, 'mask_plots', str(color_i)) os.makedirs(dir_plot, exist_ok=True) fullname_volume = os.path.join(dir_volume, p.volume_names[0]) lx, ly, lz = load_volume(fullname_volume + ali + hdf).T.shape class accum_param(pyspark.accumulators.AccumulatorParam): '''define accumulator class''' def zero(self, val0): return np.zeros(val0.shape, dtype='float64') def addInPlace(self, val1, val2): return np.add(val1, val2, dtype='float64') # geometric mean volume_accum = sc.accumulator(np.zeros((lx, ly, lz), dtype='float64'), accum_param()) def add_volume(tuple_name_volume): name_volume = tuple_name_volume[1] fullname_volume = os.path.join(dir_volume, name_volume) volume_accum.add( np.log10(load_volume(fullname_volume + ali + hdf).T)) evenly_parallelize(p.volume_names[timepoints]).foreach(add_volume) volume_mean = 10**(volume_accum.value / p.lt) # get peaks by comparing to a median-smoothed volume ball_radi = ball(0.5 * p.diam_cell, p.affine_mat)[0] volume_peak = volume_mean > median_filter(volume_mean, footprint=ball_radi) # compute power and probability voxel_intensity = np.percentile(volume_mean[volume_mean > 0], np.r_[5:95:0.001])[:, None] gmm = mixture.GaussianMixture(n_components=2, max_iter=100, n_init=100).fit(voxel_intensity) voxel_probability = gmm.predict_proba(voxel_intensity) voxel_probability = voxel_probability[:, np. argmax(voxel_intensity[np.argmax( voxel_probability, 0)])] # compute intensity threshold if (p.thr_mask > 0) and (p.thr_mask <= 1): thr_probability = p.thr_mask ix = np.argmin(np.abs(voxel_probability - thr_probability)) thr_intensity = voxel_intensity[ix][0] elif p.thr_mask > 1: thr_intensity = p.thr_mask ix = np.argmin(np.abs(voxel_intensity - thr_intensity)) thr_probability = voxel_probability[ix] else: thr_intensity = -np.inf thr_probability = 0 print('using probability threshold of %f.' % (thr_probability)) print('using intensity threshold of %f.' % (thr_intensity)) # get and save brain mask fig = plt.figure(1, (18, 6)) plt.subplot(131), _ = plt.hist(voxel_intensity, 100) plt.plot(thr_intensity, 0, '|', color='r', markersize=200) plt.xlabel('voxel intensity') plt.title('intensity histogram with threshold (red)') plt.subplot(132), _ = plt.hist(voxel_probability, 100) plt.plot(thr_probability, 0, '|', color='r', markersize=200) plt.xlabel('voxel probability') plt.title('probability histogram with threshold (red)') plt.subplot(133), plt.plot(voxel_intensity, voxel_probability, linewidth=3) plt.plot(thr_intensity, thr_probability, 'x', color='r', markersize=10) plt.xlabel('voxel intensity') plt.ylabel('voxel probability') plt.title('intensity-probability plot with threshold (red)') plt.savefig(os.path.join(dir_plot, 'histogram.png')) plt.close(fig) # remove all disconnected components less than 5000 cubic microliters in size rx, ry, rz, _ = np.diag(p.affine_mat) volume_mask = (volume_mean > thr_intensity).astype('bool') thr_size = np.round(5000 * rx * ry * rz).astype(int) volume_mask = morphology.remove_small_objects(volume_mask, thr_size) # compute background fluorescence background = np.median(volume_mean[volume_mask == 0]) # save brain mask figures for i in range(lz): fig = plt.figure(1, (18, 6)) plt.subplot(131) plt.imshow(volume_mean[:, :, i].T, vmin=voxel_intensity[0], vmax=voxel_intensity[-1]) plt.title('volume intensity (plane %d)' % (i)) plt.subplot(132) plt.imshow(volume_mask[:, :, i].T) plt.title('volume mask (plane %d)' % (i)) plt.subplot(133) img = np.stack( (volume_mean[:, :, i], volume_mask[:, :, i], volume_mask[:, :, i]), axis=2) img[:, :, 0] = (img[:, :, 0] - voxel_intensity[0]) / ( voxel_intensity[-1] - voxel_intensity[0]) img[:, :, 0] = np.minimum(np.maximum(img[:, :, 0], 0), 1) plt.imshow(np.transpose(img, [1, 0, 2])) plt.title('volume mask/intensity overlay (plane %d)' % (i)) plt.savefig(os.path.join(dir_plot, 'mask_z%03d.png' % (i))) plt.close(fig) with h5py.File(fullname_volmean + hdf, 'w') as file_handle: file_handle['volume_mask'] = volume_mask.T file_handle['volume_mean'] = volume_mean.T file_handle['volume_peak'] = volume_peak.T file_handle['thr_intensity'] = thr_intensity file_handle['thr_probability'] = thr_probability file_handle['background'] = background
def process_images(parameters): '''process original images and save into nifti format''' import os import h5py import nibabel import numpy as np from scipy import interpolate from types import SimpleNamespace from skimage.external import tifffile from voluseg._tools.nii_image import nii_image from voluseg._tools.evenly_parallelize import evenly_parallelize try: import PIL import pyklb except: pass p = SimpleNamespace(**parameters) volume_nameRDD = evenly_parallelize(p.volume_names) for color_i in range(p.n_colors): if os.path.isfile( os.path.join(p.dir_output, 'volume%d.hdf5' % (color_i))): continue dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) os.makedirs(dir_volume, exist_ok=True) def initial_processing(tuple_name_volume): name_volume = tuple_name_volume[1] fullname_original = os.path.join(dir_volume, name_volume + '_original.nii.gz') fullname_aligned = os.path.join(dir_volume, name_volume + '_aligned.nii.gz') fullname_aligned_hdf = fullname_aligned.replace('.nii.gz', '.hdf5') if os.path.isfile(fullname_original): try: volume_original = nibabel.load( fullname_original).get_data() return except: pass if os.path.isfile(fullname_aligned): try: volume_aligned = nibabel.load(fullname_aligned).get_data() return except: pass if os.path.isfile(fullname_aligned_hdf): try: with h5py.File(fullname_aligned_hdf) as file_handle: volume_aligned = file_handle['V3D'][()].T return except: pass try: # load input images fullname_input = os.path.join(p.dir_input, name_volume + p.ext) if ('.tif' in p.ext) or ('.tiff' in p.ext): try: volume_input = tifffile.imread(fullname_input) except: img = PIL.Image.open(fullname_input) volume_input = [] for i in range(img.n_frames): img.seek(i) volume_input.append(np.array(img).T) volume_input = np.array(volume_input) elif ('.h5' in p.ext) or ('.hdf5' in p.ext): with h5py.File(fullname_input, 'r') as file_handle: volume_input = file_handle[list( file_handle.keys())[0]][()] elif ('.klb' in p.ext): volume_input = pyklb.readfull(fullname_input) volume_input = volume_input.transpose(0, 2, 1) if volume_input.ndim == 2: volume_input = volume_input[None, :, :] volume_input = volume_input.transpose(2, 1, 0) # get dimensions lx, ly, lz = volume_input.shape # split two-color images into two halves if p.n_colors == 2: # ensure that two-frames have even number of y-dim voxels assert (ly % 2 == 0) ly /= 2 if color_i == 0: volume_input = volume_input[:, :ly, :] elif color_i == 1: volume_input = volume_input[:, ly:, :] # downsample in the x-y if specified if p.ds > 1: if (lx % p.ds) or (ly % p.ds): lx -= (lx % p.ds) ly -= (ly % p.ds) volume_input = volume_input[:lx, :ly, :] # make grid for computing downsampled values sx_ds = np.arange(0.5, lx, p.ds) sy_ds = np.arange(0.5, ly, p.ds) xy_grid_ds = np.dstack( np.meshgrid(sx_ds, sy_ds, indexing='ij')) # get downsampled image volume_input_ds = np.zeros((len(sx_ds), len(sy_ds), lz)) for zi in np.arange(lz): interpolation_fx = interpolate.RegularGridInterpolator( (np.arange(lx), np.arange(ly)), volume_input[:, :, zi], method='linear') volume_input_ds[:, :, zi] = interpolation_fx(xy_grid_ds) volume_input = volume_input_ds # pad planes as necessary if p.registration and p.planes_pad: volume_input = np.lib.pad( volume_input, ((0, 0), (0, 0), (p.planes_pad, p.planes_pad)), 'constant', constant_values=(np.percentile(volume_input, 1), )) # save image as a nifti file nibabel.save( nii_image(volume_input.astype('float32'), p.affine_mat), (fullname_original if p.registration else fullname_aligned)) except Exception as msg: raise Exception('image %s not processed: %s.' % (name_volume, msg)) volume_nameRDD.foreach(initial_processing)
def mask_images(parameters): '''create intensity mask from the average registered image''' import os import h5py import nibabel import pyspark import numpy as np from sklearn import mixture from skimage import morphology from types import SimpleNamespace from scipy.ndimage.filters import median_filter from voluseg._tools.evenly_parallelize import evenly_parallelize from voluseg._tools.ball import ball # set up spark from pyspark.sql.session import SparkSession spark = SparkSession.builder.getOrCreate() sc = spark.sparkContext # set up matplotlib import warnings import matplotlib with warnings.catch_warnings(): warnings.simplefilter('ignore') matplotlib.use('Agg') import matplotlib.pyplot as plt p = SimpleNamespace(**parameters) volume_nameRDD = evenly_parallelize(p.volume_names) for color_i in range(p.n_colors): if os.path.isfile( os.path.join(p.dir_output, 'volume%d.hdf5' % (color_i))): continue dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) dir_plot = os.path.join(p.dir_output, 'mask_plots', str(color_i)) os.makedirs(dir_plot, exist_ok=True) def load_volume(name_volume): fullname_aligned = os.path.join(dir_volume, name_volume + '_aligned.nii.gz') fullname_aligned_hdf = fullname_aligned.replace('.nii.gz', '.hdf5') if os.path.isfile(fullname_aligned): return nibabel.load(fullname_aligned).get_data() elif os.path.isfile(fullname_aligned_hdf): with h5py.File(fullname_aligned_hdf, 'r') as file_handle: return (file_handle['V3D'][()].T) else: raise Exception('%s or %s do not exist.' % (fullname_aligned, fullname_aligned_hdf)) lx, ly, lz = load_volume(p.volume_names[0]).shape class accum_param(pyspark.accumulators.AccumulatorParam): '''define accumulator class''' def zero(self, val0): return np.zeros(val0.shape, dtype='float64') def addInPlace(self, val1, val2): return np.add(val1, val2, dtype='float64') # geometric mean volume_accum = sc.accumulator(np.zeros((lx, ly, lz), dtype='float64'), accum_param()) def add_volume(tuple_name_volume): name_volume = tuple_name_volume[1] volume_accum.add(np.log10(load_volume(name_volume))) volume_nameRDD.foreach(add_volume) volume_mean = 10**(volume_accum.value / p.lt) # get peaks by comparing to a median-smoothed volume ball_radi = ball(0.5 * p.diam_cell, p.affine_mat)[0] volume_peak = volume_mean > median_filter(volume_mean, footprint=ball_radi) # compute power and probability voxel_intensity = np.percentile(volume_mean[volume_mean > 0], np.r_[5:95:0.001])[:, None] gmm = mixture.GaussianMixture(n_components=2, max_iter=100, n_init=100).fit(voxel_intensity) voxel_probability = gmm.predict_proba(voxel_intensity) voxel_probability = voxel_probability[:, np. argmax(voxel_intensity[np.argmax( voxel_probability, 0)])] # compute intensity threshold if (p.thr_mask > 0) and (p.thr_mask <= 1): thr_probability = p.thr_mask ix = np.argmin(np.abs(voxel_probability - thr_probability)) thr_intensity = voxel_intensity[ix][0] elif p.thr_mask > 1: thr_intensity = p.thr_mask ix = np.argmin(np.abs(voxel_intensity - thr_intensity)) thr_probability = voxel_probability[ix] else: thr_intensity = -np.inf thr_probability = 0 print('using probability threshold of %f.' % (thr_probability)) print('using intensity threshold of %f.' % (thr_intensity)) # get and save brain mask fig = plt.figure(1, (18, 6)) plt.subplot(131), _ = plt.hist(voxel_intensity, 100) plt.plot(thr_intensity, 0, '|', color='r', markersize=200) plt.xlabel('voxel intensity') plt.title('intensity histogram with threshold (red)') plt.subplot(132), _ = plt.hist(voxel_probability, 100) plt.plot(thr_probability, 0, '|', color='r', markersize=200) plt.xlabel('voxel probability') plt.title('probability histogram with threshold (red)') plt.subplot(133), plt.plot(voxel_intensity, voxel_probability, linewidth=3) plt.plot(thr_intensity, thr_probability, 'x', color='r', markersize=10) plt.xlabel('voxel intensity') plt.ylabel('voxel probability') plt.title('intensity-probability plot with threshold (red)') plt.savefig(os.path.join(dir_plot, 'histogram.png')) plt.close(fig) # remove all disconnected components less than 5000 cubic microliters in size rx, ry, rz, _ = np.diag(p.affine_mat) volume_mask = (volume_mean > thr_intensity).astype('bool') thr_size = np.round(5000 * rx * ry * rz).astype(int) volume_mask = morphology.remove_small_objects(volume_mask, thr_size) # compute background fluorescence background = np.median(volume_mean[volume_mask == 0]) # compute mean timeseries bvolume_mask = sc.broadcast(volume_mask) def masked_mean(tuple_name_volume): name_volume = tuple_name_volume[1] return np.mean(load_volume(name_volume)[bvolume_mask.value], dtype='float64') timeseries_mean = np.array(volume_nameRDD.map(masked_mean).collect()) # save brain mask figures for i in range(lz): fig = plt.figure(1, (18, 6)) plt.subplot(131) plt.imshow(volume_mean[:, :, i].T, vmin=voxel_intensity[0], vmax=voxel_intensity[-1]) plt.title('volume intensity (plane %d)' % (i)) plt.subplot(132) plt.imshow(volume_mask[:, :, i].T) plt.title('volume mask (plane %d)' % (i)) plt.subplot(133) img = np.stack( (volume_mean[:, :, i], volume_mask[:, :, i], volume_mask[:, :, i]), axis=2) img[:, :, 0] = (img[:, :, 0] - voxel_intensity[0]) / ( voxel_intensity[-1] - voxel_intensity[0]) img[:, :, 0] = np.minimum(np.maximum(img[:, :, 0], 0), 1) plt.imshow(np.transpose(img, [1, 0, 2])) plt.title('volume mask/intensity overlay (plane %d)' % (i)) plt.savefig(os.path.join(dir_plot, 'mask_z%03d.png' % (i))) plt.close(fig) with h5py.File(os.path.join(p.dir_output, 'volume%s.hdf5' % (color_i)), 'w') as file_handle: file_handle['volume_mask'] = volume_mask.T file_handle['volume_mean'] = volume_mean.T file_handle['volume_peak'] = volume_peak.T file_handle['timeseries_mean'] = timeseries_mean file_handle['thr_intensity'] = thr_intensity file_handle['thr_probability'] = thr_probability file_handle['background'] = background # convert nifti images to hdf5 files def nii2hdf(tuple_name_volume): name_volume = tuple_name_volume[1] fullname_aligned = os.path.join(dir_volume, name_volume + '_aligned.nii.gz') fullname_aligned_hdf = fullname_aligned.replace('.nii.gz', '.hdf5') if not os.path.isfile(fullname_aligned_hdf): with h5py.File(fullname_aligned_hdf, 'w') as file_handle: volume_aligned = nibabel.load( fullname_aligned).get_data().T.astype('float32') file_handle.create_dataset('V3D', data=volume_aligned, compression='gzip') try: os.remove(fullname_aligned) except: pass volume_nameRDD.foreach(nii2hdf)
def clean_cells(parameters): '''remove noise cells, detrend and detect baseline''' import os import h5py import shutil import numpy as np from types import SimpleNamespace from itertools import combinations from voluseg._steps.step4e import collect_blocks from voluseg._tools.constants import hdf, dtype from voluseg._tools.clean_signal import clean_signal from voluseg._tools.evenly_parallelize import evenly_parallelize # set up spark from pyspark.sql.session import SparkSession spark = SparkSession.builder.getOrCreate() sc = spark.sparkContext p = SimpleNamespace(**parameters) thr_similarity = 0.5 for color_i in range(p.n_colors): fullname_cells = os.path.join(p.dir_output, 'cells%s_clean' % (color_i)) if os.path.isfile(fullname_cells + hdf): continue block_id, cell_xyz, cell_weights, cell_timeseries, cell_lengths = collect_blocks( color_i, parameters) fullname_volmean = os.path.join(p.dir_output, 'volume%d' % (color_i)) with h5py.File(fullname_volmean + hdf, 'r') as file_handle: volume_mask = file_handle['volume_mask'][()].T x, y, z = volume_mask.shape cell_x = cell_xyz[:, :, 0] cell_y = cell_xyz[:, :, 1] cell_z = cell_xyz[:, :, 2] cell_w = np.nansum(cell_weights, 1) ix = np.any(np.isnan(cell_timeseries), 1) if np.any(ix): print('nans (to be removed): %d' % np.count_nonzero(ix)) cell_timeseries[ix] = 0 cell_valids = np.zeros(len(cell_w), dtype=bool) for i, (li, xi, yi, zi) in enumerate(zip(cell_lengths, cell_x, cell_y, cell_z)): cell_valids[i] = np.mean(volume_mask[xi[:li], yi[:li], zi[:li]]) > p.thr_mask # brain mask array volume_list = [[[[] for zi in range(z)] for yi in range(y)] for xi in range(x)] volume_cell_n = np.zeros((x, y, z), dtype='int') for i, (li, vi) in enumerate(zip(cell_lengths, cell_valids)): for j in range(li if vi else 0): xij, yij, zij = cell_x[i, j], cell_y[i, j], cell_z[i, j] volume_list[xij][yij][zij].append(i) volume_cell_n[xij, yij, zij] += 1 pair_cells = [ pi for a in volume_list for b in a for c in b for pi in combinations(c, 2) ] assert (len(pair_cells) == np.sum(volume_cell_n * (volume_cell_n - 1) / 2)) # remove duplicate cells pair_id, pair_count = np.unique(pair_cells, axis=0, return_counts=True) for pi, fi in zip(pair_id, pair_count): pair_overlap = (fi / np.mean(cell_lengths[pi])) > thr_similarity pair_correlation = np.corrcoef( cell_timeseries[pi])[0, 1] > thr_similarity if (pair_overlap and pair_correlation): cell_valids[pi[np.argmin(cell_w[pi])]] = 0 ## get valid version of cells block_id = block_id[cell_valids] cell_weights = cell_weights[cell_valids].astype(dtype) cell_timeseries = cell_timeseries[cell_valids].astype(dtype) cell_lengths = cell_lengths[cell_valids] cell_x = cell_x[cell_valids] cell_y = cell_y[cell_valids] cell_z = cell_z[cell_valids] cell_w = cell_w[cell_valids] ## end get valid version of cells bparameters = sc.broadcast(parameters) def get_timebase(timeseries_tuple): timeseries = timeseries_tuple[1] return clean_signal(bparameters.value, timeseries) if p.parallel_clean: print('Computing baseline in parallel mode... ', end='') timebase = evenly_parallelize(cell_timeseries).map( get_timebase).collect() else: print('Computing baseline in serial mode... ', end='') timeseries_tuple = zip([[]] * len(cell_timeseries), cell_timeseries) timebase = map(get_timebase, timeseries_tuple) print('done.') cell_timeseries1, cell_baseline1 = list(zip(*timebase)) # convert to arrays cell_timeseries1 = np.array(cell_timeseries1) cell_baseline1 = np.array(cell_baseline1) # check that all series are in single precision assert (cell_weights.dtype == dtype) assert (cell_timeseries.dtype == dtype) assert (cell_timeseries1.dtype == dtype) assert (cell_baseline1.dtype == dtype) n = np.count_nonzero(cell_valids) volume_id = -1 + np.zeros((x, y, z)) volume_weight = np.zeros((x, y, z)) for i, li in enumerate(cell_lengths): for j in range(li): xij, yij, zij = cell_x[i, j], cell_y[i, j], cell_z[i, j] if cell_weights[i, j] > volume_weight[xij, yij, zij]: volume_id[xij, yij, zij] = i volume_weight[xij, yij, zij] = cell_weights[i, j] with h5py.File(fullname_volmean + hdf, 'r') as file_handle: background = file_handle['background'][()] with h5py.File(fullname_cells + hdf, 'w') as file_handle: file_handle['n'] = n file_handle['t'] = p.lt file_handle['x'] = x file_handle['y'] = y file_handle['z'] = z file_handle['cell_x'] = cell_x file_handle['cell_y'] = cell_y file_handle['cell_z'] = cell_z file_handle['block_id'] = block_id file_handle['volume_id'] = volume_id file_handle['volume_weight'] = volume_weight file_handle['cell_weights'] = cell_weights file_handle['cell_timeseries_raw'] = cell_timeseries file_handle['cell_timeseries'] = cell_timeseries1 file_handle['cell_baseline'] = cell_baseline1 file_handle['background'] = background # clean up completion = 1 for color_i in range(p.n_colors): fullname_cells = os.path.join(p.dir_output, 'cells%s_clean' % (color_i)) if not os.path.isfile(fullname_cells + hdf): completion = 0 if completion: try: shutil.rmtree(os.path.join(p.dir_output, 'volumes')) shutil.rmtree(os.path.join(p.dir_output, 'cells')) except: pass
def align_volumes(parameters): '''register volumes to a single middle volume''' # do not run if registration is set to none if not parameters['registration']: return import os import shutil from types import SimpleNamespace from voluseg._tools.load_volume import load_volume from voluseg._tools.save_volume import save_volume from voluseg._tools.constants import ori, ali, nii, hdf from voluseg._tools.ants_registration import ants_registration from voluseg._tools.evenly_parallelize import evenly_parallelize p = SimpleNamespace(**parameters) volume_nameRDD = evenly_parallelize(p.volume_names) for color_i in range(p.n_colors): fullname_volmean = os.path.join(p.dir_output, 'volume%d' % (color_i)) if os.path.isfile(fullname_volmean + hdf): continue dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) fullname_reference = os.path.join(dir_volume, 'reference') if load_volume(fullname_reference + nii) is None: fullname_median = os.path.join(dir_volume, p.volume_names[p.lt // 2]) shutil.copyfile(fullname_median + ori + nii, fullname_reference + nii) dir_transform = os.path.join(p.dir_output, 'transforms', str(color_i)) os.makedirs(dir_transform, exist_ok=True) def register_volume(tuple_name_volume): os.environ['ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS'] = '1' name_volume = tuple_name_volume[1] fullname_volume = os.path.join(dir_volume, name_volume) # skip processing if aligned volume exists if load_volume(fullname_volume + ali + hdf) is not None: return # setup registration cmd = ants_registration(dir_ants=p.dir_ants, in_nii=fullname_volume + ori + nii, ref_nii=fullname_reference + nii, out_nii=fullname_volume + ali + nii, prefix_out_tform=os.path.join( dir_transform, name_volume + '_tform_'), typ='r') if p.registration == 'high': pass elif p.registration == 'medium': cmd = cmd.replace('[1000x500x250x125]', '[1000x500x250]')\ .replace('12x8x4x2', '12x8x4')\ .replace('4x3x2x1vox', '4x3x2vox') elif p.registration == 'low': cmd = cmd.replace('[1000x500x250x125]', '[1000x500]')\ .replace('12x8x4x2', '12x8')\ .replace('4x3x2x1vox', '4x3vox') # run registration flag = os.system(cmd) if flag: # if breaks change initialization flag = os.system(cmd.replace(nii + ',1]', nii + ',0]')) if flag and load_volume(fullname_volume + ori + nii).shape[2] == 1: # if breaks change dimensionality flag = os.system( cmd.replace('--dimensionality 3', '--dimensionality 2')) if not flag: volume = load_volume(fullname_volume + ali + nii)[:, :, None] save_volume(fullname_volume + ali + nii, volume, p.affine_mat) if flag: raise Exception('volume %s not registered: flag %d.' % (name_volume, flag)) # load aligned volume volume = load_volume(fullname_volume + ali + nii) # remove padding if p.planes_pad: volume = volume[:, :, p.planes_pad:-p.planes_pad] # save as hdf5 volume = volume.T save_volume(fullname_volume + ali + hdf, volume) # remove nifti files if load_volume(fullname_volume + ali + hdf) is not None: try: os.remove(fullname_volume + ori + nii) os.remove(fullname_volume + ali + nii) except: pass volume_nameRDD.foreach(register_volume)
def detect_cells(parameters): '''detect cells in images''' import os import h5py import time import numpy as np from types import SimpleNamespace from voluseg._steps.step4a import define_blocks from voluseg._steps.step4b import process_block_data from voluseg._steps.step4c import initialize_block_cells from voluseg._steps.step4d import nnmf_sparse # from voluseg._steps.step4e import collect_blocks from voluseg._tools.evenly_parallelize import evenly_parallelize from voluseg._tools.clean_signal import clean_signal from voluseg._tools.ball import ball # set up spark from pyspark.sql.session import SparkSession spark = SparkSession.builder.getOrCreate() sc = spark.sparkContext p = SimpleNamespace(**parameters) ball_diam, ball_diam_xyz0 = ball(1.0 * p.diam_cell, p.affine_mat) # load plane filename for color_i in range(p.n_colors): if os.path.isfile( os.path.join(p.dir_output, 'cells%s_raw.hdf5' % (color_i))): continue dir_cell = os.path.join(p.dir_output, 'cells', str(color_i)) os.makedirs(dir_cell, exist_ok=True) with h5py.File(os.path.join(p.dir_output, 'volume%s.hdf5' % (color_i)), 'r') as file_handle: volume_mean = file_handle['volume_mean'][()].T volume_mask = file_handle['volume_mask'][()].T volume_peak = file_handle['volume_peak'][()].T timeseries_mean = file_handle['timeseries_mean'][()] if 'n_blocks' in file_handle.keys(): flag = 0 n_voxels_cell = file_handle['n_voxels_cell'][()] n_blocks = file_handle['n_blocks'][()] block_valids = file_handle['block_valids'][()] xyz0 = file_handle['block_xyz0'][()] xyz1 = file_handle['block_xyz1'][()] timepoints = file_handle['timepoints'][()] else: flag = 1 # broadcast image peaks (for initialization) and volume_mean (for renormalization) bvolume_peak = sc.broadcast(volume_peak) bvolume_mean = sc.broadcast(volume_mean) # dimensions and resolution lxyz = volume_mean.shape rxyz = np.diag(p.affine_mat)[:3] # compute number of blocks (do only once) if flag: lx, ly, lz = lxyz rx, ry, rz = rxyz # get number of voxels in cell if (lz == 1) or (rz >= p.diam_cell): # area of a circle n_voxels_cell = np.pi * ((p.diam_cell / 2.0)**2) / (rx * ry) else: # volume of a cylinder (change to sphere later) n_voxels_cell = p.diam_cell * np.pi * ( (p.diam_cell / 2.0)**2) / (rx * ry * rz) n_voxels_cell = np.round(n_voxels_cell).astype(int) # get number of voxels in each cell n_blocks, block_valids, xyz0, xyz1 = \ define_blocks(lx, ly, lz, p.n_cells_block, n_voxels_cell, volume_mask) # get timepoints for cell detection if not p.nt: timepoints = np.range(p.lt) else: timeseries1, baseline1 = clean_signal(parameters, timeseries_mean) timepoints = np.sort( np.argsort( (timeseries1 - baseline1) / timeseries1)[::-1][:p.nt]) # save number and indices of blocks with h5py.File( os.path.join(p.dir_output, 'volume%s.hdf5' % (color_i)), 'r+') as file_handle: file_handle['n_voxels_cell'] = n_voxels_cell file_handle['n_blocks'] = n_blocks file_handle['block_valids'] = block_valids file_handle['block_xyz0'] = xyz0 file_handle['block_xyz1'] = xyz1 if not 'timepoints' in file_handle: file_handle['timepoints'] = timepoints print('number of blocks, total: %d.' % (block_valids.sum())) for ii in np.where(block_valids)[0]: try: with h5py.File(os.path.join(dir_cell, 'block%05d.hdf5' % (ii)), 'r') as file_handle: if ('completion' in file_handle.keys() ) and file_handle['completion'][()]: block_valids[ii] = 0 except (NameError, OSError): pass print('number of blocks, remaining: %d.' % (block_valids.sum())) ix = np.where(block_valids)[0] block_ixyz01 = list(zip(ix, xyz0[ix], xyz1[ix])) # detect individual cells with sparse nnmf algorithm def detect_cells_block(tuple_i_xyz0_xyz1): os.environ['MKL_NUM_THREADS'] = '1' ii, xyz0, xyz1 = tuple_i_xyz0_xyz1[1] voxel_xyz, voxel_timeseries, peak_idx, voxel_similarity_peak = \ process_block_data(xyz0, xyz1, parameters, color_i, lxyz, rxyz, \ ball_diam, bvolume_mean, bvolume_peak, timepoints) n_voxels_block = len(voxel_xyz) # number of voxels in block voxel_fraction_peak = np.argsort( ((voxel_timeseries[peak_idx])**2).mean(1)) / len(peak_idx) for fraction in np.r_[1:0:-0.05]: try: peak_valids = (voxel_fraction_peak >= (1 - fraction)) # valid voxel indices n_cells = np.round(peak_valids.sum() / (0.5 * n_voxels_cell)).astype(int) print((fraction, n_cells)) tic = time.time() voxel_timeseries_valid, voxel_xyz_valid, cell_weight_init_valid, \ cell_neighborhood_valid, cell_sparseness = \ initialize_block_cells( n_voxels_cell, n_voxels_block, n_cells, \ voxel_xyz, voxel_timeseries, peak_idx, peak_valids, voxel_similarity_peak, \ lxyz, rxyz, ball_diam, ball_diam_xyz0) print('cell initialization: %.1f minutes.\n' % ((time.time() - tic) / 60)) tic = time.time() cell_weights_valid, cell_timeseries_valid, d = nnmf_sparse( voxel_timeseries_valid, voxel_xyz_valid, cell_weight_init_valid, cell_neighborhood_valid, cell_sparseness, timepoints=timepoints, miniter=10, maxiter=100, tolfun=1e-3) success = 1 print('cell factorization: %.1f minutes.\n' % ((time.time() - tic) / 60)) break except ValueError as msg: print('retrying factorization of block %d: %s' % (ii, msg)) success = 0 # get cell positions and timeseries, and save cell data with h5py.File(os.path.join(dir_cell, 'block%05d.hdf5' % (ii)), 'w') as file_handle: if success: for ci in range(n_cells): ix = cell_weights_valid[:, ci] > 0 xyzi = voxel_xyz_valid[ix] wi = cell_weights_valid[ix, ci] bi = np.sum(wi * bvolume_mean.value[list( zip(*xyzi))]) / np.sum(wi) ti = bi * cell_timeseries_valid[ci] / np.mean( cell_timeseries_valid[ci]) file_handle['/cell/%05d/xyz' % (ci)] = xyzi file_handle['/cell/%05d/weights' % (ci)] = wi file_handle['/cell/%05d/timeseries' % (ci)] = ti file_handle['n_cells'] = n_cells file_handle['completion'] = 1 if block_valids.any(): evenly_parallelize(block_ixyz01).foreach(detect_cells_block)
def align_images(parameters): '''register images to a single middle image''' # do not run if registration is set to none if not parameters['registration']: return import os import h5py import shutil import nibabel from types import SimpleNamespace from voluseg._tools.nii_image import nii_image from voluseg._tools.ants_registration import ants_registration from voluseg._tools.evenly_parallelize import evenly_parallelize p = SimpleNamespace(**parameters) volume_nameRDD = evenly_parallelize(p.volume_names) for color_i in range(p.n_colors): if os.path.isfile( os.path.join(p.dir_output, 'volume%d.hdf5' % (color_i))): continue dir_volume = os.path.join(p.dir_output, 'volumes', str(color_i)) fullname_reference = os.path.join(dir_volume, 'reference_original.nii.gz') if not os.path.isfile(fullname_reference): fullname_lt_2 = os.path.join( dir_volume, p.volume_names[p.lt // 2] + '_original.nii.gz') shutil.copyfile(fullname_lt_2, fullname_reference) dir_transform = os.path.join(p.dir_output, 'transforms', str(color_i)) os.makedirs(dir_transform, exist_ok=True) def register_volume(tuple_name_volume): os.environ['ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS'] = '1' name_volume = tuple_name_volume[1] fullname_original = os.path.join(dir_volume, name_volume + '_original.nii.gz') fullname_aligned = os.path.join(dir_volume, name_volume + '_aligned.nii.gz') fullname_aligned_hdf = fullname_aligned.replace('.nii.gz', '.hdf5') if os.path.isfile(fullname_aligned): try: volume_aligned = nibabel.load(fullname_aligned).get_data() return except: pass if os.path.isfile(fullname_aligned_hdf): try: with h5py.File(fullname_aligned_hdf) as file_handle: volume_aligned = file_handle['V3D'][()].T return except: pass cmd = ants_registration(dir_ants=p.dir_ants, in_nii=fullname_original, ref_nii=fullname_reference, out_nii=fullname_aligned, prefix_out_tform=os.path.join( dir_transform, name_volume + '_tform_'), typ='r') if p.registration == 'high': pass elif p.registration == 'medium': cmd = cmd.replace('[1000x500x250x125]','[1000x500x250]')\ .replace('12x8x4x2', '12x8x4')\ .replace('4x3x2x1vox', '4x3x2vox') elif p.registration == 'low': cmd = cmd.replace('[1000x500x250x125]','[1000x500]')\ .replace('12x8x4x2', '12x8')\ .replace('4x3x2x1vox', '4x3vox') else: raise Exception('unknown registration type.') flag = os.system(cmd) if flag: flag = os.system(cmd.replace('.nii.gz,1]', '.nii.gz,0]')) if flag and nibabel.load(fullname_original).shape[2] == 1: os.system( cmd.replace('--dimensionality 3', '--dimensionality 2')) volume_input = nibabel.load(fullname_aligned).get_data()[:, :, None] nibabel.save(nii_image(volume_input, p.affine_mat), fullname_aligned) # remove padding if p.planes_pad: volume_aligned = nibabel.load(fullname_aligned).get_data() volume_aligned = volume_aligned[:, :, p.planes_pad:-p.planes_pad] nibabel.save(nii_image(volume_aligned, p.affine_mat), fullname_aligned) if os.path.isfile(fullname_aligned): try: volume_aligned = nibabel.load(fullname_aligned).get_data() os.remove(fullname_original) except: pass else: raise Exception('image %s not registered: flag %d.' % (name_volume, flag)) volume_nameRDD.foreach(register_volume)
def collect_blocks(color_i, parameters): '''collect cells across all blocks''' import os import h5py import numpy as np from types import SimpleNamespace from voluseg._tools.evenly_parallelize import evenly_parallelize # set up spark import pyspark from pyspark.sql.session import SparkSession spark = SparkSession.builder.getOrCreate() sc = spark.sparkContext p = SimpleNamespace(**parameters) dir_cell = os.path.join(p.dir_output, 'cells', str(color_i)) with h5py.File(os.path.join(p.dir_output, 'volume%d.hdf5' % (color_i)), 'r') as file_handle: block_valids = file_handle['block_valids'][()] class accum_data(pyspark.accumulators.AccumulatorParam): '''define accumulator class''' def zero(self, val0): return [[]] * 3 def addInPlace(self, val1, val2): return [val1[i] + val2[i] for i in range(3)] # cumulate collected cells cell_data = sc.accumulator([[]] * 3, accum_data()) def add_data(tuple_ii): ii = tuple_ii[1] try: cell_xyz = [] cell_weights = [] cell_timeseries = [] with h5py.File(os.path.join(dir_cell, 'block%05d.hdf5' % (ii)), 'r') as file_handle: for ci in range(file_handle['n_cells'][()]): cell_xyz.append(file_handle['/cell/%05d/xyz' % (ci)][()]) cell_weights.append(file_handle['/cell/%05d/weights' % (ci)][()]) cell_timeseries.append( file_handle['/cell/%05d/timeseries' % (ci)][()]) cell_data.add([cell_xyz, cell_weights, cell_timeseries]) except KeyError: print('block %d is empty.' % ii) except IOError: print('block %d does not exist.' % ii) evenly_parallelize(np.argwhere(block_valids).T[0]).foreach(add_data) cell_xyz, cell_weights, cell_timeseries = cell_data.value # convert lists to arrays cn = len(cell_xyz) cell_lengths = np.array([len(i) for i in cell_weights]) cell_xyz_array = np.full((cn, np.max(cell_lengths), 3), -1, dtype=int) cell_weights_array = np.full((cn, np.max(cell_lengths)), np.nan) for ci, li in enumerate(cell_lengths): cell_xyz_array[ci, :li] = cell_xyz[ci] cell_weights_array[ci, :li] = cell_weights[ci] cell_timeseries_array = np.array(cell_timeseries) return cell_xyz_array, cell_weights_array, cell_timeseries_array, cell_lengths
def process_parameters(parameters0=None): '''process parameters and create parameter file''' import os import copy import pickle import numpy as np from voluseg._tools.load_volume import load_volume from voluseg._tools.plane_name import plane_name from voluseg._tools.parameter_dictionary import parameter_dictionary from voluseg._tools.evenly_parallelize import evenly_parallelize parameters = copy.deepcopy(parameters0) ## general checks # check that parameter input is a dictionary if not type(parameters) == dict: print('error: specify parameter dictionary as input.') return # check if any parameters are missing missing_parameters = set(parameter_dictionary()) - set(parameters) if missing_parameters: print('error: missing parameters \'%s\'.' % ('\', \''.join(missing_parameters))) return # get input and output directories, and parameter filename dir_input = parameters['dir_input'] dir_output = parameters['dir_output'] filename_parameters = os.path.join(dir_output, 'parameters.pickle') # load parameters from file, if it already exists if os.path.isfile(filename_parameters): print('exiting, parameter file exists: %s.' % (filename_parameters)) return ## specific checks # check directory names for i in ['dir_ants', 'dir_input', 'dir_output', 'registration']: pi = parameters[i] if not (isinstance(pi, str) and (' ' not in pi)): print('error: \'%s\' must be a string without spaces.' % (i)) return # check booleans for i in ['parallel_clean', 'planes_packed']: pi = parameters[i] if not isinstance(pi, bool): print('error: \'%s\' must be a boolean.' % (i)) return # check integers for i in ['ds', 'n_cells_block', 'n_colors', 'nt', 'planes_pad']: pi = parameters[i] if not (np.isscalar(pi) and (pi >= 0) and (pi == np.round(pi))): print('error: \'%s\' must be a nonnegative or positive integer.' % (i)) return # check non-negative real numbers: for i in [ 'diam_cell', 'f_hipass', 'f_volume', 'res_x', 'res_y', 'res_z', 't_baseline', 't_section', 'thr_mask' ]: pi = parameters[i] if not (np.isscalar(pi) and (pi >= 0) and np.isreal(pi)): print( 'error: \'%s\' must be a nonnegative or positive real number.' % (i)) return # check registration if parameters['registration']: parameters['registration'] = parameters['registration'].lower() if parameters['registration'] == 'none': parameters['registration'] = None elif not parameters['registration'] in ['high', 'medium', 'low']: print( 'error: \'registration\' must be \'high\', \'medium\', \'low\', or \'none\'.' ) return # check plane padding if (not parameters['registration']) and not ( (parameters['planes_pad'] == 0)): print('error: \'planes_pad\' must be 0 if \'registration\' is None.') return # get volume extension, volume names and number of segmentation timepoints file_names = [i.split('.', 1) for i in os.listdir(dir_input) if '.' in i] file_exts, counts = np.unique(list(zip(*file_names))[1], return_counts=True) ext = '.' + file_exts[np.argmax(counts)] volume_names = np.sort([i for i, j in file_names if '.' + j == ext]) lt = len(volume_names) # adjust parameters for packed planes data if parameters['planes_packed']: volume_names0 = copy.deepcopy(volume_names) parameters['volume_names0'] = volume_names0 parameters['res_z'] = parameters['diam_cell'] def volume_plane_names(tuple_name_volume): name_volume = tuple_name_volume[1] fullname_volume = os.path.join(dir_input, name_volume) lp = len(load_volume(fullname_volume + ext)) return [plane_name(name_volume, pi) for pi in range(lp)] volume_names = evenly_parallelize(volume_names0).map( volume_plane_names).collect() volume_names = np.sort([pi for ni in volume_names for pi in ni]) lt = len(volume_names) # affine matrix affine_mat = np.diag([ parameters['res_x'] * parameters['ds'], parameters['res_y'] * parameters['ds'], parameters['res_z'], 1 ]) # save parameters parameters['volume_names'] = volume_names parameters['ext'] = ext parameters['lt'] = lt parameters['affine_mat'] = affine_mat try: os.makedirs(dir_output, exist_ok=True) with open(filename_parameters, 'wb') as file_handle: pickle.dump(parameters, file_handle) print('parameter file successfully saved.') except Exception as msg: print('parameter file not saved: %s.' % (msg))
def collect_blocks(color_i, parameters): '''collect cells across all blocks''' import os import h5py import numpy as np from types import SimpleNamespace from voluseg._tools.constants import hdf from voluseg._tools.evenly_parallelize import evenly_parallelize # set up spark import pyspark from pyspark.sql.session import SparkSession spark = SparkSession.builder.getOrCreate() sc = spark.sparkContext p = SimpleNamespace(**parameters) dir_cell = os.path.join(p.dir_output, 'cells', str(color_i)) fullname_volmean = os.path.join(p.dir_output, 'volume%d' % (color_i)) with h5py.File(fullname_volmean + hdf, 'r') as file_handle: block_valids = file_handle['block_valids'][()] class accum_data(pyspark.accumulators.AccumulatorParam): '''define accumulator class''' def zero(self, val0): return [[]] * 4 def addInPlace(self, val1, val2): return [val1[i] + val2[i] for i in range(4)] # cumulate collected cells if p.parallel_clean: cell_data = sc.accumulator([[]] * 4, accum_data()) def add_data(tuple_ii): ii = tuple_ii[1] try: block_id = [] cell_xyz = [] cell_weights = [] cell_timeseries = [] fullname_block = os.path.join(dir_cell, 'block%05d' % (ii)) with h5py.File(fullname_block + hdf, 'r') as file_handle: for ci in range(file_handle['n_cells'][()]): block_id.append(ii) cell_xyz.append(file_handle['/cell/%05d/xyz' % (ci)][()]) cell_weights.append(file_handle['/cell/%05d/weights' % (ci)][()]) cell_timeseries.append( file_handle['/cell/%05d/timeseries' % (ci)][()]) if p.parallel_clean: cell_data.add( [block_id, cell_xyz, cell_weights, cell_timeseries]) else: return [block_id, cell_xyz, cell_weights, cell_timeseries] except KeyError: print('block %d is empty.' % ii) except IOError: print('block %d does not exist.' % ii) if p.parallel_clean: evenly_parallelize(np.argwhere(block_valids).T[0]).foreach(add_data) block_id, cell_xyz, cell_weights, cell_timeseries = cell_data.value else: idx_block_valids = np.argwhere(block_valids).T[0] valids_tuple = zip([[]] * len(idx_block_valids), idx_block_valids) block_id, cell_xyz, cell_weights, cell_timeseries = list( zip(*map(add_data, valids_tuple))) block_id = [ii for bi in block_id for ii in bi] cell_xyz = [xyzi for ci in cell_xyz for xyzi in ci] cell_weights = [wi for ci in cell_weights for wi in ci] cell_timeseries = [ti for ci in cell_timeseries for ti in ci] # convert lists to arrays cn = len(cell_xyz) block_id = np.array(block_id) cell_lengths = np.array([len(i) for i in cell_weights]) cell_xyz_array = np.full((cn, np.max(cell_lengths), 3), -1, dtype=int) cell_weights_array = np.full((cn, np.max(cell_lengths)), np.nan) for ci, li in enumerate(cell_lengths): cell_xyz_array[ci, :li] = cell_xyz[ci] cell_weights_array[ci, :li] = cell_weights[ci] cell_timeseries_array = np.array(cell_timeseries) return block_id, cell_xyz_array, cell_weights_array, cell_timeseries_array, cell_lengths