def struct_radius_scatter(psi_arr, cur, den):
hess = hessian(psi_arr)
ntheta = 2
nx, ny = psi_arr.shape
psi_grad_x, psi_grad_y = gradient(psi_arr)
bmag = np.sqrt(psi_grad_x**2 + psi_grad_y**2)
den_x, den_y = gradient(den)
den_grad_mag = np.sqrt(den_x**2 + den_y**2)
surf = _cp.TopoSurface(psi_arr)
regions = surf.get_minmax_regions()
res = {}
for (minmax, pss, type), region in regions.items():
if len(region) < 50: continue
name_2_field = {'bmag': [bmag],
'den': [den],
'cur': [cur],
'psi': [psi_arr],
'hess': [hess],
'den_grad': [den_grad_mag],
}
for name, (field,) in name_2_field.items():
th_sector = theta_sectors(field, minmax[0], minmax[1], region, ntheta=2)[0]
field_radius = r_moment(th_sector['r'],th_sector['val'], th_sector['valstd'])
res.setdefault(name, []).append(field_radius)
return res
def mag_shear(psi_arr):
r'''
computes \grad_{\perp} \bvec{B} =
\frac{\grad{\psi}}{2 (\grad{\psi})^2} \cdot \grad{(\grad{\psi})^2}
'''
psi_grad_x, psi_grad_y = gradient(psi_arr)
psi_grad_2 = psi_grad_x**2 + psi_grad_y**2
psi_grad_grad_x, psi_grad_grad_y = gradient(psi_grad_2)
mag_shear = psi_grad_x * psi_grad_grad_x + psi_grad_y * psi_grad_grad_y
return mag_shear / (2.0 * psi_grad_2)
def __init__(self, alpha, psi, phi, den):
self.phi = phi
self.psi = psi
self.den = den
self.alpha = alpha
self.by, self.bx = gradient(psi)
self.by *= -1
self.vy, self.vx = gradient(phi)
self.vy *= -1
self.J = laplacian(psi)
self.vor = laplacian(phi)
def mag_shear_rad_profs(psi_arr):
psi_x, psi_y = gradient(psi_arr)
bmag = np.sqrt(psi_x**2 + psi_y**2)
psi_arr = gaussian_filter(psi_arr, sigma=2.0, mode='wrap')
mshear = mag_shear(psi_arr)
surf = _cp.TopoSurface(psi_arr)
regions = surf.get_minmax_regions()
for (minmax, pss, type), region in regions:
pass
rprofs = radial_profiles(surf, threshold=25, other_arr=bmag)
pl.ion()
pl.figure()
for minmax, (rprof, region) in rprofs.items():
# minmax_flux = arr_div[minmax]
pts, extremum_val, avg_bmags, avg_bmags_errs, avg_dists, avg_dists_errs = \
zip(*rprof)
# fluxes = np.abs(np.array(fluxes) - minmax_flux)
# avg_fluxes = np.abs(np.array(avg_fluxes) - minmax_flux)
pl.subplot(211)
pl.plot(avg_dists, avg_bmags, 'd-')
pl.subplot(212)
pl.imshow(bmag)
# pl.plot(avg_dists, avg_shear, 'd-')
raw_input('enter to continue')
pl.clf()
import pdb; pdb.set_trace()
pl.close('all')
def total_shear(bx, by, x0, y0):
raise RuntimeError("doesn't work; need better treatment of gradients near "
"R=0. Currently get wild oscillations near R=0, since gradient uses "
"fourier transforms.")
nx, ny = bx.shape
br, btheta = cartesian_to_polar(bx, by, x0, y0)
X, Y, R, theta = cartesian_to_polar_coords(x0, y0, nx, ny)
br_by_r = br/R
br_by_r[x0, y0] = 0.0
br_by_r_x, br_by_r_y = gradient(br_by_r)
bt_by_r = btheta/R
bt_by_r[x0, y0] = 0.0
btheta_by_r_x, btheta_by_r_y = gradient(bt_by_r)
shear_1, _ = cartesian_to_polar(btheta_by_r_x, btheta_by_r_y, x0, y0)
_, shear_2 = cartesian_to_polar(br_by_r_x, br_by_r_y, x0, y0)
return shear_1 + shear_2
def mag_shear_rad_scatter(psi_arr):
nx, ny = psi_arr.shape
# mshear = mag_shear(psi_arr)
psi_grad_x, psi_grad_y = gradient(psi_arr)
bmag = np.sqrt(psi_grad_x**2 + psi_grad_y**2)
pl.ion()
pl.figure()
pl.imshow(bmag, interpolation='nearest', cmap='hot')
surf = _cp.TopoSurface(psi_arr)
regions = surf.get_minmax_regions()
pl.figure()
for (minmax, pss, type), region in regions.items():
br, btheta = cartesian_to_polar(-psi_grad_y, psi_grad_x, minmax[0], minmax[1])
if len(region) < 100:
continue
# region = expand_region(psi_arr, region, ntimes=2)
region = expand_region_circ(psi_arr, region, minmax, extra=10.0)
# dists, shears = \
# flux_tube_radial_scatter(region, minmax, mshear)
dists, bthetas = \
flux_tube_radial_scatter(region, minmax, btheta)
dists, bmags = \
flux_tube_radial_scatter(region, minmax, bmag)
# dists, fluxes = \
# flux_tube_radial_scatter(region, minmax, psi_arr)
pl.subplot(221)
pl.scatter(dists, bthetas, c='b', marker='s')
pl.grid()
pl.title(r'$B_{\theta}$ vs. $r$')
pl.subplot(222)
pl.scatter(dists, bthetas/dists, c='g', marker='d')
pl.grid()
pl.title(r'$B_{\theta}/r$ vs. $r$')
# nonlin_shear = (shears - bmags / dists) / dists
# pl.scatter(dists, nonlin_shear, c='b', marker='s', label=r'$\nabla_{\perp}B$')
pl.subplot(224)
# psi_cpy = np.zeros_like(psi_arr)
btheta_cpy = np.zeros_like(btheta)
xs, ys = zip(*region)
btheta_cpy[xs, ys] = btheta[xs, ys]
pl.imshow(btheta_cpy, interpolation='nearest', cmap='hot')
pl.title(r'$B_{\theta}$')
pl.subplot(223)
bmag_cpy = np.zeros_like(bmag)
bmag_cpy[xs, ys] = bmag[xs, ys]
pl.imshow(bmag_cpy, interpolation='nearest', cmap='hot')
pl.title(r'$|B|$')
raw_input('enter to continue')
pl.clf()
pl.close('all')
def mag_shear_theta_sectors(psi_arr, cur, den, thresh=100, save_basename='field-vs-r'):
hess = hessian(psi_arr)
ntheta = 2
nx, ny = psi_arr.shape
psi_grad_x, psi_grad_y = gradient(psi_arr)
bmag = np.sqrt(psi_grad_x**2 + psi_grad_y**2)
den_x, den_y = gradient(den)
den_grad_mag = np.sqrt(den_x**2 + den_y**2)
# pl.ion()
# pl.figure()
# pl.imshow(bmag, interpolation='nearest', cmap='hot')
surf = _cp.TopoSurface(psi_arr)
regions = surf.get_minmax_regions()
pl.figure(figsize=(9,12))
nr = 6
nc = 2
ctr = 0
for (minmax, pss, type), region in regions.items():
if len(region) < thresh:
continue
ctr += 1
print ctr
gs = gridspec.GridSpec(nr, nc, width_ratios=[1,2], hspace=0.0)
nbr_func = lambda t: _cp.neighbors6(t[0], t[1], nx, ny)
lset = field_trace._level_set(psi_arr, level_val=psi_arr[pss],
position=pss, neighbors_func=nbr_func)
bx = -psi_grad_y
by = psi_grad_x
br, btheta = cartesian_to_polar(bx, by, minmax[0], minmax[1])
region = expand_region_circ(psi_arr, region, minmax, extra=2.0)
lset = lset.intersection(region)
# psi
ax = pl.subplot(gs[0])
ax.set_title("Field & separatrix", size='x-large')
field_region_image(ax, psi_arr, region, minmax, lset, title=r'$\psi$')
# ax = pl.subplot2grid((nr, nc), (0, 1))
ax = pl.subplot(gs[1])
ax.set_title(r"Field vs. $r/\rho_s$ (id={0})".format(ctr), size='x-large')
plot_theta_sectors(ax, psi_arr, minmax, region, ntheta, title=r'$\psi$ vs. $r/\rho_s$')
# b_theta
ax = pl.subplot(gs[2])
# field_region_image(ax, btheta, region, minmax, lset, title=r'$B_{\theta}$')
field_region_image(ax, bmag, region, minmax, lset, title=r'$|B|$')
ax = pl.subplot(gs[3])
# btheta_sectors = plot_theta_sectors(ax, btheta, minmax, region, ntheta, title=r'$B_{\theta}$ vs. $r/\rho_s$')
plot_theta_sectors(ax, bmag, minmax, region, ntheta, title=r'$|B|$ vs. $r/\rho_s$')
# b_shear
# pl.subplot(nr, nc, 5)
# field_region_image(btheta, region, minmax, lset, title=r'$B_{\theta}$')
# pl.subplot(nr, nc, 6)
# plot_sect_derivs_by_r(btheta_sectors, title=r'$\partial_{r} \frac{B_{\theta}}{r}$ vs. $r/\rho_s$')
# den
ax = pl.subplot(gs[4])
field_region_image(ax, den, region, minmax, lset, title=r'$n$')
ax = pl.subplot(gs[5])
plot_theta_sectors(ax, den, minmax, region, ntheta, title=r'$n$ vs. $r/\rho_s$')
# den_grad
ax = pl.subplot(gs[6])
field_region_image(ax, den_grad_mag, region, minmax, lset, title=r'$|\nabla n|$')
ax = pl.subplot(gs[7])
plot_theta_sectors(ax, den_grad_mag, minmax, region, ntheta, title=r'$|\nabla n|$ vs. $r/\rho_s$')
# cur
ax = pl.subplot(gs[8])
field_region_image(ax, cur, region, minmax, lset, title=r'$J$')
ax = pl.subplot(gs[9])
plot_theta_sectors(ax, cur, minmax, region, ntheta, title=r'$J$ vs. $r/\rho_s$')
# hessian
ax = pl.subplot(gs[10])
field_region_image(ax, hess, region, minmax, lset, title=r'$H(\psi)$', xvis=True)
ax = pl.subplot(gs[11])
ax.set_xlabel(r'$r/\rho_s$', size='x-large')
plot_theta_sectors(ax, hess, minmax, region, ntheta, title=r'$H(\psi)$ vs. $r/\rho_s$', xvis=True)
# raw_input('enter to continue')
pl.savefig('{0}_{1:04d}.pdf'.format(save_basename, ctr))
pl.clf()
pl.close('all')
def mag_shear_rad_spokes(psi_arr):
nx, ny = psi_arr.shape
wdist = wraparound_dist_vec(nx, ny)
# mshear = mag_shear(psi_arr)
psi_grad_x, psi_grad_y = gradient(psi_arr)
bmag = np.sqrt(psi_grad_x**2 + psi_grad_y**2)
surf = _cp.TopoSurface(psi_arr)
regions = surf.get_minmax_regions()
pl.ion()
pl.figure()
for (minmax, pss, type), region in regions.items():
br, btheta = cartesian_to_polar(-psi_grad_y, psi_grad_x, minmax[0], minmax[1])
if len(region) < 100:
continue
region = expand_region_circ(psi_arr, region, minmax, extra=10.0)
spokes = flux_tube_radial_spokes(region, minmax, psi_arr, type)
pl.clf()
# max_2_bmags = []
max_2_btheta = []
for spoke in spokes:
sp_arr = np.array(spoke)
xs, ys = sp_arr[:,0], sp_arr[:,1]
# bmags = bmag[xs, ys]
bthetas = btheta[xs, ys]
max_2_btheta.append((bthetas.max(), bthetas, spoke))
max_2_btheta.sort(key=lambda x: x[0], reverse=True)
spokes = [sp for (m, bthetas, sp) in max_2_btheta]
btheta_cpy = np.zeros_like(btheta)
bmag_cpy = np.zeros_like(bmag)
br_cpy = np.zeros_like(br)
xs, ys = zip(*region)
btheta_cpy[xs, ys] = btheta[xs, ys]
bmag_cpy[xs, ys] = bmag[xs, ys]
br_cpy[xs, ys] = br[xs, ys]
for spoke in spokes:
sp_arr = np.array(spoke)
xs, ys = sp_arr[:,0], sp_arr[:,1]
dists = wdist(minmax[0], minmax[1], xs, ys)
bthetas = btheta[xs, ys]
# mshears = mshear[xs, ys]
pl.subplot(221)
pl.plot(dists, bthetas, 'o-')
# pl.subplot(222)
# pl.plot(dists, mshears, 's-')
pl.subplot(222)
pl.plot(dists[1:], bthetas[1:]/dists[1:], 'd-')
pl.subplot(221)
pl.grid()
pl.title(r'$B_{\theta}$ vs. $r$')
# pl.subplot(222)
# pl.grid()
# pl.title(r'$\nabla_{\perp} B$ vs. $r$')
pl.subplot(222)
pl.grid()
pl.title(r'$B_{\theta}/r$ vs. $r$')
pl.subplot(234)
pl.imshow(btheta_cpy, interpolation='nearest', cmap='hot')
pl.title(r'$B_{\theta}$')
pl.subplot(235)
pl.imshow(br_cpy, interpolation='nearest', cmap='hot')
pl.title(r'$B_{r}$')
pl.subplot(236)
pl.imshow(bmag_cpy, interpolation='nearest', cmap='hot')
pl.title(r'$|B|$')
raw_input('enter to continue')
pl.clf()
pl.close('all')
def grad_par_v(self, arr):
arr_x, arr_y = gradient(arr)
return self.vx * arr_x + self.vy *arr_y
def grad_par_b(self, arr):
arr_x, arr_y = gradient(arr)
return self.bx * arr_x + self.by * arr_y