diff --git a/brainpy/_src/math/random.py b/brainpy/_src/math/random.py index 180a023f5..0238f81aa 100644 --- a/brainpy/_src/math/random.py +++ b/brainpy/_src/math/random.py @@ -42,6 +42,7 @@ # taichi.func random generator implementation 'taichi_lcg_rand', 'taichi_uniform_int_distribution', 'taichi_uniform_real_distribution', 'taichi_normal_distribution', + 'taichi_lfsr88', ] @@ -2423,9 +2424,7 @@ def randint_like(input, low=0, high=None, *, dtype=None, key=None): @ti.func def _lcg_rand(state: ti.types.ndarray(ndim=1)): # LCG constants - a = ti.u32(1664525) - c = ti.u32(1013904223) - state[0] = a * state[0] + c + state[0] = ti.u32(1664525) * state[0] + ti.u32(1013904223) return state[0] @ti.func @@ -2439,9 +2438,42 @@ def taichi_lcg_rand(seed: ti.types.ndarray(ndim=1)): Returns: float: A random number between 0 and 1. """ - m = ti.u32(2**32 - 1) - return float(_lcg_rand(seed)) / m + return float(_lcg_rand(seed)) / ti.u32(2**32 - 1) + + @ti.func + def taichi_lfsr88(s1: ti.u32, s2: ti.u32, s3: ti.u32, b: ti.u32): + """ + 32-bits Random number generator U[0,1): lfsr88 + Author: Pierre L'Ecuyer, + Source: + https://github.com/cmcqueen/simplerandom/blob/main/c/lecuyer/lfsr88.c + /**** VERY IMPORTANT **** : + The initial seeds s1, s2, s3 MUST be larger than + 1, 7, and 15 respectively. + */ + ```cpp + double taus88_double () + { /* Generates numbers between 0 and 1. */ + b = (((s1 << 13) ^ s1) >> 19); + s1 = (((s1 & 4294967294) << 12) ^ b); + b = (((s2 << 2) ^ s2) >> 25); + s2 = (((s2 & 4294967288) << 4) ^ b); + b = (((s3 << 3) ^ s3) >> 11); + s3 = (((s3 & 4294967280) << 17) ^ b); + return ((s1 ^ s2 ^ s3) * 2.3283064365386963e-10); + } + ``` + + """ + b = (((s1 << 13) ^ s1) >> 19); + s1 = (((s1 & ti.u32(4294967294)) << 12) ^ b) + b = (((s2 << 2) ^ s2) >> 25) + s2 = (((s2 & ti.u32(4294967288)) << 4) ^ b) + b = (((s3 << 3) ^ s3) >> 11) + s3 = (((s3 & ti.u32(4294967280)) << 17) ^ b) + return s1, s2, s3, b, ((s1 ^ s2 ^ s3) * ti.f32(2.3283064365386963e-10)) + @ti.func diff --git a/brainpy/_src/math/tests/test_taichi_random.py b/brainpy/_src/math/tests/test_taichi_random.py index c4e10f4cb..ea6b559ea 100644 --- a/brainpy/_src/math/tests/test_taichi_random.py +++ b/brainpy/_src/math/tests/test_taichi_random.py @@ -15,11 +15,23 @@ from brainpy.math.random import (taichi_lcg_rand as rand, taichi_uniform_int_distribution as randint, taichi_uniform_real_distribution as uniform, - taichi_normal_distribution as normal,) + taichi_normal_distribution as normal, + taichi_lfsr88) bm.set_platform('cpu') def test_taichi_random(): + @ti.kernel + def test_taichi_lfsr88(seed: ti.types.ndarray(ndim=1, dtype=ti.u32), + out: ti.types.ndarray(ndim=1, dtype=ti.f32)): + s1 = seed[0] + 1 + s2 = seed[0] + 7 + s3 = seed[0] + 15 + b = ti.u32(0) + for i in range(out.shape[0]): + s1, s2, s3, b, result = taichi_lfsr88(s1, s2, s3, b) + out[i] = result + @ti.kernel def test_taichi_lcg_rand(seed: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): @@ -30,33 +42,44 @@ def test_taichi_lcg_rand(seed: ti.types.ndarray(ndim=1), def test_taichi_uniform_int_distribution(seed: ti.types.ndarray(ndim=1), low_high: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): + s1 = seed[0] + 1 + s2 = seed[0] + 7 + s3 = seed[0] + 15 + b = ti.u32(0) low = low_high[0] high = low_high[1] - n = out.shape[0] - for i in range(n): - out[i] = randint(rand(seed), low, high) + for i in range(out.shape[0]): + s1, s2, s3, b, result = taichi_lfsr88(s1, s2, s3, b) + out[i] = randint(result, low, high) @ti.kernel def test_taichi_uniform_real_distribution(seed: ti.types.ndarray(ndim=1), low_high: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): + s1 = seed[0] + 1 + s2 = seed[0] + 7 + s3 = seed[0] + 15 + b = ti.u32(0) low = low_high[0] high = low_high[1] - n = out.shape[0] - for i in range(n): - out[i] = uniform(rand(seed), low, high) + for i in range(out.shape[0]): + s1, s2, s3, b, result = taichi_lfsr88(s1, s2, s3, b) + out[i] = uniform(result, low, high) @ti.kernel def test_taichi_normal_distribution(seed: ti.types.ndarray(ndim=1), mu_sigma: ti.types.ndarray(ndim=1), out: ti.types.ndarray(ndim=1)): + s1 = seed[0] + 1 + s2 = seed[0] + 7 + s3 = seed[0] + 15 + b = ti.u32(0) mu = mu_sigma[0] sigma = mu_sigma[1] - n = out.shape[0] - for i in range(n): - r1 = rand(seed) - r2 = rand(seed) + for i in range(out.shape[0]): + s1, s2, s3, b, r1 = taichi_lfsr88(s1, s2, s3, b) + s1, s2, s3, b, r2 = taichi_lfsr88(s1, s2, s3, b) out[i] = normal(r1, r2 , mu, sigma) @@ -65,6 +88,9 @@ def test_taichi_normal_distribution(seed: ti.types.ndarray(ndim=1), low_high = jnp.array([0, 10]) mu_sigma = jnp.array([0, 1]) + prim_lfsr88 = bm.XLACustomOp(cpu_kernel=test_taichi_lfsr88, + gpu_kernel=test_taichi_lfsr88) + prim_lcg_rand = bm.XLACustomOp(cpu_kernel=test_taichi_lcg_rand, gpu_kernel=test_taichi_lcg_rand) prim_uniform_int_distribution = bm.XLACustomOp(cpu_kernel=test_taichi_uniform_int_distribution, @@ -76,6 +102,13 @@ def test_taichi_normal_distribution(seed: ti.types.ndarray(ndim=1), file_path = os.path.dirname(os.path.abspath(__file__)) + out = prim_lfsr88(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)]) + # show the distribution of out + plt.hist(out, bins=100) + plt.title("LFSR88 random number generator") + plt.savefig(file_path + "/lfsr88.png") + plt.close() + out = prim_lcg_rand(seed, outs=[jax.ShapeDtypeStruct((n,), jnp.float32)]) # show the distribution of out diff --git a/brainpy/math/random.py b/brainpy/math/random.py index c1239c2f8..9308e2deb 100644 --- a/brainpy/math/random.py +++ b/brainpy/math/random.py @@ -75,5 +75,6 @@ taichi_uniform_int_distribution as taichi_uniform_int_distribution, taichi_uniform_real_distribution as taichi_uniform_real_distribution, taichi_normal_distribution as taichi_normal_distribution, + taichi_lfsr88 as taichi_lfsr88, )