Skip to content

File randomgaussianfield.h

File List > algebra > random > randomgaussianfield.h

Go to the documentation of this file

#ifndef TEMPLAT_LATTICE_ALGEBRA_RANDOM_RANDOMGAUSSIANFIELD_H
#define TEMPLAT_LATTICE_ALGEBRA_RANDOM_RANDOMGAUSSIANFIELD_H

/* This file is part of TempLat, available at https://cosmolattice.github.io/templat .
   Copyright 2021-2026 The TempLat authors, see AUTHORS.md.
   Released under the MIT license, see LICENSE.md. */

// File info: Main contributor(s): Wessel Valkenburg, Franz R. Sattler, Year: 2025

#include "TempLat/lattice/algebra/coordinates/dimensioncountrecorder.h"
#include "TempLat/util/constexpr_for.h"
#include "TempLat/util/random/randomgaussian.h"
#include "TempLat/parallel/device.h"

#include <tuple>

namespace TempLat
{
  MakeException(RandomGaussianFieldNegativeFrequencyException);

  template <typename T, size_t NDim, SpaceStateType Space, bool Real, bool Unitary>
  class RandomGaussianFieldHelper : public DimensionCountRecorder<NDim>
  {
    using RNGInteger = typename RandomGaussian<T>::IntegerType;
    static_assert(NDim != 0, "NDim template parameter is required.");
    static_assert(Space == SpaceStateType::Fourier || Space == SpaceStateType::Configuration,
                  "Space must be Fourier or Configuration.");
    static_assert(Space == SpaceStateType::Fourier || (!Real && !Unitary),
                  "Real/Unitary shape the Fourier complex pair; they are meaningless in configuration space.");

  public:
    // Put public methods here. These should change very little over time.
    RandomGaussianFieldHelper(std::string baseSeed, device::memory::host_ptr<MemoryToolBox<NDim>> pToolBox)
        : DimensionCountRecorder<NDim>(Space), prng(baseSeed), mToolBox(pToolBox), mLayout([&]() -> LayoutStruct<NDim> {
            if constexpr (Space == SpaceStateType::Fourier)
              return pToolBox->mLayouts.getFourierSpaceLayout();
            else
              return pToolBox->mLayouts.getConfigSpaceLayout();
          }()),
          generation(RNGInteger(0)), mGenerationSnapshot(0), mGlobalSizes(mLayout.getGlobalSizes())
    {
      DimensionCountRecorder<NDim>::confirmSpace(mLayout, Space);
    }

    void reset() { *generation = 0; }

    void preGet()
    {
      // `generation` is a host-only host_ptr (its device copy is null). Snapshot its value on the host into a
      // plain, device-copyable member so eval() -- which runs on-device -- never dereferences the host_ptr.
      mGenerationSnapshot = *generation;
    }

    void postGet()
    {
      // This is called after the get, so we can increase the generation.
      (*generation)++;
    }

    std::string saveState() const
    {
      std::ostringstream oss;
      oss << prng.saveState() << "\n"; // Underlying RNG state
      oss << *generation << " ";       // Generation counter
      return oss.str();
    }

    void loadState(const std::string &state)
    {
      std::istringstream iss(state);
      std::string uniformState;
      std::getline(iss, uniformState);
      prng.loadState(uniformState);
      iss >> *generation;
    }

    DEVICE_INLINE_FUNCTION std::tuple<RNGInteger, RNGInteger> gidx_to_idx2(const device::IdxArray<NDim> &gidx) const
    {
      constexpr size_t nd1 = NDim / 2;

      std::tuple<RNGInteger, RNGInteger> result;
      auto &r = std::get<0>(result);
      auto &c = std::get<1>(result);

      RNGInteger dim_length = 1;
      constexpr_for<0, NDim>([&](const auto _i) {
        // We go from the last dimension to the first, so we need to reverse the index.
        constexpr size_t i = NDim - 1 - decltype(_i)::value;
        // If we are in the second half of the dimensions, we sum the index to c
        if constexpr (i > nd1) c += gidx[i] * dim_length;
        // As soon, as we go into the first half, reset the dim_length to 1
        if constexpr (i == nd1) dim_length = 1;
        // If we are in the first half of the dimensions, we sum the index to r
        if constexpr (i <= nd1) r += gidx[i] * dim_length;
        dim_length *= mGlobalSizes[i];
      });

      return result;
    }

    template <typename R> DEVICE_INLINE_FUNCTION complex<T> to_complex(const device::array<R, 2> &pair) const
    {
      return complex<T>(static_cast<T>(pair[0]), static_cast<T>(pair[1]));
    }

    template <typename... IDX>
      requires IsVariadicNDIndex<NDim, IDX...>
    DEVICE_INLINE_FUNCTION auto eval(const IDX &...idx) const
    {
      if constexpr (Space == SpaceStateType::Fourier) {
        device::IdxArray<NDim> global_coord;
        mLayout.putSpatialLocationFromMemoryIndexInto(global_coord, idx...);

        device::IdxArray<NDim> hermitianPartner;
        auto hermitianType =
            DimensionCountRecorder<NDim>::getCurrentLayout().getHermitianPartners().putHermitianPartner(
                global_coord, hermitianPartner);

        // We do not need coordinates actually, but rather (positive!) global indices.
        for (size_t d = 0; d < NDim; ++d) {
          if (global_coord[d] < 0) global_coord[d] += mGlobalSizes[d];
          if (hermitianPartner[d] < 0) hermitianPartner[d] += mGlobalSizes[d];
        }

        if (hermitianType == HermitianRedundancy::none) {
          const auto [r, c] = gidx_to_idx2(global_coord);
          const complex<T> val = to_complex(prng.getPair(r, c, mGenerationSnapshot, Real, Unitary));
          return val;
        } else {
          const auto [r, c] = gidx_to_idx2(hermitianPartner);
          const complex<T> val = to_complex(prng.getPair(r, c, mGenerationSnapshot, Real, Unitary));
          return (hermitianType == HermitianRedundancy::positivePartner)   ? val
                 : (hermitianType == HermitianRedundancy::negativePartner) ? device::conj(val)
                 : (hermitianType == HermitianRedundancy::realValued)      ? complex<T>(device::real(val))
                                                                           : complex<T>(0.0, 0.0);
        }
      } else {
        // Configuration space: independent real N(0,1) white noise per site. No Hermitian structure.
        // Keyed on the GLOBAL coordinate, so the field is rank-independent and reproducible.
        device::IdxArray<NDim> global_coord;
        mLayout.putSpatialLocationFromMemoryIndexInto(global_coord, idx...);
        for (size_t d = 0; d < NDim; ++d)
          if (global_coord[d] < 0) global_coord[d] += mGlobalSizes[d];
        const auto [r, c] = gidx_to_idx2(global_coord);
        return prng.get(r, c, mGenerationSnapshot);
      }
    }

    std::string toString() const { return "Random gaussian field with seed: \"" + prng.getSeedString() + "\""; }

    auto getCurrentSeed() const { return prng.getSeed(); }

  private:
    /* Put all member variables and private methods here. These may change arbitrarily. */
    RandomGaussian<T> prng;
    device::memory::host_ptr<MemoryToolBox<NDim>> mToolBox;
    LayoutStruct<NDim> mLayout;
    device::memory::host_ptr<RNGInteger> generation;
    RNGInteger mGenerationSnapshot; // device-copyable snapshot of *generation, set on host in preGet()
    device::IdxArray<NDim> mGlobalSizes;
  };

  template <typename T, size_t NDim = 0>
  using RandomGaussianField = RandomGaussianFieldHelper<T, NDim, SpaceStateType::Fourier, false, false>;

  template <typename T, size_t NDim = 0>
  using RandomRayleighField = RandomGaussianFieldHelper<T, NDim, SpaceStateType::Fourier, true, false>;

  template <typename T, size_t NDim = 0>
  using RandomUniformUnitaryField = RandomGaussianFieldHelper<T, NDim, SpaceStateType::Fourier, false, true>;

  template <typename T, size_t NDim = 0>
  using RandomGaussianFieldConfig = RandomGaussianFieldHelper<T, NDim, SpaceStateType::Configuration, false, false>;

} // namespace TempLat

#endif