Associated Legendre polynomials via recursion relations

pmn.fut

let len_q (N:i64): i64 = (N + 1) * (N + 2) // 2

let gen_ml (N:i64): [](i64,i64) =
    loop ml = [(0,0)] for i < ((len_q N)-1) do
        let (m,l) = ml[i]
        let nl = if N == l then m + 1 else l + 1
        let nm = if N == l then m + 1 else m
        in ml ++ [(nm,nl)]

entry gen_ml_arr (N:i64): [][]i64 =
    let (m,l) = gen_ml N |> unzip in [m,l]

entry all_amm (N:i64): []f32 = iota N 
    |> map (\i -> f32.i64 (i + 1))
    |> map (\k -> (2*k+1)/(2*k))
    |> ([1f32]++)
    |> scan (*) 1f32
    |> map (\e -> f32.sqrt(e/(4*f32.pi)))

entry amn (m:f32) (n:f32): f32 = f32.sqrt((4*n*n - 1)/(n*n - m*m))

entry bmn (m:f32) (n:f32): f32 =
    let l = (2*n + 1)/(2*n - 3)
    let r = ((n - 1)*(n - 1) - m*m)/(n*n - m*m)
    in (-f32.sqrt(l * r))

entry lat_grid (nlat:i64): []f32 = iota nlat
    |> map f32.i64 
    |> map (\x -> f32.cos (x/(f32.i64 nlat)*f32.pi))


-- traverse Pmn for given m and lmax
let Pmn1 (m:i64) (n:i64) (amm:f32) (cx:f32): f32 =
    let m' = f32.i64 m
    -- P^m_m
    let p0 = amm*(1 - cx*cx)**(m'/2)*(-1)**m'
    -- P^m_(m + 1)
    let p1 = (amn m' (m' + 1))*cx*p0
    -- P^m_n -> P^m_n+1 -> P^m_n+2
    let p2 n p1 p0 = (amn m' n)*cx*p1 + (bmn m' n)*p0
    -- P^m_n
    let (pn, _) = match (n-m)
        case 0 -> (p0, 0f32)
        case 1 -> (p1, p0)
        case _ -> loop (p1,p0) 
            for i < (n-m-1) do (p2 (m'+2+f32.i64 i) p1 p0, p1)
    in pn

entry Pmn [nlat] (m:i64) (n:i64) (amm:f32) (cx:[nlat]f32): [nlat]f32 =
    map (Pmn1 m n amm) cx

-- == 
-- input { 50i64 128i64 0 }
-- input { 50i64 128i64 1 }
-- input { 50i64 128i64 2 }
entry main (lmax:i64) (nlat:i64) (grad:i32) =
    let amm = all_amm lmax
    let cx = lat_grid nlat
    let f x = map (\m -> Pmn m lmax amm[m] x) (iota lmax)
    let gr = f cx
    let gf = cx
    in
    match grad
        case 0 -> (f cx)[0]
        case 1 -> vjp f cx gr
        case 2 -> (jvp f cx gf)[0]

Pmn.py

import numpy as np
from scipy.special import sph_harm, lpmv
import shtns
import os
os.system('futhark c --library pmn.fut')
os.system('build_futhark_ffi pmn')
from futhark_ffi import Futhark
import _pmn
pmn = Futhark(_pmn)


# cf https://arxiv.org/abs/1202.6522

def gen_ml(N):
    return np.array([[m, l] for m in range(N+1) for l in range(N+1) if m<=l])

def all_amm(N):
    amm = []
    for m in range(N+1):
        els = []
        for k in range(1,m+1):
            els.append((2*k+1)/(2*k))
        el = np.prod(els)/(4*np.pi)
        amm.append(np.sqrt(el))
    amm = np.array(amm)
    return amm

def amn(m, n):
    nom = 4*n*n - 1
    den = n*n - m*m
    return np.sqrt(nom/den)

def bmn(m, n):
    l = (2*n + 1)/(2*n - 3)
    r = ((n - 1)*(n - 1) - m*m)/(n*n - m*m)
    return - np.sqrt(l*r)

def lat_grid(nlat):
    x = np.r_[:nlat] / nlat * np.pi
    cx = np.cos(x)
    return cx

check_Pmn = lambda p, m, n, cx: \
    np.testing.assert_allclose(p, 
        sph_harm(m, n, 0, np.arccos(cx)).real, rtol=1e-5, atol=1e-6)

# TODO separate algo from checks?
def test_Pmn_recursion(N=10, nlat=32):
    "check all generated P^m_n(x) coefficients"
    amm = all_amm(N)
    cx = lat_grid(nlat)
    ml = set([(m_, l_) for (m_, l_) in gen_ml(N)])

    for m in range(N+1):

        n = m
        ml.remove((m,n))
        p0 = amm[m]*(1 - cx*cx)**(m/2)*(-1)**m # eq 13
        check_Pmn(p0, m, m, cx)
        if n == N:
            break

        n += 1
        ml.remove((m,n))
        p1 = amn(m, n) * cx * p0 # eq 14
        check_Pmn(p1, m, n, cx)
        if n == N:
            continue

        n += 1
        ml.remove((m,n))
        p2 = amn(m, n)*cx*p1 + bmn(m, n)*p0 # eq 15
        check_Pmn(p2, m, n, cx)
        if n == N:
            continue

        while n < N:
            p0, p1 = p1, p2
            n += 1
            ml.remove((m,n))
            p2 = amn(m, n)*cx*p1 + bmn(m, n)*p0 # eq 15
            check_Pmn(p2, m, n, cx)

    assert(len(ml) == 0)

def test_ml():
    for N in range(8,32):
        sht = shtns.sht(N)
        # sht.set_grid(nlat, nlat*2)
        ml2 = gen_ml(N)
        ml3 = np.c_[sht.m, sht.l]
        np.testing.assert_allclose(ml2, ml3)
        mf, lf = pmn.from_futhark(pmn.gen_ml(N))
        np.testing.assert_allclose(ml2, np.c_[mf, lf])       

def test_amm():
    N = 20
    np_amm = all_amm(N)
    ft_amm = pmn.from_futhark(pmn.all_amm(N))
    np.testing.assert_allclose(ft_amm, np_amm, rtol=1e-6, atol=1e-7)

def test_amn():
    for m, l in gen_ml(10):
        np.testing.assert_allclose(
            pmn.amn(m,l),
            amn(m,l))

def test_bmn():
    for m, l in gen_ml(10):
        np.testing.assert_allclose(
            pmn.bmn(m,l),
            bmn(m,l))

def test_lat_grid():
    nlat = 64
    np.testing.assert_allclose(
        pmn.from_futhark(pmn.lat_grid(nlat)) + 1,
        lat_grid(nlat) + 1,
        rtol=1e-6, atol=1e-6
    )

def test_Pm_rec_ft(N=10, nlat=32):
    amm = all_amm(N)
    cx = lat_grid(nlat)
    for m, l in gen_ml(N):
        p = pmn.from_futhark(pmn.Pmn(m, l, amm[m], cx))
        check_Pmn(p, m, l, cx)

if __name__ == '__main__':
    # test_ml()
    test_amm()
    test_amn()
    test_bmn()
    test_lat_grid()
    test_Pmn_recursion()
    test_Pm_rec_ft()
    print('all done')