o :ήc~ ã@s$ddlZ       ddd„ZdS) éNççH¯¼šò×z>é'Fc s¸ˆdur dd„} n tˆƒrˆ} n‡fdd„} |durt |¡n| ¡} tj |¡} t|ƒr4||| ƒ} n|| | ¡} tj | ¡} | |ksK| || krM| S| | ƒ}| ¡}t | |¡}t|ƒD]v}|  ¡}t|ƒrn||ƒ}n| |¡}|t ||¡}| ||7} | ||8} tj | ¡} |durš||| || | | ƒ| |ks¤| || kr¨| S| | ƒ}|r¿t | ||¡|}t | |¡}nd|}t | |¡}||9}||9}||7}q_t d|ƒ‚)a£Solves a system of linear equations :math:`Ax=b` using conjugate gradient descent :cite:`hestenes1952methods` Parameters ---------- A: scipy.sparse.csr_matrix Square matrix b: numpy.ndarray Vector describing the right-hand side of the system x0: numpy.ndarray Initialization, if `None` then :code:`x=np.zeros_like(b)` atol: float Absolute tolerance. The loop terminates if the :math:`||r||` is smaller than `atol`, where :math:`r` denotes the residual of the current iterate. rtol: float Relative tolerance. The loop terminates if :math:`{||r||}/{||b||}` is smaller than `rtol`, where :math:`r` denotes the residual of the current iterate. callback: function Function :code:`callback(A, x, b, norm_b, r, norm_r)` called after each iteration, defaults to `None` M: function or scipy.sparse.csr_matrix Function that applies the preconditioner to a vector. Alternatively, `M` can be a matrix describing the precondioner. reorthogonalize: boolean Wether to apply reorthogonalization of the residuals after each update, defaults to `False` Returns ------- x: numpy.ndarray Solution of the system Example ------- >>> from pymatting import * >>> import numpy as np >>> A = np.array([[3.0, 1.0], [1.0, 2.0]]) >>> M = jacobi(A) >>> b = np.array([4.0, 3.0]) >>> cg(A, b, M=M) array([1., 1.]) NcSs|S©N©©Úxrrú:/tmp/pip-target-vg8gfxp4/lib/python/pymatting/solver/cg.pyÚ precondition6szcg..preconditioncs ˆ |¡Sr)Údotr©ÚMrr r =s gð?z@Conjugate gradient descent did not converge within %d iterations) ÚcallableÚnpÚ zeros_likeÚcopyÚlinalgÚnormr ÚinnerÚrangeÚ ValueError)ÚAÚbÚx0ÚatolÚrtolÚmaxiterÚcallbackr Úreorthogonalizer rÚnorm_bÚrÚnorm_rÚzÚpÚrzÚ iterationÚr_oldÚApÚalphaÚbetarr r ÚcgsR0             ÿr*)NrrrNNF)Únumpyrr*rrrr Ús÷