fix numba inference

Author: dfulu

pvlib/spa.py

@@ -456,7 +456,12 @@ def julian_ephemeris_millennium(julian_ephemeris_century):
 @jcompile(nopython=True)
 def sum_mult_cos_add_mult(arr, x):
     # shared calculation used for heliocentric longitude, latitude, and radius
-    s = (arr[:, [0]] * np.cos(arr[:, [1]] + arr[:, [2]] * np.expand_dims(x, axis=0))).sum(axis=0)
+    if USE_NUMBA:
+        s = 0.
+        for row in range(arr.shape[0]):
+            s += arr[row, 0] * np.cos(arr[row, 1] + arr[row, 2] * x)
+    else:
+        s = (arr[:, [0]] * np.cos(arr[:, [1]] + arr[:, [2]] * np.expand_dims(x, axis=0))).sum(axis=0)
     return s
 
 @jcompile('float64(float64)', nopython=True)
@@ -557,20 +562,41 @@ def moon_ascending_longitude(julian_ephemeris_century):
     nopython=True)
 def longitude_obliquity_nutation(julian_ephemeris_century, x0, x1, x2, x3, x4,
                                  out):
-    a = NUTATION_ABCD_ARRAY[:, [0]]
-    b = NUTATION_ABCD_ARRAY[:, [1]]
-    c = NUTATION_ABCD_ARRAY[:, [2]]
-    d = NUTATION_ABCD_ARRAY[:, [3]]
-    arg = np.radians(
-        NUTATION_YTERM_ARRAY[:, [0]]*np.expand_dims(x0, axis=0) +
-        NUTATION_YTERM_ARRAY[:, [1]]*np.expand_dims(x1, axis=0) +
-        NUTATION_YTERM_ARRAY[:, [2]]*np.expand_dims(x2, axis=0) +
-        NUTATION_YTERM_ARRAY[:, [3]]*np.expand_dims(x3, axis=0) +
-        NUTATION_YTERM_ARRAY[:, [4]]*np.expand_dims(x4, axis=0)
-    )
-
-    delta_psi_sum = ((a + b * julian_ephemeris_century) * np.sin(arg)).sum(axis=0)
-    delta_eps_sum = ((c + d * julian_ephemeris_century) * np.cos(arg)).sum(axis=0)
+    
+    if USE_NUMBA:
+        delta_psi_sum = 0.0
+        delta_eps_sum = 0.0
+        for row in range(NUTATION_YTERM_ARRAY.shape[0]):
+            a = NUTATION_ABCD_ARRAY[row, 0]
+            b = NUTATION_ABCD_ARRAY[row, 1]
+            c = NUTATION_ABCD_ARRAY[row, 2]
+            d = NUTATION_ABCD_ARRAY[row, 3]
+            arg = np.radians(
+                NUTATION_YTERM_ARRAY[row, 0]*x0 +
+                NUTATION_YTERM_ARRAY[row, 1]*x1 +
+                NUTATION_YTERM_ARRAY[row, 2]*x2 +
+                NUTATION_YTERM_ARRAY[row, 3]*x3 +
+                NUTATION_YTERM_ARRAY[row, 4]*x4
+            )
+            delta_psi_sum += (a + b * julian_ephemeris_century) * np.sin(arg)
+            delta_eps_sum += (c + d * julian_ephemeris_century) * np.cos(arg)
+
+    else:
+        a = NUTATION_ABCD_ARRAY[:, [0]]
+        b = NUTATION_ABCD_ARRAY[:, [1]]
+        c = NUTATION_ABCD_ARRAY[:, [2]]
+        d = NUTATION_ABCD_ARRAY[:, [3]]
+        arg = np.radians(
+            NUTATION_YTERM_ARRAY[:, [0]]*np.expand_dims(x0, axis=0) +
+            NUTATION_YTERM_ARRAY[:, [1]]*np.expand_dims(x1, axis=0) +
+            NUTATION_YTERM_ARRAY[:, [2]]*np.expand_dims(x2, axis=0) +
+            NUTATION_YTERM_ARRAY[:, [3]]*np.expand_dims(x3, axis=0) +
+            NUTATION_YTERM_ARRAY[:, [4]]*np.expand_dims(x4, axis=0)
+        )
+
+        delta_psi_sum = ((a + b * julian_ephemeris_century) * np.sin(arg)).sum(axis=0)
+        delta_eps_sum = ((c + d * julian_ephemeris_century) * np.cos(arg)).sum(axis=0)
+    
     delta_psi = delta_psi_sum*1.0/36000000
     delta_eps = delta_eps_sum*1.0/36000000
     # seems like we ought to be able to return a tuple here instead