def where_fits(val, arr): """ finds the elements between which val fits needs a value and a lineraly increasing array args: val: A value arr: linearly increasing numpy array returns: (lh_element, rh_element) both integers """ try: #Concise bit of clever code #which means it is difficult to explain #basically the LHS (arr <= val) makes #an array of bools which is true # to the LHS of where Val fits in arr #nwhere (taken carefully from numpy as a different name) #returns the indexes of those trues and [0][-1] (shape 0,len(arr) for some unknown reason) #indexes the RHS-most element that is true giving the index #to the left of where val fits in arr... #the second calc in the tupple does the same but mirrored mytup=(nwhere((arr <= val))[0][-1],nwhere((arr >= val))[0][0]) except IndexError: #if the code above throws an IndexError the val fits at an end of arr #determine which end if val < arr.min(): mytup=(0,0) if val > arr.max(): mytup=(len(arr)-1, len(arr)-1) return mytup
def is_sonde_good(sonde, min_ht, keylist): sonde_good=sonde['alt(m)'][-1] > min_ht if not(sonde_good): print "Sonde Burst at ", sonde['alt(m)'][-1], "m rejecting" for i in range(len(keylist)): occur=len(nwhere(array(sonde[keylist[i]])==-9999)[0]) sonde_good=sonde_good and not(-9999 in sonde[keylist[i]]) if (-9999 in sonde[keylist[i]]): print "Dirty data in ",keylist[i]," rejecting ", occur, " occurances" if occur < 5 and (-9999 in sonde[keylist[i]]): print "at ", sonde['alt(m)'][nwhere(array(sonde[keylist[i]])==-9999)[0]] return sonde_good
def dealias_berrimah_volume(filename, **kwargs):
pattern=kwargs.get('pattern', 'BerrimaVol')
deal_add=kwargs.get('deal_add', '_deal')
raw_path=kwargs.get('raw_path', '/data/uf_ber/')
deal_path=kwargs.get('deal_path', '/data/deal_ber/')
deal_files=os.listdir(deal_path)
raw_files=os.listdir(raw_path)
#deal_files=[]
#raw_files=[]
#for file in uf_files:
# if 'deal' in file:
# deal_files.append(file)
# else:
# raw_files.append(file)
#
if not(filename in raw_files):
raise IOError, 'File not there'
deal_files.sort()
raw_files.sort()
this_date_str=filename[len(pattern):-3].replace('_', ' ')
this_date_str=this_date_str[0:this_date_str.find(' ')+5]
sonde_name=read_sounding.make_deal_sonde(this_date_str)
deal_date=sonde_name[0:sonde_name.find('_')]
date_num_list=[]
for file in deal_files:
deal_str=file[len(pattern):-8].replace('_', ' ')
deal_str=deal_str[0:deal_str.find(' ')+5]
date_num_list.append(datestr2num(deal_str))
offset=array(date_num_list)-datestr2num(this_date_str)
idec_where=nwhere(offset < 0.0)[0]
offset_least=100.0
if len(idec_where)!=0: offset_least=offset[idec_where[-1]]
print "**********************"
if abs(offset_least) < 30.0/(60.0*24.0):
deal_fname=deal_files[idec_where[-1]]
print "Found a previous de-aliased file ", deal_fname
else:
print "No de-aliased file found within 30 minutes, only using sounding"
deal_fname='dummy'
print "********************"
outfile=filename[0:-3]+deal_add+".uf"
cwd=os.getcwd()
os.chdir('/home/scollis/bom_mds/dealias/')
execbl='./FourDD_berrimah '
command=execbl+deal_path+deal_fname+' '+raw_path+filename+' '+deal_path+outfile+' '+ sonde_name+' '+deal_date+' '+'0 1 1 1'
print command
os.system(command)
os.chdir(cwd)
return outfile
def dealias_single_volume(filename, **kwargs):
pattern=kwargs.get('pattern', 'Gunn_pt_')
deal_add=kwargs.get('deal_add', '_deal')
lassen_path=kwargs.get('lassen_path', '/bm/gscratch/scollis/lassen_cpol/')
uf_path=kwargs.get('lassen_path', '/bm/gscratch/scollis/uf_cpol/')
deal_files=os.listdir(uf_path)
raw_files=os.listdir(lassen_path)
#deal_files=[]
#raw_files=[]
#for file in uf_files:
# if 'deal' in file:
# deal_files.append(file)
# else:
# raw_files.append(file)
#
if not(filename in raw_files):
raise IOError, 'File not there: '+filename
deal_files.sort()
raw_files.sort()
this_date_str=filename[len(pattern):-11]
#this_date_str=this_date_str[0:this_date_str.find(' ')+5]
sonde_name=read_sounding.make_deal_sonde(this_date_str)
deal_date=sonde_name[0:sonde_name.find('_')]
date_num_list=[]
for file in deal_files:
deal_str=file[len(pattern):-8]
#print deal_str
#deal_str=deal_str[0:deal_str.find(' ')+5]
#print deal_str
date_num_list.append(datestr2num(deal_str))
offset=array(date_num_list)-datestr2num(this_date_str)
idec_where=nwhere(offset < 0.0)[0]
offset_least=100.0
if len(idec_where)!=0: offset_least=offset[idec_where[-1]]
print "**********************"
if abs(offset_least) < 30.0/(60.0*24.0):
deal_fname=deal_files[idec_where[-1]]
print "Found a previous de-aliased file ", deal_fname
else:
print "No de-aliased file found within 30 minutes, only using sounding"
deal_fname='dummy'
print "********************"
outfile=filename[0:-11]+deal_add+".uf"
cwd=os.getcwd()
os.chdir('/flurry/home/scollis/bom_mds/dealias/')
execbl='./FourDD_lassen '
command=execbl+uf_path+deal_fname+' '+lassen_path+filename+' '+uf_path+outfile+' '+ sonde_name+' '+deal_date+' '+'1 1 1 1'
print command
os.system(command)
os.chdir(cwd)
return outfile
def MakeResultPlotCFY(solver = None, freqCutoff = 2.0, datasetPath = "/", controlNumber = 0):
#Get results
if solver == None:
print "Please specify either oct object or HDF5 file."
return
elif solver.__class__ == str:
h5file = tables.openFile(solver, "r")
try:
timeGrid = h5file.getNode(datasetPath, "TimeGrid")[:]
timeGridResolution = timeGrid[1] - timeGrid[0]
controlVector = h5file.getNode(datasetPath, "FinalControl")[controlNumber]
octYield = h5file.getNode(datasetPath, "Yield")[:]
pMethod = pyprop.PenaltyMethod.Projection
forwardSolution = h5file.getNode(datasetPath, "ForwardSolution")[:]
try:
goodBadRatio = h5file.getNode(datasetPath, "GoodBadRatio")[:]
except tables.NoSuchNodeError:
pMethod = pyprop.PenaltyMethod.Energy
finally:
h5file.close()
else:
timeGrid = solver.TimeGrid
timeGridResolution = solver.TimeGridResolution
controlVector = solver.ControlVectors[controlNumber]
octYield = solver.Yield
pMethod = krotov.PenaltyMethod
if pMethod == pyprop.PenaltyMethod.Projection:
goodBadRatio = solver.GoodBadRatio[:]
if pMethod == pyprop.PenaltyMethod.Projection:
LaTeXFigureSettings(subFig=(2,3))
figure()
subplots = [subplot(321), subplot(322), subplot(323), subplot(324)]
subplots += [axes([.125, .1, .78, .22])]
else:
LaTeXFigureSettings(subFig=(2,2))
figure()
subplots = [subplot(221), subplot(222), subplot(223), subplot(224)]
#Plot final control
subplots[0].plot(timeGrid, controlVector, label="Control function")
subplots[0].set_xlabel("Time (a.u.)")
subplots[0].legend(loc="best")
#Plot control spectrum
freq, controlSpectrum = GetPulseSpectrum(controlVector, timeGridResolution)
spectrumSize = size(controlSpectrum)
freq = freq[spectrumSize/2:]
absSpectrum = abs(controlSpectrum[spectrumSize/2:])
I = nwhere(freq > freqCutoff)[0][0]
subplots[1].plot(freq[:I], absSpectrum[:I], label="Control spectrum")
subplots[1].set_xlabel("Frequency (a.u.)")
subplots[1].legend(loc="best")
#Plot yield
subplots[2].plot(octYield, label="Yield")
subplots[2].axis([0, len(octYield), 0, 1.1])
subplots[2].set_xlabel("Iteration number")
subplots[2].legend(loc="best")
#Plot bad/good ratio
if pMethod == pyprop.PenaltyMethod.Projection:
subplots[3].plot(r_[3:len(goodBadRatio)], goodBadRatio[3:], \
label=r"$\sqrt{\frac{u^T\Phi_{bad}u}{u^T\Phi_{good}u}}$")
subplots[3].set_xlabel("Iteration number")
subplots[3].legend(loc="best")
#Plot state populations
ax = subplots[-1]
ax.plot(timeGrid, abs(forwardSolution[0,:])**2, label=r"$|0\rangle$")
ax.plot(timeGrid, abs(forwardSolution[1,:])**2, label=r"$|1\rangle$")
ax.axis([timeGrid[0], timeGrid[-1], 0, 1.1])
xlabel("Time (a.u.)")
ylabel("Population")
legend(loc="best")
draw()
return subplots
emgsec = 1000
if trigsec != 1000:
if not (trigsec / 1000).is_integer():
up = 25000 / trigsec
trigdat = signal.resample_poly(trigdat.astype('float'),
up=int(up),
down=25)
else:
trigdat = signal.decimate(trigdat.astype('float'),
int(trigsec / 1000))
trigdat = trigdat[:len(scrdat)]
trigsec = 1000
# Find triggers
potential_peaks = nwhere(trigdat > 0.5 * np.max(trigdat[int(5 * sec):]))
trigloc = [int(potential_peaks[0])]
for idx, p in enumerate(potential_peaks[1:]):
if potential_peaks[idx + 1] - potential_peaks[idx + 1 - 1] > 5 * sec:
trigloc.append(int(p))
# Onset in ms
shock_ons_acq = np.round(np.asarray(trigloc) / (sec / 1000))
# ___________________________________________________________________
# Load eprime data
convep.text_to_csv(opj(rawpath, 'eprime', 'Fear3-' + sub + '-1.txt'),
out_file=opj(rawpath, 'eprime',
'Fear3-' + sub + '-1.csv'))
edata = pd.read_csv(opj(rawpath, 'eprime', 'Fear3-' + sub + '-1.csv'))
def PropagateWavePacket(**args): #Set up problem prop = SetupProblem(**args) conf = prop.Config #Setup traveling wavepacket initial state f = lambda x: conf.Wavepacket.function(conf.Wavepacket, x) bspl = prop.psi.GetRepresentation().GetRepresentation(0).GetBSplineObject() c = bspl.ExpandFunctionInBSplines(f) prop.psi.GetData()[:] = c prop.psi.Normalize() initialPsi = prop.psi.Copy() #Get x-grid subProp = prop.Propagator.SubPropagators[0] subProp.InverseTransform() grid = prop.psi.GetRepresentation().GetLocalGrid(0) subProp.ForwardTransform() #Setup equispaced x grid x_min = conf.BSplineRepresentation.xmin x_max = conf.BSplineRepresentation.xmax grid_eq = linspace(x_min, x_max, grid.size) x_spacing = grid_eq[1] - grid_eq[0] #Set up fft grid k_spacing = 1.0 / grid_eq.size k_max = pi / x_spacing k_min = -k_max k_spacing = (k_max - k_min) / grid_eq.size grid_fft = zeros(grid_eq.size, dtype=double) grid_fft[:grid_eq.size/2+1] = r_[0.0:k_max:k_spacing] grid_fft[grid_eq.size/2:] = r_[k_min:0.0:k_spacing] print "Momentum space resolution = %f a.u." % k_spacing k0 = conf.Wavepacket.k0 k0_trunk = 5 * k0 trunkIdx = list(nwhere(abs(grid_fft) <= k0_trunk)[0]) rcParams['interactive'] = True figure() p1 = subplot(211) p2 = subplot(212) p1.hold(False) psi_eq = zeros((grid.size), dtype=complex) for t in prop.Advance(40): print "t = %f, norm = %.15f, P = %.15f " % \ ( t, prop.psi.GetNorm(), abs(prop.psi.InnerProduct(initialPsi))**2 ) sys.stdout.flush() subProp.InverseTransform() p1.plot(grid, abs(prop.psi.GetData())**2) subProp.ForwardTransform() bspl.ConstructFunctionFromBSplineExpansion(prop.psi.GetData(), grid_eq, psi_eq) psi_fft = (abs(fft.fft(psi_eq))**2) psi_fft_max = nmax(psi_fft) psi_fft /= psi_fft_max #Plot momentum space |psi|**2 p2.hold(False) p2.semilogy(grid_fft[trunkIdx], psi_fft[trunkIdx] + 1e-21) p2.hold(True) p2.semilogy([0,0], [1e-20, psi_fft_max], 'r-') p2.semilogy([-k0,-k0], [1e-20, psi_fft_max], 'g--') p2.semilogy([k0,k0], [1e-20, psi_fft_max], 'g--') #Set subplot axis p1.axis([grid[0],grid[-1],0,0.1]) p2.axis([-k0_trunk, k0_trunk, 1e-20, psi_fft_max]) show() hold(True) return prop
def MakeResultPlotCFY(solver=None,
freqCutoff=2.0,
datasetPath="/",
controlNumber=0):
#Get results
if solver == None:
print "Please specify either oct object or HDF5 file."
return
elif solver.__class__ == str:
h5file = tables.openFile(solver, "r")
try:
timeGrid = h5file.getNode(datasetPath, "TimeGrid")[:]
timeGridResolution = timeGrid[1] - timeGrid[0]
controlVector = h5file.getNode(datasetPath,
"FinalControl")[controlNumber]
octYield = h5file.getNode(datasetPath, "Yield")[:]
pMethod = pyprop.PenaltyMethod.Projection
forwardSolution = h5file.getNode(datasetPath, "ForwardSolution")[:]
try:
goodBadRatio = h5file.getNode(datasetPath, "GoodBadRatio")[:]
except tables.NoSuchNodeError:
pMethod = pyprop.PenaltyMethod.Energy
finally:
h5file.close()
else:
timeGrid = solver.TimeGrid
timeGridResolution = solver.TimeGridResolution
controlVector = solver.ControlVectors[controlNumber]
octYield = solver.Yield
pMethod = krotov.PenaltyMethod
if pMethod == pyprop.PenaltyMethod.Projection:
goodBadRatio = solver.GoodBadRatio[:]
if pMethod == pyprop.PenaltyMethod.Projection:
LaTeXFigureSettings(subFig=(2, 3))
figure()
subplots = [subplot(321), subplot(322), subplot(323), subplot(324)]
subplots += [axes([.125, .1, .78, .22])]
else:
LaTeXFigureSettings(subFig=(2, 2))
figure()
subplots = [subplot(221), subplot(222), subplot(223), subplot(224)]
#Plot final control
subplots[0].plot(timeGrid, controlVector, label="Control function")
subplots[0].set_xlabel("Time (a.u.)")
subplots[0].legend(loc="best")
#Plot control spectrum
freq, controlSpectrum = GetPulseSpectrum(controlVector, timeGridResolution)
spectrumSize = size(controlSpectrum)
freq = freq[spectrumSize / 2:]
absSpectrum = abs(controlSpectrum[spectrumSize / 2:])
I = nwhere(freq > freqCutoff)[0][0]
subplots[1].plot(freq[:I], absSpectrum[:I], label="Control spectrum")
subplots[1].set_xlabel("Frequency (a.u.)")
subplots[1].legend(loc="best")
#Plot yield
subplots[2].plot(octYield, label="Yield")
subplots[2].axis([0, len(octYield), 0, 1.1])
subplots[2].set_xlabel("Iteration number")
subplots[2].legend(loc="best")
#Plot bad/good ratio
if pMethod == pyprop.PenaltyMethod.Projection:
subplots[3].plot(r_[3:len(goodBadRatio)], goodBadRatio[3:], \
label=r"$\sqrt{\frac{u^T\Phi_{bad}u}{u^T\Phi_{good}u}}$")
subplots[3].set_xlabel("Iteration number")
subplots[3].legend(loc="best")
#Plot state populations
ax = subplots[-1]
ax.plot(timeGrid, abs(forwardSolution[0, :])**2, label=r"$|0\rangle$")
ax.plot(timeGrid, abs(forwardSolution[1, :])**2, label=r"$|1\rangle$")
ax.axis([timeGrid[0], timeGrid[-1], 0, 1.1])
xlabel("Time (a.u.)")
ylabel("Population")
legend(loc="best")
draw()
return subplots
def PropagateWavePacket(**args):
#Set up problem
prop = SetupProblem(**args)
conf = prop.Config
#Setup traveling wavepacket initial state
f = lambda x: conf.Wavepacket.function(conf.Wavepacket, x)
bspl = prop.psi.GetRepresentation().GetRepresentation(0).GetBSplineObject()
c = bspl.ExpandFunctionInBSplines(f)
prop.psi.GetData()[:] = c
prop.psi.Normalize()
initialPsi = prop.psi.Copy()
#Get x-grid
subProp = prop.Propagator.SubPropagators[0]
subProp.InverseTransform()
grid = prop.psi.GetRepresentation().GetLocalGrid(0)
subProp.ForwardTransform()
#Setup equispaced x grid
x_min = conf.BSplineRepresentation.xmin
x_max = conf.BSplineRepresentation.xmax
grid_eq = linspace(x_min, x_max, grid.size)
x_spacing = grid_eq[1] - grid_eq[0]
#Set up fft grid
k_spacing = 1.0 / grid_eq.size
k_max = pi / x_spacing
k_min = -k_max
k_spacing = (k_max - k_min) / grid_eq.size
grid_fft = zeros(grid_eq.size, dtype=double)
grid_fft[:grid_eq.size / 2 + 1] = r_[0.0:k_max:k_spacing]
grid_fft[grid_eq.size / 2:] = r_[k_min:0.0:k_spacing]
print "Momentum space resolution = %f a.u." % k_spacing
k0 = conf.Wavepacket.k0
k0_trunk = 5 * k0
trunkIdx = list(nwhere(abs(grid_fft) <= k0_trunk)[0])
rcParams['interactive'] = True
figure()
p1 = subplot(211)
p2 = subplot(212)
p1.hold(False)
psi_eq = zeros((grid.size), dtype=complex)
for t in prop.Advance(40):
print "t = %f, norm = %.15f, P = %.15f " % \
( t, prop.psi.GetNorm(), abs(prop.psi.InnerProduct(initialPsi))**2 )
sys.stdout.flush()
subProp.InverseTransform()
p1.plot(grid, abs(prop.psi.GetData())**2)
subProp.ForwardTransform()
bspl.ConstructFunctionFromBSplineExpansion(prop.psi.GetData(), grid_eq,
psi_eq)
psi_fft = (abs(fft.fft(psi_eq))**2)
psi_fft_max = nmax(psi_fft)
psi_fft /= psi_fft_max
#Plot momentum space |psi|**2
p2.hold(False)
p2.semilogy(grid_fft[trunkIdx], psi_fft[trunkIdx] + 1e-21)
p2.hold(True)
p2.semilogy([0, 0], [1e-20, psi_fft_max], 'r-')
p2.semilogy([-k0, -k0], [1e-20, psi_fft_max], 'g--')
p2.semilogy([k0, k0], [1e-20, psi_fft_max], 'g--')
#Set subplot axis
p1.axis([grid[0], grid[-1], 0, 0.1])
p2.axis([-k0_trunk, k0_trunk, 1e-20, psi_fft_max])
show()
hold(True)
return prop