Minor Improvement

symengine · Oct 3, 2017 · c6fea62 · c6fea62
1 parent 1716ceb
commit c6fea62
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 6 deletions.
diff --git a/symengine/pow.cpp b/symengine/pow.cpp
@@ -168,11 +168,17 @@ RCP<const Basic> pow(const RCP<const Basic> &a, const RCP<const Basic> &b)
         if (is_a<Mul>(*b)
             and is_a_Complex(*down_cast<const Mul &>(*b).get_coef())) {
             const map_basic_basic &dict = down_cast<const Mul &>(*b).get_dict();
+            RCP<const Number> coef = down_cast<const Mul &>(*b).get_coef();
             if (dict.size() == 1) {
-                for (const auto &p : dict) {
-                    if (eq(*p.first, *pi) and eq(*p.second, *one)) {
-                        return add(cos(mul(mul(I, minus_one), b)),
-                                   mul(I, sin(mul(mul(I, minus_one), b))));
+                if ((is_a<Integer>(*mul(coef, I))
+                     or (is_a<Rational>(*mul(coef, I))
+                         and (eq(*mul(coef, I), *rational(1, 2))
+                              or eq(*mul(coef, I), *rational(-1, 2)))))) {
+                    for (const auto &p : dict) {
+                        if (eq(*p.first, *pi) and eq(*p.second, *one)) {
+                            return add(cos(mul(mul(I, minus_one), b)),
+                                       mul(I, sin(mul(mul(I, minus_one), b))));
+                        }
                     }
                 }
             }

diff --git a/symengine/tests/basic/test_arit.cpp b/symengine/tests/basic/test_arit.cpp
@@ -32,6 +32,7 @@ using SymEngine::sub;
 using SymEngine::exp;
 using SymEngine::E;
 using SymEngine::Rational;
+using SymEngine::rational;
 using SymEngine::Complex;
 using SymEngine::Number;
 using SymEngine::I;
@@ -1009,14 +1010,20 @@ TEST_CASE("Pow: arit", "[arit]")
     REQUIRE(eq(*r1, *mul(I, minus_one)));
 
     r1 = div(exp(mul(mul(I, pi), x)), exp(x));
-    REQUIRE(r1->__str__() == "exp(-x + I*x*pi)");
+    REQUIRE(is_a<Pow>(*r1));
+    REQUIRE(eq(*down_cast<const Pow &>(*r1).get_base(), *E));
+    REQUIRE(eq(*down_cast<const Pow &>(*r1).get_exp(),
+               *add(mul(minus_one, x), mul(x, mul(I, pi)))));
 
     r1 = exp(add(div(mul(I, pi), integer(2)), x));
     r2 = mul(I, exp(x));
     REQUIRE(eq(*r1, *r2));
 
     r1 = exp(add(div(mul(I, pi), integer(10)), x));
-    REQUIRE(r1->__str__() == "exp(x)*(I*sin((1/10)*pi) + cos((1/10)*pi))");
+    REQUIRE(is_a<Pow>(*r1));
+    REQUIRE(eq(*down_cast<const Pow &>(*r1).get_base(), *E));
+    REQUIRE(eq(*down_cast<const Pow &>(*r1).get_exp(),
+               *add(x, mul(rational(1, 10), mul(I, pi)))));
 }
 
 TEST_CASE("Log: arit", "[arit]")