Legendre Polynomials

#include <vector>
#include <cmath>
#include <numbers>
#include <iostream>

// Helper function: linspace
std::vector<double> linspace(double start, double end, size_t num) {
    std::vector<double> result(num);
    double dx = (end - start) / (num - 1);
    for (size_t i = 0; i < num; ++i) {
        result[i] = start + i * dx;
    }
    return result;
}

// Simple Legendre polynomial evaluator
namespace sp {
    struct legendre {
        // Evaluate P_n(x) for single x using recurrence
        double operator()(std::size_t N, double x) const {
            if (N == 0) return 1.0;
            if (N == 1) return x;

            double Pnm2 = 1.0;   // P_0
            double Pnm1 = x;     // P_1
            double Pn = 0.0;

            for (std::size_t n = 2; n <= N; ++n) {
                Pn = ((2*n-1) * x * Pnm1 - (n-1) * Pnm2) / n;
                Pnm2 = Pnm1;
                Pnm1 = Pn;
            }

            return Pn;
        }
    };
}

// Index helper for packed storage
inline std::size_t PT(std::size_t l, std::size_t m) {
    return m + (l * (l + 1)) / 2;
}

// Compute recurrence coefficients A[l,m] and B[l,m]
void computeAB(std::size_t L, std::vector<double>& A, std::vector<double>& B) {
    std::size_t maxIndex = PT(L, L) + 1;
    A.assign(maxIndex, 0.0);
    B.assign(maxIndex, 0.0);

    for (std::size_t l = 2; l <= L; l++) {
        double ls = static_cast<double>(l) * static_cast<double>(l);
        double lm1s = static_cast<double>(l - 1) * static_cast<double>(l - 1);

        for (std::size_t m = 0; m < l - 1; m++) {
            double ms = static_cast<double>(m) * static_cast<double>(m);

            A[PT(l, m)] = std::sqrt((4.0 * ls - 1.0) / (ls - ms));
            B[PT(l, m)] = -std::sqrt((lm1s - ms) / (4.0 * lm1s - 1.0));
        }
    }
}

// Compute all P_lm(cos θ) up to degree L
// Returns vector indexed by PT(l, m)
std::vector<double> computeP(std::size_t L, 
                             const std::vector<double>& A, 
                             const std::vector<double>& B, 
                             double x) {
    std::vector<double> P(PT(L, L) + 1, 0.0);

    double sintheta = std::sqrt(1.0 - x * x);
    double temp = 0.39894228040143267794;   // sqrt(0.5 / pi)
    P[PT(0, 0)] = temp;

    if (L > 0) {
        const double SQRT3 = 1.7320508075688772935;
        P[PT(1, 0)] = x * SQRT3 * temp;

        const double SQRT3DIV2 = -1.2247448713915890491;
        temp = SQRT3DIV2 * sintheta * temp;
        P[PT(1, 1)] = temp;

        for (std::size_t l = 2; l <= L; l++) {
            for (std::size_t m = 0; m < l - 1; m++) {
                P[PT(l, m)] = A[PT(l, m)] *
                    (x * P[PT(l - 1, m)] + B[PT(l, m)] * P[PT(l - 2, m)]);
            }

            P[PT(l, l - 1)] = x * std::sqrt(2.0 * (l - 1) + 3.0) * temp;

            temp = -std::sqrt(1.0 + 0.5 / static_cast<double>(l)) *
                   sintheta * temp;

            P[PT(l, l)] = temp;
        }
    }

    return P;
}

// Example usage
int main() {
    // Compute Legendre polynomials
    sp::legendre leg;
    std::vector<double> xs = linspace(-1.0, 1.0, 1000);
    
    std::vector<double> P1(xs.size());
    std::vector<double> P2(xs.size());
    
    for(size_t i = 0; i < xs.size(); i++) {
        P1[i] = leg(1, xs[i]);  // P_1(x)
        P2[i] = leg(2, xs[i]);  // P_2(x)
    }
    
    std::cout << "P_1(0) = " << P1[500] << std::endl;
    std::cout << "P_2(0) = " << P2[500] << std::endl;
    
    // Compute Associated Legendre polynomials
    std::size_t L = 10;
    std::vector<double> A, B;
    computeAB(L, A, B);
    
    double cos_theta = 0.5;
    std::vector<double> P_lm = computeP(L, A, B, cos_theta);
    
    // Access P_lm using PT(l, m)
    double P_2_1 = P_lm[PT(2, 1)];
    double P_3_2 = P_lm[PT(3, 2)];
    
    std::cout << "P_2^1(0.5) = " << P_2_1 << std::endl;
    std::cout << "P_3^2(0.5) = " << P_3_2 << std::endl;
    
    return 0;
}