| #include "src/algorithms/random/distributions.h" |
| |
| #include <gtest/gtest.h> |
| |
| #include "src/algorithms/random/random.h" |
| |
| namespace cobalt { |
| |
| class DistributionsTest : public ::testing::Test { |
| protected: |
| void SetUp() override { gen_ = std::make_unique<RandomNumberGenerator>(55); } |
| RandomNumberGenerator* GetGenerator() { return gen_.get(); } |
| |
| private: |
| std::unique_ptr<RandomNumberGenerator> gen_; |
| }; |
| |
| TEST_F(DistributionsTest, BernoulliSample) { |
| double p_0 = 0.0; |
| auto b_0 = BernoulliDistribution(GetGenerator(), p_0); |
| auto sample_0 = b_0.Sample(); |
| EXPECT_FALSE(sample_0); |
| |
| double p_1 = 1.0; |
| auto b_1 = BernoulliDistribution(GetGenerator(), p_1); |
| auto sample_1 = b_1.Sample(); |
| EXPECT_TRUE(sample_1); |
| |
| double p = 0.5; |
| auto b = BernoulliDistribution(GetGenerator(), p); |
| auto sample = b.Sample(); |
| EXPECT_TRUE(sample || !sample); |
| } |
| |
| TEST_F(DistributionsTest, BinomialSample) { |
| uint64_t num_trials = 100; |
| double p_0 = 0.0; |
| auto b_0 = BinomialDistribution(GetGenerator(), num_trials, p_0); |
| auto sample_0 = b_0.Sample(); |
| EXPECT_EQ(sample_0, 0ul); |
| |
| double p_1 = 1.0; |
| auto b_1 = BinomialDistribution(GetGenerator(), num_trials, p_1); |
| auto sample_1 = b_1.Sample(); |
| EXPECT_EQ(sample_1, 100ul); |
| |
| // Depends on the seed passed to |gen_|. |
| double p = 0.5; |
| auto b = BinomialDistribution(GetGenerator(), num_trials, p); |
| auto sample = b.Sample(); |
| EXPECT_EQ(sample, 55ul); |
| } |
| |
| TEST_F(DistributionsTest, DiscreteUniformSample) { |
| uint32_t min = 0; |
| uint32_t max = 9; |
| auto u = DiscreteUniformDistribution(GetGenerator(), min, max); |
| for (int i = 0; i < 1000; i++) { |
| auto sample = u.Sample(); |
| EXPECT_GE(sample, min); |
| EXPECT_LE(sample, max); |
| } |
| } |
| |
| TEST_F(DistributionsTest, GeometricSample) { |
| double p_1 = 1.0; |
| auto g_1 = GeometricDistribution(GetGenerator(), p_1); |
| auto sample_1 = g_1.Sample(); |
| EXPECT_EQ(sample_1, 0ul); |
| |
| // Depends on the seed passed to |gen_|. |
| double p = 0.5; |
| auto g = GeometricDistribution(GetGenerator(), p); |
| auto sample = g.Sample(); |
| EXPECT_EQ(sample, 1ul); |
| } |
| |
| TEST_F(DistributionsTest, PoissonSample) { |
| int mean = 5; |
| int sigma = mean; |
| auto u = PoissonDistribution(GetGenerator(), mean); |
| int count_more_than_2_sigma = 0; |
| for (int i = 0; i < 1000; i++) { |
| auto sample = u.Sample(); |
| if (sample > mean + 2 * sigma || sample < mean - 2 * sigma) { |
| count_more_than_2_sigma++; |
| } |
| } |
| EXPECT_LT(count_more_than_2_sigma, 50); |
| } |
| |
| } // namespace cobalt |