def hillshade(data,scale=10.0,azdeg=165.0,altdeg=45.0): ''' This code thanks to Ran Novitsky Nof http://rnovitsky.blogspot.co.uk/2010/04/using-hillshade-image-as-intensity.html Repeated here to make my cyclopean uk_map code prettier. convert data to hillshade based on matplotlib.colors.LightSource class. input: data - a 2-d array of data scale - scaling value of the data. higher number = lower gradient azdeg - where the light comes from: 0 south ; 90 east ; 180 north ; 270 west altdeg - where the light comes from: 0 horison ; 90 zenith output: a 2-d array of normalized hilshade ''' from pylab import pi, gradient, arctan, hypot, arctan2, sin, cos # convert alt, az to radians az = azdeg*pi/180.0 alt = altdeg*pi/180.0 # gradient in x and y directions dx, dy = gradient(data/float(scale)) slope = 0.5*pi - arctan(hypot(dx, dy)) aspect = arctan2(dx, dy) intensity = sin(alt)*sin(slope) + cos(alt)*cos(slope)*cos(-az - aspect - 0.5*pi) intensity = (intensity - intensity.min())/(intensity.max() - intensity.min()) return intensity
def hillshade(data, scale=10.0, azdeg=165.0, altdeg=45.0): ''' Convert data to hillshade based on matplotlib.colors.LightSource class. Args: data - a 2-d array of data scale - scaling value of the data. higher number = lower gradient azdeg - where the light comes from: 0 south ; 90 east ; 180 north ; 270 west altdeg - where the light comes from: 0 horison ; 90 zenith Returns: a 2-d array of normalized hilshade ''' # convert alt, az to radians az = azdeg*pi/180.0 alt = altdeg*pi/180.0 # gradient in x and y directions dx, dy = gradient(data/float(scale)) slope = 0.5*pi - arctan(hypot(dx, dy)) aspect = arctan2(dx, dy) az = -az - aspect - 0.5*pi intensity = sin(alt)*sin(slope) + cos(alt)*cos(slope)*cos(az) mi, ma = intensity.min(), intensity.max() intensity = (intensity - mi)/(ma - mi) return intensity
def findNoisyArea(self): """ Argument : - None Return : - None Use to detect the noisy area (the area without atoms) by an edge detection technique. """ OD = self.computeFirstOD() yProfile = py.sum(OD, axis=1) derivative = py.gradient(yProfile) N = 10 # because the derivative is usually very noisy, a sliding average is # performed in order to smooth the signal. This is done by a # convolution with a gate function of size "N". res = py.convolve(derivative, py.ones((N, )) / N, mode='valid') mean = res.mean() # index of the maximum value of the signal. i = res.argmax() while res[i] >= mean: # once "i" is greater or equal to the mean of the derivative, # we have found the upper bound of the noisy area. i -= 1 # index of the minimum value of the signal. upBound = i - 50 i = res.argmin() while res[i] < mean: # once "i" is smaller or equal to the mean of the derivative, # we have found the lower bound of the noisy area. i += 1 downBound = i + 50 self.setNoisyArea((upBound, downBound)) return OD
def get_kin(self, fps=250.): """ returns a list of the selected kinematics (one list item for each repetition) :args: self: kin object fps (float, default 250): sampling frequency. Required to correctly compute the velocities. :returns: a list. Each element contains the selected (-> self.selection) data with corresponding velocities (i.e. 2d x n elements per item) """ # walk through each element of "selection" all_pos = [] all_vel = [] for raw in self.raw_dat: curr_pos = [] curr_vel = [] for elem in self.selection: items = [x.strip() for x in elem.split('-')] # 1 item if no "-" present dims = [] markers = [] for item in items: if item.endswith('_x'): dims.append(0) elif item.endswith('_y'): dims.append(1) elif item.endswith('_z'): dims.append(2) else: print "invalid marker suffix: ", item continue markers.append(item[:-2]) if len(items) == 1: # single marker curr_pos.append(raw[markers[0]][:, dims[0]]) curr_vel.append(gradient(raw[markers[0]][:, dims[0]]) * fps) else: # difference between two markers curr_pos.append(raw[markers[0]][:, dims[0]] - raw[markers[1]][:, dims[1]]) curr_vel.append(gradient(raw[markers[0]][:, dims[0]] - raw[markers[1]][:, dims[1]]) * fps) all_pos.append(vstack(curr_pos + curr_vel)) all_vel.append(vstack(curr_vel)) return all_pos
def findAreaOfInterest(self): """ Argument : - None Return : - None Use to detect the AOI by an edge detection technique. """ OD = self.findNoisyArea() max = 0 bestAngle = None for angle in range(-10, 10, 1): # the best angle is the one for wich the maximum peak is reached # on the yProfile (integral on the horizontal direction) ... image = rotate(OD, angle) yProfile = py.sum(image, axis=1) newMax = yProfile.max() if newMax > max: max = newMax bestAngle = angle # ... once found, the resulting OD image is kept and used to find the # the top and bottom bounds by edge detection. bestOD = rotate(OD, bestAngle) YProfile = py.sum(bestOD, axis=1) derivative = py.gradient(YProfile) N = 10 # because the derivative is usually very noisy, a sliding average is # performed in order to smooth the signal. This is done by a # convolution with a gate function of size "N". res = py.convolve(derivative, py.ones((N, )) / N, mode='valid') mean = res.mean() # index of the maximum value of the signal. i = res.argmax() while res[i] > mean: # once "i" is greater or equal to the mean of the derivative, # we have found the upper bound of the AOI. i -= 1 # for security we take an extra 50 pixels. y0 = int(i - 50) # index of the minimum value of the signal. i = res.argmin() while res[i] < mean: # once "i" is smaller or equal to the mean of the derivative, # we have found the lower bound of the AOI. i += 1 # Again, for security, we take an extra 50 pixels. y1 = int(i + 50) # The horizontal bound are taken to be maximal, but for security, we # take an extra 50 pixels. x0 = 50 x1 = py.shape(OD)[0] - 50 self.setAreaOfInterest(area=(x0, x1, y0, y1), angle=bestAngle) return bestOD[y0:y1, x0:x1]
ax.set_ylim([20,0]) ax.set_xlabel('Distance (km)', fontsize=14) ax.set_ylabel('Distance (km)', fontsize=14) im.set_clim([0,psv.max()]) # ax.set_xlabel('Distance (km)', fontsize=14) # ax.set_ylabel('Distance (km)', fontsize=14) print 'mean psv: ', np.mean(psv[np.where(psv > 0.01)]) print 'max psv: ', psv.max() np.savefig('psv1.pdf') # plot vrup material = np.loadtxt('bbp1d_1250_dx_25.asc') vs = material[:,2] cs = ml.repmat(vs[:-1],nx,1).T * 1e3 trup_ma = np.ma.masked_values(trup,1e9) gy, gx = np.absolute(np.gradient(trup_ma)) ttime = np.sqrt(gy**2 + gx**2) vrup = dx / ttime fig = plt.figure() ax = fig.gca() im = ax.imshow(vrup/cs, extent=(0, ex, 0, ez), origin='normal', cmap='viridis') print len(np.where(vrup/cs >= 1.0)[0]) print 'median vrup: ', np.median(vrup.ravel()) # plt.scatter(1201*dx*1e-3,401*dx*1e-3, marker='*', s=150, color='k') divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cbar = np.colorbar(im, cax=cax) cbar.solids.set_rasterized(True) cbar.set_label(label=r'$V_{rup}/c_s$', size=14) cbar.solids.set_edgecolor("face")
def analyse_equil ( F, R, Z): s = numpy.shape(F) nx = s[0] ny = s[1] #;;;;;;;;;;;;;;; Find critical points ;;;;;;;;;;;;; # # Need to find starting locations for O-points (minima/maxima) # and X-points (saddle points) # Rr=numpy.tile(R,nx).reshape(nx,ny).T # needed for contour Zz=numpy.tile(Z,ny).reshape(nx,ny) contour1=contour(Rr,Zz,gradient(F)[0], levels=[0.0], colors='r') contour2=contour(Rr,Zz,gradient(F)[1], levels=[0.0], colors='r') draw() ### --- line crossings --------------------------- res=find_inter( contour1, contour2) rex=numpy.interp(res[0], R, numpy.arange(R.size)).astype(int) zex=numpy.interp(res[1], Z, numpy.arange(Z.size)).astype(int) w=numpy.where((rex > 2) & (rex < nx-3) & (zex >2) & (zex < nx-3)) nextrema = numpy.size(w) rex=rex[w].flatten() zex=zex[w].flatten() #;;;;;;;;;;;;;; Characterise extrema ;;;;;;;;;;;;;;;;; # Fit a surface through local points using 6x6 matrix # This is to determine the type of extrema, and to # refine the location # n_opoint = 0 n_xpoint = 0 # Calculate inverse matrix rio = numpy.array([-1, 0, 0, 0, 1, 1]) # R index offsets zio = numpy.array([ 0,-1, 0, 1, 0, 1]) # Z index offsets # Fitting a + br + cz + drz + er^2 + fz^2 A = numpy.transpose([[numpy.zeros(6,numpy.int)+1], [rio], [zio], [rio*zio], [rio**2], [zio**2]]) A=A[:,0,:] for e in range (nextrema) : # Fit in index space so result is index number print("Critical point "+str(e)) valid = 1 localf = numpy.zeros(6) for i in range (6) : # Get the f value in a stencil around this point xi = (rex[e]+rio[i]) #> 0) < (nx-1) # Zero-gradient at edges yi = (zex[e]+zio[i]) #> 0) < (ny-1) localf[i] = F[xi, yi] res, _, _, _ = numpy.linalg.lstsq(A,localf) # Res now contains [a,b,c,d,e,f] # [0,1,2,3,4,5] # This determines whether saddle or extremum det = 4.*res[4]*res[5] - res[3]**2 if det < 0.0 : print(" X-point") else: print(" O-point") # Get location (2x2 matrix of coefficients) rinew = old_div((res[3]*res[2] - 2.*res[1]*res[5]), det) zinew = old_div((res[3]*res[1] - 2.*res[4]*res[2]), det) if (numpy.abs(rinew) > 1.) or (numpy.abs(zinew) > 1.0) : #; Method has gone slightly wrong. Try a different method. #; Get a contour line starting at this point. Should #; produce a circle around the real o-point. print(" Fitted location deviates too much") if det < 0.0 : print(" => X-point probably not valid") print(" deviation = "+numpy.str(rinew)+","+numpy.str(zinew)) #valid = 0 # else: # # contour_lines, F, findgen(nx), findgen(ny), levels=[F[rex[e], zex[e]]], \ # path_info=info, path_xy=xy # # if np.size(info) > 1 : # # More than one contour. Select the one closest # ind = closest_line(info, xy, rex[e], zex[e]) # info = info[ind] # else: info = info[0] # # rinew = 0.5*(MAX(xy[0, info.offset:(info.offset + info.n - 1)]) + \ # MIN(xy[0, info.offset:(info.offset + info.n - 1)])) - rex[e] # zinew = 0.5*(MAX(xy[1, info.offset:(info.offset + info.n - 1)]) + \ # MIN(xy[1, info.offset:(info.offset + info.n - 1)])) - zex[e] # # if (np.abs(rinew) > 2.) or (np.abs(zinew) > 2.0) : # print " Backup method also failed. Keeping initial guess" # rinew = 0. # zinew = 0. # # if valid : fnew = (res[0] + res[1]*rinew + res[2]*zinew + res[3]*rinew*zinew + res[4]*rinew**2 + res[5]*zinew**2) rinew = rinew + rex[e] zinew = zinew + zex[e] print(" Starting index: " + str(rex[e])+", "+str(zex[e])) print(" Refined index: " + str(rinew)+", "+str(zinew)) x=numpy.arange(numpy.size(R)) y=numpy.arange(numpy.size(Z)) rnew = numpy.interp(rinew,x,R) znew = numpy.interp(zinew, y, Z) print(" Position: " + str(rnew)+", "+str(znew)) print(" F = "+str(fnew)) if det < 0.0 : if n_xpoint == 0 : xpt_ri = [rinew] xpt_zi = [zinew] xpt_f = [fnew] n_xpoint = n_xpoint + 1 else: # Check if this duplicates an existing point if rinew in xpt_ri and zinew in xpt_zi : print(" Duplicates existing X-point.") else: xpt_ri = numpy.append(xpt_ri, rinew) xpt_zi = numpy.append(xpt_zi, zinew) xpt_f = numpy.append(xpt_f, fnew) n_xpoint = n_xpoint + 1 scatter(rnew,znew,s=100, marker='x', color='r') annotate(numpy.str(n_xpoint-1), xy = (rnew, znew), xytext = (10, 10), textcoords = 'offset points',size='large', color='r') draw() else: if n_opoint == 0 : opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 else: # Check if this duplicates an existing point if rinew in opt_ri and zinew in opt_zi : print(" Duplicates existing O-point") else: opt_ri = numpy.append(opt_ri, rinew) opt_zi = numpy.append(opt_zi, zinew) opt_f = numpy.append(opt_f, fnew) n_opoint = n_opoint + 1 scatter(rnew,znew,s=100, marker='o',color='r') annotate(numpy.str(n_opoint-1), xy = (rnew, znew), xytext = (10, 10), textcoords = 'offset points', size='large', color='b') draw() print("Number of O-points: "+numpy.str(n_opoint)) print("Number of X-points: "+numpy.str(n_xpoint)) if n_opoint == 0 : opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 print("Number of O-points: "+str(n_opoint)) if n_opoint == 0 : print("No O-points! Giving up on this equilibrium") return Bunch(n_opoint=0, n_xpoint=0, primary_opt=-1) #;;;;;;;;;;;;;; Find plasma centre ;;;;;;;;;;;;;;;;;;; # Find the O-point closest to the middle of the grid mind = (opt_ri[0] - (old_div(numpy.float(nx),2.)))**2 + (opt_zi[0] - (old_div(numpy.float(ny),2.)))**2 ind = 0 for i in range (1, n_opoint) : d = (opt_ri[i] - (old_div(numpy.float(nx),2.)))**2 + (opt_zi[i] - (old_div(numpy.float(ny),2.)))**2 if d < mind : ind = i mind = d primary_opt = ind print("Primary O-point is at "+str(numpy.interp(opt_ri[ind],x,R)) + ", " + str(numpy.interp(opt_zi[ind],y,Z))) print("") if n_xpoint > 0 : # Find the primary separatrix # First remove non-monotonic separatrices nkeep = 0 for i in range (n_xpoint) : # Draw a line between the O-point and X-point n = 100 # Number of points farr = numpy.zeros(n) dr = old_div((xpt_ri[i] - opt_ri[ind]), numpy.float(n)) dz = old_div((xpt_zi[i] - opt_zi[ind]), numpy.float(n)) for j in range (n) : # interpolate f at this location func = RectBivariateSpline(x, y, F) farr[j] = func(opt_ri[ind] + dr*numpy.float(j), opt_zi[ind] + dz*numpy.float(j)) # farr should be monotonic, and shouldn't cross any other separatrices maxind = numpy.argmax(farr) minind = numpy.argmin(farr) if (maxind < minind) : maxind, minind = minind, maxind # Allow a little leeway to account for errors # NOTE: This needs a bit of refining if (maxind > (n-3)) and (minind < 3) : # Monotonic, so add this to a list of x-points to keep if nkeep == 0 : keep = [i] else: keep = numpy.append(keep, i) nkeep = nkeep + 1 if nkeep > 0 : print("Keeping x-points ", keep) xpt_ri = xpt_ri[keep] xpt_zi = xpt_zi[keep] xpt_f = xpt_f[keep] else: "No x-points kept" n_xpoint = nkeep # Now find x-point closest to primary O-point s = numpy.argsort(numpy.abs(opt_f[ind] - xpt_f)) xpt_ri = xpt_ri[s] xpt_zi = xpt_zi[s] xpt_f = xpt_f[s] inner_sep = 0 else: # No x-points. Pick mid-point in f xpt_f = 0.5*(numpy.max(F) + numpy.min(F)) print("WARNING: No X-points. Setting separatrix to F = "+str(xpt_f)) xpt_ri = 0 xpt_zi = 0 inner_sep = 0 #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Put results into a structure result = Bunch(n_opoint=n_opoint, n_xpoint=n_xpoint, # Number of O- and X-points primary_opt=primary_opt, # Which O-point is the plasma centre inner_sep=inner_sep, #Innermost X-point separatrix opt_ri=opt_ri, opt_zi=opt_zi, opt_f=opt_f, # O-point location (indices) and psi values xpt_ri=xpt_ri, xpt_zi=xpt_zi, xpt_f=xpt_f) # X-point locations and psi values return result
def analyse_equil ( F, R, Z): s = numpy.shape(F) nx = s[0] ny = s[1] #;;;;;;;;;;;;;;; Find critical points ;;;;;;;;;;;;; # # Need to find starting locations for O-points (minima/maxima) # and X-points (saddle points) # Rr=numpy.tile(R,nx).reshape(nx,ny).T Zz=numpy.tile(Z,ny).reshape(nx,ny) contour1=contour(Rr,Zz,gradient(F)[0], levels=[0.0], colors='r') contour2=contour(Rr,Zz,gradient(F)[1], levels=[0.0], colors='r') draw() ### 1st method - line crossings --------------------------- res=find_inter( contour1, contour2) #rex1=numpy.interp(res[0], R, numpy.arange(R.size)).astype(int) #zex1=numpy.interp(res[1], Z, numpy.arange(Z.size)).astype(int) rex1=res[0] zex1=res[1] w=numpy.where((rex1 > R[2]) & (rex1 < R[nx-3]) & (zex1 > Z[2]) & (zex1 < Z[nx-3])) nextrema = numpy.size(w) rex1=rex1[w].flatten() zex1=zex1[w].flatten() ### 2nd method - local maxima_minima ----------------------- res1=local_min_max.detect_local_minima(F) res2=local_min_max.detect_local_maxima(F) res=numpy.append(res1,res2,1) rex2=res[0,:].flatten() zex2=res[1,:].flatten() w=numpy.where((rex2 > 2) & (rex2 < nx-3) & (zex2 >2) & (zex2 < nx-3)) nextrema = numpy.size(w) rex2=rex2[w].flatten() zex2=zex2[w].flatten() n_opoint=nextrema n_xpoint=numpy.size(rex1)-n_opoint # Needed for interp below Rx=numpy.arange(numpy.size(R)) Zx=numpy.arange(numpy.size(Z)) print "Number of O-points: "+numpy.str(n_opoint) print "Number of X-points: "+numpy.str(n_xpoint) # Deduce the O & X points x=R[rex2] y=Z[zex2] dr=(R[numpy.size(R)-1]-R[0])/numpy.size(R) dz=(Z[numpy.size(Z)-1]-Z[0])/numpy.size(Z) repeated=set() for i in xrange(numpy.size(rex1)): for j in xrange(numpy.size(x)): if numpy.abs(rex1[i]-x[j]) < 2*dr and numpy.abs(zex1[i]-y[j]) < 2*dz : repeated.add(i) # o-points o_ri=numpy.take(rex1,numpy.array(list(repeated))) opt_ri=numpy.interp(o_ri,R,Rx) o_zi=numpy.take(zex1,numpy.array(list(repeated))) opt_zi=numpy.interp(o_zi,Z,Zx) opt_f=numpy.zeros(numpy.size(opt_ri)) func = RectBivariateSpline(Rx, Zx, F) for i in xrange(numpy.size(opt_ri)): opt_f[i]=func(opt_ri[i], opt_zi[i]) n_opoint=numpy.size(opt_ri) # x-points x_ri=numpy.delete(rex1, numpy.array(list(repeated))) xpt_ri=numpy.interp(x_ri,R,Rx) x_zi=numpy.delete(zex1, numpy.array(list(repeated))) xpt_zi=numpy.interp(x_zi,Z,Zx) xpt_f=numpy.zeros(numpy.size(xpt_ri)) func = RectBivariateSpline(Rx, Zx, F) for i in xrange(numpy.size(xpt_ri)): xpt_f[i]=func(xpt_ri[i], xpt_zi[i]) n_xpoint=numpy.size(xpt_ri) # plot o-points plot(o_ri,o_zi,'o', markersize=10) labels = ['{0}'.format(i) for i in range(o_ri.size)] for label, xp, yp in zip(labels, o_ri, o_zi): annotate(label, xy = (xp, yp), xytext = (10, 10), textcoords = 'offset points',size='large', color='b') draw() # plot x-points plot(x_ri,x_zi,'x', markersize=10) labels = ['{0}'.format(i) for i in range(x_ri.size)] for label, xp, yp in zip(labels, x_ri, x_zi): annotate(label, xy = (xp, yp), xytext = (10, 10), textcoords = 'offset points',size='large', color='r') draw() print "Number of O-points: "+str(n_opoint) if n_opoint == 0 : print "No O-points! Giving up on this equilibrium" return Bunch(n_opoint=0, n_xpoint=0, primary_opt=-1) #;;;;;;;;;;;;;; Find plasma centre ;;;;;;;;;;;;;;;;;;; # Find the O-point closest to the middle of the grid mind = (opt_ri[0] - (numpy.float(nx)/2.))**2 + (opt_zi[0] - (numpy.float(ny)/2.))**2 ind = 0 for i in range (1, n_opoint) : d = (opt_ri[i] - (numpy.float(nx)/2.))**2 + (opt_zi[i] - (numpy.float(ny)/2.))**2 if d < mind : ind = i mind = d primary_opt = ind print "Primary O-point is at "+ numpy.str(numpy.interp(opt_ri[ind],numpy.arange(numpy.size(R)),R)) + ", " + numpy.str(numpy.interp(opt_zi[ind],numpy.arange(numpy.size(Z)),Z)) print "" if n_xpoint > 0 : # Find the primary separatrix # First remove non-monotonic separatrices nkeep = 0 for i in xrange (n_xpoint) : # Draw a line between the O-point and X-point n = 100 # Number of points farr = numpy.zeros(n) dr = (xpt_ri[i] - opt_ri[ind]) / numpy.float(n) dz = (xpt_zi[i] - opt_zi[ind]) / numpy.float(n) for j in xrange (n) : # interpolate f at this location func = RectBivariateSpline(Rx, Zx, F) farr[j] = func(opt_ri[ind] + dr*numpy.float(j), opt_zi[ind] + dz*numpy.float(j)) # farr should be monotonic, and shouldn't cross any other separatrices maxind = numpy.argmax(farr) minind = numpy.argmin(farr) if (maxind < minind) : maxind, minind = minind, maxind # Allow a little leeway to account for errors # NOTE: This needs a bit of refining if (maxind > (n-3)) and (minind < 3) : # Monotonic, so add this to a list of x-points to keep if nkeep == 0 : keep = [i] else: keep = numpy.append(keep, i) nkeep = nkeep + 1 if nkeep > 0 : print "Keeping x-points ", keep xpt_ri = xpt_ri[keep] xpt_zi = xpt_zi[keep] xpt_f = xpt_f[keep] else: "No x-points kept" n_xpoint = nkeep # Now find x-point closest to primary O-point s = numpy.argsort(numpy.abs(opt_f[ind] - xpt_f)) xpt_ri = xpt_ri[s] xpt_zi = xpt_zi[s] xpt_f = xpt_f[s] inner_sep = 0 else: # No x-points. Pick mid-point in f xpt_f = 0.5*(numpy.max(F) + numpy.min(F)) print "WARNING: No X-points. Setting separatrix to F = "+str(xpt_f) xpt_ri = 0 xpt_zi = 0 inner_sep = 0 #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Put results into a structure result = Bunch(n_opoint=n_opoint, n_xpoint=n_xpoint, # Number of O- and X-points primary_opt=primary_opt, # Which O-point is the plasma centre inner_sep=inner_sep, #Innermost X-point separatrix opt_ri=opt_ri, opt_zi=opt_zi, opt_f=opt_f, # O-point location (indices) and psi values xpt_ri=xpt_ri, xpt_zi=xpt_zi, xpt_f=xpt_f) # X-point locations and psi values return result
def ComputeCoM(kinData, forceData, kin_est=None, mass=None, f_thresh=1., steepness=10., use_Fint=False, return_mass=False, adapt_kin_mean=True): """ Computes the CoM motion using a complementary filter as described in Maus et al, J Exp. Biol. (2011). :args: kinData: dict or mutils.io.saveable object with kinematic data. Should contain "fs" field forceData: dict or mutils.io.saveable object with force data. Should contain "fs" field. kin_est (optional, d-by-1 array): if present, the kinematic estimate of the CoM (x,y,z direction). Overrides missing kinData (can be empty, i.e. {}, then). Data will be (Fourier-)interpolated, so can have a different sampling frequency than the force data. mass (float): the subject's mass. If omitted, it is automatically determined. f_thresh (float): threshold frequency (recommended: slightly below dominant frequency) steepness (float): steepness of the frequency threshold (1/Hz). use_Fint (bool, default=False): whether or not to use a fourier-based integration scheme (zero-lag) return_mass(bool): whether or not to additionally return the mass adapt_kin_mean (bool): wether or not to adapt the combined mean to the kinematic mean :returns: Force, CoM: Arrays which contain physically consistent GRF and CoM data. """ if type(kinData) == dict: kd = mio.saveable(kinData) elif type(kinData) == mio.saveable: kd = kinData if type(forceData) == dict: fd = mio.saveable(forceData) elif type(forceData) == mio.saveable: fd = forceData if fd.fs < 1: warnings.warn( "warning: fs is the sampling FREQUENCY, not the sampling time" + "\nAre you sure your sampling frequency is really that low?") def weighting_function(freq_vec, f_changeover, sharpness): """ another weighting function, using tanh to prevent exp. overflow """ weight = .5 - .5 * tanh( (freq_vec.squeeze()[1:] - f_changeover) * sharpness) weight = (weight + weight[::-1]).squeeze() weight = np.hstack([1., weight]) return weight def kin_estimate(selectedData): """ calculates the kinematic CoM estimate from "selectedData" """ # anthropometry from Dempster # format: [prox.marker, dist. marker, rel. weight (%), # segment CoM pos rel. to prox. marker (%) ] aData = [('R_Hea', 'L_Hea', 8.26, 50.), ('L_Acr', 'R_Trc', 46.84 / 2., 63.), ('R_Acr', 'L_Trc', 46.84 / 2., 63.), ('R_Acr', 'R_Elb', 3.25, 43.6), ('R_Elb', 'R_WrL', 1.87 + 0.65, 43. + 25.), ('L_Acr', 'L_Elb', 3.25, 43.6), ('L_Elb', 'L_WrL', 1.87 + 0.65, 43. + 25.), ('R_Trc', 'R_Kne', 10.5, 43.3), ('R_Kne', 'R_AnL', 4.75, 43.4), ('R_Hee', 'R_Mtv', 1.43, 50. + 5.), ('L_Trc', 'L_Kne', 10.5, 43.3), ('L_Kne', 'L_AnL', 4.75, 43.4), ('L_Hee', 'L_Mtv', 1.43, 50. + 5.)] # adaptation to dataformat when extracted from database: lowercase aData = [(x[0].lower(), x[1].lower(), x[2], x[3]) for x in aData] CoM = np.zeros((len(kd.sacr[:, 0]), 3)) for segment in aData: CoM += segment[2] / 100. * ( (getattr(selectedData, segment[1]) - getattr(selectedData, segment[0])) * segment[3] / 100. + getattr(selectedData, segment[0])) return CoM elems = dir(fd) # get convenient names for forces if 'fx' in elems: Fx = fd.fx.squeeze() elif 'Fx' in elems: Fx = fd.Fx.squeeze() elif 'fx1' in elems and 'fx2' in elems: Fx = (fd.fx1 + fd.fx2).squeeze() else: raise ValueError("Error: Fx field not in forces (fx or fx1 and fx2)") if 'fy' in elems: Fy = fd.fy.squeeze() elif 'Fy' in elems: Fy = fd.Fy.squeeze() elif 'fy1' in elems and 'fy2' in elems: Fy = (fd.fy1 + fd.fy2).squeeze() else: raise ValueError("Error: Fy field not in forces (fy or fy1 and fy2)") if 'fz' in elems: Fz = fd.fz.squeeze() elif 'Fz' in elems: Fz = fd.Fz.squeeze() elif all([x in elems for x in ['fz1', 'fz2', 'fz3', 'fz4']]): Fz = (fd.fz1 + fd.fz2 + fd.fz3 + fd.fz4).squeeze() elif all([ x in elems for x in [ 'fzr1', 'fzr2', 'fzr3', 'fzr4', 'fzl1', 'fzl2', 'fzl3', 'fzl4', ] ]): Fz = (fd.fzr1 + fd.fzr2 + fd.fzr3 + fd.fzr4 + fd.fzl1 + fd.fzl2 + fd.fzl3 + fd.fzl4).squeeze() # remove bodyweight later else: raise ValueError( "Error: Fz field not in forces (fz or fz1..4 or fzr1..4 and fzl1...4)" ) if kin_est is None: kin_est = kin_estimate(kd) kin_est_i = vstack([ mi.interp_f(kin_est[:, col].copy(), len(Fz), detrend=True) for col in range(3) ]).T if use_Fint: diff_op = lambda x: mi.diff_f(x.squeeze(), fs=fd.fs) int_op = lambda x, x0: mi.int_f(x.squeeze(), fs=fd.fs) + x0 else: diff_op = lambda x: gradient(x.squeeze()) * fd.fs int_op = lambda x, x0: cumtrapz(x.squeeze(), initial=0) / fd.fs + x0 vd = vstack([diff_op(kin_est_i[:, col]) for col in range(3)]).T # correct mean forces, pt 1 n = 2 # for higher reliably use multiple points T = len(Fz) / fd.fs dvx = ((kin_est_i[-1, 0] - kin_est_i[-(n + 1), 0]) - (kin_est_i[n, 0] - kin_est_i[0, 0])) * (fd.fs / n) dvy = ((kin_est_i[-1, 1] - kin_est_i[-(n + 1), 1]) - (kin_est_i[n, 1] - kin_est_i[0, 1])) * (fd.fs / n) dvz = ((kin_est_i[-1, 2] - kin_est_i[-(n + 1), 2]) - (kin_est_i[n, 2] - kin_est_i[0, 2])) * (fd.fs / n) ax = dvx / T ay = dvy / T az = dvz / T + 9.81 #print "dvz = ", dvz #print "az = ", az if mass == None: # determine mass mass = array(mean(Fz) / az).squeeze() Fz -= mass * 9.81 else: Fz -= mean(Fz) Fz += (az - 9.81) * mass Fx = Fx - mean(Fx) + ax * mass Fy = Fy - mean(Fy) + ay * mass #print "estimated mass:", mass vi = vstack( [int_op(F / mass, v0) for F, v0 in zip([Fx, Fy, Fz], vd[0, :])]).T spect_diff = vstack([fftpack.fft(vd[:, col]) for col in range(3)]).T spect_int = vstack([fftpack.fft(vi[:, col]) for col in range(3)]).T freq_vec = linspace(0, fd.fs, len(Fz), endpoint=False) wv = weighting_function(freq_vec, f_thresh, steepness) spect_combine = vstack([ wv * spect_diff[:, col] + (1. - wv) * spect_int[:, col] for col in range(3) ]).T v_combine = vstack( [fftpack.ifft(spect_combine[:, col]).real for col in range(3)]).T x_combine = vstack( [int_op(v_combine[:, col], kin_est[0, col]) for col in range(3)]).T if adapt_kin_mean: x_combine = x_combine - mean(x_combine, axis=0) for dim in range(x_combine.shape[1]): x_combine[:, dim] += mean(kin_est[:, dim]) f_combine = vstack([diff_op(v_combine[:, col]) * mass for col in range(3)]).T f_combine[:, 2] += mass * 9.81 if return_mass: return f_combine, x_combine, mass return f_combine, x_combine
def analyse_equil(F, R, Z): s = numpy.shape(F) nx = s[0] ny = s[1] #;;;;;;;;;;;;;;; Find critical points ;;;;;;;;;;;;; # # Need to find starting locations for O-points (minima/maxima) # and X-points (saddle points) # Rr = numpy.tile(R, nx).reshape(nx, ny).T # needed for contour Zz = numpy.tile(Z, ny).reshape(nx, ny) contour1 = contour(Rr, Zz, gradient(F)[0], levels=[0.0], colors='r') contour2 = contour(Rr, Zz, gradient(F)[1], levels=[0.0], colors='r') draw() ### --- line crossings --------------------------- res = find_inter(contour1, contour2) rex = numpy.interp(res[0], R, numpy.arange(R.size)).astype(int) zex = numpy.interp(res[1], Z, numpy.arange(Z.size)).astype(int) w = numpy.where((rex > 2) & (rex < nx - 3) & (zex > 2) & (zex < nx - 3)) nextrema = numpy.size(w) rex = rex[w].flatten() zex = zex[w].flatten() #;;;;;;;;;;;;;; Characterise extrema ;;;;;;;;;;;;;;;;; # Fit a surface through local points using 6x6 matrix # This is to determine the type of extrema, and to # refine the location # n_opoint = 0 n_xpoint = 0 # Calculate inverse matrix rio = numpy.array([-1, 0, 0, 0, 1, 1]) # R index offsets zio = numpy.array([0, -1, 0, 1, 0, 1]) # Z index offsets # Fitting a + br + cz + drz + er^2 + fz^2 A = numpy.transpose([[numpy.zeros(6, numpy.int) + 1], [rio], [zio], [rio * zio], [rio**2], [zio**2]]) A = A[:, 0, :] for e in range(nextrema): # Fit in index space so result is index number print("Critical point " + str(e)) valid = 1 localf = numpy.zeros(6) for i in range(6): # Get the f value in a stencil around this point xi = (rex[e] + rio[i]) #> 0) < (nx-1) # Zero-gradient at edges yi = (zex[e] + zio[i]) #> 0) < (ny-1) localf[i] = F[xi, yi] res, _, _, _ = numpy.linalg.lstsq(A, localf) # Res now contains [a,b,c,d,e,f] # [0,1,2,3,4,5] # This determines whether saddle or extremum det = 4. * res[4] * res[5] - res[3]**2 if det < 0.0: print(" X-point") else: print(" O-point") # Get location (2x2 matrix of coefficients) rinew = old_div((res[3] * res[2] - 2. * res[1] * res[5]), det) zinew = old_div((res[3] * res[1] - 2. * res[4] * res[2]), det) if (numpy.abs(rinew) > 1.) or (numpy.abs(zinew) > 1.0): #; Method has gone slightly wrong. Try a different method. #; Get a contour line starting at this point. Should #; produce a circle around the real o-point. print(" Fitted location deviates too much") if det < 0.0: print(" => X-point probably not valid") print(" deviation = " + numpy.str(rinew) + "," + numpy.str(zinew)) #valid = 0 # else: # # contour_lines, F, findgen(nx), findgen(ny), levels=[F[rex[e], zex[e]]], \ # path_info=info, path_xy=xy # # if np.size(info) > 1 : # # More than one contour. Select the one closest # ind = closest_line(info, xy, rex[e], zex[e]) # info = info[ind] # else: info = info[0] # # rinew = 0.5*(MAX(xy[0, info.offset:(info.offset + info.n - 1)]) + \ # MIN(xy[0, info.offset:(info.offset + info.n - 1)])) - rex[e] # zinew = 0.5*(MAX(xy[1, info.offset:(info.offset + info.n - 1)]) + \ # MIN(xy[1, info.offset:(info.offset + info.n - 1)])) - zex[e] # # if (np.abs(rinew) > 2.) or (np.abs(zinew) > 2.0) : # print " Backup method also failed. Keeping initial guess" # rinew = 0. # zinew = 0. # # if valid: fnew = (res[0] + res[1] * rinew + res[2] * zinew + res[3] * rinew * zinew + res[4] * rinew**2 + res[5] * zinew**2) rinew = rinew + rex[e] zinew = zinew + zex[e] print(" Starting index: " + str(rex[e]) + ", " + str(zex[e])) print(" Refined index: " + str(rinew) + ", " + str(zinew)) x = numpy.arange(numpy.size(R)) y = numpy.arange(numpy.size(Z)) rnew = numpy.interp(rinew, x, R) znew = numpy.interp(zinew, y, Z) print(" Position: " + str(rnew) + ", " + str(znew)) print(" F = " + str(fnew)) if det < 0.0: if n_xpoint == 0: xpt_ri = [rinew] xpt_zi = [zinew] xpt_f = [fnew] n_xpoint = n_xpoint + 1 else: # Check if this duplicates an existing point if rinew in xpt_ri and zinew in xpt_zi: print(" Duplicates existing X-point.") else: xpt_ri = numpy.append(xpt_ri, rinew) xpt_zi = numpy.append(xpt_zi, zinew) xpt_f = numpy.append(xpt_f, fnew) n_xpoint = n_xpoint + 1 scatter(rnew, znew, s=100, marker='x', color='r') annotate(numpy.str(n_xpoint - 1), xy=(rnew, znew), xytext=(10, 10), textcoords='offset points', size='large', color='r') draw() else: if n_opoint == 0: opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 else: # Check if this duplicates an existing point if rinew in opt_ri and zinew in opt_zi: print(" Duplicates existing O-point") else: opt_ri = numpy.append(opt_ri, rinew) opt_zi = numpy.append(opt_zi, zinew) opt_f = numpy.append(opt_f, fnew) n_opoint = n_opoint + 1 scatter(rnew, znew, s=100, marker='o', color='r') annotate(numpy.str(n_opoint - 1), xy=(rnew, znew), xytext=(10, 10), textcoords='offset points', size='large', color='b') draw() print("Number of O-points: " + numpy.str(n_opoint)) print("Number of X-points: " + numpy.str(n_xpoint)) if n_opoint == 0: opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 print("Number of O-points: " + str(n_opoint)) if n_opoint == 0: print("No O-points! Giving up on this equilibrium") return Bunch(n_opoint=0, n_xpoint=0, primary_opt=-1) #;;;;;;;;;;;;;; Find plasma centre ;;;;;;;;;;;;;;;;;;; # Find the O-point closest to the middle of the grid mind = (opt_ri[0] - (old_div(numpy.float(nx), 2.)))**2 + ( opt_zi[0] - (old_div(numpy.float(ny), 2.)))**2 ind = 0 for i in range(1, n_opoint): d = (opt_ri[i] - (old_div(numpy.float(nx), 2.)))**2 + ( opt_zi[i] - (old_div(numpy.float(ny), 2.)))**2 if d < mind: ind = i mind = d primary_opt = ind print("Primary O-point is at " + str(numpy.interp(opt_ri[ind], x, R)) + ", " + str(numpy.interp(opt_zi[ind], y, Z))) print("") if n_xpoint > 0: # Find the primary separatrix # First remove non-monotonic separatrices nkeep = 0 for i in range(n_xpoint): # Draw a line between the O-point and X-point n = 100 # Number of points farr = numpy.zeros(n) dr = old_div((xpt_ri[i] - opt_ri[ind]), numpy.float(n)) dz = old_div((xpt_zi[i] - opt_zi[ind]), numpy.float(n)) for j in range(n): # interpolate f at this location func = RectBivariateSpline(x, y, F) farr[j] = func(opt_ri[ind] + dr * numpy.float(j), opt_zi[ind] + dz * numpy.float(j)) # farr should be monotonic, and shouldn't cross any other separatrices maxind = numpy.argmax(farr) minind = numpy.argmin(farr) if (maxind < minind): maxind, minind = minind, maxind # Allow a little leeway to account for errors # NOTE: This needs a bit of refining if (maxind > (n - 3)) and (minind < 3): # Monotonic, so add this to a list of x-points to keep if nkeep == 0: keep = [i] else: keep = numpy.append(keep, i) nkeep = nkeep + 1 if nkeep > 0: print("Keeping x-points ", keep) xpt_ri = xpt_ri[keep] xpt_zi = xpt_zi[keep] xpt_f = xpt_f[keep] else: "No x-points kept" n_xpoint = nkeep # Now find x-point closest to primary O-point s = numpy.argsort(numpy.abs(opt_f[ind] - xpt_f)) xpt_ri = xpt_ri[s] xpt_zi = xpt_zi[s] xpt_f = xpt_f[s] inner_sep = 0 else: # No x-points. Pick mid-point in f xpt_f = 0.5 * (numpy.max(F) + numpy.min(F)) print("WARNING: No X-points. Setting separatrix to F = " + str(xpt_f)) xpt_ri = 0 xpt_zi = 0 inner_sep = 0 #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Put results into a structure result = Bunch( n_opoint=n_opoint, n_xpoint=n_xpoint, # Number of O- and X-points primary_opt=primary_opt, # Which O-point is the plasma centre inner_sep=inner_sep, #Innermost X-point separatrix opt_ri=opt_ri, opt_zi=opt_zi, opt_f=opt_f, # O-point location (indices) and psi values xpt_ri=xpt_ri, xpt_zi=xpt_zi, xpt_f=xpt_f) # X-point locations and psi values return result
def analyse_equil ( F, R, Z): s = numpy.shape(F) nx = s[0] ny = s[1] #;;;;;;;;;;;;;;; Find critical points ;;;;;;;;;;;;; # # Need to find starting locations for O-points (minima/maxima) # and X-points (saddle points) # Rr=numpy.tile(R,nx).reshape(nx,ny).T Zz=numpy.tile(Z,ny).reshape(nx,ny) contour1=contour(Rr,Zz,gradient(F)[0], levels=[0.0], colors='r') contour2=contour(Rr,Zz,gradient(F)[1], levels=[0.0], colors='r') draw() ### 1st method - line crossings --------------------------- res=find_inter( contour1, contour2) #rex1=numpy.interp(res[0], R, numpy.arange(R.size)).astype(int) #zex1=numpy.interp(res[1], Z, numpy.arange(Z.size)).astype(int) rex1=res[0] zex1=res[1] w=numpy.where((rex1 > R[2]) & (rex1 < R[nx-3]) & (zex1 > Z[2]) & (zex1 < Z[nx-3])) nextrema = numpy.size(w) rex1=rex1[w].flatten() zex1=zex1[w].flatten() ### 2nd method - local maxima_minima ----------------------- res1=local_min_max.detect_local_minima(F) res2=local_min_max.detect_local_maxima(F) res=numpy.append(res1,res2,1) rex2=res[0,:].flatten() zex2=res[1,:].flatten() w=numpy.where((rex2 > 2) & (rex2 < nx-3) & (zex2 >2) & (zex2 < nx-3)) nextrema = numpy.size(w) rex2=rex2[w].flatten() zex2=zex2[w].flatten() n_opoint=nextrema n_xpoint=numpy.size(rex1)-n_opoint # Needed for interp below Rx=numpy.arange(numpy.size(R)) Zx=numpy.arange(numpy.size(Z)) print("Number of O-points: "+numpy.str(n_opoint)) print("Number of X-points: "+numpy.str(n_xpoint)) # Deduce the O & X points x=R[rex2] y=Z[zex2] dr=old_div((R[numpy.size(R)-1]-R[0]),numpy.size(R)) dz=old_div((Z[numpy.size(Z)-1]-Z[0]),numpy.size(Z)) repeated=set() for i in range(numpy.size(rex1)): for j in range(numpy.size(x)): if numpy.abs(rex1[i]-x[j]) < 2*dr and numpy.abs(zex1[i]-y[j]) < 2*dz : repeated.add(i) # o-points o_ri=numpy.take(rex1,numpy.array(list(repeated))) opt_ri=numpy.interp(o_ri,R,Rx) o_zi=numpy.take(zex1,numpy.array(list(repeated))) opt_zi=numpy.interp(o_zi,Z,Zx) opt_f=numpy.zeros(numpy.size(opt_ri)) func = RectBivariateSpline(Rx, Zx, F) for i in range(numpy.size(opt_ri)): opt_f[i]=func(opt_ri[i], opt_zi[i]) n_opoint=numpy.size(opt_ri) # x-points x_ri=numpy.delete(rex1, numpy.array(list(repeated))) xpt_ri=numpy.interp(x_ri,R,Rx) x_zi=numpy.delete(zex1, numpy.array(list(repeated))) xpt_zi=numpy.interp(x_zi,Z,Zx) xpt_f=numpy.zeros(numpy.size(xpt_ri)) func = RectBivariateSpline(Rx, Zx, F) for i in range(numpy.size(xpt_ri)): xpt_f[i]=func(xpt_ri[i], xpt_zi[i]) n_xpoint=numpy.size(xpt_ri) # plot o-points plot(o_ri,o_zi,'o', markersize=10) labels = ['{0}'.format(i) for i in range(o_ri.size)] for label, xp, yp in zip(labels, o_ri, o_zi): annotate(label, xy = (xp, yp), xytext = (10, 10), textcoords = 'offset points',size='large', color='b') draw() # plot x-points plot(x_ri,x_zi,'x', markersize=10) labels = ['{0}'.format(i) for i in range(x_ri.size)] for label, xp, yp in zip(labels, x_ri, x_zi): annotate(label, xy = (xp, yp), xytext = (10, 10), textcoords = 'offset points',size='large', color='r') draw() print("Number of O-points: "+str(n_opoint)) if n_opoint == 0 : print("No O-points! Giving up on this equilibrium") return Bunch(n_opoint=0, n_xpoint=0, primary_opt=-1) #;;;;;;;;;;;;;; Find plasma centre ;;;;;;;;;;;;;;;;;;; # Find the O-point closest to the middle of the grid mind = (opt_ri[0] - (old_div(numpy.float(nx),2.)))**2 + (opt_zi[0] - (old_div(numpy.float(ny),2.)))**2 ind = 0 for i in range (1, n_opoint) : d = (opt_ri[i] - (old_div(numpy.float(nx),2.)))**2 + (opt_zi[i] - (old_div(numpy.float(ny),2.)))**2 if d < mind : ind = i mind = d primary_opt = ind print("Primary O-point is at "+ numpy.str(numpy.interp(opt_ri[ind],numpy.arange(numpy.size(R)),R)) + ", " + numpy.str(numpy.interp(opt_zi[ind],numpy.arange(numpy.size(Z)),Z))) print("") if n_xpoint > 0 : # Find the primary separatrix # First remove non-monotonic separatrices nkeep = 0 for i in range (n_xpoint) : # Draw a line between the O-point and X-point n = 100 # Number of points farr = numpy.zeros(n) dr = old_div((xpt_ri[i] - opt_ri[ind]), numpy.float(n)) dz = old_div((xpt_zi[i] - opt_zi[ind]), numpy.float(n)) for j in range (n) : # interpolate f at this location func = RectBivariateSpline(Rx, Zx, F) farr[j] = func(opt_ri[ind] + dr*numpy.float(j), opt_zi[ind] + dz*numpy.float(j)) # farr should be monotonic, and shouldn't cross any other separatrices maxind = numpy.argmax(farr) minind = numpy.argmin(farr) if (maxind < minind) : maxind, minind = minind, maxind # Allow a little leeway to account for errors # NOTE: This needs a bit of refining if (maxind > (n-3)) and (minind < 3) : # Monotonic, so add this to a list of x-points to keep if nkeep == 0 : keep = [i] else: keep = numpy.append(keep, i) nkeep = nkeep + 1 if nkeep > 0 : print("Keeping x-points ", keep) xpt_ri = xpt_ri[keep] xpt_zi = xpt_zi[keep] xpt_f = xpt_f[keep] else: "No x-points kept" n_xpoint = nkeep # Now find x-point closest to primary O-point s = numpy.argsort(numpy.abs(opt_f[ind] - xpt_f)) xpt_ri = xpt_ri[s] xpt_zi = xpt_zi[s] xpt_f = xpt_f[s] inner_sep = 0 else: # No x-points. Pick mid-point in f xpt_f = 0.5*(numpy.max(F) + numpy.min(F)) print("WARNING: No X-points. Setting separatrix to F = "+str(xpt_f)) xpt_ri = 0 xpt_zi = 0 inner_sep = 0 #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Put results into a structure result = Bunch(n_opoint=n_opoint, n_xpoint=n_xpoint, # Number of O- and X-points primary_opt=primary_opt, # Which O-point is the plasma centre inner_sep=inner_sep, #Innermost X-point separatrix opt_ri=opt_ri, opt_zi=opt_zi, opt_f=opt_f, # O-point location (indices) and psi values xpt_ri=xpt_ri, xpt_zi=xpt_zi, xpt_f=xpt_f) # X-point locations and psi values return result
ax.set_ylim([20, 0]) ax.set_xlabel('Distance (km)', fontsize=14) ax.set_ylabel('Distance (km)', fontsize=14) im.set_clim([0, psv.max()]) # ax.set_xlabel('Distance (km)', fontsize=14) # ax.set_ylabel('Distance (km)', fontsize=14) print 'mean psv: ', np.mean(psv[np.where(psv > 0.01)]) print 'max psv: ', psv.max() np.savefig('psv1.pdf') # plot vrup material = np.loadtxt('bbp1d_1250_dx_25.asc') vs = material[:, 2] cs = ml.repmat(vs[:-1], nx, 1).T * 1e3 trup_ma = np.ma.masked_values(trup, 1e9) gy, gx = np.absolute(np.gradient(trup_ma)) ttime = np.sqrt(gy**2 + gx**2) vrup = dx / ttime fig = plt.figure() ax = fig.gca() im = ax.imshow(vrup / cs, extent=(0, ex, 0, ez), origin='normal', cmap='viridis') print len(np.where(vrup / cs >= 1.0)[0]) print 'median vrup: ', np.median(vrup.ravel()) # plt.scatter(1201*dx*1e-3,401*dx*1e-3, marker='*', s=150, color='k') divider = make_axes_locatable(ax) cax = divider.append_axes("right", size="5%", pad=0.05) cbar = np.colorbar(im, cax=cax)
def make_1D(self, nps, phase='phi2', fps=250, phases_list=None): """ interpolate the data with <nps> frames per stride *Note* The <object>.selection attribute must contain a valid list of selection items. A selection item can either be a single coordinate (e.g. 'com_z', 'r_anl_y') or the difference between two coordinates (e.g. 'l_anl_y - com_y') :args: nps (int): how many frames per stride should be sampled phase (str): 'phi1' or 'phi2' : which phase to use. Note: previous data mostly used 'phi2'. fps (int): how many frames per second do the original data have? This is only required for correct scaling of the velocities. phases_list (list of list of int): If present, use given phases instead of predefined sections. Data will *not* be sampled exactly at the given phases; instead there will be _nps_ sections per stride, and strides start (on average only) at the given phases. """ if len(self.raw_dat) == 0: print "No data loaded." return if len(self.selection) == 0: print "No data selected. Set <object>.selection to something!" # gather phases for each trial: i_phases =[] if not phases_list: # typical call: omit first and last strides; # strides go from phase 0 to phase 2pi (exluding 2pi) for raw in self.raw_dat: phi = raw[phase].squeeze() # cut lower phase and upper phase by ~4pi each first_step = (phi[0] + 6. * pi) // (2. * pi) last_step = (phi[-1] - 4. * pi) // (2. * pi) i_phases.append(linspace(first_step, last_step + 1, (last_step - first_step + 1) * nps, endpoint=False) * 2. *pi) else: # phases are explicitely given avg_phasemod = mean([mean(mod(elem, 2.*pi)) for elem in phases_list]) for elem in phases_list: firstphase = (elem[0] // (2. * pi)) * 2.*pi + avg_phasemod lastphase = (elem[-1] // (2. * pi)) * 2.*pi + avg_phasemod i_phases.append(linspace(firstphase, lastphase + 2*pi, len(elem) * nps, endpoint=False)) self.i_phases = i_phases # walk through each element of "selection" all_pos = [] all_vel = [] for elem in self.selection: items = [x.strip() for x in elem.split('-')] # 1 item if no "-" present dims = [] markers = [] for item in items: if item.endswith('_x'): dims.append(0) elif item.endswith('_y'): dims.append(1) elif item.endswith('_z'): dims.append(2) else: print "invalid marker suffix: ", item continue markers.append(item[:-2]) all_elem_pos = [] all_elem_vel = [] for raw, phi in zip(self.raw_dat, self.i_phases): if len(items) == 1: # only a single marker dat = raw[markers[0]][:, dims[0]] else: # differences between two markers dat = raw[markers[0]][:, dims[0]] - raw[markers[1]][:, dims[1]] all_elem_pos.append(interp(phi, raw[phase].squeeze(), dat)) all_elem_vel.append(interp(phi, raw[phase].squeeze(), gradient(dat) * fps)) all_pos.append(hstack(all_elem_pos)) all_vel.append(hstack(all_elem_vel)) dat_2D = vstack([all_pos, all_vel]) return mi.twoD_oneD(dat_2D, nps)
def analyse_equil(F, R, Z): s = numpy.shape(F) nx = s[0] ny = s[1] # ;;;;;;;;;;;;;;; Find critical points ;;;;;;;;;;;;; # # Need to find starting locations for O-points (minima/maxima) # and X-points (saddle points) # Rr = numpy.tile(R, nx).reshape(nx, ny).T # needed for contour Zz = numpy.tile(Z, ny).reshape(nx, ny) contour1 = contour(Rr, Zz, gradient(F)[0], levels=[0.0], colors="r") contour2 = contour(Rr, Zz, gradient(F)[1], levels=[0.0], colors="r") draw() ### --- line crossings --------------------------- lines = find_inter(contour1, contour2) rex = numpy.interp(lines[0], R, numpy.arange(R.size)).astype(int) zex = numpy.interp(lines[1], Z, numpy.arange(Z.size)).astype(int) w = numpy.where((rex > 2) & (rex < nx - 3) & (zex > 2) & (zex < nx - 3)) nextrema = numpy.size(w) rex = rex[w].flatten() zex = zex[w].flatten() rp = lines[0][w] zp = lines[1][w] # ;;;;;;;;;;;;;; Characterise extrema ;;;;;;;;;;;;;;;;; # Fit a surface through local points using 6x6 matrix # This is to determine the type of extrema, and to # refine the location # n_opoint = 0 n_xpoint = 0 # Calculate inverse matrix rio = numpy.array([-1, 0, 0, 0, 1, 1]) # R index offsets zio = numpy.array([0, -1, 0, 1, 0, 1]) # Z index offsets # Fitting a + br + cz + drz + er^2 + fz^2 A = numpy.transpose([[numpy.zeros(6, numpy.int) + 1], [rio], [zio], [rio * zio], [rio ** 2], [zio ** 2]]) A = A[:, 0, :] # Needed for interp below Rx = numpy.linspace(R[0], R[numpy.size(R) - 1], numpy.size(R)) Zx = numpy.linspace(Z[0], Z[numpy.size(Z) - 1], numpy.size(Z)) for e in range(nextrema): # Fit in index space so result is index number print("Critical point " + str(e)) localf = numpy.zeros(6) for i in range(6): # Get the f value in a stencil around this point xi = rex[e] + rio[i] # > 0) < (nx-1) # Zero-gradient at edges yi = zex[e] + zio[i] # > 0) < (ny-1) localf[i] = F[xi, yi] res, _, _, _ = numpy.linalg.lstsq(A, localf) # Res now contains [a,b,c,d,e,f] # [0,1,2,3,4,5] # This determines whether saddle or extremum det = 4.0 * res[4] * res[5] - res[3] ** 2 if det < 0.0: print(" X-point") else: print(" O-point") rnew = rp[e] znew = zp[e] func = RectBivariateSpline(Rx, Zx, F) fnew = func(rnew, znew) print(rnew, znew, fnew) x = numpy.arange(numpy.size(R)) y = numpy.arange(numpy.size(Z)) rinew = numpy.interp(rnew, R, x) zinew = numpy.interp(znew, Z, y) print(" Position: " + str(rnew) + ", " + str(znew)) print(" F = " + str(fnew)) if det < 0.0: if n_xpoint == 0: xpt_ri = [rinew] xpt_zi = [zinew] xpt_f = [fnew] n_xpoint = n_xpoint + 1 else: # Check if this duplicates an existing point if rinew in xpt_ri and zinew in xpt_zi: print(" Duplicates existing X-point.") else: xpt_ri = numpy.append(xpt_ri, rinew) xpt_zi = numpy.append(xpt_zi, zinew) xpt_f = numpy.append(xpt_f, fnew) n_xpoint = n_xpoint + 1 scatter(rnew, znew, s=100, marker="x", color="r") annotate( numpy.str(n_xpoint - 1), xy=(rnew, znew), xytext=(10, 10), textcoords="offset points", size="large", color="r", ) draw() else: if n_opoint == 0: opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 else: # Check if this duplicates an existing point if rinew in opt_ri and zinew in opt_zi: print(" Duplicates existing O-point") else: opt_ri = numpy.append(opt_ri, rinew) opt_zi = numpy.append(opt_zi, zinew) opt_f = numpy.append(opt_f, fnew) n_opoint = n_opoint + 1 scatter(rnew, znew, s=100, marker="o", color="r") annotate( numpy.str(n_opoint - 1), xy=(rnew, znew), xytext=(10, 10), textcoords="offset points", size="large", color="b", ) draw() print("Number of O-points: " + numpy.str(n_opoint)) print("Number of X-points: " + numpy.str(n_xpoint)) if n_opoint == 0: opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 print("Number of O-points: " + str(n_opoint)) if n_opoint == 0: print("No O-points! Giving up on this equilibrium") return Bunch(n_opoint=0, n_xpoint=0, primary_opt=-1) # ;;;;;;;;;;;;;; Find plasma centre ;;;;;;;;;;;;;;;;;;; # Find the O-point closest to the middle of the grid mind = (opt_ri[0] - (old_div(numpy.float(nx), 2.0))) ** 2 + (opt_zi[0] - (old_div(numpy.float(ny), 2.0))) ** 2 ind = 0 for i in range(1, n_opoint): d = (opt_ri[i] - (old_div(numpy.float(nx), 2.0))) ** 2 + (opt_zi[i] - (old_div(numpy.float(ny), 2.0))) ** 2 if d < mind: ind = i mind = d primary_opt = ind print("Primary O-point is at " + str(numpy.interp(opt_ri[ind], x, R)) + ", " + str(numpy.interp(opt_zi[ind], y, Z))) print("") if n_xpoint > 0: # Find the primary separatrix # First remove non-monotonic separatrices nkeep = 0 for i in range(n_xpoint): # Draw a line between the O-point and X-point n = 100 # Number of points farr = numpy.zeros(n) dr = old_div((xpt_ri[i] - opt_ri[ind]), numpy.float(n)) dz = old_div((xpt_zi[i] - opt_zi[ind]), numpy.float(n)) for j in range(n): # interpolate f at this location func = RectBivariateSpline(x, y, F) farr[j] = func(opt_ri[ind] + dr * numpy.float(j), opt_zi[ind] + dz * numpy.float(j)) # farr should be monotonic, and shouldn't cross any other separatrices maxind = numpy.argmax(farr) minind = numpy.argmin(farr) if maxind < minind: maxind, minind = minind, maxind # Allow a little leeway to account for errors # NOTE: This needs a bit of refining if (maxind > (n - 3)) and (minind < 3): # Monotonic, so add this to a list of x-points to keep if nkeep == 0: keep = [i] else: keep = numpy.append(keep, i) nkeep = nkeep + 1 if nkeep > 0: print("Keeping x-points ", keep) xpt_ri = xpt_ri[keep] xpt_zi = xpt_zi[keep] xpt_f = xpt_f[keep] else: "No x-points kept" n_xpoint = nkeep # Now find x-point closest to primary O-point s = numpy.argsort(numpy.abs(opt_f[ind] - xpt_f)) xpt_ri = xpt_ri[s] xpt_zi = xpt_zi[s] xpt_f = xpt_f[s] inner_sep = 0 else: # No x-points. Pick mid-point in f xpt_f = 0.5 * (numpy.max(F) + numpy.min(F)) print("WARNING: No X-points. Setting separatrix to F = " + str(xpt_f)) xpt_ri = 0 xpt_zi = 0 inner_sep = 0 # ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Put results into a structure result = Bunch( n_opoint=n_opoint, n_xpoint=n_xpoint, # Number of O- and X-points primary_opt=primary_opt, # Which O-point is the plasma centre inner_sep=inner_sep, # Innermost X-point separatrix opt_ri=opt_ri, opt_zi=opt_zi, opt_f=opt_f, # O-point location (indices) and psi values xpt_ri=xpt_ri, xpt_zi=xpt_zi, xpt_f=xpt_f, ) # X-point locations and psi values return result
def analyse_equil ( F, R, Z): s = numpy.shape(F) nx = s[0] ny = s[1] #;;;;;;;;;;;;;;; Find critical points ;;;;;;;;;;;;; # # Need to find starting locations for O-points (minima/maxima) # and X-points (saddle points) # Rr=numpy.tile(R,ny).reshape(ny,nx).T # needed for contour Zz=numpy.tile(Z,nx).reshape(nx,ny) contour1=contour(Rr,Zz,gradient(F)[0], levels=[0.0], colors='r') contour2=contour(Rr,Zz,gradient(F)[1], levels=[0.0], colors='r') draw() ### --- line crossings --------------------------- lines=find_inter( contour1, contour2) rex=numpy.interp(lines[0], R, numpy.arange(R.size)).astype(int) zex=numpy.interp(lines[1], Z, numpy.arange(Z.size)).astype(int) w=numpy.where((rex > 2) & (rex < nx-3) & (zex >2) & (zex < nx-3)) nextrema = numpy.size(w) rex=rex[w].flatten() zex=zex[w].flatten() rp=lines[0][w] zp=lines[1][w] #;;;;;;;;;;;;;; Characterise extrema ;;;;;;;;;;;;;;;;; # Fit a surface through local points using 6x6 matrix # This is to determine the type of extrema, and to # refine the location # n_opoint = 0 n_xpoint = 0 # Calculate inverse matrix rio = numpy.array([-1, 0, 0, 0, 1, 1]) # R index offsets zio = numpy.array([ 0,-1, 0, 1, 0, 1]) # Z index offsets # Fitting a + br + cz + drz + er^2 + fz^2 A = numpy.transpose([[numpy.zeros(6,numpy.int)+1], [rio], [zio], [rio*zio], [rio**2], [zio**2]]) A=A[:,0,:] # Needed for interp below Rx=numpy.linspace(R[0],R[numpy.size(R)-1],numpy.size(R)) Zx=numpy.linspace(Z[0],Z[numpy.size(Z)-1],numpy.size(Z)) for e in range (nextrema) : # Fit in index space so result is index number print("Critical point "+str(e)) localf = numpy.zeros(6) for i in range (6) : # Get the f value in a stencil around this point xi = (rex[e]+rio[i]) #> 0) < (nx-1) # Zero-gradient at edges yi = (zex[e]+zio[i]) #> 0) < (ny-1) localf[i] = F[xi, yi] res, _, _, _ = numpy.linalg.lstsq(A,localf) # Res now contains [a,b,c,d,e,f] # [0,1,2,3,4,5] # This determines whether saddle or extremum det = 4.*res[4]*res[5] - res[3]**2 if det < 0.0 : print(" X-point") else: print(" O-point") rnew = rp[e] znew = zp[e] func = RectBivariateSpline(Rx, Zx, F) fnew=func(rnew, znew) print(rnew, znew, fnew) x=numpy.arange(numpy.size(R)) y=numpy.arange(numpy.size(Z)) rinew = numpy.interp(rnew,R,x) zinew = numpy.interp(znew, Z, y) print(" Position: " + str(rnew)+", "+str(znew)) print(" F = "+str(fnew)) if det < 0.0 : if n_xpoint == 0 : xpt_ri = [rinew] xpt_zi = [zinew] xpt_f = [fnew] n_xpoint = n_xpoint + 1 else: # Check if this duplicates an existing point if rinew in xpt_ri and zinew in xpt_zi : print(" Duplicates existing X-point.") else: xpt_ri = numpy.append(xpt_ri, rinew) xpt_zi = numpy.append(xpt_zi, zinew) xpt_f = numpy.append(xpt_f, fnew) n_xpoint = n_xpoint + 1 scatter(rnew,znew,s=100, marker='x', color='r') annotate(numpy.str(n_xpoint-1), xy = (rnew, znew), xytext = (10, 10), textcoords = 'offset points',size='large', color='r') draw() else: if n_opoint == 0 : opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 else: # Check if this duplicates an existing point if rinew in opt_ri and zinew in opt_zi : print(" Duplicates existing O-point") else: opt_ri = numpy.append(opt_ri, rinew) opt_zi = numpy.append(opt_zi, zinew) opt_f = numpy.append(opt_f, fnew) n_opoint = n_opoint + 1 scatter(rnew,znew,s=100, marker='o',color='r') annotate(numpy.str(n_opoint-1), xy = (rnew, znew), xytext = (10, 10), textcoords = 'offset points', size='large', color='b') draw() print("Number of O-points: "+numpy.str(n_opoint)) print("Number of X-points: "+numpy.str(n_xpoint)) if n_opoint == 0 : opt_ri = [rinew] opt_zi = [zinew] opt_f = [fnew] n_opoint = n_opoint + 1 print("Number of O-points: "+str(n_opoint)) if n_opoint == 0 : print("No O-points! Giving up on this equilibrium") return Bunch(n_opoint=0, n_xpoint=0, primary_opt=-1) #;;;;;;;;;;;;;; Find plasma centre ;;;;;;;;;;;;;;;;;;; # Find the O-point closest to the middle of the grid mind = (opt_ri[0] - (old_div(numpy.float(nx),2.)))**2 + (opt_zi[0] - (old_div(numpy.float(ny),2.)))**2 ind = 0 for i in range (1, n_opoint) : d = (opt_ri[i] - (old_div(numpy.float(nx),2.)))**2 + (opt_zi[i] - (old_div(numpy.float(ny),2.)))**2 if d < mind : ind = i mind = d primary_opt = ind print("Primary O-point is at "+str(numpy.interp(opt_ri[ind],x,R)) + ", " + str(numpy.interp(opt_zi[ind],y,Z))) print("") if n_xpoint > 0 : # Find the primary separatrix # First remove non-monotonic separatrices nkeep = 0 for i in range (n_xpoint) : # Draw a line between the O-point and X-point n = 100 # Number of points farr = numpy.zeros(n) dr = old_div((xpt_ri[i] - opt_ri[ind]), numpy.float(n)) dz = old_div((xpt_zi[i] - opt_zi[ind]), numpy.float(n)) for j in range (n) : # interpolate f at this location func = RectBivariateSpline(x, y, F) farr[j] = func(opt_ri[ind] + dr*numpy.float(j), opt_zi[ind] + dz*numpy.float(j)) # farr should be monotonic, and shouldn't cross any other separatrices maxind = numpy.argmax(farr) minind = numpy.argmin(farr) if (maxind < minind) : maxind, minind = minind, maxind # Allow a little leeway to account for errors # NOTE: This needs a bit of refining if (maxind > (n-3)) and (minind < 3) : # Monotonic, so add this to a list of x-points to keep if nkeep == 0 : keep = [i] else: keep = numpy.append(keep, i) nkeep = nkeep + 1 if nkeep > 0 : print("Keeping x-points ", keep) xpt_ri = xpt_ri[keep] xpt_zi = xpt_zi[keep] xpt_f = xpt_f[keep] else: "No x-points kept" n_xpoint = nkeep # Now find x-point closest to primary O-point s = numpy.argsort(numpy.abs(opt_f[ind] - xpt_f)) xpt_ri = xpt_ri[s] xpt_zi = xpt_zi[s] xpt_f = xpt_f[s] inner_sep = 0 else: # No x-points. Pick mid-point in f xpt_f = 0.5*(numpy.max(F) + numpy.min(F)) print("WARNING: No X-points. Setting separatrix to F = "+str(xpt_f)) xpt_ri = 0 xpt_zi = 0 inner_sep = 0 #;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;; # Put results into a structure result = Bunch(n_opoint=n_opoint, n_xpoint=n_xpoint, # Number of O- and X-points primary_opt=primary_opt, # Which O-point is the plasma centre inner_sep=inner_sep, #Innermost X-point separatrix opt_ri=opt_ri, opt_zi=opt_zi, opt_f=opt_f, # O-point location (indices) and psi values xpt_ri=xpt_ri, xpt_zi=xpt_zi, xpt_f=xpt_f) # X-point locations and psi values return result
pl.savefig(label + ".plots/" + label + ".fils_branches.png") #================================================================ mom1 = fits.getdata(mom1file) ly, lx = mom1.shape x, y = range(0, lx), range(0, ly) xi, yi = pl.meshgrid(x, y) wid = 9 from astropy.convolution import Gaussian2DKernel from astropy.convolution import convolve kernel = Gaussian2DKernel(stddev=wid / 2.354) vsmooth = convolve(mom1, kernel) vgrad = pl.gradient(vsmooth) vsmooth[pl.where(pl.isnan(mom1))] = pl.nan for i in [0, 1]: vgrad[i][pl.where(pl.isnan(mom1))] = pl.nan vgrad[i][pl.where(pl.absolute(vgrad[i]) > 0.1)] = pl.nan vv = (vgrad[0]**2 + vgrad[1]**2)**0.6 #pl.streamplot(xi, yi, vgrad[0], vgrad[1]) s = 7 # for MC s = 10 # for dor - finer pixellation skip = (slice(None, None, s), slice(None, None, s)) pl.clf() pl.quiver(xi[skip], yi[skip],
def ComputeCoM(kinData, forceData, kin_est=None, mass=None, f_thresh=1., steepness=10., use_Fint=False, return_mass=False, adapt_kin_mean=True): """ Computes the CoM motion using a complementary filter as described in Maus et al, J Exp. Biol. (2011). :args: kinData: dict or mutils.io.saveable object with kinematic data. Should contain "fs" field forceData: dict or mutils.io.saveable object with force data. Should contain "fs" field. kin_est (optional, d-by-1 array): if present, the kinematic estimate of the CoM (x,y,z direction). Overrides missing kinData (can be empty, i.e. {}, then). Data will be (Fourier-)interpolated, so can have a different sampling frequency than the force data. mass (float): the subject's mass. If omitted, it is automatically determined. f_thresh (float): threshold frequency (recommended: slightly below dominant frequency) steepness (float): steepness of the frequency threshold (1/Hz). use_Fint (bool, default=False): whether or not to use a fourier-based integration scheme (zero-lag) return_mass(bool): whether or not to additionally return the mass adapt_kin_mean (bool): wether or not to adapt the combined mean to the kinematic mean :returns: Force, CoM: Arrays which contain physically consistent GRF and CoM data. """ if type(kinData) == dict: kd = mio.saveable(kinData) elif type(kinData) == mio.saveable: kd = kinData if type(forceData) == dict: fd = mio.saveable(forceData) elif type(forceData) == mio.saveable: fd = forceData if fd.fs < 1: warnings.warn( "warning: fs is the sampling FREQUENCY, not the sampling time" + "\nAre you sure your sampling frequency is really that low?") def weighting_function(freq_vec, f_changeover, sharpness): """ another weighting function, using tanh to prevent exp. overflow """ weight = .5 - .5*tanh((freq_vec.squeeze()[1:] - f_changeover) * sharpness) weight = (weight + weight[::-1]).squeeze() weight = np.hstack([1., weight]) return weight def kin_estimate(selectedData): """ calculates the kinematic CoM estimate from "selectedData" """ # anthropometry from Dempster # format: [prox.marker, dist. marker, rel. weight (%), # segment CoM pos rel. to prox. marker (%) ] aData = [ ('R_Hea','L_Hea',8.26,50.), ('L_Acr','R_Trc',46.84/2.,63.), ('R_Acr','L_Trc',46.84/2.,63.), ('R_Acr','R_Elb',3.25,43.6), ('R_Elb','R_WrL',1.87 + 0.65,43. + 25.), ('L_Acr','L_Elb',3.25,43.6), ('L_Elb','L_WrL',1.87 + 0.65,43. + 25.), ('R_Trc','R_Kne',10.5,43.3), ('R_Kne','R_AnL',4.75,43.4), ('R_Hee','R_Mtv',1.43,50. + 5.), ('L_Trc','L_Kne',10.5,43.3), ('L_Kne','L_AnL',4.75,43.4), ('L_Hee','L_Mtv',1.43,50. + 5.) ] # adaptation to dataformat when extracted from database: lowercase aData = [(x[0].lower(), x[1].lower(), x[2], x[3]) for x in aData] CoM = np.zeros((len(kd.sacr[:,0]),3)) for segment in aData: CoM += segment[2]/100.* ( (getattr(selectedData, segment[1]) - getattr(selectedData, segment[0])) * segment[3]/100. + getattr(selectedData, segment[0]) ) return CoM elems = dir(fd) # get convenient names for forces if 'fx' in elems: Fx = fd.fx.squeeze() elif 'Fx' in elems: Fx = fd.Fx.squeeze() elif 'fx1' in elems and 'fx2' in elems: Fx = (fd.fx1 + fd.fx2).squeeze() else: raise ValueError("Error: Fx field not in forces (fx or fx1 and fx2)") if 'fy' in elems: Fy = fd.fy.squeeze() elif 'Fy' in elems: Fy = fd.Fy.squeeze() elif 'fy1' in elems and 'fy2' in elems: Fy = (fd.fy1 + fd.fy2).squeeze() else: raise ValueError("Error: Fy field not in forces (fy or fy1 and fy2)") if 'fz' in elems: Fz = fd.fz.squeeze() elif 'Fz' in elems: Fz = fd.Fz.squeeze() elif all([x in elems for x in ['fz1', 'fz2', 'fz3', 'fz4']]): Fz = (fd.fz1 + fd.fz2 + fd.fz3 + fd.fz4).squeeze() elif all([x in elems for x in ['fzr1', 'fzr2', 'fzr3', 'fzr4', 'fzl1', 'fzl2', 'fzl3', 'fzl4',]]): Fz = (fd.fzr1 + fd.fzr2 + fd.fzr3 + fd.fzr4 + fd.fzl1 + fd.fzl2 + fd.fzl3 + fd.fzl4).squeeze() # remove bodyweight later else: raise ValueError( "Error: Fz field not in forces (fz or fz1..4 or fzr1..4 and fzl1...4)" ) if kin_est is None: kin_est = kin_estimate(kd) kin_est_i = vstack([mi.interp_f(kin_est[:, col].copy(), len(Fz), detrend=True) for col in range(3)]).T if use_Fint: diff_op = lambda x: mi.diff_f(x.squeeze(), fs=fd.fs) int_op = lambda x, x0: mi.int_f(x.squeeze(), fs=fd.fs) + x0 else: diff_op = lambda x: gradient(x.squeeze()) * fd.fs int_op = lambda x, x0: cumtrapz(x.squeeze(), initial=0) / fd.fs + x0 vd = vstack([diff_op(kin_est_i[:,col]) for col in range(3)]).T # correct mean forces, pt 1 n = 2 # for higher reliably use multiple points T = len(Fz) / fd.fs dvx = ((kin_est_i[-1,0] - kin_est_i[-(n+1),0]) - (kin_est_i[n,0] - kin_est_i[0,0])) * (fd.fs / n) dvy = ((kin_est_i[-1,1] - kin_est_i[-(n+1),1]) - (kin_est_i[n,1] - kin_est_i[0,1])) * (fd.fs / n) dvz = ((kin_est_i[-1,2] - kin_est_i[-(n+1),2]) - (kin_est_i[n,2] - kin_est_i[0,2])) * (fd.fs / n) ax = dvx/T ay = dvy/T az = dvz/T + 9.81 #print "dvz = ", dvz #print "az = ", az if mass == None: # determine mass mass = array(mean(Fz)/az).squeeze() Fz -= mass*9.81 else: Fz -= mean(Fz) Fz += (az - 9.81)*mass Fx = Fx - mean(Fx) + ax*mass Fy = Fy - mean(Fy) + ay*mass #print "estimated mass:", mass vi = vstack([int_op(F / mass, v0 ) for F, v0 in zip([Fx, Fy, Fz], vd[0,:])]).T spect_diff = vstack([fftpack.fft(vd[:, col]) for col in range(3)]).T spect_int = vstack([fftpack.fft(vi[:, col]) for col in range(3)]).T freq_vec = linspace(0, fd.fs, len(Fz), endpoint=False) wv = weighting_function(freq_vec, f_thresh, steepness) spect_combine = vstack([wv*spect_diff[:, col] + (1.-wv)*spect_int[:, col] for col in range(3)]).T v_combine = vstack([fftpack.ifft(spect_combine[:, col]).real for col in range(3)]).T x_combine = vstack([int_op(v_combine[:, col], kin_est[0, col]) for col in range(3)]).T if adapt_kin_mean: x_combine = x_combine - mean(x_combine, axis=0) for dim in range(x_combine.shape[1]): x_combine[:,dim] += mean(kin_est[:,dim]) f_combine = vstack([diff_op(v_combine[:, col])*mass for col in range(3)]).T f_combine[:, 2] += mass*9.81 if return_mass: return f_combine, x_combine, mass return f_combine, x_combine