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
