def iterate(self, iloop, save_wisdom=1, verbose=1):
dat = self.dat
ran = self.ran
smooth = self.smooth
binsize = self.binsize
beta = self.beta
bias = self.bias
f = self.f
nbins = self.nbins
print("Loop %d" % iloop)
#-- Creating arrays for FFTW
if iloop == 0:
delta = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
deltak = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
rho = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
rhok = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
psi_x = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
psi_y = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
psi_z = pyfftw.empty_aligned((nbins, nbins, nbins),
dtype='complex128')
#-- Initialize FFT objects and load wisdom if available
wisdom_file = "wisdom." + str(nbins) + "." + str(
self.nthreads) + '.npy'
if os.path.isfile(wisdom_file):
print('Reading wisdom from ', wisdom_file)
wisd = tuple(np.load(wisdom_file))
print('Status of importing wisdom', pyfftw.import_wisdom(wisd))
sys.stdout.flush()
print('Creating FFTW objects...')
fft_obj = pyfftw.FFTW(delta,
delta,
axes=[0, 1, 2],
threads=self.nthreads)
ifft_obj = pyfftw.FFTW(deltak, psi_x, axes=[0, 1, 2], \
threads=self.nthreads, \
direction='FFTW_BACKWARD')
kr = fftfreq(nbins, d=binsize) * 2 * np.pi * self.smooth
norm = np.exp(-0.5 * ( kr[:, None, None] ** 2 \
+ kr[None, :, None] ** 2 \
+ kr[None, None, :] ** 2))
if verbose:
print('Allocating randoms...')
sys.stdout.flush()
deltar = np.zeros((nbins, nbins, nbins), dtype='float64')
fastmodules.allocate_gal_cic(deltar, ran.x, ran.y, ran.z, ran.we,
ran.size, self.xmin, self.ymin,
self.zmin, self.box, nbins, 1)
if verbose:
print('Smoothing...')
sys.stdout.flush()
# We do the smoothing via FFTs rather than scipy's gaussian_filter
# because if using several threads for pyfftw it is much faster this way
# (if only using 1 thread gains are negligible)
rho = deltar + 0.0j
fft_obj(input_array=rho, output_array=rhok)
fastmodules.mult_norm(rhok, rhok, norm)
ifft_obj(input_array=rhok, output_array=rho)
deltar = rho.real
else:
delta = self.delta
deltak = self.deltak
deltar = self.deltar
rho = self.rho
rhok = self.rhok
psi_x = self.psi_x
psi_y = self.psi_y
psi_z = self.psi_z
fft_obj = self.fft_obj
ifft_obj = self.ifft_obj
norm = self.norm
#fft_obj = pyfftw.FFTW(delta, delta, threads=self.nthreads, axes=[0, 1, 2])
#-- Allocate galaxies and randoms to grid with CIC method
#-- using new positions
if verbose:
print('Allocating galaxies in cells...')
sys.stdout.flush()
deltag = np.zeros((nbins, nbins, nbins), dtype='float64')
fastmodules.allocate_gal_cic(deltag, dat.newx, dat.newy, dat.newz,
dat.we, dat.size, self.xmin, self.ymin,
self.zmin, self.box, nbins, 1)
#deltag = self.allocate_gal_cic(dat)
if verbose:
print('Smoothing...')
sys.stdout.flush()
#deltag = gaussian_filter(deltag, smooth/binsize)
##-- Smoothing via FFTs
rho = deltag + 0.0j
fft_obj(input_array=rho, output_array=rhok)
fastmodules.mult_norm(rhok, rhok, norm)
ifft_obj(input_array=rhok, output_array=rho)
deltag = rho.real
if verbose:
print('Computing density fluctuations, delta...')
sys.stdout.flush()
# normalize using the randoms, avoiding possible divide-by-zero errors
fastmodules.normalize_delta_survey(delta, deltag, deltar, self.alpha,
self.ran_min)
del (deltag) # deltag no longer required anywhere
if verbose:
print('Fourier transforming delta field...')
sys.stdout.flush()
fft_obj(input_array=delta, output_array=delta)
## -- delta/k**2
k = fftfreq(self.nbins, d=binsize) * 2 * np.pi
fastmodules.divide_k2(delta, delta, k)
# now solve the basic building block: IFFT[-i k delta(k)/(b k^2)]
if verbose:
print('Inverse Fourier transforming to get psi...')
sys.stdout.flush()
fastmodules.mult_kx(deltak, delta, k, bias)
ifft_obj(input_array=deltak, output_array=psi_x)
fastmodules.mult_ky(deltak, delta, k, bias)
ifft_obj(input_array=deltak, output_array=psi_y)
fastmodules.mult_kz(deltak, delta, k, bias)
ifft_obj(input_array=deltak, output_array=psi_z)
# from grid values of Psi_est = IFFT[-i k delta(k)/(b k^2)], compute the values at the galaxy positions
if verbose:
print('Calculating shifts...')
sys.stdout.flush()
shift_x, shift_y, shift_z = self.get_shift(dat.newx, dat.newy,
dat.newz, psi_x.real,
psi_y.real, psi_z.real)
#-- for first loop need to approximately remove RSD component
#-- from Psi to speed up calculation
#-- first loop so want this on original positions (cp),
#-- not final ones (np) - doesn't actualy matter
if iloop == 0:
psi_dot_rhat = (shift_x*dat.x + \
shift_y*dat.y + \
shift_z*dat.z ) /dat.dist
shift_x -= beta / (1 + beta) * psi_dot_rhat * dat.x / dat.dist
shift_y -= beta / (1 + beta) * psi_dot_rhat * dat.y / dat.dist
shift_z -= beta / (1 + beta) * psi_dot_rhat * dat.z / dat.dist
#-- remove RSD from original positions (cp) of
#-- galaxies to give new positions (np)
#-- these positions are then used in next determination of Psi,
#-- assumed to not have RSD.
#-- the iterative procued then uses the new positions as
#-- if they'd been read in from the start
psi_dot_rhat = (shift_x * dat.x + shift_y * dat.y +
shift_z * dat.z) / dat.dist
dat.newx = dat.x + f * psi_dot_rhat * dat.x / dat.dist
dat.newy = dat.y + f * psi_dot_rhat * dat.y / dat.dist
dat.newz = dat.z + f * psi_dot_rhat * dat.z / dat.dist
if verbose:
print(
'Debug: first 10 x,y,z shifts and old and new observer distances'
)
for i in range(10):
oldr = np.sqrt(dat.x[i]**2 + dat.y[i]**2 + dat.z[i]**2)
newr = np.sqrt(dat.newx[i]**2 + dat.newy[i]**2 +
dat.newz[i]**2)
print('%.3f %.3f %.3f %.3f %.3f' %
(shift_x[i], shift_y[i], shift_z[i], oldr, newr))
self.deltar = deltar
self.delta = delta
self.deltak = deltak
self.rho = rho
self.rhok = rhok
self.psi_x = psi_x
self.psi_y = psi_y
self.psi_z = psi_z
self.fft_obj = fft_obj
self.ifft_obj = ifft_obj
self.norm = norm
#-- save wisdom
wisdom_file = "wisdom." + str(nbins) + "." + str(
self.nthreads) + '.npy'
if iloop == 0 and save_wisdom and not os.path.isfile(wisdom_file):
wisd = pyfftw.export_wisdom()
np.save(wisdom_file, wisd)
print('Wisdom saved at', wisdom_file)
#delta[w3] = 0.
#del(w)
#del(w2)
#del(w3)
del (deltag)
if verbose:
print('Fourier transforming delta field...')
fft_obj(input_array=delta, output_array=delta)
## -- delta/k**2
k = fftfreq(nbins, d=binsize) * 2 * np.pi
fastmodules.divide_k2(delta, delta, k)
if verbose:
print('Inverse Fourier transforming to get psi...')
fastmodules.mult_kx(deltak, delta, k, bias)
ifft_obj(input_array=deltak, output_array=psi_x)
fastmodules.mult_ky(deltak, delta, k, bias)
ifft_obj(input_array=deltak, output_array=psi_y)
fastmodules.mult_kz(deltak, delta, k, bias)
ifft_obj(input_array=deltak, output_array=psi_z)
# from grid values of Psi_est = IFFT[-i k delta(k)/(b k^2)], compute the values at the galaxy positions
if verbose:
print('Calculating shifts...')
#print('Shape of psi_x:',np.shape(psi_x))
shift_x, shift_y, shift_z = get_shift(dat.newx, dat.newy, dat.newz,
psi_x.real, psi_y.real, psi_z.real,
nbins, binsize)
dat.newx += shift_x
dat.newy += shift_y
dat.newz += shift_z