Initial import

extra/hades/libLSS/samplers/bias_model_params.cpp

@@ -0,0 +1,99 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/bias_model_params.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/physics/forward_model.hpp"
+#include <vector>
+#include "libLSS/samplers/bias_model_params.hpp"
+#include "libLSS/samplers/rgen/slice_sweep.hpp"
+
+using namespace LibLSS;
+
+BiasModelParamsSampler::BiasModelParamsSampler(
+    MPI_Communication *comm_,
+    std::shared_ptr<GridDensityLikelihoodBase<3>> likelihood_,
+    std::shared_ptr<BORGForwardModel> model_, int numBias_,
+    std::string const &prefix_)
+    : MarkovSampler(), comm(comm_), likelihood(likelihood_), model(model_),
+      numBias(numBias_), biasElement(0), prefix(prefix_) {}
+
+BiasModelParamsSampler::~BiasModelParamsSampler() {}
+
+void BiasModelParamsSampler::initialize(MarkovState &state) { restore(state); }
+
+void BiasModelParamsSampler::restore(MarkovState &state) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  if (!model) {
+    error_helper<ErrorBadState>("Model for BiasModelParams is invalid");
+  }
+  state.newElement(
+      prefix + "model_params_bias",
+      biasElement = new ArrayType1d(boost::extents[numBias]), true);
+  auto default_values = boost::any_cast<LibLSS::multi_array<double, 1>>(
+      model->getModelParam("bias", "biasParameters"));
+  fwrap(*biasElement->array) = default_values;
+  biasElement->subscribeLoaded([this]() {
+    model->setModelParams({{"biasParameters", *biasElement->array}});
+  });
+}
+
+void BiasModelParamsSampler::sample(MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+  auto const &x_density = *state.get<CArrayType>("s_hat_field")->array;
+  double const ares_heat = state.getScalar<double>("ares_heat");
+
+  RandomGen *rng = state.get<RandomGen>("random_generator");
+
+  if (pre_hook)
+    pre_hook();
+
+  likelihood->updateMetaParameters(state);
+
+  auto &biasParams = *biasElement->array;
+
+  for (size_t i = 0; i < numBias; i++) {
+    if (bias_to_freeze.find(i) != bias_to_freeze.end())
+      continue;
+
+    auto tmpParams = biasParams;
+    ctx.format2<LOG_VERBOSE>("Sampling bias model %d", i);
+    biasParams[i] = slice_sweep_double(
+        comm, rng->get(),
+        [&](double x) {
+          ModelDictionnary this_param_map;
+          tmpParams[i] = x;
+          this_param_map["biasParameters"] = tmpParams;
+          try {
+            model->setModelParams(this_param_map);
+          } catch (outOfBoundParam const &) {
+            return -std::numeric_limits<double>::infinity();
+          }
+          double log_L = -likelihood->logLikelihood(x_density, false);
+          ctx.format2<LOG_VERBOSE>("log_L = %g", log_L);
+          return ares_heat * log_L;
+        },
+        biasParams[i], 1.);
+
+    comm->broadcast_t(&biasParams[i], 1, 0);
+  }
+
+  ModelDictionnary this_param_map;
+  this_param_map["biasParameters"] = biasParams;
+  model->setModelParams(this_param_map);
+  likelihood->logLikelihood(x_density, false);
+  likelihood->commitAuxiliaryFields(state);
+
+  if (post_hook)
+    post_hook();
+}

extra/hades/libLSS/samplers/bias_model_params.hpp

@@ -0,0 +1,65 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/bias_model_params.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#pragma once
+#ifndef __LIBLSS_HADES_BIAS_MODEL_PARAMS_SAMPLER_HPP
+#  define __LIBLSS_HADES_BIAS_MODEL_PARAMS_SAMPLER_HPP
+
+#  include "libLSS/mpi/generic_mpi.hpp"
+#  include "libLSS/mcmc/global_state.hpp"
+#  include "libLSS/tools/fftw_allocator.hpp"
+#  include "libLSS/samplers/core/markov.hpp"
+#  include "libLSS/samplers/core/random_number.hpp"
+#  include "libLSS/samplers/core/types_samplers.hpp"
+#  include "libLSS/physics/forward_model.hpp"
+#  include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#  include <vector>
+
+namespace LibLSS {
+
+  class BiasModelParamsSampler : public MarkovSampler {
+  protected:
+    MPI_Communication *comm;
+    std::vector<std::string> paramsToSample;
+    std::shared_ptr<BORGForwardModel> model;
+    std::shared_ptr<GridDensityLikelihoodBase<3>> likelihood;
+    ModelDictionnary init_state;
+    int numBias;
+    ArrayType1d *biasElement;
+    std::set<int> bias_to_freeze;
+    typedef std::function<void()> HookType;
+
+    HookType pre_hook, post_hook;
+    std::string prefix;
+
+  public:
+    BiasModelParamsSampler(
+        MPI_Communication *comm_,
+        std::shared_ptr<GridDensityLikelihoodBase<3>> likelihood_,
+        std::shared_ptr<BORGForwardModel> model_, int numBias_,
+        std::string const &prefix);
+    virtual ~BiasModelParamsSampler();
+
+    void freezeSet(std::set<int> bias_to_freeze) {
+      this->bias_to_freeze = bias_to_freeze;
+    }
+
+    void initialize(MarkovState &state) override;
+    void restore(MarkovState &state) override;
+    void sample(MarkovState &state) override;
+
+    void setLimiterHooks(HookType pre_hook_, HookType post_hook_) {
+      pre_hook = pre_hook_;
+      post_hook = post_hook_;
+    }
+  };
+
+} // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/core/generate_random_field.cpp

@@ -0,0 +1,60 @@
+#include "libLSS/samplers/core/generate_random_field.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/tools/fusewrapper.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+
+using namespace LibLSS;
+
+void LibLSS::generateRandomField(MPI_Communication *comm, MarkovState &state) {
+  auto &s_hat = *state.get<CArrayType>("s_hat_field")->array;
+  auto &s = *state.get<ArrayType>("s_field")->array;
+  auto &rgen = state.get<RandomGen>("random_generator")->get();
+  std::array<long, 3> N;
+  std::array<double, 3> L;
+
+  state.getScalarArray<long, 3>("N", N);
+  state.getScalarArray<double, 3>("L", L);
+
+  FFTW_Manager<double, 3> mgr(N[0], N[1], N[2], comm);
+  double volume = array::product(L);
+  double volNorm = volume / array::product(N);
+  double invN3 = 1.0 / array::product(N);
+  double sqrt_invN3 = std::sqrt(invN3);
+
+  //  auto s_hat_p = base_mgr->allocate_complex_array();
+  //  auto &s_hat = s_hat_p.get_array();
+  auto s_hat_w = fwrap(s_hat);
+  auto s_w = fwrap(s);
+  auto tmp_real_field_p = mgr.allocate_array();
+  auto &tmp_real_field = tmp_real_field_p.get_array();
+  auto tmp_complex_field_p = mgr.allocate_complex_array();
+  auto &tmp_complex_field = tmp_complex_field_p.get_array();
+
+  auto synthesis_plan =
+      mgr.create_c2r_plan(tmp_complex_field.data(), tmp_real_field.data());
+  auto analysis_plan =
+      mgr.create_r2c_plan(tmp_real_field.data(), tmp_complex_field.data());
+
+  fwrap(tmp_real_field) = rgen.gaussian(fwrap(b_fused_idx<double, 3>(
+      [](size_t, size_t, size_t) -> double { return 1; },
+      mgr.extents_real_strict())));
+
+  mgr.execute_r2c(
+      analysis_plan, tmp_real_field.data(), tmp_complex_field.data());
+
+  // Convolve with sqrt(P(k))
+  s_hat_w = fwrap(tmp_complex_field) *
+            b_fused_idx<double, 3>(
+                [sqrt_invN3](size_t i, size_t j, size_t k) -> double {
+                  return sqrt_invN3;
+                });
+  {
+    // Get back s_field now
+    fwrap(tmp_complex_field) = s_hat_w * (1 / volume);
+    mgr.execute_c2r(
+        synthesis_plan, tmp_complex_field.data(), tmp_real_field.data());
+  }
+  fwrap(s[mgr.strict_range()]) = fwrap(tmp_real_field[mgr.strict_range()]);
+}

extra/hades/libLSS/samplers/core/generate_random_field.hpp

@@ -0,0 +1,13 @@
+#ifndef __LIBLSS_GENERATE_RANDOM_FIELD_HPP
+#define __LIBLSS_GENERATE_RANDOM_FIELD_HPP
+
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+
+namespace LibLSS {
+
+  void generateRandomField(MPI_Communication *comm, MarkovState &state);
+
+}
+
+#endif

extra/hades/libLSS/samplers/core/gridLikelihoodBase.cpp

@@ -0,0 +1,111 @@
+#include <functional>
+#include "libLSS/tools/console.hpp"
+#include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#include <utility>
+
+using namespace LibLSS;
+
+namespace ph = std::placeholders;
+
+template <typename T, std::size_t Dims, std::size_t... Is>
+std::unique_ptr<FFTW_Manager<T, Dims>> makeMgr(
+    MPI_Communication *comm, std::array<size_t, Dims> const &d,
+    std::index_sequence<Is...>) {
+  typedef FFTW_Manager<T, Dims> Mgr;
+  return std::unique_ptr<Mgr>(new Mgr(d[Is]..., comm));
+}
+
+template <int Dims>
+GridDensityLikelihoodBase<Dims>::GridDensityLikelihoodBase(
+    MPI_Communication *comm_, GridSizes const &N_, GridLengths const &L_)
+    : LikelihoodBase(), comm(comm_), N(N_), L(L_), volume(array::product(L_)) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+
+  // FIXME: Is that the only preserved symmetry ?
+  //        This one is enforced because average(s_field) == 0 at all time
+  //        So the gradient of the likelihood may not create any mean or risk
+  //        imbalance.
+  // FIXME: This is only valid for 3D. It has to be generalized to N-d
+  //special_cases.push_back(std::make_tuple(Index{0, 0, 0}, 0));
+  special_cases.push_back(std::make_tuple(
+      Index{ssize_t(N[0] / 2), ssize_t(N[1] / 2), ssize_t(N[2] / 2)}, 0));
+  for (size_t i = 1; i < 7; i++) {
+    auto j = Index{0, 0, 0};
+
+    j[0] = (i & 1) * (N[0] / 2);
+    j[1] = ((i & 2) >> 1) * (N[1] / 2);
+    j[2] = ((i & 4) >> 2) * (N[2] / 2);
+    special_cases.push_back(std::make_tuple(j, 0.5));
+  }
+
+  mgr = makeMgr<double, Dims>(comm_, N_, std::make_index_sequence<Dims>());
+  auto tmp_real_field = mgr->allocate_array();
+  auto tmp_complex_field = mgr->allocate_complex_array();
+  analysis_plan = mgr->create_r2c_plan(
+      tmp_real_field.get_array().data(), tmp_complex_field.get_array().data());
+}
+
+template <int Dims>
+GridDensityLikelihoodBase<Dims>::~GridDensityLikelihoodBase() {}
+
+template <>
+void GridDensityLikelihoodBase<3>::computeFourierSpace_GradientPsi(
+    ArrayRef &real_gradient, CArrayRef &grad_array, bool accumulate,
+    double scaling) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+  double normalizer =
+      1 /
+      volume; // Normalization is like synthesis here, we consider the transpose synthesis
+  auto tmp_complex_field_p = mgr->allocate_complex_array();
+  auto plane_set = std::array<size_t, 2>{0, N[2] / 2};
+  using boost::multi_array_types::index_range;
+
+  // BEWARE: real_gradient is destroyed!
+  mgr->execute_r2c(
+      analysis_plan, real_gradient.data(),
+      tmp_complex_field_p.get_array().data());
+
+  auto &c_field = tmp_complex_field_p.get_array();
+  auto e = fwrap(c_field) * (2 * normalizer); // Factor 2 owed to hermiticity
+  auto output = fwrap(grad_array);
+  size_t startN0 = mgr->startN0;
+  size_t endN0 = mgr->startN0 + mgr->localN0;
+
+  // FIXME: Why do we do this already ?
+  ctx.print(" Handle special cases");
+  for (auto &h : special_cases) {
+    ssize_t i0 = std::get<0>(h)[0];
+    if (i0 >= startN0 && i0 < endN0)
+      c_field(std::get<0>(h)) *= std::get<1>(h);
+  }
+
+  for (auto p : plane_set) {
+    size_t i0_start = std::max(startN0, N[0] / 2 + 1);
+    size_t i0_end = std::min(endN0, N[0]);
+    ctx.print(
+        "Fix plane " + std::to_string(p) + " " + std::to_string(i0_start) +
+        " => " + std::to_string(i0_end));
+    if (i0_start < i0_end)
+      fwrap(c_field[boost::indices[index_range(i0_start, i0_end)][index_range()]
+                                  [p]]) = 0;
+    for (auto q : {size_t(0), N[0] / 2}) {
+      ctx.print("Fix plane (2) " + std::to_string(q));
+      if (mgr->on_core(q)) {
+        fwrap(c_field[boost::indices[q][index_range(N[1] / 2 + 1, N[1])][p]]) =
+            0;
+      }
+    }
+  }
+
+  if (mgr->on_core(0))
+    c_field[0][0][0] = 0;
+
+  if (accumulate)
+    output = output + scaling * e;
+  else
+    output = e * scaling;
+}
+
+namespace LibLSS {
+  template class GridDensityLikelihoodBase<3>;
+}

extra/hades/libLSS/samplers/core/gridLikelihoodBase.hpp

@@ -0,0 +1,117 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/core/gridLikelihoodBase.hpp
+    Copyright (C) 2018 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_SAMPLERS_CORE_GRID_LIKELIHOOD_BASE_HPP
+#  define __LIBLSS_SAMPLERS_CORE_GRID_LIKELIHOOD_BASE_HPP
+
+#  include <array>
+#  include <tuple>
+#  include <vector>
+#  include "libLSS/mpi/generic_mpi.hpp"
+#  include "libLSS/samplers/core/types_samplers.hpp"
+#  include "libLSS/tools/mpi_fftw_helper.hpp"
+#  include "libLSS/mcmc/global_state.hpp"
+#  include "libLSS/physics/forward_model.hpp"
+#  include "libLSS/samplers/core/likelihood.hpp"
+
+namespace LibLSS {
+
+  template <int Dims>
+  class GridDensityLikelihoodBase : virtual public LikelihoodBase {
+  public:
+    typedef std::array<size_t, Dims> GridSizes;
+    typedef std::array<double, Dims> GridLengths;
+    typedef std::array<ssize_t, Dims> Index;
+
+    typedef boost::multi_array_ref<double, Dims> ArrayRef;
+    typedef boost::multi_array_ref<std::complex<double>, Dims> CArrayRef;
+
+    GridDensityLikelihoodBase(
+        MPI_Communication *comm, GridSizes const &N, GridLengths const &L);
+    virtual ~GridDensityLikelihoodBase();
+
+    void commitAuxiliaryFields(MarkovState &state) override {}
+
+    virtual void
+    generateMockData(CArrayRef const &parameters, MarkovState &state) = 0;
+
+    /*
+     * This is the opposite of log likelihood, Fourier representation. The same
+     * rules applies as for the real space variant. The modes must be
+     * in unit of Volume as per the continuous Fourier transform.
+     */
+    virtual double
+    logLikelihood(CArrayRef const &parameters, bool gradientIsNext = false) = 0;
+
+    /*
+     * This computes the opposite of the log likelihood. If gradientIsnext is
+     * true the model has to prepare itself for a gradientLikelihood call.
+     * Otherwise it can free temporary memory used to compute the forward model.
+     * This variant takes the image in real space representation. The input is
+     * preserved as indicated by the const.
+     */
+    virtual double
+    logLikelihood(ArrayRef const &parameters, bool gradientIsNext = false) = 0;
+
+    /*
+     * This is the gradient of the logarithm of the opposite of the log
+     * likelihood (Fourier representation variant).
+     * The same rules applies as for the real space variant.
+     */
+    virtual void gradientLikelihood(
+        CArrayRef const &parameters, CArrayRef &gradient_parameters,
+        bool accumulate = false, double scaling = 1.0) = 0;
+
+    /*
+    * This is the gradient of the opposite of the log likelihood. It
+    * returns the gradient in real space representation.
+    * You must have called logLikelihood with gradientIsNext=true first.
+    * If accumulate is set to true, the gradient will be summed with existing values.
+    * 'scaling' indicates by how much the gradient must be scaled before accumulation.
+    */
+    virtual void gradientLikelihood(
+        ArrayRef const &parameters, ArrayRef &gradient_parameters,
+        bool accumulate = false, double scaling = 1.0) = 0;
+
+  protected:
+    typedef std::tuple<std::array<ssize_t, Dims>, double> SpecialType;
+    typedef FFTW_Manager<double, Dims> Mgr;
+    std::vector<SpecialType> special_cases;
+    GridSizes N;
+    GridLengths L;
+    long Ncat;
+    MPI_Communication *comm;
+    double volume;
+    std::shared_ptr<Mgr> mgr;
+    typename Mgr::plan_type analysis_plan;
+
+    void computeFourierSpace_GradientPsi(
+        ArrayRef &real_gradient, CArrayRef &grad_array, bool accumulate,
+        double scaling);
+
+  public:
+    std::shared_ptr<Mgr> getManager() { return mgr; }
+  };
+
+  class ForwardModelBasedLikelihood : public GridDensityLikelihoodBase<3> {
+  public:
+    ForwardModelBasedLikelihood(
+        MPI_Communication *comm, GridSizes const &N, GridLengths const &L)
+        : GridDensityLikelihoodBase<3>(comm, N, L) {}
+
+    virtual std::shared_ptr<BORGForwardModel> getForwardModel() = 0;
+  };
+
+} // namespace LibLSS
+
+#endif
+
+// ARES TAG: authors_num = 1
+// ARES TAG: name(0) = Guilhem Lavaux
+// ARES TAG: email(0) = guilhem.lavaux@iap.fr
+// ARES TAG: year(0) = 2018

extra/hades/libLSS/samplers/core/likelihood.hpp

@@ -0,0 +1,47 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/core/likelihood.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#pragma once
+#ifndef __LIBLSS_CORE_LIKELIHOOD_BASE_HPP
+#  define __LIBLSS_CORE_LIKELIHOOD_BASE_HPP
+
+#  include "libLSS/physics/cosmo.hpp"
+#  include "libLSS/mcmc/global_state.hpp"
+#  include "libLSS/mcmc/state_element.hpp"
+
+namespace LibLSS {
+
+  /**
+   * @brief A type to provide name parameter space mapping to the state element in the MarkovChain
+   * 
+   */
+  typedef std::map<std::string, std::shared_ptr<StateElement>> ParameterSpace;
+
+  /**
+   * @brief This is the fundamental likelihood class. 
+   * 
+   * It does not express log likelihood evaluation but provides entry points to initialize, obtain
+   * cosmology and manipulate a MarkovState. 
+   * 
+   */
+  class LikelihoodBase {
+  public:
+    LikelihoodBase() = default;
+    virtual ~LikelihoodBase() {}
+
+    virtual void initializeLikelihood(MarkovState &state) = 0;
+    virtual void updateMetaParameters(MarkovState &state) = 0;
+    virtual void setupDefaultParameters(MarkovState &state, int catalog) = 0;
+    virtual void updateCosmology(CosmologicalParameters const &params) = 0;
+    virtual void commitAuxiliaryFields(MarkovState &state) = 0;
+  };
+
+} // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/core/simpleLikelihood.cpp

@@ -0,0 +1,16 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/core/simpleLikelihood.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include "libLSS/samplers/core/simpleLikelihood.hpp"
+
+using namespace LibLSS;
+
+SimpleLikelihood::SimpleLikelihood(LikelihoodInfo const &info) {}
+
+SimpleLikelihood::~SimpleLikelihood() {}

extra/hades/libLSS/samplers/core/simpleLikelihood.hpp

@@ -0,0 +1,74 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/core/simpleLikelihood.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#pragma once
+#ifndef __HADES_SIMPLE_LIKELIHOOD_HPP
+#  define __HADES_SIMPLE_LIKELIHOOD_HPP
+
+#  include "libLSS/physics/likelihoods/base.hpp"
+#  include "libLSS/samplers/core/types_samplers.hpp"
+#  include "libLSS/samplers/core/likelihood.hpp"
+
+namespace LibLSS {
+
+  /**
+   * @brief Likelihood that has a simple structure, parallelization wise. 
+   * 
+   * It assumes that all nodes are going to obtain the same copy of the parameters
+   * and produce something collectively that may depend on other parameters that
+   * are local to the MPI nodes. For single task node, the difference to the GridDensityLikelihoodBase
+   * is only the shape of the parameter space. However for parallel jobs it can be
+   * become very different.
+   * 
+   * We note that the proposed_params arguments are unnamed, thus the behaviour will heavily
+   * depend on the semantic of those parameters. Some synchronization of the behaviour may
+   * be required with the user of such a likelihood. 
+   */
+  class SimpleLikelihood : virtual public LikelihoodBase {
+  public:
+    SimpleLikelihood(LikelihoodInfo const &info);
+    virtual ~SimpleLikelihood();
+
+    /**
+     * @brief Compute the log likelihood assuming the proposed parameters.
+     * 
+     * @param global_params     a global parameter state that is not changing with the current transition kernel
+     * @param proposed_params   the set of parameters that are provided to get a new evaluation of the likelihood
+     * @return double    the result
+     */
+    virtual double logLikelihoodSimple(
+        ParameterSpace const &global_params,
+        LibLSS::const_multi_array_ref<double, 1> const &proposed_params) = 0;
+
+    /**
+     * @brief Compute the gradient of the log-likelihood w.r.t to the proposed_params
+     * 
+     * @param global_params      a global parameter state that is not changing with the current transition kernel
+     * @param proposed_params    the set of parameters that are provided to get a new evaluation of the likelihood
+     * @param gradient           the gradient w.r.t proposed_params
+     */
+    virtual void gradientLikelihoodSimple(
+        ParameterSpace const &global_params,
+        LibLSS::const_multi_array_ref<double, 1> const &proposed_params,
+        LibLSS::multi_array_ref<double, 1> &gradient) = 0;
+
+    /**
+     * @brief Returns the number of dimensions required by this likelihood.
+     * 
+     * Note: as SimpleLikelihood is global. That number must be the same on all
+     * MPI nodes.
+     * 
+     * @return unsigned int 
+     */
+    virtual unsigned int numDimensions() const = 0;
+  };
+
+} // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/core/splitLikelihood.cpp

@@ -0,0 +1,164 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/core/splitLikelihood.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include "libLSS/tools/console.hpp"
+#include "libLSS/samplers/core/splitLikelihood.hpp"
+#include "libLSS/tools/overload.hpp"
+#include "libLSS/tools/fusewrapper.hpp"
+#include <boost/variant.hpp>
+#include <functional>
+
+using namespace LibLSS;
+
+SplitAndReshapeLikelihood::SplitAndReshapeLikelihood(MPI_Communication *comm_)
+    : comm(comm_) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+}
+
+SplitAndReshapeLikelihood::~SplitAndReshapeLikelihood() {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+}
+
+void SplitAndReshapeLikelihood::initializeLikelihood(MarkovState &state) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  for (auto &l : parameterLikelihood) {
+    boost::apply_visitor(
+        [&state](std::shared_ptr<LikelihoodBase> &likelihood) {
+          likelihood->initializeLikelihood(state);
+        },
+        std::get<1>(l));
+  }
+}
+
+void SplitAndReshapeLikelihood::updateMetaParameters(MarkovState &state) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  for (auto &l : parameterLikelihood) {
+    boost::apply_visitor(
+        [&state](std::shared_ptr<LikelihoodBase> &likelihood) {
+          likelihood->updateMetaParameters(state);
+        },
+        std::get<1>(l));
+  }
+}
+void SplitAndReshapeLikelihood::setupDefaultParameters(
+    MarkovState &state, int catalog) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  for (auto &l : parameterLikelihood) {
+    boost::apply_visitor(
+        [&state, catalog](std::shared_ptr<LikelihoodBase> &likelihood) {
+          likelihood->setupDefaultParameters(state, catalog);
+        },
+        std::get<1>(l));
+  }
+}
+
+void SplitAndReshapeLikelihood::updateCosmology(
+    CosmologicalParameters const &params) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  for (auto &l : parameterLikelihood) {
+    boost::apply_visitor(
+        [&params](std::shared_ptr<LikelihoodBase> &likelihood) {
+          likelihood->updateCosmology(params);
+        },
+        std::get<1>(l));
+  }
+}
+
+void SplitAndReshapeLikelihood::commitAuxiliaryFields(MarkovState &state) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  for (auto &l : parameterLikelihood) {
+    boost::apply_visitor(
+        [&state](std::shared_ptr<LikelihoodBase> &likelihood) {
+          likelihood->commitAuxiliaryFields(state);
+        },
+        std::get<1>(l));
+  }
+}
+
+/*
+ * ---------------------------------------------------------------------
+ * getTotalNumberOfParameters
+ *
+ * getTotalNumberOfParameters needs to implement explicitly the function
+ * for each variant of likelihood.
+ */
+
+static unsigned int computeNumberOfParametersGrid(
+    std::shared_ptr<GridDensityLikelihoodBase<3>> grid) {
+  auto mgr = grid->getManager();
+
+  return mgr->localN0 * mgr->N1 * mgr->N2_HC * 2;
+}
+
+static unsigned int computeNumberOfParametersSimple(
+    MPI_Communication *comm, std::shared_ptr<SimpleLikelihood> simple) {
+  // The node with rank 0 will be in charge of generating the required
+  // random numbers/proposals.
+  return comm->rank() == 0 ? simple->numDimensions() : 0;
+}
+
+unsigned int SplitAndReshapeLikelihood::getTotalNumberOfParameters() {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+
+  unsigned int total = 0;
+  for (auto &l : parameterLikelihood) {
+    total += boost::apply_visitor(
+        overload(
+            computeNumberOfParametersGrid,
+            std::bind(
+                &computeNumberOfParametersSimple, comm, std::placeholders::_1)),
+        std::get<1>(l));
+  }
+  return total;
+}
+
+// ---------------------------------------------------------------------
+
+void SplitAndReshapeLikelihood::addNamedParameter(
+    std::string parameter, LikelihoodVariant likelihood) {
+  parameterLikelihood.push_back(std::make_tuple(parameter, likelihood));
+}
+
+/*
+ * ---------------------------------------------------------------------
+ * logLikelihood needs a separate treatment based on each variant
+ * 
+ */
+
+static double logLikelihoodSimple(
+    MPI_Communication *comm, ParameterSpace &global,
+    std::shared_ptr<SimpleLikelihood> simple, unsigned int start,
+    boost::const_multi_array_ref<double, 1> const &params) {
+  unsigned int num = simple->numDimensions();
+  boost::multi_array<double, 1> paramsSubSet(boost::extents[num]);
+  typedef boost::multi_array_types::index_range range;
+
+  fwrap(paramsSubSet) = fwrap(params[range(start, start + num)]);
+  double loc_L = simple->logLikelihoodSimple(global, params);
+  double L = 0;
+
+  comm->all_reduce_t(&loc_L, &L, 1, MPI_SUM);
+  return L;
+}
+
+double SplitAndReshapeLikelihood::logLikelihood(
+    LibLSS::const_multi_array_ref<double, 1> &params) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  unsigned int start = 0;
+  ParameterSpace global;
+  double L = 0;
+
+  for (auto &l : parameterLikelihood) {
+    L += boost::apply_visitor(
+        overload(std::bind(&logLikelihoodSimple, comm, global, start, params)),
+        std::get<1>(l));
+  }
+
+  return L;
+}

extra/hades/libLSS/samplers/core/splitLikelihood.hpp

@@ -0,0 +1,103 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/core/splitLikelihood.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#pragma once
+#ifndef __LIBLSS_HADES_SPLIT_LIKELIHOOD_HPP
+#  define __LIBLSS_HADES_SPLIT_LIKELIHOOD_HPP
+
+#  include "libLSS/samplers/core/likelihood.hpp"
+#  include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#  include "libLSS/samplers/core/simpleLikelihood.hpp"
+#  include <boost/variant.hpp>
+#  include <list>
+#  include <memory>
+#  include <tuple>
+
+namespace LibLSS {
+
+  /**
+   * @brief 
+   * 
+   */
+  class SplitAndReshapeLikelihood : public LikelihoodBase {
+  protected:
+    typedef boost::variant<
+        std::shared_ptr<GridDensityLikelihoodBase<3>>,
+        std::shared_ptr<SimpleLikelihood>>
+        LikelihoodVariant;
+
+    MPI_Communication *comm;
+    std::list<std::tuple<std::string, LikelihoodVariant>> parameterLikelihood;
+
+  public:
+    SplitAndReshapeLikelihood(MPI_Communication *comm_);
+    ~SplitAndReshapeLikelihood();
+
+    void initializeLikelihood(MarkovState &state) override;
+    void updateMetaParameters(MarkovState &state) override;
+    void setupDefaultParameters(MarkovState &state, int catalog) override;
+    void updateCosmology(CosmologicalParameters const &params) override;
+    void commitAuxiliaryFields(MarkovState &state) override;
+
+    /**
+     * @brief Add that this parameter must be evaluated with a specific log-likelihood object
+     * The parameter needs to exist later in the MarkovState. 
+     * 
+     * Depending on the likelihood variant, the behaviour will be different:
+     *    - GridDensityLikelihoodBase<3> is for likelihoods
+     *      that are grid based and thus split over MPI nodes with FFTW slabbing.
+     *    - SimpleLikelihoodBase is for likelihoods that have their parameters identical
+     *      and replicated on all MPI nodes
+     * 
+     * @param parameter 
+     * @param likelihood
+     */
+    void addNamedParameter(std::string parameter, LikelihoodVariant likelihood);
+
+    /**
+     * @brief Get the total number Of parameters that are resident on the
+     * current MPI node. 
+     * 
+     * The sampler/optimizer will need allocate only this
+     * portion.
+     * 
+     * @return unsigned int
+     */
+    unsigned int getTotalNumberOfParameters();
+
+    /**
+     * @brief Compute the log likelihood using the provided array. 
+     * 
+     * WARNING!! The parameters provided here depends heavily on the inner structure
+     * of the split likelihood parameter space. Some parameters need be replicated over MPI nodes and 
+     * some others not.
+     * 
+     * @param params 
+     * @return double 
+     */
+    double logLikelihood(LibLSS::const_multi_array_ref<double, 1> &params);
+
+    /**
+     * @brief Compute the gradient of the log likelihood using the provided array to evaluate it. 
+     * 
+     * WARNING!! The parameters provided here depends heavily on the inner structure
+     * of the split likelihood parameter space. Some parameters need be replicated over MPI nodes and 
+     * some others not.
+     * 
+     * @param params 
+     * @return double 
+     */
+    double gradientLikelihood(
+        LibLSS::const_multi_array_ref<double, 1> &params,
+        LibLSS::multi_array_ref<double, 1> &gradient_params);
+  };
+
+} // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/hades/base_likelihood.cpp

@@ -0,0 +1,344 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/hades/base_likelihood.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include "libLSS/cconfig.h"
+#include <CosmoTool/algo.hpp>
+#include <cmath>
+#include <Eigen/Core>
+#include <boost/format.hpp>
+#include <boost/bind.hpp>
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+#include "libLSS/tools/array_tools.hpp"
+#include "libLSS/samplers/rgen/hmc/hmc_density_sampler.hpp"
+#include "libLSS/samplers/hades/base_likelihood.hpp"
+#include "libLSS/tools/fusewrapper.hpp"
+#include "libLSS/samplers/rgen/slice_sweep.hpp"
+
+using namespace LibLSS;
+using boost::extents;
+using boost::format;
+
+using CosmoTool::square;
+
+using boost::c_storage_order;
+typedef boost::multi_array_types::extent_range range;
+
+static const int ROOT_RANK = 0;
+
+namespace LI = LibLSS::Likelihood;
+
+HadesBaseDensityLikelihood::HadesBaseDensityLikelihood(
+    LikelihoodInfo &info, size_t numBiasParams_)
+    : super_t(LI::getMPI(info), LI::gridResolution(info), LI::gridSide(info)),
+      SimpleLikelihood(info), corners(LI::gridCorners(info)),
+      numBiasParams(numBiasParams_), volume(array::product(L)) {}
+
+HadesBaseDensityLikelihood::~HadesBaseDensityLikelihood() {}
+
+void HadesBaseDensityLikelihood::updateCosmology(
+    CosmologicalParameters const &cosmo_params) {
+  cosmology = std::unique_ptr<Cosmology>(new Cosmology(cosmo_params));
+  model->setCosmoParams(cosmo_params);
+}
+
+void HadesBaseDensityLikelihood::updateMetaParameters(MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_VERBOSE, ctx);
+  auto cosmo_params = state.getScalar<CosmologicalParameters>("cosmology");
+
+  Ncat = state.getScalar<long>("NCAT");
+
+  //initialize model uncertainty
+  model = state.get<BorgModelElement>("BORG_model")->obj;
+
+  ai = state.getScalar<double>("borg_a_initial");
+
+  // Update forward model for maybe new cosmo params
+  updateCosmology(cosmo_params);
+
+  auto e_Ncat = boost::extents[Ncat];
+  nmean.resize(e_Ncat);
+  biasRef.resize(Ncat);
+  data.resize(Ncat);
+  sel_field.resize(Ncat);
+  bias.resize(Ncat);
+
+  for (int c = 0; c < Ncat; c++) {
+    biasRef[c] = state.getScalar<bool>(format("galaxy_bias_ref_%d") % c);
+    nmean[c] = state.getScalar<double>(format("galaxy_nmean_%d") % c);
+    auto &stateBias =
+        *state.get<ArrayType1d>(format("galaxy_bias_%d") % c)->array;
+    if (stateBias.size() < numBiasParams) {
+      stateBias.resize(boost::extents[numBiasParams]);
+      Console::instance().print<LOG_WARNING>("I will pick default bias parameters as the present array is too short.");
+      setupDefaultParameters(state, c);
+    }
+
+    bias[c] = std::make_unique<BiasArray>(boost::extents[stateBias.size()]);
+
+    fwrap(*bias[c]) = fwrap(stateBias);
+    ctx.print(format(" b0=%g") % (*bias[c])[0]);
+    auto data_name = str(format("galaxy_data_%d") % c);
+    if (state.exists(data_name)) {
+      data[c] = state.get<ArrayType>(data_name)->array;
+      sel_field[c] =
+          state.get<ArrayType>(format("galaxy_synthetic_sel_window_%d") % c)
+              ->array;
+    }
+  }
+}
+
+double HadesBaseDensityLikelihood::logLikelihood(
+    ArrayRef const &s_array, bool final_call) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+  using CosmoTool::square;
+  typedef ArrayType::ArrayType::element ElementType;
+  double L = 0;
+  Cosmology &cosmo = *cosmology.get();
+
+  ctx.print("Run forward");
+  // Simulate forward model
+  auto box = model->get_box_model();
+  auto out_density_p = mgr->allocate_array();
+  auto &out_density = out_density_p.get_array();
+  model->setAdjointRequired(false);
+  model->forwardModel_v2(ModelInput<3>(mgr, box, s_array));
+  model->getDensityFinal(ModelOutput<3>(mgr, box, out_density));
+
+  L = logLikelihoodSpecific(out_density);
+  comm->all_reduce_t(MPI_IN_PLACE, &L, 1, MPI_SUM);
+
+  return L;
+}
+
+double HadesBaseDensityLikelihood::logLikelihood(
+    CArrayRef const &s_array, bool final_call) {
+  using CosmoTool::square;
+  typedef ArrayType::ArrayType::element ElementType;
+  double L = 0;
+  auto &out_density = final_density_field->get_array();
+
+  auto box = model->get_box_model();
+  model->setAdjointRequired(false);
+  model->forwardModel_v2(ModelInput<3>(mgr, box, s_array));
+  model->getDensityFinal(ModelOutput<3>(mgr, box, out_density));
+
+  L = logLikelihoodSpecific(out_density);
+  comm->all_reduce_t(MPI_IN_PLACE, &L, 1, MPI_SUM);
+
+  return L;
+}
+
+void HadesBaseDensityLikelihood::gradientLikelihood_internal(
+    ModelInput<3> input_field, ModelOutputAdjoint<3> grad_array) {
+  using CosmoTool::square;
+  typedef ArrayType::ArrayType::element ElementType;
+  double Epoisson = 0;
+
+  auto box = model->get_box_model();
+  auto &out_density = final_density_field->get_array();
+
+  model->setAdjointRequired(true);
+  model->forwardModel_v2(std::move(input_field));
+  model->getDensityFinal(ModelOutput<3>(mgr, box, out_density));
+
+  auto tmp_grad_p = model->out_mgr->allocate_array();
+  auto &tmp_grad = tmp_grad_p.get_array();
+  gradientLikelihoodSpecific(out_density, tmp_grad);
+
+  model->adjointModel_v2(ModelInputAdjoint<3>(mgr, box, tmp_grad));
+  model->getAdjointModelOutput(std::move(grad_array));
+}
+
+unsigned int HadesBaseDensityLikelihood::numDimensions() const { return 0; }
+
+double HadesBaseDensityLikelihood::logLikelihoodSimple(
+    ParameterSpace const &global,
+    LibLSS::const_multi_array_ref<double, 1> const &params) {
+  return 0;
+}
+
+void HadesBaseDensityLikelihood::gradientLikelihoodSimple(
+    ParameterSpace const &global,
+    LibLSS::const_multi_array_ref<double, 1> const &params,
+    LibLSS::multi_array_ref<double, 1> &gradient) {}
+
+void HadesBaseDensityLikelihood::gradientLikelihood(
+    ArrayRef const &s_array, ArrayRef &grad_array, bool accumulate,
+    double scaling) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+
+  auto box = model->get_box_model();
+
+  if (!accumulate) {
+    gradientLikelihood_internal(
+        ModelInput<3>(mgr, box, s_array),
+        ModelOutputAdjoint<3>(mgr, box, grad_array));
+    fwrap(grad_array) = fwrap(grad_array) * scaling;
+  } else {
+    auto real_gradient_p = mgr->allocate_array();
+    auto &real_gradient = real_gradient_p.get_array();
+
+    gradientLikelihood_internal(
+        ModelInput<3>(mgr, box, s_array),
+        ModelOutputAdjoint<3>(mgr, box, real_gradient));
+    fwrap(grad_array) = fwrap(grad_array) + scaling * fwrap(real_gradient);
+  }
+}
+
+void HadesBaseDensityLikelihood::gradientLikelihood(
+    CArrayRef const &parameters, CArrayRef &grad_array, bool accumulate,
+    double scaling) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+  //  auto real_gradient_p = mgr->allocate_array();
+  //  auto &real_gradient = real_gradient_p.get_array();
+  auto box = model->get_box_model();
+
+  {
+    if (!accumulate) {
+      gradientLikelihood_internal(
+          ModelInput<3>(mgr, box, parameters),
+          ModelOutputAdjoint<3>(mgr, box, grad_array));
+      fwrap(grad_array) = fwrap(grad_array) * scaling;
+    } else {
+      auto tmp_complex_field = mgr->allocate_complex_array();
+      gradientLikelihood_internal(
+          ModelInput<3>(mgr, box, parameters),
+          ModelOutputAdjoint<3>(mgr, box, tmp_complex_field.get_array()));
+
+      fwrap(grad_array) =
+          fwrap(grad_array) + fwrap(tmp_complex_field.get_array()) * scaling;
+    }
+  }
+}
+
+void HadesBaseDensityLikelihood::generateMockData(
+    CArrayRef const &parameters, MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_INFO, ctx);
+
+  auto &out_density = final_density_field->get_array();
+  auto box = model->get_box_model();
+
+  model->setAdjointRequired(false);
+  model->forwardModel_v2(ModelInput<3>(mgr, box, parameters));
+  model->getDensityFinal(ModelOutput<3>(
+      model->out_mgr, model->get_box_model_output(), out_density));
+
+  generateMockSpecific(out_density, state);
+  commitAuxiliaryFields(state);
+}
+
+void HadesBaseDensityLikelihood::initializeLikelihood(MarkovState &state) {
+  Ncat = state.getScalar<long>("NCAT");
+  model = state.get<BorgModelElement>("BORG_model")->obj;
+  borg_final_density = new ArrayType(model->out_mgr->extents_real_strict());
+  final_density_field = model->out_mgr->allocate_ptr_array();
+
+  std::array<ssize_t, 3> out_N, local_N;
+  model->get_box_model_output().fill(out_N);
+  borg_final_density->setRealDims(out_N);
+
+  std::copy(
+      borg_final_density->array->shape(),
+      borg_final_density->array->shape() + 3, local_N.begin());
+
+  for (size_t c = 0; c < Ncat; c++) {
+    auto &stateBias =
+        *state.get<ArrayType1d>(format("galaxy_bias_%d") % c)->array;
+    if (stateBias.size() < numBiasParams) {
+      stateBias.resize(boost::extents[numBiasParams]);
+      Console::instance().print<LOG_WARNING>("I will pick default bias parameters as the present array is too short.");
+      setupDefaultParameters(state, c);
+    }
+
+    auto &data = *state.get<ArrayType>(format("galaxy_data_%d") % c)->array;
+    // Check data has correct shape
+    if (std::mismatch(local_N.begin(), local_N.end(), data.shape()).first !=
+        local_N.end())
+      error_helper<ErrorParams>(boost::str(
+          boost::format(
+              "Invalid output size (%dx%dx%d) vs data size (%dx%dx%d)") %
+          local_N[0] % local_N[1] % local_N[2] % data.shape()[0] %
+          data.shape()[1] % data.shape()[2]));
+  }
+
+  state.newElement("BORG_final_density", borg_final_density, true);
+}
+
+void HadesBaseDensityLikelihood::commitAuxiliaryFields(MarkovState &state) {
+  array::scaleAndCopyArray3d(
+      *borg_final_density->array, final_density_field->get_array(), 1, true);
+}
+
+void HadesBaseDensityLikelihood::updateNmean(int catalog, double nmean_) {
+  nmean[catalog] = nmean_;
+}
+
+void HadesBaseDensityLikelihood::updateBiasParameters(
+    int catalog, BiasArray const &params) {
+  fwrap(*bias[catalog]) = fwrap(params);
+}
+
+void HadesMetaSampler::initialize(MarkovState &state) {
+  Ncat = state.getScalar<long>("NCAT");
+}
+
+void HadesMetaSampler::restore(MarkovState &state) {
+  Ncat = state.getScalar<long>("NCAT");
+}
+
+void HadesMetaSampler::sample(MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+  RandomGen *rng = state.get<RandomGen>("random_generator");
+
+  if (state.getScalar<bool>("bias_sampler_blocked"))
+    return;
+
+  auto const &density = *state.get<ArrayType>("BORG_final_density")->array;
+  double const ares_heat = state.getScalar<double>("ares_heat");
+
+  likelihood->updateMetaParameters(state);
+
+  for (int c = 0; c < Ncat; c++) {
+    bool biasRef = state.getScalar<bool>(format("galaxy_bias_ref_%d") % c);
+    auto const &sel_array =
+        *state.get<SelArrayType>(format("galaxy_sel_window_%d") % c)->array;
+    double &nmean = state.getScalar<double>(format("galaxy_nmean_%d") % c);
+    auto &bias = *state.get<ArrayType1d>(format("galaxy_bias_%d") % c)->array;
+
+    if (biasRef)
+      continue;
+
+    ctx.print(format("considering catalog %d") % c);
+
+    nmean = slice_sweep_double(
+        comm, rng->get(),
+        [&](double x) -> double {
+          likelihood->updateNmean(c, x);
+          return -ares_heat * likelihood->logLikelihoodSpecific(density);
+        },
+        nmean, 0.1);
+
+    likelihood->updateNmean(c, nmean);
+
+    for (int ib = 0; ib < bias.size(); ib++) {
+      bias[ib] = slice_sweep(
+          comm, rng->get(),
+          [&](double x) -> double {
+            boost::multi_array<double, 1> loc_bias = bias;
+            loc_bias[ib] = x;
+            likelihood->updateBiasParameters(c, loc_bias);
+            return -ares_heat * likelihood->logLikelihoodSpecific(density);
+          },
+          bias[ib], 0.1);
+      likelihood->updateBiasParameters(c, bias);
+    }
+  }
+}

extra/hades/libLSS/samplers/hades/base_likelihood.hpp

@@ -0,0 +1,165 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/hades/base_likelihood.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_HADES_BASE_LIKELIHOOD_HPP
+#define __LIBLSS_HADES_BASE_LIKELIHOOD_HPP
+
+#include "libLSS/cconfig.h"
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#include "libLSS/samplers/core/likelihood.hpp"
+#include "libLSS/samplers/core/simpleLikelihood.hpp"
+#include "libLSS/physics/forward_model.hpp"
+#include "libLSS/physics/likelihoods/base.hpp"
+
+namespace LibLSS {
+
+  /**
+   * This is a base class for simple bias models. It inherits both from forward model based likelihood
+   * and simple likelihood as two kind of gradients will be produced. The ForwardModelBasedLikelihood produces
+   * gradients w.r.t the density field, whereas the SimpleLikelihood produces gradients w.r.t to the bias parameters.
+   */
+  class HadesBaseDensityLikelihood : public ForwardModelBasedLikelihood,
+                                     public SimpleLikelihood {
+  public:
+    //typedef GridDensityLikelihoodBase<3> super_t;
+    typedef ForwardModelBasedLikelihood super_t;
+
+  protected:
+    GridLengths corners;
+    ArrayType1d *vobs;
+    ArrayType *borg_final_density;
+
+    std::shared_ptr<BORGForwardModel> model;
+
+    std::unique_ptr<Cosmology> cosmology;
+    std::shared_ptr<Mgr::U_ArrayReal> final_density_field;
+    double ai;
+    double volume;
+
+    boost::multi_array<double, 1> nmean;
+    std::vector<std::shared_ptr<ArrayType::ArrayType>> data, sel_field;
+
+    typedef ArrayType1d::ArrayType BiasArray;
+    std::vector<std::unique_ptr<BiasArray>> bias;
+    std::vector<bool> biasRef;
+    size_t numBiasParams;
+
+    void gradientLikelihood_internal(
+        ModelInput<3> parameters, ModelOutputAdjoint<3> gradient_parameters);
+
+  public:
+    HadesBaseDensityLikelihood(LikelihoodInfo &info, size_t numBiasParams);
+    virtual ~HadesBaseDensityLikelihood();
+
+    void
+    generateMockData(CArrayRef const &parameters, MarkovState &state) override;
+
+    std::shared_ptr<BORGForwardModel> getForwardModel() override {
+      return model;
+    }
+
+    void initializeLikelihood(MarkovState &state) override;
+    void updateMetaParameters(MarkovState &state) override;
+    void updateCosmology(CosmologicalParameters const &params) override;
+    void updateBiasParameters(int catalog, BiasArray const &params);
+    void updateNmean(int catalog, double nmean);
+    void commitAuxiliaryFields(MarkovState &state) override;
+
+    /**
+     * This is for computing a log likelihood with a provided dark matter density field
+     */
+    virtual double logLikelihoodSpecific(ArrayRef const &density) = 0;
+
+    /**
+     * This is for computing a gradient log likelihood with a provided dark matter density field.
+     * The output gradient must be stored in "grad_density".
+     */
+    virtual void gradientLikelihoodSpecific(
+        ArrayRef const &density, ArrayRef &grad_density) = 0;
+
+    virtual void
+    generateMockSpecific(ArrayRef const &density, MarkovState &state) = 0;
+
+    /*
+     * This computes the opposite of the log likelihood. If gradientIsnext is
+     * true the model has to prepare itself for a gradientLikelihood call.
+     * Otherwise it can free temporary memory used to compute the forward model.
+     * This variant takes the image in real space representation. The input is
+     * preserved as indicated by the const.
+     */
+    double logLikelihood(
+        ArrayRef const &parameters, bool gradientIsNext = false) override;
+
+    /*
+     * This is the gradient of the opposite of the log likelihood. It
+     * returns the gradient in real space representation.
+     * You must have called logLikelihood with gradientIsNext=true first.
+     * If accumulate is set to true, the gradient will be summed with existing values.
+     * 'scaling' indicates by how much the gradient must be scaled before accumulation.
+     */
+    void gradientLikelihood(
+        ArrayRef const &parameters, ArrayRef &gradient_parameters,
+        bool accumulate, double scaling) override;
+
+    /*
+     * This is the opposite of log likelihood, Fourier representation. The same
+     * rules applies as for the real space variant. The modes must be
+     * in unit of Volume as per the continuous Fourier transform.
+     */
+    double logLikelihood(
+        CArrayRef const &parameters, bool gradientIsNext = false) override;
+
+    /*
+     * This is the gradient of the logarithm of the opposite of the log
+     * likelihood (Fourier representation variant).
+     * The same rules applies as for the real space variant.
+     */
+    void gradientLikelihood(
+        CArrayRef const &parameters, CArrayRef &gradient_parameters,
+        bool accumulate, double scaling) override;
+
+    double logLikelihoodSimple(
+        ParameterSpace const &global,
+        LibLSS::const_multi_array_ref<double, 1> const &params) override;
+
+    void gradientLikelihoodSimple(
+        ParameterSpace const &global,
+        LibLSS::const_multi_array_ref<double, 1> const &params,
+        LibLSS::multi_array_ref<double, 1> &gradient) override;
+
+    unsigned int numDimensions() const override;
+  };
+
+  class HadesMetaSampler : public MarkovSampler {
+  protected:
+    typedef std::shared_ptr<HadesBaseDensityLikelihood> likelihood_t;
+    int Ncat;
+    MPI_Communication *comm;
+    likelihood_t likelihood;
+
+  public:
+    HadesMetaSampler(MPI_Communication *comm_, likelihood_t likelihood_)
+        : comm(comm_), likelihood(likelihood_) {}
+    virtual ~HadesMetaSampler() {}
+
+    void initialize(MarkovState &state) override;
+    void restore(MarkovState &state) override;
+    void sample(MarkovState &state) override;
+  };
+
+}; // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/hades/hades_linear_likelihood.cpp

@@ -0,0 +1,172 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/hades/hades_linear_likelihood.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include <CosmoTool/algo.hpp>
+#include <cmath>
+#include <Eigen/Core>
+#include <boost/format.hpp>
+#include <boost/bind.hpp>
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+#include "libLSS/tools/array_tools.hpp"
+#include "libLSS/samplers/hades/hades_linear_likelihood.hpp"
+#include "libLSS/tools/hermiticity_fixup.hpp"
+
+using namespace LibLSS;
+using boost::extents;
+using boost::format;
+
+using CosmoTool::square;
+
+using boost::c_storage_order;
+typedef boost::multi_array_types::extent_range range;
+
+typedef Eigen::Map<Eigen::ArrayXd, Eigen::Aligned> MappedArray;
+
+static const int ROOT_RANK = 0;
+
+HadesLinearDensityLikelihood::HadesLinearDensityLikelihood(LikelihoodInfo &info)
+    : super_t(info, 1) {}
+
+HadesLinearDensityLikelihood::~HadesLinearDensityLikelihood() {}
+
+double
+HadesLinearDensityLikelihood::logLikelihoodSpecific(ArrayRef const &delta) {
+  double L = 0;
+  size_t const startN0 = model->out_mgr->startN0;
+  size_t const endN0 = startN0 + model->out_mgr->localN0;
+  size_t const N1 = model->out_mgr->N1;
+  size_t const N2 = model->out_mgr->N2;
+
+  for (int c = 0; c < Ncat; c++) {
+    auto &sel_array = *(sel_field[c]);
+    double nmean_c = nmean[c];
+    double bias_c = (*(bias[c]))[0];
+    auto &data_c = *(data[c]);
+
+#pragma omp parallel for schedule(static) collapse(3) reduction(+ : L)
+    for (size_t n0 = startN0; n0 < endN0; n0++) {
+      for (size_t n1 = 0; n1 < N1; n1++) {
+        for (size_t n2 = 0; n2 < N2; n2++) {
+          double selection = sel_array[n0][n1][n2];
+
+          if (selection > 0) {
+            double Nobs = data_c[n0][n1][n2];
+            double d_galaxy = bias_c * delta[n0][n1][n2];
+
+            assert(!isnan(Nobs));
+            assert(!isnan(d_galaxy));
+            L += square(selection * nmean_c * (1 + d_galaxy) - Nobs) /
+                 (selection * nmean_c);
+            if (std::isnan(L)) {
+              error_helper<ErrorBadState>("NaN in Likelihood");
+            }
+            if (std::isinf(L)) {
+              error_helper<ErrorBadState>(
+                  format("Inf in hamiltonian at n0=%d n1=%d n2=%d d_g=%lg "
+                         "Nobs=%lg") %
+                  n0 % n1 % n2 % d_galaxy % Nobs);
+            }
+          }
+        }
+      }
+    }
+  }
+
+  L *= 0.5;
+
+  return L;
+}
+
+void HadesLinearDensityLikelihood::gradientLikelihoodSpecific(
+    ArrayRef const &delta, ArrayRef &grad_array) {
+
+  size_t const startN0 = model->out_mgr->startN0;
+  size_t const endN0 = startN0 + model->out_mgr->localN0;
+  size_t const N1 = model->out_mgr->N1;
+  size_t const N2 = model->out_mgr->N2;
+
+  fwrap(grad_array) = 0;
+
+  for (int c = 0; c < Ncat; c++) {
+    auto &sel_array = *(sel_field[c]);
+    auto &data_c = *(data[c]);
+    double bias_c = (*bias[c])[0];
+    double nmean_c = nmean[c];
+
+#pragma omp parallel for collapse(3)
+    for (size_t n0 = startN0; n0 < endN0; n0++) {
+      for (size_t n1 = 0; n1 < N1; n1++) {
+        for (size_t n2 = 0; n2 < N2; n2++) {
+          double deltaElement = delta[n0][n1][n2];
+          double d_galaxy = bias_c * deltaElement;
+          double d_galaxy_prime = bias_c;
+          double response = sel_array[n0][n1][n2];
+          double Nobs = data_c[n0][n1][n2];
+
+          if (response == 0)
+            continue;
+
+          grad_array[n0][n1][n2] +=
+              (nmean_c * response * (1 + d_galaxy) - Nobs) * d_galaxy_prime;
+          assert(!isnan(grad_array[n0][n1][n2]));
+        }
+      }
+    }
+  }
+}
+
+void HadesLinearDensityLikelihood::generateMockSpecific(
+    ArrayRef const &delta, MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_INFO, ctx);
+  RandomGen *rgen = state.get<RandomGen>("random_generator");
+  size_t const startN0 = model->out_mgr->startN0;
+  size_t const endN0 = startN0 + model->out_mgr->localN0;
+  size_t const N1 = model->out_mgr->N1;
+  size_t const N2 = model->out_mgr->N2;
+
+  for (int c = 0; c < Ncat; c++) {
+    ctx.print(format("Generating mock data %d") % c);
+    auto &sel_array = *sel_field[c];
+    auto &g_field = *state.get<ArrayType>(format("galaxy_data_%d") % c)->array;
+    double nmean_c = nmean[c];
+    double bias_c = (*bias[c])[0];
+
+#pragma omp parallel for schedule(static) collapse(3)
+    for (size_t n0 = startN0; n0 < endN0; n0++) {
+      for (size_t n1 = 0; n1 < N1; n1++) {
+        for (size_t n2 = 0; n2 < N2; n2++) {
+          double R = nmean_c * sel_array[n0][n1][n2];
+          double gmean = R * (1 + bias_c * delta[n0][n1][n2]);
+          assert(!isnan(gmean));
+          assert(!isnan(R));
+          assert(R >= 0);
+          g_field[n0][n1][n2] = rgen->get().gaussian() * sqrt(R) + gmean;
+          assert(!isnan(g_field[n0][n1][n2]));
+        }
+      }
+    }
+  }
+}
+
+void HadesLinearDensityLikelihood::initializeLikelihood(MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_INFO, ctx);
+  super_t::initializeLikelihood(state);
+}
+
+void HadesLinearDensityLikelihood::setupDefaultParameters(
+    MarkovState &state, int catalog) {
+  auto &bias_c =
+      *state.formatGet<ArrayType1d>("galaxy_bias_%d", catalog)->array;
+  auto &nmean_c = state.formatGetScalar<double>("galaxy_nmean_%d", catalog);
+
+  bias_c[0] = 1.0;
+  nmean_c = 1;
+}

extra/hades/libLSS/samplers/hades/hades_linear_likelihood.hpp

@@ -0,0 +1,46 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/hades/hades_linear_likelihood.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_HADES_LINEAR_DENSITY_HPP
+#define __LIBLSS_HADES_LINEAR_DENSITY_HPP
+
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#include "libLSS/physics/forward_model.hpp"
+#include "libLSS/samplers/hades/base_likelihood.hpp"
+
+namespace LibLSS {
+
+  class HadesLinearDensityLikelihood : public HadesBaseDensityLikelihood {
+  public:
+    typedef HadesBaseDensityLikelihood super_t;
+    typedef HadesBaseDensityLikelihood::super_t grid_t;
+
+  public:
+    HadesLinearDensityLikelihood(LikelihoodInfo &info);
+    virtual ~HadesLinearDensityLikelihood();
+
+    virtual void setupDefaultParameters(MarkovState &state, int catalog);
+    virtual void
+    generateMockSpecific(ArrayRef const &parameters, MarkovState &state);
+    virtual double logLikelihoodSpecific(ArrayRef const &parameters);
+    virtual void gradientLikelihoodSpecific(
+        ArrayRef const &parameters, ArrayRef &gradient_parameters);
+    virtual void initializeLikelihood(MarkovState &state);
+  };
+
+}; // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/model_params.cpp

@@ -0,0 +1,81 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/model_params.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/physics/forward_model.hpp"
+#include <vector>
+#include "libLSS/samplers/model_params.hpp"
+#include "libLSS/samplers/rgen/slice_sweep.hpp"
+
+using namespace LibLSS;
+
+ModelParamsSampler::ModelParamsSampler(
+    MPI_Communication *comm_, std::string const& prefix_, std::vector<std::string> const &params,
+    std::shared_ptr<GridDensityLikelihoodBase<3>> likelihood_,
+    std::shared_ptr<BORGForwardModel> model_, ModelDictionnary init)
+    : MarkovSampler(), comm(comm_), paramsToSample(params),
+      likelihood(likelihood_), model(model_), init_state(init), prefix(prefix_) {}
+
+ModelParamsSampler::~ModelParamsSampler() {}
+
+void ModelParamsSampler::initialize(MarkovState &state) { restore(state); }
+
+void ModelParamsSampler::restore(MarkovState &state) {
+  for (auto const &p : paramsToSample) {
+    std::string pName = std::string("model_params_") + prefix + p;
+    state
+        .newScalar<double>(pName, boost::any_cast<double>(init_state[p]), true)
+        ->subscribeLoaded([this, p, pName, &state]() {
+          model->setModelParams({{p, state.getScalar<double>(pName)}});
+        });
+  }
+  model->setModelParams(init_state);
+}
+
+void ModelParamsSampler::sample(MarkovState &state) {
+  LIBLSS_AUTO_CONTEXT(LOG_DEBUG, ctx);
+  auto const &x_density = *state.get<CArrayType>("s_hat_field")->array;
+  double const ares_heat = state.getScalar<double>("ares_heat");
+
+  RandomGen *rng = state.get<RandomGen>("random_generator");
+
+  likelihood->updateMetaParameters(state);
+
+  for (auto const &p : paramsToSample) {
+    double &value = state.getScalar<double>(std::string("model_params_") + prefix + p);
+    ctx.format2<LOG_VERBOSE>("Sampling model parameter '%s'", p);
+    value = slice_sweep_double(
+        comm, rng->get(),
+        [&](double x) {
+          ctx.format2<LOG_VERBOSE>("try x[%s] = %g", p, x);
+          if (x < 0)
+            return -std::numeric_limits<double>::infinity();
+          ModelDictionnary this_param_map;
+          this_param_map[p] = x;
+          model->setModelParams(this_param_map);
+          double log_L = -likelihood->logLikelihood(x_density, false);
+          ctx.format2<LOG_VERBOSE>("log_L = %g", log_L);
+          return ares_heat * log_L;
+        },
+        value, 1.);
+
+    comm->broadcast_t(&value, 1, 0);
+
+    ModelDictionnary this_param_map;
+    this_param_map[p] = value;
+    model->setModelParams(this_param_map);
+  }
+  likelihood->logLikelihood(x_density, false);
+  likelihood->commitAuxiliaryFields(state);
+}

extra/hades/libLSS/samplers/model_params.hpp

@@ -0,0 +1,51 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/model_params.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#pragma once
+#ifndef __LIBLSS_HADES_MODEL_PARAMS_SAMPLER_HPP
+#  define __LIBLSS_HADES_MODEL_PARAMS_SAMPLER_HPP
+
+#  include "libLSS/mpi/generic_mpi.hpp"
+#  include "libLSS/mcmc/global_state.hpp"
+#  include "libLSS/tools/fftw_allocator.hpp"
+#  include "libLSS/samplers/core/markov.hpp"
+#  include "libLSS/samplers/core/random_number.hpp"
+#  include "libLSS/samplers/core/types_samplers.hpp"
+#  include "libLSS/physics/forward_model.hpp"
+#  include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#  include <vector>
+
+namespace LibLSS {
+
+  class ModelParamsSampler : public MarkovSampler {
+  protected:
+    MPI_Communication *comm;
+    std::vector<std::string> paramsToSample;
+    std::shared_ptr<BORGForwardModel> model;
+    std::shared_ptr<GridDensityLikelihoodBase<3>> likelihood;
+    ModelDictionnary init_state;
+    std::string prefix;
+
+  public:
+    ModelParamsSampler(
+        MPI_Communication *comm_, std::string const& prefix, std::vector<std::string> const &params,
+        std::shared_ptr<GridDensityLikelihoodBase<3>> likelihood_,
+        std::shared_ptr<BORGForwardModel> model_,
+        ModelDictionnary init = ModelDictionnary());
+    virtual ~ModelParamsSampler();
+
+    void initialize(MarkovState &state) override;
+    void restore(MarkovState &state) override;
+    void sample(MarkovState &state) override;
+  };
+
+} // namespace LibLSS
+
+#endif
+

extra/hades/libLSS/samplers/rgen/density_sampler.hpp

@@ -0,0 +1,25 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/rgen/density_sampler.hpp
+    Copyright (C) 2019 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_SAMPLERS_DENSITY_SAMPLER_HPP
+#define __LIBLSS_SAMPLERS_DENSITY_SAMPLER_HPP
+
+namespace LibLSS {
+
+  class GenericDensitySampler: public MarkovSampler {
+  public:
+      virtual void generateMockData(MarkovState& state) = 0;
+  };
+
+}
+
+#endif
+// ARES TAG: authors_num = 1
+// ARES TAG: name(0) = Guilhem Lavaux
+// ARES TAG: email(0) = guilhem.lavaux@iap.fr
+// ARES TAG: year(0) = 2019

extra/hades/libLSS/samplers/rgen/frozen/frozen_phase_density_sampler.cpp

@@ -0,0 +1,188 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/rgen/frozen/frozen_phase_density_sampler.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include <CosmoTool/algo.hpp>
+#include <cmath>
+#include <Eigen/Core>
+#include <boost/format.hpp>
+#include <functional>
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+#include "libLSS/tools/array_tools.hpp"
+#include "libLSS/samplers/rgen/frozen/frozen_phase_density_sampler.hpp"
+#include "libLSS/tools/uninitialized_type.hpp"
+#include "libLSS/tools/hermiticity_fixup.hpp"
+#include <fstream>
+#include <iostream>
+#include <H5Cpp.h>
+
+static const bool ULTRA_VERBOSE = false;
+static const bool HMC_PERF_TEST = true;
+static const bool FIXED_INTEGRATION_PATH = false;
+
+using namespace LibLSS;
+using boost::extents;
+using boost::format;
+
+using CosmoTool::square;
+
+using boost::c_storage_order;
+typedef boost::multi_array_types::extent_range range;
+
+typedef Eigen::Map<Eigen::ArrayXd, Eigen::Aligned> MappedArray;
+
+static const int ROOT_RANK = 0;
+
+namespace ph = std::placeholders;
+
+FrozenPhaseDensitySampler::FrozenPhaseDensitySampler(
+    MPI_Communication *comm, Likelihood_t likelihood)
+    : analysis_plan(0), synthesis_plan(0), comm(comm) {
+  this->likelihood = likelihood;
+}
+
+void FrozenPhaseDensitySampler::generateMockData(MarkovState &state) {
+  likelihood->updateMetaParameters(state);
+
+  auto &rgen = state.get<RandomGen>("random_generator")->get();
+
+  {
+    auto tmp_real = base_mgr->allocate_array();
+    double const inv_sqN = 1.0 / std::sqrt(N0 * N1 * N2);
+
+    // Generate a bunch of gaussian distributed random number (variance 1)
+    fwrap(tmp_real.get_array()) =
+        rgen.gaussian(ones<double, 3>(base_mgr->extents_real()) * inv_sqN);
+
+    base_mgr->execute_r2c(
+        analysis_plan, tmp_real.get_array().data(), x_hat_field->array->data());
+
+    fwrap(*s_hat_field->array) = fwrap(*x_hat_field->array);
+  }
+  auto tmp_complex = base_mgr->allocate_complex_array();
+  fwrap(tmp_complex.get_array()) = fwrap(*s_hat_field->array) / volume;
+
+  base_mgr->execute_c2r(
+      synthesis_plan, tmp_complex.get_array().data(), s_field->array->data());
+
+  likelihood->generateMockData(*s_hat_field->array, state);
+}
+
+void FrozenPhaseDensitySampler::initialize(MarkovState &state) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  restore(state);
+
+  // Now load phases
+  if (phaseFilename) {
+    H5::H5File f(*phaseFilename, H5F_ACC_RDONLY);
+
+    ctx.print("Read-in phase data");
+    {
+      U_Array<double, 3> tmp_x(
+          base_mgr
+              ->extents_real_strict()); // We need to allocate this temporary array to adapt shape.
+      CosmoTool::hdf5_read_array(f, dataName, tmp_x.get_array(), false, true);
+
+      // FS: updated according to GL, 6/23/20
+      fwrap(array::slice_array(*x_field->array, base_mgr->strict_range())) = fwrap(tmp_x.get_array());
+    }
+
+    auto tmp_field = base_mgr->allocate_array();
+    fwrap(tmp_field.get_array()) = fwrap(*x_field->array) * volNorm;
+
+    ctx.print("Fourier transform");
+    base_mgr->execute_r2c(
+        analysis_plan, tmp_field.get_array().data(),
+        x_hat_field->array->data());
+
+    fwrap(*s_hat_field->array) = fwrap(*x_hat_field->array);
+    // WARNING: s_field/s_hat_field are not consistent at that moment. They will become at the first
+    //          call to sample here.
+  }
+}
+
+void FrozenPhaseDensitySampler::restore(MarkovState &state) {
+  Console &cons = Console::instance();
+  ConsoleContext<LOG_DEBUG> ctx("Initialize frozen density sampler");
+
+  N0 = state.getScalar<long>("N0");
+  N1 = state.getScalar<long>("N1");
+  N2 = state.getScalar<long>("N2");
+
+  base_mgr = std::unique_ptr<DFT_Manager>(new DFT_Manager(N0, N1, N2, comm));
+  size_t Ntot = N0 * N1 * N2;
+
+  L0 = state.getScalar<double>("L0");
+  L1 = state.getScalar<double>("L1");
+  L2 = state.getScalar<double>("L2");
+
+  startN0 = base_mgr->startN0;
+  localN0 = base_mgr->localN0;
+  endN0 = startN0 + localN0;
+
+  cons.print<LOG_DEBUG>("Allocating s field");
+  s_hat_field =
+      new CArrayType(base_mgr->extents_complex(), base_mgr->allocator_complex);
+  s_hat_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2_HC));
+  x_hat_field =
+      new CArrayType(base_mgr->extents_complex(), base_mgr->allocator_complex);
+  x_hat_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2_HC));
+  x_field = new ArrayType(base_mgr->extents_real(), base_mgr->allocator_real);
+  x_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2real));
+  s_field = new ArrayType(base_mgr->extents_real(), base_mgr->allocator_real);
+  s_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2real));
+
+  cons.print<LOG_DEBUG>("Allocating momentum field");
+
+  // Pass the ownership to state
+  state.newElement("s_hat_field", s_hat_field, true);
+  state.newElement("s_field", s_field);
+  state.newElement("x_hat_field", x_hat_field);
+  state.newElement("x_field", x_field);
+
+  fwrap(*x_hat_field->array) = 0;
+  fwrap(*x_field->array) = 0;
+
+  volume = L0 * L1 * L2;
+  volNorm = volume / Ntot;
+
+  state.newScalar("hmc_force_save_final", true);
+
+  auto tmp_field = base_mgr->allocate_array();
+  synthesis_plan = base_mgr->create_c2r_plan(
+      x_hat_field->array->data(), tmp_field.get_array().data());
+  analysis_plan = base_mgr->create_r2c_plan(
+      tmp_field.get_array().data(), x_hat_field->array->data());
+
+  likelihood->initializeLikelihood(state);
+}
+
+FrozenPhaseDensitySampler::~FrozenPhaseDensitySampler() {
+  if (base_mgr) {
+    MFCalls::destroy_plan(analysis_plan);
+    MFCalls::destroy_plan(synthesis_plan);
+  }
+}
+
+void FrozenPhaseDensitySampler::sample(MarkovState &state) {
+  ConsoleContext<LOG_INFO_SINGLE> ctx("hades density field sampler");
+
+  fwrap(*s_hat_field->array) =
+      fwrap(*state.get<CArrayType>("x_hat_field")->array);
+
+  auto tmp_complex = base_mgr->allocate_complex_array();
+  fwrap(tmp_complex.get_array()) = fwrap(*s_hat_field->array) / volume;
+
+  base_mgr->execute_c2r(
+      synthesis_plan, tmp_complex.get_array().data(), s_field->array->data());
+
+  likelihood->updateMetaParameters(state);
+  likelihood->logLikelihood(*s_hat_field->array, false);
+  likelihood->commitAuxiliaryFields(state);
+}

extra/hades/libLSS/samplers/rgen/frozen/frozen_phase_density_sampler.hpp

@@ -0,0 +1,80 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/rgen/frozen/frozen_phase_density_sampler.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_FROZEN_DENSITY_HPP
+#define __LIBLSS_FROZEN_DENSITY_HPP
+
+#include <boost/optional.hpp>
+#include <functional>
+#include <memory>
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/tools/symplectic_integrator.hpp"
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+#include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#include "libLSS/samplers/rgen/density_sampler.hpp"
+
+namespace LibLSS {
+
+  class FrozenPhaseDensitySampler : public GenericDensitySampler {
+  public:
+    typedef ArrayType::ArrayType Array;
+    typedef ArrayType::RefArrayType ArrayRef;
+    typedef CArrayType::ArrayType CArray;
+    typedef CArrayType::RefArrayType CArrayRef;
+    typedef IArrayType::ArrayType IArray;
+
+    typedef std::shared_ptr<GridDensityLikelihoodBase<3>> Likelihood_t;
+
+  protected:
+    typedef FFTW_Manager_3d<double> DFT_Manager;
+
+    MPI_Communication *comm;
+    std::shared_ptr<DFT_Manager> base_mgr;
+    Likelihood_t likelihood;
+
+    FCalls::plan_type analysis_plan, synthesis_plan;
+    size_t N0, N1, N2;
+    size_t startN0, localN0, endN0;
+    double L0, L1, L2, volume, volNorm;
+
+    ArrayType *x_field, *s_field;
+    CArrayType *x_hat_field, *s_hat_field;
+
+    boost::optional<std::string> phaseFilename;
+    std::string dataName;
+
+    auto sqrt_Pk(MarkovState &state);
+
+  public:
+    FrozenPhaseDensitySampler(MPI_Communication *comm, Likelihood_t likelihood);
+    virtual ~FrozenPhaseDensitySampler();
+
+    virtual void generateMockData(MarkovState &state);
+    void generateRandomField(MarkovState &state);
+
+    void
+    setPhaseFile(std::string const &filename, std::string const &objectName) {
+      phaseFilename = filename;
+      dataName = objectName;
+    }
+
+    void restore(MarkovState &state);
+    void initialize(MarkovState &state);
+
+    virtual void sample(MarkovState &state);
+  };
+}; // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/rgen/hmc/hmc_density_sampler.cpp

@@ -0,0 +1,720 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/rgen/hmc/hmc_density_sampler.cpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#include "libLSS/cconfig.h"
+#include <CosmoTool/algo.hpp>
+#include <cmath>
+#include <Eigen/Core>
+#include <boost/format.hpp>
+#include <functional>
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+#include "libLSS/tools/array_tools.hpp"
+#include "libLSS/samplers/rgen/hmc/hmc_density_sampler.hpp"
+#include "libLSS/tools/uninitialized_type.hpp"
+#include <fstream>
+#include <iostream>
+#include "libLSS/samplers/core/generate_random_field.hpp"
+#include "libLSS/samplers/core/powerspec_tools.hpp"
+
+static const bool ULTRA_VERBOSE = false;
+static const bool HMC_PERF_TEST = true;
+static const bool FIXED_INTEGRATION_PATH = false;
+
+using namespace LibLSS;
+using boost::extents;
+using boost::format;
+
+using CosmoTool::square;
+
+using boost::c_storage_order;
+typedef boost::multi_array_types::extent_range range;
+
+typedef Eigen::Map<Eigen::ArrayXd, Eigen::Aligned> MappedArray;
+
+static const int ROOT_RANK = 0;
+
+namespace ph = std::placeholders;
+
+HMCDensitySampler::HMCDensitySampler(
+    MPI_Communication *comm, Likelihood_t likelihood, double k_max_,
+    std::string const &prefix)
+    : momentum_field(0), analysis_plan(0), synthesis_plan(0), attempt_field(0),
+      accept_field(0), bad_sample(0), comm(comm), k_max(k_max_) {
+  this->maxTime = 50;
+  this->maxEpsilon = 0.01;
+  this->likelihood = likelihood;
+  setIntegratorScheme(HMCOption::SI_2A);
+  setupNames(prefix);
+}
+
+void HMCDensitySampler::setupNames(std::string const &prefix) {
+  momentum_field_name = prefix + "momentum_field";
+  s_hat_field_name = prefix + "s_hat_field";
+  s_field_name = prefix + "s_field";
+  hades_attempt_count_name = prefix + "hades_attempt_count";
+  hades_accept_count_name = prefix + "hades_accept_count";
+  hmc_bad_sample_name = prefix + "hmc_bad_sample";
+  hmc_force_save_final_name = prefix + "hmc_force_save_final";
+  hmc_Elh_name = prefix + "hmc_Elh";
+  hmc_Eprior_name = prefix + "hmc_Eprior";
+}
+
+void HMCDensitySampler::generateMockData(MarkovState &state) {
+  likelihood->updateMetaParameters(state);
+  if (!phaseFilename) {
+    generateRandomField(state);
+  }
+  Console::instance().print<LOG_VERBOSE>(
+      format("Max of s_field = %g") % fwrap(*s_field->array).max());
+  likelihood->generateMockData(*s_hat_field->array, state);
+}
+
+void HMCDensitySampler::computeGradientPsi_Likelihood(
+    MarkovState &state, CArrayRef const &s, CArrayRef &grad_array,
+    bool accumulate) {
+  double temp = state.getScalar<double>("ares_heat");
+  if (posttransform) {
+    BoxModel box = posttransform->get_box_model();
+    posttransform->forwardModel_v2(ModelInput<3>(base_mgr, box, s));
+    auto tmp_s = base_mgr->allocate_ptr_complex_array();
+    posttransform->getDensityFinal(
+        ModelOutput<3>(base_mgr, box, tmp_s->get_array()));
+
+    likelihood->gradientLikelihood(*tmp_s, grad_array, accumulate, temp);
+    tmp_s.reset();
+
+    posttransform->adjointModel_v2(
+        ModelInputAdjoint<3>(base_mgr, box, grad_array));
+    posttransform->getAdjointModelOutput(
+        ModelOutputAdjoint<3>(base_mgr, box, grad_array));
+  } else {
+    likelihood->gradientLikelihood(s, grad_array, accumulate, temp);
+  }
+}
+
+double HMCDensitySampler::computeHamiltonian_Likelihood(
+    MarkovState &state, CArrayRef const &s_hat, bool final_call) {
+  double temp = state.getScalar<double>("ares_heat");
+  Console::instance().print<LOG_VERBOSE>(
+      format("[LIKELIHOOD] Temperature is %lg") % temp);
+
+  if (posttransform) {
+    BoxModel box = posttransform->get_box_model();
+    posttransform->forwardModel_v2(ModelInput<3>(base_mgr, box, s_hat));
+    auto tmp_s = base_mgr->allocate_ptr_complex_array();
+    posttransform->getDensityFinal(
+        ModelOutput<3>(base_mgr, box, tmp_s->get_array()));
+    return likelihood->logLikelihood(tmp_s->get_array(), !final_call) * temp;
+  } else
+
+    return likelihood->logLikelihood(s_hat, !final_call) * temp;
+}
+
+void HMCDensitySampler::restore(MarkovState &state) { initialize(state); }
+
+void HMCDensitySampler::initialize(MarkovState &state) {
+  Console &cons = Console::instance();
+  ConsoleContext<LOG_INFO_SINGLE> ctx("Initialize hades density sampler");
+
+  N0 = state.getScalar<long>("N0");
+  N1 = state.getScalar<long>("N1");
+  N2 = state.getScalar<long>("N2");
+
+  base_mgr = std::make_shared<DFT_Manager>(N0, N1, N2, comm);
+  size_t Ntot = N0 * N1 * N2;
+
+  L0 = state.getScalar<double>("L0");
+  L1 = state.getScalar<double>("L1");
+  L2 = state.getScalar<double>("L2");
+  Ncat = state.getScalar<long>("NCAT");
+
+  startN0 = base_mgr->startN0;
+  localN0 = base_mgr->localN0;
+  endN0 = startN0 + localN0;
+
+  fixer = std::make_shared<Hermiticity_fixer<double, 3>>(base_mgr);
+
+  ctx.format("Allocating s_hat_field field: %dx%dx%d", N0, N1, base_mgr->N2_HC);
+  s_hat_field =
+      new CArrayType(base_mgr->extents_complex(), base_mgr->allocator_complex);
+  s_hat_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2_HC));
+  s_field = new ArrayType(extents[range(startN0, startN0 + localN0)][N1][N2]);
+  s_field->setRealDims(ArrayDimension(N0, N1, N2));
+  cons.print<LOG_DEBUG>("Allocating momentum field");
+  momentum_field =
+      new CArrayType(base_mgr->extents_complex(), base_mgr->allocator_complex);
+  momentum_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2_HC));
+
+  // Pass the ownership to state
+  state.newElement(momentum_field_name, momentum_field);
+  state.newElement(s_hat_field_name, s_hat_field, true);
+  state.newElement(s_field_name, s_field, true);
+  state.newElement(hades_attempt_count_name, attempt_field = new SLong(), true);
+  state.newElement(hades_accept_count_name, accept_field = new SLong(), true);
+  bad_sample = state.newScalar<int>(hmc_bad_sample_name, 0);
+
+  attempt_field->value = 0;
+  accept_field->value = 0;
+  attempt_field->setResetOnSave(0);
+  accept_field->setResetOnSave(0);
+  bad_sample->setResetOnSave(0);
+
+  s_hat_field->eigen().fill(0);
+  s_field->eigen().fill(0);
+  momentum_field->eigen().fill(0);
+
+  volume = L0 * L1 * L2;
+  volNorm = volume / Ntot;
+
+  mass_field =
+      new ArrayType(base_mgr->extents_complex(), base_mgr->allocator_real);
+  mass_field->setRealDims(ArrayDimension(N0, N1, base_mgr->N2_HC));
+  mass_field->eigen().fill(0);
+  state.newElement("hades_mass", mass_field);
+
+  state.newScalar(hmc_force_save_final_name, true);
+  state.newScalar(hmc_Elh_name, 0.0, true);
+  state.newScalar(hmc_Eprior_name, 0.0, true);
+
+  auto tmp_field = base_mgr->allocate_array();
+  synthesis_plan = base_mgr->create_c2r_plan(
+      s_hat_field->array->data(), tmp_field.get_array().data());
+  analysis_plan = base_mgr->create_r2c_plan(
+      tmp_field.get_array().data(), s_hat_field->array->data());
+
+  likelihood->initializeLikelihood(state);
+
+  // Now load phases
+  if (phaseFilename) {
+    H5::H5File f(*phaseFilename, H5F_ACC_RDONLY);
+
+    ctx.print("Read-in phase data");
+    {
+      U_Array<double, 3> tmp_x(
+          base_mgr
+              ->extents_real_strict()); // We need to allocate this temporary array to adapt shape.
+      CosmoTool::hdf5_read_array(f, dataName, tmp_x.get_array(), false, true);
+
+      ctx.print("Saving and Rescaling");
+      fwrap(array::slice_array(*s_field->array, base_mgr->strict_range())) =
+          fwrap(tmp_x.get_array());
+    }
+
+    auto tmp_field = base_mgr->allocate_array();
+    fwrap(array::slice_array(tmp_field.get_array(), base_mgr->strict_range())) =
+        fwrap(*s_field->array) * volNorm;
+
+    ctx.print("Fourier transform");
+    base_mgr->execute_r2c(
+        analysis_plan, tmp_field.get_array().data(),
+        s_hat_field->array->data());
+  }
+}
+
+HMCDensitySampler::~HMCDensitySampler() {
+  if (base_mgr) {
+    Console::instance().print<LOG_INFO_SINGLE>("Cleaning up HMCDensitySampler");
+
+    MFCalls::destroy_plan(analysis_plan);
+    MFCalls::destroy_plan(synthesis_plan);
+  }
+}
+
+auto HMCDensitySampler::free_phase_mask() {
+  double kmax2 = k_max * k_max;
+
+  return fwrap(b_fused_idx<double, 3>([this, kmax2](int a, int b, int c) {
+    double kx = kmode(a, N0, L0);
+    double ky = kmode(b, N1, L1);
+    double kz = kmode(c, N2, L2);
+
+    return (kx * kx + ky * ky + kz * kz) > kmax2;
+  }));
+}
+
+void HMCDensitySampler::initializeMomenta(MarkovState &state) {
+  CArrayType::ArrayType &momentum_array = *momentum_field->array;
+  RandomGen *rgen = state.get<RandomGen>("random_generator");
+  fwrap(momentum_array) = make_complex(
+      rgen->get().gaussian(
+          constant<double, 3>(M_SQRT2, base_mgr->extents_complex())),
+      rgen->get().gaussian(
+          constant<double, 3>(M_SQRT2, base_mgr->extents_complex())));
+
+  fwrap(momentum_array) = fwrap(momentum_array) * free_phase_mask();
+}
+
+#include "hmc_kinetic.tcc"
+#include "hmc_prior.tcc"
+
+void HMCDensitySampler::computeGradientPsi(
+    MarkovState &state, CArrayRef const &s, CArrayRef &grad_array) {
+  array::fill(grad_array, 0);
+  computeGradientPsi_Prior(state, s, grad_array);
+  computeGradientPsi_Likelihood(state, s, grad_array, true);
+  fwrap(grad_array) = fwrap(grad_array) * free_phase_mask();
+}
+
+HMCDensitySampler::HamiltonianType HMCDensitySampler::computeHamiltonian(
+    MarkovState &state, CArrayRef const &s_hat, bool final_call) {
+  ConsoleContext<LOG_DEBUG> ctx("hamiltonian computation");
+
+  HamiltonianType Ekin = computeHamiltonian_Kinetic();
+  ctx.print(format("Ekin = %lg") % double(Ekin));
+  HamiltonianType Eprior = computeHamiltonian_Prior(state, s_hat);
+  ctx.print(format("Eprior = %lg") % double(Eprior));
+  HamiltonianType Elh = computeHamiltonian_Likelihood(state, s_hat, final_call);
+  ctx.print(format("Elh = %lg") % double(Elh));
+
+  return Ekin + Eprior + Elh;
+}
+
+void HMCDensitySampler::setTransforms(
+    Model_t pretransform_, Model_t posttransform_) {
+  pretransform = pretransform_;
+  posttransform = posttransform_;
+}
+
+void HMCDensitySampler::setIntegratorScheme(IntegratorScheme scheme) {
+  current_scheme = scheme;
+  symp.setIntegratorScheme(scheme);
+}
+
+void HMCDensitySampler::doSympInt(MarkovState &state, CArrayRef &s_hat) {
+  ConsoleContext<LOG_INFO_SINGLE> ctx("Symplectic integration");
+  RandomGen *rgen = state.get<RandomGen>("random_generator");
+  int Ntime;
+  double epsilon;
+  ArrayType::ArrayType &mass = *mass_field->array;
+  CArrayType::ArrayType &momentum_array = *momentum_field->array;
+
+  if (comm->rank() == ROOT_RANK) {
+    lastEpsilon = epsilon = maxEpsilon * rgen->get().uniform();
+    do {
+      lastTime = Ntime = 1 + int(std::floor(rgen->get().uniform() * maxTime));
+    } while (Ntime == 0);
+  }
+
+  if (HMC_PERF_TEST && FIXED_INTEGRATION_PATH) {
+    epsilon = maxEpsilon;
+    Ntime = maxTime;
+  }
+
+  comm->broadcast_t(&epsilon, 1, ROOT_RANK);
+  comm->broadcast_t(&Ntime, 1, ROOT_RANK);
+
+  ctx.print(format("epsilon = %lg, Ntime = %d") % epsilon % Ntime);
+
+  auto gradient_psi_p = base_mgr->allocate_complex_array();
+  auto &gradient_psi = gradient_psi_p.get_array();
+
+  symp.integrate(
+      std::bind(
+          &HMCDensitySampler::computeGradientPsi, this, std::ref(state), ph::_1,
+          ph::_2),
+      mass, epsilon, Ntime, s_hat, momentum_array, gradient_psi);
+}
+
+void HMCDensitySampler::updateMass(MarkovState &state) {
+  LIBLSS_AUTO_DEBUG_CONTEXT(ctx);
+  ArrayType::ArrayType &mass = *mass_field->array;
+
+  fwrap(mass) = 1 / 2.0 * free_phase_mask();
+}
+
+void HMCDensitySampler::sample(MarkovState &state) {
+  RandomGen *rgen = state.get<RandomGen>("random_generator");
+  double HamiltonianInit, HamiltonianFinal, deltaH;
+  ConsoleContext<LOG_INFO_SINGLE> ctx("hades density field sampler");
+  //adjust_field = state.get<IArrayType>("adjust_mode_multiplier");
+
+  auto s_hat_p = base_mgr->allocate_complex_array();
+  auto &s_hat = s_hat_p.get_array();
+
+  if (pretransform) {
+    BoxModel box = pretransform->get_box_model();
+    // We explicitly protect the input here.
+    pretransform->forwardModel_v2(ModelInput<3>(
+        base_mgr, box,
+        (CArrayRef const &)*state.get<CArrayType>("s_hat_field")->array));
+    pretransform->getDensityFinal(ModelOutput<3>(base_mgr, box, s_hat));
+  } else
+    array::copyArray3d(s_hat, *state.get<CArrayType>("s_hat_field")->array);
+
+  updateMass(state);
+
+  likelihood->updateMetaParameters(state);
+  initializeMomenta(state);
+
+  HamiltonianType init_Ekin = computeHamiltonian_Kinetic();
+  HamiltonianType init_Eprior = computeHamiltonian_Prior(state, s_hat);
+  HamiltonianType init_Elh = computeHamiltonian_Likelihood(state, s_hat, false);
+  HamiltonianInit = init_Ekin + init_Eprior +
+                    init_Elh; //computeHamiltonian(state, s_hat, false);
+  // If we are the very first step, save the result of the forward model for the other samplers.
+  bool &force_save = state.getScalar<bool>("hmc_force_save_final");
+  if (state.get<SLong>("MCMC_STEP")->value == 0 || force_save) {
+    likelihood->commitAuxiliaryFields(state);
+    force_save = false;
+  }
+
+  if (state.get<SBool>("hades_sampler_blocked")->value)
+    return;
+
+  boost::chrono::system_clock::time_point time_start;
+  if (HMC_PERF_TEST) {
+    time_start = boost::chrono::system_clock::now();
+  }
+
+  HamiltonianType final_Ekin = 0;
+  HamiltonianType final_Eprior = 0;
+  HamiltonianType final_Elh = 0;
+
+  try {
+    if (std::isnan(HamiltonianInit)) {
+      error_helper<ErrorBadState>("NaN in hamiltonian initial");
+    }
+    doSympInt(state, s_hat);
+    final_Ekin = computeHamiltonian_Kinetic();
+    final_Eprior = computeHamiltonian_Prior(state, s_hat);
+    final_Elh = computeHamiltonian_Likelihood(state, s_hat, true);
+    HamiltonianFinal = final_Ekin + final_Eprior + final_Elh;
+    //if (std::isnan(HamiltonianFinal)) { error_helper<ErrorBadState>("NaN in hamiltonian final"); }
+
+  } catch (const ErrorLoadBalance &) {
+    // Stop everything now
+    state.getScalar<int>(hmc_bad_sample_name)++;
+    if (HMC_PERF_TEST && comm->rank() == ROOT_RANK) {
+      std::ofstream f("hmc_performance.txt", std::ios::app);
+      boost::chrono::duration<double> compute_time =
+          boost::chrono::system_clock::now() - time_start;
+
+      f << format("% 10.5le % 6d % 15.15le % 15.15le %d 0 0") % lastEpsilon %
+               lastTime % 0 % compute_time.count() % int(current_scheme)
+        << std::endl;
+    }
+    return;
+  }
+  double log_u;
+
+  ctx.format(
+      "init_Ekin=%g, init_Eprior=%g, init_Elh=%g, final_Ekin=%g, "
+      "final_Eprior=%g, final_Elh=%g",
+      init_Ekin, init_Eprior, init_Elh, final_Ekin, final_Eprior, final_Elh);
+
+  deltaH = HamiltonianFinal - HamiltonianInit;
+  if (comm->rank() == ROOT_RANK)
+    log_u = log(rgen->get().uniform());
+
+  comm->broadcast_t(&log_u, 1, ROOT_RANK);
+
+  if (attempt_field)
+    attempt_field->value++;
+  ctx.print2<LOG_INFO_SINGLE>(
+      format("log_u = %lg, deltaH = %lg") % log_u % deltaH);
+  if (HMC_PERF_TEST && comm->rank() == ROOT_RANK) {
+    std::ofstream f("hmc_performance.txt", std::ios::app);
+    boost::chrono::duration<double> compute_time =
+        boost::chrono::system_clock::now() - time_start;
+
+    f << format("% 10.5le % 6d % 15.15le % 15.15le %d % 15.15le 0 %d") %
+             lastEpsilon % lastTime % deltaH % compute_time.count() %
+             int(current_scheme) % HamiltonianFinal % (log_u <= -deltaH)
+      << std::endl;
+  }
+
+  if (log_u <= -deltaH) {
+    // Accept the move
+    if (comm->rank() == ROOT_RANK)
+      ctx.print("accepting the move");
+
+    likelihood->commitAuxiliaryFields(state);
+
+    state.getScalar<double>(hmc_Elh_name) = final_Elh;
+    state.getScalar<double>(hmc_Eprior_name) = final_Eprior;
+
+    auto tmp_real_field = base_mgr->allocate_array();
+    fixer->forward(s_hat);
+    if (posttransform) {
+      BoxModel box = posttransform->get_box_model();
+      posttransform->forwardModel_v2(ModelInput<3>(base_mgr, box, s_hat));
+      posttransform->getDensityFinal(ModelOutput<3>(
+          base_mgr, box, *state.get<CArrayType>(s_hat_field_name)->array));
+      LibLSS::copy_array(
+          s_hat, *state.get<CArrayType>(s_hat_field_name)->array);
+    } else {
+      LibLSS::copy_array(
+          *state.get<CArrayType>(s_hat_field_name)->array, s_hat);
+    }
+    base_mgr->execute_c2r(
+        synthesis_plan, s_hat.data(), tmp_real_field.get_array().data());
+
+    // This one handles padded and unpadded data through multi_array
+    array::scaleAndCopyArray3d(
+        *state.get<ArrayType>(s_field_name)->array, tmp_real_field.get_array(),
+        1. / volume, true);
+
+    if (accept_field)
+      accept_field->value++;
+  } else {
+    state.getScalar<double>(hmc_Elh_name) = init_Elh;
+    state.getScalar<double>(hmc_Eprior_name) = init_Eprior;
+  }
+}
+
+void HMCDensitySampler::checkGradientReal(MarkovState &state, int step) {
+  ConsoleContext<LOG_DEBUG> ctx("checkGradientReal");
+  HamiltonianType H0, H1, H2;
+  HamiltonianType H0prior, H1prior, H2prior;
+  HamiltonianType H0poisson, H1poisson, H2poisson;
+  auto s_hat_p = base_mgr->allocate_complex_array();
+  auto &s_hat = s_hat_p.get_array();
+  double epsilon = 0.001;
+
+  ArrayType *gradient_field_lh, *gradient_field_lh_ref;
+  ArrayType *s_field = state.get<ArrayType>(s_field_name);
+
+  if (state.exists("gradient_array_lh_real")) {
+    gradient_field_lh = state.get<ArrayType>("gradient_array_lh_real");
+    gradient_field_lh_ref = state.get<ArrayType>("gradient_array_lh_ref_real");
+  } else {
+    auto real_dim =
+        ArrayDimension(base_mgr->N0, base_mgr->N1, base_mgr->N2real);
+
+    state.newElement(
+        "gradient_array_lh_real",
+        gradient_field_lh =
+            new ArrayType(base_mgr->extents_real(), base_mgr->allocator_real));
+    gradient_field_lh->setRealDims(real_dim);
+
+    state.newElement(
+        "gradient_array_lh_ref_real",
+        gradient_field_lh_ref =
+            new ArrayType(base_mgr->extents_real(), base_mgr->allocator_real));
+    gradient_field_lh_ref->setRealDims(real_dim);
+  }
+
+  auto &gradient_array_lh = *gradient_field_lh->array;
+  auto &gradient_array_lh_ref = *gradient_field_lh_ref->array;
+  auto &s = *s_field->array;
+
+  double log_L0, log_L1, log_L2;
+  likelihood->gradientLikelihood(s, gradient_array_lh, false, 1.0);
+
+  log_L0 = likelihood->logLikelihood(s, false);
+  Progress<LOG_INFO_SINGLE> &progress =
+      Console::instance().start_progress<LOG_INFO_SINGLE>(
+          "doing numerical gradient (real)",
+          base_mgr->N0 * base_mgr->N1 * base_mgr->N2, 5);
+
+  for (int n0 = 0; n0 < N0; n0 += step) {
+    for (int n1 = 0; n1 < N1; n1++) {
+      for (int n2 = 0; n2 < base_mgr->N2; n2++) {
+        double backup;
+        double pert;
+        double n_backup;
+
+        bool oncore = base_mgr->on_core(n0);
+
+        if (n0 == 0 && n1 == 0 && n2 == 0)
+          continue;
+
+        if (oncore) {
+          backup = s[n0][n1][n2];
+          n_backup = abs(backup);
+          if (n_backup == 0)
+            n_backup = 1;
+          pert = backup + n_backup * epsilon;
+          s[n0][n1][n2] = pert;
+        }
+
+        log_L1 = likelihood->logLikelihood(s, false);
+
+        if (oncore) {
+          gradient_array_lh_ref[n0][n1][n2] =
+              (log_L1 - log_L0) / (n_backup * epsilon);
+          s[n0][n1][n2] = backup;
+        }
+
+        long n = ((n0 - startN0) * N1 + n1) * base_mgr->N2 + n2;
+        progress.update(n);
+      }
+    }
+  }
+}
+
+void HMCDensitySampler::checkGradient(MarkovState &state, int step) {
+  ConsoleContext<LOG_DEBUG> ctx("checkGradient");
+  HamiltonianType H0, H1, H2;
+  HamiltonianType H0prior, H1prior, H2prior;
+  HamiltonianType H0poisson, H1poisson, H2poisson;
+  auto s_hat_p = base_mgr->allocate_complex_array();
+  auto &s_hat = s_hat_p.get_array();
+  double epsilon = 0.01;
+
+  CArrayType *gradient_field, *gradient_field_ref;
+  CArrayType *gradient_field_prior, *gradient_field_prior_ref;
+  CArrayType *gradient_field_poisson, *gradient_field_poisson_ref;
+  CArrayType *s_hat_field = state.get<CArrayType>(s_hat_field_name);
+
+  if (state.exists("gradient_array")) {
+    gradient_field = state.get<CArrayType>("gradient_array");
+    gradient_field_prior = state.get<CArrayType>("gradient_array_prior");
+    gradient_field_poisson = state.get<CArrayType>("gradient_array_lh");
+    gradient_field_ref = state.get<CArrayType>("gradient_array_ref");
+    gradient_field_prior_ref =
+        state.get<CArrayType>("gradient_array_prior_ref");
+    gradient_field_poisson_ref = state.get<CArrayType>("gradient_array_lh_ref");
+  } else {
+    auto complex_dim =
+        ArrayDimension(base_mgr->N0, base_mgr->N1, base_mgr->N2_HC);
+    state.newElement(
+        "gradient_array",
+        gradient_field = new CArrayType(
+            base_mgr->extents_complex(), base_mgr->allocator_complex));
+    gradient_field->setRealDims(complex_dim);
+
+    state.newElement(
+        "gradient_array_ref",
+        gradient_field_ref = new CArrayType(
+            base_mgr->extents_complex(), base_mgr->allocator_complex));
+    gradient_field_ref->setRealDims(complex_dim);
+
+    state.newElement(
+        "gradient_array_prior",
+        gradient_field_prior = new CArrayType(
+            base_mgr->extents_complex(), base_mgr->allocator_complex));
+    gradient_field_prior->setRealDims(complex_dim);
+
+    state.newElement(
+        "gradient_array_prior_ref",
+        gradient_field_prior_ref = new CArrayType(
+            base_mgr->extents_complex(), base_mgr->allocator_complex));
+    gradient_field_prior_ref->setRealDims(complex_dim);
+
+    state.newElement(
+        "gradient_array_lh",
+        gradient_field_poisson = new CArrayType(
+            base_mgr->extents_complex(), base_mgr->allocator_complex));
+    gradient_field_poisson->setRealDims(complex_dim);
+
+    state.newElement(
+        "gradient_array_lh_ref",
+        gradient_field_poisson_ref = new CArrayType(
+            base_mgr->extents_complex(), base_mgr->allocator_complex));
+    gradient_field_poisson_ref->setRealDims(complex_dim);
+  }
+  FFTW_Complex_Array &grad_array = *gradient_field->array;
+  FFTW_Complex_Array &grad_array_ref = *gradient_field_ref->array;
+  FFTW_Complex_Array &grad_array_prior = *gradient_field_prior->array;
+  FFTW_Complex_Array &grad_array_prior_ref = *gradient_field_prior_ref->array;
+  FFTW_Complex_Array &grad_array_poisson = *gradient_field_poisson->array;
+  FFTW_Complex_Array &grad_array_poisson_ref =
+      *gradient_field_poisson_ref->array;
+  //adjust_field = state.get<IArrayType>("adjust_mode_multiplier");
+
+  updateMass(state);
+  if (startN0 == 0 && localN0 > 0)
+    grad_array[0][0][0] = 0;
+  ctx.print("Initialize momenta");
+  initializeMomenta(state);
+
+  ctx.print("Compute for prior");
+  computeGradientPsi_Prior(state, *s_hat_field->array, grad_array_prior);
+  ctx.print("Compute for likelihood");
+  computeGradientPsi_Likelihood(
+      state, *s_hat_field->array, grad_array_poisson, false);
+  ctx.print("Compute for both");
+  computeGradientPsi(state, *s_hat_field->array, grad_array);
+
+  gradient_field_ref->eigen().fill(0);
+  gradient_field_prior_ref->eigen().fill(0);
+  gradient_field_poisson_ref->eigen().fill(0);
+
+  Progress<LOG_INFO_SINGLE> &progress =
+      Console::instance().start_progress<LOG_INFO_SINGLE>(
+          "doing numerical gradient",
+          base_mgr->N0 * base_mgr->N1 * base_mgr->N2_HC, 5);
+  array::copyArray3d(s_hat, *s_hat_field->array);
+  H0 = computeHamiltonian(state, s_hat, false);
+  H0prior = computeHamiltonian_Prior(state, s_hat);
+  H0poisson = computeHamiltonian_Likelihood(state, s_hat, false);
+  for (int n0 = 0; n0 < N0; n0 += step) {
+    for (int n1 = 0; n1 < N1; n1++) {
+      for (int n2 = 0; n2 < base_mgr->N2_HC; n2++) {
+        FFTW_Complex_Array::element backup;
+        std::complex<double> pert_r, pert_i;
+        double n_backup;
+
+        bool oncore = (n0 >= startN0 && n0 < startN0 + localN0);
+
+        if (n0 == 0 && n1 == 0 && n2 == 0)
+          continue;
+
+        if (oncore) {
+          backup = s_hat[n0][n1][n2];
+          n_backup = abs(backup);
+          if (n_backup == 0)
+            n_backup = 1;
+          Console::instance().print<LOG_DEBUG>(
+              format("n_backup=%lg") % n_backup);
+          pert_r = backup + std::complex<double>(n_backup * epsilon, 0);
+          s_hat[n0][n1][n2] = pert_r;
+        }
+
+        H1 = computeHamiltonian(state, s_hat, false);
+        H1prior = computeHamiltonian_Prior(state, s_hat);
+        H1poisson = computeHamiltonian_Likelihood(state, s_hat, false);
+
+        if (oncore) {
+          pert_i = backup + std::complex<double>(0, n_backup * epsilon);
+          s_hat[n0][n1][n2] = pert_i;
+        }
+
+        H2 = computeHamiltonian(state, s_hat, false);
+        H2prior = computeHamiltonian_Prior(state, s_hat);
+        H2poisson = computeHamiltonian_Likelihood(state, s_hat, false);
+
+        if (oncore) {
+          grad_array_ref[n0][n1][n2] =
+              std::complex<double>((H1), (H0)) / (n_backup * epsilon);
+          grad_array_prior_ref[n0][n1][n2] =
+              std::complex<double>((H1prior - H0prior), (H2prior - H0prior)) /
+              (n_backup * epsilon);
+          grad_array_poisson_ref[n0][n1][n2] =
+              std::complex<double>(
+                  (H1poisson - H0poisson), (H2poisson - H0poisson)) /
+              (n_backup * epsilon);
+          s_hat[n0][n1][n2] = backup;
+        }
+
+        long n = ((n0 - startN0) * N1 + n1) * base_mgr->N2_HC + n2;
+        progress.update(n);
+      }
+    }
+  }
+
+  progress.destroy();
+}
+
+void HMCDensitySampler::generateRandomField(MarkovState &state) {
+  LibLSS::generateRandomField(comm, state);
+}
+
+HMCDensitySampler::HamiltonianType
+HMCDensitySampler::computeHamiltonian(MarkovState &state, bool gradient_next) {
+  auto s_hat_p = base_mgr->allocate_complex_array();
+  auto &s_hat = s_hat_p.get_array();
+
+  array::copyArray3d(s_hat, *state.get<CArrayType>(s_hat_field_name)->array);
+
+  return computeHamiltonian_Likelihood(state, s_hat, gradient_next);
+}

extra/hades/libLSS/samplers/rgen/hmc/hmc_density_sampler.hpp

@@ -0,0 +1,247 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/rgen/hmc/hmc_density_sampler.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_HMC_DENSITY_HPP
+#define __LIBLSS_HMC_DENSITY_HPP
+
+#include <functional>
+#include <memory>
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/tools/symplectic_integrator.hpp"
+#include "libLSS/tools/mpi_fftw_helper.hpp"
+#include "libLSS/samplers/core/gridLikelihoodBase.hpp"
+#include "libLSS/samplers/rgen/density_sampler.hpp"
+#include "libLSS/tools/hermiticity_fixup.hpp"
+
+namespace LibLSS {
+
+  namespace HMC_Details {
+
+    template <typename Function>
+    void accumulateHamiltonian(
+        int n0, int n1, int n2, const IArrayType::ArrayType &adjust_array,
+        const Function &f, double &E) {
+      E += f(n0, n1, n2, adjust_array);
+    }
+
+  } // namespace HMC_Details
+
+  namespace HMCOption {
+    using namespace SymplecticOption;
+  };
+
+  class HMCDensitySampler : public GenericDensitySampler {
+  public:
+    typedef ArrayType::ArrayType Array;
+    typedef ArrayType::RefArrayType ArrayRef;
+    typedef CArrayType::ArrayType CArray;
+    typedef CArrayType::RefArrayType CArrayRef;
+    typedef IArrayType::ArrayType IArray;
+
+    typedef double HamiltonianType;
+
+    typedef std::shared_ptr<GridDensityLikelihoodBase<3>> Likelihood_t;
+    typedef std::shared_ptr<BORGForwardModel> Model_t;
+
+  protected:
+    typedef boost::multi_array_ref<IArrayType::ArrayType::element, 1>
+        FlatIntType;
+    typedef FFTW_Manager_3d<double> DFT_Manager;
+
+    MPI_Communication *comm;
+
+    std::shared_ptr<DFT_Manager> base_mgr;
+    Model_t pretransform, posttransform;
+    Likelihood_t likelihood;
+
+    long fourierLocalSize;
+    FCalls::plan_type analysis_plan, synthesis_plan;
+    size_t N0, N1, N2;
+    size_t startN0, localN0, endN0;
+    double L0, L1, L2, volume, volNorm;
+    int Ncat;
+    SLong *attempt_field, *accept_field;
+    ScalarStateElement<int> *bad_sample;
+
+    ArrayType *mass_field, *s_field;
+    IArrayType *adjust_field;
+    CArrayType *momentum_field, *s_hat_field;
+
+    int maxTime, lastTime;
+    double lastEpsilon, maxEpsilon;
+    SymplecticIntegrators symp;
+
+    boost::optional<std::string> phaseFilename;
+    std::string dataName;
+
+    std::string momentum_field_name, s_hat_field_name, s_field_name,
+        hades_attempt_count_name, hades_accept_count_name, hmc_bad_sample_name,
+        hmc_force_save_final_name, hmc_Elh_name, hmc_Eprior_name;
+
+    void setupNames(std::string const &prefix);
+
+    void doSympInt(MarkovState &state, CArrayRef &s_hat);
+    HamiltonianType computeHamiltonian(
+        MarkovState &state, CArrayRef const &s_hat, bool final_call);
+    HamiltonianType
+    computeHamiltonian_Prior(MarkovState &state, CArrayRef const &s_hat);
+    HamiltonianType computeHamiltonian_Kinetic();
+    HamiltonianType computeHamiltonian_Likelihood(
+        MarkovState &state, CArrayRef const &s_hat, bool final_call);
+    void initializeMomenta(MarkovState &state);
+    void computeGradientPsi(
+        MarkovState &state, CArrayRef const &s, CArrayRef &grad_array);
+    void computeGradientPsi_Prior(
+        MarkovState &state, CArrayRef const &s, CArrayRef &grad_array);
+
+    void updateMass(MarkovState &state);
+
+    void computeGradientPsi_Likelihood(
+        MarkovState &state, CArrayRef const &s, CArrayRef &grad_array,
+        bool accumulate);
+
+    typedef boost::multi_array<double, 2> IntegratorCoefficients;
+    typedef HMCOption::IntegratorScheme IntegratorScheme;
+
+    IntegratorScheme current_scheme;
+    std::shared_ptr<Hermiticity_fixer<double, 3>> fixer;
+    double k_max;
+
+    auto free_phase_mask();
+
+  public:
+    HMCDensitySampler(
+        MPI_Communication *comm, Likelihood_t likelihood, double k_max_ = 1000,
+        std::string const &prefix = std::string());
+    virtual ~HMCDensitySampler();
+
+    void
+    setPhaseFile(std::string const &filename, std::string const &objectName) {
+      phaseFilename = filename;
+      dataName = objectName;
+    }
+
+    virtual void generateMockData(MarkovState &state);
+    void generateRandomField(MarkovState &state);
+
+    void setTransforms(Model_t pretransform_, Model_t posttransform_);
+    void setIntegratorScheme(IntegratorScheme scheme);
+
+    double computeHamiltonian(MarkovState &state, bool gradient_next = false);
+
+    void restore(MarkovState &state);
+    void initialize(MarkovState &state);
+
+    virtual void sample(MarkovState &state);
+
+    void checkGradient(MarkovState &state, int step = 10);
+    void checkGradientReal(MarkovState &state, int step = 10);
+
+    void setMaxEpsilon(double eps) { this->maxEpsilon = eps; }
+    void setMaxTimeSteps(int Ntime) { this->maxTime = Ntime; }
+
+    // Codelet generic functions
+    template <typename Function>
+    void codeletGeneral(IArray &adjust_array, Function codelet) {
+      size_t N2_HC = base_mgr->N2_HC;
+#pragma omp parallel for schedule(static) collapse(3)
+      for (size_t n0 = startN0; n0 < endN0; n0++) {
+        for (size_t n1 = 0; n1 < N1; n1++) {
+          for (size_t n2 = 1; n2 < N2_HC - 1; n2++) {
+            codelet(n0, n1, n2, adjust_array);
+          }
+        }
+      }
+
+      codeletNyquist(0, adjust_array, codelet);
+      codeletNyquist(base_mgr->N2_HC - 1, adjust_array, codelet);
+    }
+
+    template <typename Function>
+    void codeletNyquist(int n2, const IArray &adjust_array, Function codelet) {
+      size_t N0end = std::min(endN0, N0);
+      size_t N0start = std::max(startN0, size_t(0));
+#pragma omp parallel for schedule(static) collapse(2)
+      for (size_t n0 = N0start; n0 < N0end; n0++) {
+        for (size_t n1 = 0; n1 < N1; n1++) {
+          codelet(n0, n1, n2, adjust_array);
+        }
+      }
+
+      if (startN0 == 0 && localN0 > 0)
+        for (size_t n1 = 0; n1 < N1; n1++)
+          codelet(0, n1, n2, adjust_array);
+
+      if (startN0 <= N0 / 2 && startN0 + localN0 > N0 / 2)
+        for (size_t n1 = 0; n1 < N1; n1++)
+          codelet(N0 / 2, n1, n2, adjust_array);
+    }
+
+    template <typename Function>
+    HamiltonianType codeletNyquistHamiltonian(
+        int n2, const IArray &adjust_array, Function codelet) {
+      size_t N0end = std::min(endN0, N0);
+      size_t N0start = std::max(startN0, size_t(0));
+      HamiltonianType E = 0;
+
+#pragma omp parallel for schedule(static) collapse(2) reduction(+ : E)
+      for (size_t n0 = N0start; n0 < N0end; n0++) {
+        for (size_t n1 = 0; n1 < N1; n1++) {
+          E += codelet(n0, n1, n2, adjust_array);
+        }
+      }
+
+      if (startN0 == 0 && localN0 > 0)
+#pragma omp parallel for schedule(static) reduction(+ : E)
+        for (size_t n1 = 0; n1 < N1; n1++)
+          E += codelet(0, n1, n2, adjust_array);
+
+      if (startN0 <= N0 / 2 && startN0 + localN0 > N0 / 2)
+#pragma omp parallel for schedule(static) reduction(+ : E)
+        for (size_t n1 = 0; n1 < N1; n1++)
+          E += codelet(N0 / 2, n1, n2, adjust_array);
+
+      return E;
+    }
+
+    template <typename Function>
+    HamiltonianType
+    codeletGeneralHamiltonian(const IArray &adjust_array, Function codelet) {
+      HamiltonianType E = 0;
+      namespace ph = std::placeholders;
+      size_t N2_HC = base_mgr->N2_HC;
+
+#pragma omp parallel for schedule(static) collapse(3) reduction(+ : E)
+      for (size_t n0 = startN0; n0 < startN0 + localN0; n0++) {
+        for (size_t n1 = 0; n1 < N1; n1++) {
+          for (size_t n2 = 1; n2 < N2_HC - 1; n2++) {
+            E += codelet(n0, n1, n2, adjust_array);
+          }
+        }
+      }
+
+      E += codeletNyquistHamiltonian(
+          0, adjust_array, std::bind(codelet, ph::_1, ph::_2, ph::_3, ph::_4));
+      E += codeletNyquistHamiltonian(
+          base_mgr->N2_HC - 1, adjust_array,
+          std::bind(codelet, ph::_1, ph::_2, ph::_3, ph::_4));
+
+      return E;
+    }
+  };
+
+}; // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/rgen/hmc/hmc_kinetic.tcc

@@ -0,0 +1,13 @@
+HMCDensitySampler::HamiltonianType
+HMCDensitySampler::computeHamiltonian_Kinetic() {
+  CArrayType::ArrayType &momentum_array = *momentum_field->array;
+
+
+  auto m_w = fwrap(momentum_array);
+  auto r = std::real(m_w);
+  auto i = std::imag(m_w);
+
+  double Ekin = (r*r+i*i).sum() / (2.);
+  comm->all_reduce_t(MPI_IN_PLACE, &Ekin, 1, MPI_SUM);
+  return Ekin*0.5;
+}

extra/hades/libLSS/samplers/rgen/hmc/hmc_prior.tcc

@@ -0,0 +1,17 @@
+HMCDensitySampler::HamiltonianType HMCDensitySampler::computeHamiltonian_Prior(
+    MarkovState &state, CArrayRef const &s_array) {
+  auto s_w = fwrap(s_array);
+  auto r = std::real(s_w);
+  auto i = std::imag(s_w);
+  double Eprior = (r*r+i*i).sum() / (0.5);
+
+  comm->all_reduce_t(MPI_IN_PLACE, &Eprior, 1, MPI_SUM);
+
+  return Eprior*0.5;
+}
+
+void HMCDensitySampler::computeGradientPsi_Prior(
+    MarkovState &state, CArrayRef const &s, CArrayRef &grad_array) {
+  ConsoleContext<LOG_DEBUG> ctx("gradient psi prior");
+  fwrap(grad_array) = fwrap(s)*2.0;
+}

extra/hades/libLSS/samplers/rgen/nuts/nuts_density_sampler.hpp

@@ -0,0 +1,244 @@
+/*+
+    ARES/HADES/BORG Package -- ./extra/hades/libLSS/samplers/rgen/nuts/nuts_density_sampler.hpp
+    Copyright (C) 2014-2020 Guilhem Lavaux <guilhem.lavaux@iap.fr>
+    Copyright (C) 2009-2020 Jens Jasche <jens.jasche@fysik.su.se>
+
+    Additional contributions from:
+       Guilhem Lavaux <guilhem.lavaux@iap.fr> (2023)
+    
++*/
+#ifndef __LIBLSS_NUTS_DENSITY_HPP
+#define __LIBLSS_NUTS_DENSITY_HPP
+
+#include <CosmoTool/fourier/fft/fftw_calls.hpp>
+#include "libLSS/mpi/generic_mpi.hpp"
+#include "libLSS/mcmc/global_state.hpp"
+#include "libLSS/tools/fftw_allocator.hpp"
+#include "libLSS/samplers/core/markov.hpp"
+#include "libLSS/samplers/core/random_number.hpp"
+#include "libLSS/samplers/core/types_samplers.hpp"
+#include "libLSS/tools/symplectic_integrator.hpp"
+
+namespace LibLSS {
+
+  namespace NUTS_step {
+
+    //1) store some constants that are required during recursion
+    struct nuts_util {
+      // Constants through each recursion
+      double log_u;
+      double H0;
+      int sign;
+
+      // Aggregators through each recursion
+      int n_tree;
+      double sum_prob;
+      bool criterion;
+    };
+  } // namespace NUTS_step
+
+  namespace NUTS_Details {
+
+    template <typename Function>
+    void accumulateHamiltonian(
+        int n0, int n1, int n2, const IArrayType::ArrayType &adjust_array,
+        const Function &f, double &E) {
+      E += f(n0, n1, n2, adjust_array);
+    }
+
+  } // namespace NUTS_Details
+
+  namespace NUTSOption {
+    using namespace SymplecticOption;
+  };
+
+  class NUTSDensitySampler : public MarkovSampler {
+  public:
+    typedef ArrayType::ArrayType Array;
+    typedef ArrayType::RefArrayType ArrayRef;
+    typedef CArrayType::ArrayType CArray;
+    typedef CArrayType::RefArrayType CArrayRef;
+    typedef IArrayType::ArrayType IArray;
+
+    typedef double HamiltonianType;
+
+  protected:
+    typedef boost::multi_array_ref<IArrayType::ArrayType::element, 1>
+        FlatIntType;
+
+    MPI_Communication *comm;
+
+    long fourierLocalSize;
+    FCalls::plan_type analysis_plan, synthesis_plan;
+    FlatIntType *flat_key;
+    long N0, N1, N2, N2real, Ntot, Ntot_k, N2_HC;
+    long startN0, localN0, localNtot, localNtot_k;
+    double L0, L1, L2, volume, volNorm;
+    int Ncat;
+    SLong *attempt_field, *accept_field;
+
+    FFTW_Allocator<double> allocator_real;
+    FFTW_Allocator<std::complex<double>> allocator_complex;
+
+    ArrayType *sqrt_mass_field, *s_field;
+    IArrayType *adjust_field;
+    CArrayType *momentum_field, *s_hat_field;
+    FFTW_Real_Array *tmp_real_field;
+    FFTW_Complex_Array *tmp_complex_field;
+    boost::multi_array<int, 1> peer;
+
+    int maxTime;
+    double maxEpsilon;
+    SymplecticIntegrators symp;
+
+    void doSympInt(MarkovState &state, MarkovState &state, CArray &s_hat);
+    void updateMomentum(MarkovState &state, double dt, CArrayRef &force);
+    void updatePosition(double dt, CArray &s_hat);
+    HamiltonianType computeHamiltonian(
+        MarkovState &state, MarkovState &state, CArray &s_hat, bool final_call);
+    HamiltonianType computeHamiltonian_Prior(
+        MarkovState &state, MarkovState &state, CArray &s_hat);
+    HamiltonianType computeHamiltonian_Kinetic();
+    void initializeMomenta(MarkovState &state);
+    void computeGradientPsi(
+        MarkovState &state, MarkovState &state, CArray &s,
+        CArrayRef &grad_array);
+    void computeGradientPsi_Prior(
+        MarkovState &state, MarkovState &state, CArray &s,
+        CArrayRef &grad_array);
+
+    void updateMass(MarkovState &state, MarkovState &state);
+    void Hermiticity_fixup(CArrayType::ArrayType &a);
+    void Hermiticity_fixup_plane(int Nplane, CArrayType::ArrayType &a);
+
+    virtual void computeGradientPsi_Likelihood(
+        MarkovState &state, MarkovState &state, CArray &s,
+        CArrayRef &grad_array, bool accumulate) = 0;
+    virtual HamiltonianType computeHamiltonian_Likelihood(
+        MarkovState &state, MarkovState &state, CArray &s_hat,
+        bool final_call) = 0;
+
+    virtual void
+    saveAuxiliaryAcceptedFields(MarkovState &state, MarkovState &state) {}
+
+    void computeFourierSpace_GradientPsi(
+        ArrayRef &real_gradient, CArrayRef &grad_array, bool accumulate);
+
+    typedef boost::multi_array<double, 2> IntegratorCoefficients;
+    typedef NUTSOption::IntegratorScheme IntegratorScheme;
+
+    IntegratorCoefficients I_coefs;
+
+  public:
+    NUTSDensitySampler(
+        MPI_Communication *comm, int maxTimeIntegration, double maxEpsilon);
+    virtual ~NUTSDensitySampler();
+
+    virtual void generateMockData(MarkovState &state, MarkovState &state) = 0;
+    void generateRandomField(MarkovState &state, MarkovState &state);
+
+    void setIntegratorScheme(IntegratorScheme scheme);
+    //        void setIntegratorScheme(const IntegratorCoefficients& coefs);
+
+    HamiltonianType computeHamiltonian(MarkovState &state, MarkovState &state);
+
+    void restore_NUTS(MarkovState &state, MarkovState &state);
+    void initialize_NUTS(MarkovState &state, MarkovState &state);
+
+    virtual void sample(MarkovState &state, MarkovState &state);
+
+    void checkGradient(MarkovState &state, MarkovState &state, int step = 10);
+    void checkHermiticityFixup(MarkovState &state, MarkovState &state);
+
+    void setMaxEpsilon(double eps) { this->maxEpsilon = eps; }
+    void setMaxTimeSteps(int Ntime) { this->maxTime = Ntime; }
+
+    // Codelet generic functions
+    template <typename Function>
+    void codeletGeneral(IArray &adjust_array, Function codelet) {
+#pragma omp parallel for schedule(static)
+      for (int n0 = startN0; n0 < startN0 + localN0; n0++) {
+        for (int n1 = 0; n1 < N1; n1++) {
+          for (int n2 = 1; n2 < N2_HC - 1; n2++) {
+            codelet(n0, n1, n2, adjust_array);
+          }
+        }
+      }
+
+      codeletNyquist(0, adjust_array, codelet);
+      codeletNyquist(N2_HC - 1, adjust_array, codelet);
+    }
+
+    template <typename Function>
+    void codeletNyquist(int n2, const IArray &adjust_array, Function codelet) {
+      int N0end = std::min(startN0 + localN0, N0 / 2);
+      int N0start = std::max(startN0, 1L);
+#pragma omp parallel for schedule(static)
+      for (int n0 = N0start; n0 < N0end; n0++) {
+        for (int n1 = 0; n1 < N1; n1++) {
+          codelet(n0, n1, n2, adjust_array);
+        }
+      }
+
+      if (startN0 == 0 && localN0 > 0)
+        for (int n1 = 0; n1 <= N1 / 2; n1++)
+          codelet(0, n1, n2, adjust_array);
+
+      if (startN0 <= N0 / 2 && startN0 + localN0 > N0 / 2)
+        for (int n1 = 0; n1 <= N1 / 2; n1++)
+          codelet(N0 / 2, n1, n2, adjust_array);
+    }
+
+    template <typename Function>
+    HamiltonianType codeletNyquistHamiltonian(
+        int n2, const IArray &adjust_array, Function codelet) {
+      int N0end = std::min(startN0 + localN0, N0 / 2);
+      int N0start = std::max(startN0, 1L);
+      HamiltonianType E = 0;
+
+#pragma omp parallel for schedule(static) reduction(+ : E)
+      for (int n0 = N0start; n0 < N0end; n0++) {
+        for (int n1 = 0; n1 < N1; n1++) {
+          E += codelet(n0, n1, n2, adjust_array);
+        }
+      }
+
+      if (startN0 == 0 && localN0 > 0)
+#pragma omp parallel for schedule(static) reduction(+ : E)
+        for (int n1 = 0; n1 <= N1 / 2; n1++)
+          E += codelet(0, n1, n2, adjust_array);
+
+      if (startN0 <= N0 / 2 && startN0 + localN0 > N0 / 2)
+#pragma omp parallel for schedule(static) reduction(+ : E)
+        for (int n1 = 0; n1 <= N1 / 2; n1++)
+          E += codelet(N0 / 2, n1, n2, adjust_array);
+
+      return E;
+    }
+
+    template <typename Function>
+    HamiltonianType
+    codeletGeneralHamiltonian(const IArray &adjust_array, Function codelet) {
+      HamiltonianType E = 0;
+
+#pragma omp parallel for schedule(static) reduction(+ : E)
+      for (int n0 = startN0; n0 < startN0 + localN0; n0++) {
+        for (int n1 = 0; n1 < N1; n1++) {
+          for (int n2 = 1; n2 < N2_HC - 1; n2++) {
+            E += codelet(n0, n1, n2, adjust_array);
+          }
+        }
+      }
+
+      E += codeletNyquistHamiltonian(
+          0, adjust_array, boost::bind(codelet, _1, _2, _3, _4));
+      E += codeletNyquistHamiltonian(
+          N2_HC - 1, adjust_array, boost::bind(codelet, _1, _2, _3, _4));
+
+      return E;
+    }
+  };
+
+}; // namespace LibLSS
+
+#endif

extra/hades/libLSS/samplers/rgen/nuts/nuts_kinetic.tcc

@@ -0,0 +1,34 @@
+inline
+HMCDensitySampler::HamiltonianType codeletHamiltonianKinetic(int n0, int n1, int n2,
+                               const IArrayType::ArrayType& adjust_array,
+                               const ArrayType::ArrayType& sqM,
+                               const CArrayType::ArrayType& momentum_array
+                            )
+{
+    const CArrayType::ArrayType::element& e = momentum_array[n0][n1][n2];
+    int adjust = adjust_array[n0][n1][n2];
+    double Amplitude = (sqM[n0][n1][n2]);
+    if (Amplitude == 0 || adjust == 0)
+        return 0;
+    
+    double Ekin = (square(e.real()) + square(e.imag()))/Amplitude;
+    if (isnan(Ekin)) { 
+        error_helper<ErrorBadState>(format("NaN in kinetic hamiltonian n0=%d n1=%d n2=%d Mass=%lg") % n0 % n1 % n2 % Amplitude); 
+    }
+    
+    return Ekin;
+}
+
+HMCDensitySampler::HamiltonianType HMCDensitySampler::computeHamiltonian_Kinetic()
+{
+    CArrayType::ArrayType& momentum_array = *momentum_field->array;
+    ArrayType::ArrayType& sqrt_mass = *sqrt_mass_field->array;
+    IArrayType::ArrayType& adjust_array = *adjust_field->array;
+
+    double Ekin = codeletGeneralHamiltonian(adjust_array, 
+            boost::bind(codeletHamiltonianKinetic, _1, _2, _3, _4, boost::cref(sqrt_mass), boost::cref(momentum_array)));
+
+    Ekin *= 0.5;
+    comm->all_reduce_t(MPI_IN_PLACE, &Ekin, 1, MPI_SUM);
+    return Ekin;
+}

extra/hades/libLSS/samplers/rgen/nuts/nuts_prior.tcc

@@ -0,0 +1,78 @@
+inline
+HMCDensitySampler::HamiltonianType codeletHamiltonianPrior(int n0, int n1, int n2,
+                               const IArrayType::ArrayType& adjust_array,
+                               const ArrayType1d::ArrayType& pspec,
+                               const IArrayType::ArrayType& key_array,
+                               const CArrayType::ArrayType& s_array
+                            )
+{
+    const CArrayType::ArrayType::element& e = s_array[n0][n1][n2];
+    long powerPosition = key_array[n0][n1][n2];
+    double Amplitude = pspec.data()[powerPosition];
+    double adjust = adjust_array[n0][n1][n2];
+    
+    if (Amplitude == 0)
+        return 0;
+    
+    double Eprior = adjust * (square(e.real()) + square(e.imag()))/Amplitude;
+
+    if (isnan(Eprior)) { 
+        error_helper<ErrorBadState>("NaN in hamiltonian"); 
+    }
+    
+    return Eprior;
+}
+
+inline
+HMCDensitySampler::HamiltonianType codeletGradientPrior(int n0, int n1, int n2,
+                            const IArrayType::ArrayType& adjust_array,
+                            double volume,
+                            const ArrayType1d::ArrayType& pspec,
+                            const IArrayType::ArrayType& key_array,
+                            const CArrayType::ArrayType& s,
+                            CArrayType::RefArrayType& grad_array    
+                            )
+{
+    const CArrayType::ArrayType::element& e = s[n0][n1][n2];
+    long powerPosition = key_array[n0][n1][n2];
+    double Amplitude = pspec.data()[powerPosition] * volume;
+    CArrayType::ArrayType::element& gradient = grad_array[n0][n1][n2];
+    double adjust = adjust_array[n0][n1][n2];
+
+    if (Amplitude == 0 || adjust == 0)
+        gradient = 0;
+    else
+        gradient = adjust * e / Amplitude;
+}
+
+
+HMCDensitySampler::HamiltonianType HMCDensitySampler::computeHamiltonian_Prior(MarkovState& state, CArray& s_array)
+{
+    ArrayType1d::ArrayType& pspec = *state.get<ArrayType1d>("powerspectrum")->array;
+    IArrayType::ArrayType& adjust_array = *state.get<IArrayType>("adjust_mode_multiplier")->array;
+    IArrayType::ArrayType& key_array = *state.get<IArrayType>("k_keys")->array;
+    
+    double Eprior = 0;
+    
+    Eprior = codeletGeneralHamiltonian(adjust_array, 
+                                       boost::bind(codeletHamiltonianPrior, _1, _2, _3, _4, boost::cref(pspec), boost::cref(key_array), boost::cref(s_array)));
+    
+    Eprior *= 0.5/volume;
+    comm->all_reduce_t(MPI_IN_PLACE, &Eprior, 1, MPI_SUM);
+    
+    return Eprior;
+}
+
+
+void HMCDensitySampler::computeGradientPsi_Prior(MarkovState& state, CArray& s, CArrayRef& grad_array)
+{
+  IArrayType::ArrayType& adjust_array = *adjust_field->array;
+  IArrayType::ArrayType& key_array = *state.get<IArrayType>("k_keys")->array;
+  ArrayType1d::ArrayType& pspec = *state.get<ArrayType1d>("powerspectrum")->array;
+
+  for (long n = 0 ; n < grad_array.num_elements(); n++)
+    grad_array.data()[n] = 0;
+
+  codeletGeneral(adjust_array,
+                 boost::bind(&codeletGradientPrior, _1, _2, _3, _4, volume, boost::cref(pspec), boost::cref(key_array), boost::cref(s), boost::ref(grad_array)));
+}