<s> """<STR_LIT>""" <EOL> from __future__ import print_function <EOL> import sys <EOL> import os <EOL> import re <EOL> import warnings <EOL> from collections import OrderedDict <EOL> from itertools import chain <EOL> from six import iteritems , itervalues , iterkeys <EOL> import numpy as np <EOL> from openmdao . core . basic_impl import BasicImpl <EOL> from openmdao . core . system import System <EOL> from openmdao . core . mpi_wrap import MPI <EOL> from openmdao . core . vec_wrapper import _ByObjWrapper <EOL> from openmdao . core . vec_wrapper_complex_step import ComplexStepSrcVecWrapper , ComplexStepTgtVecWrapper <EOL> from openmdao . core . fileref import FileRef <EOL> from openmdao . util . type_util import is_differentiable <EOL> _NotSet = object ( ) <EOL> namecheck_rgx = re . compile ( <EOL> '<STR_LIT>' ) <EOL> trace = os . environ . get ( '<STR_LIT>' ) <EOL> if trace : <EOL> from openmdao . core . mpi_wrap import debug <EOL> empty_arr = np . zeros ( <NUM_LIT:0> ) <EOL> class Component ( System ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> super ( Component , self ) . __init__ ( ) <EOL> self . _post_setup_vars = False <EOL> self . _jacobian_cache = { } <EOL> self . _init_params_dict = OrderedDict ( ) <EOL> self . _init_unknowns_dict = OrderedDict ( ) <EOL> self . _pbo_warns = [ ] <EOL> self . _run_apply = False <EOL> def _get_initial_val ( self , val , shape ) : <EOL> """<STR_LIT>""" <EOL> if val is _NotSet : <EOL> if shape == <NUM_LIT:1> : <EOL> return <NUM_LIT:0.> <EOL> return np . zeros ( shape ) <EOL> return val <EOL> def _check_val ( self , name , var_type , val , shape ) : <EOL> """<STR_LIT>""" <EOL> if val is _NotSet and shape is None : <EOL> msg = ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> msg = msg . format ( var_type = var_type , name = name ) <EOL> raise ValueError ( msg ) <EOL> def _add_variable ( self , name , val , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> shape = kwargs . get ( '<STR_LIT>' ) <EOL> self . _check_varname ( name ) <EOL> meta = kwargs . copy ( ) <EOL> if isinstance ( val , FileRef ) : <EOL> val . _set_meta ( kwargs ) <EOL> meta [ '<STR_LIT>' ] = val = self . _get_initial_val ( val , shape ) <EOL> if is_differentiable ( val ) and not meta . get ( '<STR_LIT>' ) : <EOL> if isinstance ( val , np . ndarray ) : <EOL> meta [ '<STR_LIT:size>' ] = val . size <EOL> meta [ '<STR_LIT>' ] = val . shape <EOL> else : <EOL> meta [ '<STR_LIT:size>' ] = <NUM_LIT:1> <EOL> meta [ '<STR_LIT>' ] = <NUM_LIT:1> <EOL> else : <EOL> if not meta . get ( '<STR_LIT>' ) : <EOL> self . _pbo_warns . append ( ( name , val ) ) <EOL> meta [ '<STR_LIT:size>' ] = <NUM_LIT:0> <EOL> meta [ '<STR_LIT>' ] = True <EOL> if isinstance ( shape , int ) and shape > <NUM_LIT:1> : <EOL> meta [ '<STR_LIT>' ] = ( shape , ) <EOL> return meta <EOL> def add_param ( self , name , val = _NotSet , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> self . _init_params_dict [ name ] = self . _add_variable ( name , val , ** kwargs ) <EOL> def add_output ( self , name , val = _NotSet , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> shape = kwargs . get ( '<STR_LIT>' ) <EOL> self . _check_val ( name , '<STR_LIT>' , val , shape ) <EOL> self . _init_unknowns_dict [ name ] = self . _add_variable ( name , val , ** kwargs ) <EOL> def add_state ( self , name , val = _NotSet , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> shape = kwargs . get ( '<STR_LIT>' ) <EOL> self . _check_val ( name , '<STR_LIT:state>' , val , shape ) <EOL> args = self . _add_variable ( name , val , ** kwargs ) <EOL> args [ '<STR_LIT:state>' ] = True <EOL> self . _init_unknowns_dict [ name ] = args <EOL> self . _run_apply = True <EOL> def set_var_indices ( self , name , val = _NotSet , shape = None , <EOL> src_indices = None ) : <EOL> """<STR_LIT>""" <EOL> meta = self . _init_params_dict . get ( name ) <EOL> if meta is None : <EOL> meta = self . _init_unknowns_dict [ name ] <EOL> if src_indices is None : <EOL> raise ValueError ( "<STR_LIT>" % <EOL> name ) <EOL> if not isinstance ( meta [ '<STR_LIT>' ] , np . ndarray ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( name , type ( meta [ '<STR_LIT>' ] ) ) ) <EOL> if val is _NotSet : <EOL> if shape : <EOL> val = np . zeros ( shape , dtype = meta [ '<STR_LIT>' ] . dtype ) <EOL> else : <EOL> val = np . zeros ( ( len ( src_indices ) , ) , dtype = meta [ '<STR_LIT>' ] . dtype ) <EOL> if val . size != len ( src_indices ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( val . size , name , len ( src_indices ) ) ) <EOL> meta [ '<STR_LIT>' ] = val <EOL> meta [ '<STR_LIT>' ] = val . shape <EOL> meta [ '<STR_LIT:size>' ] = val . size <EOL> meta [ '<STR_LIT>' ] = src_indices <EOL> def _check_varname ( self , name ) : <EOL> """<STR_LIT>""" <EOL> if self . _post_setup_vars : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> ( self . pathname , name ) ) <EOL> if name in self . _init_params_dict or name in self . _init_unknowns_dict : <EOL> raise RuntimeError ( "<STR_LIT>" % <EOL> ( self . pathname , name ) ) <EOL> match = namecheck_rgx . match ( name ) <EOL> if match is None or match . group ( ) != name : <EOL> raise NameError ( "<STR_LIT>" % <EOL> ( self . pathname , name ) ) <EOL> def setup_distrib ( self ) : <EOL> """<STR_LIT>""" <EOL> pass <EOL> def _get_fd_params ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _fd_params is None : <EOL> self . _fd_params = [ k for k , acc in iteritems ( self . params . _dat ) <EOL> if not acc . pbo ] <EOL> return self . _fd_params <EOL> def _get_fd_unknowns ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ k for k , acc in iteritems ( self . unknowns . _dat ) <EOL> if not acc . pbo ] <EOL> def _setup_variables ( self , compute_indices = False ) : <EOL> """<STR_LIT>""" <EOL> to_prom_name = self . _sysdata . to_prom_name = { } <EOL> to_abs_uname = self . _sysdata . to_abs_uname = { } <EOL> to_abs_pnames = self . _sysdata . to_abs_pnames = OrderedDict ( ) <EOL> to_prom_uname = self . _sysdata . to_prom_uname = OrderedDict ( ) <EOL> to_prom_pname = self . _sysdata . to_prom_pname = OrderedDict ( ) <EOL> if MPI and compute_indices and self . is_active ( ) : <EOL> if hasattr ( self , '<STR_LIT>' ) : <EOL> warnings . simplefilter ( '<STR_LIT>' , DeprecationWarning ) <EOL> warnings . warn ( "<STR_LIT>" , <EOL> DeprecationWarning , stacklevel = <NUM_LIT:2> ) <EOL> warnings . simplefilter ( '<STR_LIT:ignore>' , DeprecationWarning ) <EOL> self . setup_distrib_idxs ( ) <EOL> else : <EOL> self . setup_distrib ( ) <EOL> sizes = [ ] <EOL> names = [ ] <EOL> for name , meta in iteritems ( self . _init_unknowns_dict ) : <EOL> if '<STR_LIT>' in meta : <EOL> sizes . append ( len ( meta [ '<STR_LIT>' ] ) ) <EOL> names . append ( name ) <EOL> if sizes : <EOL> if trace : <EOL> debug ( "<STR_LIT>" ) <EOL> allsizes = np . zeros ( ( self . comm . size , len ( sizes ) ) , dtype = int ) <EOL> self . comm . Allgather ( np . array ( sizes , dtype = int ) , allsizes ) <EOL> for i , name in enumerate ( names ) : <EOL> self . _init_unknowns_dict [ name ] [ '<STR_LIT>' ] = np . sum ( allsizes [ : , i ] ) <EOL> self . _params_dict = OrderedDict ( ) <EOL> for name , meta in iteritems ( self . _init_params_dict ) : <EOL> pathname = self . _get_var_pathname ( name ) <EOL> self . _params_dict [ pathname ] = meta <EOL> meta [ '<STR_LIT>' ] = pathname <EOL> to_prom_pname [ pathname ] = name <EOL> to_abs_pnames [ name ] = ( pathname , ) <EOL> self . _unknowns_dict = OrderedDict ( ) <EOL> for name , meta in iteritems ( self . _init_unknowns_dict ) : <EOL> pathname = self . _get_var_pathname ( name ) <EOL> self . _unknowns_dict [ pathname ] = meta <EOL> meta [ '<STR_LIT>' ] = pathname <EOL> to_prom_uname [ pathname ] = name <EOL> to_abs_uname [ name ] = pathname <EOL> to_prom_name . update ( to_prom_uname ) <EOL> to_prom_name . update ( to_prom_pname ) <EOL> self . _post_setup_vars = True <EOL> self . _sysdata . _params_dict = self . _params_dict <EOL> self . _sysdata . _unknowns_dict = self . _unknowns_dict <EOL> return self . _params_dict , self . _unknowns_dict <EOL> def _setup_communicators ( self , comm , parent_dir ) : <EOL> """<STR_LIT>""" <EOL> super ( Component , self ) . _setup_communicators ( comm , parent_dir ) <EOL> for meta in chain ( itervalues ( self . _init_unknowns_dict ) , <EOL> itervalues ( self . _init_params_dict ) ) : <EOL> val = meta [ '<STR_LIT>' ] <EOL> if isinstance ( val , FileRef ) : <EOL> self . _fileref_setup ( val ) <EOL> if not self . is_active ( ) : <EOL> for meta in itervalues ( self . _init_params_dict ) : <EOL> meta [ '<STR_LIT>' ] = True <EOL> for meta in itervalues ( self . _init_unknowns_dict ) : <EOL> meta [ '<STR_LIT>' ] = True <EOL> def _fileref_setup ( self , fref ) : <EOL> fref . parent_dir = self . _sysdata . absdir <EOL> d = fref . _abspath ( ) <EOL> if self . is_active ( ) and not os . path . exists ( os . path . dirname ( d ) ) : <EOL> if self . create_dirs : <EOL> os . makedirs ( os . path . dirname ( d ) ) <EOL> else : <EOL> raise RuntimeError ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( os . path . dirname ( d ) , <EOL> fref . fname , self . pathname ) ) <EOL> def _setup_vectors ( self , param_owners , parent , <EOL> top_unknowns = None , impl = None , alloc_derivs = True ) : <EOL> """<STR_LIT>""" <EOL> self . params = self . unknowns = self . resids = None <EOL> self . dumat , self . dpmat , self . drmat = OrderedDict ( ) , OrderedDict ( ) , OrderedDict ( ) <EOL> self . connections = self . _probdata . connections <EOL> relevance = self . _probdata . relevance <EOL> if not self . is_active ( ) : <EOL> return <EOL> self . _setup_prom_map ( ) <EOL> self . _impl = impl <EOL> self . _relname_map = self . _get_relname_map ( parent . _sysdata . to_prom_name ) <EOL> self . _shared_dp_vec = empty_arr <EOL> self . _shared_p_offsets = { None : <NUM_LIT:0> } <EOL> for vois in chain ( relevance . inputs , relevance . outputs ) : <EOL> for voi in vois : <EOL> self . _shared_p_offsets [ voi ] = <NUM_LIT:0> <EOL> self . _create_views ( top_unknowns , parent , [ ] , None ) <EOL> all_vois = set ( [ None ] ) <EOL> if self . _probdata . top_lin_gs : <EOL> for vois in relevance . groups : <EOL> all_vois . update ( vois ) <EOL> for voi in vois : <EOL> self . _create_views ( top_unknowns , parent , [ ] , voi ) <EOL> for meta in itervalues ( self . _params_dict ) : <EOL> pathname = meta [ '<STR_LIT>' ] <EOL> name = self . _sysdata . _scoped_abs_name ( pathname ) <EOL> if name not in self . params : <EOL> self . params . _add_unconnected_var ( pathname , meta ) <EOL> def apply_nonlinear ( self , params , unknowns , resids ) : <EOL> """<STR_LIT>""" <EOL> resids . vec [ : ] = - unknowns . vec <EOL> self . solve_nonlinear ( params , unknowns , resids ) <EOL> resids . vec [ : ] += unknowns . vec <EOL> unknowns . vec [ : ] -= resids . vec <EOL> def solve_nonlinear ( self , params , unknowns , resids ) : <EOL> """<STR_LIT>""" <EOL> msg = "<STR_LIT>" <EOL> raise NotImplementedError ( msg % self . __class__ . __name__ ) <EOL> def linearize ( self , params , unknowns , resids ) : <EOL> """<STR_LIT>""" <EOL> return None <EOL> def apply_linear ( self , params , unknowns , dparams , dunknowns , dresids , mode ) : <EOL> """<STR_LIT>""" <EOL> self . _apply_linear_jac ( params , unknowns , dparams , dunknowns , dresids , <EOL> mode ) <EOL> def solve_linear ( self , dumat , drmat , vois , mode = None ) : <EOL> """<STR_LIT>""" <EOL> if mode == '<STR_LIT>' : <EOL> sol_vec , rhs_vec = self . dumat , self . drmat <EOL> else : <EOL> sol_vec , rhs_vec = self . drmat , self . dumat <EOL> for voi in vois : <EOL> sol_vec [ voi ] . vec [ : ] = - rhs_vec [ voi ] . vec <EOL> def dump ( self , nest = <NUM_LIT:0> , out_stream = sys . stdout , verbose = False , dvecs = False , <EOL> sizes = False ) : <EOL> """<STR_LIT>""" <EOL> klass = self . __class__ . __name__ <EOL> if dvecs : <EOL> ulabel , plabel , uvecname , pvecname = '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' <EOL> else : <EOL> ulabel , plabel , uvecname , pvecname = '<STR_LIT:u>' , '<STR_LIT:p>' , '<STR_LIT>' , '<STR_LIT>' <EOL> uvec = getattr ( self , uvecname ) <EOL> pvec = getattr ( self , pvecname ) <EOL> template = "<STR_LIT>" <EOL> out_stream . write ( template % ( "<STR_LIT:U+0020>" * nest , klass , self . name ) ) <EOL> if sizes : <EOL> commsz = self . comm . size if hasattr ( self . comm , '<STR_LIT:size>' ) else <NUM_LIT:0> <EOL> template = "<STR_LIT>" <EOL> out_stream . write ( template % ( self . get_req_procs ( ) , <EOL> uvec . vec . size , <EOL> pvec . vec . size , <EOL> commsz ) ) <EOL> out_stream . write ( "<STR_LIT:\n>" ) <EOL> if verbose : <EOL> lens = [ len ( n ) for n in uvec ] <EOL> nwid = max ( lens ) if lens else <NUM_LIT:12> <EOL> for v in uvec : <EOL> if v in uvec . _dat and uvec . _dat [ v ] . slice is not None : <EOL> uslice = '<STR_LIT>' . format ( ulabel , <EOL> uvec . _dat [ v ] . slice ) <EOL> tem = "<STR_LIT>" <EOL> out_stream . write ( tem . format ( "<STR_LIT:U+0020>" * ( nest + <NUM_LIT:8> ) , v , uslice , <EOL> repr ( uvec [ v ] ) , nwid = nwid ) ) <EOL> elif not dvecs : <EOL> tem = "<STR_LIT>" <EOL> out_stream . write ( tem . format ( "<STR_LIT:U+0020>" * ( nest + <NUM_LIT:8> ) , v , repr ( uvec [ v ] ) , <EOL> nwid = nwid ) ) <EOL> out_stream . flush ( ) <EOL> def _get_relname_map ( self , parent_proms ) : <EOL> """<STR_LIT>""" <EOL> umap = OrderedDict ( ) <EOL> for key , meta in iteritems ( self . _init_unknowns_dict ) : <EOL> umap [ parent_proms [ '<STR_LIT:.>' . join ( ( self . pathname , key ) ) ] ] = key <EOL> return umap <EOL> def complex_step_jacobian ( self , params , unknowns , resids , total_derivs = False , <EOL> fd_params = None , fd_states = None , fd_unknowns = None , <EOL> poi_indices = None , qoi_indices = None ) : <EOL> """<STR_LIT>""" <EOL> if fd_params is None : <EOL> fd_params = self . _get_fd_params ( ) <EOL> if fd_unknowns is None : <EOL> fd_unknowns = self . _get_fd_unknowns ( ) <EOL> step_size = self . fd_options . get ( '<STR_LIT>' , <NUM_LIT> ) <EOL> jac = { } <EOL> csparams = ComplexStepTgtVecWrapper ( params ) <EOL> csunknowns = ComplexStepSrcVecWrapper ( unknowns ) <EOL> csresids = ComplexStepSrcVecWrapper ( resids ) <EOL> states = self . states <EOL> if len ( states ) > <NUM_LIT:0> : <EOL> resultvec = csresids <EOL> else : <EOL> resultvec = csunknowns <EOL> if fd_states is not None : <EOL> states = fd_states <EOL> for p_name in chain ( fd_params , states ) : <EOL> if p_name in states : <EOL> stepvec = csunknowns <EOL> target_input = unknowns . _dat [ p_name ] . val <EOL> else : <EOL> stepvec = csparams <EOL> target_input = params . _dat [ p_name ] . val <EOL> stepvec . set_complex_var ( p_name ) <EOL> mydict = self . _init_params_dict . get ( p_name , { } ) <EOL> fdstep = mydict . get ( '<STR_LIT>' , step_size ) <EOL> p_size = np . size ( target_input ) <EOL> p_idxs = range ( p_size ) <EOL> for u_name in fd_unknowns : <EOL> u_size = np . size ( unknowns [ u_name ] ) <EOL> jac [ u_name , p_name ] = np . zeros ( ( u_size , p_size ) ) <EOL> for j , idx in enumerate ( p_idxs ) : <EOL> stepvec . step_complex ( idx , fdstep ) <EOL> self . apply_nonlinear ( csparams , csunknowns , csresids ) <EOL> stepvec . step_complex ( idx , - fdstep ) <EOL> for u_name in fd_unknowns : <EOL> result = resultvec . flat ( u_name ) <EOL> jac [ u_name , p_name ] [ : , j ] = result . imag / fdstep <EOL> stepvec . set_complex_var ( None ) <EOL> return jac <EOL> def alloc_jacobian ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . _jacobian_cache is not None and len ( self . _jacobian_cache ) > <NUM_LIT:0> : <EOL> return self . _jacobian_cache <EOL> self . _jacobian_cache = jac = { } <EOL> u_vec = self . unknowns <EOL> p_vec = self . params <EOL> states = self . states <EOL> p_size_storage = [ ( n , m [ '<STR_LIT:size>' ] ) for n , m in iteritems ( p_vec ) <EOL> if not m . get ( '<STR_LIT>' ) and not m . get ( '<STR_LIT>' ) ] <EOL> s_size_storage = [ ] <EOL> u_size_storage = [ ] <EOL> for n , meta in iteritems ( u_vec ) : <EOL> if meta . get ( '<STR_LIT>' ) or meta . get ( '<STR_LIT>' ) : <EOL> continue <EOL> if meta . get ( '<STR_LIT:state>' ) : <EOL> s_size_storage . append ( ( n , meta [ '<STR_LIT:size>' ] ) ) <EOL> u_size_storage . append ( ( n , meta [ '<STR_LIT:size>' ] ) ) <EOL> for u_var , u_size in u_size_storage : <EOL> for p_var , p_size in p_size_storage : <EOL> jac [ u_var , p_var ] = np . zeros ( ( u_size , p_size ) ) <EOL> for s_var , s_size in s_size_storage : <EOL> jac [ u_var , s_var ] = np . zeros ( ( u_size , s_size ) ) <EOL> return jac </s>
