class TestLaminoFourier(unittest.TestCase, OperatorTests):
"""Test the Laminography operator."""
def setUp(self, n=16, ntheta=8, tilt=np.pi / 3, eps=1e-6):
self.operator = Lamino(
n=n,
tilt=tilt,
eps=eps,
)
self.operator.__enter__()
self.xp = self.operator.xp
np.random.seed(0)
self.m = self.xp.asarray(random_complex(n, n, n), dtype='complex64')
self.m_name = 'u'
self.d = self.xp.asarray(random_complex(ntheta, n, n),
dtype='complex64')
self.d_name = 'data'
self.kwargs = {
'theta': self.xp.linspace(0, 2 * np.pi, ntheta).astype('float32')
}
print(self.operator)
@unittest.skip('FIXME: This operator is not scaled.')
def test_scaled(self):
pass
def test_adjoint(self):
"""Check that the adjoint operator is correct."""
np.random.seed(0)
obj = random_complex(self.n, self.n, self.n)
data = random_complex(self.ntheta, self.n, self.n)
with Lamino(n=self.n, theta=self.theta, tilt=self.tilt,
eps=self.eps) as op:
obj = op.asarray(obj.astype('complex64'))
data = op.asarray(data.astype('complex64'))
d = op.fwd(obj)
assert d.shape == data.shape
o = op.adj(data)
assert obj.shape == o.shape
a = inner_complex(d, data)
b = inner_complex(obj, o)
print()
print('<Lobj, data> = {:.6f}{:+.6f}j'.format(
a.real.item(), a.imag.item()))
print('<obj , L*data> = {:.6f}{:+.6f}j'.format(
b.real.item(), b.imag.item()))
# Test whether Adjoint fixed probe operator is correct
op.xp.testing.assert_allclose(a.real, b.real, rtol=1e-2)
op.xp.testing.assert_allclose(a.imag, b.imag, rtol=1e-2)
def setUp(self, n=16, ntheta=8, tilt=np.pi / 3, eps=1e-6):
self.operator = Lamino(
n=n,
theta=np.linspace(0, 2 * np.pi, ntheta),
tilt=tilt,
eps=eps,
)
self.operator.__enter__()
self.xp = self.operator.xp
np.random.seed(0)
self.m = self.xp.asarray(random_complex(n, n, n), dtype='complex64')
self.m_name = 'u'
self.d = self.xp.asarray(random_complex(ntheta, n, n),
dtype='complex64')
self.d_name = 'data'
self.kwargs = {}
print(self.operator)
def simulate(
obj,
theta,
tilt,
**kwargs
): # yapf: disable
"""Return complex values of simulated laminography data."""
assert obj.ndim == 3
assert theta.ndim == 1
with Lamino(
n=obj.shape[-1],
theta=theta,
tilt=tilt,
**kwargs,
) as operator:
data = operator.fwd(u=operator.asarray(obj, dtype='complex64'))
assert data.dtype == 'complex64', data.dtype
return operator.asnumpy(data)
def simulate(
obj,
theta,
tilt,
**kwargs
): # yapf: disable
"""Return complex values of simulated laminography data."""
assert obj.ndim == 3
assert theta.ndim == 1
with Lamino(
n=obj.shape[-1],
tilt=tilt,
**kwargs,
) as operator:
grid = operator._make_grid().reshape(obj.shape[-1]**3, 3)
data = operator.fwd(
u=operator.asarray(obj, dtype='complex64'),
theta=operator.asarray(theta, dtype='float32'),
grid=operator.asarray(grid, dtype='int16'),
)
assert data.dtype == 'complex64', data.dtype
return operator.asnumpy(data)
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,
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,
theta,
tilt,
algorithm,
obj=None, num_iter=1, rtol=-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.
"""
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],
theta=theta,
tilt=tilt,
eps=1e-3,
**kwargs,
) as operator:
# send any array-likes to device
data = operator.asarray(data, dtype='complex64')
result = {
'obj': operator.asarray(obj, dtype='complex64'),
}
for key, value in kwargs.items():
if np.ndim(value) > 0:
kwargs[key] = operator.asarray(value)
logger.info("{} on {:,d} by {:,d} by {:,d} volume for {:,d} "
"iterations.".format(algorithm, *obj.shape, num_iter))
cost = 0
for i in range(num_iter):
kwargs.update(result)
result = getattr(solvers, algorithm)(
operator,
data=data,
**kwargs,
)
# Check for early termination
if i > 0 and abs((result['cost'] - cost) / cost) < rtol:
logger.info(
"Cost function rtol < %g reached at %d "
"iterations.", rtol, i)
break
cost = result['cost']
return {k: operator.asnumpy(v) for k, v in result.items()}
else:
raise ValueError(
"The '{}' algorithm is not an available.".format(algorithm))