Skip to content

Commit

Permalink
Implement lfsr88 random generator algorithm
Browse files Browse the repository at this point in the history
  • Loading branch information
Routhleck committed Dec 8, 2023
1 parent 8ec4ce8 commit 81c6722
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 16 deletions.
42 changes: 37 additions & 5 deletions brainpy/_src/math/random.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
]


Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down
55 changes: 44 additions & 11 deletions brainpy/_src/math/tests/test_taichi_random.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)):
Expand All @@ -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)


Expand All @@ -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,
Expand All @@ -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
Expand Down
1 change: 1 addition & 0 deletions brainpy/math/random.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,

)

0 comments on commit 81c6722

Please sign in to comment.