# -*- coding: utf-8 -*- """ Created on Fri Sep 28 17:16:12 2012 @author: Edmondo """ import numpy import matplotlib.pyplot as plt from scipy.signal import hilbert, resample, medfilt, lfilter, butter, buttord, filtfilt from scipy.integrate import cumtrapz from scipy.interpolate import RectBivariateSpline, InterpolatedUnivariateSpline #from ftubase import smoothd import inspect class _AxInfo(object): """Keeps the axis label and unit in a single class. """ def __init__(self, name = '', label=None, unit=None): self.name = name self.label = label self.unit = unit def addlu(self, label, unit): """Change the label and the unit to the current object """ self.label = label self.unit = unit def parse1d(self, exprl): """The expression coming from a 1d FTU channel is decomposed in label and unit """ xxl = str(exprl).split('(') self.label = xxl[0] self.unit = xxl[1].split(')')[0] def _exposednames(self): """ expose the name of the label and unit. It basically correspond to self.name+'label and self.name+'unit' A dictionary is returned """ return {'labelname':self.name+'label', 'unitname':self.name+'unit'} def getnames(self): """Return a dictionary whith the label and utni name associated with self. It should be used in the methid __getattr__ of Ftudata in order to understand which _AxInfo should be called to get the proper label """ names = self._exposednames() return {names['labelname']:self, names['unitname']:self} def torstr(self): """A list of string representations are returned. To be used in the __repr__() method of Ftudata """ rstr = [] if self.label is not None: names = self._exposednames() rstr += ["%8s: %s" % (names['labelname'], self.label)] rstr += ["%8s: %s" % (names['unitname'], self.unit)] return rstr def __getitem__(self, name): names = self._exposednames() if name == names['labelname']: return self.label elif name == names['unitname']: return self.unit raise AttributeError('No {0} attribute is present'.format(name)) class V2d(object): """It defines object with an R a Z and a V fields, various operation are defined on them. """ def __init__(self, r=None, t=None, v=None, ier=None): self.r = r self.t = t self.v = v if ier is None: ier = 0 self.ier = ier def copy(self): """return a copy of the given object """ res = self._getacopy() return res def _getacopy(self, copydata=True, inplace=False): """Return a copy of the given object. copydata: if True also the field v is copied inplace: if True the actual object is reeturned isntead of a copy """ if inplace: return self other = self.__class__() if self.r is not None: other.r = self.r.copy() if self.t is not None: other.t = self.t.copy() other.ier = self.ier if copydata and self.v is not None: other.v = self.v.copy() return other def __repr__(self): try: nr = len(self.r) except TypeError: nr = 0 try: nt = len(self.t) except TypeError: nt = 0 rstr = [] rstr += [" v: Float %dx%d" % (nt, nr)] rstr += [" t: Float %d" % nt] rstr += [" r: Float %d" % nr] rstr += [" ier: %d" % self.ier] return "\n".join(rstr) def ok(self): """Return True if the object contains correct data """ return self.ier == 0 and self.v.size > 0 def vtline(self, *arg, **args): """Plot a vertical line for each time which is present """ for t in self.t: plt.axvline(t, *arg, **args) def plot(self, t = numpy.nan, r = numpy.nan, func=None, ax=None, **args): """Plot the data, various keyword are defined: t: if not is NAN the time slice closest to the given time is plotted """ if ax is None: ax = plt if func is None: func = lambda x:x if (numpy.isnan(t) and numpy.isnan(r)): temp = ax.plot(self.t, func(self.v), **args) elif (numpy.isnan(t)): ir = numpy.argmin(abs(self.r - r)) temp = ax.plot(self.t, func(self.v[:,ir]), **args) elif (numpy.isnan(r)): it = numpy.argmin(abs(self.t - t)) temp = ax.plot(self.r, func(self.v[it,:]), **args) return temp def get(self, r=numpy.nan, t=numpy.nan): if (numpy.isnan(t) and numpy.isnan(r)): return (self.r, self.t, self.v) elif (numpy.isnan(t)): ir = numpy.argmin(abs(self.r - r)) if len(self.v.shape) == 1: return (self.t,self.v) else: return (self.t, self.v[:,ir]) elif (numpy.isnan(r)): it = numpy.argmin(abs(self.t - t)) if len(self.v.shape) == 1: return (self.r, self.v[it]) else: return (self.r, self.v[it,:]) else: ir = numpy.argmin(abs(self.r - r)) it = numpy.argmin(abs(self.t - t)) if len(self.v.shape) == 1: return self.v[it] else: return self.v[it,ir] def slice(self, r=numpy.nan, t=numpy.nan): result = self._getacopy(copydata=False) if not self.ok(): return result if (numpy.isnan(t) and numpy.isnan(r)): return result elif (numpy.isnan(t)): ir = numpy.argmin(abs(self.r - r)) result.r = self.r[ir][numpy.newaxis] result.v = self.v[:,ir][:,numpy.newaxis] elif (numpy.isnan(r)): it = numpy.argmin(abs(self.t - t)) result.t = self.t[it][numpy.newaxis] result.v = self.v[it,:][numpy.newaxis,:] else: ir = numpy.argmin(abs(self.r - r)) it = numpy.argmin(abs(self.t - t)) result.r = self.r[ir][numpy.newaxis] result.t = self.t[it][numpy.newaxis] result.v = self.v[it,ir][numpy.newaxis,numpy.newaxis] return result def get_r_profile(self, t = numpy.nan): it = numpy.argmin(abs(self.t - t)) return (self.r, self.v[it,:]) def get_t_profile(self, r = numpy.nan): ir = numpy.argmin(abs(self.r - r)) return (self.t, self.v[:,ir]) def trange(self, tlim, inplace=False): """Restrict the time axis in the specified range """ result = self._getacopy(copydata=False, inplace=inplace) if not self.ok(): return result tlimm = numpy.asarray(tlim, dtype=float) if tlimm.size == 2: idt = (self.t >= numpy.min(tlimm)) & (self.t <= numpy.max(tlimm)) result.t = self.t[idt] if len(self.v.shape)==1: result.v = self.v[idt] else: result.v = self.v[idt,:] return result def tresample(self, n, inplace=False): """Resample the data along the time axis. Calling: obj.tresample(n,inplace=[True,False]) n specify the number of time points it will get """ result = self._getacopy(copydata=False, inplace=inplace) v, t = resample(self.v, n, t=self.t, axis=0) result.v = v result.t = t return result def mean(self, trange): if not self.ok(): return None trange = numpy.asarray(trange, dtype=float) idt = (self.t >= numpy.min(trange)) & (self.t <= numpy.max(trange)) if len(self.v.shape) == 2: return numpy.mean(self.v[idt,:]) else: return numpy.mean(self.v[idt]) def smooth(self, inplace=False): if self.v is None: return result = self._getacopy(copydata=False, inplace=inplace) if len(self.v.shape) == 1: #result.v, rough = smoothd(self.v) pass elif len(self.v.shape) == 2: vv = numpy.zeros(self.v.T.shape) for v,sv in zip(vv,self.v.T): # v[:], rough = smoothd(sv) pass result.v = vv.T return result def medfilt(self, nm=5, inplace=True): if self.v is None: return result = self._getacopy(copydata=False, inplace=inplace) if len(self.v.shape) == 1: result.v = medfilt(self.v, kernel_size=nm) elif len(self.v.shape) == 2: vv = numpy.zeros(self.v.T.shape) for v,sv in zip(vv,self.v.T): v[:] = medfilt(sv, kernel_size=nm) result.v = vv.T return result def specgram(self, ax=None, *args, **kwargs): if self.v is None: return if ax is None: ax = plt dt = numpy.mean(numpy.diff(self.t)) Fs = 1.0 / dt xextent = [self.t.min(), self.t.max()] # not correct TODO ax.specgram(self.v, *args, Fs=Fs, xextent = xextent, **kwargs) def filter(self, Wn, ftype='low', inplace=False, symmetric=True): if self.v is None: return result = self._getacopy(copydata=False, inplace=inplace) dt = numpy.mean(numpy.diff(self.t)) FNy2 = 1.0 / 2.0 / dt Wno = Wn / FNy2 if ftype == 'low': bord = buttord(Wno,Wno*1.7,1,10) elif ftype == 'high': bord = buttord(Wno*0.6,Wno,1,10) else: raise Exception('Ma che fai') bf,af = butter(*bord,btype=ftype) if len(self.v.shape) == 1: if symmetric: result.v = filtfilt(bf,af, self.v) else: result.v = lfilter(bf, af, self.v) elif len(self.v.shape) == 2: result.v = lfilter(bf, af, self.v, axis=0) return result def integrate(self, Wn=None, inplace=False, symmetric=True): if self.v is None: return result = self._getacopy(copydata=False, inplace=inplace) if len(self.v.shape) == 1: ii = cumtrapz(self.v, x=self.t) result.v = numpy.hstack((numpy.zeros(1),ii)) else: assert(False) if Wn is not None: result.filter(Wn,ftype='high',inplace=True, symmetric=symmetric) return result def hilbert(self, inplace=False): """return the hilbert transform of the signal. The original signal is left unchanged. """ if self.v is None: return result = self._getacopy(copydata=False, inplace=inplace) if 'axis' in inspect.getargspec(hilbert)[0]: result.v = hilbert(self.v, axis=0) else: if len(self.v.shape) == 1: result.v = hilbert(self.v) elif len(self.v.shape) == 2: vv = numpy.zeros(self.v.T.shape,dtype=numpy.complex) for v,sv in zip(vv,self.v.T): v[:] = hilbert(sv) result.v = vv.T return result def packhilbert(self, packlen=500, inplace=False): """Return the Hilbert transform, a packet algorithm is used """ if self.v is None: return result = self._getacopy(copydata=False, inplace=inplace) nt = self.v.shape[0] npacket = ((nt - 1) // packlen) + 1 for i in range(npacket): ibase = numpy.clip(packlen*(i - 1), 0, nt-1) iplus = numpy.clip(packlen*(i + 2), 0, nt) istart = numpy.clip(packlen*i, 0, nt-1) iend = numpy.clip(packlen*(i + 1), 0, nt) def clip(self, vmin, vmax, inplace=False): result = self._getacopy(copydata=False, inplace=inplace) result.v = numpy.clip(self.v,vmin, vmax) return result def abs(self, inplace=False): result = self._getacopy(copydata=False, inplace=inplace) result.v = abs(self.v) return result def real(self, inplace=False): result = self._getacopy(copydata=False, inplace=inplace) result.v = numpy.real(self.v) return result def imag(self, inplace=False): result = self._getacopy(copydata=False, inplace=inplace) result.v = numpy.imag(self.v) return result def angle(self, inplace=False): result = self._getacopy(copydata=False, inplace=inplace) result.v = numpy.angle(self.v) return result def unwrap(self, axis=0, inplace=False): """Convenience method for numpy.unwrap. Contrary to the unwrap convention, it is applied on the first dimension (that generally correspond to the time coordinate) Calling: obj.unwrap(axis=) """ result = self._getacopy(copydata=False, inplace=inplace) result.v = numpy.unwrap(self.v, axis=axis) return result def unifunct(self, funct=None, inplace=False, **args): """Apply a function to the actual data. Calling: obj.unifunct(funct=, inplace=[True,False]) funct could be a lambda expression and if not specified it is funct=lambda x:x """ if funct is None: funct = lambda x:x result = self._getacopy(copydata=False, inplace=inplace) result.v = funct(self.v, **args) return result def difft(self, inplace=False): """Derivative along the t direction. The original signal is left unchanged. """ result = self._getacopy(copydata=False, inplace=inplace) vdiff = self.v[1:, :] - self.v[:-1, :] tdiff = self.t[1:] - self.t[:-1] tmea = 0.5*(self.t[1:] + self.t[:-1]) result.t = tmea result.v = vdiff/tdiff[:, numpy.newaxis] return result def dt(self, inplace=False): """Derivative along the t direction. The original signal is left unchanged. """ result = self._getacopy(copydata=False, inplace=inplace) if self.v.shape == 2: vdt, vdr = numpy.gradient(self.v) else: vdt = numpy.gradient(self.v) dtt = numpy.gradient(self.t) result.v = vdt/dtt return result def dr(self, inplace=False): """Gradient along the radial direction """ result = self._getacopy(copydata=False, inplace=inplace) vdt, vdr = numpy.gradient(self.v) dr = numpy.gradient(self.r) result.v = vdr/dr return result def integ_r(self, inplace=False): """Integrate along the radial direction """ result = self._getacopy(copydata=False, inplace=inplace) v = numpy.hstack((numpy.zeros((self.v.shape[0],1)),cumtrapz(self.v,x=self.r))) result.v = v return result @staticmethod def show(): """It is equivalent to plt.show(). Just for convenience """ plt.show() def movezeroto(self, tzero, inplace=False): """Move the zero of the t axis to tzero. The t coordinate is changed so that t(new) = t(old) - tzero. This is useful when something interesting is happening at different times on the same or a different channel """ result = self._getacopy(inplace=inplace) result.t = self.t - tzero return result def __add__(self, other): result = self._getacopy(copydata=False) if isinstance(other, V2d): result.v = self.v + other.v else: result.v = self.v + other return result def __radd__(self, other): result = self._getacopy(copydata=False) result.v = other + self.v return result def __iadd__(self, other): self.v = self.v + other return self def __sub__(self, other): result = self._getacopy(copydata=False) if isinstance(other, V2d): result.v = self.v - other.v else: result.v = self.v - other return result def __rsub__(self, other): result = self._getacopy(copydata=False) result.v = other - self.v return result def __isub__(self, other): self.v = self.v - other return self def __mul__(self, other): result = self._getacopy(copydata=False) if isinstance(other, V2d): result.v = self.v * other.v else: result.v = self.v * other return result def __rmul__(self, other): result = self._getacopy(copydata=False) result.v = other * self.v return result def __imul__(self, other): self.v = self.v * other return self def __div__(self, other): result = self._getacopy(copydata=False) if isinstance(other, V2d): result.v = self.v / other.v else: result.v = self.v / other return result def __rdiv__(self, other): result = self._getacopy(copydata=False) result.v = other / self.v return result def __idiv__(self, other): self.v = self.v / other return self def __pow__(self, other): result = self._getacopy(copydata=False) result.v = self.v**other return result def __ipow__(self, other): self.v = self.v ** other return self def __neg__(self): result = self._getacopy(copydata=False) result.v = -self.v return result def __pos__(self): return self def function(self, obj=None, r=None, t=None, **args): if obj is not None: r = obj.r t = obj.t rr = r tt = t if rr is None: rr = 0 if tt is None: tt = 0 if self.r is None: func_base = InterpolatedUnivariateSpline(self.t,self.v) func = lambda t, r: func_base(t)[:,numpy.newaxis] + numpy.zeros((1,numpy.asarray(r).size)) elif self.t is None: func_base = InterpolatedUnivariateSpline(self.r,self.v) func = lambda t, r: func_base(r) + numpy.zeros((numpy.asarray(t).size,1)) else: func = RectBivariateSpline(self.t, self.r, self.v, **args) return V2d(r,t,func(tt, rr)) class NamedV2d(V2d): """A V2d class with added label and unit to each axis """ def __init__(self, **args): super(NamedV2d, self).__init__(**args) self.tax = _AxInfo('t') self.rax = _AxInfo('r') self.vax = _AxInfo('') def _getacopy(self, *arg, **args): other = super(NamedV2d, self)._getacopy(*arg, **args) if other is not self: other.tax = self.tax other.rax = self.rax other.vax = self.vax return other def __repr__(self): basestr = super(NamedV2d, self).__repr__() rstr = [] rstr += self.tax.torstr() rstr += self.rax.torstr() rstr += self.vax.torstr() if len(rstr) == 0: return basestr else: return "\n".join(rstr)+"\n\n"+basestr def __getattr__(self, name): namedict = self.tax.getnames() namedict.update(self.rax.getnames()) namedict.update(self.vax.getnames()) if name in namedict: return namedict[name][name] else: raise AttributeError('No {0} attribute is present'.format(name))