def test_eigen_probe(self): leading = (2, ) wide = 18 high = 21 posi = 53 eigen = 1 comm = Comm(2, None) R = comm.pool.bcast(np.random.rand(*leading, posi, 1, 1, wide, high)) eigen_probe = comm.pool.bcast( np.random.rand(*leading, 1, eigen, 1, wide, high)) weights = np.random.rand(*leading, posi) weights -= np.mean(weights) weights = comm.pool.bcast(weights) patches = comm.pool.bcast( np.random.rand(*leading, posi, 1, 1, wide, high)) diff = comm.pool.bcast(np.random.rand(*leading, posi, 1, 1, wide, high)) new_probe, new_weights = tike.ptycho.probe.update_eigen_probe( comm=comm, R=R, eigen_probe=eigen_probe, weights=weights, patches=patches, diff=diff, ) assert eigen_probe[0].shape == new_probe[0].shape
def test_eigen_probe(self): leading = (2,) wide = 18 high = 21 posi = 53 eigen = 1 comm = Comm(2, None) R = comm.pool.bcast([np.random.rand(*leading, posi, 1, 1, wide, high)]) eigen_probe = comm.pool.bcast( [np.random.rand(*leading, 1, eigen, 1, wide, high)]) weights = np.random.rand(*leading, posi, eigen + 1, 1) weights -= np.mean(weights, axis=-3, keepdims=True) weights = comm.pool.bcast([weights]) patches = comm.pool.bcast( [np.random.rand(*leading, posi, 1, 1, wide, high)]) diff = comm.pool.bcast( [np.random.rand(*leading, posi, 1, 1, wide, high)]) new_probe, new_weights = tike.ptycho.probe.update_eigen_probe( comm=comm, R=R, eigen_probe=eigen_probe, weights=weights, patches=patches, diff=diff, c=1, m=0, ) assert eigen_probe[0].shape == new_probe[0].shape
class TestComm(unittest.TestCase): def setUp(self, workers=4): self.comm = Comm(workers) self.xp = self.comm.pool.xp def test_reduce(self): a = self.xp.ones((1, )) a_list = self.comm.pool.bcast([a]) a = a * self.comm.pool.num_workers result = self.comm.reduce(a_list, 'cpu') self.xp.testing.assert_array_equal(a, result) result = self.comm.reduce(a_list, 'gpu') self.xp.testing.assert_array_equal(a, result[0]) def test_Allreduce_reduce(self): a = self.xp.ones((1, )) a_list = self.comm.pool.bcast([a]) a = a * self.comm.pool.num_workers * self.comm.mpi.size result = self.comm.Allreduce_reduce(a_list, 'cpu') self.xp.testing.assert_array_equal(a, result) result = self.comm.Allreduce_reduce(a_list, 'gpu') self.xp.testing.assert_array_equal(a, result[0])
def setUp(self, workers=4): self.comm = Comm(workers) self.xp = self.comm.pool.xp
def reconstruct( data, theta, tilt, algorithm, obj=None, num_iter=1, rtol=-1, eps=1e-1, num_gpu=1, obj_split=1, use_mpi=False, **kwargs ): # yapf: disable """Solve the Laminography problem using the given `algorithm`. Parameters ---------- algorithm : string The name of one algorithms from :py:mod:`.lamino.solvers`. rtol : float Terminate early if the relative decrease of the cost function is less than this amount. """ n = data.shape[2] if use_mpi is True: mpi = MPIComm if obj is None: raise ValueError( "When MPI is enabled, initial object guess cannot be None.") else: mpi = None obj = np.zeros([n, n, n], dtype='complex64') if obj is None else obj if algorithm in solvers.__all__: # Initialize an operator. with Lamino( n=obj.shape[-1], tilt=tilt, eps=eps, **kwargs, ) as operator, Comm(num_gpu, mpi) as comm: # send any array-likes to device obj_split = max(1, min(comm.pool.num_workers, obj_split)) data_split = comm.pool.num_workers // obj_split data = np.array_split(data.astype('complex64'), data_split) data = comm.pool.scatter(data, obj_split) theta = np.array_split(theta.astype('float32'), data_split) theta = comm.pool.scatter(theta, obj_split) obj = np.array_split(obj.astype('complex64'), obj_split) if comm.use_mpi is True: grid = operator._make_grid(comm.mpi.size, comm.mpi.rank) else: grid = operator._make_grid() grid = np.array_split(grid.astype('int16'), obj_split) grid = [x.reshape(x.shape[0] * n * n, 3) for x in grid] grid = comm.pool.bcast(grid, obj_split) result = { 'obj': comm.pool.bcast(obj, obj_split), } for key, value in kwargs.items(): if np.ndim(value) > 0: kwargs[key] = comm.pool.bcast([value]) logger.info("{} on {:,d} by {:,d} by {:,d} volume for {:,d} " "iterations.".format(algorithm, *obj[0].shape, num_iter)) costs = [] for i in range(num_iter): kwargs.update(result) result = getattr(solvers, algorithm)( operator, comm, data=data, theta=theta, grid=grid, obj_split=obj_split, **kwargs, ) if result['cost'] is not None: costs.append(result['cost']) # Check for early termination if ( len(costs) > 1 and abs((costs[-1] - costs[-2]) / costs[-2]) < rtol ): # yapf: disable logger.info( "Cost function rtol < %g reached at %d " "iterations.", rtol, i) break result['cost'] = operator.asarray(costs) for k, v in result.items(): if isinstance(v, list): result[k] = np.concatenate( [operator.asnumpy(part) for part in v[:obj_split]], axis=0, ) elif np.ndim(v) > 0: result[k] = operator.asnumpy(v) return result else: raise ValueError( "The '{}' algorithm is not an available.".format(algorithm))
def reconstruct( data, probe, scan, algorithm, psi=None, num_gpu=1, num_iter=1, rtol=-1, model='gaussian', use_mpi=False, cost=None, times=None, eigen_probe=None, eigen_weights=None, batch_size=None, **kwargs ): # yapf: disable """Solve the ptychography problem using the given `algorithm`. Parameters ---------- data : (..., FRAME, WIDE, HIGH) float32 The intensity (square of the absolute value) of the propagated wavefront; i.e. what the detector records. eigen_probe : (..., 1, EIGEN, SHARED, WIDE, HIGH) complex64 The eigen probes for all positions. eigen_weights : (..., POSI, EIGEN, SHARED) float32 The relative intensity of the eigen probes at each position. psi : (..., WIDE, HIGH) complex64 The wavefront modulation coefficients of the object. probe : (..., 1, 1, SHARED, WIDE, HIGH) complex64 The shared complex illumination function amongst all positions. scan : (..., POSI, 2) float32 Coordinates of the minimum corner of the probe grid for each measurement in the coordinate system of psi. Coordinate order consistent with WIDE, HIGH order. algorithm : string The name of one algorithms from :py:mod:`.ptycho.solvers`. rtol : float Terminate early if the relative decrease of the cost function is less than this amount. batch_size : int The approximate number of scan positions processed by each GPU simultaneously per view. """ (psi, scan) = get_padded_object(scan, probe) if psi is None else (psi, scan) check_allowed_positions(scan, psi, probe) if use_mpi is True: mpi = MPIComm else: mpi = None if algorithm in solvers.__all__: # Initialize an operator. with Ptycho( probe_shape=probe.shape[-1], detector_shape=data.shape[-1], nz=psi.shape[-2], n=psi.shape[-1], ntheta=scan.shape[0], model=model, ) as operator, Comm(num_gpu, mpi) as comm: logger.info("{} for {:,d} - {:,d} by {:,d} frames for {:,d} " "iterations.".format(algorithm, *data.shape[-3:], num_iter)) num_batch = 1 if batch_size is None else max( 1, int(data.shape[-3] / batch_size / comm.pool.num_workers), ) # Divide the inputs into regions odd_pool = comm.pool.num_workers % 2 order, scan, data, eigen_weights = split_by_scan_grid( comm.pool, ( comm.pool.num_workers if odd_pool else comm.pool.num_workers // 2, 1 if odd_pool else 2, ), scan, data, eigen_weights, ) result = { 'psi': comm.pool.bcast(psi.astype('complex64')), 'probe': comm.pool.bcast(probe.astype('complex64')), 'eigen_probe': comm.pool.bcast(eigen_probe.astype('complex64')) if eigen_probe is not None else None, 'scan': scan, 'eigen_weights': eigen_weights, } for key, value in kwargs.items(): if np.ndim(value) > 0: kwargs[key] = comm.pool.bcast(value) result['probe'] = comm.pool.bcast( _rescale_obj_probe( operator, comm, data[0], result['psi'][0], scan[0], result['probe'][0], num_batch=num_batch, )) costs = [] times = [] start = time.perf_counter() for i in range(num_iter): logger.info(f"{algorithm} epoch {i:,d}") kwargs.update(result) result = getattr(solvers, algorithm)( operator, comm, data=data, num_batch=num_batch, **kwargs, ) if result['cost'] is not None: costs.append(result['cost']) times.append(time.perf_counter() - start) start = time.perf_counter() # Check for early termination if i > 0 and abs((costs[-1] - costs[-2]) / costs[-2]) < rtol: logger.info( "Cost function rtol < %g reached at %d " "iterations.", rtol, i) break reorder = np.argsort(np.concatenate(order)) result['scan'] = comm.pool.gather(scan, axis=1)[:, reorder] if 'eigen_weights' in result: result['eigen_weights'] = comm.pool.gather( eigen_weights, axis=1, )[:, reorder] result['eigen_probe'] = result['eigen_probe'][0] result['probe'] = result['probe'][0] result['cost'] = operator.asarray(costs) result['times'] = operator.asarray(times) for k, v in result.items(): if isinstance(v, list): result[k] = v[0] return {k: operator.asnumpy(v) for k, v in result.items()} else: raise ValueError(f"The '{algorithm}' algorithm is not an option.\n" f"\tAvailable algorithms are : {solvers.__all__}")
def reconstruct( data, theta, tilt, algorithm, obj=None, num_iter=1, rtol=-1, eps=1e-3, num_gpu=1, **kwargs ): # yapf: disable """Solve the Laminography problem using the given `algorithm`. Parameters ---------- algorithm : string The name of one algorithms from :py:mod:`.lamino.solvers`. rtol : float Terminate early if the relative decrease of the cost function is less than this amount. tilt : float32 [radians] The tilt angle; the angle between the rotation axis of the object and the light source. π / 2 for conventional tomography. 0 for a beam path along the rotation axis. obj : (nz, n, n) complex64 The complex refractive index of the object. nz is the axis corresponding to the rotation axis. data : (ntheta, n, n) complex64 The complex projection data of the object. theta : array-like float32 [radians] The projection angles; rotation around the vertical axis of the object. """ n = data.shape[2] obj = np.zeros([n, n, n], dtype='complex64') if obj is None else obj if algorithm in solvers.__all__: # Initialize an operator. with Lamino( n=obj.shape[-1], tilt=tilt, eps=eps, **kwargs, ) as operator, Comm(num_gpu, mpi=None) as comm: # send any array-likes to device data = np.array_split(data.astype('complex64'), comm.pool.num_workers) data = comm.pool.scatter(data) theta = np.array_split(theta.astype('float32'), comm.pool.num_workers) theta = comm.pool.scatter(theta) result = { 'obj': comm.pool.bcast([obj.astype('complex64')]), } for key, value in kwargs.items(): if np.ndim(value) > 0: kwargs[key] = comm.pool.bcast([value]) logger.info("{} on {:,d} by {:,d} by {:,d} volume for {:,d} " "iterations.".format(algorithm, *obj.shape, num_iter)) costs = [] for i in range(num_iter): kwargs.update(result) result = getattr(solvers, algorithm)( operator, comm, data=data, theta=theta, **kwargs, ) if result['cost'] is not None: costs.append(result['cost']) # Check for early termination if ( len(costs) > 1 and abs((costs[-1] - costs[-2]) / costs[-2]) < rtol ): # yapf: disable logger.info( "Cost function rtol < %g reached at %d " "iterations.", rtol, i) break result['cost'] = operator.asarray(costs) for k, v in result.items(): if isinstance(v, list): result[k] = v[0] return { k: operator.asnumpy(v) if np.ndim(v) > 0 else v for k, v in result.items() } else: raise ValueError( "The '{}' algorithm is not an available.".format(algorithm))
def reconstruct( data, probe, scan, algorithm, psi=None, num_gpu=1, num_iter=1, rtol=-1, model='gaussian', use_mpi=False, cost=None, times=None, batch_size=None, subset_is_random=None, eigen_probe=None, eigen_weights=None, **kwargs ): # yapf: disable """Solve the ptychography problem using the given `algorithm`. Parameters ---------- algorithm : string The name of one algorithms from :py:mod:`.ptycho.solvers`. rtol : float Terminate early if the relative decrease of the cost function is less than this amount. split : 'grid' or 'stripe' The method to use for splitting the scan positions among GPUS. """ (psi, scan) = get_padded_object(scan, probe) if psi is None else (psi, scan) check_allowed_positions(scan, psi, probe) if use_mpi is True: mpi = MPIComm else: mpi = None if algorithm in solvers.__all__: # Initialize an operator. with Ptycho( probe_shape=probe.shape[-1], detector_shape=data.shape[-1], nz=psi.shape[-2], n=psi.shape[-1], ntheta=scan.shape[0], model=model, ) as operator, Comm(num_gpu, mpi) as comm: logger.info("{} for {:,d} - {:,d} by {:,d} frames for {:,d} " "iterations.".format(algorithm, *data.shape[-3:], num_iter)) # Divide the inputs into regions and mini-batches num_batch = 1 if batch_size is not None: num_batch = max( 1, int(data.shape[1] / batch_size / comm.pool.num_workers), ) odd_pool = comm.pool.num_workers % 2 order = np.arange(data.shape[1]) order, data, scan, eigen_weights = split_by_scan_grid( order, data, scan, ( comm.pool.num_workers if odd_pool else comm.pool.num_workers // 2, 1 if odd_pool else 2, ), eigen_weights=eigen_weights, ) order, data, scan, eigen_weights = zip(*comm.pool.map( _make_mini_batches, order, data, scan, eigen_weights, num_batch=num_batch, subset_is_random=subset_is_random, )) result = { 'psi': comm.pool.bcast(psi.astype('complex64')), 'probe': comm.pool.bcast(probe.astype('complex64')), 'eigen_probe': comm.pool.bcast(eigen_probe.astype('complex64')) if eigen_probe is not None else None, } for key, value in kwargs.items(): if np.ndim(value) > 0: kwargs[key] = comm.pool.bcast(value) result['probe'] = comm.pool.bcast( _rescale_obj_probe( operator, comm, data[0][0], result['psi'][0], scan[0][0], result['probe'][0], )) costs = [] times = [] start = time.perf_counter() for i in range(num_iter): logger.info(f"{algorithm} epoch {i:,d}") for b in randomizer.permutation(num_batch): kwargs.update(result) kwargs['scan'] = [s[b] for s in scan] kwargs['eigen_weights'] = [w[b] for w in eigen_weights] result = getattr(solvers, algorithm)( operator, comm, data=[d[b] for d in data], **kwargs, ) if result['cost'] is not None: costs.append(result['cost']) for g in range(comm.pool.num_workers): scan[g][b] = result['scan'][g] eigen_weights[g][b] = result['eigen_weights'][ g] if 'eigen_weights' in result else None times.append(time.perf_counter() - start) start = time.perf_counter() # Check for early termination if i > 0 and abs((costs[-1] - costs[-2]) / costs[-2]) < rtol: logger.info( "Cost function rtol < %g reached at %d " "iterations.", rtol, i) break reorder = np.argsort( np.concatenate(list(chain.from_iterable(order)))) result['scan'] = comm.pool.gather( list(comm.pool.map(cp.concatenate, scan, axis=1)), axis=1, )[:, reorder] if 'eigen_weights' in result: result['eigen_weights'] = comm.pool.gather( list(comm.pool.map(cp.concatenate, eigen_weights, axis=1)), axis=1, )[:, reorder] result['eigen_probe'] = result['eigen_probe'][0] result['probe'] = result['probe'][0] result['cost'] = operator.asarray(costs) result['times'] = operator.asarray(times) for k, v in result.items(): if isinstance(v, list): result[k] = v[0] return {k: operator.asnumpy(v) for k, v in result.items()} else: raise ValueError(f"The '{algorithm}' algorithm is not an option.\n" f"\tAvailable algorithms are : {solvers.__all__}")
def reconstruct( data, probe, psi, scan, algorithm_options=solvers.RpieOptions(), eigen_probe=None, eigen_weights=None, model='gaussian', num_gpu=1, object_options=None, position_options=None, probe_options=None, use_mpi=False, ): """Solve the ptychography problem using the given `algorithm`. Parameters ---------- data : (FRAME, WIDE, HIGH) float32 The intensity (square of the absolute value) of the propagated wavefront; i.e. what the detector records. FFT-shifted so the diffraction peak is at the corners. probe : (1, 1, SHARED, WIDE, HIGH) complex64 The shared complex illumination function amongst all positions. scan : (POSI, 2) float32 Coordinates of the minimum corner of the probe grid for each measurement in the coordinate system of psi. Coordinate order consistent with WIDE, HIGH order. algorithm_options : :py:class:`tike.ptycho.solvers.IterativeOptions` A class containing algorithm specific parameters eigen_probe : (EIGEN, SHARED, WIDE, HIGH) complex64 The eigen probes for all positions. eigen_weights : (POSI, EIGEN, SHARED) float32 The relative intensity of the eigen probes at each position. model : "gaussian", "poisson" The noise model to use for the cost function. num_gpu : int, tuple(int) The number of GPUs to use or a tuple of the device numbers of the GPUs to use. If the number of GPUs is less than the requested number, only workers for the available GPUs are allocated. object_options : :py:class:`tike.ptycho.ObjectOptions` A class containing settings related to object updates. position_options : :py:class:`tike.ptycho.PositionOptions` A class containing settings related to position correction. probe_options : :py:class:`tike.ptycho.ProbeOptions` A class containing settings related to probe updates. psi : (WIDE, HIGH) complex64 The wavefront modulation coefficients of the object. use_mpi : bool Whether to use MPI or not. Raises ------ ValueError When shapes are incorrect for various input parameters. Returns ------- result : dict A dictionary of the above parameters that may be passed to this function to resume reconstruction from the previous state. """ if (np.any(np.asarray(data.shape) < 1) or data.ndim != 3 or data.shape[-2] != data.shape[-1]): raise ValueError( f"data shape {data.shape} is incorrect. " "It should be (N, W, H), " "where N >= 1 is the number of square diffraction patterns.") if (scan.ndim != 2 or scan.shape[1] != 2 or np.any(np.asarray(scan.shape) < 1)): raise ValueError(f"scan shape {scan.shape} is incorrect. " "It should be (N, 2) " "where N >= 1 is the number of scan positions.") if data.shape[0] != scan.shape[0]: raise ValueError( f"data shape {data.shape} and scan shape {scan.shape} " "are incompatible. They should have the same leading dimension.") if (probe.ndim != 5 or probe.shape[:2] != (1, 1) or np.any(np.asarray(probe.shape) < 1) or probe.shape[-2] != probe.shape[-1]): raise ValueError(f"probe shape {probe.shape} is incorrect. " "It should be (1, 1, S, W, H) " "where S >=1 is the number of probes, and " "W, H >= 1 are the square probe grid dimensions.") if np.any(np.asarray(probe.shape[-2:]) > np.asarray(data.shape[-2:])): raise ValueError(f"probe shape {probe.shape} is incorrect." "The probe width/height must be " f"<= the data width/height {data.shape}.") if (psi.ndim != 2 or np.any(np.asarray(psi.shape) <= np.asarray(probe.shape[-2:]))): raise ValueError(f"psi shape {psi.shape} is incorrect. " "It should be (W, H) where W, H > probe.shape[-2:].") check_allowed_positions(scan, psi, probe.shape) logger.info("{} for {:,d} - {:,d} by {:,d} frames for {:,d} " "iterations.".format( algorithm_options.name, *data.shape[-3:], algorithm_options.num_iter, )) if use_mpi is True: mpi = MPIComm if psi is None: raise ValueError( "When MPI is enabled, initial object guess cannot be None; " "automatic psi initialization is not synchronized " "across processes.") else: mpi = None with (cp.cuda.Device(num_gpu[0] if isinstance(num_gpu, tuple) else None)): with Ptycho( probe_shape=probe.shape[-1], detector_shape=data.shape[-1], nz=psi.shape[-2], n=psi.shape[-1], model=model, ) as operator, Comm(num_gpu, mpi) as comm: ( batches, data, result, scan, ) = _setup( algorithm_options, comm, data, eigen_probe, eigen_weights, object_options, operator, probe, psi, position_options, probe_options, scan, ) start = time.perf_counter() for i in range(algorithm_options.num_iter): logger.info(f"{algorithm_options.name} epoch {i:,d}") # TODO: Append new information to everything that emits from # _setup. result = _iterate( algorithm_options, batches, comm, data, operator, position_options, probe_options, result, ) # TODO: Grab intermediate psi/probe from GPU. algorithm_options.times.append(time.perf_counter() - start) start = time.perf_counter() return _teardown( algorithm_options, comm, eigen_probe, eigen_weights, object_options, position_options, probe_options, result, scan, )