AiryWavesSimd.h

#pragma once
 
#include <marenv/marenv.h>
#include <marenv/wave/LinearWaveTheory.h>
#include <marenv/wave/WaveEnergySpectrum.h>
 
#include <array>
 
namespace marenv::wave
{
 
class AiryWavesSimd : public LinearWaveTheory
{
public:
    struct WaveStaticSIMD
    {
        std::vector<xsimd::batch<double>> omega; 
        std::vector<xsimd::batch<double>> k; 
        std::vector<xsimd::batch<double>> cosT; 
        std::vector<xsimd::batch<double>> sinT; 
        std::vector<xsimd::batch<double>> eps; 
        std::vector<xsimd::batch<double>> zetaA; 
        // Pre-multiplied products (avoid recomputing each Compute* call)
        std::vector<xsimd::batch<double>> zetaAK; 
        std::vector<xsimd::batch<double>> zetaAKcosT; 
        std::vector<xsimd::batch<double>> zetaAKsinT; 
        std::vector<xsimd::batch<double>> ampVelK; 
        std::vector<xsimd::batch<double>> ampVelKcosT; 
        std::vector<xsimd::batch<double>> ampVelKsinT; 
        std::vector<xsimd::batch<double>> ampAccK; 
        std::vector<xsimd::batch<double>> ampAccKcosT; 
        std::vector<xsimd::batch<double>> ampAccKsinT; 
        std::vector<xsimd::batch<double>> totalAmplitude; 
    };
 
    struct SimdSnapshot : public WaveSnapshot
    {
        WaveStaticSIMD simdStatic; 
 
        explicit SimdSnapshot(std::shared_ptr<WaveComponents> wc, WaveStaticSIMD s)
            : WaveSnapshot(std::move(wc))
            , simdStatic(std::move(s))
        { }
    };
 
    struct WaveDynamicSIMD
    {
        std::vector<xsimd::batch<double>> angle; 
        std::vector<xsimd::batch<double>> sin_angle; 
        std::vector<xsimd::batch<double>> cos_angle; 
    };
 
    struct WaveDynamicSIMDExpAmp : public WaveDynamicSIMD
    {
        std::vector<xsimd::batch<double>> expk; 
    };
 
    struct WaveDynamicStack
    {
        static constexpr int kMaxWaves   = 52; 
        static constexpr int kBatchSize  = static_cast<int>(xsimd::batch<double>::size); 
        static constexpr int kMaxBatches = (kMaxWaves + kBatchSize - 1) / kBatchSize; 
 
        std::array<xsimd::batch<double>, kMaxBatches> angle; 
        std::array<xsimd::batch<double>, kMaxBatches> sin_angle; 
        std::array<xsimd::batch<double>, kMaxBatches> cos_angle; 
        std::array<xsimd::batch<double>, kMaxBatches> expk; 
        int nbatches = 0; 
    };
 
    explicit AiryWavesSimd(std::shared_ptr<WaveComponents> waveComponents);
    AiryWavesSimd(WaveEnergySpectrum::SpectrumParameters spectrumParameters, int numWaves, uint32_t randomSeed);
 
    ~AiryWavesSimd() override = default;
 
    bool PositionIsInWater(double time, const double pos[3]);
 
    Status ComputeSurfaceElevation(double time, const double pos[2], double& elevationOut) override;
 
    Status ComputeSurfaceUnitNormal(double time, const double pos[2], double unitNormalOut[3]) override;
 
    Status ComputeParticleVelocity(double time, const double pos[3], double velOut[3]) override;
 
    Status ComputeParticleAcceleration(double time, const double pos[3], double accOut[3]) override;
 
    Status ComputeDynamicPressure(double time, const double pos[3], double& pressureOut) override;
 
    Status ComputeTotalPressure(double time, const double pos[3], double& pressureOut) override;
 
    Status ComputePointQuery(
        double time,
        const double pos[3],
        double& surface_elevation,
        double surface_unit_normal[3],
        double particle_velocity[3],
        double particle_acceleration[3],
        double& pressure) override;
 
    void BlendInPlace(WaveComponentData& other, double alpha) override;
 
    [[nodiscard]] const WaveStaticSIMD& GetSimdStatic() const
    {
        auto snap = LoadSnapshot();
        return static_cast<SimdSnapshot*>(snap.get())->simdStatic;
    }
 
private:
    void PublishSimdSnapshot(std::shared_ptr<WaveComponents> wc);
};
} // namespace marenv::wave