Skip to content

Commit 4f59e82

Browse files
oleksandr-pavlykoleksandr-pavlyk
committed
Use public API to work with cnp.broadcast
1 parent 51f8ddc commit 4f59e82

File tree

1 file changed

+91
-38
lines changed

1 file changed

+91
-38
lines changed

‎mkl_random/mklrand.pyx‎

Lines changed: 91 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ cdef extern from "Python.h":
3939

4040
cimport numpy as cnp
4141
from libc.string cimport memset, memcpy
42+
cimport cpython.tuple
4243

4344
cdef extern from "math.h":
4445
double floor(double x)
@@ -244,6 +245,9 @@ cdef object vec_cont1_array(irk_state *state, irk_cont1_vec func, object size,
244245
cdef cnp.flatiter itera
245246
cdef cnp.broadcast multi
246247
cdef object arr_obj
248+
cdef Py_ssize_t multi_nd
249+
cdef tuple multi_shape
250+
cdef cnp.npy_intp *multi_dims
247251

248252
if size is None:
249253
array = <cnp.ndarray>cnp.PyArray_SimpleNew(cnp.PyArray_NDIM(oa),
@@ -261,20 +265,25 @@ cdef object vec_cont1_array(irk_state *state, irk_cont1_vec func, object size,
261265
array_data = <double *>cnp.PyArray_DATA(array)
262266
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(2, <void *>array, <void *>oa)
263267
res_size = cnp.PyArray_SIZE(array)
264-
if (multi.size != res_size):
268+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
265269
raise ValueError("size is not compatible with inputs")
266270

267271
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(1, <void *>oa)
268-
imax = multi.size
272+
imax = cnp.PyArray_MultiIter_SIZE(multi)
269273
n = res_size // imax
270274
with lock, nogil:
271275
for i from 0 <= i < imax:
272276
oa_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 0)
273277
func(state, n, array_data + n*i, oa_data[0])
274278
cnp.PyArray_MultiIter_NEXT(multi)
275279
arr_obj = <object>array
276-
arr_obj.shape = (multi.shape + arr_obj.shape)[:arr_obj.ndim]
277-
multi_ndim = len(multi.shape)
280+
multi_nd = cnp.PyArray_MultiIter_NDIM(multi)
281+
multi_dims = cnp.PyArray_MultiIter_DIMS(multi)
282+
multi_shape = cpython.tuple.PyTuple_New(multi_nd)
283+
for i from 0 <= i < multi_nd:
284+
cpython.tuple.PyTuple_SetItem(multi_shape, i, multi_dims[i])
285+
arr_obj.shape = (multi_shape + arr_obj.shape)[:arr_obj.ndim]
286+
multi_ndim = len(multi_shape)
278287
arr_obj = arr_obj.transpose(tuple(range(multi_ndim, arr_obj.ndim)) + tuple(range(0, multi_ndim)))
279288

280289
return arr_obj
@@ -308,13 +317,20 @@ cdef object vec_cont2_array(irk_state *state, irk_cont2_vec func, object size,
308317
cdef cnp.npy_intp i, n, imax, res_size
309318
cdef cnp.broadcast multi
310319
cdef object arr_obj
320+
cdef Py_ssize_t multi_nd
321+
cdef tuple multi_shape
322+
cdef cnp.npy_intp *multi_dims
311323

312324
if size is None:
313325
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(2, <void *>oa, <void *>ob)
314-
array = <cnp.ndarray> cnp.PyArray_SimpleNew(multi.nd, multi.dimensions, cnp.NPY_DOUBLE)
326+
array = <cnp.ndarray> cnp.PyArray_SimpleNew(
327+
cnp.PyArray_MultiIter_NDIM(multi),
328+
cnp.PyArray_MultiIter_DIMS(multi),
329+
cnp.NPY_DOUBLE
330+
)
315331
array_data = <double *>cnp.PyArray_DATA(array)
316332
with lock, nogil:
317-
for i from 0 <= i < multi.size:
333+
for i from 0 <= i < cnp.PyArray_MultiIter_SIZE(multi):
318334
oa_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 0)
319335
ob_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 1)
320336
func(state, 1, &array_data[i], oa_data[0], ob_data[0])
@@ -324,12 +340,12 @@ cdef object vec_cont2_array(irk_state *state, irk_cont2_vec func, object size,
324340
array = <cnp.ndarray>np.empty(size, np.float64)
325341
array_data = <double *>cnp.PyArray_DATA(array)
326342
multi = <cnp.broadcast >cnp.PyArray_MultiIterNew(3, <void*>array, <void *>oa, <void *>ob)
327-
res_size = cnp.PyArray_SIZE(array);
328-
if (multi.size != res_size):
343+
res_size = cnp.PyArray_SIZE(array)
344+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
329345
raise ValueError("size is not compatible with inputs")
330346

331347
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(2, <void *>oa, <void *>ob)
332-
imax = multi.size
348+
imax = cnp.PyArray_MultiIter_SIZE(multi)
333349
n = res_size // imax
334350
with lock, nogil:
335351
for i from 0 <= i < imax:
@@ -338,8 +354,13 @@ cdef object vec_cont2_array(irk_state *state, irk_cont2_vec func, object size,
338354
func(state, n, array_data + n*i, oa_data[0], ob_data[0])
339355
cnp.PyArray_MultiIter_NEXT(multi)
340356
arr_obj = <object> array
341-
arr_obj.shape = (multi.shape + arr_obj.shape)[:arr_obj.ndim]
342-
multi_ndim = len(multi.shape)
357+
multi_nd = cnp.PyArray_MultiIter_NDIM(multi)
358+
multi_dims = cnp.PyArray_MultiIter_DIMS(multi)
359+
multi_shape = cpython.tuple.PyTuple_New(multi_nd)
360+
for i from 0 <= i < multi_nd:
361+
cpython.tuple.PyTuple_SetItem(multi_shape, i, multi_dims[i])
362+
arr_obj.shape = (multi_shape + arr_obj.shape)[:arr_obj.ndim]
363+
multi_ndim = len(multi_shape)
343364
arr_obj = arr_obj.transpose(tuple(range(multi_ndim, arr_obj.ndim)) + tuple(range(0, multi_ndim)))
344365

345366
return arr_obj
@@ -375,13 +396,16 @@ cdef object vec_cont3_array(irk_state *state, irk_cont3_vec func, object size,
375396
cdef cnp.npy_intp i, res_size, n, imax
376397
cdef cnp.broadcast multi
377398
cdef object arr_obj
399+
cdef Py_ssize_t multi_nd
400+
cdef tuple multi_shape
401+
cdef cnp.npy_intp *multi_dims
378402

379403
if size is None:
380404
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(3, <void *>oa, <void *>ob, <void *>oc)
381-
array = <cnp.ndarray> cnp.PyArray_SimpleNew(multi.nd, multi.dimensions, cnp.NPY_DOUBLE)
405+
array = <cnp.ndarray> cnp.PyArray_SimpleNew(cnp.PyArray_MultiIter_NDIM(multi), cnp.PyArray_MultiIter_DIMS(multi), cnp.NPY_DOUBLE)
382406
array_data = <double *>cnp.PyArray_DATA(array)
383407
with lock, nogil:
384-
for i from 0 <= i < multi.size:
408+
for i from 0 <= i < cnp.PyArray_MultiIter_SIZE(multi):
385409
oa_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 0)
386410
ob_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 1)
387411
oc_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 2)
@@ -394,11 +418,11 @@ cdef object vec_cont3_array(irk_state *state, irk_cont3_vec func, object size,
394418
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(4, <void*>array, <void *>oa,
395419
<void *>ob, <void *>oc)
396420
res_size = cnp.PyArray_SIZE(array)
397-
if (multi.size != res_size):
421+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
398422
raise ValueError("size is not compatible with inputs")
399423

400424
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(3, <void *>oa, <void *>ob, <void *>oc)
401-
imax = multi.size
425+
imax = cnp.PyArray_MultiIter_SIZE(multi)
402426
n = res_size // imax
403427
with lock, nogil:
404428
for i from 0 <= i < imax:
@@ -408,8 +432,13 @@ cdef object vec_cont3_array(irk_state *state, irk_cont3_vec func, object size,
408432
func(state, n, array_data + n*i, oa_data[0], ob_data[0], oc_data[0])
409433
cnp.PyArray_MultiIter_NEXT(multi)
410434
arr_obj = <object>array
411-
arr_obj.shape = (multi.shape + arr_obj.shape)[:arr_obj.ndim]
412-
multi_ndim = len(multi.shape)
435+
multi_nd = cnp.PyArray_MultiIter_NDIM(multi)
436+
multi_dims = cnp.PyArray_MultiIter_DIMS(multi)
437+
multi_shape = cpython.tuple.PyTuple_New(multi_nd)
438+
for i from 0 <= i < multi_nd:
439+
cpython.tuple.PyTuple_SetItem(multi_shape, i, multi_dims[i])
440+
arr_obj.shape = (multi_shape + arr_obj.shape)[:arr_obj.ndim]
441+
multi_ndim = len(multi_shape)
413442
arr_obj = arr_obj.transpose(tuple(range(multi_ndim, arr_obj.ndim)) + tuple(range(0, multi_ndim)))
414443

415444
return arr_obj
@@ -489,13 +518,16 @@ cdef object vec_discnp_array(irk_state *state, irk_discnp_vec func, object size,
489518
cdef int *on_data
490519
cdef cnp.broadcast multi
491520
cdef object arr_obj
521+
cdef Py_ssize_t multi_nd
522+
cdef tuple multi_shape
523+
cdef cnp.npy_intp *multi_dims
492524

493525
if size is None:
494526
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(2, <void *>on, <void *>op)
495-
array = <cnp.ndarray> cnp.PyArray_SimpleNew(multi.nd, multi.dimensions, cnp.NPY_INT)
527+
array = <cnp.ndarray> cnp.PyArray_SimpleNew(cnp.PyArray_MultiIter_NDIM(multi), cnp.PyArray_MultiIter_DIMS(multi), cnp.NPY_INT)
496528
array_data = <int *>cnp.PyArray_DATA(array)
497529
with lock, nogil:
498-
for i from 0 <= i < multi.size:
530+
for i from 0 <= i < cnp.PyArray_MultiIter_SIZE(multi):
499531
on_data = <int *>cnp.PyArray_MultiIter_DATA(multi, 0)
500532
op_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 1)
501533
func(state, 1, &array_data[i], on_data[0], op_data[0])
@@ -506,11 +538,11 @@ cdef object vec_discnp_array(irk_state *state, irk_discnp_vec func, object size,
506538
array_data = <int *>cnp.PyArray_DATA(array)
507539
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(3, <void*>array, <void *>on, <void *>op)
508540
res_size = cnp.PyArray_SIZE(array)
509-
if (multi.size != res_size):
541+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
510542
raise ValueError("size is not compatible with inputs")
511543

512544
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(2, <void *>on, <void *>op)
513-
imax = multi.size
545+
imax = cnp.PyArray_MultiIter_SIZE(multi)
514546
n = res_size // imax
515547
with lock, nogil:
516548
for i from 0 <= i < imax:
@@ -519,8 +551,13 @@ cdef object vec_discnp_array(irk_state *state, irk_discnp_vec func, object size,
519551
func(state, n, array_data + n * i, on_data[0], op_data[0])
520552
cnp.PyArray_MultiIter_NEXT(multi)
521553
arr_obj = <object>array
522-
arr_obj.shape = (multi.shape + arr_obj.shape)[:arr_obj.ndim]
523-
multi_ndim = len(multi.shape)
554+
multi_nd = cnp.PyArray_MultiIter_NDIM(multi)
555+
multi_dims = cnp.PyArray_MultiIter_DIMS(multi)
556+
multi_shape = cpython.tuple.PyTuple_New(multi_nd)
557+
for i from 0 <= i < multi_nd:
558+
cpython.tuple.PyTuple_SetItem(multi_shape, i, multi_dims[i])
559+
arr_obj.shape = (multi_shape + arr_obj.shape)[:arr_obj.ndim]
560+
multi_ndim = len(multi_shape)
524561
arr_obj = arr_obj.transpose(tuple(range(multi_ndim, arr_obj.ndim)) + tuple(range(0, multi_ndim)))
525562

526563
return arr_obj
@@ -556,13 +593,16 @@ cdef object vec_discdd_array(irk_state *state, irk_discdd_vec func, object size,
556593
cdef double *on_data
557594
cdef cnp.broadcast multi
558595
cdef object arr_obj
596+
cdef Py_ssize_t multi_nd
597+
cdef tuple multi_shape
598+
cdef cnp.npy_intp *multi_dims
559599

560600
if size is None:
561601
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(2, <void *>on, <void *>op)
562-
array = <cnp.ndarray> cnp.PyArray_SimpleNew(multi.nd, multi.dimensions, cnp.NPY_INT)
602+
array = <cnp.ndarray> cnp.PyArray_SimpleNew(cnp.PyArray_MultiIter_NDIM(multi), cnp.PyArray_MultiIter_DIMS(multi), cnp.NPY_INT)
563603
array_data = <int *>cnp.PyArray_DATA(array)
564604
with lock, nogil:
565-
for i from 0 <= i < multi.size:
605+
for i from 0 <= i < cnp.PyArray_MultiIter_SIZE(multi):
566606
on_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 0)
567607
op_data = <double *>cnp.PyArray_MultiIter_DATA(multi, 1)
568608
func(state, 1, &array_data[i], on_data[0], op_data[0])
@@ -573,11 +613,11 @@ cdef object vec_discdd_array(irk_state *state, irk_discdd_vec func, object size,
573613
array_data = <int *>cnp.PyArray_DATA(array)
574614
res_size = cnp.PyArray_SIZE(array)
575615
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(3, <void*>array, <void *>on, <void *>op)
576-
if (multi.size != res_size):
616+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
577617
raise ValueError("size is not compatible with inputs")
578618

579619
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(2, <void *>on, <void *>op)
580-
imax = multi.size
620+
imax = cnp.PyArray_MultiIter_SIZE(multi)
581621
n = res_size // imax
582622
with lock, nogil:
583623
for i from 0 <= i < imax:
@@ -586,8 +626,13 @@ cdef object vec_discdd_array(irk_state *state, irk_discdd_vec func, object size,
586626
func(state, n, array_data + n * i, on_data[0], op_data[0])
587627
cnp.PyArray_MultiIter_NEXT(multi)
588628
arr_obj = <object>array
589-
arr_obj.shape = (multi.shape + arr_obj.shape)[:arr_obj.ndim]
590-
multi_ndim = len(multi.shape)
629+
multi_nd = cnp.PyArray_MultiIter_NDIM(multi)
630+
multi_dims = cnp.PyArray_MultiIter_DIMS(multi)
631+
multi_shape = cpython.tuple.PyTuple_New(multi_nd)
632+
for i from 0 <= i < multi_nd:
633+
cpython.tuple.PyTuple_SetItem(multi_shape, i, multi_dims[i])
634+
arr_obj.shape = (multi_shape + arr_obj.shape)[:arr_obj.ndim]
635+
multi_ndim = len(multi_shape)
591636
arr_obj = arr_obj.transpose(tuple(range(multi_ndim, arr_obj.ndim)) + tuple(range(0, multi_ndim)))
592637

593638
return arr_obj
@@ -624,13 +669,16 @@ cdef object vec_discnmN_array(irk_state *state, irk_discnmN_vec func, object siz
624669
cdef cnp.broadcast multi, multi2
625670
cdef cnp.npy_intp imax, n, res_size
626671
cdef object arr_obj
672+
cdef Py_ssize_t multi_nd
673+
cdef tuple multi_shape
674+
cdef cnp.npy_intp *multi_dims
627675

628676
if size is None:
629677
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(3, <void *>on, <void *>om, <void *>oN)
630-
array = <cnp.ndarray> cnp.PyArray_SimpleNew(multi.nd, multi.dimensions, cnp.NPY_INT)
678+
array = <cnp.ndarray> cnp.PyArray_SimpleNew(cnp.PyArray_MultiIter_NDIM(multi), cnp.PyArray_MultiIter_DIMS(multi), cnp.NPY_INT)
631679
array_data = <int *>cnp.PyArray_DATA(array)
632680
with lock, nogil:
633-
for i from 0 <= i < multi.size:
681+
for i from 0 <= i < cnp.PyArray_MultiIter_SIZE(multi):
634682
on_data = <int *>cnp.PyArray_MultiIter_DATA(multi, 0)
635683
om_data = <int *>cnp.PyArray_MultiIter_DATA(multi, 1)
636684
oN_data = <int *>cnp.PyArray_MultiIter_DATA(multi, 2)
@@ -643,11 +691,11 @@ cdef object vec_discnmN_array(irk_state *state, irk_discnmN_vec func, object siz
643691
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(4, <void*>array, <void *>on, <void *>om,
644692
<void *>oN)
645693
res_size = cnp.PyArray_SIZE(array)
646-
if (multi.size != res_size):
694+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
647695
raise ValueError("size is not compatible with inputs")
648696

649697
multi = <cnp.broadcast> cnp.PyArray_MultiIterNew(3, <void *>on, <void *>om, <void *>oN)
650-
imax = multi.size
698+
imax = cnp.PyArray_MultiIter_SIZE(multi)
651699
n = res_size // imax
652700
with lock, nogil:
653701
for i from 0 <= i < imax:
@@ -657,8 +705,13 @@ cdef object vec_discnmN_array(irk_state *state, irk_discnmN_vec func, object siz
657705
func(state, n, array_data + n*i, on_data[0], om_data[0], oN_data[0])
658706
cnp.PyArray_MultiIter_NEXT(multi)
659707
arr_obj = <object>array
660-
arr_obj.shape = (multi.shape + arr_obj.shape)[:arr_obj.ndim]
661-
multi_ndim = len(multi.shape)
708+
multi_nd = cnp.PyArray_MultiIter_NDIM(multi)
709+
multi_dims = cnp.PyArray_MultiIter_DIMS(multi)
710+
multi_shape = cpython.tuple.PyTuple_New(multi_nd)
711+
for i from 0 <= i < multi_nd:
712+
cpython.tuple.PyTuple_SetItem(multi_shape, i, multi_dims[i])
713+
arr_obj.shape = (multi_shape + arr_obj.shape)[:arr_obj.ndim]
714+
multi_ndim = len(multi_shape)
662715
arr_obj = arr_obj.transpose(tuple(range(multi_ndim, arr_obj.ndim)) + tuple(range(0, multi_ndim)))
663716

664717
return arr_obj
@@ -730,7 +783,7 @@ cdef object vec_discd_array(irk_state *state, irk_discd_vec func, object size, c
730783
array_data = <int *>cnp.PyArray_DATA(array)
731784
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(2, <void *>array, <void *>oa)
732785
res_size = cnp.PyArray_SIZE(array)
733-
if (multi.size != res_size):
786+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
734787
raise ValueError("size is not compatible with inputs")
735788

736789
imax = oa.size
@@ -773,7 +826,7 @@ cdef object vec_long_discd_array(irk_state *state, irk_discd_long_vec func, obje
773826
array_data = <long *>cnp.PyArray_DATA(array)
774827
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(2, <void *>array, <void *>oa)
775828
res_size = cnp.PyArray_SIZE(array)
776-
if (multi.size != res_size):
829+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
777830
raise ValueError("size is not compatible with inputs")
778831

779832
imax = oa.size
@@ -814,7 +867,7 @@ cdef object vec_Poisson_array(irk_state *state, irk_discdptr_vec func1, irk_disc
814867
array_data = <int *>cnp.PyArray_DATA(array)
815868
multi = <cnp.broadcast>cnp.PyArray_MultiIterNew(2, <void *>array, <void *>olambda)
816869
res_size = cnp.PyArray_SIZE(array)
817-
if (multi.size != res_size):
870+
if (cnp.PyArray_MultiIter_SIZE(multi) != res_size):
818871
raise ValueError("size is not compatible with inputs")
819872

820873
imax = olambda.size

0 commit comments

Comments
 (0)