def cryo_pft(p, n_r, n_theta):
"""
Compute the polar Fourier transform of projections with resolution n_r in the radial direction
and resolution n_theta in the angular direction.
:param p:
:param n_r: Number of samples along each ray (in the radial direction).
:param n_theta: Angular resolution. Number of Fourier rays computed for each projection.
:return:
"""
if n_theta % 2 == 1:
raise ValueError('n_theta must be even')
n_projs = p.shape[2]
omega0 = 2 * np.pi / (2 * n_r - 1)
dtheta = 2 * np.pi / n_theta
freqs = np.zeros((2, n_r * n_theta // 2))
for i in range(n_theta // 2):
freqs[0, i * n_r:(i + 1) * n_r] = np.arange(n_r) * np.sin(i * dtheta)
freqs[1, i * n_r:(i + 1) * n_r] = np.arange(n_r) * np.cos(i * dtheta)
freqs *= omega0
# finufftpy require it to be aligned in fortran order
pf = np.empty((n_r * n_theta // 2, n_projs), dtype='complex128', order='F')
finufftpy.nufft2d2many(freqs[0], freqs[1], pf, 1, 1e-15, p)
pf = pf.reshape((n_r, n_theta // 2, n_projs), order='F')
pf = np.concatenate((pf, pf.conj()), axis=1).copy()
return pf, freqs
def evaluate_t(self, x):
"""
Evaluate coefficient in polar Fourier grid from those in standard 2D coordinate basis
:param x: The coefficient array in the standard 2D coordinate basis to be
evaluated. The first two dimensions must equal `self.sz`.
:return v: The evaluation of the coefficient array `v` in the polar Fourier grid.
This is an array of vectors whose first dimension is `self.count` and
whose remaining dimensions correspond to higher dimensions of `x`.
"""
# ensure the first two dimensions with size of self.sz
x, sz_roll = unroll_dim(x, self.ndim + 1)
nimgs = x.shape[2]
# finufftpy require it to be aligned in fortran order
half_size = self.nrad * self.ntheta // 2
pf = np.empty((half_size, nimgs), dtype='complex128', order='F')
finufftpy.nufft2d2many(self.freqs[0, :half_size], self.freqs[1, :half_size], pf, 1, 1e-15, x)
pf = m_reshape(pf, (self.nrad, self.ntheta // 2, nimgs))
v = np.concatenate((pf, pf.conj()), axis=1)
# return v coefficients with the first dimension size of self.count
v = m_reshape(v, (self.nrad * self.ntheta, nimgs))
v = roll_dim(v, sz_roll)
return v
def cryo_pft_nfft(projections, precomp):
freqs = precomp.freqs
n_theta = precomp.n_theta
n_r = precomp.n_r
num_projections = projections.shape[2]
x = -2 * np.pi * freqs.T
x = x.copy()
pf = np.empty((x.shape[1], num_projections), dtype='complex128', order='F')
finufftpy.nufft2d2many(x[0], x[1], pf, -1, 1e-15, projections)
pf = pf.reshape((n_r, n_theta, num_projections), order='F')
return pf
def transformed_mapping_matrices_from_mapping_matrix(self, mapping_matrix):
mapping_matrix_reshaped = self.reshape_mapping_matrix(
mapping_matrix=mapping_matrix, shape_2d=self.grid.shape_2d)
visibilities = np.zeros(shape=(self.uv_wavelengths.shape[0],
mapping_matrix.shape[1]),
order='F',
dtype=np.complex)
finufftpy.nufft2d2many(self.uv_wavelengths_normalized[:, 1] * np.pi,
self.uv_wavelengths_normalized[:, 0] * np.pi,
visibilities, 1, self.eps,
mapping_matrix_reshaped)
for i in range(visibilities.shape[1]):
visibilities[:, i] *= self.shift
return [visibilities.real, visibilities.imag]
def vel_convolution_nufft(scalar, source_strenth, num_modes, L, eps=1e-8, method=1, *args, **kwargs):
if('kernel_hat' in kwargs):
kernel_hat = kwargs['kernel_hat']
assert isinstance(kernel_hat, np.ndarray)
else:
raise Exception(
'kernel_hat, fourier transform of kernel should be given')
Nshape = scalar.shape
dim = Nshape[0]
assert dim == 2, "Dim should be 2!"
nx, ny = num_modes[:]
if method == 1:
xj, yj = scalar
nj = len(xj)
xmin = xj.min()
xmax = xj.max()
ymin = yj.min()
ymax = yj.max()
Lx, Ly = L
assert ((xmax-xmin) <= Lx) & ((ymax-ymin) <= Ly), "All particles should be limieted to given [Lx,Ly] box!"
# double zero-padding the delta function and scale it into [-pi,pi)
xj = (xj-(xmax+xmin)/2)/Lx*np.pi
yj = (yj-(ymax+ymin)/2)/Ly*np.pi
fk = np.zeros(num_modes, dtype=np.complex128, order='F')
# particle locations change at every iteration, no need to reuse fftw object
ret = finufftpy.nufft2d1(xj, yj, source_strenth, -1, eps, nx, ny, fk, modeord=1)
assert ret == 0, "NUFFT not successful!"
v_hat = np.zeros((nx, ny, 2), order='F', dtype=np.complex128)
v_hat[:, :, 0] = fk * kernel_hat[0]
v_hat[:, :, 1] = fk * kernel_hat[1]
v = np.zeros((nj, 2), order='F', dtype=np.complex128)
ret = finufftpy.nufft2d2many(xj, yj, v, 1, eps, v_hat, modeord=1)
assert ret == 0, "NUFFT not successful!"
v = np.real(v)/(2*Lx*2*Ly) # Real?
return v
def accuracy_speed_tests(num_nonuniform_points,num_uniform_points,eps):
nj,nk = int(num_nonuniform_points),int(num_nonuniform_points)
iflag=1
num_samples=int(np.minimum(5,num_uniform_points*0.5+1)) # number of outputs used for estimating accuracy; is small for speed
print('Accuracy and speed tests for %d nonuniform points and eps=%g (error estimates use %d samples per run)' % (num_nonuniform_points,eps,num_samples))
# for doing the error estimates
Xest=np.zeros(num_samples,dtype=np.complex128)
Xtrue=np.zeros(num_samples,dtype=np.complex128)
###### 1-d cases ........................................................
ms=int(num_uniform_points)
xj=np.random.rand(nj)*2*math.pi-math.pi
cj=np.random.rand(nj)+1j*np.random.rand(nj);
fk=np.zeros([ms],dtype=np.complex128)
timer=time.time()
ret=finufftpy.nufft1d1(xj,cj,iflag,eps,ms,fk)
elapsed=time.time()-timer
k=np.arange(-np.floor(ms/2),np.floor((ms-1)/2+1))
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(cj * np.exp(1j*k[ii]*xj))
Xtrue[ii]=fk[ii]
print_report('finufft1d1',elapsed,Xest,Xtrue,nj)
xj=np.random.rand(nj)*2*math.pi-math.pi
cj=np.zeros([nj],dtype=np.complex128);
fk=np.random.rand(ms)+1j*np.random.rand(ms);
timer=time.time()
ret=finufftpy.nufft1d2(xj,cj,iflag,eps,fk)
elapsed=time.time()-timer
k=np.arange(-np.floor(ms/2),np.floor((ms-1)/2+1))
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(fk * np.exp(1j*k*xj[ii]))
Xtrue[ii]=cj[ii]
print_report('finufft1d2',elapsed,Xest,Xtrue,nj)
x=np.random.rand(nj)*2*math.pi-math.pi
c=np.random.rand(nj)+1j*np.random.rand(nj);
s=np.random.rand(nk)*2*math.pi-math.pi
f=np.zeros([nk],dtype=np.complex128)
timer=time.time()
ret=finufftpy.nufft1d3(x,c,iflag,eps,s,f)
elapsed=time.time()-timer
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(c * np.exp(1j*s[ii]*x))
Xtrue[ii]=f[ii]
print_report('finufft1d3',elapsed,Xest,Xtrue,nj+nk)
###### 2-d cases ....................................................
ms=int(np.ceil(np.sqrt(num_uniform_points)))
mt=ms
xj=np.random.rand(nj)*2*math.pi-math.pi
yj=np.random.rand(nj)*2*math.pi-math.pi
cj=np.random.rand(nj)+1j*np.random.rand(nj)
fk=np.zeros([ms,mt],dtype=np.complex128,order='F')
timer=time.time()
ret=finufftpy.nufft2d1(xj,yj,cj,iflag,eps,ms,mt,fk)
elapsed=time.time()-timer
Ks,Kt=np.mgrid[-np.floor(ms/2):np.floor((ms-1)/2+1),-np.floor(mt/2):np.floor((mt-1)/2+1)]
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(cj * np.exp(1j*(Ks.ravel()[ii]*xj+Kt.ravel()[ii]*yj)))
Xtrue[ii]=fk.ravel()[ii]
print_report('finufft2d1',elapsed,Xest,Xtrue,nj)
## 2d1many:
ndata = 5 # how many vectors to do
cj=np.array(np.random.rand(nj,ndata)+1j*np.random.rand(nj,ndata),order='F')
fk=np.zeros([ms,mt,ndata],dtype=np.complex128,order='F')
timer=time.time()
ret=finufftpy.nufft2d1many(xj,yj,cj,iflag,eps,ms,mt,fk)
elapsed=time.time()-timer
dtest = ndata-1 # which of the ndata to test (in 0,..,ndata-1)
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(cj[:,dtest] * np.exp(1j*(Ks.ravel(order='F')[ii]*xj+Kt.ravel(order='F')[ii]*yj))) # note fortran-ravel-order needed throughout - mess.
Xtrue[ii]=fk.ravel(order='F')[ii + dtest*ms*mt] # hack the offset in fk array - has to be better way
print_report('finufft2d1many',elapsed,Xest,Xtrue,ndata*nj)
# 2d2
xj=np.random.rand(nj)*2*math.pi-math.pi
yj=np.random.rand(nj)*2*math.pi-math.pi
cj=np.zeros([nj],dtype=np.complex128);
fk=np.random.rand(ms,mt)+1j*np.random.rand(ms,mt);
timer=time.time()
ret=finufftpy.nufft2d2(xj,yj,cj,iflag,eps,fk)
elapsed=time.time()-timer
Ks,Kt=np.mgrid[-np.floor(ms/2):np.floor((ms-1)/2+1),-np.floor(mt/2):np.floor((mt-1)/2+1)]
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(fk * np.exp(1j*(Ks*xj[ii]+Kt*yj[ii])))
Xtrue[ii]=cj[ii]
print_report('finufft2d2',elapsed,Xest,Xtrue,nj)
# 2d2many (using same ndata and dtest as 2d1many; see above)
cj=np.zeros([nj,ndata],order='F',dtype=np.complex128);
fk=np.array(np.random.rand(ms,mt,ndata)+1j*np.random.rand(ms,mt,ndata),order='F')
timer=time.time()
ret=finufftpy.nufft2d2many(xj,yj,cj,iflag,eps,fk)
elapsed=time.time()-timer
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(fk[:,:,dtest] * np.exp(1j*(Ks*xj[ii]+Kt*yj[ii])))
Xtrue[ii]=cj[ii,dtest]
print_report('finufft2d2many',elapsed,Xest,Xtrue,ndata*nj)
# 2d3
x=np.random.rand(nj)*2*math.pi-math.pi
y=np.random.rand(nj)*2*math.pi-math.pi
c=np.random.rand(nj)+1j*np.random.rand(nj);
s=np.random.rand(nk)*2*math.pi-math.pi
t=np.random.rand(nk)*2*math.pi-math.pi
f=np.zeros([nk],dtype=np.complex128)
timer=time.time()
ret=finufftpy.nufft2d3(x,y,c,iflag,eps,s,t,f)
elapsed=time.time()-timer
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(c * np.exp(1j*(s[ii]*x+t[ii]*y)))
Xtrue[ii]=f[ii]
print_report('finufft2d3',elapsed,Xest,Xtrue,nj+nk)
###### 3-d cases ............................................................
ms=int(np.ceil(num_uniform_points**(1.0/3)))
mt=ms
mu=ms
xj=np.random.rand(nj)*2*math.pi-math.pi
yj=np.random.rand(nj)*2*math.pi-math.pi
zj=np.random.rand(nj)*2*math.pi-math.pi
cj=np.random.rand(nj)+1j*np.random.rand(nj);
fk=np.zeros([ms,mt,mu],dtype=np.complex128,order='F')
timer=time.time()
ret=finufftpy.nufft3d1(xj,yj,zj,cj,iflag,eps,ms,mt,mu,fk)
elapsed=time.time()-timer
Ks,Kt,Ku=np.mgrid[-np.floor(ms/2):np.floor((ms-1)/2+1),-np.floor(mt/2):np.floor((mt-1)/2+1),-np.floor(mu/2):np.floor((mu-1)/2+1)]
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(cj * np.exp(1j*(Ks.ravel()[ii]*xj+Kt.ravel()[ii]*yj+Ku.ravel()[ii]*zj)))
Xtrue[ii]=fk.ravel()[ii]
print_report('finufft3d1',elapsed,Xest,Xtrue,nj)
xj=np.random.rand(nj)*2*math.pi-math.pi
yj=np.random.rand(nj)*2*math.pi-math.pi
zj=np.random.rand(nj)*2*math.pi-math.pi
cj=np.zeros([nj],dtype=np.complex128);
fk=np.random.rand(ms,mt,mu)+1j*np.random.rand(ms,mt,mu);
timer=time.time()
ret=finufftpy.nufft3d2(xj,yj,zj,cj,iflag,eps,fk)
elapsed=time.time()-timer
Ks,Kt,Ku=np.mgrid[-np.floor(ms/2):np.floor((ms-1)/2+1),-np.floor(mt/2):np.floor((mt-1)/2+1),-np.floor(mu/2):np.floor((mu-1)/2+1)]
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(fk * np.exp(1j*(Ks*xj[ii]+Kt*yj[ii]+Ku*zj[ii])))
Xtrue[ii]=cj[ii]
print_report('finufft3d2',elapsed,Xest,Xtrue,nj)
x=np.random.rand(nj)*2*math.pi-math.pi
y=np.random.rand(nj)*2*math.pi-math.pi
z=np.random.rand(nj)*2*math.pi-math.pi
c=np.random.rand(nj)+1j*np.random.rand(nj);
s=np.random.rand(nk)*2*math.pi-math.pi
t=np.random.rand(nk)*2*math.pi-math.pi
u=np.random.rand(nk)*2*math.pi-math.pi
f=np.zeros([nk],dtype=np.complex128)
timer=time.time()
ret=finufftpy.nufft3d3(x,y,z,c,iflag,eps,s,t,u,f)
elapsed=time.time()-timer
for ii in np.arange(0,num_samples):
Xest[ii]=np.sum(c * np.exp(1j*(s[ii]*x+t[ii]*y+u[ii]*z)))
Xtrue[ii]=f[ii]
print_report('finufft3d3',elapsed,Xest,Xtrue,nj+nk)