def __init__(self, filename):
self.filename = filename
self.label = filename # May need to shorten
self.open(filename)
self.dimensions = self.getDimensions() # A dictionary of XPadDataDim objects
self.varNames = self.handle.variables.keys() # A list of variable names
for i,v in enumerate(self.varNames):
try:
# Python 2
if isinstance(v, unicode):
v = v.encode('utf-8')
v = str(v).translate(None, '\0')
except NameError:
# Python 3
if isinstance(v, str):
v = v.encode('utf-8')
self.varNames[i] = v
self.variables = {} # A dictionary of XPadDataItem objects with empty data
for name, var in self.handle.variables.items():
item = XPadDataItem()
item.name = name
item.source = self.filename
item.dim = map(lambda d: self.dimensions[d], var.dimensions)
self.variables[name] = item
self.close()
def read(self, name, shot):
"""Read a variable from the file."""
self.open()
if self.handle == None: return None
try:
var = self.handle.variables[name]
except KeyError:
# Not found. Try to find using case-insensitive search
var = None
for n in self.handle.variables.keys():
if n.lower() == name.lower():
print("WARNING: Reading '"+n+"' instead of '"+name+"'")
var = self.handle.variables[n]
if var == None:
return None
ndims = len(var.dimensions)
if ndims == 0:
data = var.getValue()
else:
data = var[:]
item = XPadDataItem()
item.name = name
item.source = self.filename
item.data = data
item.dim = map(lambda d: self.dimensions[d], var.dimensions)
self.close()
return item
def newfunc(data):
result = XPadDataItem()
if data.name != "":
result.name = name + "( " + data.name + " )"
result.source = data.source
if data.label != "":
result.label = name + "( " + data.label + " )"
result.data = func(data.data)
result.dim = data.dim
result.order = data.order
result.time = data.time
return result
def chop(item, t_min, t_max):
"""
>>> from user_functions import *
>>> a = chop(XMC_OMV_110, 0.274, 0.276)
>>> a_amp,a_phase = fftp(a)
>>> b = chop(a_phase, 0.0, 100.0)
>>> plot(b)
"""
if len(item.dim) != 1:
raise ValueError("chop can only operate on 1D traces currently")
if t_max < t_min or t_max < item.dim[0].data[0] or t_min > item.dim[0].data[-1]:
raise ValueError("New time-range not defined correctly")
idx = np.where(np.logical_and(item.dim[0].data >= t_min, item.dim[0].data <= t_max))
if len(idx[0]) == 0:
raise ValueError("No data in time-range specified")
# Calculate the phase
chopped = XPadDataItem()
if item.name != "":
chopped.name = "CHOP( "+item.name+", "+str(t_min)+", "+str(t_max)+" )"
chopped.source = item.source
if item.label != "":
chopped.label = "CHOP( "+item.label+", "+str(t_min)+", "+str(t_max)+" )"
chopped.units = item.units
chopped.data = item.data[idx]
# Create a dimension
dim = XPadDataDim()
dim.name = item.dim[0].name
dim.label = item.dim[0].label
dim.units = item.dim[0].units
dim.data = item.dim[0].data[idx]
chopped.dim = [dim]
if chopped.dim[0].units in ["s", "S", "sec", "Sec", "SEC"]:
chopped.order = 0
chopped.time = chopped.dim[0].data
return chopped
def integrate(item):
"""
Integrate the given trace
Inputs
------
item - an XPadDataItem object (or equivalent)
Returns
-------
an XPadDataItem object
"""
if len(item.dim) != 1:
raise ValueError("Can only integrate 1D traces currently")
# Create a result
result = XPadDataItem()
if item.name != "":
result.name = "INTG( "+item.name+" )"
result.source = item.source
if item.label != "":
result.label = "INTG( "+item.label+" )"
if item.units != "":
result.units = item.units+"*"+item.dim[0].units
result.data = zeros(item.data.shape)
time = item.dim[0].data
result.data[1:] = cumsum(0.5*(time[1:]-time[0:-1])*(item.data[1:] + item.data[0:-1]))
result.dim = item.dim
result.order = item.order
result.time = item.time
return result
def fftp(item):
"""
Calculate amplitude and phase as a function of frequency
Inputs
------
item - an XPadDataItem object
Returns
-------
amplitude, phase pair of XPadDataItem objects
"""
if len(item.dim) != 1:
raise ValueError("fftp can only operate on 1D traces currently")
# Calculate FFT
data = rfft(item.data)*(1./len(item.data))
# Create a dimension
dim = XPadDataDim()
dim.name = "Frequency"
step = (item.dim[0].data[1] - item.dim[0].data[0])
dim.data = rfftfreq(len(item.data), step)
dim.units = "1/"+item.dim[0].units
if item.dim[0].units in ["s", "S", "sec", "Sec", "SEC"]:
dim.data /= 1000.
dim.units = "kHz"
# Calculate the amplitude
amp = XPadDataItem()
if item.name != "":
amp.name = "AMP( "+item.name+" )"
amp.source = item.source
if item.label != "":
amp.label = "AMP( "+item.label+" )"
amp.units = item.units
amp.data = abs(data)
amp.dim = [dim]
# Calculate the phase
phase = XPadDataItem()
if item.name != "":
phase.name = "PHASE( "+item.name+" )"
phase.source = item.source
if item.label != "":
phase.label = "PHASE( "+item.label+" )"
phase.units = "Radians"
phase.data = arctan2(data.real, data.imag)
a = phase.data - 2*pi
for i in range(1,len(phase.data)):
if abs(phase.data[i-1] - a[i]) < 1.:
phase.data[i] = a[i]
a = phase.data + 2*pi
for i in range(1,len(phase.data)):
if abs(phase.data[i-1] - a[i]) < 1.:
phase.data[i] = a[i]
phase.dim = [dim]
return amp, phase
