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