Python codegen.codegen函数代码示例

def test_jl_InOutArgument_order():
    # can specify the order as (x, y)
    expr = Equality(x, x**2 + y)
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Julia", header=False,
                      empty=False, argument_sequence=(x,y))
    source = result[1]
    expected = (
        "function test(x, y)\n"
        "    x = x.^2 + y\n"
        "    return x\n"
        "end\n"
    )
    assert source == expected
    # make sure it gives (x, y) not (y, x)
    expr = Equality(x, x**2 + y)
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Julia", header=False, empty=False)
    source = result[1]
    expected = (
        "function test(x, y)\n"
        "    x = x.^2 + y\n"
        "    return x\n"
        "end\n"
    )
    assert source == expected

def test_InOutArgument_order():
    # can specify the order as (x, y)
    expr = Equality(x, x**2 + y)
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Rust", header=False,
                      empty=False, argument_sequence=(x,y))
    source = result[1]
    expected = (
        "fn test(x: f64, y: f64) -> f64 {\n"
        "    let x = x.powi(2) + y;\n"
        "    x\n"
        "}\n"
    )
    assert source == expected
    # make sure it gives (x, y) not (y, x)
    expr = Equality(x, x**2 + y)
    name_expr = ("test", expr)
    result, = codegen(name_expr, "Rust", header=False, empty=False)
    source = result[1]
    expected = (
        "fn test(x: f64, y: f64) -> f64 {\n"
        "    let x = x.powi(2) + y;\n"
        "    x\n"
        "}\n"
    )
    assert source == expected

def test_global_vars():
    x, y, z, t = symbols("x y z t")
    result = codegen(('f', x*y), "F95", header=False, empty=False,
                     global_vars=(y,))
    source = result[0][1]
    expected = (
        "REAL*8 function f(x)\n"
        "implicit none\n"
        "REAL*8, intent(in) :: x\n"
        "f = x*y\n"
        "end function\n"
        )
    assert source == expected

    result = codegen(('f', x*y+z), "C", header=False, empty=False,
                     global_vars=(z, t))
    source = result[0][1]
    expected = (
        '#include "f.h"\n'
        '#include <math.h>\n'
        'double f(double x, double y) {\n'
        '   double f_result;\n'
        '   f_result = x*y + z;\n'
        '   return f_result;\n'
        '}\n'
    )
    assert source == expected

def tensorized_basis_3D(order):
    total_integration_points = (order + 1) * (order + 1)
    r,s,t = sym.symbols('r,s,t')
    r_gll = sym.symbols('r_0:%d' % (order + 1))
    s_gll = sym.symbols('s_0:%d' % (order + 1))
    t_gll = sym.symbols('t_0:%d' % (order + 1))

    # Get N + 1 lagrange polynomials in each direction.
    generator_r = generating_polynomial_lagrange(order, 'r', r_gll)
    generator_s = generating_polynomial_lagrange(order, 's', s_gll)
    generator_t = generating_polynomial_lagrange(order, 't', t_gll)

    # Get tensorized basis.
    basis = TensorProduct(generator_t, generator_s, generator_r)
    gll_coordinates, gll_weights = gauss_lobatto_legendre_quadruature_points_weights(order + 1)

    basis = basis.subs([(v, c) for v, c in zip(r_gll, gll_coordinates)])
    basis = basis.subs([(v, c) for v, c in zip(s_gll, gll_coordinates)])
    basis = basis.subs([(v, c) for v, c in zip(t_gll, gll_coordinates)])

    # Get gradient of basis functions.
    basis_gradient_r = sym.Matrix([sym.diff(i, r) for i in basis])
    basis_gradient_s = sym.Matrix([sym.diff(i, s) for i in basis])
    basis_gradient_t = sym.Matrix([sym.diff(i, t) for i in basis])

    routines = [('interpolate_order{}_hex'.format(order),basis),
                ('interpolate_r_derivative_order{}_hex'.format(order), basis_gradient_r),
                ('interpolate_s_derivative_order{}_hex'.format(order), basis_gradient_s),
                ('interpolate_t_derivative_order{}_hex'.format(order), basis_gradient_t)]
    codegen(routines, "C", "order{}_hex".format(order), to_files=True, project="SALVUS")

    # fix headers (#include "order4_hex.h" -> #include <Element/HyperCube/Autogen/order3_hex.h>)
    fixHeader(order,"hex",".")

def test_global_vars():
    x, y, z, t = symbols("x y z t")
    result = codegen(('f', x*y), "F95", header=False, empty=False,
                     global_vars=(y,))
    source = result[0][1]
    expected = (
        "REAL*8 function f(x)\n"
        "implicit none\n"
        "REAL*8, intent(in) :: x\n"
        "f = x*y\n"
        "end function\n"
        )
    assert source == expected

    expected = (
        '#include "f.h"\n'
        '#include <math.h>\n'
        'double f(double x, double y) {\n'
        '   double f_result;\n'
        '   f_result = x*y + z;\n'
        '   return f_result;\n'
        '}\n'
    )
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", category=SymPyDeprecationWarning)
        result = codegen(('f', x*y+z), "C", header=False, empty=False,
                         global_vars=(z, t))
        source = result[0][1]
        assert source == expected

 def __init__(self,init):
     from numpy import zeros, f2py
     from sympy.utilities.codegen import codegen
     import sympy as sp
     self.type = init.type
     self.bounding_points = np.array(init.bounding_points)
     self.bounding_points_xi = None
     self.boundary_surface = init.boundary_surface
     self.flow_state = init.flow_state
     if self.boundary_surface:
         gradxsrc = codegen(
             ('gradx'+str(id(self)),
              sp.diff(sp.sympify(self.boundary_surface.split('=')[1].strip()),
                                 sp.Symbol('x'))),'F95','junk')[0]
         gradysrc = codegen(
             ('grady'+str(id(self)),
              sp.diff(sp.sympify(self.boundary_surface.split('=')[1].strip()),
                                 sp.Symbol('y'))),'F95','junk')[0]
         gradzsrc = codegen(
             ('gradz'+str(id(self)),
              sp.diff(sp.sympify(self.boundary_surface.split('=')[1].strip()),
                                 sp.Symbol('z'))),'F95','junk')[0]
         self.gradsrc = gradxsrc[1]+'\n'+gradysrc[1]+'\n'+gradzsrc[1]+'\n'
     try:
         iftest = bool(self.flow_state)
     except ValueError: # This means that the truth value is ambiguous (array).
         iftest = bool(self.flow_state.ndim == 1)
     if iftest:
         temp = zeros((points.shape[1],points.shape[2]))
         for i in range(points.shape[1]):
             for j in range(points.shape[2]):
                 temp[:,i,j] = self.flow_state
         self.flow_state = temp

def generate_from_sympy(expr, derivative=False, wrt=None):
    """ given a SymPy expression 'expr', generates C-code for that expression
        or its derivative (if 'derivative' == True); if returning
        derivative, wrt should also be specified
    """
    global _expr_count
    _expr_count += 1
    name = 'expr_%d' % _expr_count
    if derivative and (wrt is not None):
        return codegen((name, sympy.diff(expr, wrt)), 'C', name)[0][1]
    else:
        return codegen((name, expr), 'C', name)[0][1]

def test_check_case_false_positive():
    # The upper case/lower case exception should not be triggered by SymPy
    # objects that differ only because of assumptions.  (It may be useful to
    # have a check for that as well, but here we only want to test against
    # false positives with respect to case checking.)
    x1 = symbols('x')
    x2 = symbols('x', my_assumption=True)
    try:
        codegen(('test', x1*x2), 'f95', 'prefix')
    except CodeGenError as e:
        if e.args[0].startswith("Fortran ignores case."):
            raise AssertionError("This exception should not be raised!")

def test_check_case_false_positive():
    # The upper case/lower case exception should not be triggered by Sympy
    # objects that differ only because of assumptions.  (It may be useful to
    # have a check for that as well, but here we only want to test against
    # false positives with respect to case checking.)
    x1 = symbols("x")
    x2 = symbols("x", my_assumption=True)
    try:
        codegen(("test", x1 * x2), "f95", "prefix")
    except CodeGenError, e:
        if e.args[0][0:21] == "Fortran ignores case.":
            raise AssertionError("This exception should not be raised!")

def test_multifcns_per_file_w_header():
    name_expr = [ ("foo", [2*x, 3*y]), ("bar", [y**2, 4*y]) ]
    result = codegen(name_expr, "Rust", header=True, empty=False)
    assert result[0][0] == "foo.rs"
    source = result[0][1];
    version_str = "Code generated with sympy %s" % sympy.__version__
    version_line = version_str.center(76).rstrip()
    expected = (
        "/*\n"
        " *%(version_line)s\n"
        " *\n"
        " *              See http://www.sympy.org/ for more information.\n"
        " *\n"
        " *                       This file is part of 'project'\n"
        " */\n"
        "fn foo(x: f64, y: f64) -> (f64, f64) {\n"
        "    let out1 = 2*x;\n"
        "    let out2 = 3*y;\n"
        "    (out1, out2)\n"
        "}\n"
        "fn bar(y: f64) -> (f64, f64) {\n"
        "    let out1 = y.powi(2);\n"
        "    let out2 = 4*y;\n"
        "    (out1, out2)\n"
        "}\n"
    ) % {'version_line': version_line}
    assert source == expected

def test_complicated_codegen_f95():
    from sympy import sin, cos, tan
    x, y, z = symbols('x,y,z')
    name_expr = [
        ("test1", ((sin(x) + cos(y) + tan(z))**7).expand()),
        ("test2", cos(cos(cos(cos(cos(cos(cos(cos(x + y + z))))))))),
    ]
    result = codegen(name_expr, "F95", "file", header=False, empty=False)
    assert result[0][0] == "file.f90"
    expected = (
        'REAL*8 function test1(x, y, z)\n'
        'implicit none\n'
        'REAL*8, intent(in) :: x\n'
        'REAL*8, intent(in) :: y\n'
        'REAL*8, intent(in) :: z\n'
        'test1 = sin(x)**7 + 7*sin(x)**6*cos(y) + 7*sin(x)**6*tan(z) + 21*sin(x) &\n'
        '      **5*cos(y)**2 + 42*sin(x)**5*cos(y)*tan(z) + 21*sin(x)**5*tan(z) &\n'
        '      **2 + 35*sin(x)**4*cos(y)**3 + 105*sin(x)**4*cos(y)**2*tan(z) + &\n'
        '      105*sin(x)**4*cos(y)*tan(z)**2 + 35*sin(x)**4*tan(z)**3 + 35*sin( &\n'
        '      x)**3*cos(y)**4 + 140*sin(x)**3*cos(y)**3*tan(z) + 210*sin(x)**3* &\n'
        '      cos(y)**2*tan(z)**2 + 140*sin(x)**3*cos(y)*tan(z)**3 + 35*sin(x) &\n'
        '      **3*tan(z)**4 + 21*sin(x)**2*cos(y)**5 + 105*sin(x)**2*cos(y)**4* &\n'
        '      tan(z) + 210*sin(x)**2*cos(y)**3*tan(z)**2 + 210*sin(x)**2*cos(y) &\n'
        '      **2*tan(z)**3 + 105*sin(x)**2*cos(y)*tan(z)**4 + 21*sin(x)**2*tan &\n'
        '      (z)**5 + 7*sin(x)*cos(y)**6 + 42*sin(x)*cos(y)**5*tan(z) + 105* &\n'
        '      sin(x)*cos(y)**4*tan(z)**2 + 140*sin(x)*cos(y)**3*tan(z)**3 + 105 &\n'
        '      *sin(x)*cos(y)**2*tan(z)**4 + 42*sin(x)*cos(y)*tan(z)**5 + 7*sin( &\n'
        '      x)*tan(z)**6 + cos(y)**7 + 7*cos(y)**6*tan(z) + 21*cos(y)**5*tan( &\n'
        '      z)**2 + 35*cos(y)**4*tan(z)**3 + 35*cos(y)**3*tan(z)**4 + 21*cos( &\n'
        '      y)**2*tan(z)**5 + 7*cos(y)*tan(z)**6 + tan(z)**7\n'
        'end function\n'
        'REAL*8 function test2(x, y, z)\n'
        'implicit none\n'
        'REAL*8, intent(in) :: x\n'
        'REAL*8, intent(in) :: y\n'
        'REAL*8, intent(in) :: z\n'
        'test2 = cos(cos(cos(cos(cos(cos(cos(cos(x + y + z))))))))\n'
        'end function\n'
    )
    assert result[0][1] == expected
    assert result[1][0] == "file.h"
    expected = (
        'interface\n'
        'REAL*8 function test1(x, y, z)\n'
        'implicit none\n'
        'REAL*8, intent(in) :: x\n'
        'REAL*8, intent(in) :: y\n'
        'REAL*8, intent(in) :: z\n'
        'end function\n'
        'end interface\n'
        'interface\n'
        'REAL*8 function test2(x, y, z)\n'
        'implicit none\n'
        'REAL*8, intent(in) :: x\n'
        'REAL*8, intent(in) :: y\n'
        'REAL*8, intent(in) :: z\n'
        'end function\n'
        'end interface\n'
    )
    assert result[1][1] == expected

def test_fcode_matrix_output():
    x, y, z = symbols('x,y,z')
    e1 = x + y
    e2 = Matrix([[x, y], [z, 16]])
    name_expr = ("test", (e1, e2))
    result = codegen(name_expr, "f95", "test", header=False, empty=False)
    source = result[0][1]
    expected = (
        "REAL*8 function test(x, y, z, out_%(hash)s)\n"
        "implicit none\n"
        "REAL*8, intent(in) :: x\n"
        "REAL*8, intent(in) :: y\n"
        "REAL*8, intent(in) :: z\n"
        "REAL*8, intent(out), dimension(1:2, 1:2) :: out_%(hash)s\n"
        "out_%(hash)s(1, 1) = x\n"
        "out_%(hash)s(2, 1) = z\n"
        "out_%(hash)s(1, 2) = y\n"
        "out_%(hash)s(2, 2) = 16\n"
        "test = x + y\n"
        "end function\n"
    )
    # look for the magic number
    a = source.splitlines()[5]
    b = a.split('_')
    out = b[1]
    expected = expected % {'hash': out}
    assert source == expected

def test_jl_tensor_loops_multiple_contractions():
    # see comments in previous test about vectorizing
    from sympy.tensor import IndexedBase, Idx
    from sympy import symbols
    n, m, o, p = symbols('n m o p', integer=True)
    A = IndexedBase('A')
    B = IndexedBase('B')
    y = IndexedBase('y')
    i = Idx('i', m)
    j = Idx('j', n)
    k = Idx('k', o)
    l = Idx('l', p)
    result, = codegen(('tensorthing', Eq(y[i], B[j, k, l]*A[i, j, k, l])),
                      "Julia", header=False, empty=False)
    source = result[1]
    expected = (
        'function tensorthing(y, A, B, m, n, o, p)\n'
        '    for i = 1:m\n'
        '        y[i] = 0\n'
        '    end\n'
        '    for i = 1:m\n'
        '        for j = 1:n\n'
        '            for k = 1:o\n'
        '                for l = 1:p\n'
        '                    y[i] = y[i] + B[j,k,l].*A[i,j,k,l]\n'
        '                end\n'
        '            end\n'
        '        end\n'
        '    end\n'
        '    return y\n'
        'end\n'
    )
    assert source == expected

def test_fcode_matrixsymbol_slice():
    A = MatrixSymbol('A', 2, 3)
    B = MatrixSymbol('B', 1, 3)
    C = MatrixSymbol('C', 1, 3)
    D = MatrixSymbol('D', 2, 1)
    name_expr = ("test", [Equality(B, A[0, :]),
                          Equality(C, A[1, :]),
                          Equality(D, A[:, 2])])
    result = codegen(name_expr, "f95", "test", header=False, empty=False)
    source = result[0][1]
    expected = (
        "subroutine test(A, B, C, D)\n"
        "implicit none\n"
        "REAL*8, intent(in), dimension(1:2, 1:3) :: A\n"
        "REAL*8, intent(out), dimension(1:1, 1:3) :: B\n"
        "REAL*8, intent(out), dimension(1:1, 1:3) :: C\n"
        "REAL*8, intent(out), dimension(1:2, 1:1) :: D\n"
        "B(1, 1) = A(1, 1)\n"
        "B(1, 2) = A(1, 2)\n"
        "B(1, 3) = A(1, 3)\n"
        "C(1, 1) = A(2, 1)\n"
        "C(1, 2) = A(2, 2)\n"
        "C(1, 3) = A(2, 3)\n"
        "D(1, 1) = A(1, 3)\n"
        "D(2, 1) = A(2, 3)\n"
        "end subroutine\n"
    )
    assert source == expected

def test_complicated_codegen():
    from sympy import sin, cos, tan, N
    x,y,z = symbols('xyz')
    name_expr = [
        ("test1", ((sin(x)+cos(y)+tan(z))**7).expand()),
        ("test2", cos(cos(cos(cos(cos(cos(cos(cos(x+y+z))))))))),
    ]
    result = codegen(name_expr, "C", "file", header=False, empty=False)
    assert result[0][0] == "file.c"
    assert result[0][1] == (
        '#include "file.h"\n#include <math.h>\n'
        'double test1(double x, double y, double z) {\n'
        '   return '
        '7*pow(cos(y), 6)*sin(x) + '
        '7*pow(cos(y), 6)*tan(z) + '
        '7*pow(sin(x), 6)*cos(y) + '
        '7*pow(sin(x), 6)*tan(z) + '
        '7*pow(tan(z), 6)*cos(y) + '
        '7*pow(tan(z), 6)*sin(x) + '
        '42*pow(cos(y), 5)*sin(x)*tan(z) + '
        '42*pow(sin(x), 5)*cos(y)*tan(z) + '
        '42*pow(tan(z), 5)*cos(y)*sin(x) + '
        '105*pow(cos(y), 2)*pow(sin(x), 4)*tan(z) + '
        '105*pow(cos(y), 2)*pow(tan(z), 4)*sin(x) + '
        '105*pow(cos(y), 4)*pow(sin(x), 2)*tan(z) + '
        '105*pow(cos(y), 4)*pow(tan(z), 2)*sin(x) + '
        '105*pow(sin(x), 2)*pow(tan(z), 4)*cos(y) + '
        '105*pow(sin(x), 4)*pow(tan(z), 2)*cos(y) + '
        '140*pow(cos(y), 3)*pow(sin(x), 3)*tan(z) + '
        '140*pow(cos(y), 3)*pow(tan(z), 3)*sin(x) + '
        '140*pow(sin(x), 3)*pow(tan(z), 3)*cos(y) + '
        '21*pow(cos(y), 5)*pow(sin(x), 2) + '
        '21*pow(cos(y), 5)*pow(tan(z), 2) + '
        '21*pow(sin(x), 5)*pow(tan(z), 2) + '
        '210*pow(cos(y), 2)*pow(sin(x), 3)*pow(tan(z), 2) + '
        '210*pow(cos(y), 3)*pow(sin(x), 2)*pow(tan(z), 2) + '
        '35*pow(cos(y), 4)*pow(sin(x), 3) + '
        '35*pow(cos(y), 4)*pow(tan(z), 3) + '
        '35*pow(sin(x), 4)*pow(tan(z), 3) + '
        '210*pow(cos(y), 2)*pow(sin(x), 2)*pow(tan(z), 3) + '
        '35*pow(cos(y), 3)*pow(sin(x), 4) + '
        '35*pow(cos(y), 3)*pow(tan(z), 4) + '
        '35*pow(sin(x), 3)*pow(tan(z), 4) + '
        '21*pow(cos(y), 2)*pow(sin(x), 5) + '
        '21*pow(cos(y), 2)*pow(tan(z), 5) + '
        '21*pow(sin(x), 2)*pow(tan(z), 5) + '
        'pow(cos(y), 7) + pow(sin(x), 7) + pow(tan(z), 7);\n'
        '}\n'
        'double test2(double x, double y, double z) {\n'
        '   return cos(cos(cos(cos(cos(cos(cos(cos(x + y + z))))))));\n'
        '}\n'
    )
    assert result[1][0] == "file.h"
    assert result[1][1] == (
        '#ifndef PROJECT__FILE__H\n'
        '#define PROJECT__FILE__H\n'
        'double test1(double x, double y, double z);\n'
        'double test2(double x, double y, double z);\n'
        '#endif\n'
    )