From 2fb7511a714c7a6916918b703af84d8770db0cd2 Mon Sep 17 00:00:00 2001 From: math-vrn Date: Sun, 12 Jul 2020 15:31:24 -0500 Subject: [PATCH 1/6] add fwd and adj gradient operators --- src/tike/operators/cupy/reg.py | 10 ++++++++++ src/tike/operators/numpy/reg.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+) create mode 100644 src/tike/operators/cupy/reg.py create mode 100644 src/tike/operators/numpy/reg.py diff --git a/src/tike/operators/cupy/reg.py b/src/tike/operators/cupy/reg.py new file mode 100644 index 00000000..95c407cc --- /dev/null +++ b/src/tike/operators/cupy/reg.py @@ -0,0 +1,10 @@ +__author__ = "Daniel Ching, Viktor Nikitin" +__copyright__ = "Copyright (c) 2020, UChicago Argonne, LLC." + +import cupy as cp +from tike.operators import numpy +from .operator import Operator + + +class Reg(Operator, numpy.Reg): + pass diff --git a/src/tike/operators/numpy/reg.py b/src/tike/operators/numpy/reg.py new file mode 100644 index 00000000..2f44f7ad --- /dev/null +++ b/src/tike/operators/numpy/reg.py @@ -0,0 +1,33 @@ +__author__ = "Daniel Ching, Viktor Nikitin" +__copyright__ = "Copyright (c) 2020, UChicago Argonne, LLC." + +import numpy as np +from numpy.fft import fft2, ifft2 + +from .operator import Operator + + +class Reg(Operator): + """3D Gradient and divergence operators for regularization.""" + + def fwd(self, u): + """3D gradient operator""" + res = self.xp.zeros([3, *u.shape], dtype='float32') + res[0, :, :, :-1] = u[:, :, 1:]-u[:, :, :-1] + res[1, :, :-1, :] = u[:, 1:, :]-u[:, :-1, :] + res[2, :-1, :, :] = u[1:, :, :]-u[:-1, :, :] + return res + + def adj(self, gr): + """3D negative divergence operator""" + res = self.xp.zeros(gr.shape[1:], dtype='float32') + res[:, :, 1:] = gr[0, :, :, 1:]-gr[0, :, :, :-1] + res[:, :, 0] = gr[0, :, :, 0] + res[:, 1:, :] += gr[1, :, 1:, :]-gr[1, :, :-1, :] + res[:, 0, :] += gr[1, :, 0, :] + res[1:, :, :] += gr[2, 1:, :, :]-gr[2, :-1, :, :] + res[0, :, :] += gr[2, 0, :, :] + return -res + + + \ No newline at end of file From 97b4418aa0e7d5631c95ef656882119ebcd5a5b0 Mon Sep 17 00:00:00 2001 From: math-vrn Date: Sun, 12 Jul 2020 15:43:39 -0500 Subject: [PATCH 2/6] remove white spaces --- src/tike/operators/numpy/reg.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/tike/operators/numpy/reg.py b/src/tike/operators/numpy/reg.py index 2f44f7ad..ef03bc79 100644 --- a/src/tike/operators/numpy/reg.py +++ b/src/tike/operators/numpy/reg.py @@ -27,7 +27,4 @@ def adj(self, gr): res[:, 0, :] += gr[1, :, 0, :] res[1:, :, :] += gr[2, 1:, :, :]-gr[2, :-1, :, :] res[0, :, :] += gr[2, 0, :, :] - return -res - - - \ No newline at end of file + return -res \ No newline at end of file From 0f5e85e184110e5200a9b5f36bb259343cf5522f Mon Sep 17 00:00:00 2001 From: math-vrn Date: Sun, 12 Jul 2020 15:44:31 -0500 Subject: [PATCH 3/6] add empty line --- src/tike/operators/numpy/reg.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/tike/operators/numpy/reg.py b/src/tike/operators/numpy/reg.py index ef03bc79..e89c6014 100644 --- a/src/tike/operators/numpy/reg.py +++ b/src/tike/operators/numpy/reg.py @@ -27,4 +27,5 @@ def adj(self, gr): res[:, 0, :] += gr[1, :, 0, :] res[1:, :, :] += gr[2, 1:, :, :]-gr[2, :-1, :, :] res[0, :, :] += gr[2, 0, :, :] - return -res \ No newline at end of file + return -res + \ No newline at end of file From 80a948ddec500c410d602ced4991dc4d06e337fe Mon Sep 17 00:00:00 2001 From: math-vrn Date: Sun, 12 Jul 2020 15:45:51 -0500 Subject: [PATCH 4/6] autopep8 --- src/tike/operators/numpy/reg.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/tike/operators/numpy/reg.py b/src/tike/operators/numpy/reg.py index e89c6014..5245ef36 100644 --- a/src/tike/operators/numpy/reg.py +++ b/src/tike/operators/numpy/reg.py @@ -28,4 +28,3 @@ def adj(self, gr): res[1:, :, :] += gr[2, 1:, :, :]-gr[2, :-1, :, :] res[0, :, :] += gr[2, 0, :, :] return -res - \ No newline at end of file From d601256e8c2e727b7872e8966d4e03159b40f975 Mon Sep 17 00:00:00 2001 From: math-vrn Date: Thu, 13 Aug 2020 10:40:58 -0500 Subject: [PATCH 5/6] add 2d usfft functions --- src/tike/operators/cupy/usfft.py | 121 ++++++++++++++++++++++++++++++- 1 file changed, 120 insertions(+), 1 deletion(-) diff --git a/src/tike/operators/cupy/usfft.py b/src/tike/operators/cupy/usfft.py index 1d35411d..9b0fbffc 100644 --- a/src/tike/operators/cupy/usfft.py +++ b/src/tike/operators/cupy/usfft.py @@ -13,6 +13,11 @@ def _get_kernel(xp, pad, mu): xeq = xp.mgrid[-pad:pad, -pad:pad, -pad:pad] return xp.exp(-mu * xp.sum(xeq**2, axis=0)).astype('float32') +def _get_kernel2d(xp, pad, mu): + """Return the interpolation kernel for the 2d USFFT.""" + xeq = xp.mgrid[-pad:pad, -pad:pad] + return xp.exp(-mu * xp.sum(xeq**2, axis=0)).astype('float32') + def vector_gather(xp, Fe, x, n, m, mu): """A faster implementation of sequential_gather""" @@ -35,6 +40,23 @@ def delta(l, i, x): (n + ell[:, 2] + i2) % (2 * n)] * Fkernel return F +def vector_gather2d(xp, Fe, x, n, m, mu): + """A faster implementation of sequential_gather""" + cons = [xp.sqrt(xp.pi / mu)**2, -xp.pi**2 / mu] + + def delta(l, i, x): + return ((l + i).astype('float32') / (2 * n) - x)**2 + + F = xp.zeros(x.shape[0], dtype="complex64") + ell = ((2 * n * x) // 1).astype(xp.int32) # nearest grid to x + for i0 in range(-m, m): + delta0 = delta(ell[:, 0], i0, x[:, 0]) + for i1 in range(-m, m): + delta1 = delta(ell[:, 1], i1, x[:, 1]) + Fkernel = cons[0] * xp.exp(cons[1] * (delta0 + delta1)) + F += Fe[(n + ell[:, 0] + i0) % (2 * n), + (n + ell[:, 1] + i1) % (2 * n)] * Fkernel + return F def sequential_gather(xp, Fe, x, n, m, mu): """Gather F from the regular grid. @@ -104,6 +126,38 @@ def eq2us(f, x, n, eps, xp, gather=vector_gather, fftn=None): return F +def eq2us2d(f, x, n, eps, xp, gather=vector_gather2d, fftn=None): + """2d USFFT from equally-spaced grid to unequally-spaced grid. + + Parameters + ---------- + f : (n, n) complex64 + The function to be transformed on a regular-grid of size n. + x : (N, 2) + The sampled frequencies on unequally-spaced grid. + eps : float + The desired relative accuracy of the USFFT. + """ + fftn = xp.fft.fftn if fftn is None else fftn + ndim = f.ndim + pad = n // 2 # where zero-padding stops + end = pad + n # where f stops + + # parameters for the USFFT transform + mu = -xp.log(eps) / (2 * n**2) + Te = 1 / xp.pi * xp.sqrt(-mu * xp.log(eps) + (mu * n)**2 / 4) + m = xp.int(xp.ceil(2 * n * Te)) + + # smearing kernel (kernel) + kernel = _get_kernel2d(xp, pad, mu) + + # FFT and compesantion for smearing + fe = xp.zeros([2 * n] * ndim, dtype="complex64") + fe[pad:end, pad:end] = f / ((2 * n)**ndim * kernel) + Fe = checkerboard(xp, fftn(checkerboard(xp, fe)), inverse=True) + F = gather(xp, Fe, x, n, m, mu) + + return F def sequential_scatter(xp, f, x, n, m, mu): """Scatter f to the regular grid. @@ -171,6 +225,34 @@ def delta(l, i, x): G[ids] += vals[ids] return G.reshape([2 * n] * ndim) +def vector_scatter2d(xp, f, x, n, m, mu): + """A faster implemenation of sequential_scatter.""" + cons = [xp.sqrt(xp.pi / mu)**2, -xp.pi**2 / mu] + + def delta(l, i, x): + return ((l + i).astype('float32') / (2 * n) - x)**2 + + G = xp.zeros([(2 * n)**2], dtype="complex64") + ell = ((2 * n * x) // 1).astype(xp.int32) # nearest grid to x + stride = ((2 * n)**2, 2 * n) + for i0 in range(-m, m): + delta0 = delta(ell[:, 0], i0, x[:, 0]) + for i1 in range(-m, m): + delta1 = delta(ell[:, 1], i1, x[:, 1]) + Fkernel = cons[0] * xp.exp(cons[1] * (delta0 + delta1)) + ids = ( + ((n + ell[:, 2] + i2) % (2 * n)) + + stride[1] * ((n + ell[:, 1] + i1) % (2 * n)) + ) # yapf: disable + vals = f * Fkernel + # accumulate by indexes (with possible index intersections), + # TODO acceleration of bincount!! + vals = (xp.bincount(ids, weights=vals.real) + + 1j * xp.bincount(ids, weights=vals.imag)) + ids = xp.nonzero(vals)[0] + G[ids] += vals[ids] + return G.reshape([2 * n] * 2) + def us2eq(f, x, n, eps, xp, scatter=vector_scatter, fftn=None): """USFFT from unequally-spaced grid to equally-spaced grid. @@ -198,7 +280,7 @@ def us2eq(f, x, n, eps, xp, scatter=vector_scatter, fftn=None): m = xp.int(xp.ceil(2 * n * Te)) # smearing kernel (ker) - kernel = _get_kernel(xp, pad, mu) + kernel = _get_kernel2d(xp, pad, mu) G = scatter(xp, f, x, n, m, mu) @@ -209,6 +291,43 @@ def us2eq(f, x, n, eps, xp, scatter=vector_scatter, fftn=None): return F +def us2eq2d(f, x, n, eps, xp, scatter=vector_scatter2d, fftn=None): + """2d USFFT from unequally-spaced grid to equally-spaced grid. + + Parameters + ---------- + f : (n**2) complex64 + Values of unequally-spaced function on the grid x + x : (n**2) float + The frequencies on the unequally-spaced grid + n : int + The size of the equall spaced grid. + eps : float + The accuracy of computing USFFT + scatter : function + The scatter function to use. + """ + fftn = xp.fft.fftn if fftn is None else fftn + pad = n // 2 # where zero-padding stops + end = pad + n # where f stops + + # parameters for the USFFT transform + mu = -xp.log(eps) / (2 * n**2) + Te = 1 / xp.pi * xp.sqrt(-mu * xp.log(eps) + (mu * n)**2 / 4) + m = xp.int(xp.ceil(2 * n * Te)) + + # smearing kernel (ker) + kernel = _get_kernel(xp, pad, mu) + + G = scatter(xp, f, x, n, m, mu) + + # FFT and compesantion for smearing + F = checkerboard(xp, fftn(checkerboard(xp, G)), inverse=True) + F = F[pad:end, pad:end] / ((2 * n)**2 * kernel) + + return F + + def _unpad(array, width, mode='wrap'): """Remove padding from an array in-place. From 4a6ed4793bfbc57e54215f06e1f568fb7a326540 Mon Sep 17 00:00:00 2001 From: math-vrn Date: Thu, 13 Aug 2020 10:44:02 -0500 Subject: [PATCH 6/6] fix 2d usfft --- src/tike/operators/cupy/usfft.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/tike/operators/cupy/usfft.py b/src/tike/operators/cupy/usfft.py index 9b0fbffc..28dc0573 100644 --- a/src/tike/operators/cupy/usfft.py +++ b/src/tike/operators/cupy/usfft.py @@ -280,7 +280,7 @@ def us2eq(f, x, n, eps, xp, scatter=vector_scatter, fftn=None): m = xp.int(xp.ceil(2 * n * Te)) # smearing kernel (ker) - kernel = _get_kernel2d(xp, pad, mu) + kernel = _get_kernel(xp, pad, mu) G = scatter(xp, f, x, n, m, mu) @@ -317,7 +317,7 @@ def us2eq2d(f, x, n, eps, xp, scatter=vector_scatter2d, fftn=None): m = xp.int(xp.ceil(2 * n * Te)) # smearing kernel (ker) - kernel = _get_kernel(xp, pad, mu) + kernel = _get_kernel2d(xp, pad, mu) G = scatter(xp, f, x, n, m, mu)