<s> from sympy . external import import_module <EOL> from sympy . utilities . pytest import raises <EOL> import ctypes <EOL> if import_module ( '<STR_LIT>' ) : <EOL> import sympy . printing . llvmjitcode as g <EOL> else : <EOL> disabled = True <EOL> import sympy <EOL> from sympy . abc import a , b , n <EOL> def isclose ( a , b ) : <EOL> rtol = <NUM_LIT> <EOL> atol = <NUM_LIT> <EOL> return abs ( a - b ) <= atol + rtol * abs ( b ) <EOL> def test_simple_expr ( ) : <EOL> e = a + <NUM_LIT:1.0> <EOL> f = g . llvm_callable ( [ a ] , e ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_two_arg ( ) : <EOL> e = <NUM_LIT> * a + b + <NUM_LIT> <EOL> f = g . llvm_callable ( [ a , b ] , e ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> , b : <NUM_LIT> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> , <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_func ( ) : <EOL> e = <NUM_LIT> * sympy . exp ( - a ) <EOL> f = g . llvm_callable ( [ a ] , e ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_two_func ( ) : <EOL> e = <NUM_LIT> * sympy . exp ( - a ) + sympy . exp ( b ) <EOL> f = g . llvm_callable ( [ a , b ] , e ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> , b : <NUM_LIT> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> , <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_two_sqrt ( ) : <EOL> e = <NUM_LIT> * sympy . sqrt ( a ) + sympy . sqrt ( b ) <EOL> f = g . llvm_callable ( [ a , b ] , e ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> , b : <NUM_LIT> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> , <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_two_pow ( ) : <EOL> e = a ** <NUM_LIT> + b ** <NUM_LIT:7> <EOL> f = g . llvm_callable ( [ a , b ] , e ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> , b : <NUM_LIT> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> , <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_callback ( ) : <EOL> e = a + <NUM_LIT> <EOL> f = g . llvm_callable ( [ a ] , e , callback_type = '<STR_LIT>' ) <EOL> m = ctypes . c_int ( <NUM_LIT:1> ) <EOL> array_type = ctypes . c_double * <NUM_LIT:1> <EOL> inp = { a : <NUM_LIT> } <EOL> array = array_type ( inp [ a ] ) <EOL> jit_res = f ( m , array ) <EOL> res = float ( e . subs ( inp ) . evalf ( ) ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_callback_cubature ( ) : <EOL> e = a + <NUM_LIT> <EOL> f = g . llvm_callable ( [ a ] , e , callback_type = '<STR_LIT>' ) <EOL> m = ctypes . c_int ( <NUM_LIT:1> ) <EOL> array_type = ctypes . c_double * <NUM_LIT:1> <EOL> inp = { a : <NUM_LIT> } <EOL> array = array_type ( inp [ a ] ) <EOL> out_array = array_type ( <NUM_LIT:0.0> ) <EOL> jit_ret = f ( m , array , None , m , out_array ) <EOL> assert jit_ret == <NUM_LIT:0> <EOL> res = float ( e . subs ( inp ) . evalf ( ) ) <EOL> assert isclose ( out_array [ <NUM_LIT:0> ] , res ) <EOL> def test_callback_two ( ) : <EOL> e = <NUM_LIT:3> * a * b <EOL> f = g . llvm_callable ( [ a , b ] , e , callback_type = '<STR_LIT>' ) <EOL> m = ctypes . c_int ( <NUM_LIT:2> ) <EOL> array_type = ctypes . c_double * <NUM_LIT:2> <EOL> inp = { a : <NUM_LIT> , b : <NUM_LIT> } <EOL> array = array_type ( inp [ a ] , inp [ b ] ) <EOL> jit_res = f ( m , array ) <EOL> res = float ( e . subs ( inp ) . evalf ( ) ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_callback_alt_two ( ) : <EOL> d = sympy . IndexedBase ( '<STR_LIT:d>' ) <EOL> e = <NUM_LIT:3> * d [ <NUM_LIT:0> ] * d [ <NUM_LIT:1> ] <EOL> f = g . llvm_callable ( [ n , d ] , e , callback_type = '<STR_LIT>' ) <EOL> m = ctypes . c_int ( <NUM_LIT:2> ) <EOL> array_type = ctypes . c_double * <NUM_LIT:2> <EOL> inp = { d [ <NUM_LIT:0> ] : <NUM_LIT> , d [ <NUM_LIT:1> ] : <NUM_LIT> } <EOL> array = array_type ( inp [ d [ <NUM_LIT:0> ] ] , inp [ d [ <NUM_LIT:1> ] ] ) <EOL> jit_res = f ( m , array ) <EOL> res = float ( e . subs ( inp ) . evalf ( ) ) <EOL> assert isclose ( jit_res , res ) <EOL> def test_multiple_statements ( ) : <EOL> e = [ [ ( b , <NUM_LIT> * a ) ] , [ b + <NUM_LIT:5> ] ] <EOL> f = g . llvm_callable ( [ a ] , e ) <EOL> b_val = e [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] [ <NUM_LIT:1> ] . subs ( { a : <NUM_LIT> } ) <EOL> res = float ( e [ <NUM_LIT:1> ] [ <NUM_LIT:0> ] . subs ( { b : b_val } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> ) <EOL> assert isclose ( jit_res , res ) <EOL> f_callback = g . llvm_callable ( [ a ] , e , callback_type = '<STR_LIT>' ) <EOL> m = ctypes . c_int ( <NUM_LIT:1> ) <EOL> array_type = ctypes . c_double * <NUM_LIT:1> <EOL> array = array_type ( <NUM_LIT> ) <EOL> jit_callback_res = f_callback ( m , array ) <EOL> assert isclose ( jit_callback_res , res ) <EOL> def test_cse ( ) : <EOL> e = a * a + b * b + sympy . exp ( - a * a - b * b ) <EOL> e2 = sympy . cse ( e ) <EOL> f = g . llvm_callable ( [ a , b ] , e2 ) <EOL> res = float ( e . subs ( { a : <NUM_LIT> , b : <NUM_LIT:0.1> } ) . evalf ( ) ) <EOL> jit_res = f ( <NUM_LIT> , <NUM_LIT:0.1> ) <EOL> assert isclose ( jit_res , res ) <EOL> def eval_cse ( e , sub_dict ) : <EOL> tmp_dict = dict ( ) <EOL> for tmp_name , tmp_expr in e [ <NUM_LIT:0> ] : <EOL> e2 = tmp_expr . subs ( sub_dict ) <EOL> e3 = e2 . subs ( tmp_dict ) <EOL> tmp_dict [ tmp_name ] = e3 <EOL> return [ e . subs ( sub_dict ) . subs ( tmp_dict ) for e in e [ <NUM_LIT:1> ] ] <EOL> def test_cse_multiple ( ) : <EOL> e1 = a * a <EOL> e2 = a * a + b * b <EOL> e3 = sympy . cse ( [ e1 , e2 ] ) <EOL> raises ( NotImplementedError , <EOL> lambda : g . llvm_callable ( [ a , b ] , e3 , callback_type = '<STR_LIT>' ) ) <EOL> f = g . llvm_callable ( [ a , b ] , e3 ) <EOL> jit_res = f ( <NUM_LIT:0.1> , <NUM_LIT> ) <EOL> assert len ( jit_res ) == <NUM_LIT:2> <EOL> res = eval_cse ( e3 , { a : <NUM_LIT:0.1> , b : <NUM_LIT> } ) <EOL> assert isclose ( res [ <NUM_LIT:0> ] , jit_res [ <NUM_LIT:0> ] ) <EOL> assert isclose ( res [ <NUM_LIT:1> ] , jit_res [ <NUM_LIT:1> ] ) <EOL> def test_callback_cubature_multiple ( ) : <EOL> e1 = a * a <EOL> e2 = a * a + b * b <EOL> e3 = sympy . cse ( [ e1 , e2 , <NUM_LIT:4> * e2 ] ) <EOL> f = g . llvm_callable ( [ a , b ] , e3 , callback_type = '<STR_LIT>' ) <EOL> ndim = <NUM_LIT:2> <EOL> outdim = <NUM_LIT:3> <EOL> m = ctypes . c_int ( ndim ) <EOL> fdim = ctypes . c_int ( outdim ) <EOL> array_type = ctypes . c_double * ndim <EOL> out_array_type = ctypes . c_double * outdim <EOL> inp = { a : <NUM_LIT> , b : <NUM_LIT> } <EOL> array = array_type ( inp [ a ] , inp [ b ] ) <EOL> out_array = out_array_type ( ) <EOL> jit_ret = f ( m , array , None , fdim , out_array ) <EOL> assert jit_ret == <NUM_LIT:0> <EOL> res = eval_cse ( e3 , inp ) <EOL> assert isclose ( out_array [ <NUM_LIT:0> ] , res [ <NUM_LIT:0> ] ) <EOL> assert isclose ( out_array [ <NUM_LIT:1> ] , res [ <NUM_LIT:1> ] ) <EOL> assert isclose ( out_array [ <NUM_LIT:2> ] , res [ <NUM_LIT:2> ] ) <EOL> def test_symbol_not_found ( ) : <EOL> e = a * a + b <EOL> raises ( LookupError , lambda : g . llvm_callable ( [ a ] , e ) ) <EOL> def test_bad_callback ( ) : <EOL> e = a <EOL> raises ( ValueError , lambda : g . llvm_callable ( [ a ] , e , callback_type = '<STR_LIT>' ) ) </s>
