Add optimizations for integral powers of complex numbers to nc_pow.

1 files changed, 30 insertions, 8 deletions

diff --git a/numpy/core/src/umathmodule.c.src b/numpy/core/src/umathmodule.c.src
index 896371e65..11d7c0c7b 100644
--- a/numpy/core/src/umathmodule.c.src
+++ b/numpy/core/src/umathmodule.c.src
@@ -644,22 +644,44 @@ nc_expm1@c@(c@typ@ *x, c@typ@ *r)
 static void
 nc_pow@c@(c@typ@ *a, c@typ@ *b, c@typ@ *r)
 {
-	@typ@ ar=a->real, br=b->real, ai=a->imag, bi=b->imag;
+        intp n;
+        @typ@ ar=a->real, br=b->real, ai=a->imag, bi=b->imag;
 
 	if (br == 0. && bi == 0.) {
 		r->real = 1.;
 		r->imag = 0.;
+                return;
 	}
-	else if (ar == 0. && ai == 0.) {
+	if (ar == 0. && ai == 0.) {
 		r->real = 0.;
 		r->imag = 0.;
+                return;
 	}
-	else {
-		nc_log@c@(a, r);
-		nc_prod@c@(r, b, r);
-		nc_exp@c@(r, r);
-	}
-	return;
+        if (bi == 0 && (n=(intp)br) == br) {
+            if (n > -100 && n < 100) {
+                c@typ@ p, a;
+                intp mask = 1;
+                if (n < 0) n = -n;
+                a = nc_1@c@;
+                p.real = ar; p.imag = ai;
+                while (1) {
+                        if (n & mask)
+                            nc_prod@c@(&a,&p,&a);
+                        mask <<= 1;
+                        if (n < mask || mask <= 0) break;
+                        nc_prod@c@(&p,&p,&p);
+                }
+                r->real = a.real; r->imag = a.imag;
+                if (bi < 0) nc_quot@c@(&nc_1@c@, r, r);
+                return;
+            }
+        }
+        /* complexobect.c uses an inline version of this formula
+           investigate whether this had better performance or accuracy */
+        nc_log@c@(a, r);
+        nc_prod@c@(r, b, r);
+        nc_exp@c@(r, r);
+        return;
 }