def contour(A, d1, d2, thr=None):
# Adopted from https://github.com/agiovann/Constrained_NMF
from scipy.sparse import issparse
if issparse(A):
A = np.array(A.todense())
else:
A = np.array(A)
d, nr = np.shape(A)
# x, y = np.mgrid[:d1:1, :d2:1]
x = np.arange(d1)
y = np.arange(d2)
coordinates = []
for i in range(nr):
indx = np.argsort(A[:, i], axis=None)[::-1]
cumEn = np.cumsum(A[:, i].flatten()[indx]**2)
cumEn /= cumEn[-1]
Bvec = np.zeros(d)
Bvec[indx] = cumEn
Bmat = np.reshape(Bvec, (d1, d2), order='F')
Bmat = 1 - Bmat
c = QuadContourGenerator.from_rectilinear(y, x, Bmat, shapely_fmt)
cs = c.filled_contour(min=1 - thr, max=None)
# cs = cntr.trace(thr)
if len(cs) > 0:
coordinates.append(cs)
else:
sys.stdout.write("No polygon found\n")
coordinates.append([[0, 0, 0], [0, 0, 0], [0, 0, 0]])
return coordinates
def draw_contour_line(R, Z, array, val, pathnum):
c = QuadContourGenerator.from_rectilinear(R[0], Z[:, 0], array)
res = c.contour(val)[pathnum]
x = res[:, 0]
y = res[:, 1]
return x, y
def pfr_lines(self, inp, core):
c = QuadContourGenerator.from_rectilinear(core.psi_data.R[0], core.psi_data.Z[:, 0], core.psi_data.psi_norm)
contours = c.contour(0.999)
if len(contours) == 1:
# then we're definitely dealing with a surface inside the seperatrix
raise ValueError("Did not find PFR flux surface. Stopping.")
else:
# we need to find the surface that is contained within the private flux region
for j, contour in enumerate(contours):
# if the contour is entirely below the x-point and contour extends to both the left and the
# right of the x-point horizontally, then the contour is almost certainly the desired PFR contour
if np.amax(contour[:, 1]) < core.pts.xpt[1] and \
np.amin(contour[:, 0]) < core.pts.xpt[0] < np.amax(contour[:, 0]):
# then it's a probably a pfr flux surface, might need to add additional checks later
# make sure the line goes from inboard to outboard
if contour[-1, 0] < contour[0, 0]:
contour = np.flipud(contour)
# find cut points
cut_pt_ib = LineString(contour).intersection(inp.wall_line)[0]
cut_pt_ob = LineString(contour).intersection(inp.wall_line)[1]
dist1 = LineString(contour).project(cut_pt_ib, normalized=True)
cutline_temp = cut(LineString(contour), dist1)[1]
# reverse line point order so we can reliably find the second intersection point
cutline_temp_rev = LineString(np.flipud(np.asarray(cutline_temp.coords)))
dist2 = cutline_temp_rev.project(cut_pt_ob, normalized=True)
cutline_final_rev = cut(cutline_temp_rev, dist2)[1]
# reverse again for final pfr flux line
pfr_flux_line = LineString(np.flipud(np.asarray(cutline_final_rev.coords)))
# add pfr_line intersection points on inboard side
# for some reason, union freaks out when I try to do inboard and outboard
# at the same time.
union = inp.wall_line.union(cut(LineString(contour), 0.5)[0])
result = [geom for geom in polygonize(union)][0]
inp.wall_line = LineString(result.exterior.coords)
# add pfr line intersection points on outboard side
union = inp.wall_line.union(cut(LineString(contour), 0.5)[1])
result = [geom for geom in polygonize(union)][0]
inp.wall_line = LineString(result.exterior.coords)
# cut out pfr section of wall line
wall_pts = np.asarray(inp.wall_line.xy).T
wall_start_pos = np.where((wall_pts == cut_pt_ob).all(axis=1))[0][0]
wall_line_rolled = LineString(np.roll(wall_pts, -wall_start_pos, axis=0))
wall_line_cut_pfr = cut(wall_line_rolled,
wall_line_rolled.project(cut_pt_ib, normalized=True))[0]
# create LineString with pfr line and section of wall line
self.pfr_line = linemerge((pfr_flux_line, wall_line_cut_pfr))
break
def getZeroContour(self):
SDF = np.copy(self.SDF)
# compute contour filled from -inf to zero (in shapely)
if self.zerocontour is None:
c = QuadContourGenerator.from_rectilinear(self.xpts, self.ypts,
SDF, shapely_fmt)
# print(c.filled_contour(max=0.0), type(c.filled_contour(max=0.0)), len(c.filled_contour(max=0.0)))
cutoff_len = 0
poly = c.filled_contour(max=cutoff_len)[0]
nm = np.max(SDF[self.getInsidePoints()])
min_dist = min(self.delta_x, self.delta_y)
smooth_at_end = True
if smooth_at_end and nm <= min_dist * 2:
# smooth polygon if close to end
poly = poly.buffer(1, resolution=16,
join_style=1).buffer(-1,
resolution=16,
join_style=1)
self.zerocontour = poly
return self.zerocontour
def draw_core_line(R, Z, psinorm, psi_val, sep_pts):
# create contour generator
c = QuadContourGenerator.from_rectilinear(R[0], Z[:, 0], psinorm)
# draw contours with psi_val
contours = c.contour(psi_val)
if len(contours) == 0 and isclose(psi_val, 0, abs_tol=1E-9):
# This is probably the magnetic axis
fs_line = None
fs_inn_pt = None
fs_out_pt = None
fs_top_pt = None
fs_bot_pt = None
fs_axis = None
elif len(contours) == 0 and not isclose(psi_val, 0, abs_tol=1E-9):
# This either means that either:
# A) psi is a value not present in the psi data or
# B) psi is very close to the magnetic axis, but is too small for the contours
# package to give us a contour. This can happen if you you have an very fine
# radial mesh in the vicinity of the magnetic axis.
# Passing back None for now, but should probably raise an exception in the future # TODO
fs_line = None
fs_inn_pt = None
fs_out_pt = None
fs_top_pt = None
fs_bot_pt = None
fs_axis = None
else:
if len(contours) == 1:
# then we're definitely dealing with a surface inside the seperatrix
x, y = draw_contour_line(R, Z, psinorm, psi_val, 0)
else:
# we need to find which of the surfaces is inside the seperatrix
for j, line in enumerate(contours):
x, y = draw_contour_line(R, Z, psinorm, psi_val, j)
if (np.amax(x) < np.amax(sep_pts[:, 0]) and
np.amin(x) > np.amin(sep_pts[:, 0]) and
np.amax(y) < np.amax(sep_pts[:, 1]) and
np.amin(y) > np.amin(sep_pts[:, 1])):
# then it's an internal flux surface
break
pts = np.column_stack((x, y))
fs_line = LineString(pts)
fs_out_pt = pts[np.argmax(pts, axis=0)[0]]
fs_inn_pt = pts[np.argmin(pts, axis=0)[0]]
fs_top_pt = pts[np.argmax(pts, axis=0)[1]]
fs_bot_pt = pts[np.argmin(pts, axis=0)[1]]
fs_axis = [(fs_out_pt[0]+fs_inn_pt[0])/2, (fs_out_pt[1]+fs_inn_pt[1])/2]
fs_pts = namedtuple('fs_pts', 'inn out top bot axis')(
fs_inn_pt,
fs_out_pt,
fs_top_pt,
fs_bot_pt,
fs_axis
)
return fs_line, fs_pts
def draw_core_line(R, Z, psi, psi_val, sep_pts):
# create contour generator
c = QuadContourGenerator.from_rectilinear(R[0], Z[:, 0], psi)
# draw contours with psi_val
contours = c.contour(psi_val)
if len(contours) == 1:
# then we're definitely dealing with a surface inside the seperatrix
x, y = draw_contour_line(R, Z, psi, psi_val, 0)
else:
# we need to find which of the surfaces is inside the seperatrix
for j, line in enumerate(contours):
x, y = draw_contour_line(R, Z, psi, psi_val, j)
if (np.amax(x) < np.amax(sep_pts[:, 0])
and np.amin(x) > np.amin(sep_pts[:, 0])
and np.amax(y) < np.amax(sep_pts[:, 1])
and np.amin(y) > np.amin(sep_pts[:, 1])):
# then it's an internal flux surface
break
pts = np.column_stack((x, y))
line = LineString(pts)
out_pt = pts[np.argmax(pts, axis=0)[0]]
in_pt = pts[np.argmin(pts, axis=0)[0]]
top_pt = pts[np.argmax(pts, axis=0)[1]]
bot_pt = pts[np.argmin(pts, axis=0)[1]]
fs_axis = [(out_pt[0] + in_pt[0]) / 2, (out_pt[1] + in_pt[1]) / 2]
return line, fs_axis
def pdf_contour_coords(xx, yy, pdf, level):
c = QuadContourGenerator.from_rectilinear(xx, yy, pdf.T)
levelc = c.contour(level)
cx = []
cy = []
for v in levelc:
cx.extend(v[:, 0])
cy.extend(v[:, 1])
return cx, cy
def _extract_polygon(src_data, dx, dy, minval, maxval):
"""
"""
src_gauss = gaussian(np.where((src_data > minval) * (src_data <= maxval), 1, 0),
sigma=0.75,
preserve_range=True)
c = QuadContourGenerator.from_rectilinear(dx,
dy,
np.flipud(src_gauss),
shapely_fmt)
contour = c.filled_contour(min=0.7, max=2.0)
results = ({'properties': {'value': minval}, 'geometry': s}
for i, s in enumerate(contour))
geoms = list(results)
return geoms
def calc_sol_lines(self, inp, core):
c = QuadContourGenerator.from_rectilinear(core.psi_data.R[0], core.psi_data.Z[:, 0], core.psi_data.psi_norm)
self.sol_lines = []
self.sol_lines_cut = []
#
# first we need to make sure that the inp.sollines_psi_max specified in the input file doesn't
# go outside of the first wall. For now, we will check for this as follows:
# 1. draw the contours for psi_norm = inp.sollines_psi_max
# 2. if there is only 1, make sure it only intersects the first wall no more than twice
# 3. if there is only 1 and it intersects the first wall more than twice, raise an error
# that tells the user to reduce their sollines_psi_max value in the input file
# 4. if there is more than 1 contour, then we need to determine which one is the correct one.
# for now, we will check that it's largest y value is higher than the magnetic axis,
# it's largest x value is to the right of the magnetic axis, it's lowest y value is lower
# than the magnetic axis, and it's lowest x value is to the left of the magnetic axis.
# 5. Once the correct contour has been identified, do steps 2 and 3.
#
sollines_psi_max_contours = c.contour(inp.sollines_psi_max)
num_lines = len(sollines_psi_max_contours)
print 'num_lines = ', num_lines
if num_lines == 1:
# then this is probably the correct line. Check to see how many times it intersects
# with the first wall
num_wall_ints = len(LineString(sollines_psi_max_contours[0]).intersection(inp.wall_line))
print 'num_wall_ints = ',num_wall_ints
if num_wall_ints > 2:
print 'It looks like your sollines_psi_max value might be intersecting the wall.' \
'Try reducing it. Stopping.'
sys.exit()
else:
for i, line in enumerate(sollines_psi_max_contours):
max_x = np.amax(line[:,0])
min_x = np.amin(line[:,0])
max_y = np.amax(line[:,1])
min_y = np.amin(line[:,1])
if min_x < core.pts.axis.mag[0] < max_x and min_y < core.pts.axis.mag[1] < max_y:
# then this is probably the correct line. Check to see how many times it intersects
# with the first wall
num_wall_ints = len(LineString(line).intersection(inp.wall_line))
if num_wall_ints > 2:
print 'It looks like your sollines_psi_max value might be intersecting the wall.' \
'Try reducing it. Stopping.'
sys.exit()
else:
break
psi_pts = np.linspace(1, inp.sollines_psi_max, inp.num_sollines+1, endpoint=True)[1:]
for i, v in enumerate(psi_pts):
num_lines = len(c.contour(v))
if num_lines == 1:
# this is probably the correct line
self.sol_lines.append(LineString(c.contour(v)[0]))
else:
# we need to determine which contour to use
for line in c.contour(v):
max_x = np.amax(line[:, 0])
min_x = np.amin(line[:, 0])
max_y = np.amax(line[:, 1])
min_y = np.amin(line[:, 1])
if min_x < core.pts.axis.mag[0] < max_x and min_y < core.pts.axis.mag[1] < max_y:
# then this is probably the correct line.
self.sol_lines.append(LineString(line))
for line in self.sol_lines:
# find intersection points with the wall
int_pts = line.intersection(inp.wall_line)
# cut line at intersection points
cut_line = cut(line, line.project(int_pts[0], normalized=True))[1]
cut_line = cut(cut_line, cut_line.project(int_pts[1], normalized=True))[0]
self.sol_lines_cut.append(cut_line)
# add wall intersection points from divertor legs and sol lines to wall_line.
# This is necessary to prevent thousands of tiny triangles from forming if the
# end of the flux line isn't exactly on top of the wall line.
# add inboard seperatrix strike point
#plt.plot(np.asarray(inp.wall_line)[:,0], np.asarray(inp.wall_line)[:,1])
#plt.plot(np.asarray(core.lines.div.ib_long)[:,0], np.asarray(core.lines.div.ib_long)[:, 1])
union = inp.wall_line.union(core.lines.div.ib_long)
result = [geom for geom in polygonize(union)][0]
inp.wall_line = LineString(result.exterior.coords)
# add outboard seperatrix strike point
union = inp.wall_line.union(core.lines.div.ob_long)
result = [geom for geom in polygonize(union)][0]
inp.wall_line = LineString(result.exterior.coords)
# add sol line intersection points on inboard side
# for some reason, union freaks out when I try to do inboard and outboard
# at the same time.
for num, line in enumerate(self.sol_lines):
union = inp.wall_line.union(cut(line, 0.5)[0])
result = [geom for geom in polygonize(union)][0]
inp.wall_line = LineString(result.exterior.coords)
# add sol line intersection points on outboard side
for num, line in enumerate(self.sol_lines):
union = inp.wall_line.union(cut(line, 0.5)[1])
result = [geom for geom in polygonize(union)][0]
inp.wall_line = LineString(result.exterior.coords)
def fit_plots(x, y, minimizer, minimizer_result, residual,
model, title='',
contour_level=0.6827, cmap=plt.cm.coolwarm,
xlabel='x', ylabel='y', xlim=None, ylim=None):
params = minimizer_result.params
ci, trace = lmfit.conf_interval(minimizer, minimizer_result, trace=True)
param_names = list(params.keys())
figs = []
##################################
# Figure 1: Fit Parameters as text
##################################
fig = plt.figure()
s = '%s\n\n' % model
if title:
s += title + '\n'
for ndx, k in enumerate(param_names):
s += '%s: %.3f ± %.3f' % (k, minimizer_result.params[k].value, minimizer_result.params[k].stderr)
s += '\n'
plt.text(0.5, 0.5, s, fontsize=12, ha='center', va='center')
plt.axis('off')
figs.append(fig)
#############################
# Figure 2: Fit and residuals
#############################
fig, (ax1, ax2) = plt.subplots(2, 1, sharex=True, gridspec_kw={'height_ratios': [3, 1]})
ax1.plot(x, y, '.')
xs = np.linspace(x[0], x[-1])
ax1.plot(xs, residual(minimizer_result.params, xs), '-')
ax1.set_ylabel(ylabel)
if ylim:
ax1.set_ylim(ylim)
r = residual(minimizer_result.params, x, data=y)
ax2.plot(x, r)
ax2.axhline(y=0, color='k', linestyle=':')
mx = np.max(np.abs(r))
ax2.set_ylim([-mx, mx])
ax2.set_ylabel('R')
ax2.set_xlabel(xlabel)
if xlim:
ax2.set_xlim(xlim)
figs.append(fig)
#############################
# Figure 3: Probability plots
#############################
contours = {}
fig, axs = plt.subplots(len(param_names), len(param_names))
for ndx1 in range(len(param_names)):
for ndx2 in range(len(param_names)):
ax = axs[ndx2][ndx1]
if ndx1 > ndx2:
ax.set_axis_off()
continue
if ndx1 == ndx2:
x = trace[param_names[ndx1]][param_names[ndx1]]
y = trace[param_names[ndx1]]['prob']
t, s = np.unique(x, True)
f = interp1d(t, y[s], 'slinear')
xn = np.linspace(x.min(), x.max(), 50)
ax.plot(xn, f(xn), 'g', lw=1)
contours[ndx1] = (x, y)
else:
x, y, m = lmfit.conf_interval2d(minimizer, minimizer_result, param_names[ndx1], param_names[ndx2], 20, 20)
ax.contourf(x, y, m, np.linspace(0, 1, 10), cmap=cmap)
ch = QuadContourGenerator.from_rectilinear(x, y, m, numpy_formatter)
contours[(ndx1, ndx2)] = ch.contour(contour_level)
if ndx1 == 0:
if ndx2 > 0:
ax.set_ylabel(param_names[ndx2])
else:
ax.set_ylabel('prob')
else:
ax.set_yticks([])
if ndx2 == len(param_names) - 1:
ax.set_xlabel(param_names[ndx1])
else:
ax.set_xticks([])
figs.append(fig)
return figs, contours