def plug_and_play(x, b, tp, param):
## Parameters
n = param.npix_op
nz = param.nz
lambd1 = param.lambd1
lambd2 = param.lambd2
mu = param.mu
n2 = param.n_iter_pcg
it = 0
err = 1e20
info = []
## Set stopping criterion
if param.eps is None:
eps = 1e-10
else:
eps = param.eps
if param.n_iter is None:
n1 = 200
else:
n1 = param.n_iter
## Reconstruction indices
if param.edge_padding != 0.0:
i1 = param.index_start
i2 = param.index_end
else:
i1 = 0
i2 = n
u = np.zeros((nz, n, n), dtype=myfloat)
## Initialize plot
if param.plot is True:
fig = plt.figure()
ax = fig.add_subplot(111)
if nz == 1:
Ln = ax.imshow(x[0, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
else:
nzz = np.int(nz * 0.5)
Ln = ax.imshow(x[nzz, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
plt.ion()
## Start main loop
while it < n1 and err > eps:
print('\nADMM-PP --- iteraz. n.: ', it)
# --- step 0
if it == 0:
#u = np.zeros( ( nz , n , n ) , dtype=myfloat )
alpha = u.copy()
x_old = x.copy()
## --- step 1
x_old[:] = x.copy()
if param.num_cores == -1:
pool = mproc.Pool()
else:
pool = mproc.Pool(param.num_cores)
results = [ pool.apply_async( cg.cg_plug_and_play , args=( x[i,:,:] , b[i,:,:] , tp , n2 ,
( u - alpha )[i,:,:] , lambd1 , mu , i , ) ) \
for i in range( nz ) ]
x[:] = np.array([res.get() for res in results])
pool.close()
pool.join()
if param.mask is True:
for i in range(nz):
x[i, param.mask[i] == 0.0] = 0.0
elif param.pc >= 0.0:
if param.pc >= 0.0:
if param.pc == 1.0:
x[:, i1:i2,
i1:i2] = utils.phys_constr(x[:, i1:i2, i1:i2], 0.0)
x[:, i1:i2, i1:i2] = utils.resol_circle_constr(x[:, i1:i2,
i1:i2])
if param.pc == 2.0:
x[:, i1:i2, i1:i2] = utils.resol_circle_constr(x[:, i1:i2,
i1:i2])
else:
x[:, i1:i2,
i1:i2] = utils.phys_constr(x[:, i1:i2, i1:i2], param.pc)
## --- step 2
if nz > 1:
if param.reg == 'pp-breg':
u[:, i1:i2, i1:i2] = den.tv_breg(
x[:, i1:i2, i1:i2] + alpha[:, i1:i2, i1:i2], lambd2 / mu)
elif param.reg == 'pp-chamb':
u[:, i1:i2, i1:i2] = den.tv_chamb(
x[:, i1:i2, i1:i2] + alpha[:, i1:i2, i1:i2], lambd2 / mu)
elif param.reg == 'pp-nlmeans':
u[:, i1:i2, i1:i2] = den.nl_means(
x[:, i1:i2, i1:i2] + alpha[:, i1:i2, i1:i2], lambd2 / mu)
elif param.reg == 'pp-tgv':
x_aux = x + alpha
param1 = lambd2 / mu
param2 = beta = lambd2 / mu
if param.num_cores == -1:
pool = mproc.Pool()
else:
pool = mproc.Pool(param.num_cores)
results = [ pool.apply_async( plug_and_play_tgv_fast , args=( x_aux[i,:,:] , param1 , param2 ) ) \
for i in range( nz ) ]
x[:] = np.array([res.get() for res in results])
pool.close()
pool.join()
elif param.reg == 'pp-nltv':
x_aux = x + alpha
if param.num_cores == -1:
pool = mproc.Pool()
else:
pool = mproc.Pool(param.num_cores)
results = [ pool.apply_async( plug_and_play_nltv , args=( x_aux[i,:,:] , lambd2 , mu ) ) \
for i in range( nz ) ]
x[:] = np.array([res.get() for res in results])
pool.close()
pool.join()
elif nz == 1:
if param.reg == 'pp-breg':
u[0, i1:i2, i1:i2] = den.tv_breg((x + alpha)[0, i1:i2, i1:i2],
lambd2 / mu)
elif param.reg == 'pp-chamb':
u[0, i1:i2, i1:i2] = den.tv_chamb((x + alpha)[0, i1:i2, i1:i2],
lambd2 / mu)
elif param.reg == 'pp-nlmeans':
u[0, i1:i2, i1:i2] = den.nl_means((x + alpha)[0, i1:i2, i1:i2],
h=lambd2 / mu)
elif param.reg == 'pp-tgv':
u[0, i1:i2, i1:i2] = den.tgv_fast((x + alpha)[0, i1:i2, i1:i2],
alpha=lambd2 / mu,
beta=lambd2 / mu,
niter=5)
elif param.reg == 'pp-nltv':
u[0, i1:i2, i1:i2] = den.nltv_sb((x + alpha)[0, i1:i2, i1:i2],
lambd=1.0,
ps=5,
ws=11,
h=lambd2 / mu,
nn=10,
icn=1,
mu=1.0,
niter1=4,
niter2=2)
## --- step 3
r = x - u
diff = x - x_old
alpha += r
## --- step 4
diff = np.linalg.norm(x[0, :, :] - x_old[0, :, :])
err = np.linalg.norm(r)
tmp = np.sqrt(cg.G(x[0, :, :])[0]**2 + cg.G(x[0, :, :])[1]**2)
obj_0 = np.linalg.norm(tmp.reshape(-1), 0)
obj_1 = np.linalg.norm(tp.A(x[0, :, :]) - b[0, :, :])**2
obj_2 = np.linalg.norm(x[0, :, :])
obj_3 = np.linalg.norm(tmp, 1)
obj = 0.5 * obj_1 + 0.5 * lambd1 * obj_2 + lambd2 * obj_3
if param.checkit is True:
if param.projector == 'grid-pswf' or param.projector == 'grid-kb':
x_aux = x[0, i1:i2, i1:i2].copy()
elif param.projector == 'radon':
x_aux = x[0, i2:i1:-1, i2:i1:-1].copy()
elif param.projector == 'bspline':
x_aux = bfun.convert_from_bspline_to_pixel_basis(
x[0, i2:i1:-1, i2:i1:-1], 3)
if it < 10:
niter = '00' + str(it)
elif it < 100:
niter = '0' + str(it)
else:
niter = str(it)
io.writeImage(param.path_rmse + 'reco_iter' + niter + '.DMP',
x_aux)
info.append([err, diff, obj])
print('Multiplier update: %.4e' % err)
print('Relative difference: %.4e' % diff)
print('Fidelity term: %.4e' % obj_1, ' Second term: %.4e' % obj_2,
' TV term: %.4e' % obj_3, ' || G(x) ||_{0}: %.4e' % obj_0)
print('Function score: %.4e' % obj)
## Plot intermediate reconstruction as check
if param.plot is True:
if nz == 1:
ax.imshow(x[0, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
else:
nzz = np.int(nz * 0.5)
ax.imshow(x[nzz, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
plt.draw()
plt.pause(1)
it += 1
if param.plot is True:
py.ioff()
return x, info
def update_lasso_tv(x, u):
return cg.G(x) - u
def lasso_tv(x, b, tp, param):
## Parameters
n = param.npix_op
nz = param.nz
lambd1 = param.lambd1
lambd2 = param.lambd2
mu = param.mu
n2 = param.n_iter_pcg
it = 0
err = 1e20
info = []
## Set stopping criterion
if param.eps is None:
eps = 1e-10
else:
eps = param.eps
if param.n_iter is None:
n1 = 200
else:
n1 = param.n_iter
## Reconstruction indices
if param.edge_padding != 0.0:
i1 = param.index_start
i2 = param.index_end
else:
i1 = 0
i2 = n
## Initialize plot
if param.plot is True:
fig = plt.figure()
ax = fig.add_subplot(111)
if nz == 1:
Ln = ax.imshow(x[0, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
else:
nzz = np.int(nz * 0.5)
Ln = ax.imshow(x[nzz, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
plt.ion()
## Start main loop
while it < n1 and err > eps:
print('\nADMM Lasso TV --- iteraz. n.: ', it)
# --- step 0
if it == 0:
u = np.zeros((nz, 2, n, n), dtype=myfloat)
alpha = u.copy()
x_old = x.copy()
## --- step 1
x_old[:] = x.copy()
if param.num_cores == -1:
pool = mproc.Pool()
else:
pool = mproc.Pool(param.num_cores)
results = [ pool.apply_async( cg.cg_lasso_tv , args=( x[i,:,:] , b[i,:,:] , tp , n2 ,
mu * u[i,:,:] - alpha[i,:,:] ,
lambd1 , mu , i , ) ) \
for i in range( nz ) ]
x[:] = np.array([res.get() for res in results])
pool.close()
pool.join()
if param.mask is not None:
for i in range(nz):
x[i, i1:i2, i1:i2] = utils.supp_constr(x[i, i1:i2, i1:i2],
param.mask)
if param.mask_add is not None:
for i in range(nz):
aux = x[i, i1:i2, i1:i2].copy()
for j in range(param.mask_add_n):
aux[:] = utils.mask_constr(aux, param.mask_add[j])
x[i, i1:i2, i1:i2] = aux
if param.pc >= 0.0:
if param.pc >= 0.0:
if param.pc == 1.0:
x[:, i1:i2,
i1:i2] = utils.phys_constr(x[:, i1:i2, i1:i2], 0.0)
x[:, i1:i2, i1:i2] = utils.resol_circle_constr(x[:, i1:i2,
i1:i2])
if param.pc == 2.0:
x[:, i1:i2, i1:i2] = utils.resol_circle_constr(x[:, i1:i2,
i1:i2])
else:
x[:, i1:i2,
i1:i2] = utils.phys_constr(x[:, i1:i2, i1:i2], param.pc)
## --- step 2
if param.num_cores == -1:
pool = mproc.Pool()
else:
pool = mproc.Pool(param.num_cores)
results = [ pool.apply_async( hard_thres_multi , args=( x[i,:,:] , alpha[i,:,:,:] , lambd2 , mu ) ) \
for i in range( nz ) ]
u[:] = np.array([res.get() for res in results])
pool.close()
pool.join()
## --- step 3
r = u.copy()
if param.num_cores == -1:
pool = mproc.Pool()
else:
pool = mproc.Pool(param.num_cores)
results = [ pool.apply_async( update_lasso_tv , args=( x[i,:,:] , u[i,:,:,:] ) ) \
for i in range( nz ) ]
r[:] = np.array([res.get() for res in results])
pool.close()
pool.join()
diff = x - x_old
alpha += r
## --- step 4
diff = np.linalg.norm(x[0, :, :] - x_old[0, :, :])
err = np.linalg.norm(r)
tmp = np.sqrt(cg.G(x[0, :, :])[0]**2 + cg.G(x[0, :, :])[1]**2)
obj_0 = np.linalg.norm(tmp.reshape(-1), 0)
obj_1 = np.linalg.norm(tp.A(x[0, :, :]) - b[0, :, :])**2
obj_2 = np.linalg.norm(x[0, :, :])
obj_3 = np.linalg.norm(tmp, 1)
obj = 0.5 * obj_1 + 0.5 * lambd1 * obj_2 + lambd2 * obj_3
if param.checkit is True:
if param.projector == 'grid-pswf' or param.projector == 'grid-kb':
x_aux = x[0, i1:i2, i1:i2].copy()
elif param.projector == 'radon':
x_aux = x[0, i2:i1:-1, i2:i1:-1].copy()
elif param.projector == 'bspline':
x_aux = bfun.convert_from_bspline_to_pixel_basis(
x[0, i2:i1:-1, i2:i1:-1], 3)
if it < 10:
niter = '00' + str(it)
elif it < 100:
niter = '0' + str(it)
else:
niter = str(it)
io.writeImage(param.path_rmse + 'reco_iter' + niter + '.DMP',
x_aux)
info.append([err, diff, obj])
print('Multiplier update: %.4e' % err)
print('Relative difference: %.4e' % diff)
print('Fidelity term: %.4e' % obj_1, ' Second term: %.4e' % obj_2,
' TV term: %.4e' % obj_3, ' || G(x) ||_{0}: %.4e' % obj_0)
print('Function score: %.4e' % obj)
## Plot intermediate reconstruction as check
if param.plot is True:
if nz == 1:
ax.imshow(x[0, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
else:
nzz = np.int(nz * 0.5)
ax.imshow(x[nzz, i2:i1:-1, i1:i2],
animated=True,
cmap=cm.Greys_r)
plt.draw()
plt.pause(1)
it += 1
if param.plot is True:
py.ioff()
return x, info
def hard_thres_multi(x, alpha, lambd2, mu):
return utils.hard_thres(cg.G(x) + alpha / mu, lambd2 / mu)