@@ -16,7 +16,7 @@ ctypedef uint32_t (* random_uint_1_i_32)(aug_state* state, uint32_t a) nogil
16
16
ctypedef int32_t (* random_int_2_i_32)(aug_state* state, int32_t a, int32_t b) nogil
17
17
ctypedef int64_t (* random_int_2_i)(aug_state* state, int64_t a, int64_t b) nogil
18
18
19
- ctypedef void (* random_double_fill)(aug_state* state, int count, double * out) nogil
19
+ ctypedef void (* random_double_fill)(aug_state* state, np.npy_intp count, double * out) nogil
20
20
21
21
cdef Py_ssize_t compute_numel(size):
22
22
cdef Py_ssize_t i, n = 1
@@ -104,25 +104,28 @@ cdef object cont_broadcast_1(aug_state* state, void* func, object size, object l
104
104
object a, object a_name, constraint_type a_constraint):
105
105
106
106
cdef np.ndarray a_arr, randoms
107
+ cdef double * randoms_data
107
108
cdef np.broadcast it
108
109
cdef random_double_1 f = (< random_double_1> func)
110
+ cdef np.npy_intp i, n
109
111
110
112
a_arr = < np.ndarray> np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
111
113
if a_constraint != CONS_NONE:
112
114
check_array_constraint(a_arr, a_name, a_constraint)
113
115
114
116
if size is not None :
115
- randoms = np.empty(size, np.double)
117
+ randoms = < np.ndarray > np.empty(size, np.double)
116
118
else :
117
- # randoms = np.empty(np.shape(a_arr), np.double)
118
119
randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_DOUBLE)
119
120
120
-
121
+ randoms_data = < double * > np.PyArray_DATA(randoms)
122
+ n = np.PyArray_SIZE(randoms)
121
123
it = np.broadcast(randoms, a_arr)
124
+
122
125
with lock, nogil:
123
- while np.PyArray_MultiIter_NOTDONE(it ):
126
+ for i in range (n ):
124
127
a_val = (< double * > np.PyArray_MultiIter_DATA(it, 1 ))[0 ]
125
- ( < double * > np.PyArray_MultiIter_DATA(it, 0 ))[ 0 ] = f(state, a_val)
128
+ randoms_data[i ] = f(state, a_val)
126
129
127
130
np.PyArray_MultiIter_NEXT(it)
128
131
@@ -132,8 +135,10 @@ cdef object cont_broadcast_2(aug_state* state, void* func, object size, object l
132
135
object a, object a_name, constraint_type a_constraint,
133
136
object b, object b_name, constraint_type b_constraint):
134
137
cdef np.ndarray a_arr, b_arr, randoms
138
+ cdef double * randoms_data
135
139
cdef np.broadcast it
136
140
cdef random_double_2 f = (< random_double_2> func)
141
+ cdef np.npy_intp i, n
137
142
138
143
a_arr = < np.ndarray> np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
139
144
if a_constraint != CONS_NONE:
@@ -144,18 +149,22 @@ cdef object cont_broadcast_2(aug_state* state, void* func, object size, object l
144
149
check_array_constraint(b_arr, b_name, b_constraint)
145
150
146
151
if size is not None :
147
- randoms = np.empty(size, np.double)
152
+ randoms = < np.ndarray > np.empty(size, np.double)
148
153
else :
149
154
it = np.PyArray_MultiIterNew2(a_arr, b_arr)
150
- randoms = np.empty(it.shape, np.double)
155
+ randoms = < np.ndarray > np.empty(it.shape, np.double)
151
156
# randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE)
157
+ # TODO: These are needed when using a for loop
158
+
159
+ randoms_data = < double * > np.PyArray_DATA(randoms)
160
+ n = np.PyArray_SIZE(randoms)
152
161
153
162
it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
154
163
with lock, nogil:
155
- while np.PyArray_MultiIter_NOTDONE(it ):
164
+ for i in range (n ):
156
165
a_val = (< double * > np.PyArray_MultiIter_DATA(it, 1 ))[0 ]
157
166
b_val = (< double * > np.PyArray_MultiIter_DATA(it, 2 ))[0 ]
158
- ( < double * > np.PyArray_MultiIter_DATA(it, 0 ))[ 0 ] = f(state, a_val, b_val)
167
+ randoms_data[i ] = f(state, a_val, b_val)
159
168
160
169
np.PyArray_MultiIter_NEXT(it)
161
170
0 commit comments