extending_distributions.pyx 3.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116
  1. #cython: language_level=3
  2. """
  3. This file shows how the to use a BitGenerator to create a distribution.
  4. """
  5. import numpy as np
  6. cimport numpy as np
  7. cimport cython
  8. from cpython.pycapsule cimport PyCapsule_IsValid, PyCapsule_GetPointer
  9. from libc.stdint cimport uint16_t, uint64_t
  10. from numpy.random cimport bitgen_t
  11. from numpy.random import PCG64
  12. from numpy.random.c_distributions cimport (
  13. random_standard_uniform_fill, random_standard_uniform_fill_f)
  14. @cython.boundscheck(False)
  15. @cython.wraparound(False)
  16. def uniforms(Py_ssize_t n):
  17. """
  18. Create an array of `n` uniformly distributed doubles.
  19. A 'real' distribution would want to process the values into
  20. some non-uniform distribution
  21. """
  22. cdef Py_ssize_t i
  23. cdef bitgen_t *rng
  24. cdef const char *capsule_name = "BitGenerator"
  25. cdef double[::1] random_values
  26. x = PCG64()
  27. capsule = x.capsule
  28. # Optional check that the capsule if from a BitGenerator
  29. if not PyCapsule_IsValid(capsule, capsule_name):
  30. raise ValueError("Invalid pointer to anon_func_state")
  31. # Cast the pointer
  32. rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
  33. random_values = np.empty(n, dtype='float64')
  34. with x.lock, nogil:
  35. for i in range(n):
  36. # Call the function
  37. random_values[i] = rng.next_double(rng.state)
  38. randoms = np.asarray(random_values)
  39. return randoms
  40. # cython example 2
  41. @cython.boundscheck(False)
  42. @cython.wraparound(False)
  43. def uint10_uniforms(Py_ssize_t n):
  44. """Uniform 10 bit integers stored as 16-bit unsigned integers"""
  45. cdef Py_ssize_t i
  46. cdef bitgen_t *rng
  47. cdef const char *capsule_name = "BitGenerator"
  48. cdef uint16_t[::1] random_values
  49. cdef int bits_remaining
  50. cdef int width = 10
  51. cdef uint64_t buff, mask = 0x3FF
  52. x = PCG64()
  53. capsule = x.capsule
  54. if not PyCapsule_IsValid(capsule, capsule_name):
  55. raise ValueError("Invalid pointer to anon_func_state")
  56. rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
  57. random_values = np.empty(n, dtype='uint16')
  58. # Best practice is to release GIL and acquire the lock
  59. bits_remaining = 0
  60. with x.lock, nogil:
  61. for i in range(n):
  62. if bits_remaining < width:
  63. buff = rng.next_uint64(rng.state)
  64. random_values[i] = buff & mask
  65. buff >>= width
  66. randoms = np.asarray(random_values)
  67. return randoms
  68. # cython example 3
  69. def uniforms_ex(bit_generator, Py_ssize_t n, dtype=np.float64):
  70. """
  71. Create an array of `n` uniformly distributed doubles via a "fill" function.
  72. A 'real' distribution would want to process the values into
  73. some non-uniform distribution
  74. Parameters
  75. ----------
  76. bit_generator: BitGenerator instance
  77. n: int
  78. Output vector length
  79. dtype: {str, dtype}, optional
  80. Desired dtype, either 'd' (or 'float64') or 'f' (or 'float32'). The
  81. default dtype value is 'd'
  82. """
  83. cdef Py_ssize_t i
  84. cdef bitgen_t *rng
  85. cdef const char *capsule_name = "BitGenerator"
  86. cdef np.ndarray randoms
  87. capsule = bit_generator.capsule
  88. # Optional check that the capsule if from a BitGenerator
  89. if not PyCapsule_IsValid(capsule, capsule_name):
  90. raise ValueError("Invalid pointer to anon_func_state")
  91. # Cast the pointer
  92. rng = <bitgen_t *> PyCapsule_GetPointer(capsule, capsule_name)
  93. _dtype = np.dtype(dtype)
  94. randoms = np.empty(n, dtype=_dtype)
  95. if _dtype == np.float32:
  96. with bit_generator.lock:
  97. random_standard_uniform_fill_f(rng, n, <float*>np.PyArray_DATA(randoms))
  98. elif _dtype == np.float64:
  99. with bit_generator.lock:
  100. random_standard_uniform_fill(rng, n, <double*>np.PyArray_DATA(randoms))
  101. else:
  102. raise TypeError('Unsupported dtype %r for random' % _dtype)
  103. return randoms