def test_loops():
m, n = symbols('m n', integer=True)
A = IndexedBase('A')
x = IndexedBase('x')
y = IndexedBase('y')
z = IndexedBase('z')
i = Idx('i', m)
j = Idx('j', n)
assert rust_code(A[i, j]*x[j], assign_to=y[i]) == (
"for i in 0..m {\n"
" y[i] = 0;\n"
"}\n"
"for i in 0..m {\n"
" for j in 0..n {\n"
" y[i] = A[n*i + j]*x[j] + y[i];\n"
" }\n"
"}")
assert rust_code(A[i, j]*x[j] + x[i] + z[i], assign_to=y[i]) == (
"for i in 0..m {\n"
" y[i] = x[i] + z[i];\n"
"}\n"
"for i in 0..m {\n"
" for j in 0..n {\n"
" y[i] = A[n*i + j]*x[j] + y[i];\n"
" }\n"
"}")
def test_printmethod():
class fabs(Abs):
def _rust_code(self, printer):
return "%s.fabs()" % printer._print(self.args[0])
assert rust_code(fabs(x)) == "x.fabs()"
a = MatrixSymbol("a", 1, 3)
assert rust_code(a[0,0]) == 'a[0]'
def test_user_functions():
x = symbols('x', integer=False)
n = symbols('n', integer=True)
custom_functions = {
"ceiling": "ceil",
"Abs": [(lambda x: not x.is_integer, "fabs", 4), (lambda x: x.is_integer, "abs", 4)],
}
assert rust_code(ceiling(x), user_functions=custom_functions) == "x.ceil()"
assert rust_code(Abs(x), user_functions=custom_functions) == "fabs(x)"
assert rust_code(Abs(n), user_functions=custom_functions) == "abs(n)"
def test_reserved_words():
x, y = symbols("x if")
expr = sin(y)
assert rust_code(expr) == "if_.sin()"
assert rust_code(expr, dereference=[y]) == "(*if_).sin()"
assert rust_code(expr, reserved_word_suffix='_unreserved') == "if_unreserved.sin()"
with raises(ValueError):
rust_code(expr, error_on_reserved=True)
def test_sign():
expr = sign(x) * y
assert rust_code(expr) == "y*x.signum()"
assert rust_code(expr, assign_to='r') == "r = y*x.signum();"
expr = sign(x + y) + 42
assert rust_code(expr) == "(x + y).signum() + 42"
assert rust_code(expr, assign_to='r') == "r = (x + y).signum() + 42;"
expr = sign(cos(x))
assert rust_code(expr) == "x.cos().signum()"
def test_Indexed():
n, m, o = symbols('n m o', integer=True)
i, j, k = Idx('i', n), Idx('j', m), Idx('k', o)
x = IndexedBase('x')[j]
assert rust_code(x) == "x[j]"
A = IndexedBase('A')[i, j]
assert rust_code(A) == "A[m*i + j]"
B = IndexedBase('B')[i, j, k]
assert rust_code(B) == "B[m*o*i + o*j + k]"
def test_inline_function():
x = symbols('x')
g = implemented_function('g', Lambda(x, 2*x))
assert rust_code(g(x)) == "2*x"
g = implemented_function('g', Lambda(x, 2*x/Catalan))
assert rust_code(g(x)) == (
"const Catalan: f64 = %s;\n2*x/Catalan" % Catalan.evalf(17))
A = IndexedBase('A')
i = Idx('i', symbols('n', integer=True))
g = implemented_function('g', Lambda(x, x*(1 + x)*(2 + x)))
assert rust_code(g(A[i]), assign_to=A[i]) == (
"for i in 0..n {\n"
" A[i] = (A[i] + 1)*(A[i] + 2)*A[i];\n"
"}")
def test_loops_addfactor():
m, n, o, p = symbols('m n o p', integer=True)
a = IndexedBase('a')
b = IndexedBase('b')
c = IndexedBase('c')
y = IndexedBase('y')
i = Idx('i', m)
j = Idx('j', n)
k = Idx('k', o)
l = Idx('l', p)
code = rust_code((a[i, j, k, l] + b[i, j, k, l])*c[j, k, l], assign_to=y[i])
assert code == (
"for i in 0..m {\n"
" y[i] = 0;\n"
"}\n"
"for i in 0..m {\n"
" for j in 0..n {\n"
" for k in 0..o {\n"
" for l in 0..p {\n"
" y[i] = (a[%s] + b[%s])*c[%s] + y[i];\n" % (i*n*o*p + j*o*p + k*p + l, i*n*o*p + j*o*p + k*p + l, j*o*p + k*p + l) +\
" }\n"
" }\n"
" }\n"
"}")
def test_ITE():
expr = ITE(x < 1, y, z)
assert rust_code(expr) == (
"if (x < 1) {\n"
" y\n"
"} else {\n"
" z\n"
"}")
def test_dummy_loops():
i, m = symbols('i m', integer=True, cls=Dummy)
x = IndexedBase('x')
y = IndexedBase('y')
i = Idx(i, m)
assert rust_code(x[i], assign_to=y[i]) == (
"for i in 0..m {\n"
" y[i] = x[i];\n"
"}")
def test_Rational():
assert rust_code(Rational(3, 7)) == "3_f64/7.0"
assert rust_code(Rational(18, 9)) == "2"
assert rust_code(Rational(3, -7)) == "-3_f64/7.0"
assert rust_code(Rational(-3, -7)) == "3_f64/7.0"
assert rust_code(x + Rational(3, 7)) == "x + 3_f64/7.0"
assert rust_code(Rational(3, 7)*x) == "(3_f64/7.0)*x"
def test_Relational():
assert rust_code(Eq(x, y)) == "x == y"
assert rust_code(Ne(x, y)) == "x != y"
assert rust_code(Le(x, y)) == "x <= y"
assert rust_code(Lt(x, y)) == "x < y"
assert rust_code(Gt(x, y)) == "x > y"
assert rust_code(Ge(x, y)) == "x >= y"
def test_Piecewise():
expr = Piecewise((x, x < 1), (x + 2, True))
assert rust_code(expr) == (
"if (x < 1) {\n"
" x\n"
"} else {\n"
" x + 2\n"
"}")
assert rust_code(expr, assign_to="r") == (
"r = if (x < 1) {\n"
" x\n"
"} else {\n"
" x + 2\n"
"};")
assert rust_code(expr, assign_to="r", inline=True) == (
"r = if (x < 1) { x } else { x + 2 };")
expr = Piecewise((x, x < 1), (x + 1, x < 5), (x + 2, True))
assert rust_code(expr, inline=True) == (
"if (x < 1) { x } else if (x < 5) { x + 1 } else { x + 2 }")
assert rust_code(expr, assign_to="r", inline=True) == (
"r = if (x < 1) { x } else if (x < 5) { x + 1 } else { x + 2 };")
assert rust_code(expr, assign_to="r") == (
"r = if (x < 1) {\n"
" x\n"
"} else if (x < 5) {\n"
" x + 1\n"
"} else {\n"
" x + 2\n"
"};")
expr = 2*Piecewise((x, x < 1), (x + 1, x < 5), (x + 2, True))
assert rust_code(expr, inline=True) == (
"2*if (x < 1) { x } else if (x < 5) { x + 1 } else { x + 2 }")
expr = 2*Piecewise((x, x < 1), (x + 1, x < 5), (x + 2, True)) - 42
assert rust_code(expr, inline=True) == (
"2*if (x < 1) { x } else if (x < 5) { x + 1 } else { x + 2 } - 42")
# Check that Piecewise without a True (default) condition error
expr = Piecewise((x, x < 1), (x**2, x > 1), (sin(x), x > 0))
raises(ValueError, lambda: rust_code(expr))
def test_loops_multiple_contractions():
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)
assert rust_code(b[j, k, l]*a[i, j, k, l], assign_to=y[i]) == (
"for i in 0..m {\n"
" y[i] = 0;\n"
"}\n"
"for i in 0..m {\n"
" for j in 0..n {\n"
" for k in 0..o {\n"
" for l in 0..p {\n"
" y[i] = a[%s]*b[%s] + y[i];\n" % (i*n*o*p + j*o*p + k*p + l, j*o*p + k*p + l) +\
" }\n"
" }\n"
" }\n"
"}")
def test_constants_other():
assert rust_code(2*GoldenRatio) == "const GoldenRatio: f64 = %s;\n2*GoldenRatio" % GoldenRatio.evalf(17)
assert rust_code(
2*Catalan) == "const Catalan: f64 = %s;\n2*Catalan" % Catalan.evalf(17)
assert rust_code(2*EulerGamma) == "const EulerGamma: f64 = %s;\n2*EulerGamma" % EulerGamma.evalf(17)
def test_sparse_matrix():
# gh-15791
assert 'Not supported in Rust' in rust_code(SparseMatrix([[1, 2, 3]]))
def test_basic_ops():
assert rust_code(x + y) == "x + y"
assert rust_code(x - y) == "x - y"
assert rust_code(x * y) == "x*y"
assert rust_code(x / y) == "x/y"
assert rust_code(-x) == "-x"
def test_matrix():
assert rust_code(Matrix([1, 2, 3])) == '[1, 2, 3]'
with raises(ValueError):
rust_code(Matrix([[1, 2, 3]]))
def test_dereference_printing():
expr = x + y + sin(z) + z
assert rust_code(expr, dereference=[z]) == "x + y + (*z) + (*z).sin()"
def test_boolean():
assert rust_code(True) == "true"
assert rust_code(S.true) == "true"
assert rust_code(False) == "false"
assert rust_code(S.false) == "false"
assert rust_code(x & y) == "x && y"
assert rust_code(x | y) == "x || y"
assert rust_code(~x) == "!x"
assert rust_code(x & y & z) == "x && y && z"
assert rust_code(x | y | z) == "x || y || z"
assert rust_code((x & y) | z) == "z || x && y"
assert rust_code((x | y) & z) == "z && (x || y)"
def test_settings():
raises(TypeError, lambda: rust_code(sin(x), method="garbage"))
def test_Integer():
assert rust_code(Integer(42)) == "42"
assert rust_code(Integer(-56)) == "-56"
