test_direct.py 17 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518
  1. import os
  2. from os.path import join
  3. import sys
  4. import numpy as np
  5. from numpy.testing import (assert_equal, assert_allclose, assert_array_equal,
  6. assert_raises)
  7. import pytest
  8. from numpy.random import (
  9. Generator, MT19937, PCG64, PCG64DXSM, Philox, RandomState, SeedSequence,
  10. SFC64, default_rng
  11. )
  12. from numpy.random._common import interface
  13. try:
  14. import cffi # noqa: F401
  15. MISSING_CFFI = False
  16. except ImportError:
  17. MISSING_CFFI = True
  18. try:
  19. import ctypes # noqa: F401
  20. MISSING_CTYPES = False
  21. except ImportError:
  22. MISSING_CTYPES = False
  23. if sys.flags.optimize > 1:
  24. # no docstrings present to inspect when PYTHONOPTIMIZE/Py_OptimizeFlag > 1
  25. # cffi cannot succeed
  26. MISSING_CFFI = True
  27. pwd = os.path.dirname(os.path.abspath(__file__))
  28. def assert_state_equal(actual, target):
  29. for key in actual:
  30. if isinstance(actual[key], dict):
  31. assert_state_equal(actual[key], target[key])
  32. elif isinstance(actual[key], np.ndarray):
  33. assert_array_equal(actual[key], target[key])
  34. else:
  35. assert actual[key] == target[key]
  36. def uint32_to_float32(u):
  37. return ((u >> np.uint32(8)) * (1.0 / 2**24)).astype(np.float32)
  38. def uniform32_from_uint64(x):
  39. x = np.uint64(x)
  40. upper = np.array(x >> np.uint64(32), dtype=np.uint32)
  41. lower = np.uint64(0xffffffff)
  42. lower = np.array(x & lower, dtype=np.uint32)
  43. joined = np.column_stack([lower, upper]).ravel()
  44. return uint32_to_float32(joined)
  45. def uniform32_from_uint53(x):
  46. x = np.uint64(x) >> np.uint64(16)
  47. x = np.uint32(x & np.uint64(0xffffffff))
  48. return uint32_to_float32(x)
  49. def uniform32_from_uint32(x):
  50. return uint32_to_float32(x)
  51. def uniform32_from_uint(x, bits):
  52. if bits == 64:
  53. return uniform32_from_uint64(x)
  54. elif bits == 53:
  55. return uniform32_from_uint53(x)
  56. elif bits == 32:
  57. return uniform32_from_uint32(x)
  58. else:
  59. raise NotImplementedError
  60. def uniform_from_uint(x, bits):
  61. if bits in (64, 63, 53):
  62. return uniform_from_uint64(x)
  63. elif bits == 32:
  64. return uniform_from_uint32(x)
  65. def uniform_from_uint64(x):
  66. return (x >> np.uint64(11)) * (1.0 / 9007199254740992.0)
  67. def uniform_from_uint32(x):
  68. out = np.empty(len(x) // 2)
  69. for i in range(0, len(x), 2):
  70. a = x[i] >> 5
  71. b = x[i + 1] >> 6
  72. out[i // 2] = (a * 67108864.0 + b) / 9007199254740992.0
  73. return out
  74. def uniform_from_dsfmt(x):
  75. return x.view(np.double) - 1.0
  76. def gauss_from_uint(x, n, bits):
  77. if bits in (64, 63):
  78. doubles = uniform_from_uint64(x)
  79. elif bits == 32:
  80. doubles = uniform_from_uint32(x)
  81. else: # bits == 'dsfmt'
  82. doubles = uniform_from_dsfmt(x)
  83. gauss = []
  84. loc = 0
  85. x1 = x2 = 0.0
  86. while len(gauss) < n:
  87. r2 = 2
  88. while r2 >= 1.0 or r2 == 0.0:
  89. x1 = 2.0 * doubles[loc] - 1.0
  90. x2 = 2.0 * doubles[loc + 1] - 1.0
  91. r2 = x1 * x1 + x2 * x2
  92. loc += 2
  93. f = np.sqrt(-2.0 * np.log(r2) / r2)
  94. gauss.append(f * x2)
  95. gauss.append(f * x1)
  96. return gauss[:n]
  97. def test_seedsequence():
  98. from numpy.random.bit_generator import (ISeedSequence,
  99. ISpawnableSeedSequence,
  100. SeedlessSeedSequence)
  101. s1 = SeedSequence(range(10), spawn_key=(1, 2), pool_size=6)
  102. s1.spawn(10)
  103. s2 = SeedSequence(**s1.state)
  104. assert_equal(s1.state, s2.state)
  105. assert_equal(s1.n_children_spawned, s2.n_children_spawned)
  106. # The interfaces cannot be instantiated themselves.
  107. assert_raises(TypeError, ISeedSequence)
  108. assert_raises(TypeError, ISpawnableSeedSequence)
  109. dummy = SeedlessSeedSequence()
  110. assert_raises(NotImplementedError, dummy.generate_state, 10)
  111. assert len(dummy.spawn(10)) == 10
  112. def test_generator_spawning():
  113. """ Test spawning new generators and bit_generators directly.
  114. """
  115. rng = np.random.default_rng()
  116. seq = rng.bit_generator.seed_seq
  117. new_ss = seq.spawn(5)
  118. expected_keys = [seq.spawn_key + (i,) for i in range(5)]
  119. assert [c.spawn_key for c in new_ss] == expected_keys
  120. new_bgs = rng.bit_generator.spawn(5)
  121. expected_keys = [seq.spawn_key + (i,) for i in range(5, 10)]
  122. assert [bg.seed_seq.spawn_key for bg in new_bgs] == expected_keys
  123. new_rngs = rng.spawn(5)
  124. expected_keys = [seq.spawn_key + (i,) for i in range(10, 15)]
  125. found_keys = [rng.bit_generator.seed_seq.spawn_key for rng in new_rngs]
  126. assert found_keys == expected_keys
  127. # Sanity check that streams are actually different:
  128. assert new_rngs[0].uniform() != new_rngs[1].uniform()
  129. def test_non_spawnable():
  130. from numpy.random.bit_generator import ISeedSequence
  131. class FakeSeedSequence:
  132. def generate_state(self, n_words, dtype=np.uint32):
  133. return np.zeros(n_words, dtype=dtype)
  134. ISeedSequence.register(FakeSeedSequence)
  135. rng = np.random.default_rng(FakeSeedSequence())
  136. with pytest.raises(TypeError, match="The underlying SeedSequence"):
  137. rng.spawn(5)
  138. with pytest.raises(TypeError, match="The underlying SeedSequence"):
  139. rng.bit_generator.spawn(5)
  140. class Base:
  141. dtype = np.uint64
  142. data2 = data1 = {}
  143. @classmethod
  144. def setup_class(cls):
  145. cls.bit_generator = PCG64
  146. cls.bits = 64
  147. cls.dtype = np.uint64
  148. cls.seed_error_type = TypeError
  149. cls.invalid_init_types = []
  150. cls.invalid_init_values = []
  151. @classmethod
  152. def _read_csv(cls, filename):
  153. with open(filename) as csv:
  154. seed = csv.readline()
  155. seed = seed.split(',')
  156. seed = [int(s.strip(), 0) for s in seed[1:]]
  157. data = []
  158. for line in csv:
  159. data.append(int(line.split(',')[-1].strip(), 0))
  160. return {'seed': seed, 'data': np.array(data, dtype=cls.dtype)}
  161. def test_raw(self):
  162. bit_generator = self.bit_generator(*self.data1['seed'])
  163. uints = bit_generator.random_raw(1000)
  164. assert_equal(uints, self.data1['data'])
  165. bit_generator = self.bit_generator(*self.data1['seed'])
  166. uints = bit_generator.random_raw()
  167. assert_equal(uints, self.data1['data'][0])
  168. bit_generator = self.bit_generator(*self.data2['seed'])
  169. uints = bit_generator.random_raw(1000)
  170. assert_equal(uints, self.data2['data'])
  171. def test_random_raw(self):
  172. bit_generator = self.bit_generator(*self.data1['seed'])
  173. uints = bit_generator.random_raw(output=False)
  174. assert uints is None
  175. uints = bit_generator.random_raw(1000, output=False)
  176. assert uints is None
  177. def test_gauss_inv(self):
  178. n = 25
  179. rs = RandomState(self.bit_generator(*self.data1['seed']))
  180. gauss = rs.standard_normal(n)
  181. assert_allclose(gauss,
  182. gauss_from_uint(self.data1['data'], n, self.bits))
  183. rs = RandomState(self.bit_generator(*self.data2['seed']))
  184. gauss = rs.standard_normal(25)
  185. assert_allclose(gauss,
  186. gauss_from_uint(self.data2['data'], n, self.bits))
  187. def test_uniform_double(self):
  188. rs = Generator(self.bit_generator(*self.data1['seed']))
  189. vals = uniform_from_uint(self.data1['data'], self.bits)
  190. uniforms = rs.random(len(vals))
  191. assert_allclose(uniforms, vals)
  192. assert_equal(uniforms.dtype, np.float64)
  193. rs = Generator(self.bit_generator(*self.data2['seed']))
  194. vals = uniform_from_uint(self.data2['data'], self.bits)
  195. uniforms = rs.random(len(vals))
  196. assert_allclose(uniforms, vals)
  197. assert_equal(uniforms.dtype, np.float64)
  198. def test_uniform_float(self):
  199. rs = Generator(self.bit_generator(*self.data1['seed']))
  200. vals = uniform32_from_uint(self.data1['data'], self.bits)
  201. uniforms = rs.random(len(vals), dtype=np.float32)
  202. assert_allclose(uniforms, vals)
  203. assert_equal(uniforms.dtype, np.float32)
  204. rs = Generator(self.bit_generator(*self.data2['seed']))
  205. vals = uniform32_from_uint(self.data2['data'], self.bits)
  206. uniforms = rs.random(len(vals), dtype=np.float32)
  207. assert_allclose(uniforms, vals)
  208. assert_equal(uniforms.dtype, np.float32)
  209. def test_repr(self):
  210. rs = Generator(self.bit_generator(*self.data1['seed']))
  211. assert 'Generator' in repr(rs)
  212. assert f'{id(rs):#x}'.upper().replace('X', 'x') in repr(rs)
  213. def test_str(self):
  214. rs = Generator(self.bit_generator(*self.data1['seed']))
  215. assert 'Generator' in str(rs)
  216. assert str(self.bit_generator.__name__) in str(rs)
  217. assert f'{id(rs):#x}'.upper().replace('X', 'x') not in str(rs)
  218. def test_pickle(self):
  219. import pickle
  220. bit_generator = self.bit_generator(*self.data1['seed'])
  221. state = bit_generator.state
  222. bitgen_pkl = pickle.dumps(bit_generator)
  223. reloaded = pickle.loads(bitgen_pkl)
  224. reloaded_state = reloaded.state
  225. assert_array_equal(Generator(bit_generator).standard_normal(1000),
  226. Generator(reloaded).standard_normal(1000))
  227. assert bit_generator is not reloaded
  228. assert_state_equal(reloaded_state, state)
  229. ss = SeedSequence(100)
  230. aa = pickle.loads(pickle.dumps(ss))
  231. assert_equal(ss.state, aa.state)
  232. def test_invalid_state_type(self):
  233. bit_generator = self.bit_generator(*self.data1['seed'])
  234. with pytest.raises(TypeError):
  235. bit_generator.state = {'1'}
  236. def test_invalid_state_value(self):
  237. bit_generator = self.bit_generator(*self.data1['seed'])
  238. state = bit_generator.state
  239. state['bit_generator'] = 'otherBitGenerator'
  240. with pytest.raises(ValueError):
  241. bit_generator.state = state
  242. def test_invalid_init_type(self):
  243. bit_generator = self.bit_generator
  244. for st in self.invalid_init_types:
  245. with pytest.raises(TypeError):
  246. bit_generator(*st)
  247. def test_invalid_init_values(self):
  248. bit_generator = self.bit_generator
  249. for st in self.invalid_init_values:
  250. with pytest.raises((ValueError, OverflowError)):
  251. bit_generator(*st)
  252. def test_benchmark(self):
  253. bit_generator = self.bit_generator(*self.data1['seed'])
  254. bit_generator._benchmark(1)
  255. bit_generator._benchmark(1, 'double')
  256. with pytest.raises(ValueError):
  257. bit_generator._benchmark(1, 'int32')
  258. @pytest.mark.skipif(MISSING_CFFI, reason='cffi not available')
  259. def test_cffi(self):
  260. bit_generator = self.bit_generator(*self.data1['seed'])
  261. cffi_interface = bit_generator.cffi
  262. assert isinstance(cffi_interface, interface)
  263. other_cffi_interface = bit_generator.cffi
  264. assert other_cffi_interface is cffi_interface
  265. @pytest.mark.skipif(MISSING_CTYPES, reason='ctypes not available')
  266. def test_ctypes(self):
  267. bit_generator = self.bit_generator(*self.data1['seed'])
  268. ctypes_interface = bit_generator.ctypes
  269. assert isinstance(ctypes_interface, interface)
  270. other_ctypes_interface = bit_generator.ctypes
  271. assert other_ctypes_interface is ctypes_interface
  272. def test_getstate(self):
  273. bit_generator = self.bit_generator(*self.data1['seed'])
  274. state = bit_generator.state
  275. alt_state = bit_generator.__getstate__()
  276. assert_state_equal(state, alt_state)
  277. class TestPhilox(Base):
  278. @classmethod
  279. def setup_class(cls):
  280. cls.bit_generator = Philox
  281. cls.bits = 64
  282. cls.dtype = np.uint64
  283. cls.data1 = cls._read_csv(
  284. join(pwd, './data/philox-testset-1.csv'))
  285. cls.data2 = cls._read_csv(
  286. join(pwd, './data/philox-testset-2.csv'))
  287. cls.seed_error_type = TypeError
  288. cls.invalid_init_types = []
  289. cls.invalid_init_values = [(1, None, 1), (-1,), (None, None, 2 ** 257 + 1)]
  290. def test_set_key(self):
  291. bit_generator = self.bit_generator(*self.data1['seed'])
  292. state = bit_generator.state
  293. keyed = self.bit_generator(counter=state['state']['counter'],
  294. key=state['state']['key'])
  295. assert_state_equal(bit_generator.state, keyed.state)
  296. class TestPCG64(Base):
  297. @classmethod
  298. def setup_class(cls):
  299. cls.bit_generator = PCG64
  300. cls.bits = 64
  301. cls.dtype = np.uint64
  302. cls.data1 = cls._read_csv(join(pwd, './data/pcg64-testset-1.csv'))
  303. cls.data2 = cls._read_csv(join(pwd, './data/pcg64-testset-2.csv'))
  304. cls.seed_error_type = (ValueError, TypeError)
  305. cls.invalid_init_types = [(3.2,), ([None],), (1, None)]
  306. cls.invalid_init_values = [(-1,)]
  307. def test_advance_symmetry(self):
  308. rs = Generator(self.bit_generator(*self.data1['seed']))
  309. state = rs.bit_generator.state
  310. step = -0x9e3779b97f4a7c150000000000000000
  311. rs.bit_generator.advance(step)
  312. val_neg = rs.integers(10)
  313. rs.bit_generator.state = state
  314. rs.bit_generator.advance(2**128 + step)
  315. val_pos = rs.integers(10)
  316. rs.bit_generator.state = state
  317. rs.bit_generator.advance(10 * 2**128 + step)
  318. val_big = rs.integers(10)
  319. assert val_neg == val_pos
  320. assert val_big == val_pos
  321. def test_advange_large(self):
  322. rs = Generator(self.bit_generator(38219308213743))
  323. pcg = rs.bit_generator
  324. state = pcg.state["state"]
  325. initial_state = 287608843259529770491897792873167516365
  326. assert state["state"] == initial_state
  327. pcg.advance(sum(2**i for i in (96, 64, 32, 16, 8, 4, 2, 1)))
  328. state = pcg.state["state"]
  329. advanced_state = 135275564607035429730177404003164635391
  330. assert state["state"] == advanced_state
  331. class TestPCG64DXSM(Base):
  332. @classmethod
  333. def setup_class(cls):
  334. cls.bit_generator = PCG64DXSM
  335. cls.bits = 64
  336. cls.dtype = np.uint64
  337. cls.data1 = cls._read_csv(join(pwd, './data/pcg64dxsm-testset-1.csv'))
  338. cls.data2 = cls._read_csv(join(pwd, './data/pcg64dxsm-testset-2.csv'))
  339. cls.seed_error_type = (ValueError, TypeError)
  340. cls.invalid_init_types = [(3.2,), ([None],), (1, None)]
  341. cls.invalid_init_values = [(-1,)]
  342. def test_advance_symmetry(self):
  343. rs = Generator(self.bit_generator(*self.data1['seed']))
  344. state = rs.bit_generator.state
  345. step = -0x9e3779b97f4a7c150000000000000000
  346. rs.bit_generator.advance(step)
  347. val_neg = rs.integers(10)
  348. rs.bit_generator.state = state
  349. rs.bit_generator.advance(2**128 + step)
  350. val_pos = rs.integers(10)
  351. rs.bit_generator.state = state
  352. rs.bit_generator.advance(10 * 2**128 + step)
  353. val_big = rs.integers(10)
  354. assert val_neg == val_pos
  355. assert val_big == val_pos
  356. def test_advange_large(self):
  357. rs = Generator(self.bit_generator(38219308213743))
  358. pcg = rs.bit_generator
  359. state = pcg.state
  360. initial_state = 287608843259529770491897792873167516365
  361. assert state["state"]["state"] == initial_state
  362. pcg.advance(sum(2**i for i in (96, 64, 32, 16, 8, 4, 2, 1)))
  363. state = pcg.state["state"]
  364. advanced_state = 277778083536782149546677086420637664879
  365. assert state["state"] == advanced_state
  366. class TestMT19937(Base):
  367. @classmethod
  368. def setup_class(cls):
  369. cls.bit_generator = MT19937
  370. cls.bits = 32
  371. cls.dtype = np.uint32
  372. cls.data1 = cls._read_csv(join(pwd, './data/mt19937-testset-1.csv'))
  373. cls.data2 = cls._read_csv(join(pwd, './data/mt19937-testset-2.csv'))
  374. cls.seed_error_type = ValueError
  375. cls.invalid_init_types = []
  376. cls.invalid_init_values = [(-1,)]
  377. def test_seed_float_array(self):
  378. assert_raises(TypeError, self.bit_generator, np.array([np.pi]))
  379. assert_raises(TypeError, self.bit_generator, np.array([-np.pi]))
  380. assert_raises(TypeError, self.bit_generator, np.array([np.pi, -np.pi]))
  381. assert_raises(TypeError, self.bit_generator, np.array([0, np.pi]))
  382. assert_raises(TypeError, self.bit_generator, [np.pi])
  383. assert_raises(TypeError, self.bit_generator, [0, np.pi])
  384. def test_state_tuple(self):
  385. rs = Generator(self.bit_generator(*self.data1['seed']))
  386. bit_generator = rs.bit_generator
  387. state = bit_generator.state
  388. desired = rs.integers(2 ** 16)
  389. tup = (state['bit_generator'], state['state']['key'],
  390. state['state']['pos'])
  391. bit_generator.state = tup
  392. actual = rs.integers(2 ** 16)
  393. assert_equal(actual, desired)
  394. tup = tup + (0, 0.0)
  395. bit_generator.state = tup
  396. actual = rs.integers(2 ** 16)
  397. assert_equal(actual, desired)
  398. class TestSFC64(Base):
  399. @classmethod
  400. def setup_class(cls):
  401. cls.bit_generator = SFC64
  402. cls.bits = 64
  403. cls.dtype = np.uint64
  404. cls.data1 = cls._read_csv(
  405. join(pwd, './data/sfc64-testset-1.csv'))
  406. cls.data2 = cls._read_csv(
  407. join(pwd, './data/sfc64-testset-2.csv'))
  408. cls.seed_error_type = (ValueError, TypeError)
  409. cls.invalid_init_types = [(3.2,), ([None],), (1, None)]
  410. cls.invalid_init_values = [(-1,)]
  411. class TestDefaultRNG:
  412. def test_seed(self):
  413. for args in [(), (None,), (1234,), ([1234, 5678],)]:
  414. rg = default_rng(*args)
  415. assert isinstance(rg.bit_generator, PCG64)
  416. def test_passthrough(self):
  417. bg = Philox()
  418. rg = default_rng(bg)
  419. assert rg.bit_generator is bg
  420. rg2 = default_rng(rg)
  421. assert rg2 is rg
  422. assert rg2.bit_generator is bg