def test_export(self): forget_wisdom() before_wisdom = export_wisdom() self.generate_wisdom() after_wisdom = export_wisdom() self.compare(before_wisdom, after_wisdom)
def test_export(self): forget_wisdom() before_wisdom = export_wisdom() self.generate_wisdom() after_wisdom = export_wisdom() for n in range(0,2): self.assertNotEqual(before_wisdom[n], after_wisdom[n])
def store_plan_hints(filename, locking=True, reload_first=True): """Store data about the best FFT plans for this computer. FFT planning can take quite a while. After planning, the knowledge about the best plan for a given computer and given transform parameters can be written to disk so that the next time, planning can make use of that knowledge. Parameters: filename: file to write hints to. locking: if True, attempt to acquire an exclusive lock before writing which can otherwise cause problems if multiple processes are attempting to write to the same plan hints file. reload_first: if True, if the file exists, load the plan hints before storing them back. Safer in a multi-process setting where the hints may be written by a different process. """ filename = pathlib.Path(filename) if not filename.exists(): filename.touch() # can't open a file for read/write updating if it doesn't exist... with filename.open('r+b') as f: if locking: import fcntl fcntl.flock(f, fcntl.LOCK_EX) if reload_first: try: pyfftw.import_wisdom(pickle.load(f)) except: pass f.seek(0) pickle.dump(pyfftw.export_wisdom(), f) if locking: fcntl.flock(f, fcntl.LOCK_UN)
def compute_pow2_complex_wisdom(path, pow2=range(20), rigor='measure', threads=16): """ ??? If you plan with FFTW_PATIENT, it will automatically disable threads for sizes that don't benefit from parallelization. """ flags = [RIGOR_MAP[rigor], 'FFTW_DESTROY_INPUT'] wisdom_fnames = [] for pow2_i in pow2: N = 2**pow2_i x_input = pyfftw.empty_aligned(N, dtype='complex128') x_output = pyfftw.empty_aligned(N, dtype='complex128') for direction in ['forward', 'backward']: logger.info('building wisdom for complex 2**{} {}'.format( pow2_i, direction)) plan = pyfftw.FFTW(x_input, x_output, direction=DIRECTION_MAP[direction], flags=flags, threads=threads) wisdom_fnames.append( wisdom_fname(path, 'complex', pow2_i, threads, direction, rigor)) logger.info('writing to {}'.format(wisdom_fnames[-1])) with open(wisdom_fnames[-1], 'w') as fid: dump(pyfftw.export_wisdom(), fid, -1) pyfftw.forget_wisdom() return wisdom_fnames
def export_wisdom(self, wisdom_file='./fftw_wisdom.npy'): '''Parameters: ---------- wisdom_file: string,optional Wisdom filename ''' np.save(wisdom_file, pyfftw.export_wisdom())
def test_import(self): forget_wisdom() self.generate_wisdom() after_wisdom = export_wisdom() forget_wisdom() before_wisdom = export_wisdom() success = import_wisdom(after_wisdom) self.compare(before_wisdom, after_wisdom) self.assertEqual(success, tuple([x in _supported_types for x in ['64', '32', 'ld']]))
def _savewisdom(self, outfile): if outfile is None: return if outfile: pickle.dump(pyfftw.export_wisdom(), open(outfile, "wb"), protocol=0)
def save_wisdom(self): """Will overwrite.""" self._wisdom64, self._wisdom32, _ = pyfftw.export_wisdom() for name in ('wisdom32', 'wisdom64'): path = getattr(self, f"_{name}_path") with open(path, 'wb') as f: f.write(getattr(self, f"_{name}"))
def fftw_save_wisdom(filename=None): """ Save accumulated FFTW wisdom to a file By default this file will be in the user's astropy configuration directory. (Another location could be chosen - this is simple and works easily cross-platform.) Parameters ------------ filename : string, optional Filename to use (instead of the default, poppy_fftw_wisdom.txt) """ import os import astropy.config try: import pyfftw except: return # FFTW is not present, therefore this is a null op if filename is None: filename=os.path.join( astropy.config.get_config_dir(), "poppy_fftw_wisdom.txt") double, single, longdouble = pyfftw.export_wisdom() f = open(filename, 'w') f.write("# FFTW wisdom information saved by POPPY\n") f.write("# Do not edit this file by hand; that will likely break it.\n") f.write("# See http://hgomersall.github.io/pyFFTW/pyfftw/pyfftw.html?#wisdom-functions for more information\n") f.write('# the following three lines are the double, single, and longdouble wisdoms\n') f.write(double.replace('\n','\\n')+'\n') f.write(single.replace('\n','\\n')+'\n') f.write(longdouble.replace('\n','\\n')+'\n') _log.debug("FFTW wisdom saved to "+filename)
def __init__(self, n = 1364, N = 2048, iter0 = 138000, nfiles = 64, fft_threads = 32, out_threads = 4, src_dir = './', dst_dir = './', src_format = 'K{0:0>6}QNP{1:0>3}', dst_format = 'RMHD_{0}_t{1:0>4x}_z{2:0>7x}'): self.src_format = src_format self.dst_format = dst_format self.src_dir = src_dir self.dst_dir = dst_dir self.n = n self.N = N self.iter0 = iter0 self.nfiles = nfiles self.kdata = pyfftw.n_byte_align_empty( (self.N//2+1, 2, self.N, self.N), pyfftw.simd_alignment, dtype = np.complex64) self.rrdata = pyfftw.n_byte_align_empty( (self.N, self.N, self.N, 2), pyfftw.simd_alignment, dtype = np.float32) self.rzdata = pyfftw.n_byte_align_empty( ((self.N//8) * (self.N//8) * (self.N//8), 8*8*8*2), pyfftw.simd_alignment, dtype = np.float32) if type(self.dst_dir) == type([]): self.zdir = np.array( range(0, self.rzdata.shape[0], self.rzdata.shape[0] // len(self.dst_dir))) self.cubbies_per_file = self.rzdata.shape[0] // self.nfiles if (os.path.isfile('fftw_wisdom.pickle.gz')): pyfftw.import_wisdom( pickle.load(gzip.open('fftw_wisdom.pickle.gz', 'rb'))) print('about to initialize the fftw plan, which can take a while') self.plan = pyfftw.FFTW( self.kdata.transpose(3, 2, 0, 1), self.rrdata, axes = (0, 1, 2), direction = 'FFTW_BACKWARD', flags = ('FFTW_MEASURE', 'FFTW_DESTROY_INPUT'), threads = fft_threads) print('finalized fftw initialization') bla = pyfftw.export_wisdom() pickle.dump(bla, gzip.open('fftw_wisdom.pickle.gz', 'wb')) self.fft_threads = fft_threads self.out_threads = out_threads self.shuffle_lib = np.ctypeslib.load_library( 'libzshuffle.so', os.path.abspath(os.path.join( os.path.expanduser('~'), 'repos/RMHD_converter/C-shuffle'))) return None
def test_import(self): forget_wisdom() self.generate_wisdom() after_wisdom = export_wisdom() forget_wisdom() before_wisdom = export_wisdom() success = import_wisdom(after_wisdom) for n in range(0,2): self.assertNotEqual(before_wisdom[n], after_wisdom[n]) self.assertEqual(success, (True, True, True))
def save(cls): try: import pickle wisdom = pyfftw.export_wisdom() with pyfs.aopen(cls.paths[0], 'wb') as dst: pickle.dump(wisdom, dst) except IOError: pass
def test_import(self): forget_wisdom() self.generate_wisdom() after_wisdom = export_wisdom() forget_wisdom() before_wisdom = export_wisdom() success = import_wisdom(after_wisdom) self.compare(before_wisdom, after_wisdom) self.assertEqual( success, tuple([x in _supported_types for x in ['64', '32', 'ld']]))
def save_plan(self, plan_filename): if not (self.use_fftw): raise RuntimeError("FFTW is not used") T = pyfftw.export_wisdom() try: np.savez_compressed(plan_filename, plan=T) except IOError: print( "Error: could not save plan %s for some reason (possibly permission denied)" % plan_filename)
def __init__(self, n=1364, N=2048, iter0=138000, nfiles=64, fft_threads=32, out_threads=4, src_dir='./', dst_dir='./', src_format='K{0:0>6}QNP{1:0>3}', dst_format='RMHD_{0}_t{1:0>4x}_z{2:0>7x}'): self.src_format = src_format self.dst_format = dst_format self.src_dir = src_dir self.dst_dir = dst_dir self.n = n self.N = N self.iter0 = iter0 self.nfiles = nfiles self.kdata = pyfftw.n_byte_align_empty( (self.N // 2 + 1, 2, self.N, self.N), pyfftw.simd_alignment, dtype=np.complex64) self.rrdata = pyfftw.n_byte_align_empty((self.N, self.N, self.N, 2), pyfftw.simd_alignment, dtype=np.float32) self.rzdata = pyfftw.n_byte_align_empty( ((self.N // 8) * (self.N // 8) * (self.N // 8), 8 * 8 * 8 * 2), pyfftw.simd_alignment, dtype=np.float32) if type(self.dst_dir) == type([]): self.zdir = np.array( range(0, self.rzdata.shape[0], self.rzdata.shape[0] // len(self.dst_dir))) self.cubbies_per_file = self.rzdata.shape[0] // self.nfiles if (os.path.isfile('fftw_wisdom.pickle.gz')): pyfftw.import_wisdom( pickle.load(gzip.open('fftw_wisdom.pickle.gz', 'rb'))) print('about to initialize the fftw plan, which can take a while') self.plan = pyfftw.FFTW(self.kdata.transpose(3, 2, 0, 1), self.rrdata, axes=(0, 1, 2), direction='FFTW_BACKWARD', flags=('FFTW_MEASURE', 'FFTW_DESTROY_INPUT'), threads=fft_threads) print('finalized fftw initialization') bla = pyfftw.export_wisdom() pickle.dump(bla, gzip.open('fftw_wisdom.pickle.gz', 'wb')) self.fft_threads = fft_threads self.out_threads = out_threads self.shuffle_lib = np.ctypeslib.load_library( 'libzshuffle.so', os.path.abspath( os.path.join(os.path.expanduser('~'), 'repos/RMHD_converter/C-shuffle'))) return None
def export_wisdom(self): wis = pyfftw.export_wisdom() basedir = os.path.dirname(MyFFTW._WISDOM_FILE) if not os.path.exists(basedir): os.makedirs(basedir) with open(MyFFTW._WISDOM_FILE,"w") as f: pickle.dump(wis,f)
def save_wisdom(self,path_wisdom=None): if path_wisdom is not None: self.attrs['path_wisdom'] = path_wisdom path_wisdom = self.attrs['path_wisdom'].format(self.Nmesh,self.attrs['nthreads']) if not os.path.isfile(path_wisdom): self.logger.info('Saving wisdom to {}.'.format(path_wisdom)) wisdom = pyfftw.export_wisdom() file = open(path_wisdom,'w') json.dump(wisdom,file) file.close()
def set_wisdom(npixx, npixy): """ Run single 2d ifft like image to prep fftw wisdom in worker cache """ logger.info('Calculating FFT wisdom...') arr = pyfftw.empty_aligned((npixx, npixy), dtype='complex64', n=16) arr[:] = np.random.randn(*arr.shape) + 1j*np.random.randn(*arr.shape) fft_arr = pyfftw.interfaces.numpy_fft.ifft2(arr, auto_align_input=True, auto_contiguous=True, planner_effort='FFTW_MEASURE') return pyfftw.export_wisdom()
def save_wisdom(): """Save generated wisdom to file specified when configuring the project. """ if(wisdom_file is not None): wisdom = pyfftw.export_wisdom() with file(wisdom_file, mode="w") as f: for wisdom_bit in wisdom: f.write(wisdom_bit) else: raise Exception("Configured not to use any FFTW wisdom!")
def save_wisdom(): """Save generated wisdom to file specified when configuring the project. """ if (wisdom_file is not None): wisdom = pyfftw.export_wisdom() with file(wisdom_file, mode="w") as f: for wisdom_bit in wisdom: f.write(wisdom_bit) else: raise Exception("Configured not to use any FFTW wisdom!")
def _save_wisdom(): """ Save wisdom as 3 files. """ wisdom = pyfftw.export_wisdom() for filename, w in zip(FFTW_WISDOM_FILES, wisdom): try: os.remove(filename) except OSError: pass if len(w) == 0: continue with open(filename, 'w') as f: f.write(w)
def save_wisdom(wisdomfile): """ Save the acquired 'wisdom' generated by FFTW to file so that future initializations of FFTW will be faster. """ if wisdomfile is None: return if wisdomfile: pickle.dump(pyfftw.export_wisdom(), open(wisdomfile, 'wb'), protocol=-1)
def __init__(self): """ Setup the frequency vu """ self.bar = LightBarController() self.max_vol = 1 logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.INFO) self.pallett = RGB.REVERSE_RAINBOW pyfftw.interfaces.cache.enable() self.audio_16 = pyfftw.n_byte_align_empty(4096, 16, 'float64') # TODO: Play around with saving and changing planner effort self.fftw = pyfftw.builders.fft(self.audio_16, overwrite_input=True, planner_effort='FFTW_PATIENT', avoid_copy=True, threads=1, auto_align_input=True, auto_contiguous=True) logging.info("Wisdom: " + str(pyfftw.export_wisdom())) logging.warn("Updated wisdom, restart audio for accurate data") self.audio_16 = np.zeros_like(self.fftw.get_input_array())
def set_wisdom(npixx, npixy=None): """ Run single ifft to prep fftw wisdom in worker cache Supports 1d and 2d ifft. """ logger.info('Calculating FFT wisdom...') if npixy is not None: arr = pyfftw.empty_aligned((npixx, npixy), dtype='complex64', n=16) fft_arr = pyfftw.interfaces.numpy_fft.ifft2(arr, auto_align_input=True, auto_contiguous=True, planner_effort='FFTW_MEASURE') else: arr = pyfftw.empty_aligned((npixx), dtype='complex64', n=16) fft_arr = pyfftw.interfaces.numpy_fft.ifft(arr, auto_align_input=True, auto_contiguous=True, planner_effort='FFTW_MEASURE') return pyfftw.export_wisdom()
def hilbert(x, N=None): """ Compute the analytic signal, using the Hilbert transform. The transformation is done along the last axis by default. -> From scipy but replace fft with fftw Note: Inpute must be (N_channels x Time)! Note: Input array is destroyed! """ axis = 1 x = asarray(x) if iscomplexobj(x): raise ValueError("x must be real.") if N is None: N = x.shape[axis] if N <= 0: raise ValueError("N must be positive.") ax = [0] * len(x.shape) xshape, xdim = x.shape, x.ndim Xf = fft(x) del x h = zeros(N) if N % 2 == 0: h[0] = h[N // 2] = 1 h[1:N // 2] = 2 else: h[0] = 1 h[1:(N + 1) // 2] = 2 if len(xshape) > 1: ind = [newaxis] * xdim ind[axis] = slice(None) h = h[ind] x = ifft(Xf * h) del Xf try: cPickle.dump(pyfftw.export_wisdom(), open(cache, 'w')) except: print('Did not save wisdom cache') return x
def __init__(self, ellmat, filt_func=lambda ell: ell > 0, num_threads=4, flags_init=('FFTW_MEASURE', )): super(ffs_alm_pyFFTW, self).__init__(ellmat, filt_func=filt_func) # FIXME : This can be tricky in in hybrid MPI-OPENMP # Builds FFTW Wisdom : wisdom_fname = self.ell_mat.lib_dir + '/FFTW_wisdom_%s_%s.npy' % ( num_threads, ''.join(flags_init)) if not os.path.exists(wisdom_fname): print "++ ffs_alm_pyFFTW :: building and caching FFTW wisdom, this might take a little while..." if pbs.rank == 0: inpt = pyfftw.empty_aligned(self.ell_mat.shape, dtype='float64') oupt = pyfftw.empty_aligned(self.ell_mat.rshape, dtype='complex128') fft = pyfftw.FFTW(inpt, oupt, axes=(0, 1), direction='FFTW_FORWARD', flags=flags_init, threads=num_threads) ifft = pyfftw.FFTW(oupt, inpt, axes=(0, 1), direction='FFTW_BACKWARD', flags=flags_init, threads=num_threads) wisdom = pyfftw.export_wisdom() np.save(wisdom_fname, wisdom) del inpt, oupt, fft, ifft pbs.barrier() pyfftw.import_wisdom(np.load(wisdom_fname)) # print "++ ffs_alm_pyFFTW :: loaded widsom ", wisdom_fname self.flags = ( 'FFTW_WISDOM_ONLY', ) # This will make the code crash if arrays are not properly aligned. # self.flags = ('FFTW_MEASURE',) self.threads = num_threads
def compute_pow2_real_wisdom(path, pow2=range(20), rigor='measure', threads=16): """ ??? If you plan with FFTW_PATIENT, it will automatically disable threads for sizes that don't benefit from parallelization. """ flags = [RIGOR_MAP[rigor], 'FFTW_DESTROY_INPUT'] wisdom_fnames = [] for pow2_i in pow2: N = 2**pow2_i for direction in ['forward', 'backward']: logger.info('building wisdom for real 2**{} {}'.format(pow2_i, direction)) if direction == 'forward': x_input = pyfftw.empty_aligned(N, dtype='float64') x_output = pyfftw.empty_aligned(int(N // 2) + 1, dtype='complex128') else: x_output = pyfftw.empty_aligned(N, dtype='float64') x_input = pyfftw.empty_aligned(int(N // 2) + 1, dtype='complex128') plan = pyfftw.FFTW(x_input, x_output, direction=DIRECTION_MAP[direction], flags=flags, threads=threads) wisdom_fnames.append(wisdom_fname(path, 'real', pow2_i, threads, direction, rigor)) logger.info('writing to {}'.format(wisdom_fnames[-1])) with open(wisdom_fnames[-1], 'w') as fid: dump(pyfftw.export_wisdom(), fid, -1) pyfftw.forget_wisdom() return wisdom_fnames
def save_wisdom(filenames=default_filenames): """ Save the wisdom accumulated since the last load_wisdom, which occurs when pyfusion is imported Best to run typical codes before saving. To see the current state of wisdom !fftw-wisdom - or look at the file. I guess these are just timings for the routines tested e.g extensive tests of a 32 element complex forward out of place transform. !fftw-wisdom -x cof32 (fftw-3.3.3 fftw_wisdom #x458a31c8 #x92381c4c #x4f974889 #xcd46f97e (fftw_codelet_t2fv_4_avx 0 #x1040 #x1040 #x0 #x12dc3d8e #xe40293c9 #x508e7f21 #x18911bc9) (fftw_codelet_n2fv_8_avx 0 #x1040 #x1040 #x0 #xb916fb98 #xf394dae5 #xe9a593f6 #x4ce2ac3f) ) for in place, just (fftw-3.3.3 fftw_wisdom #x458a31c8 #x92381c4c #x4f974889 #xcd46f97e (fftw_codelet_n2fv_32_sse2 0 #x1040 #x1040 #x0 #x8ee86d9a #x3bf651fe #x73d5cbe4 #xfd6f3826) ) """ wisdom = pyfftw.export_wisdom() for (fn,w) in zip(filenames,wisdom): f = open(fn,'w') f.write(w) f.close()
def save_wisdom(filenames=default_filenames): """ Save the wisdom accumulated since the last load_wisdom, which occurs when pyfusion is imported Best to run typical codes before saving. To see the current state of wisdom !fftw-wisdom - or look at the file. I guess these are just timings for the routines tested e.g extensive tests of a 32 element complex forward out of place transform. !fftw-wisdom -x cof32 (fftw-3.3.3 fftw_wisdom #x458a31c8 #x92381c4c #x4f974889 #xcd46f97e (fftw_codelet_t2fv_4_avx 0 #x1040 #x1040 #x0 #x12dc3d8e #xe40293c9 #x508e7f21 #x18911bc9) (fftw_codelet_n2fv_8_avx 0 #x1040 #x1040 #x0 #xb916fb98 #xf394dae5 #xe9a593f6 #x4ce2ac3f) ) for in place, just (fftw-3.3.3 fftw_wisdom #x458a31c8 #x92381c4c #x4f974889 #xcd46f97e (fftw_codelet_n2fv_32_sse2 0 #x1040 #x1040 #x0 #x8ee86d9a #x3bf651fe #x73d5cbe4 #xfd6f3826) ) """ wisdom = pyfftw.export_wisdom() for (fn, w) in zip(filenames, wisdom): f = open(fn, 'w') f.write(w) f.close()
def cb(**kwargs): sig = kwargs['value'] cur = pv.curt.get() global pars if not q['update_conf'].empty(): print('updating...') q['update_conf'].get() pars = inittswa() try: plotdata, resultsdata = tswa(sig, cur, pars, calib) except Exception as e: showerror(title='analysis error', message=e.message) q['run'].put(False) q['update_plots'].put(plotdata) root.event_generate('<<update_plots>>', when='tail') q['update_results'].put(resultsdata) root.event_generate('<<update_results>>', when='tail') if not q['run'].empty(): cbvar.clear_callbacks() save('wisdom.npy', pyfftw.export_wisdom()) print('wisdom saved')
x = cv2.getOptimalDFTSize(rows) y = cv2.getOptimalDFTSize(cols) # learn 2D wisdom start1 = time.time() bb = pyfftw.zeros_aligned((x, y), dtype='float32', n=8) # numpy array storing the real-space data bf = pyfftw.zeros_aligned( (x, y // 2 + 1), dtype='complex64', n=8) # numpy array storing the Fourier-space data fft_object_b = pyfftw.FFTW(bb, bf, axes=(-2,-1), flags=('FFTW_MEASURE',), \ direction='FFTW_FORWARD',threads=mp.cpu_count()) fft_object_c = pyfftw.FFTW(bf, bb, axes=(-2,-1), flags=('FFTW_MEASURE',), \ direction='FFTW_BACKWARD',threads=mp.cpu_count()) end1 = time.time() # Save the learned 2D wisdom result to "wisdom" folder in three ".txt" files bb = pyfftw.export_wisdom() print(bb) Length_data = str((x, y)) file = open(dir_wisdom + Length_data + "x1.txt", "wb") file.write(bb[0]) file.close file = open(dir_wisdom + Length_data + "x2.txt", "wb") file.write(bb[1]) file.close file = open(dir_wisdom + Length_data + "x3.txt", "wb") file.write(bb[2]) file.close() print(end1 - start1, ' s')
def __init__(self, spatialGrid, kGrid, damping='default', FFTW_METHOD='FFTW_PATIENT', N_THREADS=6): """ Initialise an instance of a GPESolver. Parameters: spatialGrid: A tuple (x, y) representing the spatial grid that the simulation will be performed on. This grid should be scaled to units of the characteristic length defined in ParameterContainer. kGrid: A tuple (k_x, k_y) representing the k-space grid corresponding to the (x, y) grid. The scaling of this grid should correspond to that of the (x, y) grid. That is, it should be scaled to units of the inverse of the characterestic length defined in ParameterContainer. damping: The damping method to use in order to suppress the implicit periodic boundary conditions. Default is a tanh function that drops from 1.0 to 0 over the last 10% of each dimension. Presently, the only other option is "None", which means no damping, and hence periodic boundary conditions FFTW_METHOD: The method for FFTW to use when planning the transforms FFTW_PATIENT, FFTW_EXHAUSTIVE and FFTW_MEASURE will result in faster transforms, but may take a significant amount of time to plan N_THREADS: The number of threads for FFTW to use. Currently 2 threads seems to give the lowest time per step (8 core computer). Increasing this may greatly increase the time that FFTW requires to plan the transforms. """ # Load any existing wisdom try: wisdomFile = open(WISDOM_LOCATION, 'r+') importStatus = pyfftw.import_wisdom(json.load(wisdomFile)) print "Wisdom found" if not np.array(importStatus).all(): print "Wisdom not loaded correctly" # raise Warning("Wisdom not loaded correctly.") wisdomFile.close() except IOError: print "Wisdom not present" # raise Warning("Wisdom not present.") self.x, self.y = spatialGrid self.kx, self.ky = kGrid self.K = self.kx ** 2 + self.ky ** 2 # This is already scaled because we obtained it from the scaled grid. self.max_XY = np.abs(self.x[-1, -1]) self.N = self.x.shape[0] self.N_THREADS = N_THREADS self.time = 0 # TODO: Allow for rectangular grids. assert self.x.shape == self.y.shape, "Spatial grids are not the same\ shape" assert self.kx.shape == self.ky.shape, "k grids are not the same shape" assert self.x.shape == self.kx.shape, "Spatial grids are not the same\ shape as k grid" assert self.x.shape == (self.N, self.N), 'Grid must be square.' if damping == 'default': tanhDamping = UtilityFunctions.RadialTanhDamping(self.max_XY) # We can use the unscaled function here because max_XY is already # scaled. # A damping mask self.damping = tanhDamping.unscaledFunction()(self.x, self.y) # Set up fftw objects # Optimal alignment for this CPU self.al = pyfftw.simd_alignment self.psi = pyfftw.n_byte_align_empty((self.N, self.N), self.al, 'complex128') flag = FFTW_METHOD self.fft_object = pyfftw.FFTW(self.psi, self.psi, flags=[flag], axes=(0, 1), threads=N_THREADS) self.ifft_object = pyfftw.FFTW(self.psi, self.psi, flags=[flag], axes=(0, 1), threads=N_THREADS, direction='FFTW_BACKWARD') # Save pyfftw's newfound wisdom f = open(WISDOM_LOCATION, 'w+') json.dump(pyfftw.export_wisdom(), f) f.close()
def lens_glm_GLth_sym_timed(spin, dlm, glm, lmax_target, nband=16, facres=0, clm=None, olm=None, rotpol=True): """ Same as lens_alm but lens simultnously a North and South colatitude band, to make profit of the symmetries of the spherical harmonics. """ assert spin >= 0, spin times = {} t0 = time.time() tGL, wg = gauleg.get_xgwg(lmax_target + 2) times['GL points and weights'] = time.time() - t0 target_nt = 3**1 * 2**(11 + facres ) # on one hemisphere (0.87 arcmin spacing) th1s = np.arange(nband) * (np.pi * 0.5 / nband) th2s = np.concatenate((th1s[1:], [np.pi * 0.5])) Nt = target_nt / nband tGL = np.arccos(tGL) tGL = np.sort(tGL) wg = wg[np.argsort(tGL)] times['pol. rot.'] = 0. times['vtm2defl2ang'] = 0. times['vtmdefl'] = 0. def coadd_times(tim): for _k, _t in tim.iteritems(): if _k not in times: times[_k] = _t else: times[_k] += _t shapes = [] shapes_d = [] tGLNs = [] tGLSs = [] wgs = [] # Collects (Nt,Nphi) per band and prepare wisdom wisdomhash = str(lmax_target) + '_' + str(nband) + '_' + str(facres + 1000) + '.npy' assert os.path.exists( os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/') t0 = time.time() print "building and caching FFTW wisdom, this might take a while" for ib, th1, th2 in zip(range(nband), th1s, th2s): Np = get_Nphi(th1, th2, facres=facres, target_amin=60. * 90. / target_nt) # same spacing as theta grid Np_d = min(get_Nphi(th1, th2, target_amin=180. * 60. / lmax_target), 2 * lmax_target) #Equator point density pixN, = np.where((tGL >= th1) & (tGL <= th2)) pixS, = np.where((tGL >= (np.pi - th2)) & (tGL <= (np.pi - th1))) assert np.all(pixN[::-1] == len(tGL) - 1 - pixS), 'symmetry of GL points' shapes_d.append((len(pixN), Np_d)) shapes.append((Nt, Np)) tGLNs.append(tGL[pixN]) tGLSs.append(tGL[pixS]) wgs.append(np.concatenate([wg[pixN], wg[pixS]])) print "BAND %s in %s. deflection (%s x %s) pts " % (ib, nband, len(pixN), Np_d) print " interpolation (%s x %s) pts " % (Nt, Np) #==== For each block we have the following ffts: # (Np_d) complex to complex (deflection map) BACKWARD (vtm2map) # (Nt,Np) complex to complex (bicubic prefiltering) BACKWARD (vt2mmap) (4 threads) # (Nt) complex to complex (bicubic prefiltering) FORWARD (vt2map) # (Np_d) complex to complex FORWARD (map2vtm) # Could rather do a try with FFTW_WISDOM_ONLY if not os.path.exists( os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' + wisdomhash): a = pyfftw.empty_aligned(Np_d, dtype='complex128') b = pyfftw.empty_aligned(Np_d, dtype='complex128') fft = pyfftw.FFTW(a, b, direction='FFTW_FORWARD', threads=1) fft = pyfftw.FFTW(a, b, direction='FFTW_BACKWARD', threads=1) a = pyfftw.empty_aligned(Nt, dtype='complex128') b = pyfftw.empty_aligned(Nt, dtype='complex128') fft = pyfftw.FFTW(a, b, direction='FFTW_FORWARD', threads=1) a = pyfftw.empty_aligned((Nt, Np), dtype='complex128') b = pyfftw.empty_aligned((Nt, Np), dtype='complex128') fft = pyfftw.FFTW(a, b, direction='FFTW_BACKWARD', axes=(0, 1), threads=4) if not os.path.exists( os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' + wisdomhash): np.save( os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' + wisdomhash, pyfftw.export_wisdom()) pyfftw.import_wisdom( np.load( os.path.dirname(os.path.realpath(__file__)) + '/pyfftw_cache/' + wisdomhash)) shts.PYFFTWFLAGS = ['FFTW_WISDOM_ONLY'] times['pyfftw_caches'] = time.time() - t0 print "Total number of interpo points: %s = %s ** 2" % (np.sum( [np.prod(s) for s in shapes]), np.sqrt(1. * np.sum([np.prod(s) for s in shapes]))) print "Total number of deflect points: %s = %s ** 2" % (np.sum([ np.prod(s) for s in shapes_d ]), np.sqrt(1. * np.sum([np.prod(s) for s in shapes_d]))) glmout = np.zeros(shts.util.lmax2nlm(lmax_target), dtype=np.complex) clmout = np.zeros(shts.util.lmax2nlm(lmax_target), dtype=np.complex) for ib, th1, th2 in zip(range(nband), th1s, th2s): Nt_d, Np_d = shapes_d[ib] Nt, Np = shapes[ib] t0 = time.time() vtm_def = shts.vlm2vtm_sym(1, _th2colat(tGLNs[ib]), shts.util.alm2vlm(dlm, clm=olm)) times['vtmdefl'] += time.time() - t0 #==== gettting deflected positions # NB: forward slice to keep theta -> pi - theta correspondance. t0 = time.time() dmapN = shts.vtm2map(1, vtm_def[:Nt_d, :], Np_d).flatten() dmapS = shts.vtm2map(1, vtm_def[slice(Nt_d, 2 * Nt_d), :], Np_d).flatten() told = np.outer(tGLNs[ib], np.ones(Np_d)).flatten() phiold = np.outer(np.ones(Nt_d), np.arange(Np_d) * (2. * np.pi / Np_d)).flatten() tnewN, phinewN = _buildangles((told, phiold), dmapN.real, dmapN.imag) tnewS, phinewS = _buildangles(((np.pi - told)[::-1], phiold), dmapS.real, dmapS.imag) del vtm_def times['vtm2defl2ang'] += time.time() - t0 #===== Adding a 10 pixels buffer for new angles to be safely inside interval. # th1,th2 is mapped onto pi - th2,pi -th1 so we need to make sure to cover both buffers matnewN = np.max(tnewN) mitnewN = np.min(tnewN) matnewS = np.max(tnewS) mitnewS = np.min(tnewS) buffN = 10 * (matnewN - mitnewN) / (Nt - 1) / (1. - 2. * 10. / (Nt - 1)) buffS = 10 * (matnewS - mitnewS) / (Nt - 1) / (1. - 2. * 10. / (Nt - 1)) _thup = min(np.pi - (matnewS + buffS), mitnewN - buffN) _thdown = max(np.pi - (mitnewS - buffS), matnewN + buffN) #==== these are the theta and limits. It is ok to go negative or > 180 dphi_patch = (2. * np.pi) / Np * max(np.sin(_thup), np.sin(_thdown)) dth_patch = (_thdown - _thup) / (Nt - 1) print 'input t1,t2 %.3f %.3f in degrees' % (_thup / np.pi * 180, _thdown / np.pi * 180.) print 'North %.3f and South %.3f buffers in amin' % ( buffN / np.pi * 180 * 60, buffS / np.pi * 180. * 60.) print "cell (theta,phi) in amin (%.3f,%.3f)" % ( dth_patch / np.pi * 60. * 180, dphi_patch / np.pi * 60. * 180) if spin == 0: lenN, lenS, tim = lens_band_sym_timed(glm, _thup, _thdown, Nt, (tnewN, phinewN), (tnewS, phinewS), Nphi=Np) ret = np.zeros((2 * Nt_d, Np_d), dtype=complex) ret[:Nt_d, :] = lenN.reshape((Nt_d, Np_d)) ret[Nt_d:, :] = lenS.reshape((Nt_d, Np_d)) vtm = shts.map2vtm(spin, lmax_target, ret) glmout -= shts.vtm2tlm_sym( np.concatenate([tGLNs[ib], tGLSs[ib]]), vtm * np.outer(wgs[ib], np.ones(vtm.shape[1]))) else: assert 0, 'fix this' lenNR, lenNI, lenSR, lenSI, tim = gclm2lensmap_symband_timed( spin, glm, _thup, _thdown, Nt, (tnewN, phinewN), (tnewS, phinewS), Nphi=Nphi, clm=clm) retN = (lenNR + 1j * lenNI).reshape((len(pixN), Np_d)) retS = (lenSR + 1j * lenSI).reshape((len(pixN), Np_d)) glm, clm = shts.util.vlm2alm( shts.vtm2vlm(spin, tGL, vtm * np.outer(wg, np.ones(vtm.shape[1])))) t0 = time.time() if rotpol and spin > 0: ret[pixN, :] *= polrot(spin, retN.flatten(), tnewN, dmapN.real, dmapN.imag) ret[pixS, :] *= polrot(spin, retS.flatten(), tnewS, dmapS.real, dmapS.imag) times['pol. rot.'] += time.time() - t0 coadd_times(tim) t0 = time.time() print "STATS for lmax tlm %s lmax dlm %s" % (hp.Alm.getlmax( glm.size), hp.Alm.getlmax(dlm.size)) tot = 0. for _k, _t in times.iteritems(): print '%20s: %.2f' % (_k, _t) tot += _t print "%20s: %2.f sec." % ('tot', tot) return glmout, clmout, ret
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)
plt.legend(loc='upper left') plt.xlabel('time $t$ (a.u.)') plt.subplot(132) plt.title("Ehrenfest 2") plt.plot(times, np.gradient(quant_sys.P_average, dt), 'r-', label='$d\\langle p \\rangle/dt$') plt.plot( times, quant_sys.P_average_RHS, 'b--', label='$\\langle -\\partial V/\\partial x + \\gamma p \\rangle$' ) plt.legend(loc='upper left') plt.xlabel('time $t$ (a.u.)') plt.subplot(133) plt.title('Hamiltonian') plt.plot(times, quant_sys.hamiltonian_average) plt.xlabel('time $t$ (a.u.)') plt.show() ################################################################# # # Save FFTW wisdom # ################################################################# with open('fftw_wisdom', 'wb') as f: pickle.dump(pyfftw.export_wisdom(), f)
def save_wisdom(self): with open(_wisdom_filename(), 'wb') as fout: pickle.dump(pyfftw.export_wisdom(), fout, protocol=pickle.HIGHEST_PROTOCOL)
def __init__(self, **kwargs): """ The following parameters must be specified X_gridDIM - specifying the grid size X_amplitude - maximum value of the coordinates V - potential energy (as a string to be evaluated by numexpr) K - momentum dependent part of the hamiltonian (as a string to be evaluated by numexpr) A - a coordinate dependent Lindblad dissipator (as a string to be evaluated by numexpr) RHS_P_A (optional) -- the correction to the second Ehrenfest theorem due to A B - a momentum dependent Lindblad dissipator (as a string to be evaluated by numexpr) RHS_X_B (optional) -- the correction to the first Ehrenfest theorem due to B diff_V (optional) -- the derivative of the potential energy for the Ehrenfest theorem calculations diff_K (optional) -- the derivative of the kinetic energy for the Ehrenfest theorem calculations t (optional) - initial value of time abs_boundary (optional) -- absorbing boundary (as a string to be evaluated by numexpr) """ # save all attributes for name, value in kwargs.items(): # if the value supplied is a function, then dynamically assign it as a method; # otherwise bind it a property if isinstance(value, FunctionType): setattr(self, name, MethodType(value, self, self.__class__)) else: setattr(self, name, value) # Check that all attributes were specified try: # make sure self.X_gridDIM has a value of power of 2 assert 2 ** int(np.log2(self.X_gridDIM)) == self.X_gridDIM, \ "A value of the grid size (X_gridDIM) must be a power of 2" except AttributeError: raise AttributeError("Grid size (X_gridDIM) was not specified") try: self.X_amplitude except AttributeError: raise AttributeError("Coordinate range (X_amplitude) was not specified") try: self.V except AttributeError: raise AttributeError("Potential energy (V) was not specified") try: self.K except AttributeError: raise AttributeError("Momentum dependence (K) was not specified") try: self.A except AttributeError: self.A = self.RHS_P_A = "0." print("Warning: Coordinate dependent Lindblad dissipator (A) was not specified so it is set to zero") try: self.B except AttributeError: self.B = self.RHS_X_B = "0." print("Warning: Momentum dependent Lindblad dissipator (B) was not specified so it is set to zero") try: self.dt except AttributeError: raise AttributeError("Time-step (dt) was not specified") try: self.t except AttributeError: print("Warning: Initial time (t) was not specified, thus it is set to zero.") self.t = 0. try: self.abs_boundary except AttributeError: print("Warning: Absorbing boundary (abs_boundary) was not specified, thus it is turned off") self.abs_boundary = "1." ######################################################################################## # # Initialize Fourier transform for efficient calculations # ######################################################################################## # Load FFTW wisdom if saved before try: with open('fftw_wisdom', 'rb') as f: pyfftw.import_wisdom(pickle.load(f)) print("\nFFTW wisdom has been loaded\n") except IOError: pass # allocate the array for density matrix self.rho = pyfftw.empty_aligned((self.X_gridDIM, self.X_gridDIM), dtype=np.complex) # FFTW settings to achive good performace. For details see # https://hgomersall.github.io/pyFFTW/pyfftw/pyfftw.html#pyfftw.FFTW fftw_flags = ('FFTW_MEASURE','FFTW_DESTROY_INPUT') # how many threads to use for parallelized calculation of FFT. # Use the same number of threads as in numexpr fftw_nthreads = ne.nthreads # Create plan to pefrom FFT over the zeroth axis. It is equivalent to # fftpack.fft(self.rho, axis=0, overwrite_x=True) self.rho_fft_ax0 = pyfftw.FFTW( self.rho, self.rho, axes=(0,), direction='FFTW_FORWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_fft_ax1 = pyfftw.FFTW( self.rho, self.rho, axes=(1,), direction='FFTW_FORWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_ifft_ax0 = pyfftw.FFTW( self.rho, self.rho, axes=(0,), direction='FFTW_BACKWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_ifft_ax1 = pyfftw.FFTW( self.rho, self.rho, axes=(1,), direction='FFTW_BACKWARD', flags=fftw_flags, threads=fftw_nthreads ) # Save FFTW wisdom with open('fftw_wisdom', 'wb') as f: pickle.dump(pyfftw.export_wisdom(), f) ######################################################################################## # get coordinate step size self.dX = 2. * self.X_amplitude / self.X_gridDIM # generate coordinate range k = np.arange(self.X_gridDIM) self.k = k[:, np.newaxis] self.k_prime = k[np.newaxis, :] X = (k - self.X_gridDIM / 2) * self.dX self.X = X[:, np.newaxis] self.X_prime = X[np.newaxis, :] # generate momentum range self.dP = np.pi / self.X_amplitude P = (k - self.X_gridDIM / 2) * self.dP self.P = P[:, np.newaxis] self.P_prime = P[np.newaxis, :] # allocate an axillary array needed for propagation self.expV = np.zeros_like(self.rho) # construct the coordinate dependent phase containing the dissipator as well as coherent propagator phase_X = "1j * (({V_X_prime}) - ({V_X})) " \ "+ ({A_X}) * conj({A_X_prime}) - 0.5 * abs({A_X}) ** 2 - 0.5 * abs({A_X_prime}) ** 2".format( V_X_prime=self.V.format(X="X_prime"), V_X=self.V.format(X="X"), A_X_prime=self.A.format(X="X_prime"), A_X=self.A.format(X="X"), ) # numexpr code to calculate (-)**(k + k_prime) * exp(0.5 * dt * F) self.code_expV = "(%s) * (%s) * (-1) ** (k + k_prime) * exp(0.5 * dt * (%s))" % ( self.abs_boundary.format(X="X"), self.abs_boundary.format(X="X_prime"), phase_X ) # construct the coordinate dependent phase containing the dissipator as well as coherent propagator phase_P = "1j * (({K_P_prime}) - ({K_P})) " \ "+ ({B_P}) * conj({B_P_prime}) - 0.5 * abs({B_P}) ** 2 - 0.5 * abs({B_P_prime}) ** 2".format( K_P_prime=self.K.format(P="P_prime"), K_P=self.K.format(P="P"), B_P_prime=self.B.format(P="P_prime"), B_P=self.B.format(P="P"), ) # numexpr code to calculate rho * exp(1j * dt * K) self.code_expK = "rho * exp(dt * (%s))" % phase_P # Check whether the necessary terms are specified to calculate the first-order Ehrenfest theorems try: # Allocate a copy of the wavefunction for storing the density matrix in the momentum representation self.rho_p = pyfftw.empty_aligned(self.rho.shape, dtype=self.rho.dtype) # Create FFT plans to operate on self.rho_p self.rho_p_fft_ax0 = pyfftw.FFTW( self.rho_p, self.rho_p, axes=(0,), direction='FFTW_FORWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_p_ifft_ax1 = pyfftw.FFTW( self.rho_p, self.rho_p, axes=(1,), direction='FFTW_BACKWARD', flags=fftw_flags, threads=fftw_nthreads ) # numexpr codes to calculate the First Ehrenfest theorems self.code_V_average = "sum((%s) * density)" % self.V.format(X="X") self.code_K_average = "sum((%s) * density)" % self.K.format(P="P") self.code_X_average = "sum(X * density)" self.code_P_average = "sum(P * density)" self.code_P_average_RHS = "sum((-(%s) + (%s)) * density)" % ( self.diff_V.format(X="X"), self.RHS_P_A.format(X="X") ) self.code_X_average_RHS = "sum(((%s) + (%s)) * density)" % ( self.diff_K.format(P="P"), self.RHS_X_B.format(P="P") ) # Lists where the expectation values of X and P self.X_average = [] self.P_average = [] # Lists where the right hand sides of the Ehrenfest theorems for X and P self.X_average_RHS = [] self.P_average_RHS = [] # List where the expectation value of the Hamiltonian will be calculated self.hamiltonian_average = [] # Flag requesting tha the Ehrenfest theorem calculations self.isEhrenfest = True except AttributeError: # Since self.diff_V and self.diff_K are not specified, # the first Ehrenfest theorem will not be calculated self.isEhrenfest = False
def pyFFTWPlanner( realMage, fouMage=None, wisdomFile = None, effort = 'FFTW_MEASURE', n_threads = None, doForward = True, doReverse = True ): """ Appends an FFTW plan for the given realMage to a text file stored in the same directory as RAMutil, which can then be loaded in the future with pyFFTWLoadWisdom. NOTE: realMage should be typecast to 'complex64' normally. NOTE: planning pickle files are hardware dependant, so don't copy them from one machine to another. wisdomFile allows you to specify a .pkl file with the wisdom tuple written to it. The wisdomFile is never updated, whereas the default wisdom _is_ updated with each call. For multiprocessing, it's important to let FFTW generate its plan from an ideal processor state. TODO: implement real, half-space fourier transforms rfft2 and irfft2 as built """ import pyfftw import pickle import os.path from multiprocessing import cpu_count utilpath = os.path.dirname(os.path.realpath(__file__)) # First import whatever we already have if wisdomFile is None: wisdomFile = os.path.join( utilpath, "pyFFTW_wisdom.pkl" ) if os.path.isfile(wisdomFile): try: fh = open( wisdomFile, 'rb') except: print( "Util: pyFFTW wisdom plan file: " + str(wisdomFile) + " invalid/unreadable" ) try: pyfftw.import_wisdom( pickle.load( fh ) ) except: # THis is not normally a problem, it might be empty? print( "Util: pickle failed to import FFTW wisdom" ) pass try: fh.close() except: pass else: # Touch the file os.umask(0000) # Everyone should be able to delete scratch files with open( wisdomFile, 'wb') as fh: pass # I think the fouMage array has to be smaller to do the real -> complex FFT? if fouMage is None: if realMage.dtype.name == 'float32': print( "pyFFTW is recommended to work on purely complex data" ) fouShape = realMage.shape fouShape.shape[-1] = realMage.shape[-1]//2 + 1 fouDtype = 'complex64' fouMage = np.empty( fouShape, dtype=fouDtype ) elif realMage.dtype.name == 'float64': print( "pyFFTW is recommended to work on purely complex data" ) fouShape = realMage.shape fouShape.shape[-1] = realMage.shape[-1]//2 + 1 fouDtype = 'complex128' fouMage = np.empty( fouShape, dtype=fouDtype ) else: # Assume dtype is complexXX fouDtype = realMage.dtype.name fouMage = np.zeros( realMage.shape, dtype=fouDtype ) if n_threads is None: n_threads = cpu_count() print( "FFTW using " + str(n_threads) + " threads" ) if bool(doForward): #print( "Planning forward pyFFTW for shape: " + str( realMage.shape ) ) FFT2 = pyfftw.builders.fft2( realMage, planner_effort=effort, threads=n_threads, auto_align_input=True ) else: FFT2 = None if bool(doReverse): #print( "Planning reverse pyFFTW for shape: " + str( realMage.shape ) ) IFFT2 = pyfftw.builders.ifft2( fouMage, planner_effort=effort, threads=n_threads, auto_align_input=True ) else: IFFT2 = None # Setup so that we can call .execute on each one without re-copying arrays # if FFT2 is not None and IFFT2 is not None: # FFT2.update_arrays( FFT2.get_input_array(), IFFT2.get_input_array() ) # IFFT2.update_arrays( IFFT2.get_input_array(), FFT2.get_input_array() ) # Something is different in the builders compared to FFTW directly. # Can also repeat this for pyfftw.builders.rfft2 and .irfft2 if desired, but # generally it seems slower. # Opening a file for writing is supposed to truncate it # if bool(savePlan): #if wisdomFile is None: # with open( utilpath + "/pyFFTW_wisdom.pkl", 'wb') as fh: with open( wisdomFile, 'wb' ) as fh: pickle.dump( pyfftw.export_wisdom(), fh ) return FFT2, IFFT2
def save_wisdom(self, wisdom_file=DEFAULT_WISDOM_FILE): """Save FFTW wisdom to file.""" with open(wisdom_file, 'wb') as f: cPickle.dump(pyfftw.export_wisdom(), f, protocol=cPickle.HIGHEST_PROTOCOL)
def pyfftw_call(array_in, array_out, direction='forward', axes=None, halfcomplex=False, **kwargs): """Calculate the DFT with pyfftw. The discrete Fourier (forward) transform calcuates the sum:: f_hat[k] = sum_j( f[j] * exp(-2*pi*1j * j*k/N) ) where the summation is taken over all indices ``j = (j[0], ..., j[d-1])`` in the range ``0 <= j < N`` (component-wise), with ``N`` being the shape of the input array. The output indices ``k`` lie in the same range, except for half-complex transforms, where the last axis ``i`` in ``axes`` is shortened to ``0 <= k[i] < floor(N[i]/2) + 1``. In the backward transform, sign of the the exponential argument is flipped. Parameters ---------- array_in : `numpy.ndarray` Array to be transformed array_out : `numpy.ndarray` Output array storing the transformed values, may be aliased with ``array_in``. direction : {'forward', 'backward'}, optional Direction of the transform axes : int or sequence of ints, optional Dimensions along which to take the transform. ``None`` means using all axes and is equivalent to ``np.arange(ndim)``. halfcomplex : bool, optional If ``True``, calculate only the negative frequency part along the last axis. If ``False``, calculate the full complex FFT. This option can only be used with real input data. Other Parameters ---------------- fftw_plan : ``pyfftw.FFTW``, optional Use this plan instead of calculating a new one. If specified, the options ``planning_effort``, ``planning_timelimit`` and ``threads`` have no effect. planning_effort : str, optional Flag for the amount of effort put into finding an optimal FFTW plan. See the `FFTW doc on planner flags <http://www.fftw.org/fftw3_doc/Planner-Flags.html>`_. Available options: {'estimate', 'measure', 'patient', 'exhaustive'} Default: 'estimate' planning_timelimit : float or ``None``, optional Limit planning time to roughly this many seconds. Default: ``None`` (no limit) threads : int, optional Number of threads to use. Default: Number of CPUs if the number of data points is larger than 4096, else 1. normalise_idft : bool, optional If ``True``, the result of the backward transform is divided by ``1 / N``, where ``N`` is the total number of points in ``array_in[axes]``. This ensures that the IDFT is the true inverse of the forward DFT. Default: ``False`` import_wisdom : filename or file handle, optional File to load FFTW wisdom from. If the file does not exist, it is ignored. export_wisdom : filename or file handle, optional File to append the accumulated FFTW wisdom to Returns ------- fftw_plan : ``pyfftw.FFTW`` The plan object created from the input arguments. It can be reused for transforms of the same size with the same data types. Note that reuse only gives a speedup if the initial plan used a planner flag other than ``'estimate'``. If ``fftw_plan`` was specified, the returned object is a reference to it. Notes ----- * The planning and direction flags can also be specified as capitalized and prepended by ``'FFTW_'``, i.e. in the original FFTW form. * For a ``halfcomplex`` forward transform, the arrays must fulfill ``array_out.shape[axes[-1]] == array_in.shape[axes[-1]] // 2 + 1``, and vice versa for backward transforms. * All planning schemes except ``'estimate'`` require an internal copy of the input array but are often several times faster after the first call (measuring results are cached). Typically, 'measure' is a good compromise. If you cannot afford the copy, use ``'estimate'``. * If a plan is provided via the ``fftw_plan`` parameter, no copy is needed internally. """ import pickle if not array_in.flags.aligned: raise ValueError('input array not aligned') if not array_out.flags.aligned: raise ValueError('output array not aligned') if axes is None: axes = tuple(range(array_in.ndim)) axes = normalized_axes_tuple(axes, array_in.ndim) direction = _flag_pyfftw_to_odl(direction) fftw_plan_in = kwargs.pop('fftw_plan', None) planning_effort = _flag_pyfftw_to_odl( kwargs.pop('planning_effort', 'estimate') ) planning_timelimit = kwargs.pop('planning_timelimit', None) threads = kwargs.pop('threads', None) normalise_idft = kwargs.pop('normalise_idft', False) wimport = kwargs.pop('import_wisdom', '') wexport = kwargs.pop('export_wisdom', '') # Cast input to complex if necessary array_in_copied = False if is_real_dtype(array_in.dtype) and not halfcomplex: # Need to cast array_in to complex dtype array_in = array_in.astype(complex_dtype(array_in.dtype)) array_in_copied = True # Do consistency checks on the arguments _pyfftw_check_args(array_in, array_out, axes, halfcomplex, direction) # Import wisdom if possible if wimport: try: with open(wimport, 'rb') as wfile: wisdom = pickle.load(wfile) except IOError: wisdom = [] except TypeError: # Got file handle wisdom = pickle.load(wimport) if wisdom: pyfftw.import_wisdom(wisdom) # Copy input array if it hasn't been done yet and the planner is likely # to destroy it. If we already have a plan, we don't have to worry. planner_destroys = _pyfftw_destroys_input( [planning_effort], direction, halfcomplex, array_in.ndim) must_copy_array_in = fftw_plan_in is None and planner_destroys if must_copy_array_in and not array_in_copied: plan_arr_in = np.empty_like(array_in) flags = [_flag_odl_to_pyfftw(planning_effort), 'FFTW_DESTROY_INPUT'] else: plan_arr_in = array_in flags = [_flag_odl_to_pyfftw(planning_effort)] if fftw_plan_in is None: if threads is None: if plan_arr_in.size <= 4096: # Trade-off wrt threading overhead threads = 1 else: threads = cpu_count() fftw_plan = pyfftw.FFTW( plan_arr_in, array_out, direction=_flag_odl_to_pyfftw(direction), flags=flags, planning_timelimit=planning_timelimit, threads=threads, axes=axes) else: fftw_plan = fftw_plan_in fftw_plan(array_in, array_out, normalise_idft=normalise_idft) if wexport: try: with open(wexport, 'ab') as wfile: pickle.dump(pyfftw.export_wisdom(), wfile) except TypeError: # Got file handle pickle.dump(pyfftw.export_wisdom(), wexport) return fftw_plan
def save_wisdom(filenames=default_filenames): wisdom = pyfftw.export_wisdom() for (fn,w) in zip(filenames,wisdom): f = open(fn,'w') f.write(w) f.close()
def save_wisdom(self, wisdom_file=DEFAULT_WISDOM_FILE): """Save FFTW wisdom to file.""" with open(wisdom_file, 'wb') as f: cPickle.dump( pyfftw.export_wisdom(), f, protocol=cPickle.HIGHEST_PROTOCOL)
def __init__(self, **kwargs): """ The following parameters must be specified X_gridDIM - specifying the grid size X_amplitude - maximum value of the coordinates V - potential energy (as a string to be evaluated by numexpr) K - momentum dependent part of the hamiltonian (as a string to be evaluated by numexpr) A - a coordinate dependent Lindblad dissipator (as a string to be evaluated by numexpr) RHS_P_A (optional) -- the correction to the second Ehrenfest theorem due to A B - a momentum dependent Lindblad dissipator (as a string to be evaluated by numexpr) RHS_X_B (optional) -- the correction to the first Ehrenfest theorem due to B diff_V (optional) -- the derivative of the potential energy for the Ehrenfest theorem calculations diff_K (optional) -- the derivative of the kinetic energy for the Ehrenfest theorem calculations t (optional) - initial value of time abs_boundary (optional) -- absorbing boundary (as a string to be evaluated by numexpr) """ # save all attributes for name, value in kwargs.items(): # if the value supplied is a function, then dynamically assign it as a method; if isinstance(value, FunctionType): setattr(self, name, MethodType(value, self)) # otherwise bind it as a property else: setattr(self, name, value) # Check that all attributes were specified try: # make sure self.X_gridDIM has a value of power of 2 assert 2 ** int(np.log2(self.X_gridDIM)) == self.X_gridDIM, \ "A value of the grid size (X_gridDIM) must be a power of 2" except AttributeError: raise AttributeError("Grid size (X_gridDIM) was not specified") try: self.X_amplitude except AttributeError: raise AttributeError("Coordinate range (X_amplitude) was not specified") try: self.V except AttributeError: raise AttributeError("Potential energy (V) was not specified") try: self.K except AttributeError: raise AttributeError("Momentum dependence (K) was not specified") try: self.A except AttributeError: self.A = self.RHS_P_A = "0." warnings.warn("coordinate dependent Lindblad dissipator (A) was not specified so it is set to zero") try: self.B except AttributeError: self.B = self.RHS_X_B = "0." warnings.warn("momentum dependent Lindblad dissipator (B) was not specified so it is set to zero") try: self.dt except AttributeError: raise AttributeError("Time-step (dt) was not specified") try: self.t except AttributeError: warnings.warn("initial time (t) was not specified, thus it is set to zero.") self.t = 0. try: self.abs_boundary except AttributeError: warnings.warn("absorbing boundary (abs_boundary) was not specified, thus it is turned off") self.abs_boundary = "1." ######################################################################################## # # Initialize Fourier transform for efficient calculations # ######################################################################################## # Load FFTW wisdom if saved before try: with open('fftw_wisdom', 'rb') as f: pyfftw.import_wisdom(pickle.load(f)) print("\nFFTW wisdom has been loaded\n") except IOError: pass # allocate the array for density matrix self.rho = pyfftw.empty_aligned((self.X_gridDIM, self.X_gridDIM), dtype=np.complex) # FFTW settings to achive good performace. For details see # https://hgomersall.github.io/pyFFTW/pyfftw/pyfftw.html#pyfftw.FFTW fftw_flags = ('FFTW_MEASURE','FFTW_DESTROY_INPUT') # how many threads to use for parallelized calculation of FFT. # Use the same number of threads as in numexpr fftw_nthreads = ne.nthreads # Create plan to pefrom FFT over the zeroth axis. It is equivalent to # fftpack.fft(self.rho, axis=0, overwrite_x=True) self.rho_fft_ax0 = pyfftw.FFTW( self.rho, self.rho, axes=(0,), direction='FFTW_FORWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_fft_ax1 = pyfftw.FFTW( self.rho, self.rho, axes=(1,), direction='FFTW_FORWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_ifft_ax0 = pyfftw.FFTW( self.rho, self.rho, axes=(0,), direction='FFTW_BACKWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_ifft_ax1 = pyfftw.FFTW( self.rho, self.rho, axes=(1,), direction='FFTW_BACKWARD', flags=fftw_flags, threads=fftw_nthreads ) # Save FFTW wisdom with open('fftw_wisdom', 'wb') as f: pickle.dump(pyfftw.export_wisdom(), f) ######################################################################################## # get coordinate step size self.dX = 2. * self.X_amplitude / self.X_gridDIM # generate coordinate range k = np.arange(self.X_gridDIM) self.k = k[:, np.newaxis] self.k_prime = k[np.newaxis, :] X = (k - self.X_gridDIM / 2) * self.dX self.X = X[:, np.newaxis] self.X_prime = X[np.newaxis, :] # generate momentum range self.dP = np.pi / self.X_amplitude P = (k - self.X_gridDIM / 2) * self.dP self.P = P[:, np.newaxis] self.P_prime = P[np.newaxis, :] # allocate an axillary array needed for propagation self.expV = np.zeros_like(self.rho) # construct the coordinate dependent phase containing the dissipator as well as coherent propagator phase_X = "1j * (({V_X_prime}) - ({V_X})) " \ "+ ({A_X}) * conj({A_X_prime}) - 0.5 * abs({A_X}) ** 2 - 0.5 * abs({A_X_prime}) ** 2".format( V_X_prime=self.V.format(X="X_prime"), V_X=self.V.format(X="X"), A_X_prime=self.A.format(X="X_prime"), A_X=self.A.format(X="X"), ) # numexpr code to calculate (-)**(k + k_prime) * exp(0.5 * dt * F) self.code_expV = "(%s) * (%s) * (-1) ** (k + k_prime) * exp(0.5 * dt * (%s))" % ( self.abs_boundary.format(X="X"), self.abs_boundary.format(X="X_prime"), phase_X ) # construct the coordinate dependent phase containing the dissipator as well as coherent propagator phase_P = "1j * (({K_P_prime}) - ({K_P})) " \ "+ ({B_P}) * conj({B_P_prime}) - 0.5 * abs({B_P}) ** 2 - 0.5 * abs({B_P_prime}) ** 2".format( K_P_prime=self.K.format(P="P_prime"), K_P=self.K.format(P="P"), B_P_prime=self.B.format(P="P_prime"), B_P=self.B.format(P="P"), ) # numexpr code to calculate rho * exp(1j * dt * K) self.code_expK = "rho * exp(dt * (%s))" % phase_P # Check whether the necessary terms are specified to calculate the first-order Ehrenfest theorems try: # Allocate a copy of the wavefunction for storing the density matrix in the momentum representation self.rho_p = pyfftw.empty_aligned(self.rho.shape, dtype=self.rho.dtype) # Create FFT plans to operate on self.rho_p self.rho_p_fft_ax0 = pyfftw.FFTW( self.rho_p, self.rho_p, axes=(0,), direction='FFTW_FORWARD', flags=fftw_flags, threads=fftw_nthreads ) self.rho_p_ifft_ax1 = pyfftw.FFTW( self.rho_p, self.rho_p, axes=(1,), direction='FFTW_BACKWARD', flags=fftw_flags, threads=fftw_nthreads ) # numexpr codes to calculate the First Ehrenfest theorems self.code_V_average = "sum((%s) * density)" % self.V.format(X="X") self.code_K_average = "sum((%s) * density)" % self.K.format(P="P") self.code_X_average = "sum(X * density)" self.code_P_average = "sum(P * density)" self.code_P_average_RHS = "sum((-(%s) + (%s)) * density)" % ( self.diff_V.format(X="X"), self.RHS_P_A.format(X="X") ) self.code_X_average_RHS = "sum(((%s) + (%s)) * density)" % ( self.diff_K.format(P="P"), self.RHS_X_B.format(P="P") ) # Lists where the expectation values of X and P self.X_average = [] self.P_average = [] # Lists where the right hand sides of the Ehrenfest theorems for X and P self.X_average_RHS = [] self.P_average_RHS = [] # List where the expectation value of the Hamiltonian will be calculated self.hamiltonian_average = [] # Flag requesting tha the Ehrenfest theorem calculations self.isEhrenfest = True except AttributeError: # Since self.diff_V and self.diff_K are not specified, # the first Ehrenfest theorem will not be calculated self.isEhrenfest = False
def hilbert_fftw(s, debug=False, dtype = 'complex64'): ''' fftw drop replacement for scipy_fftpack.hilbert beware of sign of returned seq. is '-' http://au.mathworks.com/help/signal/ref/hilbert.html The analytic signal for a sequence x has a one-sided Fourier transform. That is, the transform vanishes for negative frequencies. To approximate the analytic signal, hilbert calculates the FFT of the input sequence, replaces those FFT coefficients that correspond to negative frequencies with zeros, and calculates the inverse FFT of the result. In detail, hilbert uses a four-step algorithm: It calculates the FFT of the input sequence, storing the result in a vector x. It creates a vector h whose elements h(i) have the values: 1 for i = 1, (n/2)+1 2 for i = 2, 3, ... , (n/2) 0 for i = (n/2)+2, ... , n It calculates the element-wise product of x and h. It calculates the inverse FFT of the sequence obtained in step 3 and returns the first n elements of the result. This algorithm was first introduced in [2]. The technique assumes that the input signal, x, is a finite block of data. This assumption allows the function to remove the spectral redundancy in x exactly. Methods based on FIR filtering can only approximate the analytic signal, but they have the advantage that they operate continuously on the data. See Single-Sideband Amplitude Modulation for another example of a Hilbert transform computed with an FIR filter.''' n = len(s) pyfftw.interfaces.cache.enable() pyfftw.interfaces.cache.set_keepalive_time(50.0) align = pyfftw.simd_alignment write_wisdom=False try: wisdom = pickle.load( open( "wisdom_hilbert.wis", "rb" ) ) pyfftw.import_wisdom(wisdom) except: write_wisdom = True print 'no wisdom file' fft_in = pyfftw.empty_aligned(n, dtype=dtype, n=align) fft_out = pyfftw.empty_aligned(n, dtype=dtype, n=align) ifft_in = pyfftw.empty_aligned(n, dtype=dtype, n=align) ifft_out = pyfftw.empty_aligned(n, dtype=dtype, n=align) fft_machine = pyfftw.FFTW(fft_in, fft_out, direction='FFTW_FORWARD', flags=('FFTW_ESTIMATE',), threads=8) ifft_machine = pyfftw.FFTW(ifft_in, ifft_out, direction='FFTW_BACKWARD', flags=('FFTW_ESTIMATE',), threads=8) if write_wisdom: wisdom = pyfftw.export_wisdom() pickle.dump( wisdom, open( "wisdom_hilbert.wis", "wb" ) ) s_0 = time.time() fft_in[:] = s S = fft_machine() if debug: print 'fft', time.time()-s_0 s_1 = time.time() h = np.zeros(n) h[0] = 1. h[n/2] = 1. h[1:n/2] = 2. if debug: print 'setup', time.time()-s_1 s_2 = time.time() ifft_in[:] = h*S ret = ifft_machine() if debug: print 'ifft', time.time()-s_2 return -ret[:n].imag
def pyfftw_call(array_in, array_out, direction='forward', axes=None, halfcomplex=False, **kwargs): """Calculate the DFT with pyfftw. The discrete Fourier (forward) transform calcuates the sum:: f_hat[k] = sum_j( f[j] * exp(-2*pi*1j * j*k/N) ) where the summation is taken over all indices ``j = (j[0], ..., j[d-1])`` in the range ``0 <= j < N`` (component-wise), with ``N`` being the shape of the input array. The output indices ``k`` lie in the same range, except for half-complex transforms, where the last axis ``i`` in ``axes`` is shortened to ``0 <= k[i] < floor(N[i]/2) + 1``. In the backward transform, sign of the the exponential argument is flipped. Parameters ---------- array_in : `numpy.ndarray` Array to be transformed array_out : `numpy.ndarray` Output array storing the transformed values, may be aliased with ``array_in``. direction : {'forward', 'backward'} Direction of the transform axes : int or sequence of ints, optional Dimensions along which to take the transform. ``None`` means using all axes and is equivalent to ``np.arange(ndim)``. halfcomplex : bool, optional If ``True``, calculate only the negative frequency part along the last axis. If ``False``, calculate the full complex FFT. This option can only be used with real input data. Other Parameters ---------------- fftw_plan : ``pyfftw.FFTW``, optional Use this plan instead of calculating a new one. If specified, the options ``planning_effort``, ``planning_timelimit`` and ``threads`` have no effect. planning_effort : {'estimate', 'measure', 'patient', 'exhaustive'} Flag for the amount of effort put into finding an optimal FFTW plan. See the `FFTW doc on planner flags <http://www.fftw.org/fftw3_doc/Planner-Flags.html>`_. Default: 'estimate' planning_timelimit : float or ``None``, optional Limit planning time to roughly this many seconds. Default: ``None`` (no limit) threads : int, optional Number of threads to use. Default: Number of CPUs if the number of data points is larger than 4096, else 1. normalise_idft : bool, optional If ``True``, the result of the backward transform is divided by ``1 / N``, where ``N`` is the total number of points in ``array_in[axes]``. This ensures that the IDFT is the true inverse of the forward DFT. Default: ``False`` import_wisdom : filename or file handle, optional File to load FFTW wisdom from. If the file does not exist, it is ignored. export_wisdom : filename or file handle, optional File to append the accumulated FFTW wisdom to Returns ------- fftw_plan : ``pyfftw.FFTW`` The plan object created from the input arguments. It can be reused for transforms of the same size with the same data types. Note that reuse only gives a speedup if the initial plan used a planner flag other than ``'estimate'``. If ``fftw_plan`` was specified, the returned object is a reference to it. Notes ----- * The planning and direction flags can also be specified as capitalized and prepended by ``'FFTW_'``, i.e. in the original FFTW form. * For a ``halfcomplex`` forward transform, the arrays must fulfill ``array_out.shape[axes[-1]] == array_in.shape[axes[-1]] // 2 + 1``, and vice versa for backward transforms. * All planning schemes except ``'estimate'`` require an internal copy of the input array but are often several times faster after the first call (measuring results are cached). Typically, 'measure' is a good compromise. If you cannot afford the copy, use ``'estimate'``. * If a plan is provided via the ``fftw_plan`` parameter, no copy is needed internally. """ import pickle if not array_in.flags.aligned: raise ValueError('input array not aligned') if not array_out.flags.aligned: raise ValueError('output array not aligned') if axes is None: axes = tuple(range(array_in.ndim)) axes = normalized_axes_tuple(axes, array_in.ndim) direction = _pyfftw_to_local(direction) fftw_plan_in = kwargs.pop('fftw_plan', None) planning_effort = _pyfftw_to_local(kwargs.pop('planning_effort', 'estimate')) planning_timelimit = kwargs.pop('planning_timelimit', None) threads = kwargs.pop('threads', None) normalise_idft = kwargs.pop('normalise_idft', False) wimport = kwargs.pop('import_wisdom', '') wexport = kwargs.pop('export_wisdom', '') # Cast input to complex if necessary array_in_copied = False if is_real_dtype(array_in.dtype) and not halfcomplex: # Need to cast array_in to complex dtype array_in = array_in.astype(complex_dtype(array_in.dtype)) array_in_copied = True # Do consistency checks on the arguments _pyfftw_check_args(array_in, array_out, axes, halfcomplex, direction) # Import wisdom if possible if wimport: try: with open(wimport, 'rb') as wfile: wisdom = pickle.load(wfile) except IOError: wisdom = [] except TypeError: # Got file handle wisdom = pickle.load(wimport) if wisdom: pyfftw.import_wisdom(wisdom) # Copy input array if it hasn't been done yet and the planner is likely # to destroy it. If we already have a plan, we don't have to worry. planner_destroys = _pyfftw_destroys_input( [planning_effort], direction, halfcomplex, array_in.ndim) must_copy_array_in = fftw_plan_in is None and planner_destroys if must_copy_array_in and not array_in_copied: plan_arr_in = np.empty_like(array_in) flags = [_local_to_pyfftw(planning_effort), 'FFTW_DESTROY_INPUT'] else: plan_arr_in = array_in flags = [_local_to_pyfftw(planning_effort)] if fftw_plan_in is None: if threads is None: if plan_arr_in.size <= 4096: # Trade-off wrt threading overhead threads = 1 else: threads = cpu_count() fftw_plan = pyfftw.FFTW( plan_arr_in, array_out, direction=_local_to_pyfftw(direction), flags=flags, planning_timelimit=planning_timelimit, threads=threads, axes=axes) else: fftw_plan = fftw_plan_in fftw_plan(array_in, array_out, normalise_idft=normalise_idft) if wexport: try: with open(wexport, 'ab') as wfile: pickle.dump(pyfftw.export_wisdom(), wfile) except TypeError: # Got file handle pickle.dump(pyfftw.export_wisdom(), wexport) return fftw_plan
def save_wisdom(self, wisdom_file=DEFAULT_WISDOM_FILE): """Save FFTW wisdom to file.""" with open(wisdom_file, 'wb') as f: pickle.dump(pyfftw.export_wisdom(), f)