<s> import unittest <EOL> import random <EOL> import itertools <EOL> from numpy import array , linspace , sin , cos , pi <EOL> from openmdao . api import KrigingSurrogate <EOL> from openmdao . test . util import assert_rel_error <EOL> from six . moves import zip <EOL> def branin ( x ) : <EOL> y = ( x [ <NUM_LIT:1> ] - ( <NUM_LIT> / ( <NUM_LIT> * pi ** <NUM_LIT> ) ) * x [ <NUM_LIT:0> ] ** <NUM_LIT> + <NUM_LIT> * x [ <NUM_LIT:0> ] / pi - <NUM_LIT> ) ** <NUM_LIT> + <NUM_LIT> * ( <NUM_LIT:1.> - <NUM_LIT:1.> / ( <NUM_LIT> * pi ) ) * cos ( <EOL> x [ <NUM_LIT:0> ] ) + <NUM_LIT> <EOL> return y <EOL> def branin_1d ( x ) : <EOL> return branin ( array ( [ x [ <NUM_LIT:0> ] , <NUM_LIT> ] ) ) <EOL> class TestKrigingSurrogate ( unittest . TestCase ) : <EOL> def test_1d_training ( self ) : <EOL> x = array ( [ [ <NUM_LIT:0.0> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] ] ) <EOL> y = array ( [ [ branin_1d ( case ) ] for case in x ] ) <EOL> surrogate = KrigingSurrogate ( ) <EOL> surrogate . train ( x , y ) <EOL> for x0 , y0 in zip ( x , y ) : <EOL> mu , sigma = surrogate . predict ( x0 ) <EOL> assert_rel_error ( self , mu , y0 , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0> , <NUM_LIT> ) <EOL> def test_1d_predictor ( self ) : <EOL> x = array ( [ [ <NUM_LIT:0.0> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] ] ) <EOL> y = array ( [ [ branin_1d ( case ) ] for case in x ] ) <EOL> surrogate = KrigingSurrogate ( ) <EOL> surrogate . train ( x , y ) <EOL> new_x = array ( [ pi ] ) <EOL> mu , sigma = surrogate . predict ( new_x ) <EOL> assert_rel_error ( self , mu , <NUM_LIT> , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT> ) <EOL> def test_1d_ill_conditioned ( self ) : <EOL> x = array ( [ [ case ] for case in linspace ( <NUM_LIT:0.> , <NUM_LIT:1.> , <NUM_LIT> ) ] ) <EOL> y = sin ( x ) <EOL> surrogate = KrigingSurrogate ( ) <EOL> surrogate . train ( x , y ) <EOL> new_x = array ( [ <NUM_LIT:0.5> ] ) <EOL> mu , sigma = surrogate . predict ( new_x ) <EOL> self . assertTrue ( sigma < <NUM_LIT> ) <EOL> assert_rel_error ( self , mu , sin ( <NUM_LIT:0.5> ) , <NUM_LIT> ) <EOL> def test_2d ( self ) : <EOL> x = array ( [ [ - <NUM_LIT> , <NUM_LIT:0.> ] , [ - <NUM_LIT:0.5> , <NUM_LIT> ] , [ <NUM_LIT:1.> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] , [ - <NUM_LIT> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] , [ - <NUM_LIT> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] ] ) <EOL> y = array ( [ [ branin ( case ) ] for case in x ] ) <EOL> surrogate = KrigingSurrogate ( ) <EOL> surrogate . train ( x , y ) <EOL> for x0 , y0 in zip ( x , y ) : <EOL> mu , sigma = surrogate . predict ( x0 ) <EOL> assert_rel_error ( self , mu , y0 , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0> , <NUM_LIT> ) <EOL> mu , sigma = surrogate . predict ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> assert_rel_error ( self , mu , branin ( [ <NUM_LIT> , <NUM_LIT> ] ) , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT> ) <EOL> def test_no_training_data ( self ) : <EOL> surrogate = KrigingSurrogate ( ) <EOL> try : <EOL> surrogate . predict ( [ <NUM_LIT:0.> , <NUM_LIT:1.> ] ) <EOL> except RuntimeError as err : <EOL> self . assertEqual ( str ( err ) , <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> else : <EOL> self . fail ( "<STR_LIT>" ) <EOL> def test_one_pt ( self ) : <EOL> surrogate = KrigingSurrogate ( ) <EOL> x = [ [ <NUM_LIT:0.> ] ] <EOL> y = [ [ <NUM_LIT:1.> ] ] <EOL> with self . assertRaises ( ValueError ) as cm : <EOL> surrogate . train ( x , y ) <EOL> self . assertEqual ( str ( cm . exception ) , '<STR_LIT>' <EOL> '<STR_LIT>' ) <EOL> def test_vector_input ( self ) : <EOL> surrogate = KrigingSurrogate ( ) <EOL> x = array ( [ [ <NUM_LIT:0.> , <NUM_LIT:0.> , <NUM_LIT:0.> ] , [ <NUM_LIT:1.> , <NUM_LIT:1.> , <NUM_LIT:1.> ] ] ) <EOL> y = array ( [ [ <NUM_LIT:0.> ] , [ <NUM_LIT> ] ] ) <EOL> surrogate . train ( x , y ) <EOL> for x0 , y0 in zip ( x , y ) : <EOL> mu , sigma = surrogate . predict ( x0 ) <EOL> assert_rel_error ( self , mu , y0 , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0> , <NUM_LIT> ) <EOL> def test_vector_output ( self ) : <EOL> surrogate = KrigingSurrogate ( ) <EOL> y = array ( [ [ <NUM_LIT:0.> , <NUM_LIT:0.> ] , [ <NUM_LIT:1.> , <NUM_LIT:1.> ] , [ <NUM_LIT> , <NUM_LIT:0.> ] ] ) <EOL> x = array ( [ [ <NUM_LIT:0.> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] ] ) <EOL> surrogate . train ( x , y ) <EOL> for x0 , y0 in zip ( x , y ) : <EOL> mu , sigma = surrogate . predict ( x0 ) <EOL> assert_rel_error ( self , mu , y0 , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0> , <NUM_LIT> ) <EOL> def test_scalar_derivs ( self ) : <EOL> surrogate = KrigingSurrogate ( ) <EOL> x = array ( [ [ <NUM_LIT:0.> ] , [ <NUM_LIT:1.> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] ] ) <EOL> y = x . copy ( ) <EOL> surrogate . train ( x , y ) <EOL> jac = surrogate . linearize ( array ( [ [ <NUM_LIT:0.> ] ] ) ) <EOL> assert_rel_error ( self , jac [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] , <NUM_LIT:1.> , <NUM_LIT> ) <EOL> def test_vector_derivs ( self ) : <EOL> surrogate = KrigingSurrogate ( ) <EOL> x = array ( [ [ a , b ] for a , b in <EOL> itertools . product ( linspace ( <NUM_LIT:0> , <NUM_LIT:1> , <NUM_LIT:10> ) , repeat = <NUM_LIT:2> ) ] ) <EOL> y = array ( [ [ a + b , a - b ] for a , b in x ] ) <EOL> surrogate . train ( x , y ) <EOL> jac = surrogate . linearize ( array ( [ [ <NUM_LIT:0.5> , <NUM_LIT:0.5> ] ] ) ) <EOL> assert_rel_error ( self , jac , array ( [ [ <NUM_LIT:1> , <NUM_LIT:1> ] , [ <NUM_LIT:1> , - <NUM_LIT:1> ] ] ) , <NUM_LIT> ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> unittest . main ( ) </s>
