MooseRandomStateless.h

Go to the documentation of this file.

//* This file is part of the MOOSE framework
//* https://mooseframework.inl.gov
//*
//* All rights reserved, see COPYRIGHT for full restrictions
//* https://github.com/idaholab/moose/blob/master/COPYRIGHT
//*
//* Licensed under LGPL 2.1, please see LICENSE for details
//* https://www.gnu.org/licenses/lgpl-2.1.html

#pragma once

// MOOSE includes
#include "MooseError.h"
#include "DataIO.h"

// External library includes
#include "randistrs.h"

// STL includes
#include <unordered_map>

class MooseRandomStateless
{
public:
  MooseRandomStateless(const MooseRandomStateless &) = delete;
  MooseRandomStateless & operator=(const MooseRandomStateless &) = delete;

  MooseRandomStateless(unsigned int seed)
    : _seed(seed), _rand_generator(mts_ldrand, seed), _randl_generator(mts_lrand, seed)
  {
  }

  unsigned int getSeed() const { return _seed; }

  std::size_t getAdvanceCount() const { return _advance_count; }

  Real rand(std::size_t n) const { return _rand_generator.evaluate(n); }

  unsigned int randl(std::size_t n) const { return _randl_generator.evaluate(n); }

  unsigned int randl(std::size_t n, unsigned int lower, unsigned int upper) const
  {
    mooseAssert(upper >= lower, "randl: upper < lower");
    const auto range = upper - lower;

    auto [it, is_new] = _randlb_generators.try_emplace(
        range, [&range](mt_state * state) { return rds_iuniform(state, 0, range); }, _seed);

    if (is_new)
      it->second.advance(_advance_count);

    return lower + it->second.evaluate(n);
  }

  void advance(std::size_t count)
  {
    _rand_generator.advance(count);
    _randl_generator.advance(count);
    for (auto & [range, gen] : _randlb_generators)
      gen.advance(count);
    _advance_count += count;
  }

  template <typename T>
  class Generator
  {
  public:
    Generator(std::function<T(mt_state *)> rng_func, unsigned int seed)
      : _rng_func(std::move(rng_func))
    {
      // Create a new state with the given seed
      mts_seed32new(&_initial_state, seed);
      _current_state = _initial_state;
      _current_index = 0;
    }

    T evaluate(std::size_t n) const
    {
      if (n < _current_index)
      {
        _current_state = _initial_state;
        _current_index = 0;
      }

      T val = T();
      for (; _current_index <= n; ++_current_index)
        val = _rng_func(&_current_state);

      return val;
    }

    void advance(std::size_t count)
    {
      advance(_initial_state, count);
      _current_state = _initial_state;
      _current_index = 0;
    }

  private:
    void advance(mt_state & state, std::size_t count)
    {
      for (std::size_t i = 0; i < count; ++i)
        _rng_func(&state);
    }

    const std::function<T(mt_state *)> _rng_func;
    mt_state _initial_state;
    mutable mt_state _current_state;
    mutable std::size_t _current_index;
  };

private:
  const unsigned int _seed;
  Generator<Real> _rand_generator;
  Generator<unsigned int> _randl_generator;
  mutable std::unordered_map<unsigned int, Generator<unsigned int>> _randlb_generators;

  std::size_t _advance_count = 0;
};

template <>
inline void
dataStore(std::ostream & stream, MooseRandomStateless & v, void * context)
{
  std::size_t count = v.getAdvanceCount();
  storeHelper(stream, count, context);
}

template <>
inline void
dataLoad(std::istream & stream, MooseRandomStateless & v, void * context)
{
  std::size_t count;
  loadHelper(stream, count, context);
  v.advance(count);
}

template <>
inline void
dataStore(std::ostream & stream, std::unique_ptr<MooseRandomStateless> & v, void * context)
{
  unsigned int seed = v->getSeed();
  storeHelper(stream, seed, context);
  storeHelper(stream, *v, context);
}

template <>
inline void
dataLoad(std::istream & stream, std::unique_ptr<MooseRandomStateless> & v, void * context)
{
  unsigned int seed;
  loadHelper(stream, seed, context);
  v = std::make_unique<MooseRandomStateless>(seed);
  loadHelper(stream, *v, context);
}