<s> import unittest <EOL> from numpy import array , sin , cos , pi , ones <EOL> from openmdao . api import MultiFiCoKrigingSurrogate <EOL> from openmdao . test . util import assert_rel_error <EOL> class CoKrigingSurrogateTest ( unittest . TestCase ) : <EOL> def test_1d_1fi_cokriging ( self ) : <EOL> x = array ( [ [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] ] ) <EOL> y = array ( [ <NUM_LIT> , - <NUM_LIT> , - <NUM_LIT> , <NUM_LIT> ] ) <EOL> krig1 = MultiFiCoKrigingSurrogate ( ) <EOL> krig1 . train ( x , y ) <EOL> new_x = array ( [ <NUM_LIT:0.5> ] ) <EOL> mu , sigma = krig1 . predict ( x [ <NUM_LIT:0> ] ) <EOL> assert_rel_error ( self , mu , y [ <NUM_LIT:0> ] , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0.> , <NUM_LIT> ) <EOL> mu , sigma = krig1 . predict ( new_x ) <EOL> assert_rel_error ( self , mu , - <NUM_LIT> , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT> ) <EOL> krig2 = MultiFiCoKrigingSurrogate ( theta = <NUM_LIT:0.1> ) <EOL> krig2 . train ( x , y ) <EOL> mu , sigma = krig2 . predict ( x [ <NUM_LIT:0> ] ) <EOL> assert_rel_error ( self , mu , y [ <NUM_LIT:0> ] , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT> ) <EOL> mu , sigma = krig2 . 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_2fi_cokriging ( self ) : <EOL> def f_expensive ( x ) : <EOL> return ( ( x * <NUM_LIT:6> - <NUM_LIT:2> ) ** <NUM_LIT:2> ) * sin ( ( x * <NUM_LIT:6> - <NUM_LIT:2> ) * <NUM_LIT:2> ) <EOL> def f_cheap ( x ) : <EOL> return <NUM_LIT:0.5> * ( ( x * <NUM_LIT:6> - <NUM_LIT:2> ) ** <NUM_LIT:2> ) * sin ( ( x * <NUM_LIT:6> - <NUM_LIT:2> ) * <NUM_LIT:2> ) + ( x - <NUM_LIT:0.5> ) * <NUM_LIT> - <NUM_LIT:5> <EOL> x = array ( [ [ [ <NUM_LIT:0.0> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT:1.0> ] ] , <EOL> [ [ <NUM_LIT:0.1> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT:0.5> ] , [ <NUM_LIT> ] , <EOL> [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT:0.0> ] , [ <NUM_LIT> ] , [ <NUM_LIT> ] , [ <NUM_LIT:1.0> ] ] ] ) <EOL> y = array ( [ [ f_expensive ( v ) for v in array ( x [ <NUM_LIT:0> ] ) . ravel ( ) ] , <EOL> [ f_cheap ( v ) for v in array ( x [ <NUM_LIT:1> ] ) . ravel ( ) ] ] ) <EOL> cokrig = MultiFiCoKrigingSurrogate ( ) <EOL> cokrig . train_multifi ( x , y ) <EOL> new_x = array ( [ <NUM_LIT> ] ) <EOL> mu , sigma = cokrig . predict ( new_x ) <EOL> assert_rel_error ( self , mu , f_expensive ( new_x [ <NUM_LIT:0> ] ) , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0.> , <NUM_LIT> ) <EOL> def test_2d_1fi_cokriging ( self ) : <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 ( x [ <NUM_LIT:0> ] ) + <NUM_LIT> <EOL> return y <EOL> x = array ( [ [ - <NUM_LIT> , <NUM_LIT:0.> ] , [ - <NUM_LIT:0.5> , <NUM_LIT> ] , [ <NUM_LIT:1.> , <NUM_LIT> ] , [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ - <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> krig1 = MultiFiCoKrigingSurrogate ( ) <EOL> krig1 . train ( x , y ) <EOL> mu , sigma = krig1 . predict ( [ - <NUM_LIT> , <NUM_LIT:0.> ] ) <EOL> assert_rel_error ( self , mu , branin ( x [ <NUM_LIT:0> ] ) , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0.> , <NUM_LIT> ) <EOL> mu , sigma = krig1 . predict ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> assert_rel_error ( self , mu , <NUM_LIT> , <NUM_LIT:1> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT:1> ) <EOL> krig2 = MultiFiCoKrigingSurrogate ( theta = [ <NUM_LIT:0.1> ] ) <EOL> krig1 . train ( x , y ) <EOL> mu , sigma = krig1 . predict ( [ - <NUM_LIT> , <NUM_LIT:0.> ] ) <EOL> assert_rel_error ( self , mu , branin ( x [ <NUM_LIT:0> ] ) , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT:0.> , <NUM_LIT> ) <EOL> mu , sigma = krig1 . predict ( [ <NUM_LIT> , <NUM_LIT> ] ) <EOL> assert_rel_error ( self , mu , <NUM_LIT> , <NUM_LIT:1> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT:1> ) <EOL> def test_2d_2fi_cokriging ( self ) : <EOL> def branin ( x ) : <EOL> x1 = <NUM_LIT:15> * x [ <NUM_LIT:0> ] - <NUM_LIT:5> <EOL> x2 = <NUM_LIT:15> * x [ <NUM_LIT:1> ] <EOL> return ( x2 - ( <NUM_LIT> / ( <NUM_LIT> * pi ** <NUM_LIT> ) ) * x1 ** <NUM_LIT> + <NUM_LIT> * x1 / pi - <NUM_LIT> ) ** <NUM_LIT> + <NUM_LIT> * ( <NUM_LIT:1.> - <NUM_LIT:1.> / ( <NUM_LIT> * pi ) ) * cos ( x1 ) + <NUM_LIT> <EOL> def branin_low_fidelity ( x ) : <EOL> return branin ( x ) + <NUM_LIT> * x [ <NUM_LIT:1> ] + <NUM_LIT> <EOL> x = [ [ [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] ] , <EOL> [ [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] , <EOL> [ <NUM_LIT> , <NUM_LIT> ] ] ] <EOL> y = array ( [ [ branin ( case ) for case in x [ <NUM_LIT:0> ] ] , <EOL> [ branin_low_fidelity ( case ) for case in x [ <NUM_LIT:1> ] ] ] ) <EOL> nfi = <NUM_LIT:2> <EOL> cokrig = MultiFiCoKrigingSurrogate ( ) <EOL> cokrig . train_multifi ( x , y ) <EOL> mu , sigma = cokrig . predict ( [ <NUM_LIT> / <NUM_LIT> , <NUM_LIT:1> / <NUM_LIT> ] ) <EOL> assert_rel_error ( self , mu , <NUM_LIT> , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT> ) <EOL> cokrig2 = MultiFiCoKrigingSurrogate ( theta = <NUM_LIT:0.1> ) <EOL> cokrig2 . train_multifi ( x , y ) <EOL> mu , sigma = cokrig2 . predict ( [ <NUM_LIT> / <NUM_LIT> , <NUM_LIT:1> / <NUM_LIT> ] ) <EOL> assert_rel_error ( self , mu , <NUM_LIT> , <NUM_LIT:0.1> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT:0.1> ) <EOL> cokrig2 = MultiFiCoKrigingSurrogate ( theta = [ <NUM_LIT:0.1> , <NUM_LIT:10> ] ) <EOL> cokrig2 . train_multifi ( x , y ) <EOL> mu , sigma = cokrig2 . predict ( [ <NUM_LIT> / <NUM_LIT> , <NUM_LIT:1> / <NUM_LIT> ] ) <EOL> assert_rel_error ( self , mu , <NUM_LIT> , <NUM_LIT> ) <EOL> assert_rel_error ( self , sigma , <NUM_LIT> , <NUM_LIT> ) <EOL> cokrig3 = MultiFiCoKrigingSurrogate ( theta = [ <NUM_LIT:0.1> ] ) <EOL> try : <EOL> cokrig3 . train_multifi ( x , y ) <EOL> except ValueError as err : <EOL> self . assertEqual ( str ( err ) , <EOL> "<STR_LIT>" ) <EOL> else : <EOL> self . fail ( "<STR_LIT>" ) <EOL> if __name__ == "<STR_LIT:__main__>" : <EOL> unittest . main ( ) </s>
