test_direct.py 20 KB

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