<s> """<STR_LIT>""" <EOL> import sys <EOL> import traceback <EOL> import numpy as np <EOL> from scipy . sparse . linalg import gmres , LinearOperator <EOL> from openmdao . main . mpiwrap import MPI , PETSc , get_norm <EOL> from openmdao . util . graph import fix_single_tuple <EOL> from openmdao . util . log import logger <EOL> class LinearSolver ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , system ) : <EOL> """<STR_LIT>""" <EOL> self . _system = system <EOL> self . options = system . options <EOL> self . custom_jacs = { } <EOL> level = <NUM_LIT:0> <EOL> if hasattr ( system , '<STR_LIT>' ) and system . _parent_system is not None and hasattr ( system . _parent_system , '<STR_LIT>' ) : <EOL> drv = system . _parent_system . _comp <EOL> self . drv_name = drv . name <EOL> if drv . itername == '<STR_LIT>' : <EOL> level = <NUM_LIT:0> <EOL> else : <EOL> level = drv . itername . count ( '<STR_LIT:.>' ) + <NUM_LIT:1> <EOL> if hasattr ( drv , '<STR_LIT>' ) : <EOL> for constraint in drv . get_constraints ( ) . values ( ) : <EOL> if constraint . jacs is not None : <EOL> key = '<STR_LIT>' % constraint . pcomp_name <EOL> self . custom_jacs [ key ] = constraint . jacs <EOL> if hasattr ( drv , '<STR_LIT>' ) : <EOL> for constraint in drv . get_2sided_constraints ( ) . values ( ) : <EOL> if constraint . jacs is not None : <EOL> key = '<STR_LIT>' % constraint . pcomp_name <EOL> self . custom_jacs [ key ] = constraint . jacs <EOL> else : <EOL> self . drv_name = system . name <EOL> self . indent = '<STR_LIT:U+0020>' * level <EOL> def print_norm ( self , driver_string , iteration , res , res0 , msg = None , solver = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> if msg is not None : <EOL> form = self . indent + '<STR_LIT>' <EOL> print form % ( self . drv_name , solver , self . ln_string , iteration , msg ) <EOL> return <EOL> form = self . indent + '<STR_LIT>' <EOL> print form % ( self . drv_name , solver , self . ln_string , iteration , res , res / res0 ) <EOL> def _norm ( self ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> system . rhs_vec . array [ : ] = <NUM_LIT:0.0> <EOL> system . applyJ ( system . vector_vars . keys ( ) ) <EOL> system . rhs_vec . array [ : ] *= - <NUM_LIT:1.0> <EOL> system . rhs_vec . array [ : ] += system . rhs_buf [ : ] <EOL> return get_norm ( system . rhs_vec ) <EOL> def user_defined_jacobian ( self , con , params , J ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( J , dict ) : <EOL> msg = '<STR_LIT>' <EOL> raise RuntimeError ( msg ) <EOL> jacs = self . custom_jacs [ con ] ( ) <EOL> for param in params : <EOL> if param in jacs : <EOL> J [ con ] [ param ] = jacs [ param ] <EOL> else : <EOL> psize = self . _system . get_size ( param ) <EOL> csize = self . _system . get_size ( con ) <EOL> J [ con ] [ param ] = np . zeros ( ( csize , psize ) ) <EOL> class ScipyGMRES ( LinearSolver ) : <EOL> """<STR_LIT>""" <EOL> ln_string = '<STR_LIT>' <EOL> def __init__ ( self , system ) : <EOL> """<STR_LIT>""" <EOL> super ( ScipyGMRES , self ) . __init__ ( system ) <EOL> n_edge = system . vec [ '<STR_LIT:f>' ] . array . size <EOL> system . rhs_buf = np . zeros ( ( n_edge , ) ) <EOL> system . sol_buf = np . zeros ( ( n_edge , ) ) <EOL> self . A = LinearOperator ( ( n_edge , n_edge ) , <EOL> matvec = self . mult , <EOL> dtype = float ) <EOL> def calc_gradient ( self , inputs , outputs , return_format = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> RHS = system . rhs_buf <EOL> A = self . A <EOL> if return_format == '<STR_LIT>' : <EOL> J = { } <EOL> for okey in outputs : <EOL> J [ okey ] = { } <EOL> for ikey in inputs : <EOL> if isinstance ( ikey , tuple ) : <EOL> ikey = ikey [ <NUM_LIT:0> ] <EOL> J [ okey ] [ ikey ] = None <EOL> else : <EOL> num_input = system . get_size ( inputs ) <EOL> num_output = system . get_size ( outputs ) <EOL> J = np . zeros ( ( num_output , num_input ) ) <EOL> if system . mode == '<STR_LIT>' : <EOL> outputs , inputs = inputs , outputs <EOL> j = <NUM_LIT:0> <EOL> for param in inputs : <EOL> if isinstance ( param , tuple ) : <EOL> param = param [ <NUM_LIT:0> ] <EOL> in_indices = system . vec [ '<STR_LIT:u>' ] . indices ( system , param ) <EOL> jbase = j <EOL> if system . mode == '<STR_LIT>' and param in self . custom_jacs : <EOL> self . user_defined_jacobian ( param , outputs , J ) <EOL> j += len ( in_indices ) <EOL> continue <EOL> for irhs in in_indices : <EOL> RHS [ irhs ] = <NUM_LIT:1.0> <EOL> dx = self . solve ( RHS ) <EOL> RHS [ irhs ] = <NUM_LIT:0.0> <EOL> i = <NUM_LIT:0> <EOL> for item in outputs : <EOL> if isinstance ( item , tuple ) : <EOL> item = item [ <NUM_LIT:0> ] <EOL> out_indices = system . vec [ '<STR_LIT:u>' ] . indices ( system , item ) <EOL> nk = len ( out_indices ) <EOL> if return_format == '<STR_LIT>' : <EOL> if system . mode == '<STR_LIT>' : <EOL> if J [ item ] [ param ] is None : <EOL> J [ item ] [ param ] = np . zeros ( ( nk , len ( in_indices ) ) ) <EOL> J [ item ] [ param ] [ : , j - jbase ] = dx [ out_indices ] <EOL> else : <EOL> if J [ param ] [ item ] is None : <EOL> J [ param ] [ item ] = np . zeros ( ( len ( in_indices ) , nk ) ) <EOL> J [ param ] [ item ] [ j - jbase , : ] = dx [ out_indices ] <EOL> else : <EOL> if system . mode == '<STR_LIT>' : <EOL> J [ i : i + nk , j ] = dx [ out_indices ] <EOL> else : <EOL> J [ j , i : i + nk ] = dx [ out_indices ] <EOL> i += nk <EOL> j += <NUM_LIT:1> <EOL> return J <EOL> def solve ( self , arg ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> options = self . options <EOL> A = self . A <EOL> dx , info = gmres ( A , arg , <EOL> tol = options . atol , <EOL> maxiter = options . maxiter ) <EOL> if info > <NUM_LIT:0> : <EOL> msg = "<STR_LIT>" "<STR_LIT>" <EOL> logger . error ( msg , system . name , info ) <EOL> elif info < <NUM_LIT:0> : <EOL> msg = "<STR_LIT>" <EOL> logger . error ( msg , system . name ) <EOL> return dx <EOL> def mult ( self , arg ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> system . sol_vec . array [ : ] = arg [ : ] <EOL> system . rhs_vec . array [ : ] = <NUM_LIT:0.0> <EOL> system . clear_dp ( ) <EOL> if system . _parent_system : <EOL> vnames = system . _parent_system . _relevant_vars <EOL> else : <EOL> vnames = system . flat_vars . keys ( ) <EOL> system . applyJ ( vnames ) <EOL> return system . rhs_vec . array [ : ] <EOL> class PETSc_KSP ( LinearSolver ) : <EOL> """<STR_LIT>""" <EOL> ln_string = '<STR_LIT>' <EOL> class Monitor ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , ksp ) : <EOL> """<STR_LIT>""" <EOL> self . _ksp = ksp <EOL> self . _norm0 = <NUM_LIT:1.0> <EOL> def __call__ ( self , ksp , counter , norm ) : <EOL> """<STR_LIT>""" <EOL> if counter == <NUM_LIT:0> and norm != <NUM_LIT:0.0> : <EOL> self . _norm0 = norm <EOL> if self . _ksp . options . iprint > <NUM_LIT:0> : <EOL> self . _ksp . print_norm ( self . _ksp . ln_string , counter , norm , self . _norm0 ) <EOL> def __init__ ( self , system ) : <EOL> """<STR_LIT>""" <EOL> super ( PETSc_KSP , self ) . __init__ ( system ) <EOL> lsize = np . sum ( system . local_var_sizes [ system . mpi . rank , : ] ) <EOL> size = np . sum ( system . local_var_sizes ) <EOL> jac_mat = PETSc . Mat ( ) . createPython ( [ ( lsize , size ) , ( lsize , size ) ] , <EOL> comm = system . mpi . comm ) <EOL> jac_mat . setPythonContext ( self ) <EOL> jac_mat . setUp ( ) <EOL> self . ksp = PETSc . KSP ( ) . create ( comm = system . mpi . comm ) <EOL> self . ksp . setOperators ( jac_mat ) <EOL> self . ksp . setType ( '<STR_LIT>' ) <EOL> self . ksp . setGMRESRestart ( <NUM_LIT:1000> ) <EOL> self . ksp . setPCSide ( PETSc . PC . Side . RIGHT ) <EOL> self . ksp . setMonitor ( self . Monitor ( self ) ) <EOL> pc_mat = self . ksp . getPC ( ) <EOL> pc_mat . setType ( '<STR_LIT>' ) <EOL> pc_mat . setPythonContext ( self ) <EOL> system . rhs_buf = np . zeros ( ( lsize , ) ) <EOL> system . sol_buf = np . zeros ( ( lsize , ) ) <EOL> system . sol_buf_petsc = PETSc . Vec ( ) . createWithArray ( system . sol_buf , <EOL> comm = system . mpi . comm ) <EOL> system . rhs_buf_petsc = PETSc . Vec ( ) . createWithArray ( system . rhs_buf , <EOL> comm = system . mpi . comm ) <EOL> def calc_gradient ( self , inputs , outputs , return_format = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> options = self . options <EOL> name2collapsed = system . scope . name2collapsed <EOL> inputs = [ fix_single_tuple ( x ) for x in inputs ] <EOL> outputs = [ fix_single_tuple ( x ) for x in outputs ] <EOL> if return_format == '<STR_LIT>' : <EOL> J = { } <EOL> for okey in outputs : <EOL> J [ okey ] = { } <EOL> for ikey in inputs : <EOL> if isinstance ( ikey , tuple ) : <EOL> ikey = ikey [ <NUM_LIT:0> ] <EOL> J [ okey ] [ ikey ] = None <EOL> else : <EOL> num_input = system . get_size ( inputs ) <EOL> num_output = system . get_size ( outputs ) <EOL> J = np . zeros ( ( num_output , num_input ) ) <EOL> if system . mode == '<STR_LIT>' : <EOL> outputs , inputs = inputs , outputs <EOL> self . ksp . setTolerances ( max_it = options . maxiter , <EOL> atol = options . atol , <EOL> rtol = options . rtol ) <EOL> j = <NUM_LIT:0> <EOL> for param in inputs : <EOL> if isinstance ( param , tuple ) : <EOL> param = param [ <NUM_LIT:0> ] <EOL> param_tup = name2collapsed [ param ] <EOL> param_size = system . get_size ( param ) <EOL> jbase = j <EOL> if system . mode == '<STR_LIT>' and param in self . custom_jacs : <EOL> self . user_defined_jacobian ( param , outputs , J ) <EOL> j += param_size <EOL> continue <EOL> for irhs in xrange ( param_size ) : <EOL> solvec = system . _compute_derivatives ( param_tup , irhs ) <EOL> i = <NUM_LIT:0> <EOL> for out in outputs : <EOL> if isinstance ( out , tuple ) : <EOL> out = out [ <NUM_LIT:0> ] <EOL> out_size = system . get_size ( out ) <EOL> if return_format == '<STR_LIT>' : <EOL> if system . mode == '<STR_LIT>' : <EOL> if out in solvec : <EOL> if J [ out ] [ param ] is None : <EOL> J [ out ] [ param ] = np . zeros ( ( out_size , param_size ) ) <EOL> J [ out ] [ param ] [ : , j - jbase ] = solvec [ out ] <EOL> else : <EOL> del J [ out ] [ param ] <EOL> else : <EOL> if out in solvec : <EOL> if J [ param ] [ out ] is None : <EOL> J [ param ] [ out ] = np . zeros ( ( param_size , out_size ) ) <EOL> J [ param ] [ out ] [ j - jbase , : ] = solvec [ out ] <EOL> else : <EOL> del J [ param ] [ out ] <EOL> else : <EOL> if out in solvec : <EOL> nk = len ( solvec [ out ] ) <EOL> if system . mode == '<STR_LIT>' : <EOL> J [ i : i + nk , j ] = solvec [ out ] <EOL> else : <EOL> J [ j , i : i + nk ] = solvec [ out ] <EOL> i += nk <EOL> j += <NUM_LIT:1> <EOL> return J <EOL> def solve ( self , arg ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> options = self . options <EOL> self . ksp . setTolerances ( max_it = options . maxiter , <EOL> atol = options . atol , <EOL> rtol = options . rtol ) <EOL> system . rhs_buf [ : ] = arg [ : ] <EOL> self . ksp . solve ( system . rhs_buf_petsc , system . sol_buf_petsc ) <EOL> return system . sol_buf [ : ] <EOL> def mult ( self , mat , sol_vec , rhs_vec ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> system . sol_vec . array [ : ] = sol_vec . array [ : ] <EOL> system . rhs_vec . array [ : ] = <NUM_LIT:0.0> <EOL> system . clear_dp ( ) <EOL> if system . _parent_system : <EOL> vnames = system . _parent_system . _relevant_vars <EOL> else : <EOL> vnames = system . flat_vars . keys ( ) <EOL> system . applyJ ( vnames ) <EOL> rhs_vec . array [ : ] = system . rhs_vec . array [ : ] <EOL> def apply ( self , mat , sol_vec , rhs_vec ) : <EOL> """<STR_LIT>""" <EOL> rhs_vec . array [ : ] = sol_vec . array [ : ] <EOL> class LinearGS ( LinearSolver ) : <EOL> """<STR_LIT>""" <EOL> ln_string = '<STR_LIT>' <EOL> def __init__ ( self , system ) : <EOL> """<STR_LIT>""" <EOL> super ( LinearGS , self ) . __init__ ( system ) <EOL> lsize = np . sum ( system . local_var_sizes [ system . mpi . rank , : ] ) <EOL> system . sol_buf = np . zeros ( lsize ) <EOL> system . rhs_buf = np . zeros ( lsize ) <EOL> def calc_gradient ( self , inputs , outputs , return_format = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> num_input = system . get_size ( inputs ) <EOL> num_output = system . get_size ( outputs ) <EOL> if return_format == '<STR_LIT>' : <EOL> J = { } <EOL> for okey in outputs : <EOL> J [ okey ] = { } <EOL> for ikey in inputs : <EOL> if isinstance ( ikey , tuple ) : <EOL> ikey = ikey [ <NUM_LIT:0> ] <EOL> J [ okey ] [ ikey ] = None <EOL> else : <EOL> J = np . zeros ( ( num_output , num_input ) ) <EOL> if system . mode == '<STR_LIT>' : <EOL> outputs , inputs = inputs , outputs <EOL> j = <NUM_LIT:0> <EOL> for param in inputs : <EOL> if isinstance ( param , tuple ) : <EOL> param = param [ <NUM_LIT:0> ] <EOL> in_indices = system . rhs_vec . indices ( system , param ) <EOL> nj = len ( in_indices ) <EOL> jbase = j <EOL> if system . mode == '<STR_LIT>' and param in self . custom_jacs : <EOL> self . user_defined_jacobian ( param , outputs , J ) <EOL> j += nj <EOL> continue <EOL> for irhs in in_indices : <EOL> system . clear_dp ( ) <EOL> system . sol_vec . array [ : ] = <NUM_LIT:0.0> <EOL> system . rhs_vec . array [ : ] = <NUM_LIT:0.0> <EOL> system . rhs_vec . array [ irhs ] = <NUM_LIT:1.0> <EOL> dx = self . solve ( system . rhs_vec . array ) <EOL> i = <NUM_LIT:0> <EOL> for item in outputs : <EOL> if isinstance ( item , tuple ) : <EOL> item = item [ <NUM_LIT:0> ] <EOL> out_indices = system . sol_vec . indices ( system , item ) <EOL> nk = len ( out_indices ) <EOL> if return_format == '<STR_LIT>' : <EOL> if system . mode == '<STR_LIT>' : <EOL> if J [ item ] [ param ] is None : <EOL> J [ item ] [ param ] = np . zeros ( ( nk , nj ) ) <EOL> J [ item ] [ param ] [ : , j - jbase ] = dx [ out_indices ] <EOL> else : <EOL> if J [ param ] [ item ] is None : <EOL> J [ param ] [ item ] = np . zeros ( ( nj , nk ) ) <EOL> J [ param ] [ item ] [ j - jbase , : ] = dx [ out_indices ] <EOL> else : <EOL> if system . mode == '<STR_LIT>' : <EOL> J [ i : i + nk , j ] = dx [ out_indices ] <EOL> else : <EOL> J [ j , i : i + nk ] = dx [ out_indices ] <EOL> i += nk <EOL> j += <NUM_LIT:1> <EOL> return J <EOL> def solve ( self , arg ) : <EOL> """<STR_LIT>""" <EOL> system = self . _system <EOL> system . rhs_buf [ : ] = arg [ : ] <EOL> system . sol_buf [ : ] = system . sol_vec . array [ : ] <EOL> options = self . options <EOL> system = self . _system <EOL> norm0 , norm = <NUM_LIT:1.0> , <NUM_LIT:1.0> <EOL> counter = <NUM_LIT:0> <EOL> if self . options . iprint > <NUM_LIT:0> : <EOL> self . print_norm ( self . ln_string , counter , norm , norm0 ) <EOL> while counter < options . maxiter and norm > options . atol and norm / norm0 > options . rtol : <EOL> if system . mode == '<STR_LIT>' : <EOL> for subsystem in system . subsystems ( local = True ) : <EOL> system . scatter ( '<STR_LIT>' , '<STR_LIT>' , subsystem = subsystem ) <EOL> system . rhs_vec . array [ : ] = <NUM_LIT:0.0> <EOL> subsystem . applyJ ( system . flat_vars . keys ( ) ) <EOL> system . rhs_vec . array [ : ] *= - <NUM_LIT:1.0> <EOL> system . rhs_vec . array [ : ] += system . rhs_buf [ : ] <EOL> sub_options = options if subsystem . options is None else subsystem . options <EOL> subsystem . solve_linear ( sub_options ) <EOL> elif system . mode == '<STR_LIT>' : <EOL> rev_systems = [ item for item in reversed ( list ( system . subsystems ( local = True ) ) ) ] <EOL> for subsystem in rev_systems : <EOL> system . sol_buf [ : ] = system . rhs_buf [ : ] <EOL> succs = [ str ( node ) for node in system . graph . successors ( subsystem . node ) ] <EOL> for subsystem2 in rev_systems : <EOL> if subsystem2 . name in succs : <EOL> system . rhs_vec . array [ : ] = <NUM_LIT:0.0> <EOL> args = subsystem . flat_vars . keys ( ) <EOL> subsystem2 . applyJ ( args ) <EOL> system . scatter ( '<STR_LIT>' , '<STR_LIT>' , subsystem = subsystem2 ) <EOL> system . sol_buf [ : ] -= system . rhs_vec . array [ : ] <EOL> system . vec [ '<STR_LIT>' ] . array [ : ] = <NUM_LIT:0.0> <EOL> system . rhs_vec . array [ : ] = system . sol_buf [ : ] <EOL> subsystem . solve_linear ( options ) <EOL> norm = self . _norm ( ) <EOL> counter += <NUM_LIT:1> <EOL> if self . options . iprint > <NUM_LIT:0> : <EOL> self . print_norm ( self . ln_string , counter , norm , norm0 ) <EOL> if len ( system . vec [ '<STR_LIT>' ] . array ) == <NUM_LIT:0> : <EOL> psys = system . _parent_system <EOL> return system . sol_vec . array </s>
