potentials_base.h

#pragma once
#include "core.h"
#include "particle.h"
#include "aux/pairmatrix.h"
#include <array>

namespace Faunus::pairpotential {

using TPairMatrix = Eigen::MatrixXd;
using TPairMatrixPtr = std::shared_ptr<TPairMatrix>;

using TExtractorFunc = std::function<double(const InteractionData&)>;

using TCombinatorFunc = std::function<double(double, double)>;

using TModifierFunc = std::function<double(double)>;

struct CustomInteractionData
{
    std::array<AtomData::index_type, 2> atom_id;
    InteractionData interaction;
};

void from_json(const json& j, CustomInteractionData& interaction);
void to_json(json& j, const CustomInteractionData& interaction);
void from_json(const json& j, std::vector<CustomInteractionData>& interactions);

enum class CombinationRuleType
{
    UNDEFINED,
    ARITHMETIC,
    GEOMETRIC,
    LORENTZ_BERTHELOT
};
NLOHMANN_JSON_SERIALIZE_ENUM(CombinationRuleType,
                             {{CombinationRuleType::UNDEFINED, "undefined"},
                              {CombinationRuleType::ARITHMETIC, "arithmetic"},
                              {CombinationRuleType::GEOMETRIC, "geometric"},
                              {CombinationRuleType::LORENTZ_BERTHELOT, "lorentz_berthelot"},
                              {CombinationRuleType::LORENTZ_BERTHELOT, "LB"}})

struct PairPotentialException : public std::runtime_error
{
    explicit PairPotentialException(const std::string& msg);
};

class PairMixer
{
    TExtractorFunc
        extractor; 
    TCombinatorFunc combinator; 
    TModifierFunc
        modifier; 

  public:
    PairMixer(TExtractorFunc extractor, TCombinatorFunc combinator,
              TModifierFunc modifier = &modIdentity);

    TPairMatrixPtr createPairMatrix(const std::vector<AtomData>& atoms);
    TPairMatrixPtr createPairMatrix(const std::vector<AtomData>& atoms,
                                    const std::vector<CustomInteractionData>& interactions);

    enum class CoefficientType
    {
        ANY,
        SIGMA,
        EPSILON
    };
    static TCombinatorFunc getCombinator(CombinationRuleType combination_rule,
                                         CoefficientType coefficient = CoefficientType::ANY);

    // when explicit custom pairs are the only option
    inline static constexpr double combUndefined(double = 0.0, double = 0.0)
    {
        return std::numeric_limits<double>::signaling_NaN();
    };

    inline static double combArithmetic(double a, double b) { return 0.5 * (a + b); }

    inline static double combGeometric(double a, double b) { return std::sqrt(a * b); }

    inline static double modIdentity(double x) { return x; }

    inline static double modSquared(double x) { return x * x; }
};

class PairPotential
{
  private:
    friend void from_json(const json&, PairPotential&);

  public:
    std::string
        name; 
    std::string cite;      
    bool isotropic = true; 
    std::function<double(const Particle&)> selfEnergy = nullptr; 
    virtual void to_json(json&) const = 0;
    virtual void from_json(const json&) = 0;
    virtual ~PairPotential() = default;

    [[nodiscard]] virtual Point force(const Particle& a, const Particle& b, double squared_distance,
                                      const Point& b_towards_a) const;

    virtual double operator()(const Particle& particle_a, const Particle& particle_b,
                              double squared_distance, const Point& b_towards_a) const = 0;

  protected:
    explicit PairPotential(std::string name = std::string(), std::string cite = std::string(),
                           bool isotropic = true);
};

void to_json(json& j, const PairPotential& base);   
void from_json(const json& j, PairPotential& base); 

template <class T>
concept RequirePairPotential = std::derived_from<T, pairpotential::PairPotential>;

template <RequirePairPotential T> auto makePairPotential(const json& j)
{
    T pair_potential;
    pairpotential::from_json(j, pair_potential);
    return pair_potential;
}

class MixerPairPotentialBase : public PairPotential
{
  protected:
    CombinationRuleType combination_rule;
    std::shared_ptr<std::vector<CustomInteractionData>> custom_pairs =
        std::make_shared<std::vector<CustomInteractionData>>();
    json json_extra_params; 
    void init(); 
    virtual void
    initPairMatrices() = 0; 
    virtual void extractorsFromJson(
        const json&); 
  public:
    explicit MixerPairPotentialBase(
        const std::string& name = std::string(), const std::string& cite = std::string(),
        CombinationRuleType combination_rule = CombinationRuleType::UNDEFINED,
        bool isotropic = true);
    ~MixerPairPotentialBase() override = default;
    void from_json(const json& j) override;
    void to_json(json& j) const override;
};

template <RequirePairPotential T1, RequirePairPotential T2>
struct CombinedPairPotential : public PairPotential
{
    T1 first;  
    T2 second; 
    explicit CombinedPairPotential(const std::string& name = "")
        : PairPotential(name) {};

    inline double operator()(const Particle& particle_a, const Particle& particle_b,
                             const double squared_distance,
                             const Point& b_towards_a = {0, 0, 0}) const override
    {
        return first(particle_a, particle_b, squared_distance, b_towards_a) +
               second(particle_a, particle_b, squared_distance, b_towards_a);
    } 

    [[nodiscard]] inline Point force(const Particle& particle_a, const Particle& particle_b,
                                     const double squared_distance,
                                     const Point& b_towards_a) const override
    {
        return first.force(particle_a, particle_b, squared_distance, b_towards_a) +
               second.force(particle_a, particle_b, squared_distance, b_towards_a);
    } 

    void from_json(const json& j) override
    {
        Faunus::pairpotential::from_json(j, first);
        Faunus::pairpotential::from_json(j, second);
        name = first.name + "/" + second.name;
        if (first.selfEnergy or second.selfEnergy) { // combine self-energies
            selfEnergy = [u1 = first.selfEnergy, u2 = second.selfEnergy](const Particle& p) {
                if (u1 and u2) {
                    return u1(p) + u2(p);
                }
                if (u1) {
                    return u1(p);
                }
                return u2(p);
            };
        }
        else {
            selfEnergy = nullptr;
        }
    }

    void to_json(json& j) const override
    {
        assert(j.is_object());
        auto& _j = j["default"] = json::array();
        _j.push_back(first);
        _j.push_back(second);
    }
};

} // namespace Faunus::pairpotential