<s> import os <EOL> import sys <EOL> import numpy <EOL> from contextlib import contextmanager <EOL> def _redirect_streams ( to_fd ) : <EOL> """<STR_LIT>""" <EOL> original_stdout_fd = sys . stdout . fileno ( ) <EOL> original_stderr_fd = sys . stderr . fileno ( ) <EOL> sys . stdout . close ( ) <EOL> sys . stderr . close ( ) <EOL> os . dup2 ( to_fd , original_stdout_fd ) <EOL> os . dup2 ( to_fd , original_stderr_fd ) <EOL> sys . stdout = os . fdopen ( original_stdout_fd , '<STR_LIT:wb>' , <NUM_LIT:0> ) <EOL> sys . stderr = os . fdopen ( original_stderr_fd , '<STR_LIT:wb>' , <NUM_LIT:0> ) <EOL> def use_proc_files ( ) : <EOL> if MPI is not None : <EOL> rank = MPI . COMM_WORLD . rank <EOL> sname = "<STR_LIT>" % rank <EOL> ofile = open ( sname , '<STR_LIT:wb>' ) <EOL> _redirect_streams ( ofile . fileno ( ) ) <EOL> def under_mpirun ( ) : <EOL> """<STR_LIT>""" <EOL> for name in os . environ . keys ( ) : <EOL> if name . startswith ( '<STR_LIT>' ) or name . startswith ( '<STR_LIT>' ) : <EOL> return True <EOL> return False <EOL> class PETSc ( object ) : <EOL> def __init__ ( self ) : <EOL> self . needs_ksp = False <EOL> self . _PETSc = None <EOL> @ property <EOL> def installed ( self ) : <EOL> try : <EOL> if self . _PETSc is None : <EOL> PETSc = _import_petsc ( ) <EOL> del sys . modules [ '<STR_LIT>' ] <EOL> self . _PETSc = PETSc <EOL> return True <EOL> except ImportError : <EOL> self . _PETSc = None <EOL> return False <EOL> def __getattr__ ( self , name ) : <EOL> if self . installed : <EOL> return getattr ( self . _PETSc , name ) <EOL> raise AttributeError ( name ) <EOL> def create_petsc_vec ( comm , arr ) : <EOL> if under_mpirun ( ) or PETSc . needs_ksp : <EOL> if PETSc . installed and ( MPI is None or comm != MPI . COMM_NULL ) : <EOL> return PETSc . Vec ( ) . createWithArray ( arr , comm = comm ) <EOL> return None <EOL> def _import_petsc ( ) : <EOL> import petsc4py <EOL> from petsc4py import PETSc <EOL> return PETSc <EOL> if under_mpirun ( ) : <EOL> from mpi4py import MPI <EOL> PETSc = _import_petsc ( ) <EOL> PETSc . installed = True <EOL> COMM_NULL = MPI . COMM_NULL <EOL> else : <EOL> MPI = None <EOL> COMM_NULL = None <EOL> PETSc = PETSc ( ) <EOL> class MPI_info ( object ) : <EOL> def __init__ ( self ) : <EOL> self . requested_cpus = ( <NUM_LIT:1> , <NUM_LIT:1> ) <EOL> self . comm = COMM_NULL <EOL> @ property <EOL> def size ( self ) : <EOL> if MPI and self . comm != COMM_NULL : <EOL> return self . comm . size <EOL> return <NUM_LIT:1> <EOL> @ property <EOL> def rank ( self ) : <EOL> if MPI : <EOL> if self . comm != COMM_NULL : <EOL> return self . comm . rank <EOL> else : <EOL> return - <NUM_LIT:1> <EOL> return <NUM_LIT:0> <EOL> def get_norm ( vec , order = None ) : <EOL> """<STR_LIT>""" <EOL> if MPI : <EOL> vec . petsc_vec . assemble ( ) <EOL> return vec . petsc_vec . norm ( ) <EOL> else : <EOL> return numpy . linalg . norm ( vec . array , ord = order ) <EOL> idx_arr_type = PETSc . IntType if MPI else '<STR_LIT:i>' <EOL> def make_idx_array ( start , end ) : <EOL> """<STR_LIT>""" <EOL> return numpy . arange ( start , end , dtype = idx_arr_type ) <EOL> def to_idx_array ( idxs ) : <EOL> """<STR_LIT>""" <EOL> return numpy . array ( idxs , dtype = idx_arr_type ) <EOL> def evenly_distrib_idxs ( num_divisions , arr_size ) : <EOL> """<STR_LIT>""" <EOL> base = arr_size / num_divisions <EOL> leftover = arr_size % num_divisions <EOL> sizes = numpy . ones ( num_divisions , dtype = "<STR_LIT:int>" ) * base <EOL> sizes [ : leftover ] += <NUM_LIT:1> <EOL> offsets = numpy . zeros ( num_divisions , dtype = "<STR_LIT:int>" ) <EOL> offsets [ <NUM_LIT:1> : ] = numpy . cumsum ( sizes ) [ : - <NUM_LIT:1> ] <EOL> return sizes , offsets <EOL> @ contextmanager <EOL> def MPIContext ( ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> yield <EOL> except : <EOL> exc_type , exc_val , exc_tb = sys . exc_info ( ) <EOL> if exc_val is not None : <EOL> fail = True <EOL> else : <EOL> fail = False <EOL> fails = MPI . COMM_WORLD . allgather ( fail ) <EOL> if fail or not any ( fails ) : <EOL> raise exc_type , exc_val , exc_tb <EOL> else : <EOL> for i , f in enumerate ( fails ) : <EOL> if f : <EOL> raise RuntimeError ( "<STR_LIT>" % i ) <EOL> if os . environ . get ( '<STR_LIT>' ) : <EOL> use_proc_files ( ) </s>
