<s> import numpy as np <EOL> from openmdao . api import Component , Problem , Group , IndepVarComp <EOL> from openmdao . core . mpi_wrap import MPI <EOL> from openmdao . util . array_util import evenly_distrib_idxs <EOL> from openmdao . test . mpi_util import MPITestCase <EOL> if MPI : <EOL> from openmdao . core . petsc_impl import PetscImpl as impl <EOL> else : <EOL> from openmdao . api import BasicImpl as impl <EOL> from openmdao . test . util import assert_rel_error <EOL> class DistributedAdder ( Component ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , size = <NUM_LIT:100> ) : <EOL> super ( DistributedAdder , self ) . __init__ ( ) <EOL> self . local_size = self . size = int ( size ) <EOL> self . add_param ( '<STR_LIT:x>' , shape = size ) <EOL> self . add_output ( '<STR_LIT:y>' , shape = size ) <EOL> def get_req_procs ( self ) : <EOL> """<STR_LIT>""" <EOL> return ( <NUM_LIT:1> , self . size ) <EOL> def setup_distrib ( self ) : <EOL> """<STR_LIT>""" <EOL> comm = self . comm <EOL> rank = comm . rank <EOL> sizes , offsets = evenly_distrib_idxs ( comm . size , self . size ) <EOL> local_size , local_offset = sizes [ rank ] , offsets [ rank ] <EOL> self . local_size = int ( local_size ) <EOL> start = local_offset <EOL> end = local_offset + local_size <EOL> self . set_var_indices ( '<STR_LIT:x>' , val = np . zeros ( local_size , float ) , <EOL> src_indices = np . arange ( start , end , dtype = int ) ) <EOL> self . set_var_indices ( '<STR_LIT:y>' , val = np . zeros ( local_size , float ) , <EOL> src_indices = np . arange ( start , end , dtype = int ) ) <EOL> def solve_nonlinear ( self , params , unknowns , resids ) : <EOL> unknowns [ '<STR_LIT:y>' ] = params [ '<STR_LIT:x>' ] + <NUM_LIT:10> <EOL> class Summer ( Component ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , size = <NUM_LIT:100> ) : <EOL> super ( Summer , self ) . __init__ ( ) <EOL> self . add_param ( '<STR_LIT:y>' , val = np . zeros ( size ) ) <EOL> self . add_output ( '<STR_LIT>' , shape = <NUM_LIT:1> ) <EOL> def solve_nonlinear ( self , params , unknowns , resids ) : <EOL> unknowns [ '<STR_LIT>' ] = np . sum ( params [ '<STR_LIT:y>' ] ) <EOL> class DistributedAdderTest ( MPITestCase ) : <EOL> N_PROCS = <NUM_LIT:10> <EOL> def test_distributed_adder ( self ) : <EOL> size = <NUM_LIT> <EOL> prob = Problem ( impl = impl ) <EOL> prob . root = Group ( ) <EOL> prob . root . add ( '<STR_LIT>' , IndepVarComp ( '<STR_LIT:x>' , np . ones ( size ) ) , promotes = [ '<STR_LIT:x>' ] ) <EOL> prob . root . add ( '<STR_LIT>' , DistributedAdder ( size ) , promotes = [ '<STR_LIT:x>' , '<STR_LIT:y>' ] ) <EOL> prob . root . add ( '<STR_LIT>' , Summer ( size ) , promotes = [ '<STR_LIT:y>' , '<STR_LIT>' ] ) <EOL> prob . setup ( check = False ) <EOL> prob [ '<STR_LIT:x>' ] = np . ones ( size ) <EOL> prob . run ( ) <EOL> assert_rel_error ( self , prob [ '<STR_LIT>' ] / size , <NUM_LIT> , <NUM_LIT> ) <EOL> if __name__ == '<STR_LIT:__main__>' : <EOL> from openmdao . test . mpi_util import mpirun_tests <EOL> mpirun_tests ( ) </s>
