def vstack(self):
"""Test vstack operator.
"""
# diagonals.
x = Variable(1)
y = Variable(1)
fn = vstack([x, y])
assert fn.is_gram_diag(freq=False)
assert fn.is_gram_diag(freq=True)
fn = vstack([fn])
assert fn.is_gram_diag(freq=False)
assert fn.is_gram_diag(freq=True)
def vstack(self):
"""Test vstack operator.
"""
# diagonals.
x = Variable(1)
y = Variable(1)
fn = vstack([x, y])
assert fn.is_gram_diag(freq=False)
assert fn.is_gram_diag(freq=True)
fn = vstack([fn])
assert fn.is_gram_diag(freq=False)
assert fn.is_gram_diag(freq=True)
def test_combo(self):
"""Test subsampling followed by convolution.
"""
# Forward.
var = Variable((2, 3))
kernel = np.array([[1, 2, 3]]) # 2x3
fn = vstack([conv(kernel, subsample(var, (2, 1)))])
fn = CompGraph(fn)
x = np.arange(6) * 1.0
x = np.reshape(x, (2, 3))
out = np.zeros(fn.output_size)
fn.forward(x.flatten(), out)
y = np.zeros((1, 3))
xsub = x[::2, ::1]
y = ndimage.convolve(xsub, kernel, mode='wrap')
self.assertItemsAlmostEqual(np.reshape(out, y.shape), y)
# Adjoint.
x = np.arange(3) * 1.0
x = np.reshape(x, (1, 3))
out = np.zeros(var.size)
fn.adjoint(x.flatten(), out)
y = ndimage.correlate(x, kernel, mode='wrap')
y2 = np.zeros((2, 3))
y2[::2, :] = y
self.assertItemsAlmostEqual(np.reshape(out, y2.shape), y2)
out = np.zeros(var.size)
fn.adjoint(x.flatten(), out)
self.assertItemsAlmostEqual(np.reshape(out, y2.shape), y2)
def partition(prox_fns, try_diagonalize=True):
"""Divide the proxable functions into sets Psi and Omega.
"""
# Omega must be a single function.
# Merge quadratic functions into the x update.
# Automatically try to split the problem.
quad_fns = max_diag_set(prox_fns)
split_fn = []
omega_fns = []
if len(quad_fns) == 0:
for fn in prox_fns:
if type(fn.lin_op) == Variable:
split_fn = [fn]
break
omega_fns = split_fn
else:
# Proximal solve for:
# G(x) + 1/(2*tau) * ||x - v||^2_2, with G containing all quadratics
quad_ops = []
const_terms = []
for fn in quad_fns:
fn = fn.absorb_params()
quad_ops.append( fn.beta*fn.lin_op )
const_terms.append( fn.b.flatten() )
stacked_ops = vstack(quad_ops)
b = np.hstack(const_terms)
# Get optimize inverse (tries spatial and frequency diagonalization)
x_update = get_least_squares_inverse(quad_ops, b, try_diagonalize)
omega_fns = [x_update]
psi_fns = [fn for fn in prox_fns if fn not in split_fn + quad_fns]
return psi_fns, omega_fns
def test_diagonalization(self):
"""Test automatic diagonalization.
"""
var = Variable((2, 5))
K = np.array([[-1, 1]])
expr = 2*vstack([conv(K, var), conv(K, var)])
assert expr.is_gram_diag(freq=True)
def get_least_squares_inverse(op_list, b, try_freq_diagonalize=True, verbose=False):
if len(op_list) == 0:
return None
# Are all the operators diagonal?
stacked = vstack(op_list)
if stacked.is_gram_diag(freq=False):
if verbose:
print 'Optimized for diagonal inverse'
diag = stacked.get_diag(freq=False).values()[0]
diag = diag*np.conj(diag)
x_update = least_squares(stacked, b, diag=diag)
# Are all the operators diagonal in the frequency domain?
elif try_freq_diagonalize and stacked.is_gram_diag(freq=True):
diag = stacked.get_diag(freq=True).values()[0]
diag = diag*np.conj(diag)
dims = get_dims(op_list)
implem = get_implem(op_list) #If any freqdiag is halide, solve with halide
if verbose:
dimstr = (' with dimensionality %d' % dims) if dims != None else ''
print 'Optimized for diagonal frequency inverse' + dimstr
x_update = least_squares(stacked, b,
freq_diag=diag, freq_dims=dims, implem = implem)
else:
x_update = least_squares(stacked, b)
return x_update
def test_combo(self):
"""Test subsampling followed by convolution.
"""
# Forward.
var = Variable((2, 3))
kernel = np.array([[1, 2, 3]]) # 2x3
fn = vstack([conv(kernel, subsample(var, (2, 1)))])
fn = CompGraph(fn)
x = np.arange(6) * 1.0
x = np.reshape(x, (2, 3))
out = np.zeros(fn.output_size)
fn.forward(x.flatten(), out)
y = np.zeros((1, 3))
xsub = x[::2, ::1]
y = ndimage.convolve(xsub, kernel, mode='wrap')
self.assertItemsAlmostEqual(np.reshape(out, y.shape), y)
# Adjoint.
x = np.arange(3) * 1.0
x = np.reshape(x, (1, 3))
out = np.zeros(var.size)
fn.adjoint(x.flatten(), out)
y = ndimage.correlate(x, kernel, mode='wrap')
y2 = np.zeros((2, 3))
y2[::2, :] = y
self.assertItemsAlmostEqual(np.reshape(out, y2.shape), y2)
out = np.zeros(var.size)
fn.adjoint(x.flatten(), out)
self.assertItemsAlmostEqual(np.reshape(out, y2.shape), y2)
def test_diagonalization(self):
"""Test automatic diagonalization.
"""
var = Variable((2, 5))
K = np.array([[-1, 1]])
expr = 2 * vstack([conv(K, var), conv(K, var)])
assert expr.is_gram_diag(freq=True)
def solve(self, solver=None, *args, **kwargs):
if solver is None:
solver = self.solver
if len(self.omega_fns + self.psi_fns) == 0:
prox_fns = self.prox_fns
else:
prox_fns = self.omega_fns + self.psi_fns
# Absorb lin ops if desired.
if self.absorb:
prox_fns = absorb.absorb_all_lin_ops(prox_fns)
# Merge prox fns.
if self.merge:
prox_fns = merge.merge_all(prox_fns)
# Absorb offsets.
prox_fns = [absorb.absorb_offset(fn) for fn in prox_fns]
# TODO more analysis of what solver to use.
# Short circuit with one function.
if len(prox_fns) == 1 and type(prox_fns[0].lin_op) == Variable:
fn = prox_fns[0]
var = fn.lin_op
var.value = fn.prox(0, np.zeros(fn.lin_op.shape))
return fn.value
elif solver in NAME_TO_SOLVER:
module = NAME_TO_SOLVER[solver]
if len(self.omega_fns + self.psi_fns) == 0:
if self.try_split:
psi_fns, omega_fns = module.partition(prox_fns,
self.try_diagonalize)
else:
psi_fns = prox_fns
omega_fns = []
# Scale the problem.
if self.scale:
K = CompGraph(vstack([fn.lin_op for fn in psi_fns]),
implem=self.implem)
Knorm = est_CompGraph_norm(K, try_fast_norm=self.try_fast_norm)
for idx, fn in enumerate(psi_fns):
psi_fns[idx] = fn.copy(fn.lin_op / Knorm,
beta=fn.beta * np.sqrt(Knorm),
implem=self.implem)
for idx, fn in enumerate(omega_fns):
omega_fns[idx] = fn.copy(beta=fn.beta / np.sqrt(Knorm),
implem=self.implem)
opt_val = module.solve(psi_fns, omega_fns,
lin_solver=self.lin_solver,
try_diagonalize=self.try_diagonalize,
try_fast_norm=self.try_fast_norm,
scaled=self.scale,
*args, **kwargs)
# Unscale the variables.
if self.scale:
for var in self.variables():
var.value /= np.sqrt(Knorm)
return opt_val
else:
raise Exception("Unknown solver.")
def solve(self, solver=None, *args, **kwargs):
if solver is None:
solver = self.solver
if len(self.omega_fns + self.psi_fns) == 0:
prox_fns = self.prox_fns
else:
prox_fns = self.omega_fns + self.psi_fns
# Absorb lin ops if desired.
if self.absorb:
prox_fns = absorb.absorb_all_lin_ops(prox_fns)
# Merge prox fns.
if self.merge:
prox_fns = merge.merge_all(prox_fns)
# Absorb offsets.
prox_fns = [absorb.absorb_offset(fn) for fn in prox_fns]
# TODO more analysis of what solver to use.
# Short circuit with one function.
if len(prox_fns) == 1 and type(prox_fns[0].lin_op) == Variable:
fn = prox_fns[0]
var = fn.lin_op
var.value = fn.prox(0, np.zeros(fn.lin_op.shape))
return fn.value
elif solver in NAME_TO_SOLVER:
module = NAME_TO_SOLVER[solver]
if len(self.omega_fns + self.psi_fns) == 0:
if self.try_split and len(prox_fns) > 1 and len(self.variables()) == 1:
psi_fns, omega_fns = module.partition(prox_fns,
self.try_diagonalize)
else:
psi_fns = prox_fns
omega_fns = []
# Scale the problem.
if self.scale:
K = CompGraph(vstack([fn.lin_op for fn in psi_fns]),
implem=self.implem)
Knorm = est_CompGraph_norm(K, try_fast_norm=self.try_fast_norm)
for idx, fn in enumerate(psi_fns):
psi_fns[idx] = fn.copy(fn.lin_op / Knorm,
beta=fn.beta * np.sqrt(Knorm),
implem=self.implem)
for idx, fn in enumerate(omega_fns):
omega_fns[idx] = fn.copy(beta=fn.beta / np.sqrt(Knorm),
implem=self.implem)
opt_val = module.solve(psi_fns, omega_fns,
lin_solver=self.lin_solver,
try_diagonalize=self.try_diagonalize,
try_fast_norm=self.try_fast_norm,
scaled=self.scale,
*args, **kwargs)
# Unscale the variables.
if self.scale:
for var in self.variables():
var.value /= np.sqrt(Knorm)
return opt_val
else:
raise Exception("Unknown solver.")
def get_least_squares_inverse(op_list,
b,
try_freq_diagonalize=True,
verbose=False):
if len(op_list) == 0:
return None
# Are all the operators diagonal?
stacked = vstack(op_list)
if stacked.is_gram_diag(freq=False):
if verbose:
print('Optimized for diagonal inverse')
diag = list(stacked.get_diag(freq=False).values())[0]
diag = diag * np.conj(diag)
x_update = least_squares(stacked, b, diag=diag)
# Are all the operators diagonal in the frequency domain?
elif try_freq_diagonalize and stacked.is_gram_diag(freq=True):
diag = list(stacked.get_diag(freq=True).values())[0]
diag = diag * np.conj(diag)
dims = get_dims(op_list)
implem = get_implem(
op_list) # If any freqdiag is halide, solve with halide
if verbose:
dimstr = (' with dimensionality %d' %
dims) if dims is not None else ''
print('Optimized for diagonal frequency inverse' + dimstr)
x_update = least_squares(stacked,
b,
freq_diag=diag,
freq_dims=dims,
implem=implem)
else:
x_update = least_squares(stacked, b)
return x_update
def solve(psi_fns,
omega_fns,
tau=None,
sigma=None,
theta=None,
max_iters=1000,
eps_abs=1e-3,
eps_rel=1e-3,
x0=None,
lin_solver="cg",
lin_solver_options=None,
conv_check=100,
try_diagonalize=True,
try_fast_norm=False,
scaled=True,
metric=None,
convlog=None,
verbose=0):
# Can only have one omega function.
assert len(omega_fns) <= 1
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
v = np.zeros(K.input_size)
# Select optimal parameters if wanted
if tau is None or sigma is None or theta is None:
tau, sigma, theta = est_params_pc(K, tau, sigma, verbose, scaled,
try_fast_norm)
# Initialize
x = np.zeros(K.input_size)
y = np.zeros(K.output_size)
xbar = np.zeros(K.input_size)
u = np.zeros(K.output_size)
z = np.zeros(K.output_size)
if x0 is not None:
x[:] = np.reshape(x0, K.input_size)
K.forward(x, y)
xbar[:] = x
# Buffers.
Kxbar = np.zeros(K.output_size)
Kx = np.zeros(K.output_size)
KTy = np.zeros(K.input_size)
KTu = np.zeros(K.input_size)
s = np.zeros(K.input_size)
prev_x = x.copy()
prev_Kx = Kx.copy()
prev_z = z.copy()
prev_u = u.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("PC iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(x)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing.tic()
if convlog is not None:
convlog.tic()
# Keep track of previous iterates
np.copyto(prev_x, x)
np.copyto(prev_z, z)
np.copyto(prev_u, u)
np.copyto(prev_Kx, Kx)
# Compute z
K.forward(xbar, Kxbar)
z = y + sigma * Kxbar
# Update y.
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
z_slc = np.reshape(z[slc], fn.lin_op.shape)
# Moreau identity: apply and time prox.
prox_log[fn].tic()
y[slc] = (z_slc -
sigma * fn.prox(sigma, z_slc / sigma, i)).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
y[offset:] = 0
# Update x
K.adjoint(y, KTy)
x -= tau * KTy
if len(omega_fns) > 0:
xtmp = np.reshape(x, omega_fns[0].lin_op.shape)
x[:] = omega_fns[0].prox(1.0 / tau,
xtmp,
x_init=prev_x,
lin_solver=lin_solver,
options=lin_solver_options).flatten()
# Update xbar
np.copyto(xbar, x)
xbar += theta * (x - prev_x)
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(x)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
""" Old convergence check
#Very basic convergence check.
r_x = np.linalg.norm(x - prev_x)
r_xbar = np.linalg.norm(xbar - prev_xbar)
r_ybar = np.linalg.norm(y - prev_y)
error = r_x + r_xbar + r_ybar
"""
# Residual based convergence check
K.forward(x, Kx)
u = 1.0 / sigma * y + theta * (Kx - prev_Kx)
z = prev_u + prev_Kx - 1.0 / sigma * y
# Iteration order is different than
# lin-admm (--> start checking at iteration 1)
if i > 0 and i % conv_check == 0:
# Check convergence
r = prev_Kx - z
K.adjoint(sigma * (z - prev_z), s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(prev_Kx), np.linalg.norm(z)])
K.adjoint(u, KTu)
eps_dual = np.sqrt(
K.input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) / sigma
# Progress
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(x)
objstr = ", obj_val = %02.03e" % sum(
[fn.value for fn in prox_fns])
""" Old convergence check
#Evaluate metric potentially
metstr = '' if metric is None else ", {}".format( metric.message(x.copy()) )
print "iter [%04d]:" \
"||x - x_prev||_2 = %02.02e " \
"||xbar - xbar_prev||_2 = %02.02e " \
"||y - y_prev||_2 = %02.02e " \
"SUM = %02.02e (eps=%02.03e)%s%s" \
% (i, r_x, r_xbar, r_ybar, error, eps, objstr, metstr)
"""
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(
metric.message(v))
print(
"iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s"
% (i, np.linalg.norm(r), eps_pri, np.linalg.norm(s),
eps_dual, objstr, metstr))
iter_timing.toc()
if np.linalg.norm(r) <= eps_pri and np.linalg.norm(s) <= eps_dual:
break
else:
iter_timing.toc()
""" Old convergence check
if error <= eps:
break
"""
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(x)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(self,
solver=None,
test_adjoints=False,
test_norm=False,
show_graph=False,
*args,
**kwargs):
if solver is None:
solver = self.solver
if len(self.omega_fns + self.psi_fns) == 0:
prox_fns = self.prox_fns
else:
prox_fns = self.omega_fns + self.psi_fns
# Absorb lin ops if desired.
if self.absorb:
prox_fns = absorb.absorb_all_lin_ops(prox_fns)
# Merge prox fns.
if self.merge:
prox_fns = merge.merge_all(prox_fns)
# Absorb offsets.
prox_fns = [absorb.absorb_offset(fn) for fn in prox_fns]
# TODO more analysis of what solver to use.
if show_graph:
print("Computational graph before optimizing:")
graph_visualize(
prox_fns,
filename=show_graph if type(show_graph) is str else None)
# Short circuit with one function.
if len(prox_fns) == 1 and type(prox_fns[0].lin_op) == Variable:
fn = prox_fns[0]
var = fn.lin_op
var.value = fn.prox(0, np.zeros(fn.lin_op.shape))
return fn.value
elif solver in NAME_TO_SOLVER:
module = NAME_TO_SOLVER[solver]
if len(self.omega_fns + self.psi_fns) == 0:
if self.try_split and len(prox_fns) > 1 and len(
self.variables()) == 1:
psi_fns, omega_fns = module.partition(
prox_fns, self.try_diagonalize)
else:
psi_fns = prox_fns
omega_fns = []
else:
psi_fns = self.psi_fns
omega_fns = self.omega_fns
if test_norm:
L = CompGraph(vstack([fn.lin_op for fn in psi_fns]))
from numpy.random import random
output_mags = [NotImplemented]
L.norm_bound(output_mags)
if not NotImplemented in output_mags:
assert len(output_mags) == 1
x = random(L.input_size)
x = x / LA.norm(x)
y = np.zeros(L.output_size)
y = L.forward(x, y)
ny = LA.norm(y)
nL2 = est_CompGraph_norm(L, try_fast_norm=False)
if ny > output_mags[0]:
raise RuntimeError("wrong implementation of norm!")
print("%.3f <= ||K|| = %.3f (%.3f)" %
(ny, output_mags[0], nL2))
# Scale the problem.
if self.scale:
K = CompGraph(vstack([fn.lin_op for fn in psi_fns]),
implem=self.implem)
Knorm = est_CompGraph_norm(K, try_fast_norm=self.try_fast_norm)
for idx, fn in enumerate(psi_fns):
psi_fns[idx] = fn.copy(fn.lin_op / Knorm,
beta=fn.beta * np.sqrt(Knorm),
implem=self.implem)
for idx, fn in enumerate(omega_fns):
omega_fns[idx] = fn.copy(beta=fn.beta / np.sqrt(Knorm),
implem=self.implem)
for v in K.orig_end.variables():
if v.initval is not None:
v.initval *= np.sqrt(Knorm)
if not test_adjoints in [False, None]:
if test_adjoints is True:
test_adjoints = 1e-6
# test adjoints
L = CompGraph(vstack([fn.lin_op for fn in psi_fns]))
from numpy.random import random
x = random(L.input_size)
yt = np.zeros(L.output_size)
#print("x=", x)
yt = L.forward(x, yt)
#print("yt=", yt)
#print("x=", x)
y = random(L.output_size)
#print("y=", y)
xt = np.zeros(L.input_size)
xt = L.adjoint(y, xt)
#print("xt=", xt)
#print("y=", y)
r = np.abs(
np.dot(np.ravel(y), np.ravel(yt)) -
np.dot(np.ravel(x), np.ravel(xt)))
#print( x.shape, y.shape, xt.shape, yt.shape)
if r > test_adjoints:
#print("yt=", yt)
#print("y =", y)
#print("xt=", xt)
#print("x =", x)
raise RuntimeError("Unmatched adjoints: " + str(r))
else:
print("Adjoint test passed.", r)
if self.implem == Impl['pycuda']:
kwargs['adapter'] = PyCudaAdapter()
opt_val = module.solve(psi_fns,
omega_fns,
lin_solver=self.lin_solver,
try_diagonalize=self.try_diagonalize,
try_fast_norm=self.try_fast_norm,
scaled=self.scale,
*args,
**kwargs)
# Unscale the variables.
if self.scale:
for var in self.variables():
var.value /= np.sqrt(Knorm)
return opt_val
else:
raise Exception("Unknown solver.")
def solve(psi_fns, omega_fns, rho=1.0,
max_iters=1000, eps_abs=1e-3, eps_rel=1e-3, x0=None,
lin_solver="cg", lin_solver_options=None,
try_diagonalize=True, try_fast_norm=False,
scaled=True, conv_check=100,
metric=None, convlog=None, verbose=0):
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Rescale so (rho/2)||x - b||^2_2
rescaling = np.sqrt(2. / rho)
quad_ops = []
const_terms = []
for fn in omega_fns:
fn = fn.absorb_params()
quad_ops.append(scale(rescaling * fn.beta, fn.lin_op))
const_terms.append(fn.b.flatten() * rescaling)
# Check for fast inverse.
op_list = [func.lin_op for func in psi_fns] + quad_ops
stacked_ops = vstack(op_list)
# Get optimize inverse (tries spatial and frequency diagonalization)
v_update = get_least_squares_inverse(op_list, None, try_diagonalize, verbose)
# Initialize everything to zero.
v = np.zeros(K.input_size)
z = np.zeros(K.output_size)
u = np.zeros(K.output_size)
# Initialize
if x0 is not None:
v[:] = np.reshape(x0, K.input_size)
K.forward(v, z)
# Buffers.
Kv = np.zeros(K.output_size)
KTu = np.zeros(K.input_size)
s = np.zeros(K.input_size)
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("ADMM iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(v)
objval = sum([func.value for func in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing.tic()
if convlog is not None:
convlog.tic()
z_prev = z.copy()
# Update v.
tmp = np.hstack([z - u] + const_terms)
v = v_update.solve(tmp, x_init=v, lin_solver=lin_solver, options=lin_solver_options)
# Update z.
K.forward(v, Kv)
Kv_u = Kv + u
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
prox_log[fn].tic()
z[slc] = fn.prox(rho, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update u.
u += Kv - z
# Check convergence.
if i % conv_check == 0:
r = Kv - z
K.adjoint(u, KTu)
K.adjoint(rho * (z - z_prev), s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(Kv), np.linalg.norm(z)])
eps_dual = np.sqrt(K.input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) * rho
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0 and i % conv_check == 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(v)
objstr = ", obj_val = %02.03e" % sum([fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(v))
print("iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s" % (
i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual, objstr, metstr))
iter_timing.toc()
# Exit if converged.
if np.linalg.norm(r) <= eps_pri and np.linalg.norm(s) <= eps_dual:
break
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(v)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(psi_fns, omega_fns, tau=None, sigma=None, theta=None,
max_iters=1000, eps_abs=1e-3, eps_rel=1e-3, x0=None,
lin_solver="cg", lin_solver_options=None, conv_check=100,
try_diagonalize=True, try_fast_norm=False, scaled=True,
metric=None, convlog=None, verbose=0, callback=None, adapter = NumpyAdapter()):
# Can only have one omega function.
assert len(omega_fns) <= 1
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
#graph_visualize(prox_fns)
if adapter.implem() == 'numpy':
K_forward = K.forward
K_adjoint = K.adjoint
prox_off_and_fac = lambda offset, factor, fn, *args, **kw: offset + factor*fn.prox(*args, **kw)
prox = lambda fn, *args, **kw: fn.prox(*args, **kw)
elif adapter.implem() == 'pycuda':
K_forward = K.forward_cuda
K_adjoint = K.adjoint_cuda
prox_off_and_fac = lambda offset, factor, fn, *args, **kw: fn.prox_cuda(*args, offset=offset, factor=factor, **kw)
prox = lambda fn, *args, **kw: fn.prox_cuda(*args, **kw)
else:
raise RuntimeError("Implementation %s unknown" % adapter.implem())
# Select optimal parameters if wanted
if tau is None or sigma is None or theta is None:
tau, sigma, theta = est_params_pc(K, tau, sigma, verbose, scaled, try_fast_norm)
elif callable(tau) or callable(sigma) or callable(theta):
if scaled:
L = 1
else:
L = est_CompGraph_norm(K, try_fast_norm)
# Initialize
x = adapter.zeros(K.input_size)
y = adapter.zeros(K.output_size)
xbar = adapter.zeros(K.input_size)
u = adapter.zeros(K.output_size)
z = adapter.zeros(K.output_size)
if x0 is not None:
x[:] = adapter.reshape(adapter.from_np(x0), K.input_size)
else:
x[:] = adapter.from_np(K.x0())
K_forward(x, y)
xbar[:] = x
# Buffers.
Kxbar = adapter.zeros(K.output_size)
Kx = adapter.zeros(K.output_size)
KTy = adapter.zeros(K.input_size)
KTu = adapter.zeros(K.input_size)
s = adapter.zeros(K.input_size)
prev_x = x.copy()
prev_Kx = Kx.copy()
prev_z = z.copy()
prev_u = u.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
prox_log_tot = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsLog(["pc_iteration_tot",
"copyprev",
"calcz",
"calcx",
"omega_fn",
"xbar",
"conv_check"])
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(adapter.to_np(x))
objval = 0.0
for f in prox_fns:
evp = f.value
#print(str(f), '->', f.value)
objval += evp
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing["pc_iteration_tot"].tic()
if convlog is not None:
convlog.tic()
if callable(sigma):
csigma = sigma(i, L)
else:
csigma = sigma
if callable(tau):
ctau = tau(i, L)
else:
ctau = tau
if callable(theta):
ctheta = theta(i, L)
else:
ctheta = theta
csigma = adapter.scalar(csigma)
ctau = adapter.scalar(ctau)
ctheta = adapter.scalar(ctheta)
# Keep track of previous iterates
iter_timing["copyprev"].tic()
adapter.copyto(prev_x, x)
adapter.copyto(prev_z, z)
adapter.copyto(prev_u, u)
adapter.copyto(prev_Kx, Kx)
iter_timing["copyprev"].toc()
# Compute z
iter_timing["calcz"].tic()
K_forward(xbar, Kxbar)
z = y + csigma * Kxbar
iter_timing["calcz"].toc()
# Update y.
offset = 0
for fn in psi_fns:
prox_log_tot[fn].tic()
slc = slice(offset, offset + fn.lin_op.size, None)
z_slc = adapter.reshape(z[slc], fn.lin_op.shape)
# Moreau identity: apply and time prox.
prox_log[fn].tic()
y[slc] = adapter.flatten( prox_off_and_fac(z_slc, -csigma, fn, csigma, z_slc / csigma, i) )
prox_log[fn].toc()
offset += fn.lin_op.size
prox_log_tot[fn].toc()
iter_timing["calcx"].tic()
if offset < y.shape[0]:
y[offset:] = 0
# Update x
K_adjoint(y, KTy)
x -= ctau * KTy
iter_timing["calcx"].toc()
iter_timing["omega_fn"].tic()
if len(omega_fns) > 0:
fn = omega_fns[0]
prox_log_tot[fn].tic()
xtmp = adapter.reshape(x, fn.lin_op.shape)
prox_log[fn].tic()
x[:] = adapter.flatten( prox(fn, adapter.scalar(1.0) / ctau, xtmp, x_init=prev_x,
lin_solver=lin_solver, options=lin_solver_options) )
prox_log[fn].toc()
prox_log_tot[fn].toc()
iter_timing["omega_fn"].toc()
iter_timing["xbar"].tic()
# Update xbar
adapter.copyto(xbar, x)
xbar += ctheta * (x - prev_x)
iter_timing["xbar"].toc()
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(adapter.to_np(x))
objval = list([fn.value for fn in prox_fns])
objval = sum(objval)
convlog.record_objective(objval)
""" Old convergence check
#Very basic convergence check.
r_x = np.linalg.norm(x - prev_x)
r_xbar = np.linalg.norm(xbar - prev_xbar)
r_ybar = np.linalg.norm(y - prev_y)
error = r_x + r_xbar + r_ybar
"""
# Residual based convergence check
if i % conv_check in [0, conv_check-1]:
iter_timing["conv_check"].tic()
K_forward(x, Kx)
u = adapter.scalar(1.0) / csigma * y + ctheta * (Kx - prev_Kx)
z = prev_u + prev_Kx - adapter.scalar(1.0) / csigma * y
iter_timing["conv_check"].toc()
# Iteration order is different than
# lin-admm (--> start checking at iteration 1)
if i > 0 and i % conv_check == 0:
# Check convergence
r = prev_Kx - z
K_adjoint(csigma * (z - prev_z), s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(adapter.to_np(prev_Kx)), np.linalg.norm(adapter.to_np(z))])
K_adjoint(u, KTu)
eps_dual = np.sqrt(K.input_size) * eps_abs + eps_rel * np.linalg.norm(adapter.to_np(KTu)) / csigma
if not callback is None or verbose == 2:
K.update_vars(adapter.to_np(x))
if not callback is None:
callback(adapter.to_np(x))
# Progress
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
ov = list([fn.value for fn in prox_fns])
objval = sum(ov)
objstr = ", obj_val = %02.03e [%s] " % (objval, ", ".join("%02.03e" % x for x in ov))
""" Old convergence check
#Evaluate metric potentially
metstr = '' if metric is None else ", {}".format( metric.message(x.copy()) )
print "iter [%04d]:" \
"||x - x_prev||_2 = %02.02e " \
"||xbar - xbar_prev||_2 = %02.02e " \
"||y - y_prev||_2 = %02.02e " \
"SUM = %02.02e (eps=%02.03e)%s%s" \
% (i, r_x, r_xbar, r_ybar, error, eps, objstr, metstr)
"""
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(v))
print(
"iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s"
% (i, np.linalg.norm(adapter.to_np(r)), eps_pri, np.linalg.norm(adapter.to_np(s)), eps_dual, objstr, metstr)
)
iter_timing["pc_iteration_tot"].toc()
if np.linalg.norm(adapter.to_np(r)) <= eps_pri and np.linalg.norm(adapter.to_np(s)) <= eps_dual:
break
else:
iter_timing["pc_iteration_tot"].toc()
""" Old convergence check
if error <= eps:
break
"""
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs total:")
print(prox_log_tot)
print("prox funcs inner:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(adapter.to_np(x))
if not callback is None:
callback(adapter.to_np(x))
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(self, solver=None, test_adjoints = False, test_norm = False, show_graph = False, *args, **kwargs):
if solver is None:
solver = self.solver
if len(self.omega_fns + self.psi_fns) == 0:
prox_fns = self.prox_fns
else:
prox_fns = self.omega_fns + self.psi_fns
# Absorb lin ops if desired.
if self.absorb:
prox_fns = absorb.absorb_all_lin_ops(prox_fns)
# Merge prox fns.
if self.merge:
prox_fns = merge.merge_all(prox_fns)
# Absorb offsets.
prox_fns = [absorb.absorb_offset(fn) for fn in prox_fns]
# TODO more analysis of what solver to use.
if show_graph:
print("Computational graph before optimizing:")
graph_visualize(prox_fns, filename = show_graph if type(show_graph) is str else None)
# Short circuit with one function.
if len(prox_fns) == 1 and type(prox_fns[0].lin_op) == Variable:
fn = prox_fns[0]
var = fn.lin_op
var.value = fn.prox(0, np.zeros(fn.lin_op.shape))
return fn.value
elif solver in NAME_TO_SOLVER:
module = NAME_TO_SOLVER[solver]
if len(self.omega_fns + self.psi_fns) == 0:
if self.try_split and len(prox_fns) > 1 and len(self.variables()) == 1:
psi_fns, omega_fns = module.partition(prox_fns,
self.try_diagonalize)
else:
psi_fns = prox_fns
omega_fns = []
else:
psi_fns = self.psi_fns
omega_fns = self.omega_fns
if test_norm:
L = CompGraph(vstack([fn.lin_op for fn in psi_fns]))
from numpy.random import random
output_mags = [NotImplemented]
L.norm_bound(output_mags)
if not NotImplemented in output_mags:
assert len(output_mags) == 1
x = random(L.input_size)
x = x / LA.norm(x)
y = np.zeros(L.output_size)
y = L.forward(x, y)
ny = LA.norm(y)
nL2 = est_CompGraph_norm(L, try_fast_norm=False)
if ny > output_mags[0]:
raise RuntimeError("wrong implementation of norm!")
print("%.3f <= ||K|| = %.3f (%.3f)" % (ny, output_mags[0], nL2))
# Scale the problem.
if self.scale:
K = CompGraph(vstack([fn.lin_op for fn in psi_fns]),
implem=self.implem)
Knorm = est_CompGraph_norm(K, try_fast_norm=self.try_fast_norm)
for idx, fn in enumerate(psi_fns):
psi_fns[idx] = fn.copy(fn.lin_op / Knorm,
beta=fn.beta * np.sqrt(Knorm),
implem=self.implem)
for idx, fn in enumerate(omega_fns):
omega_fns[idx] = fn.copy(beta=fn.beta / np.sqrt(Knorm),
implem=self.implem)
for v in K.orig_end.variables():
if v.initval is not None:
v.initval *= np.sqrt(Knorm)
if not test_adjoints in [False, None]:
if test_adjoints is True:
test_adjoints = 1e-6
# test adjoints
L = CompGraph(vstack([fn.lin_op for fn in psi_fns]))
from numpy.random import random
x = random(L.input_size)
yt = np.zeros(L.output_size)
#print("x=", x)
yt = L.forward(x, yt)
#print("yt=", yt)
#print("x=", x)
y = random(L.output_size)
#print("y=", y)
xt = np.zeros(L.input_size)
xt = L.adjoint(y, xt)
#print("xt=", xt)
#print("y=", y)
r = np.abs( np.dot(np.ravel(y), np.ravel(yt)) - np.dot(np.ravel(x), np.ravel(xt)) )
#print( x.shape, y.shape, xt.shape, yt.shape)
if r > test_adjoints:
#print("yt=", yt)
#print("y =", y)
#print("xt=", xt)
#print("x =", x)
raise RuntimeError("Unmatched adjoints: " + str(r))
else:
print("Adjoint test passed.", r)
if self.implem == Impl['pycuda']:
kwargs['adapter'] = PyCudaAdapter()
opt_val = module.solve(psi_fns, omega_fns,
lin_solver=self.lin_solver,
try_diagonalize=self.try_diagonalize,
try_fast_norm=self.try_fast_norm,
scaled=self.scale,
*args, **kwargs)
# Unscale the variables.
if self.scale:
for var in self.variables():
var.value /= np.sqrt(Knorm)
return opt_val
else:
raise Exception("Unknown solver.")
def solve(psi_fns,
omega_fns,
tau=None,
sigma=None,
theta=None,
max_iters=1000,
eps_abs=1e-3,
eps_rel=1e-3,
x0=None,
lin_solver="cg",
lin_solver_options=None,
conv_check=100,
try_diagonalize=True,
try_fast_norm=False,
scaled=True,
implem=None,
metric=None,
convlog=None,
verbose=0,
callback=None,
adapter=NumpyAdapter()):
# Can only have one omega function.
assert len(omega_fns) <= 1
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops, implem=implem)
#graph_visualize(prox_fns)
if adapter.implem() == 'numpy':
K_forward = K.forward
K_adjoint = K.adjoint
prox_off_and_fac = lambda offset, factor, fn, *args, **kw: ne.evaluate(
'x*a+b', {
'x': fn.prox(*args, **kw),
'a': factor,
'b': offset
})
prox = lambda fn, *args, **kw: fn.prox(*args, **kw)
elif adapter.implem() == 'pycuda':
K_forward = K.forward_cuda
K_adjoint = K.adjoint_cuda
prox_off_and_fac = lambda offset, factor, fn, *args, **kw: fn.prox_cuda(
*args, offset=offset, factor=factor, **kw)
prox = lambda fn, *args, **kw: fn.prox_cuda(*args, **kw)
else:
raise RuntimeError("Implementation %s unknown" % adapter.implem())
# Select optimal parameters if wanted
if tau is None or sigma is None or theta is None:
tau, sigma, theta = est_params_pc(K, tau, sigma, verbose, scaled,
try_fast_norm)
elif callable(tau) or callable(sigma) or callable(theta):
if scaled:
L = 1
else:
L = est_CompGraph_norm(K, try_fast_norm)
# Initialize
x = adapter.zeros(K.input_size)
y = adapter.zeros(K.output_size)
xbar = adapter.zeros(K.input_size)
u = adapter.zeros(K.output_size)
z = adapter.zeros(K.output_size)
if x0 is not None:
x[:] = adapter.reshape(adapter.from_np(x0), K.input_size)
else:
x[:] = adapter.from_np(K.x0())
K_forward(x, y)
xbar[:] = x
# Buffers.
Kxbar = adapter.zeros(K.output_size)
Kx = adapter.zeros(K.output_size)
KTy = adapter.zeros(K.input_size)
KTu = adapter.zeros(K.input_size)
s = adapter.zeros(K.input_size)
prev_x = x.copy()
prev_Kx = Kx.copy()
prev_z = z.copy()
prev_u = u.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
prox_log_tot = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsLog([
"pc_iteration_tot", "copyprev", "calcz", "calcx", "omega_fn", "xbar",
"conv_check"
])
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(adapter.to_np(x))
objval = 0.0
for f in prox_fns:
evp = f.value
#print(str(f), '->', f.value)
objval += evp
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing["pc_iteration_tot"].tic()
if convlog is not None:
convlog.tic()
if callable(sigma):
csigma = sigma(i, L)
else:
csigma = sigma
if callable(tau):
ctau = tau(i, L)
else:
ctau = tau
if callable(theta):
ctheta = theta(i, L)
else:
ctheta = theta
csigma = adapter.scalar(csigma)
ctau = adapter.scalar(ctau)
ctheta = adapter.scalar(ctheta)
# Keep track of previous iterates
iter_timing["copyprev"].tic()
adapter.copyto(prev_x, x)
adapter.copyto(prev_z, z)
adapter.copyto(prev_u, u)
adapter.copyto(prev_Kx, Kx)
iter_timing["copyprev"].toc()
# Compute z
iter_timing["calcz"].tic()
K_forward(xbar, Kxbar)
ne.evaluate('y + csigma * Kxbar', out=z)
iter_timing["calcz"].toc()
# Update y.
offset = 0
for fn in psi_fns:
prox_log_tot[fn].tic()
slc = slice(offset, offset + fn.lin_op.size, None)
z_slc = adapter.reshape(z[slc], fn.lin_op.shape)
# Moreau identity: apply and time prox.
prox_log[fn].tic()
y[slc] = adapter.flatten(
prox_off_and_fac(z_slc, -csigma, fn, csigma, z_slc / csigma,
i))
prox_log[fn].toc()
offset += fn.lin_op.size
prox_log_tot[fn].toc()
iter_timing["calcx"].tic()
if offset < y.shape[0]:
y[offset:] = 0
# Update x
K_adjoint(y, KTy)
ne.evaluate('x - ctau * KTy', out=x)
iter_timing["calcx"].toc()
iter_timing["omega_fn"].tic()
if len(omega_fns) > 0:
fn = omega_fns[0]
prox_log_tot[fn].tic()
xtmp = adapter.reshape(x, fn.lin_op.shape)
prox_log[fn].tic()
if adapter.implem() == 'numpy':
# ravel() avoids a redundant memcpy
x[:] = prox(fn,
1.0 / ctau,
xtmp,
x_init=prev_x,
lin_solver=lin_solver,
options=lin_solver_options).ravel()
else:
x[:] = adapter.flatten(
prox(fn,
1.0 / ctau,
xtmp,
x_init=prev_x,
lin_solver=lin_solver,
options=lin_solver_options))
prox_log[fn].toc()
prox_log_tot[fn].toc()
iter_timing["omega_fn"].toc()
iter_timing["xbar"].tic()
# Update xbar
ne.evaluate('x + ctheta * (x - prev_x)', out=xbar)
iter_timing["xbar"].toc()
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(adapter.to_np(x))
objval = list([fn.value for fn in prox_fns])
objval = sum(objval)
convlog.record_objective(objval)
# Residual based convergence check
if i % conv_check in [0, conv_check - 1]:
iter_timing["conv_check"].tic()
K_forward(x, Kx)
ne.evaluate('y / csigma + ctheta * (Kx - prev_Kx)',
out=u,
casting='unsafe')
ne.evaluate('prev_u + prev_Kx - y / csigma',
out=z,
casting='unsafe')
iter_timing["conv_check"].toc()
# Iteration order is different than
# lin-admm (--> start checking at iteration 1)
if i > 0 and i % conv_check == 0:
# Check convergence
r = ne.evaluate('prev_Kx - z')
dz = ne.evaluate('csigma * (z - prev_z)')
K_adjoint(dz, s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(prev_Kx), np.linalg.norm(z)])
K_adjoint(u, KTu)
eps_dual = np.sqrt(
K.input_size) * eps_abs + eps_rel * np.linalg.norm(
KTu) / csigma
if not callback is None or verbose == 2:
K.update_vars(adapter.to_np(x))
if not callback is None:
callback(adapter.to_np(x))
# Progress
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
ov = list([fn.value for fn in prox_fns])
objval = sum(ov)
objstr = ", obj_val = %02.03e [%s] " % (objval, ", ".join(
"%02.03e" % x for x in ov))
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(
metric.message(v))
print(
"iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s"
% (i, np.linalg.norm(adapter.to_np(r)), eps_pri,
np.linalg.norm(
adapter.to_np(s)), eps_dual, objstr, metstr))
iter_timing["pc_iteration_tot"].toc()
if np.linalg.norm(adapter.to_np(r)) <= eps_pri and np.linalg.norm(
adapter.to_np(s)) <= eps_dual:
break
else:
iter_timing["pc_iteration_tot"].toc()
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs total:")
print(prox_log_tot)
print("prox funcs inner:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(adapter.to_np(x))
if not callback is None:
callback(adapter.to_np(x))
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(psi_fns,
omega_fns,
rho_0=1.0,
rho_scale=math.sqrt(2.0) * 2.0,
rho_max=2**8,
max_iters=-1,
max_inner_iters=100,
x0=None,
eps_rel=1e-3,
eps_abs=1e-3,
lin_solver="cg",
lin_solver_options=None,
try_diagonalize=True,
scaled=False,
try_fast_norm=False,
metric=None,
convlog=None,
verbose=0):
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Rescale so (1/2)||x - b||^2_2
rescaling = np.sqrt(2.)
quad_ops = []
quad_weights = []
const_terms = []
for fn in omega_fns:
fn = fn.absorb_params()
quad_ops.append(scale(rescaling * fn.beta, fn.lin_op))
quad_weights.append(rescaling * fn.beta)
const_terms.append(fn.b.flatten() * rescaling)
# Get optimize inverse (tries spatial and frequency diagonalization)
op_list = [func.lin_op for func in psi_fns] + quad_ops
stacked_ops = vstack(op_list)
x_update = get_least_squares_inverse(op_list, None, try_diagonalize,
verbose)
# Initialize
if x0 is not None:
x = np.reshape(x0, K.input_size)
else:
x = np.zeros(K.input_size)
Kx = np.zeros(K.output_size)
w = Kx.copy()
# Temporary iteration counts
x_prev = x.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("HQS iteration")
inner_iter_timing = TimingsEntry("HQS inner iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(x)
objval = sum([func.value for func in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
# Rho scedule
rho = rho_0
i = 0
while rho < rho_max and i < max_iters:
iter_timing.tic()
if convlog is not None:
convlog.tic()
# Update rho for quadratics
for idx, op in enumerate(quad_ops):
op.scalar = quad_weights[idx] / np.sqrt(rho)
x_update = get_least_squares_inverse(op_list, CompGraph(stacked_ops),
try_diagonalize, verbose)
for ii in range(max_inner_iters):
inner_iter_timing.tic()
# Update Kx.
K.forward(x, Kx)
# Prox update to get w.
offset = 0
w_prev = w.copy()
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
# Apply and time prox.
prox_log[fn].tic()
w[slc] = fn.prox(rho, np.reshape(Kx[slc], fn.lin_op.shape),
ii).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update x.
x_prev[:] = x
tmp = np.hstack([w] +
[cterm / np.sqrt(rho) for cterm in const_terms])
x = x_update.solve(tmp,
x_init=x,
lin_solver=lin_solver,
options=lin_solver_options)
# Very basic convergence check.
r_x = np.linalg.norm(x_prev - x)
eps_x = eps_rel * np.prod(K.input_size)
r_w = np.linalg.norm(w_prev - w)
eps_w = eps_rel * np.prod(K.output_size)
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(x)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(x)
objstr = ", obj_val = %02.03e" % sum(
[fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(
metric.message(x))
print("iter [%02d (rho=%2.1e) || %02d]:"
"||w - w_prev||_2 = %02.02e (eps=%02.03e)"
"||x - x_prev||_2 = %02.02e (eps=%02.03e)%s%s" %
(i, rho, ii, r_x, eps_x, r_w, eps_w, objstr, metstr))
inner_iter_timing.toc()
if r_x < eps_x and r_w < eps_w:
break
# Update rho
rho = np.minimum(rho * rho_scale, rho_max)
i += 1
iter_timing.toc()
# Print out timings info.
if verbose > 0:
print(iter_timing)
print(inner_iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(x)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(psi_fns,
omega_fns,
rho=1.0,
max_iters=1000,
eps_abs=1e-1,
eps_rel=1e-3,
x0=None,
lin_solver="cg",
lin_solver_options=None,
try_diagonalize=True,
try_fast_norm=False,
scaled=True,
metric=None,
convlog=None,
verbose=0):
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Rescale so (rho/2)||x - b||^2_2
rescaling = np.sqrt(2. / rho)
quad_ops = []
const_terms = []
for fn in omega_fns:
fn = fn.absorb_params()
quad_ops.append(scale(rescaling * fn.beta, fn.lin_op))
const_terms.append(fn.b.flatten() * rescaling)
# Check for fast inverse.
op_list = [func.lin_op for func in psi_fns] + quad_ops
stacked_ops = vstack(op_list)
# Get optimize inverse (tries spatial and frequency diagonalization)
v_update = get_least_squares_inverse(op_list, None, try_diagonalize,
verbose)
# Initialize everything to zero.
input_size = K.input_size
output_size = K.output_size
v = np.zeros(input_size)
z = np.zeros(output_size)
u = np.zeros(output_size)
N_z = len(z[:])
print(input_size)
print(output_size)
# Initialize
if x0 is not None:
v[:] = np.reshape(x0, input_size)
K.forward(v, z)
# Buffers.
Kv = np.zeros(output_size)
KTu = np.zeros(input_size)
s = np.zeros(input_size)
Kv_pre = Kv.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("ADMM iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(v)
objval = sum([func.value for func in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
res_pre = 9e20
res = 0
curr_time = 0
total_time = []
Combine_res = []
# ------------------------------------------------------------------------------------
for i in range(max_iters):
# iter_timing.tic()
t1 = time.time()
if convlog is not None:
convlog.tic()
K.forward(v, Kv)
Kv_pre = Kv.copy()
# z_prev = z.copy()
# Update z.
K.forward(v, Kv)
Kv_u = Kv + u
offset = 0
for fn in psi_fns:
tmp = np.hstack([z - u] + const_terms)
v = v_update.solve(tmp,
x_init=v,
lin_solver=lin_solver,
options=lin_solver_options)
K.forward(v, Kv)
Kv_u = Kv + u
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
z_pre = z.copy()
prox_log[fn].tic()
z[slc] = fn.prox(rho, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update v.
# Check convergence.
r = Kv - z
# Update u.
u += r
K.adjoint(u, KTu)
# K.adjoint(rho * (z - z_prev), s)
s = z - z_pre
t2 = time.time()
curr_time += t2 - t1
res = np.linalg.norm(r)**2 + np.linalg.norm(s)**2
# K.adjoint((z-z_prev),s)
# eps_pri = np.sqrt(output_size) * eps_abs + eps_rel * \
# max([np.linalg.norm(Kv), np.linalg.norm(z)])
# eps_dual = np.sqrt(input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) / rho
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(v)
objstr = ", obj_val = %02.03e" % sum(
[fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(v))
# print("iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s" % (
# i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual, objstr, metstr))
print("iter %d: combine residual = %.8f" % (i, res))
#curr_time = curr_time + iter_timing.toc()
Combine_res.append(np.sqrt(rho * res / N_z))
total_time.append(curr_time)
# Exit if converged.
if (res) < eps_abs:
break
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(v)
# Return optimal value.
# return sum([fn.value for fn in prox_fns])
return total_time, Combine_res
def solve(psi_fns, omega_fns, tau=None, sigma=None, theta=None,
max_iters=1000, eps_abs=1e-3, eps_rel=1e-3, x0=None,
lin_solver="cg", lin_solver_options=None,
try_diagonalize=True, try_fast_norm=False, scaled=True,
metric=None, convlog=None, verbose=0):
# Can only have one omega function.
assert len(omega_fns) <= 1
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
v = np.zeros(K.input_size)
# Select optimal parameters if wanted
if tau is None or sigma is None or theta is None:
tau, sigma, theta = est_params_pc(K, tau, sigma, verbose, scaled, try_fast_norm)
# Initialize
x = np.zeros(K.input_size)
y = np.zeros(K.output_size)
xbar = np.zeros(K.input_size)
u = np.zeros(K.output_size)
z = np.zeros(K.output_size)
if x0 is not None:
x[:] = np.reshape(x0, K.input_size)
K.forward(x, y)
xbar[:] = x
# Buffers.
Kxbar = np.zeros(K.output_size)
Kx = np.zeros(K.output_size)
KTy = np.zeros(K.input_size)
KTu = np.zeros(K.input_size)
s = np.zeros(K.input_size)
prev_x = x.copy()
prev_Kx = Kx.copy()
prev_z = z.copy()
prev_u = u.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("PC iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(x)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing.tic()
if convlog is not None:
convlog.tic()
# Keep track of previous iterates
np.copyto(prev_x, x)
np.copyto(prev_z, z)
np.copyto(prev_u, u)
np.copyto(prev_Kx, Kx)
# Compute z
K.forward(xbar, Kxbar)
z = y + sigma * Kxbar
# Update y.
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
z_slc = np.reshape(z[slc], fn.lin_op.shape)
# Moreau identity: apply and time prox.
prox_log[fn].tic()
y[slc] = (z_slc - sigma * fn.prox(sigma, z_slc / sigma, i)).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
y[offset:] = 0
# Update x
K.adjoint(y, KTy)
x -= tau * KTy
if len(omega_fns) > 0:
xtmp = np.reshape(x, omega_fns[0].lin_op.shape)
x[:] = omega_fns[0].prox(1.0 / tau, xtmp, x_init=prev_x,
lin_solver=lin_solver, options=lin_solver_options).flatten()
# Update xbar
np.copyto(xbar, x)
xbar += theta * (x - prev_x)
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(x)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
""" Old convergence check
#Very basic convergence check.
r_x = np.linalg.norm(x - prev_x)
r_xbar = np.linalg.norm(xbar - prev_xbar)
r_ybar = np.linalg.norm(y - prev_y)
error = r_x + r_xbar + r_ybar
"""
# Residual based convergence check
K.forward(x, Kx)
u = 1.0 / sigma * y + theta * (Kx - prev_Kx)
z = prev_u + prev_Kx - 1.0 / sigma * y
# Iteration order is different than
# lin-admm (--> start checking at iteration 1)
if i > 0:
# Check convergence
r = prev_Kx - z
K.adjoint(sigma * (z - prev_z), s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(prev_Kx), np.linalg.norm(z)])
K.adjoint(u, KTu)
eps_dual = np.sqrt(K.input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) / sigma
# Progress
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(x)
objstr = ", obj_val = %02.03e" % sum([fn.value for fn in prox_fns])
""" Old convergence check
#Evaluate metric potentially
metstr = '' if metric is None else ", {}".format( metric.message(x.copy()) )
print "iter [%04d]:" \
"||x - x_prev||_2 = %02.02e " \
"||xbar - xbar_prev||_2 = %02.02e " \
"||y - y_prev||_2 = %02.02e " \
"SUM = %02.02e (eps=%02.03e)%s%s" \
% (i, r_x, r_xbar, r_ybar, error, eps, objstr, metstr)
"""
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(v))
print(
"iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s"
% (i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual, objstr, metstr)
)
iter_timing.toc()
if np.linalg.norm(r) <= eps_pri and np.linalg.norm(s) <= eps_dual:
break
else:
iter_timing.toc()
""" Old convergence check
if error <= eps:
break
"""
# Print out timings info.
if verbose > 0:
print iter_timing
print "prox funcs:"
print prox_log
print "K forward ops:"
print K.forward_log
print "K adjoint ops:"
print K.adjoint_log
# Assign values to variables.
K.update_vars(x)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(psi_fns, omega_fns, lmb=1.0, mu=None, quad_funcs=None,
max_iters=1000, eps_abs=1e-3, eps_rel=1e-3,
lin_solver="cg", lin_solver_options=None,
try_diagonalize=True, try_fast_norm=True, scaled=False,
metric=None, convlog=None, verbose=0):
# Can only have one omega function.
assert len(omega_fns) <= 1
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Select optimal parameters if wanted
if lmb is None or mu is None:
lmb, mu = est_params_lin_admm(K, lmb, verbose, scaled, try_fast_norm)
# Initialize everything to zero.
v = np.zeros(K.input_size)
z = np.zeros(K.output_size)
u = np.zeros(K.output_size)
# Buffers.
Kv = np.zeros(K.output_size)
KTu = np.zeros(K.input_size)
s = np.zeros(K.input_size)
Kvzu = np.zeros(K.output_size)
v_prev = np.zeros(K.input_size)
z_prev = np.zeros(K.output_size)
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("LIN-ADMM iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing.tic()
if convlog is not None:
convlog.tic()
v_prev[:] = v
z_prev[:] = z
# Update v
K.forward(v, Kv)
Kvzu[:] = Kv - z + u
K.adjoint(Kvzu, v)
v[:] = v_prev - (mu / lmb) * v
if len(omega_fns) > 0:
v[:] = omega_fns[0].prox(1.0 / mu, v, x_init=v_prev.copy(),
lin_solver=lin_solver, options=lin_solver_options)
# Update z.
K.forward(v, Kv)
Kv_u = Kv + u
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
prox_log[fn].tic()
z[slc] = fn.prox(1.0 / lmb, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update u.
u += Kv - z
K.adjoint(u, KTu)
# Check convergence.
r = Kv - z
K.adjoint((1.0 / lmb) * (z - z_prev), s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(Kv), np.linalg.norm(z)])
eps_dual = np.sqrt(K.input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) / (1.0 / lmb)
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(v)
objstr = ", obj_val = %02.03e" % sum([fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(v))
print "iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s" % (
i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual, objstr, metstr)
iter_timing.toc()
if np.linalg.norm(r) <= eps_pri and np.linalg.norm(s) <= eps_dual:
break
# Print out timings info.
if verbose > 0:
print iter_timing
print "prox funcs:"
print prox_log
print "K forward ops:"
print K.forward_log
print "K adjoint ops:"
print K.adjoint_log
# Assign values to variables.
K.update_vars(v)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(psi_fns,
omega_fns,
rho=1.0,
max_iters=1000,
eps_abs=1e-10,
eps_rel=1e-3,
x0=None,
lin_solver="cg",
lin_solver_options=None,
try_diagonalize=True,
try_fast_norm=False,
scaled=True,
metric=None,
convlog=None,
verbose=0):
# C=np.array([[1,0],[0,0]]);
# b=np.array([2,0]);
# print(np.linalg.lstsq(C,b,rcond=None)[0])
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Rescale so (rho/2)||x - b||^2_2
rescaling = np.sqrt(2. / rho)
quad_ops = []
const_terms = []
for fn in omega_fns:
fn = fn.absorb_params()
quad_ops.append(scale(rescaling * fn.beta, fn.lin_op))
const_terms.append(fn.b.flatten() * rescaling)
# Check for fast inverse.
op_list = [func.lin_op for func in psi_fns] + quad_ops
stacked_ops = vstack(op_list)
# Get optimize inverse (tries spatial and frequency diagonalization)
v_update = get_least_squares_inverse(op_list, None, try_diagonalize,
verbose)
# Initialize everything to zero.
input_size = K.input_size
output_size = K.output_size
v = np.zeros(input_size)
z = np.zeros(output_size)
u = np.zeros(output_size)
print(output_size)
# Initialize
if x0 is not None:
v[:] = np.reshape(x0, input_size)
K.forward(v, z)
# Buffers.
v0 = v.copy()
z0 = z.copy()
u0 = u.copy()
N_z = len(z[:])
Kv = np.zeros(output_size)
KTu = np.zeros(input_size)
s = np.zeros(input_size)
Kv_pre = Kv.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("ADMM iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(v)
objval = sum([func.value for func in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
# --------------------------------------------------------------------------------------------------
print("Anderson Acceleration:")
for andersonmk in range(6, 7):
v = v0.copy()
u = u0.copy()
v_d = v.copy()
u_d = u.copy()
res_pre = 9e20
total_energy = []
total_time = []
Combine_res = []
reset = False
sca_z = 1
size = v.flatten().shape[0]
total_size = (u.flatten()).shape[0] + size
print(size)
sign = 0
curr_time = 0
AA_compute_time = 0
acc1 = Anderson(
np.concatenate((v.flatten(), sca_z * u.flatten()), axis=0),
total_size, andersonmk)
for i in range(max_iters):
t1 = time.time()
if convlog is not None:
convlog.tic()
K.forward(v, Kv)
# Update z.
Kv_pre = Kv.copy()
K.forward(v, Kv)
Kv_u = Kv + u
offset = 0
for fn in psi_fns:
tmp = np.hstack([z - u] + const_terms)
v = v_update.solve(tmp,
x_init=v,
lin_solver=lin_solver,
options=lin_solver_options)
K.forward(v, Kv)
Kv_u = Kv + u
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
z_pre = z.copy()
prox_log[fn].tic()
z[slc] = fn.prox(rho, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update u.
r = Kv - z
u += r
K.adjoint(u, KTu)
# print(np.linalg.norm(u))
# Check convergence.
# K.adjoint(rho * (z - z_prev), s)
s = z - z_pre
res = np.linalg.norm(r)**2 + np.linalg.norm(s)**2
# K.adjoint((z-z_prev),s)
# eps_pri = np.sqrt(output_size) * eps_abs + eps_rel * \
# max([np.linalg.norm(Kv), np.linalg.norm(z)])
# eps_dual = np.sqrt(input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) / rho
t3 = time.time()
if res < res_pre or reset == True:
v_d = v.copy()
u_d = u.copy()
res_pre = res
reset = False
tt = acc1.compute(
np.concatenate((v.flatten(), sca_z * u.flatten()), axis=0))
v = tt[0:size].reshape(v.shape)
u = tt[size:].reshape(u.shape) / sca_z
else:
sign = sign + 1
v = v_d.copy()
u = u_d.copy()
reset = True
acc1.reset(
np.concatenate((v.flatten(), sca_z * u.flatten()), axis=0))
t4 = time.time()
AA_compute_time += t4 - t3
t2 = time.time()
curr_time += t2 - t1
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(v)
objstr = ", obj_val = %02.03e" % sum(
[fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(
metric.message(v))
# print("iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s" % (
# i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual, objstr, metstr))
print("iter %d: combine residual = %.8f" % (i, res))
Combine_res.append(np.sqrt(rho * res_pre / N_z))
total_time.append(curr_time)
# Exit if converged.
if (res) < eps_abs:
break
print("current time: %.6f, AA compute: %.6f, sign: %d" %
(curr_time, AA_compute_time, sign))
hm_src_path = 'residual-' + str(andersonmk) + '.txt'
iter_num = []
iter_num.append(len(total_time))
iter_num.append(len(Combine_res))
with open(hm_src_path, 'w') as f:
for i in range(0, min(iter_num)):
f.write('%f\t%.20f\n' % (total_time[i], Combine_res[i]))
f.close()
print("Anderson Acceleration with Douglas-Rachford splitting:")
for andersonmk in range(6, 7):
v = v0.copy()
u = u0.copy()
K.forward(v, Kv)
v_d = v.copy()
u_d = u.copy()
d_s = z0.copy()
d_u = d_s.copy()
d_s_d = d_s.copy()
d_v = d_s.copy()
d_unew = d_u.copy()
res_pre = 9e20
r_com = 0
r_com_pre = r_com
total_energy = []
total_time = []
Combine_res = []
reset = False
size = v.flatten().shape[0]
sign = 0
curr_time = 0
acc1 = Anderson(d_s.flatten(), size, andersonmk)
for i in range(max_iters):
t1 = time.time()
if convlog is not None:
convlog.tic()
# K.forward(v, Kv)
# Update v.
Kv_u = d_s.copy()
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
prox_log[fn].tic()
z[slc] = fn.prox(rho, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
d_u = z.copy()
temp = 2 * d_u - d_s
tmp = np.hstack([temp] + const_terms)
v = v_update.solve(tmp,
x_init=v,
lin_solver=lin_solver,
options=lin_solver_options)
K.forward(v, d_v)
# z_prev = z.copy()
# Update z.
# Update d_s
r = d_v - d_u
d_s += r
res = np.linalg.norm(r)**2
t2 = time.time()
curr_time += t2 - t1
# print(np.linalg.norm(u))
Kv_u = d_s.copy()
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
prox_log[fn].tic()
z[slc] = fn.prox(rho, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
d_unew = z.copy()
# Check convergence.
# K.adjoint(rho * (z - z_prev),
r_com = np.linalg.norm(d_unew - d_v)**2 + np.linalg.norm(d_unew -
d_u)**2
# K.adjoint((z-z_prev),s)
# eps_pri = np.sqrt(output_size) * eps_abs + eps_rel * \
# max([np.linalg.norm(Kv), np.linalg.norm(z)])
# eps_dual = np.sqrt(input_size) * eps_abs + eps_rel * np.linalg.norm(KTu) / rho
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
t1 = time.time()
if res < res_pre or reset == True:
d_s_d = d_s.copy()
res_pre = res
r_com_pre = r_com
reset = False
tt = acc1.compute(d_s.flatten())
d_s = tt.reshape(d_s.shape)
else:
sign = sign + 1
d_s = d_s_d.copy()
reset = True
acc1.reset(d_s.flatten())
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(v)
objstr = ", obj_val = %02.03e" % sum(
[fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(
metric.message(v))
# print("iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s" % (
# i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual, objstr, metstr))
print("iter %d: combine residual = %.8f" % (i, r_com))
t2 = time.time()
curr_time += t2 - t1
Combine_res.append(np.sqrt(rho * r_com_pre / N_z))
total_time.append(curr_time)
# Exit if converged.
# if (res) < eps_abs:
# break
hm_src_path = 'dr-' + str(andersonmk) + '.txt'
iter_num = []
iter_num.append(len(total_time))
iter_num.append(len(Combine_res))
with open(hm_src_path, 'w') as f:
for i in range(0, min(iter_num)):
f.write('%f\t%.20f\n' % (total_time[i], Combine_res[i]))
f.close()
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(v)
# Return optimal value.
# return sum([fn.value for fn in prox_fns])
return total_time, Combine_res
def solve(psi_fns,
omega_fns,
lmb=1.0,
mu=None,
quad_funcs=None,
max_iters=1000,
eps_abs=1e-3,
eps_rel=1e-3,
lin_solver="cg",
lin_solver_options=None,
try_diagonalize=True,
try_fast_norm=True,
scaled=False,
metric=None,
convlog=None,
verbose=0):
# Can only have one omega function.
assert len(omega_fns) <= 1
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Select optimal parameters if wanted
if lmb is None or mu is None:
lmb, mu = est_params_lin_admm(K, lmb, verbose, scaled, try_fast_norm)
# Initialize everything to zero.
v = np.zeros(K.input_size)
z = np.zeros(K.output_size)
u = np.zeros(K.output_size)
# Buffers.
Kv = np.zeros(K.output_size)
KTu = np.zeros(K.input_size)
s = np.zeros(K.input_size)
Kvzu = np.zeros(K.output_size)
v_prev = np.zeros(K.input_size)
z_prev = np.zeros(K.output_size)
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("LIN-ADMM iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
for i in range(max_iters):
iter_timing.tic()
if convlog is not None:
convlog.tic()
v_prev[:] = v
z_prev[:] = z
# Update v
K.forward(v, Kv)
Kvzu[:] = Kv - z + u
K.adjoint(Kvzu, v)
v[:] = v_prev - (mu / lmb) * v
if len(omega_fns) > 0:
v[:] = omega_fns[0].prox(1.0 / mu,
v,
x_init=v_prev.copy(),
lin_solver=lin_solver,
options=lin_solver_options)
# Update z.
K.forward(v, Kv)
Kv_u = Kv + u
offset = 0
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
Kv_u_slc = np.reshape(Kv_u[slc], fn.lin_op.shape)
# Apply and time prox.
prox_log[fn].tic()
z[slc] = fn.prox(1.0 / lmb, Kv_u_slc, i).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update u.
u += Kv - z
K.adjoint(u, KTu)
# Check convergence.
r = Kv - z
K.adjoint((1.0 / lmb) * (z - z_prev), s)
eps_pri = np.sqrt(K.output_size) * eps_abs + eps_rel * \
max([np.linalg.norm(Kv), np.linalg.norm(z)])
eps_dual = np.sqrt(K.input_size) * eps_abs + eps_rel * np.linalg.norm(
KTu) / (1.0 / lmb)
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(v)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(v)
objstr = ", obj_val = %02.03e" % sum(
[fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(v))
print(
"iter %d: ||r||_2 = %.3f, eps_pri = %.3f, ||s||_2 = %.3f, eps_dual = %.3f%s%s"
% (i, np.linalg.norm(r), eps_pri, np.linalg.norm(s), eps_dual,
objstr, metstr))
iter_timing.toc()
if np.linalg.norm(r) <= eps_pri and np.linalg.norm(s) <= eps_dual:
break
# Print out timings info.
if verbose > 0:
print(iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(v)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
def solve(psi_fns, omega_fns,
rho_0=1.0, rho_scale=math.sqrt(2.0) * 2.0, rho_max=2**8,
max_iters=-1, max_inner_iters=100, x0=None,
eps_rel=1e-3, eps_abs=1e-3,
lin_solver="cg", lin_solver_options=None,
try_diagonalize=True, scaled=False, try_fast_norm=False,
metric=None, convlog=None, verbose=0):
prox_fns = psi_fns + omega_fns
stacked_ops = vstack([fn.lin_op for fn in psi_fns])
K = CompGraph(stacked_ops)
# Rescale so (1/2)||x - b||^2_2
rescaling = np.sqrt(2.)
quad_ops = []
quad_weights = []
const_terms = []
for fn in omega_fns:
fn = fn.absorb_params()
quad_ops.append(scale(rescaling * fn.beta, fn.lin_op))
quad_weights.append(rescaling * fn.beta)
const_terms.append(fn.b.flatten() * rescaling)
# Get optimize inverse (tries spatial and frequency diagonalization)
op_list = [func.lin_op for func in psi_fns] + quad_ops
stacked_ops = vstack(op_list)
x_update = get_least_squares_inverse(op_list, None,
try_diagonalize, verbose)
# Initialize
if x0 is not None:
x = np.reshape(x0, K.input_size)
else:
x = np.zeros(K.input_size)
Kx = np.zeros(K.output_size)
w = Kx.copy()
# Temporary iteration counts
x_prev = x.copy()
# Log for prox ops.
prox_log = TimingsLog(prox_fns)
# Time iterations.
iter_timing = TimingsEntry("HQS iteration")
inner_iter_timing = TimingsEntry("HQS inner iteration")
# Convergence log for initial iterate
if convlog is not None:
K.update_vars(x)
objval = sum([func.value for func in prox_fns])
convlog.record_objective(objval)
convlog.record_timing(0.0)
# Rho scedule
rho = rho_0
i = 0
while rho < rho_max and i < max_iters:
iter_timing.tic()
if convlog is not None:
convlog.tic()
# Update rho for quadratics
for idx, op in enumerate(quad_ops):
op.scalar = quad_weights[idx] / np.sqrt(rho)
x_update = get_least_squares_inverse(op_list, CompGraph(stacked_ops),
try_diagonalize, verbose)
for ii in range(max_inner_iters):
inner_iter_timing.tic()
# Update Kx.
K.forward(x, Kx)
# Prox update to get w.
offset = 0
w_prev = w.copy()
for fn in psi_fns:
slc = slice(offset, offset + fn.lin_op.size, None)
# Apply and time prox.
prox_log[fn].tic()
w[slc] = fn.prox(rho, np.reshape(Kx[slc], fn.lin_op.shape), ii).flatten()
prox_log[fn].toc()
offset += fn.lin_op.size
# Update x.
x_prev[:] = x
tmp = np.hstack([w] + [cterm / np.sqrt(rho) for cterm in const_terms])
x = x_update.solve(tmp, x_init=x, lin_solver=lin_solver, options=lin_solver_options)
# Very basic convergence check.
r_x = np.linalg.norm(x_prev - x)
eps_x = eps_rel * np.prod(K.input_size)
r_w = np.linalg.norm(w_prev - w)
eps_w = eps_rel * np.prod(K.output_size)
# Convergence log
if convlog is not None:
convlog.toc()
K.update_vars(x)
objval = sum([fn.value for fn in prox_fns])
convlog.record_objective(objval)
# Show progess
if verbose > 0:
# Evaluate objective only if required (expensive !)
objstr = ''
if verbose == 2:
K.update_vars(x)
objstr = ", obj_val = %02.03e" % sum([fn.value for fn in prox_fns])
# Evaluate metric potentially
metstr = '' if metric is None else ", {}".format(metric.message(x))
print("iter [%02d (rho=%2.1e) || %02d]:"
"||w - w_prev||_2 = %02.02e (eps=%02.03e)"
"||x - x_prev||_2 = %02.02e (eps=%02.03e)%s%s"
% (i, rho, ii, r_x, eps_x, r_w, eps_w, objstr, metstr))
inner_iter_timing.toc()
if r_x < eps_x and r_w < eps_w:
break
# Update rho
rho = np.minimum(rho * rho_scale, rho_max)
i += 1
iter_timing.toc()
# Print out timings info.
if verbose > 0:
print(iter_timing)
print(inner_iter_timing)
print("prox funcs:")
print(prox_log)
print("K forward ops:")
print(K.forward_log)
print("K adjoint ops:")
print(K.adjoint_log)
# Assign values to variables.
K.update_vars(x)
# Return optimal value.
return sum([fn.value for fn in prox_fns])
