Merge pull request #351 from FourierFlows/ncc/fix-stoch-forcing

Use `randn!` for stochastic forcing implementations
FourierFlows · Mar 11, 2024 · 7693b7b · 7693b7b
2 parents 5e8891c + 9ac215a
commit 7693b7b
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 24 deletions.
diff --git a/examples/barotropicqgql_betaforced.jl b/examples/barotropicqgql_betaforced.jl
@@ -79,14 +79,12 @@ nothing #hide
 
 
 # Next we construct function `calcF!` that computes a forcing realization every timestep.
-# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random
-# numbers uniformly distributed between 0 and 1.
-random_uniform = dev==CPU() ? rand : CUDA.rand
+# For that, we call `randn!` to obtain complex numbers whose real and imaginary part
+# are normally-distributed with zero mean and variance 1/2.
 
 function calcF!(Fh, sol, t, clock, vars, params, grid)
-  T = eltype(grid)
-  @. Fh = sqrt(forcing_spectrum) * cis(2π * random_uniform(T)) / sqrt(clock.dt)
-
+  randn!(Fh)
+  @. Fh *= sqrt(forcing_spectrum) / sqrt(clock.dt)
   return nothing
 end
 nothing #hide

diff --git a/examples/singlelayerqg_betaforced.jl b/examples/singlelayerqg_betaforced.jl
@@ -80,14 +80,12 @@ nothing #hide
 
 
 # Next we construct function `calcF!` that computes a forcing realization every timestep.
-# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random
-# numbers uniformly distributed between 0 and 1.
-random_uniform = dev==CPU() ? rand : CUDA.rand
+# For that, we call `randn!` to obtain complex numbers whose real and imaginary part
+# are normally-distributed with zero mean and variance 1/2.
 
 function calcF!(Fh, sol, t, clock, vars, params, grid)
-  T = eltype(grid)
-  @. Fh = sqrt(forcing_spectrum) * cis(2π * random_uniform(T)) / sqrt(clock.dt)
-
+  randn!(Fh)
+  @. Fh *= sqrt(forcing_spectrum) / sqrt(clock.dt)
   return nothing
 end
 nothing #hide

diff --git a/examples/twodnavierstokes_stochasticforcing.jl b/examples/twodnavierstokes_stochasticforcing.jl
@@ -72,14 +72,12 @@ nothing #hide
 
 
 # Next we construct function `calcF!` that computes a forcing realization every timestep.
-# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random
-# numbers uniformly distributed between 0 and 1.
-random_uniform = dev==CPU() ? rand : CUDA.rand
+# For that, we call `randn!` to obtain complex numbers whose real and imaginary part
+# are normally-distributed with zero mean and variance 1/2.
 
 function calcF!(Fh, sol, t, clock, vars, params, grid)
-  T = eltype(grid)
-  @. Fh = sqrt(forcing_spectrum) * cis(2π * random_uniform(T)) / sqrt(clock.dt)
-
+  randn!(Fh)
+  @. Fh *= sqrt(forcing_spectrum) / sqrt(clock.dt)
   return nothing
 end
 nothing #hide

diff --git a/examples/twodnavierstokes_stochasticforcing_budgets.jl b/examples/twodnavierstokes_stochasticforcing_budgets.jl
@@ -74,14 +74,12 @@ nothing #hide
 
 
 # Next we construct function `calcF!` that computes a forcing realization every timestep.
-# First we make sure that if `dev=GPU()`, then `CUDA.rand()` function is called for random
-# numbers uniformly distributed between 0 and 1.
-random_uniform = dev==CPU() ? rand : CUDA.rand
+# For that, we call `randn!` to obtain complex numbers whose real and imaginary part
+# are normally-distributed with zero mean and variance 1/2.
 
 function calcF!(Fh, sol, t, clock, vars, params, grid)
-  T = eltype(grid)
-  @. Fh = sqrt(forcing_spectrum) * cis(2π * random_uniform(T)) / sqrt(clock.dt)
-
+  randn!(Fh)
+  @. Fh *= sqrt(forcing_spectrum) / sqrt(clock.dt)
   return nothing
 end
 nothing #hide