maedoc

using FFTW
using DiffEqSensitivity, OrdinaryDiffEq, Zygote

nlon = 64
nlat = 64
K = (1:(div(nlat,2)+1)) .* (1:nlon)' / (nlat*nlon)

function f(x,p,t)
    k = p[1]
    lc = k * irfft(K .* rfft(x), nlat)

maedoc

import numpy as np
from scipy.special import sph_harm

class ShtDiff:

    def __init__(self, lmax=16, nlat=64, D=-1e-2, nlon=None):
        self.lmax = lmax
        self.nlat = nlat
        self.D = D
        self.nlon = nlon or (2 * nlat)

maedoc

import numpy as np

# Compute 2N-point real FFT w/ a N-point complex FFT, cf. TI SPRA 291

# random 2N-point real sequence
g = np.random.randn(16)

# size of sequence
n = len(g)//2
n2 = 2 * n

maedoc

-- was not clear to me from the user guide how this is supposed to work
-- when in fact it's straightforward

-- put the needed arithmetic ops in an "interface" / "concept"
module type arith = {
    type t
    val i32 : i32 -> t
    val * : t -> t -> t
}

maedoc

-- we want to verify we can do the irregular structure of the sht
-- when mapped on the GPU.

-- if we can't, we have (lmax+1)*lmax/2 wasted space, as follows
let lt1 [nlat] (lmax:i64) (m:i64) (fm:[nlat]f32): [lmax]f32 =
    let q = tabulate lmax (\_ -> 0f32)
    let (q, _) =
    loop (q, x) = (q, 0f32) for i < (lmax - m) do
        let m' = m + i
        let q[m'] = reduce (+) x fm |> (*x)

maedoc

import os, ctypes
import numpy as np
from sympy import exp, pi, I, Symbol, re, im, cse
from sympy.printing.c import ccode

# DIF FT recursion
def ft(x):
    N = len(x)
    if N == 2:
        a, b = x

maedoc

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)

maedoc

-- this variant uses functional approach, appending new states to buffer
-- instead of in-place update

let pi = 3.141592653589793f32

-- some type abbreviations
type mpr_pars = {G: f32, I: f32, Delta: f32, eta: f32, tau: f32, J: f32}
type mpr_node = (f32, f32)
type mpr_net [n] = [n] mpr_node
-- this is tranposed from mpr-pdq to avoid tranposes in history update

maedoc

mini-example of a TF op with a custom C++ resource

There's a TF tutorial on creating a op, which covers anything written from scratch, but when you've got existing C++ classes (or anything callable from C++, really) which do the heavy lifting and are stateful, or at least require intialization/instatiation, the tutorial recommends, as a side note that ops need to be reentrant, that a ResourceMgr be used, linking to the corresponding header and no example.

That's sort of a cliff-hanger given how nice the rest of the tutorial. Two related SO questions

maedoc