Skip to content

Instantly share code, notes, and snippets.

@maedoc

maedoc/SHTMatrix2.ipynb

Last active May 1, 2021 20:34
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save maedoc/4243b83196a5cfec9ac826a1647c5bcd to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save maedoc/4243b83196a5cfec9ac826a1647c5bcd to your computer and use it in GitHub Desktop.
Download ZIP
Prototyping spherical harmonic transforms
Display the source blob
Display the rendered blob
Raw
SHTfs.ipynb
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Display the source blob
Display the rendered blob
Raw
SHTMatrix2.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"from scipy.special import sph_harm\n",
"import shtns"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Just for reference let's use a SHTns grid:"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"nlat, nlon = 128,64\n",
"\n",
"sht = shtns.sht(50)\n",
"sht.set_grid(nlon,nlat)\n",
"x = np.random.randn(nlon, nlat) # phi=azi x nlat\n",
"q = sht.analys(x)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"`theta` and `phi` are the spatial grid points"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [],
"source": [
"theta = np.r_[0:2*np.pi:(nlon+1)*1j][:-1]\n",
"phi = np.r_[0:np.pi:(nlat+1)*1j][:-1]\n",
"dt, dp = theta[1] - theta[0], phi[1] - phi[0]\n",
"theta += dt/2\n",
"phi += dp/2"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Method 1**: Compute the whole of coefficients $Y_{lm}(\\theta,\\phi)$ and $\\bar{Y}_{lm}(\\theta,\\phi) \\sin \\phi$ (where does $\\sin \\phi$ come from?) at once: "
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"the expected size of the array"
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(66, 'MB')"
]
},
"execution_count": 4,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"(sht.m.size * theta.size * phi.size * 16) >> 20, 'MB'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Compute the arrays with the SciPy `sph_harm` function,"
]
},
{
"cell_type": "code",
"execution_count": 5,
"metadata": {},
"outputs": [],
"source": [
"Y = sph_harm(\n",
" sht.m[:,None,None],sht.l[:,None,None], \n",
" theta[:,None], phi)\n",
"Yb = np.conj(Y) * np.sin(phi)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"define the synth and analys operators:"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"array([1.09780178e-14, 2.20362188e-04, 2.66357471e-03, 1.99161754e-02,\n",
" 8.37742400e-02])"
]
},
"execution_count": 6,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"def synth(q):\n",
" return (q[:,None,None] * Y).sum(axis=0)\n",
"\n",
"def analys(x):\n",
" return ((x * Yb).sum(axis=1)*dt).sum(axis=1)*dp\n",
"\n",
"re = np.abs(q - analys(synth(q))) / np.abs(q)\n",
"np.percentile(re, [50, 90, 95, 99, 100])"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"check the orthonormality for a small block"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"[[ 1. -0. 0. -0. 0.]\n",
" [ 0. 1. -0. 0. -0.]\n",
" [ 0. -0. 1. -0. 0.]\n",
" [-0. 0. -0. 1. -0.]\n",
" [ 0. -0. 0. 0. 1.]]\n",
"[[0. 0. 0. 0. 0.]\n",
" [0. 0. 0. 0. 0.]\n",
" [0. 0. 0. 0. 0.]\n",
" [0. 0. 0. 0. 0.]\n",
" [0. 0. 0. 0. 0.]]\n"
]
}
],
"source": [
"yy = lambda i,j: ((Y[i] * Yb[j]).sum(axis=1)*dt).sum(axis=-1)*dp\n",
"yy5 = np.array([[yy(i,j) for j in range(5)] for i in range(5)])\n",
"with np.printoptions(precision=2, suppress=True):\n",
" print(yy5.real)\n",
" print(yy5.imag)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"and the maximum of the diagonal"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1.00805918250734"
]
},
"execution_count": 8,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"np.array([yy(i,i) for i in range(sht.l.size)]).real.max()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"so this is not perfect."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"and time it"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"35.7 ms ± 144 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
]
}
],
"source": [
"%timeit analys(synth(q))"
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"23 ms ± 102 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
]
}
],
"source": [
"qs = q.astype(np.complex64)\n",
"Ys = Y.astype(np.complex64)\n",
"Ybs = Yb.astype(np.complex64)\n",
"\n",
"def synths(q):\n",
" return (q[:,None,None] * Ys).sum(axis=0)\n",
"\n",
"def analyss(x):\n",
" return ((x * Ybs).sum(axis=1)*dt).sum(axis=1)*dp\n",
"\n",
"%timeit analyss(synths(qs))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This likely just reflects the memory traffic, as the arithmetic is trivial. If we think about memory layout it could be done faster."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"30.8 µs ± 917 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
]
}
],
"source": [
"%timeit sht.analys(sht.synth(q))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Method 2**: factor matrix into Fourier and Legendre steps separately."
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"# synth (q -> x)\n",
"F = np.exp(1j * sht.m[:,None] * theta)\n",
"L = sph_harm(sht.m[:,None], sht.l[:,None], 0, phi)\n",
"# analys (x -> q)\n",
"Lb = np.conj(L) * np.sin(phi) * dp\n",
"Fb = np.conj(F) * dt"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(False, True)"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"np.any(L.imag), np.all(np.conj(L)==L)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"So we can see that the Legendre transform is all real-valued. The adjoint can be constructed then applying the other transform to appropriately rescaled gradients: adjoint of `Lb` computed with `L` but prescaling gradients by `sin(phi)*dp`, and adjoint of of L computed with `Lb` but postscaling VJP by `1/sin(phi)/dp`.\n",
"\n",
"The Fourier transform is the complex-valued one, which, as a quick guess, requires just a conjugate, but this needs to be checked."
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(0.05096333459619811-0.031200423682453272j)"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"F[128,16] * q[128]"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(0.05096333459619811-0.031200423682453272j)"
]
},
"execution_count": 15,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"np.conj(Fb[128,16] * np.conj(q[128] / dt))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"so it seems the adjoint of the synthesis can done this way, and the analysis,"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(0.002154798162808452-0.001990168003692288j)"
]
},
"execution_count": 16,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"Fb[128,16] * q[128]"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(0.0021547981628084524-0.001990168003692288j)"
]
},
"execution_count": 17,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"np.conj(F[128,16] * np.conj(q[128]) * dt)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"we can see the memory size for these matrices is lower than the previous method,"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(132, 'MB', 3, 'MB')"
]
},
"execution_count": 18,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"(Y.nbytes+Yb.nbytes)>>20, 'MB', sum([_.nbytes for _ in (F,L,Lb,Fb)])>>20, 'MB'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"and that when the separate arrays are multiplied, they match the full ones,"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"(True, False)"
]
},
"execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"np.allclose(L[:,None,:] * F[:,:,None], Y), np.allclose(Lb[:,None,:] * Fb[:,:,None], Yb)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now we can define the operators"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"True"
]
},
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"synth2 = lambda L, F, q: np.dot(F.T * q, L)\n",
"FT = F.T.copy()\n",
"synth2t = lambda L, FT, q: np.dot(FT * q, L)\n",
"analys2 = lambda Lb, Fb, x: np.sum(Lb * np.dot(Fb,x), axis=1)\n",
"\n",
"# based on above cells, to be checked numerically\n",
"adj_synth2t = lambda Lb, Fb, g: np.conj(analys2(Lb, Fb, np.conj(g/dt/np.sin(phi)/dp)))\n",
"adj_analys2 = lambda L, FT, g: np.conj(synth2t(L, FT, np.conj(g))*dt*np.sin(phi)*dp)\n",
"\n",
"np.allclose(synth2t(L, FT, q), synth(q))\n",
"np.allclose(analys2(Lb, Fb, x), analys(x))"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"8.19 ms ± 108 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"%timeit analys2(Lb,Fb,synth2t(L,FT,q))"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"8.36 ms ± 16.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
]
}
],
"source": [
"Lbs, Fbs, Ls, FTs = [_.astype(np.complex64) for _ in (Lb, Fb, L, FT)]\n",
"%timeit analys2(Lbs,Fbs,synth2t(Ls,FTs,qs))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"comparing this to SHTns"
]
},
{
"cell_type": "code",
"execution_count": 23,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"36.6 µs ± 7.49 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
]
}
],
"source": [
"%timeit sht.analys(sht.synth(q))"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"there are probably some tricks to use though, let's focus just on synth step"
]
},
{
"cell_type": "code",
"execution_count": 52,
"metadata": {},
"outputs": [
{
"ename": "RuntimeError",
"evalue": "spectral array has wrong size.",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m<ipython-input-52-1b0b2a867938>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mget_ipython\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_line_magic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'timeit'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'sht.synth(q)'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_line_magic\u001b[0;34m(self, magic_name, line, _stack_depth)\u001b[0m\n\u001b[1;32m 2324\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'local_ns'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_local_scope\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstack_depth\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 2325\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuiltin_trap\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2326\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 2327\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 2328\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m<decorator-gen-60>\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/magic.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(f, *a, **k)\u001b[0m\n\u001b[1;32m 185\u001b[0m \u001b[0;31m# but it's overkill for just that one bit of state.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 186\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mmagic_deco\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 187\u001b[0;31m \u001b[0mcall\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 188\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 189\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/magics/execution.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n\u001b[1;32m 1167\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mindex\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1168\u001b[0m \u001b[0mnumber\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m10\u001b[0m \u001b[0;34m**\u001b[0m \u001b[0mindex\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1169\u001b[0;31m \u001b[0mtime_number\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtimer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimeit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnumber\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1170\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtime_number\u001b[0m \u001b[0;34m>=\u001b[0m \u001b[0;36m0.2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1171\u001b[0m \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/magics/execution.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, number)\u001b[0m\n\u001b[1;32m 167\u001b[0m \u001b[0mgc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdisable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 168\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 169\u001b[0;31m \u001b[0mtiming\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 170\u001b[0m \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 171\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mgcold\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m<magic-timeit>\u001b[0m in \u001b[0;36minner\u001b[0;34m(_it, _timer)\u001b[0m\n",
"\u001b[0;32m~/src/nfs/src/shtns/shtns.py\u001b[0m in \u001b[0;36msynth\u001b[0;34m(self, *arg)\u001b[0m\n\u001b[1;32m 289\u001b[0m \u001b[0mq\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 290\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 291\u001b[0;31m \u001b[0;32mif\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnlm\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"spectral array has wrong size.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 292\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'complex128'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"spectral array should be dtype=complex.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 293\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflags\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontiguous\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# contiguous array required.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mRuntimeError\u001b[0m: spectral array has wrong size."
]
}
],
"source": [
"%timeit sht.synth(q)"
]
},
{
"cell_type": "code",
"execution_count": 25,
"metadata": {},
"outputs": [],
"source": [
"# (theta,lm)*lm) . (lm, phi)\n",
"from numba import njit, complex128, prange\n",
"@njit(boundscheck=False,parallel=True)\n",
"def synth2j(FT, LT, q, out):\n",
" assert out.shape[0] == FT.shape[0] # theta\n",
" assert out.shape[1] == LT.shape[0] # phi\n",
" assert q.shape[0] == FT.shape[1] == LT.shape[1] # lm\n",
" for i in prange(out.shape[0]):\n",
" for j in range(out.shape[1]):\n",
" acc = 0.0\n",
" for l in range(q.size):\n",
" acc += q[l] * FT[i,l] * LT[j,l]\n",
" out[i,j] = acc.real"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"2.26 ms ± 10.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
]
}
],
"source": [
"out = np.empty_like(x)\n",
"outs = np.empty_like(x).astype(np.complex64)\n",
"LT = L.T.copy()\n",
"LTs = L.T.copy().astype(np.complex64)\n",
"synth2j(FTs, LTs, qs, outs)\n",
"%timeit synth2j(FTs, LTs, qs, outs)"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Overwriting synth2c.c\n"
]
}
],
"source": [
"%%file synth2c.c\n",
"#include <complex.h>\n",
"void synth2cs(\n",
" int ntheta, int nphi, int nlm,\n",
" float complex * restrict ft, float complex * restrict lt, \n",
" float complex * restrict q, float * restrict out)\n",
"{\n",
" #pragma omp parallel for\n",
" for (int i=0; i<ntheta; i++) {\n",
" float complex *fti = ft + i*nlm;\n",
" float *outi = out + i*nphi;\n",
" #pragma omp simd\n",
" for (int j=0; j<nphi; j++) {\n",
" float complex *ltj = lt + j*nlm;\n",
" float complex acc = 0.0f;\n",
" for (int l=0; l<nlm; l++) \n",
" acc += q[l] * fti[l] * ltj[l];\n",
" outi[j] = crealf(acc);\n",
" }\n",
" }\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 28,
"metadata": {},
"outputs": [],
"source": [
"!gcc -fPIC -O3 -fopenmp -march=native -mtune=native -ftree-vectorize -mavx2 -mavx512f -c synth2c.c\n",
"!gcc -shared synth2c.o -o synth2c.1.so -lgomp"
]
},
{
"cell_type": "code",
"execution_count": 29,
"metadata": {},
"outputs": [],
"source": [
"import ctypes as ct\n",
"from numpy.ctypeslib import ndpointer\n",
"lib = ct.CDLL('./synth2c.1.so')"
]
},
{
"cell_type": "code",
"execution_count": 30,
"metadata": {},
"outputs": [],
"source": [
"lib.synth2cs.restype = None\n",
"lib.synth2cs.argtypes = [\n",
" ct.c_int, ct.c_int, ct.c_int,\n",
" ndpointer(dtype=np.complex64, ndim=2, flags='C_CONTIGUOUS'), # FT\n",
" ndpointer(dtype=np.complex64, ndim=2, flags='C_CONTIGUOUS'), # LT\n",
" ndpointer(dtype=np.complex64, ndim=1, flags='C_CONTIGUOUS'), # q\n",
" ndpointer(dtype=np.float32, ndim=2, flags='C_CONTIGUOUS'), # out\n",
"]"
]
},
{
"cell_type": "code",
"execution_count": 31,
"metadata": {},
"outputs": [],
"source": [
"ntheta, nphi, nlm = theta.size, phi.size, sht.l.size"
]
},
{
"cell_type": "code",
"execution_count": 32,
"metadata": {},
"outputs": [],
"source": [
"outs = np.empty(out.shape, np.float32)"
]
},
{
"cell_type": "code",
"execution_count": 51,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"22.3 ms ± 5.59 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
]
}
],
"source": [
"%timeit lib.synth2cs(ntheta, nphi, nlm, FTs, LTs, qs, outs)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"so ~22ms vs ~18us for SHTns. but these faster functions are doing essentially method 1 with lower memory bandwidth.\n",
"\n",
"The conclusion for now is that we could use matrices if nothing else for data fitting, though it would be terribly slow. It would be better to figure out the adjoints by obtaining Y from Yb, etc.\n",
"\n",
"The dumb approach is just to brute force it with the GPU:"
]
},
{
"cell_type": "code",
"execution_count": 34,
"metadata": {},
"outputs": [],
"source": [
"import cupy as cp"
]
},
{
"cell_type": "code",
"execution_count": 35,
"metadata": {},
"outputs": [],
"source": [
"gFTs, gLs, gqs = [cp.array(_) for _ in (FTs, Ls, qs)] # cp.array(FTs)"
]
},
{
"cell_type": "code",
"execution_count": 36,
"metadata": {},
"outputs": [],
"source": [
"synth2tg = lambda L, FT, q: cp.dot(FT * q, L)"
]
},
{
"cell_type": "code",
"execution_count": 49,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"177 µs ± 1.64 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)\n"
]
}
],
"source": [
"%timeit synth2tg(gLs, gFTs, gqs)"
]
},
{
"cell_type": "code",
"execution_count": 50,
"metadata": {},
"outputs": [
{
"ename": "RuntimeError",
"evalue": "spectral array has wrong size.",
"output_type": "error",
"traceback": [
"\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
"\u001b[0;31mRuntimeError\u001b[0m Traceback (most recent call last)",
"\u001b[0;32m<ipython-input-50-1b0b2a867938>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[0;32m----> 1\u001b[0;31m \u001b[0mget_ipython\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mrun_line_magic\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'timeit'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'sht.synth(q)'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/interactiveshell.py\u001b[0m in \u001b[0;36mrun_line_magic\u001b[0;34m(self, magic_name, line, _stack_depth)\u001b[0m\n\u001b[1;32m 2324\u001b[0m \u001b[0mkwargs\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m'local_ns'\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mget_local_scope\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mstack_depth\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 2325\u001b[0m \u001b[0;32mwith\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbuiltin_trap\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2326\u001b[0;31m \u001b[0mresult\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mfn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 2327\u001b[0m \u001b[0;32mreturn\u001b[0m \u001b[0mresult\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 2328\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m<decorator-gen-60>\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/magic.py\u001b[0m in \u001b[0;36m<lambda>\u001b[0;34m(f, *a, **k)\u001b[0m\n\u001b[1;32m 185\u001b[0m \u001b[0;31m# but it's overkill for just that one bit of state.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 186\u001b[0m \u001b[0;32mdef\u001b[0m \u001b[0mmagic_deco\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 187\u001b[0;31m \u001b[0mcall\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;32mlambda\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mf\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0ma\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mk\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 188\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 189\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mcallable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/magics/execution.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, line, cell, local_ns)\u001b[0m\n\u001b[1;32m 1167\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mindex\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;36m10\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1168\u001b[0m \u001b[0mnumber\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0;36m10\u001b[0m \u001b[0;34m**\u001b[0m \u001b[0mindex\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1169\u001b[0;31m \u001b[0mtime_number\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtimer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimeit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mnumber\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 1170\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mtime_number\u001b[0m \u001b[0;34m>=\u001b[0m \u001b[0;36m0.2\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 1171\u001b[0m \u001b[0;32mbreak\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m/usr/lib/python3/dist-packages/IPython/core/magics/execution.py\u001b[0m in \u001b[0;36mtimeit\u001b[0;34m(self, number)\u001b[0m\n\u001b[1;32m 167\u001b[0m \u001b[0mgc\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdisable\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 168\u001b[0m \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 169\u001b[0;31m \u001b[0mtiming\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0minner\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mit\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mtimer\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 170\u001b[0m \u001b[0;32mfinally\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 171\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mgcold\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;32m<magic-timeit>\u001b[0m in \u001b[0;36minner\u001b[0;34m(_it, _timer)\u001b[0m\n",
"\u001b[0;32m~/src/nfs/src/shtns/shtns.py\u001b[0m in \u001b[0;36msynth\u001b[0;34m(self, *arg)\u001b[0m\n\u001b[1;32m 289\u001b[0m \u001b[0mq\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlist\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0marg\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 290\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mi\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;36m0\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 291\u001b[0;31m \u001b[0;32mif\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnlm\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"spectral array has wrong size.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 292\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum\u001b[0m \u001b[0;34m!=\u001b[0m \u001b[0mnp\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mdtype\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'complex128'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mnum\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0;32mraise\u001b[0m \u001b[0mRuntimeError\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m\"spectral array should be dtype=complex.\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 293\u001b[0m \u001b[0;32mif\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mflags\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcontiguous\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;32mFalse\u001b[0m\u001b[0;34m:\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mq\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcopy\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# contiguous array required.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
"\u001b[0;31mRuntimeError\u001b[0m: spectral array has wrong size."
]
}
],
"source": [
"%timeit sht.synth(q)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Only 3x slowdown (probably less on better version) is enough. We'd be able to trivially put this into TF."
]
},
{
"cell_type": "code",
"execution_count": 39,
"metadata": {},
"outputs": [],
"source": [
"from numba import cuda\n",
"import numba as nb\n",
"@cuda.jit\n",
"def synth2cj(FT, LT, q, out):\n",
" i = cuda.blockIdx.x\n",
" j = cuda.blockIdx.y\n",
" nl = cuda.blockDim.x\n",
" l = cuda.threadIdx.x\n",
" qfl = cuda.shared.array((128,), nb.complex64)\n",
" if l==0: out[i,j] = 0\n",
" for offset in range(4):\n",
" o = offset * 128\n",
" qfl[l] = q[o+l] * FT[i,o+l] * LT[j,o+l]\n",
" ll = nl # 545\n",
" while ll>1:\n",
" # if odd add last to penultimate\n",
" if ll & 0x1:\n",
" if l==0: qfl[ll-2] += qfl[ll-1]\n",
" ll -= 1\n",
" # procede with even number\n",
" llh = ll >> 1\n",
" # qfl[:272] += qfl[272:]\n",
" if l < llh:\n",
" qfl[l] += qfl[llh + l]\n",
" ll >>= 1\n",
" if l==0:\n",
" out[i,j] += qfl[l].real"
]
},
{
"cell_type": "code",
"execution_count": 40,
"metadata": {},
"outputs": [],
"source": [
"gf, gl, gq, go = [cuda.to_device(_) for _ in (FTs, LTs, qs, out)]"
]
},
{
"cell_type": "code",
"execution_count": 48,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"1.25 ms ± 109 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)\n"
]
}
],
"source": [
"stream = nb.cuda.stream()\n",
"%timeit synth2cj[go.shape, 128, stream](gf, gl, gq, go)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"This could go faster if it used less shared memory. Profiler would help."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### computing equivalent lc matrix\n",
"\n",
"in the linear comparison case, \n",
"\n",
"- lmax 70\n",
"- mmax 70\n",
"- nlat 74\n",
"- nphi 148 (this is longitude in SHTns)\n"
]
},
{
"cell_type": "code",
"execution_count": 42,
"metadata": {},
"outputs": [],
"source": [
"sht = shtns.sht(lmax=70,mmax=70)\n",
"sht.set_grid(nlat=74,nphi=148)\n",
"theta = np.r_[0:2*np.pi:(sht.nphi+1)*1j][:-1]\n",
"phi = np.r_[0:np.pi:(sht.nlat+1)*1j][:-1]\n",
"F = np.exp(1j * sht.m[:,None] * theta)\n",
"L = sph_harm(sht.m[:,None], sht.l[:,None], 0, phi)\n",
"Lb = np.conj(L) * np.sin(phi) * dp\n",
"Fb = np.conj(F) * dt"
]
},
{
"cell_type": "code",
"execution_count": 43,
"metadata": {},
"outputs": [],
"source": [
"diff_term = -0.000314138 * sht.l * (sht.l + 1) # local_connectivity.py L98"
]
},
{
"cell_type": "code",
"execution_count": 44,
"metadata": {},
"outputs": [],
"source": [
"# synth2 = lambda L, F, q: np.dot(F.T * q, L)\n",
"# FT = F.T.copy()\n",
"# synth2t = lambda L, FT, q: np.dot(FT * q, L)\n",
"# analys2 = lambda Lb, Fb, x: np.sum(Lb * np.dot(Fb,x), axis=1)"
]
},
{
"cell_type": "code",
"execution_count": 45,
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"((2556, 74, 148), (2556, 74, 148))"
]
},
"execution_count": 45,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"r = diff_term[:,None,None] * Lb[:,:,None]*Fb[:,None,:]\n",
"l = F[:,None,:] * L[:,:,None]\n",
"r.shape, l.shape"
]
},
{
"cell_type": "code",
"execution_count": 46,
"metadata": {},
"outputs": [],
"source": [
"# np.einsum('ijk,ilm',r,l).shape"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "nfs",
"language": "python",
"name": "nfs"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.8.6"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment