random.cpp

// Copyright 2018 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "random.h"

#include <atomic>


    namespace {
        std::atomic<int> unique_stream_id{0};

        uint64_t ChooseSeed(uint64_t seed) {
            if (seed == 0) {
                std::random_device rd;
                seed = rd();
                if (sizeof(std::random_device::result_type) < 8) {
                    seed = (seed << 32) | rd();
                }
            }
            return seed;
        }

        int ChooseStream(int stream) {
            if (stream == 0) {
                stream = unique_stream_id.fetch_add(1);
            }
            return stream;
        }
    }  // namespace

    constexpr uint64_t Random::kLargePrime;
    constexpr uint64_t Random::kUniqueSeed;
    constexpr int Random::kUniqueStream;

    Random::Random(uint64_t seed, int stream)
            : seed_(ChooseSeed(seed)), impl_(seed_, ChooseStream(stream)) {}

    void Random::Dirichlet(float alpha, absl::Span<float> samples) {
        std::gamma_distribution<float> distribution(alpha);

        float sum = 0;
        for (float& sample : samples) {
            sample = distribution(impl_);
            sum += sample;
        }
        float norm = 1 / sum;
        for (float& sample : samples) {
            sample *= norm;
        }
    }

    float Random::Uniform(float a, float b) {
        std::uniform_real_distribution<float> distribution(a, b);
        return distribution(impl_);
    }

    void Random::Uniform(float a, float b, absl::Span<float> samples) {
        std::uniform_real_distribution<float> distribution(a, b);
        for (float& sample : samples) {
            sample = distribution(impl_);
        }
    }

    float Random::NormalDistribution(float mean, float stddev) {
        return std::normal_distribution<float>(mean, stddev)(impl_);
    }

    void Random::NormalDistribution(float mean, float stddev,
                                    absl::Span<float> samples) {
        std::normal_distribution<float> distribution(mean, stddev);
        for (float& sample : samples) {
            sample = distribution(impl_);
        }
    }

    int Random::SampleCdf(absl::Span<float> cdf) {
        // Take care to handle the case where the first elements in the CDF have zero
        // probability: discard any 0.0 values that the random number generator
        // produces. Admittedly, this isn't going to happen very often.
        float e;
        do {
            e = operator()();
        } while (e == 0);

        float x = cdf.back() * e;
        return std::distance(cdf.begin(),
                             std::lower_bound(cdf.begin(), cdf.end(), x));
    }