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()
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 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 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 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():
""" 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)
