[Lang] Add ti.randn (#2266)

* [Test] Update float64 random number test

* [LLVM] Update float64 random number generator

* [skip ci] enforce code format

* [test] simpifies test_random_f64 granularity test

* [Lang] Add ti.randn() for Gaussian random numbers

* [test] Add ti.randn() test

* [doc] add documentation for ti.randn()

* [skip ci] enforce code format

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>

docs/arithmetics.rst

@@ -123,6 +123,11 @@ Random number generator
 
 .. function:: ti.random(dtype = float)
 
+    Generates a uniform random float or integer number.
+
+.. function:: ti.randn(dtype = None)
+
+    Generates a random floating point number from the standard normal distribution.
 
 Element-wise arithmetics for vectors and matrices
 -------------------------------------------------

python/taichi/lang/__init__.py

@@ -2,6 +2,7 @@
 import os
 from copy import deepcopy as _deepcopy
 
+import taichi as ti
 from taichi.core.util import ti_core as _ti_core
 from taichi.lang import impl
 from taichi.lang.impl import *
@@ -19,8 +20,6 @@
                               taichi_scope, to_numpy_type, to_pytorch_type,
                               to_taichi_type)
 
-import taichi as ti
-
 # TODO(#2223): Remove
 core = _ti_core
 
@@ -290,6 +289,13 @@ def svd(A, dt=None):
     return svd(A, dt)
 
 
+def randn(dt=None):
+    if dt is None:
+        dt = impl.get_runtime().default_fp
+    from .random import randn
+    return randn(dt)
+
+
 determinant = deprecated('ti.determinant(a)',
                          'a.determinant()')(Matrix.determinant)
 tr = deprecated('ti.tr(a)', 'a.trace()')(Matrix.trace)

python/taichi/lang/random.py

@@ -0,0 +1,17 @@
+import math
+
+import taichi as ti
+
+
+@ti.func
+def randn(dt):
+    '''
+    Generates a random number from standard normal distribution
+    using the Box-Muller transform.
+    '''
+    assert dt == ti.f32 or dt == ti.f64
+    u1 = ti.random(dt)
+    u2 = ti.random(dt)
+    r = ti.sqrt(-2 * ti.log(u1))
+    c = ti.cos(math.tau * u2)
+    return r * c

tests/python/test_random.py

@@ -21,7 +21,7 @@ def fill():
 
         fill()
         X = x.to_numpy()
-        for i in range(4):
+        for i in range(1, 4):
             assert (X**i).mean() == approx(1 / (i + 1), rel=1e-2)
 
 
@@ -44,7 +44,7 @@ def fill():
 
         fill()
         X = x.to_numpy()
-        for i in range(4):
+        for i in range(1, 4):
             assert (X**i).mean() == approx(1 / (i + 1), rel=1e-2)
 
 
@@ -125,3 +125,28 @@ def foo() -> ti.f64:
     foo()
     frac, _ = np.modf(x.to_numpy() * 4294967296)
     assert np.max(frac) > 0
+
+
+@archs_support_random
+def test_randn():
+    '''
+    Tests the generation of Gaussian random numbers.
+    '''
+    for precision in [ti.f32, ti.f64]:
+        ti.init()
+        n = 1024
+        x = ti.field(ti.f32, shape=(n, n))
+
+        @ti.kernel
+        def fill():
+            for i in range(n):
+                for j in range(n):
+                    x[i, j] = ti.randn(precision)
+
+        fill()
+        X = x.to_numpy()
+
+        # https://en.wikipedia.org/wiki/Normal_distribution#Moments
+        moments = [0.0, 1.0, 0.0, 3.0]
+        for i in range(4):
+            assert (X**(i + 1)).mean() == approx(moments[i], abs=3e-2)