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