From 6454098c0abfa49098134cf70e06537ad1134460 Mon Sep 17 00:00:00 2001
From: Andrew Myers <atmyers2@gmail.com>
Date: Wed, 20 Nov 2024 13:42:03 -0800
Subject: [PATCH] Add portable implementation of random numbers drawn from the
 Gamma distribution.

---
 Src/Base/AMReX_Random.H   | 60 +++++++++++++++++++++++++++++++++++++++
 Src/Base/AMReX_Random.cpp |  7 +++++
 2 files changed, 67 insertions(+)

diff --git a/Src/Base/AMReX_Random.H b/Src/Base/AMReX_Random.H
index 50b2c2693b0..a8bda8b036d 100644
--- a/Src/Base/AMReX_Random.H
+++ b/Src/Base/AMReX_Random.H
@@ -7,6 +7,7 @@
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_RandomEngine.H>
 #include <limits>
+#include <cmath>
 #include <cstdint>
 
 namespace amrex
@@ -118,6 +119,65 @@ namespace amrex
 #endif
     }
 
+    /**
+    * \brief Generate a psuedo-random floating point number from the Gamma distribution
+    *
+    *  Generates one real number (single or double)
+    *  extracted from a Gamma distribution, given the Real parameters alpha and beta.
+    *  alpha and beta must both be > 0.
+    *  The CPU version of this function relies on the Standard Template Library
+    *  The GPU version of this function relies is implemented in terms of Random
+    *  and RandomNormal.
+    *
+    */
+    Real RandomGamma (Real alpha, Real beta);
+
+    AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
+    Real RandomGamma (Real alpha, Real beta, RandomEngine const& random_engine)
+    {
+        AMREX_ASSERT(alpha > 0);
+        AMREX_ASSERT(beta > 0);
+
+        AMREX_IF_ON_DEVICE((
+        if (alpha < 1)
+        {
+            Real u = amrex::Random(random_engine);
+            return RandomGamma(1.0_rt + alpha, beta, random_engine) * std::pow(u, 1.0_rt / alpha);
+        }
+
+        {
+            Real x, v, u;
+            Real d = alpha - 1.0_rt / 3.0_rt;
+            Real c = (1.0_rt / 3.0_rt) / std::sqrt(d);
+
+            while (1) {
+                do {
+                    x = amrex::RandomNormal(0.0_rt, 1.0_rt, random_engine);
+                    v = 1.0_rt + c * x;
+                } while (v <= 0.0_rt);
+
+                v = v * v * v;
+                u = amrex::Random(random_engine);
+
+                if (u < 1.0_rt - 0.0331_rt * x * x * x * x) {
+                    break;
+                }
+
+                if (std::log(u) < 0.5_rt * x * x + d * (1.0_rt - v + std::log(v))) {
+                    break;
+                }
+            }
+            return beta * d * v;
+        }
+        ))
+
+        AMREX_IF_ON_HOST((
+                amrex::ignore_unused(random_engine);
+                return RandomGamma(alpha, beta);
+        ))
+
+    }
+
     /**
     * \brief Generates one pseudorandom unsigned integer which is
     *  uniformly distributed on [0,n-1]-interval for each call.
diff --git a/Src/Base/AMReX_Random.cpp b/Src/Base/AMReX_Random.cpp
index 891a69e140a..d6d6d6805f3 100644
--- a/Src/Base/AMReX_Random.cpp
+++ b/Src/Base/AMReX_Random.cpp
@@ -134,6 +134,13 @@ unsigned int RandomPoisson (Real lambda)
     return distribution(generators[tid]);
 }
 
+Real RandomGamma (Real alpha, Real beta)
+{
+    std::gamma_distribution<Real> distribution(alpha, beta);
+    int tid = OpenMP::get_thread_num();
+    return distribution(generators[tid]);
+}
+
 unsigned int Random_int (unsigned int n)
 {
     std::uniform_int_distribution<unsigned int> distribution(0, n-1);