@@ -16,7 +16,7 @@ ctypedef uint32_t (* random_uint_1_i_32)(aug_state* state, uint32_t a) nogil
1616ctypedef int32_t (* random_int_2_i_32)(aug_state* state, int32_t a, int32_t b) nogil
1717ctypedef int64_t (* random_int_2_i)(aug_state* state, int64_t a, int64_t b) nogil
1818
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
2020
2121cdef Py_ssize_t compute_numel(size):
2222 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
104104 object a, object a_name, constraint_type a_constraint):
105105
106106 cdef np.ndarray a_arr, randoms
107+ cdef double * randoms_data
107108 cdef np.broadcast it
108109 cdef random_double_1 f = (< random_double_1> func)
110+ cdef np.npy_intp i, n
109111
110112 a_arr = < np.ndarray> np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
111113 if a_constraint != CONS_NONE:
112114 check_array_constraint(a_arr, a_name, a_constraint)
113115
114116 if size is not None :
115- randoms = np.empty(size, np.double)
117+ randoms = < np.ndarray > np.empty(size, np.double)
116118 else :
117- # randoms = np.empty(np.shape(a_arr), np.double)
118119 randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_DOUBLE)
119120
120-
121+ randoms_data = < double * > np.PyArray_DATA(randoms)
122+ n = np.PyArray_SIZE(randoms)
121123 it = np.broadcast(randoms, a_arr)
124+
122125 with lock, nogil:
123- while np.PyArray_MultiIter_NOTDONE(it ):
126+ for i in range (n ):
124127 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)
126129
127130 np.PyArray_MultiIter_NEXT(it)
128131
@@ -132,8 +135,10 @@ cdef object cont_broadcast_2(aug_state* state, void* func, object size, object l
132135 object a, object a_name, constraint_type a_constraint,
133136 object b, object b_name, constraint_type b_constraint):
134137 cdef np.ndarray a_arr, b_arr, randoms
138+ cdef double * randoms_data
135139 cdef np.broadcast it
136140 cdef random_double_2 f = (< random_double_2> func)
141+ cdef np.npy_intp i, n
137142
138143 a_arr = < np.ndarray> np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
139144 if a_constraint != CONS_NONE:
@@ -144,18 +149,22 @@ cdef object cont_broadcast_2(aug_state* state, void* func, object size, object l
144149 check_array_constraint(b_arr, b_name, b_constraint)
145150
146151 if size is not None :
147- randoms = np.empty(size, np.double)
152+ randoms = < np.ndarray > np.empty(size, np.double)
148153 else :
149154 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)
151156 # 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)
152161
153162 it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
154163 with lock, nogil:
155- while np.PyArray_MultiIter_NOTDONE(it ):
164+ for i in range (n ):
156165 a_val = (< double * > np.PyArray_MultiIter_DATA(it, 1 ))[0 ]
157166 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)
159168
160169 np.PyArray_MultiIter_NEXT(it)
161170
0 commit comments