test_sympytools.py 文件源码

def test_print_as_array():
    xvar = sympy.var('xvar')
    yvar = sympy.var('yvar')
    m = sympy.Matrix([[1, 2*xvar + yvar**3, + 2*xvar*(yvar+1)]])
    res1 = '''
# m
# m_num=3
m[pos+0] += 1
m[pos+1] += 2*xvar + (yvar*yvar*yvar)
m[pos+2] += 2*xvar*(yvar + 1)
    '''
    res1 = res1.strip()
    assert print_as_array(m, 'm') == res1

    res2 = '''
# cdefs
cdef double x0
# subs
x0 = 2*xvar
# m
# m_num=3
m[pos+0] += 1
m[pos+1] += x0 + (yvar*yvar*yvar)
m[pos+2] += x0*(yvar + 1)
'''
    res2 = res2.strip()
    assert print_as_array(m, 'm', use_cse=True) == res2