Skip to content

Commit 7ed4529

Browse files
bashtagebashtage
committed
Merge pull request #27 from bashtage/array-init
Array init
2 parents 1dce278 + 0ef4409 commit 7ed4529

File tree

7 files changed

+142
-82
lines changed

7 files changed

+142
-82
lines changed

‎randomstate/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,3 +3,5 @@
33
from randomstate.prng.mt19937 import *
44
from randomstate.entropy import random_entropy
55
import randomstate.prng
6+
7+
sample = ranf = random = random_sample

‎randomstate/array_utilities.pxi

Lines changed: 82 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,10 @@ ctypedef uint32_t (* random_uint_1_i_32)(aug_state* state, uint32_t a) nogil
1616
ctypedef int32_t (* random_int_2_i_32)(aug_state* state, int32_t a, int32_t b) nogil
1717
ctypedef 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

21-
cdef Py_ssize_t compute_numel(size):
22-
cdef Py_ssize_t i, n = 1
21+
cdef np.npy_intp compute_numel(size):
22+
cdef np.npy_intp i, n = 1
2323
if isinstance(size, tuple):
2424
for i in range(len(size)):
2525
n *= size[i]
@@ -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)
151-
#randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE)
155+
randoms = <np.ndarray>np.empty(it.shape, np.double)
156+
# randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE)
157+
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

@@ -166,8 +175,10 @@ cdef object cont_broadcast_3(aug_state* state, void* func, object size, object l
166175
object b, object b_name, constraint_type b_constraint,
167176
object c, object c_name, constraint_type c_constraint):
168177
cdef np.ndarray a_arr, b_arr, c_arr, randoms
178+
cdef double *randoms_data
169179
cdef np.broadcast it
170180
cdef random_double_3 f = (<random_double_3>func)
181+
cdef np.npy_intp i, n
171182

172183
a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
173184
if a_constraint != CONS_NONE:
@@ -182,20 +193,22 @@ cdef object cont_broadcast_3(aug_state* state, void* func, object size, object l
182193
check_array_constraint(c_arr, c_name, c_constraint)
183194

184195
if size is not None:
185-
randoms = np.empty(size, np.double)
196+
randoms = <np.ndarray>np.empty(size, np.double)
186197
else:
187198
it = np.PyArray_MultiIterNew3(a_arr, b_arr, c_arr)
188199
#randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_DOUBLE)
189-
randoms = np.empty(it.shape, np.double)
200+
randoms = <np.ndarray>np.empty(it.shape, np.double)
190201

191-
it = it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr)
192-
# np.broadcast(randoms, a_arr, b_arr, c_arr)
202+
randoms_data = <double *>np.PyArray_DATA(randoms)
203+
n = np.PyArray_SIZE(randoms)
204+
205+
it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr)
193206
with lock, nogil:
194-
while np.PyArray_MultiIter_NOTDONE(it):
207+
for i in range(n):
195208
a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
196209
b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0]
197210
c_val = (<double*>np.PyArray_MultiIter_DATA(it, 3))[0]
198-
(<double*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val, b_val, c_val)
211+
randoms_data[i] = f(state, a_val, b_val, c_val)
199212

200213
np.PyArray_MultiIter_NEXT(it)
201214

@@ -254,8 +267,9 @@ cdef object cont(aug_state* state, void* func, object size, object lock, int nar
254267
elif narg == 3:
255268
return (<random_double_3>func)(state, _a, _b, _c)
256269

257-
cdef Py_ssize_t i, n = compute_numel(size)
258-
cdef double [::1] randoms = np.empty(n, np.double)
270+
cdef np.npy_intp i, n = compute_numel(size)
271+
cdef np.ndarray randoms = np.empty(n, np.double)
272+
cdef double *randoms_data = <double *>np.PyArray_DATA(randoms)
259273
cdef random_double_0 f0;
260274
cdef random_double_1 f1;
261275
cdef random_double_2 f2;
@@ -265,28 +279,30 @@ cdef object cont(aug_state* state, void* func, object size, object lock, int nar
265279
if narg == 0:
266280
f0 = (<random_double_0>func)
267281
for i in range(n):
268-
randoms[i] = f0(state)
282+
randoms_data[i] = f0(state)
269283
elif narg == 1:
270284
f1 = (<random_double_1>func)
271285
for i in range(n):
272-
randoms[i] = f1(state, _a)
286+
randoms_data[i] = f1(state, _a)
273287
elif narg == 2:
274288
f2 = (<random_double_2>func)
275289
for i in range(n):
276-
randoms[i] = f2(state, _a, _b)
290+
randoms_data[i] = f2(state, _a, _b)
277291
elif narg == 3:
278292
f3 = (<random_double_3>func)
279293
for i in range(n):
280-
randoms[i] = f3(state, _a, _b, _c)
294+
randoms_data[i] = f3(state, _a, _b, _c)
281295

282296
return np.asarray(randoms).reshape(size)
283297

284298
cdef object discrete_broadcast_d(aug_state* state, void* func, object size, object lock,
285299
object a, object a_name, constraint_type a_constraint):
286300

287301
cdef np.ndarray a_arr, randoms
302+
cdef long *randoms_data
288303
cdef np.broadcast it
289304
cdef random_uint_d f = (<random_uint_d>func)
305+
cdef np.npy_intp i, n
290306

291307
a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
292308
if a_constraint != CONS_NONE:
@@ -298,11 +314,14 @@ cdef object discrete_broadcast_d(aug_state* state, void* func, object size, obje
298314
#randoms = np.empty(np.shape(a_arr), np.double)
299315
randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_LONG)
300316

317+
randoms_data = <long *>np.PyArray_DATA(randoms)
318+
n = np.PyArray_SIZE(randoms)
319+
301320
it = np.broadcast(randoms, a_arr)
302321
with lock, nogil:
303-
while np.PyArray_MultiIter_NOTDONE(it):
322+
for i in range(n):
304323
a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
305-
(<long*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val)
324+
randoms_data[i] = f(state, a_val)
306325

307326
np.PyArray_MultiIter_NEXT(it)
308327

@@ -312,8 +331,10 @@ cdef object discrete_broadcast_dd(aug_state* state, void* func, object size, obj
312331
object a, object a_name, constraint_type a_constraint,
313332
object b, object b_name, constraint_type b_constraint):
314333
cdef np.ndarray a_arr, b_arr, randoms
334+
cdef long *randoms_data
315335
cdef np.broadcast it
316336
cdef random_uint_dd f = (<random_uint_dd>func)
337+
cdef np.npy_intp i, n
317338

318339
a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
319340
if a_constraint != CONS_NONE:
@@ -323,19 +344,21 @@ cdef object discrete_broadcast_dd(aug_state* state, void* func, object size, obj
323344
check_array_constraint(b_arr, b_name, b_constraint)
324345

325346
if size is not None:
326-
randoms = np.empty(size, np.int)
347+
randoms = <np.ndarray>np.empty(size, np.int)
327348
else:
328349
it = np.PyArray_MultiIterNew2(a_arr, b_arr)
329-
randoms = np.empty(it.shape, np.int)
350+
randoms = <np.ndarray>np.empty(it.shape, np.int)
330351
# randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_LONG)
331352

353+
randoms_data = <long *>np.PyArray_DATA(randoms)
354+
n = np.PyArray_SIZE(randoms)
355+
332356
it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
333357
with lock, nogil:
334-
while np.PyArray_MultiIter_NOTDONE(it):
358+
for i in range(n):
335359
a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
336360
b_val = (<double*>np.PyArray_MultiIter_DATA(it, 2))[0]
337-
338-
(<long*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val, b_val)
361+
randoms_data[i] = f(state, a_val, b_val)
339362

340363
np.PyArray_MultiIter_NEXT(it)
341364

@@ -345,8 +368,11 @@ cdef object discrete_broadcast_di(aug_state* state, void* func, object size, obj
345368
object a, object a_name, constraint_type a_constraint,
346369
object b, object b_name, constraint_type b_constraint):
347370
cdef np.ndarray a_arr, b_arr, randoms
371+
cdef long *randoms_data
348372
cdef np.broadcast it
349373
cdef random_uint_di f = (<random_uint_di>func)
374+
cdef np.npy_intp i, n
375+
350376

351377
a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_DOUBLE, np.NPY_ALIGNED)
352378
if a_constraint != CONS_NONE:
@@ -357,15 +383,17 @@ cdef object discrete_broadcast_di(aug_state* state, void* func, object size, obj
357383
check_array_constraint(b_arr, b_name, b_constraint)
358384

359385
if size is not None:
360-
randoms = np.empty(size, np.int)
386+
randoms = <np.ndarray>np.empty(size, np.int)
361387
else:
362388
it = np.PyArray_MultiIterNew2(a_arr, b_arr)
363-
randoms = np.empty(it.shape, np.int)
364-
#randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_LONG)
389+
randoms = <np.ndarray>np.empty(it.shape, np.int)
390+
391+
randoms_data = <long *>np.PyArray_DATA(randoms)
392+
n = np.PyArray_SIZE(randoms)
365393

366394
it = np.PyArray_MultiIterNew3(randoms, a_arr, b_arr)
367395
with lock, nogil:
368-
while np.PyArray_MultiIter_NOTDONE(it):
396+
for i in range(n):
369397
a_val = (<double*>np.PyArray_MultiIter_DATA(it, 1))[0]
370398
b_val = (<long*>np.PyArray_MultiIter_DATA(it, 2))[0]
371399
(<long*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val, b_val)
@@ -379,8 +407,10 @@ cdef object discrete_broadcast_iii(aug_state* state, void* func, object size, ob
379407
object b, object b_name, constraint_type b_constraint,
380408
object c, object c_name, constraint_type c_constraint):
381409
cdef np.ndarray a_arr, b_arr, c_arr, randoms
410+
cdef long *randoms_data
382411
cdef np.broadcast it
383412
cdef random_uint_iii f = (<random_uint_iii>func)
413+
cdef np.npy_intp i, n
384414

385415
a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_LONG, np.NPY_ALIGNED)
386416
if a_constraint != CONS_NONE:
@@ -395,19 +425,21 @@ cdef object discrete_broadcast_iii(aug_state* state, void* func, object size, ob
395425
check_array_constraint(c_arr, c_name, c_constraint)
396426

397427
if size is not None:
398-
randoms = np.empty(size, np.int)
428+
randoms = <np.ndarray>np.empty(size, np.int)
399429
else:
400430
it = np.PyArray_MultiIterNew3(a_arr, b_arr, c_arr)
401-
randoms = np.empty(it.shape, np.int)
402-
#randoms = np.PyArray_SimpleNew(it.nd, np.PyArray_DIMS(it), np.NPY_LONG)
431+
randoms = <np.ndarray>np.empty(it.shape, np.int)
432+
433+
randoms_data = <long *>np.PyArray_DATA(randoms)
434+
n = np.PyArray_SIZE(randoms)
403435

404436
it = np.PyArray_MultiIterNew4(randoms, a_arr, b_arr, c_arr)
405437
with lock, nogil:
406-
while np.PyArray_MultiIter_NOTDONE(it):
438+
for i in range(n):
407439
a_val = (<long*>np.PyArray_MultiIter_DATA(it, 1))[0]
408440
b_val = (<long*>np.PyArray_MultiIter_DATA(it, 2))[0]
409441
c_val = (<long*>np.PyArray_MultiIter_DATA(it, 3))[0]
410-
(<long*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val, b_val, c_val)
442+
randoms_data[i] = f(state, a_val, b_val, c_val)
411443

412444
np.PyArray_MultiIter_NEXT(it)
413445

@@ -416,24 +448,28 @@ cdef object discrete_broadcast_iii(aug_state* state, void* func, object size, ob
416448
cdef object discrete_broadcast_i(aug_state* state, void* func, object size, object lock,
417449
object a, object a_name, constraint_type a_constraint):
418450
cdef np.ndarray a_arr, randoms
451+
cdef long *randoms_data
419452
cdef np.broadcast it
420453
cdef random_uint_i f = (<random_uint_i>func)
454+
cdef np.npy_intp i, n
421455

422456
a_arr = <np.ndarray>np.PyArray_FROM_OTF(a, np.NPY_LONG, np.NPY_ALIGNED)
423457
if a_constraint != CONS_NONE:
424458
check_array_constraint(a_arr, a_name, a_constraint)
425459

426460
if size is not None:
427-
randoms = np.empty(size, np.int)
461+
randoms = <np.ndarray>np.empty(size, np.int)
428462
else:
429-
#randoms = np.empty(np.shape(a_arr), np.double)
430463
randoms = np.PyArray_SimpleNew(np.PyArray_NDIM(a_arr), np.PyArray_DIMS(a_arr), np.NPY_LONG)
431464

465+
randoms_data = <long *>np.PyArray_DATA(randoms)
466+
n = np.PyArray_SIZE(randoms)
467+
432468
it = np.broadcast(randoms, a_arr)
433469
with lock, nogil:
434-
while np.PyArray_MultiIter_NOTDONE(it):
470+
for i in range(n):
435471
a_val = (<long*>np.PyArray_MultiIter_DATA(it, 1))[0]
436-
(<long*>np.PyArray_MultiIter_DATA(it, 0))[0] = f(state, a_val)
472+
randoms_data[i] = f(state, a_val)
437473

438474
np.PyArray_MultiIter_NEXT(it)
439475

@@ -531,7 +567,7 @@ cdef object disc(aug_state* state, void* func, object size, object lock,
531567
else:
532568
return (<random_uint_iii>func)(state, _ia, _ib, _ic)
533569

534-
cdef Py_ssize_t i, n = compute_numel(size)
570+
cdef np.npy_intp i, n = compute_numel(size)
535571
cdef np.int_t [::1] randoms = np.empty(n, np.int)
536572
cdef random_uint_0 f0;
537573
cdef random_uint_d fd;

0 commit comments

Comments
 (0)