def test_pylops(xp, G, b): G = xp.asarray(G) b = xp.asarray(b) svx = xp.linalg.solve(G, b) svx2 = xp.tile(svx, 2) G2 = pylops.BlockDiag([G, G]) b2 = xp.tile(b, 2) xo, *_ = lsqr(G2, b2, show=show, atol=tol, btol=tol, iter_lim=maxit) assert xp.allclose(xo, svx2, atol=tol, rtol=tol)
import numpy as np import matplotlib.pyplot as plt import pylops plt.close('all') ############################################################################### # Let's start by creating N MatrixMult operators and the BlockDiag operator N = 100 Nops = 32 Ops = [ pylops.MatrixMult(np.random.normal(0., 1., (N, N))) for _ in range(Nops) ] Op = pylops.BlockDiag(Ops, nproc=1) ############################################################################### # We can now perform a scalability test on the forward operation workers = [2, 3, 4] compute_times, speedup = \ pylops.utils.multiproc.scalability_test(Op, np.ones(Op.shape[1]), workers=workers, forward=True) plt.figure(figsize=(12, 3)) plt.plot(workers, speedup, 'ko-') plt.xlabel('# Workers') plt.ylabel('Speed Up') plt.title('Forward scalability test') plt.tight_layout() ###############################################################################
# to three different subset of the model and data # # .. math:: # \mathbf{D_{BDiag}} = # \begin{bmatrix} # \mathbf{D_v} & \mathbf{0} & \mathbf{0} \\ # \mathbf{0} & 0.5*\mathbf{D_v} & \mathbf{0} \\ # \mathbf{0} & \mathbf{0} & -\mathbf{D_h} # \end{bmatrix}, \qquad # \mathbf{y} = # \begin{bmatrix} # \mathbf{D_v} \mathbf{x_1} \\ # 0.5*\mathbf{D_v} \mathbf{x_2} \\ # -\mathbf{D_h} \mathbf{x_3} # \end{bmatrix} BD = pylops.BlockDiag([D2vop, 0.5 * D2vop, -1 * D2hop]) Y = np.reshape(BD * X.ravel(), (3 * Nv, Nh)) fig, axs = plt.subplots(1, 2, figsize=(10, 3)) fig.suptitle("Block-diagonal", fontsize=14, fontweight="bold", y=0.95) im = axs[0].imshow(X, interpolation="nearest") axs[0].axis("tight") axs[0].set_title(r"$x$") plt.colorbar(im, ax=axs[0]) im = axs[1].imshow(Y, interpolation="nearest") axs[1].axis("tight") axs[1].set_title(r"$y$") plt.colorbar(im, ax=axs[1]) plt.tight_layout() plt.subplots_adjust(top=0.8)
def BlockDiagonalOperator(*linops: LinearOperator, n_jobs: int = 1) -> PyLopLinearOperator: r""" Construct a block diagonal operator from N linear operators. Parameters ---------- linops: LinearOperator Linear operators forming the diagonal blocks. Alternatively, numpy.ndarray or scipy.sparse.spmatrix can be passed in place of one or more operators. n_jobs: int Number of processes used to evaluate the N operators in parallel using multiprocessing. If ``n_jobs=1`` (default), work in serial mode. Returns ------- PyLopLinearOperator Block diagonal linear operator. Examples -------- .. doctest:: >>> from pycsou.linop.base import BlockDiagonalOperator >>> from pycsou.linop.diff import SecondDerivative >>> Nv, Nh = 11, 21 >>> D2hop = SecondDerivative(size=Nv * Nh, shape=(Nv,Nh), axis=1) >>> D2vop = SecondDerivative(size=Nv * Nh, shape=(Nv,Nh), axis=0) >>> Dblockdiag = BlockDiagonalOperator(D2vop, 0.5 * D2vop, -1 * D2hop) >>> x = np.zeros((Nv, Nh)); x[int(Nv//2), int(Nh//2)] = 1; z = np.tile(x, (3,1)).flatten() >>> np.allclose(Dblockdiag(z), np.concatenate((D2vop(x.flatten()), 0.5 * D2vop(x.flatten()), - D2hop(x.flatten())))) True Notes ----- A block-diagonal operator composed of N linear operators is created such as its application in forward mode leads to .. math:: \begin{bmatrix} \mathbf{L_1} & \mathbf{0} & \cdots & \mathbf{0} \\ \mathbf{0} & \mathbf{L_2} & \cdots & \mathbf{0} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{0} & \mathbf{0} & \cdots & \mathbf{L_N} \end{bmatrix} \begin{bmatrix} \mathbf{x}_{1} \\ \mathbf{x}_{2} \\ \vdots \\ \mathbf{x}_{N} \end{bmatrix} = \begin{bmatrix} \mathbf{L_1} \mathbf{x}_{1} \\ \mathbf{L_2} \mathbf{x}_{2} \\ \vdots \\ \mathbf{L_N} \mathbf{x}_{N} \end{bmatrix} while its application in adjoint mode leads to .. math:: \begin{bmatrix} \mathbf{L_1}^\ast & \mathbf{0} & \cdots & \mathbf{0} \\ \mathbf{0} & \mathbf{L_2}^\ast & \cdots & \mathbf{0} \\ \vdots & \vdots & \ddots & \vdots \\ \mathbf{0} & \mathbf{0} & \cdots & \mathbf{L_N}^\ast \end{bmatrix} \begin{bmatrix} \mathbf{y}_{1} \\ \mathbf{y}_{2} \\ \vdots \\ \mathbf{y}_{N} \end{bmatrix} = \begin{bmatrix} \mathbf{L_1}^\ast \mathbf{y}_{1} \\ \mathbf{L_2}^\ast \mathbf{y}_{2} \\ \vdots \\ \mathbf{L_N}^\ast \mathbf{y}_{N} \end{bmatrix} The Lipschitz constant of the block-diagonal operator can be bounded by :math:`{\max_{i=1}^N \|\mathbf{L}_{i}\|_2}`. Warnings -------- The parameter ``n_jobs`` is currently unused and is there for compatibility with the future API of PyLops. The code should be updated when the next version on PyLops is released. See Also -------- :py:class:`~pycsou.linop.base.BlockOperator`, :py:class:`~pycsou.linop.base.LinOpStack` """ pylinops = [linop.PyLop for linop in linops] lipschitz_cst = np.array([linop.lipschitz_cst for linop in linops]).max() block_diag = pylops.BlockDiag(ops=pylinops) return PyLopLinearOperator(block_diag, lipschitz_cst=lipschitz_cst)
# \mathbf{0} & 0.5*\mathbf{D_v} & \mathbf{0} \\ # \mathbf{0} & \mathbf{0} & -\mathbf{D_h} # \end{bmatrix}, \qquad # \mathbf{y} = # \begin{bmatrix} # \mathbf{D_v} \mathbf{x_1} \\ # 0.5*\mathbf{D_v} \mathbf{x_2} \\ # -\mathbf{D_h} \mathbf{x_3} # \end{bmatrix} Nv, Nh = 11, 21 X = np.zeros((Nv * 3, Nh)) X[int(Nv / 2), int(Nh / 2)] = 1 X[int(Nv / 2) + Nv, int(Nh / 2)] = 1 X[int(Nv / 2) + 2 * Nv, int(Nh / 2)] = 1 Block = pylops.BlockDiag([D2vop, 0.5 * D2vop, -1 * D2hop]) Y = np.reshape(Block * np.ndarray.flatten(X), (11 * 3, 21)) fig, axs = plt.subplots(1, 2, figsize=(10, 3)) fig.suptitle('Block-diagonal', fontsize=14, fontweight='bold') im = axs[0].imshow(X, interpolation='nearest') axs[0].axis('tight') axs[0].set_title(r'$x$') plt.colorbar(im, ax=axs[0]) im = axs[1].imshow(Y, interpolation='nearest') axs[1].axis('tight') axs[1].set_title(r'$y$') plt.colorbar(im, ax=axs[1]) plt.tight_layout()
def solve(self, f: np.ndarray): """ Description of main primal-dual iteration. :return: None """ (primal_n, primal_m) = self.image_size v = w = 0 g = f.ravel() p = p_bar = np.zeros(primal_n*primal_m) q = q_bar = np.zeros(2*primal_n*primal_m) if self.reg_mode != 'tik': grad = pylops.Gradient(dims=(primal_n, primal_m), dtype='float64', edge=True, kind="backward") else: grad = pylops.Identity(np.prod(self.image_size)) grad1 = pylops.BlockDiag([grad, grad]) # symmetric dxdy <-> dydx not necessary (expensive) but easy and functional proj_0 = IndicatorL2((primal_n, primal_m), upper_bound=self.alpha[0]) proj_1 = IndicatorL2((2 * primal_n, primal_m), upper_bound=self.alpha[1]) if not self.silent_mode: progress = progressbar.ProgressBar(max_value=self.max_iter) k = 0 while (self.tol < self.sens or k == 0) and (k < self.max_iter): p_old = p q_old = q # Dual Update g = self.lam / (self.tau + self.lam) * (g + self.tau * (self.A*(p_bar ) - f)) #- self.alpha[0]*self.breg_p if self.reg_mode != 'tik': v = proj_0.prox(v + self.tau * (grad * p_bar - q_bar)) else: v = self.alpha[0] / (self.tau + self.alpha[0]) * \ (v + self.tau * (grad*p_bar - self.data)) if self.reg_mode == 'tgv': w = proj_1.prox(w + self.tau * grad1*q_bar) # Primal Update p = p - self.tau * (-self.alpha[0]*self.breg_p + self.A.H*g + grad.H * v) if self.reg_mode == 'tgv': q = q + self.tau * (v - grad1.H * w) # Extragradient Update p_bar = 2 * p - p_old q_bar = 2 * q - q_old if k % 50 == 0: p_gap = p - p_old self.sens = np.linalg.norm(p_gap)/np.linalg.norm(p_old) print(self.sens) if self.gamma: raise NotImplementedError("The adjustment of the step size in the " "Primal-Dual is not yet fully developed.") thetha = 1 / np.sqrt(1 + 2*self.gamma * self.G.prox_param) self.G.prox_param = thetha * self.G.prox_param self.F_star.prox_param = self.F_star.prox_param / thetha k += 1 if not self.silent_mode: progress.update(k) self.x = p if k <= self.max_iter: print(" Early stopping.") return self.x
def trap_phase_2(image_vecs_medsub, model_vecs, temporal_basis, trap_params: TrapParams): xp = core.get_array_module(image_vecs_medsub) was_gpu_array = xp is cp timers = {} flat_model_vecs = model_vecs.ravel() if trap_params.scale_model_std: model_coeff_scale = xp.std(flat_model_vecs) flat_model_vecs /= model_coeff_scale else: model_coeff_scale = 1 if trap_params.force_gpu_fit: temporal_basis = cp.asarray(temporal_basis) flat_model_vecs = cp.asarray(flat_model_vecs) operator_block_diag = [temporal_basis.T] * image_vecs_medsub.shape[0] opstack = [ pylops.BlockDiag(operator_block_diag), ] if trap_params.incorporate_offset: if trap_params.background_split_mask is not None: left_mask_vec = trap_params.background_split_mask left_mask_megavec = np.repeat(left_mask_vec[:, np.newaxis], model_vecs.shape[1]).ravel() assert len(left_mask_megavec) == len(flat_model_vecs) left_mask_megavec = left_mask_megavec[np.newaxis, :].astype( flat_model_vecs.dtype) left_mask_megavec = left_mask_megavec - left_mask_megavec.mean() left_mask_megavec /= np.linalg.norm(left_mask_megavec) # "ones" for left side pixels -> fit constant offset for left psf opstack.append(xp.asarray(left_mask_megavec)) # "ones" for right side pixels -> fit constant offset for right psf right_mask_megavec = -1 * left_mask_megavec opstack.append(xp.asarray(right_mask_megavec)) else: background_megavec = np.ones_like(flat_model_vecs[xp.newaxis, :]) background_megavec /= np.linalg.norm(background_megavec) opstack.append(xp.asarray(background_megavec)) opstack.append(flat_model_vecs[xp.newaxis, :]) op = pylops.VStack(opstack).transpose() log.debug(f"TRAP operator: {op}") image_megavec = image_vecs_medsub.ravel() log.debug( f"Performing inversion on A.shape={op.shape} and b={image_megavec.shape}" ) timers['invert'] = time.perf_counter() if trap_params.use_cgls: solver = pylops.optimization.solver.cgls log.debug(f"{solver=}") solver_kwargs = dict(damp=trap_params.damp, tol=trap_params.tol) cgls_result = solver(op, image_megavec, xp.zeros(int(op.shape[1])), **solver_kwargs) xinv = cgls_result[0] else: soln = lsqr.lsqr( op, image_megavec, x0=None, atol=trap_params.tol, damp=trap_params.damp, # show=True, # calc_var=True, ) xinv = soln[0] # import matplotlib.pyplot as plt # plt.plot(soln[-1]) # plt.yscale('log') timers['invert'] = time.perf_counter() - timers['invert'] log.debug(f"Finished RegularizedInversion in {timers['invert']} sec") if core.get_array_module(xinv) is cp: model_coeff = float(xinv.get()[-1]) else: model_coeff = float(xinv[-1]) model_coeff = model_coeff / model_coeff_scale # return model_coeff, timers if trap_params.compute_residuals: solnvec = xinv solnvec[-1] = 0 # zero planet model contribution log.debug(f"Constructing starlight estimate from fit vector") timers['subtract'] = time.perf_counter() estimate_vecs = op.dot(solnvec).reshape(image_vecs_medsub.shape) if core.get_array_module( image_vecs_medsub) is not core.get_array_module(estimate_vecs): image_vecs_medsub = core.get_array_module(estimate_vecs).asarray( image_vecs_medsub) resid_vecs = image_vecs_medsub - estimate_vecs if core.get_array_module(resid_vecs) is cp and not was_gpu_array: resid_vecs = resid_vecs.get() timers['subtract'] = time.perf_counter() - timers['subtract'] log.debug(f"Starlight subtracted in {timers['subtract']} sec") return model_coeff, timers, resid_vecs else: return model_coeff, timers, None