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,
)