scatter.h

#pragma once

#include <fstream>
#include <algorithm>
#include <cmath>
#include "atomdata.h"

namespace Faunus {

namespace Scatter {

enum Algorithm
{
    SIMD,
    EIGEN,
    GENERIC
}; 

template <std::floating_point T = float> class FormFactorSphere
{
  private:
    T j1(T x) const
    { // spherical Bessel function
        T xinv = 1 / x;
        return xinv * (sin(x) * xinv - cos(x));
    }

  public:
    template <class Tparticle> T operator()(T q, const Tparticle& a) const
    {
        assert(q > 0 && a.radius > 0 && "Particle radius and q must be positive");
        T qR = q * a.radius;
        qR = 3. / (qR * qR * qR) * (sin(qR) - qR * cos(qR));
        return qR * qR;
    }
};

template <std::floating_point T = float> struct FormFactorUnity
{
    template <class Tparticle> [[nodiscard]] constexpr T operator()(T, const Tparticle&) const
    {
        return T{1};
    }
};

struct Scatterer
{
    Point pos;
    int id = 0; 
};

template <typename T>
constexpr const auto& getPosition(const T& scatterer)
{
    if constexpr (requires { scatterer.pos; }) {
        return scatterer.pos;
    }
    else {
        return scatterer;
    }
}

template <std::floating_point T = float> struct FormFactorAtomicConstant
{
    template <class Tscatterer> [[nodiscard]] T operator()(T, const Tscatterer& scatterer) const
    {
        if (scatterer.id < 0) {
            return T{1};
        }
        assert(static_cast<size_t>(scatterer.id) < atoms.size());
        return static_cast<T>(atoms[scatterer.id].scattering_f0);
    }
};

#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wdouble-promotion"

template <class Tformfactor, std::floating_point T = float> class DebyeFormula
{
    static constexpr T r_cutoff_infty =
        1e9; //<! a cutoff distance in angstrom considered to be infinity
    T q_mesh_min, q_mesh_max,
        q_mesh_step; //<! q_mesh parameters in inverse angstrom; used for inline lambda-functions

#pragma omp declare simd uniform(this) linear(m : 1)

    inline T q_mesh(int m) { return q_mesh_min + m * q_mesh_step; }

    void init_mesh(T q_min, T q_max, T q_step)
    {
        if (q_step <= 0 || q_min <= 0 || q_max <= 0 || q_min > q_max ||
            q_step / q_max < 4 * std::numeric_limits<T>::epsilon()) {
            throw std::range_error("DebyeFormula: Invalid mesh parameters for q");
        }
        q_mesh_min = q_min < T(1e-6) ? q_step : q_min; // ensure that q > 0
        q_mesh_max = q_max;
        q_mesh_step = q_step;
        try {
            // resolution of the 1D mesh approximation of the scattering vector magnitude q
            const int q_resolution = numeric_cast<int>(1.0 + std::floor((q_max - q_min) / q_step));
            intensity.resize(q_resolution, 0.0);
            sampling.resize(q_resolution, 0.0);
        }
        catch (std::overflow_error& e) {
            throw std::range_error("DebyeFormula: Too many samples");
        }
    }

    T r_cutoff;               
    Tformfactor form_factor;  
    std::vector<T> intensity; 
    std::vector<T> sampling;  

  public:
    DebyeFormula(T q_min, T q_max, T q_step, T r_cutoff)
        : r_cutoff(r_cutoff)
    {
        init_mesh(q_min, q_max, q_step);
    };

    DebyeFormula(T q_min, T q_max, T q_step)
        : DebyeFormula(r_cutoff_infty, q_min, q_max, q_step) {};

    explicit DebyeFormula(const json& j)
        : DebyeFormula(j.at("qmin").get<double>(), j.at("qmax").get<double>(),
                       j.at("dq").get<double>(), j.value("cutoff", r_cutoff_infty)) {};

    template <class Tpvec> void sample(const Tpvec& p, const T weight = 1, const T volume = -1)
    {
        const int N = (int)p.size();         // number of particles
        const int M = (int)intensity.size(); // number of mesh points
        std::vector<T> intensity_sum(M, 0.0);

// Allow parallelization with a hand-written reduction of intensity_sum at the end.
// https://gcc.gnu.org/gcc-9/porting_to.html#ompdatasharing
// #pragma omp parallel default(none) shared(N, M) shared(geo, r_cutoff, p) shared(intensity_sum)
#pragma omp parallel default(shared) shared(intensity_sum)
        {
            std::vector<T> intensity_sum_private(M, 0.0); // a temporal private intensity_sum
#pragma omp for schedule(dynamic)
            for (int i = 0; i < N - 1; ++i) {
                for (int j = i + 1; j < N; ++j) {
                    T r = T(Faunus::Geometry::Sphere::sqdist(
                        getPosition(p[i]), getPosition(p[j]))); // the square root follows
                    if (r < r_cutoff * r_cutoff) {
                        r = std::sqrt(r);
                        // Black magic: The q_mesh function must be inlineable otherwise the loop
                        // cannot be unrolled using advanced SIMD instructions leading to a huge
                        // performance penalty (a factor of 4). The unrolled loop uses a different
                        // sin implementation, which may be spotted when profiling.
                        // TODO: Optimize also for other compilers than GCC by using a vector math
                        // library, e.g.,
                        // TODO: https://github.com/vectorclass/version2
                        // #pragma GCC unroll 16 // for diagnostics, GCC issues warning when cannot
                        // unroll
                        for (int m = 0; m < M; ++m) {
                            const T q = q_mesh(m);
                            intensity_sum_private[m] += form_factor(q, p[i]) *
                                                        form_factor(q, p[j]) * std::sin(q * r) /
                                                        (q * r);
                        }
                    }
                }
            }
// reduce intensity_sum_private into intensity_sum
#pragma omp critical
            std::transform(intensity_sum.begin(), intensity_sum.end(),
                           intensity_sum_private.begin(), intensity_sum.begin(), std::plus<T>());
        }

// https://gcc.gnu.org/gcc-9/porting_to.html#ompdatasharing
// #pragma omp parallel for default(none) shared(N, M, weight, volume) shared(p, r_cutoff,
// intensity_sum) shared(sampling, intensity)
#pragma omp parallel for shared(sampling, intensity)
        for (int m = 0; m < M; ++m) {
            const T q = q_mesh(m);
            T intensity_self_sum = 0;
            for (int i = 0; i < N; ++i) {
                intensity_self_sum += std::pow(form_factor(q, p[i]), 2);
            }
            T intensity_corr = 0;
            if (r_cutoff < r_cutoff_infty && volume > 0) {
                intensity_corr = 4 * pc::pi * N / (volume * std::pow(q, 3)) *
                                 (q * r_cutoff * std::cos(q * r_cutoff) - std::sin(q * r_cutoff));
            }
            sampling[m] += weight;
            if (intensity_self_sum != T{0}) {
                intensity[m] +=
                    ((2 * intensity_sum[m] + intensity_self_sum) / intensity_self_sum +
                     intensity_corr) *
                    weight;
            }
        }
    }

    auto getQMeshParameters() { return std::make_tuple(q_mesh_min, q_mesh_max, q_mesh_step); }

    auto getIntensity()
    {
        std::map<T, T> averaged_intensity;
        for (size_t m = 0; m < intensity.size(); ++m) {
            const T average = intensity[m] / (sampling[m] != T(0.0) ? sampling[m] : T(1.0));
            averaged_intensity.emplace(q_mesh(m), average);
        }
        return averaged_intensity;
    }
};

#pragma GCC diagnostic pop

template <std::floating_point T> class SamplingPolicy
{
  public:
    struct sampled_value
    {
        T value;
        T weight;
    };

    typedef std::map<T, sampled_value> TSampledValueMap;

  private:
    TSampledValueMap samples;
    const T precision = 10000.0; 

  public:
    std::map<T, T> getSampling() const
    {
        std::map<T, T> average;
        for (auto [key, sample] : samples) {
            average.emplace(key, sample.value / sample.weight);
        }
        return average;
    }

    void addSampling(T key_approx, T value, T weight = 1.0)
    {
        const T key =
            std::round(key_approx * precision) / precision; // round |q| for better binning
        samples[key].value += value * weight;
        samples[key].weight += weight;
    }
};

template <std::floating_point T>
std::map<T, T> averageByMagnitude(const std::vector<std::pair<T, T>>& pairs,
                                  T precision = T{10000})
{
    struct Accumulator
    {
        T sum = T{0};
        int count = 0;
    };
    std::map<T, Accumulator> bins;
    for (const auto& [key, value] : pairs) {
        const T rounded = std::round(key * precision) / precision;
        bins[rounded].sum += value;
        bins[rounded].count++;
    }
    std::map<T, T> result;
    for (const auto& [key, acc] : bins) {
        result[key] = acc.sum / static_cast<T>(acc.count);
    }
    return result;
}

template <class Tformfactor = FormFactorUnity<double>, typename T = double, Algorithm method = SIMD,
          typename TSamplingPolicy = SamplingPolicy<T>>
class StructureFactorPBC : private TSamplingPolicy
{
    const std::vector<Point> directions = {
        {1, 0, 0}, {0, 1, 0},  {0, 0, 1},                                      // 3 permutations
        {1, 1, 0}, {0, 1, 1},  {1, 0, 1},  {-1, 1, 0}, {-1, 0, 1}, {0, -1, 1}, // 6 permutations
        {1, 1, 1}, {-1, 1, 1}, {1, -1, 1}, {1, 1, -1}                          // 4 permutations
    };

    const int p_max; 
    Tformfactor form_factor; 
    using TSamplingPolicy::addSampling;

  public:
    StructureFactorPBC(int q_multiplier)
        : p_max(q_multiplier)
    {
    }

    template <typename Tscatterers>
    void sample(const Tscatterers& scatterers, const Point& boxlength)
    {
        const auto n = directions.size() * static_cast<size_t>(p_max);
        std::vector<std::pair<T, T>> q_intensity(n);

#pragma omp parallel for collapse(2) default(shared)
        for (size_t i = 0; i < directions.size(); ++i) {
            for (int p = 1; p <= p_max; ++p) {
                const Point q =
                    2.0 * pc::pi * p * directions[i].cwiseQuotient(boxlength);
                q_intensity[i * static_cast<size_t>(p_max) + static_cast<size_t>(p - 1)] =
                    {static_cast<T>(q.norm()), calculateIntensity(scatterers, q)};
            }
        }

        for (const auto& [q, intensity] : averageByMagnitude(q_intensity)) {
            addSampling(q, intensity);
        }
    }

    template <typename Tscatterers>
    T calculateIntensity(const Tscatterers& scatterers, const Point& q) const
    {
        T sum_cos = 0.0;
        T sum_sin = 0.0;
        T sum_f_squared = 0.0;
        const auto q_norm = q.norm();
        for (const auto& scatterer : scatterers) {
            const auto& pos = getPosition(scatterer);
            const auto f = form_factor(q_norm, scatterer);
            const auto qr = static_cast<T>(q.dot(pos));
            sum_cos += f * cos(qr);
            sum_sin += f * sin(qr);
            sum_f_squared += f * f;
        }
        if (sum_f_squared == T{0}) {
            return T{0};
        }
        return std::norm(std::complex<T>(sum_cos, sum_sin)) / sum_f_squared;
    }

    int getQMultiplier() { return p_max; }

    using TSamplingPolicy::getSampling;
};

template <class Tformfactor = FormFactorUnity<float>, std::floating_point T = float,
          typename TSamplingPolicy = SamplingPolicy<T>>
class StructureFactorIPBC : private TSamplingPolicy
{
    std::vector<Point> directions = {{1, 0, 0}, {1, 1, 0}, {1, 1, 1}};

    int p_max; 
    Tformfactor form_factor; 
    using TSamplingPolicy::addSampling;

  public:
    explicit StructureFactorIPBC(int q_multiplier)
        : p_max(q_multiplier)
    {
    }

    template <typename Tscatterers>
    void sample(const Tscatterers& scatterers, const Point& boxlength)
    {
        const auto n = directions.size() * static_cast<size_t>(p_max);
        std::vector<std::pair<T, T>> q_intensity(n);

// https://gcc.gnu.org/gcc-9/porting_to.html#ompdatasharing
// #pragma omp parallel for collapse(2) default(none) shared(directions, p_max, scatterers,
// boxlength)
#pragma omp parallel for collapse(2) default(shared)
        for (size_t i = 0; i < directions.size(); ++i) {
            for (int p = 1; p <= p_max; ++p) { // loop over multiples of q
                const Point q =
                    2.0 * pc::pi * p * directions[i].cwiseQuotient(boxlength); // scattering vector
                T sum_f_cos = 0;
                T sum_f_squared = 0;
                const auto q_norm = q.norm();
                for (const auto& scatterer : scatterers) {
                    const auto& r = getPosition(scatterer);
                    const auto f = form_factor(q_norm, scatterer);
                    // if q[i] == 0 then its cosine == 1 hence we can avoid cosine computation for
                    // performance reasons
                    T product = std::cos(T(q[0] * r[0]));
                    if (q[1] != 0)
                        product *= std::cos(T(q[1] * r[1]));
                    if (q[2] != 0)
                        product *= std::cos(T(q[2] * r[2]));
                    sum_f_cos += f * product;
                    sum_f_squared += f * f;
                }
                const T ipbc_factor =
                    std::pow(2, directions[i].count()); // 2 ^ number of non-zero elements
                T intensity = T{0};
                if (sum_f_squared != T{0}) {
                    intensity = (sum_f_cos * sum_f_cos) / sum_f_squared * ipbc_factor;
                }
                q_intensity[i * static_cast<size_t>(p_max) + static_cast<size_t>(p - 1)] =
                    {static_cast<T>(q_norm), intensity};
            }
        }

        for (const auto& [q, intensity] : averageByMagnitude(q_intensity)) {
            addSampling(q, intensity);
        }
    }

    int getQMultiplier() { return p_max; }

    using TSamplingPolicy::getSampling;
};

} // namespace Scatter
} // namespace Faunus