smart_montecarlo.h

#pragma once

#include "regions.h"

namespace Faunus::SmarterMonteCarlo {

template <class T>
concept GroupOrParticle = std::is_convertible_v<T, Group> || std::is_convertible_v<T, Particle>;

template <GroupOrParticle T> class Selection
{
  public:
    T* item = nullptr; 
    int n_total = 0;
    int n_inside = 0;
    bool is_inside = false;
    Selection(T& item, int number_total, int number_inside, bool item_is_inside);
};

template <GroupOrParticle T>
Selection<T>::Selection(T& item, int number_total, int number_inside, bool item_is_inside)
    : item(&item)
    , n_total(number_total)
    , n_inside(number_inside)
    , is_inside(item_is_inside)
{
}

enum class BiasDirection
{
    ENTER_REGION,
    EXIT_REGION,
    NO_CROSSING
}; 

double bias(double outside_acceptance, int n_total, int n_inside,
            BiasDirection direction); 

class RegionSampler
{
  private:
    const double outside_acceptance =
        1.0; 
    static BiasDirection getDirection(bool inside_before, bool inside_after);
    template <std::ranges::range Range> double getNumberInside(Range& range) const;

  protected:
    const std::unique_ptr<Region::RegionBase> region; 

  public:
    RegionSampler(double outside_acceptance, std::unique_ptr<Region::RegionBase> region);
    virtual ~RegionSampler() = default;
    void to_json(json& j) const; 

    template <GroupOrParticle T, std::ranges::range Range>
    std::optional<Selection<T>> select(Range& range, Random& random);

    template <GroupOrParticle T> double bias(const Selection<T>& selection);
    std::optional<double> fixed_number_inside; 
};

template <std::ranges::range Range> double RegionSampler::getNumberInside(Range& range) const
{
    if (fixed_number_inside) {
        return fixed_number_inside.value();
    }
    return static_cast<double>(std::ranges::count_if(
        range, [&](const auto& i) { return region->inside(i); })); // expensive
}

template <GroupOrParticle T, std::ranges::range Range>
std::optional<Selection<T>> RegionSampler::select(Range& range, Random& random)
{
    const auto n_total = std::ranges::distance(range.begin(), range.end());
    int max_selection_attempts = 10 * n_total;
    do {
        auto it = random.sample(range.begin(), range.end()); // random particle or group
        if (it == range.end()) {
            return std::nullopt;
        }
        const auto inside = region->inside(*it); // is element inside or outside region?
        if (not inside && outside_acceptance < random()) {
            continue;
        }
        return Selection<T>(*it, n_total, getNumberInside(range), inside);
    } while (max_selection_attempts-- > 0);
    faunus_logger->warn("Max selection attempts reached. Increase outside acceptance (p) ?");
    return std::nullopt;
}

template <GroupOrParticle T> double RegionSampler::bias(const Selection<T>& selection)
{
    const auto is_inside_after = region->inside(*(selection.item)); // may have changed due to move
    const auto direction = getDirection(selection.is_inside, is_inside_after);
    return SmarterMonteCarlo::bias(outside_acceptance, selection.n_total, selection.n_inside,
                                   direction);
}

template <GroupOrParticle T> class MoveSupport
{
  private:
    using OptionalElement =
        std::optional<std::reference_wrapper<T>>; 
    std::optional<SmarterMonteCarlo::Selection<T>>
        selection; 
    SmarterMonteCarlo::RegionSampler region_sampler; 
    Average<double> mean_bias_energy;                
    Average<double> mean_selected_inside;            
    AverageStdev<double> mean_number_inside; 
    AverageStdev<double> blocks_inside;      
    void analyseSelection(int count_inside); 

  public:
    MoveSupport(const Space& spc, const json& j);
    double bias();
    void to_json(json& j) const;
    template <std::ranges::range Range> OptionalElement select(Range& mollist, Random& random);
};

template <GroupOrParticle T>
MoveSupport<T>::MoveSupport(const Space& spc, const json& j)
    : region_sampler(j.at("p").get<double>(), Region::createRegion(spc, j))
{
}

template <GroupOrParticle T> void MoveSupport<T>::analyseSelection(int count_inside)
{
    assert(selection);
    mean_selected_inside += static_cast<double>(selection->is_inside);

    if (region_sampler.fixed_number_inside) {
        return; // skip the rest as we have already fixed n_inside
    }

    mean_number_inside += static_cast<double>(count_inside);
    if (mean_number_inside.size() % 100 == 0) {    // build up a block average
        blocks_inside += mean_number_inside.avg(); // add a new block
        if (blocks_inside.size() % 2 == 0 &&
            blocks_inside.rsd() < 0.04) { // check if block has converged
            region_sampler.fixed_number_inside =
                blocks_inside.avg(); // disable all further sampling!
            faunus_logger->info(
                "Regional <N_inside> = {:.1f} converged and settled after {} iterations",
                region_sampler.fixed_number_inside.value(), mean_number_inside.size());
        }
    }
}

template <GroupOrParticle T> double MoveSupport<T>::bias()
{
    auto bias_energy(0.0);
    if (selection) {
        analyseSelection(selection->n_inside);
        bias_energy = region_sampler.bias(*selection);
    }
    mean_bias_energy += bias_energy;
    return bias_energy;
}

template <GroupOrParticle T> void MoveSupport<T>::to_json(json& j) const
{
    region_sampler.to_json(j);
    if (mean_number_inside) {
        j["mean number inside"] = mean_number_inside.avg();
    }
    if (mean_bias_energy) {
        j["mean bias energy (kT)"] = mean_bias_energy.avg();
    }
    if (mean_selected_inside) {
        j["selected inside fraction"] = mean_selected_inside.avg();
    }
}

template <GroupOrParticle T>
template <std::ranges::range Range>
typename MoveSupport<T>::OptionalElement MoveSupport<T>::select(Range& mollist, Random& random)
{
    selection = region_sampler.select<T>(mollist, random);
    if (selection) {
        return *(selection->item);
    }
    return std::nullopt;
}

} // namespace Faunus::SmarterMonteCarlo