|
new file 100644
|
|
|
import matplotlib
|
|
|
from matplotlib.ticker import Formatter
|
|
|
matplotlib.use('Agg')
|
|
|
|
|
|
#from matplotlib import rc
|
|
|
#rc('font',**{'family':'sans-serif','sans-serif':['Helvetica']})
|
|
|
## for Palatino and other serif fonts use:
|
|
|
#rc('font',**{'family':'serif','serif':['Palatino']))
|
|
|
#rc('text', usetex=True)
|
|
|
|
|
|
from pylab import *
|
|
|
import os, os.path
|
|
|
import sys
|
|
|
import glob
|
|
|
import tarfile
|
|
|
import math
|
|
|
import scaling
|
|
|
import re
|
|
|
import numpy
|
|
|
from subinterpol import load_fine
|
|
|
|
|
|
rcParams['contour.negative_linestyle'] = 'solid'
|
|
|
|
|
|
output_dir = '.'
|
|
|
#try:
|
|
|
# output_dir = os.environ['CSTREAM_OUTPUT_DIR']
|
|
|
#except KeyError:
|
|
|
# output_dir = '/export/scratch1/luque/cstream/data/'
|
|
|
|
|
|
try:
|
|
|
BASE_PATH = os.environ['CSTREAM_DEFAULT_OUT_PATH']
|
|
|
except KeyError:
|
|
|
BASE_PATH = '/export/scratch1/luque/cstream/'
|
|
|
|
|
|
LOG_EPS = 1e-20
|
|
|
|
|
|
# Some rescalings:
|
|
|
nitrogen1atm_cm = {'r': 2.3e-4, 'z': 2.3e-4,
|
|
|
'electrons': 1.0 / (2.3e-4)**3}
|
|
|
|
|
|
nitrogen1atm_microm = {'r': 2.3, 'z': 2.3,
|
|
|
'electrons': 4.7e2,
|
|
|
'charge': 4.7e2,
|
|
|
'eabs': 200,
|
|
|
'ez': 200}
|
|
|
|
|
|
nitrogen1atm_microm_units = {'r': r'\mu m', 'z': r'\mu m',
|
|
|
'electrons': r'\mu m^{-3}',
|
|
|
'charge': r'e/\mu m^{3}',
|
|
|
'eabs': r'kV/cm'}
|
|
|
|
|
|
nitrogen1atm_mm = {'r': 2.3e-3, 'z': 2.3e-3,
|
|
|
'electrons': 4.7e2,
|
|
|
'charge': 4.7e2,
|
|
|
'eabs': 200}
|
|
|
|
|
|
nitrogen1atm_mm_units = {'r': r'mm', 'z': r'mm',
|
|
|
'electrons': r'\mu m^{-3}',
|
|
|
'charge': r'e/\mu m^{3}',
|
|
|
'eabs': r'kV/cm'}
|
|
|
|
|
|
|
|
|
|
|
|
nitrogen5atm_microm = scaling.rescale2(1, 5, nitrogen1atm_microm)
|
|
|
nitrogen5atm_microm_units = nitrogen1atm_microm_units
|
|
|
|
|
|
|
|
|
nitrogen15mbar_mm = {'r': 0.153, 'z': 0.153,
|
|
|
'electrons': 1.0 / (0.153)**3}
|
|
|
|
|
|
nitrogen15mbar_mm_units = {'x': r'mm', 'y': r'mm',
|
|
|
'electrons':r'mm^{-3}'}
|
|
|
|
|
|
air100mbar_mm = {'r': 0.0233, 'z': 0.0233,
|
|
|
'electrons': 1.0 / (0.0233)**3,
|
|
|
'ions': 1.0 / (0.0233)**3,
|
|
|
'charge': 1.0 / (0.0233)**3
|
|
|
}
|
|
|
|
|
|
air100mbar_mm_units = {'r': r'mm', 'z': r'mm',
|
|
|
'electrons':r'mm^{-3}',
|
|
|
'ions': r'mm^{-3}',
|
|
|
'charge': r'e/mm^{3}'
|
|
|
}
|
|
|
|
|
|
|
|
|
nitrogen005torr_m = {'r': 0.035, 'z': 0.035,
|
|
|
'electrons': 1.0 / (0.035)**3,
|
|
|
'charge': 1.0 / (0.035)**3,
|
|
|
}
|
|
|
|
|
|
nitrogen005torr_m_units = {'r': r'm', 'z': r'm',
|
|
|
'electrons':r'm^{-3}',
|
|
|
'charge': r'e/m^{3}',
|
|
|
}
|
|
|
|
|
|
|
|
|
micrombar = nitrogen1atm_microm
|
|
|
micrombar_units = {'pr': r'\mu m \cdot bar',
|
|
|
'px': r'\mu m \cdot bar',
|
|
|
'pz': r'\mu m \cdot bar',
|
|
|
'r': r'\mu m \cdot bar',
|
|
|
'x': r'\mu m \cdot bar',
|
|
|
'z': r'\mu m \cdot bar',
|
|
|
'electrons':r'\mu m^{-3} \cdot bar^{-3}'}
|
|
|
|
|
|
# Rescaling for sprites:
|
|
|
def get_sprite_f(dens_decay_len):
|
|
|
def _f(lz):
|
|
|
return numpy.exp(-lz / dens_decay_len)
|
|
|
return _f
|
|
|
|
|
|
DENS_DECAY_LENGTH_KM = 7.2
|
|
|
|
|
|
def dens_at(h_km):
|
|
|
""" Density at h_km kilometers relative to the density at ground level."""
|
|
|
return numpy.exp(-h_km / DENS_DECAY_LENGTH_KM)
|
|
|
|
|
|
def rescale_sprite_at(h_km):
|
|
|
dens = dens_at(h_km)
|
|
|
l0 = 2.3e-9 / dens
|
|
|
f = get_sprite_f(DENS_DECAY_LENGTH_KM / l0)
|
|
|
n0 = dens * 1.0 / (2.3e-4)**3
|
|
|
return {'r': l0, 'z': -l0, 'z0': h_km,
|
|
|
'eabs_f': f, 'ez_f': f, 'er_f': f,
|
|
|
'electrons': n0, 'ions': n0,
|
|
|
'charge': n0 }
|
|
|
|
|
|
sprite70km = rescale_sprite_at(70)
|
|
|
sprite70km_units = {'r': 'km', 'z': 'km', 'electrons': 'cm^{-3}',
|
|
|
'log_electrons': 'cm^{-3}'}
|
|
|
sprite80km = rescale_sprite_at(80)
|
|
|
sprite80km_units = sprite70km_units
|
|
|
sprite82km = rescale_sprite_at(82)
|
|
|
sprite82km_units = sprite70km_units
|
|
|
sprite90km = rescale_sprite_at(90)
|
|
|
sprite90km_units = sprite70km_units
|
|
|
|
|
|
|
|
|
PRETTY_VARS = {
|
|
|
'electrons': 'n_e',
|
|
|
'ions': 'n_+',
|
|
|
'eabs': 'E',
|
|
|
'er': 'E',
|
|
|
'_er': '-E',
|
|
|
'_ez': '$-E_z$',
|
|
|
'ez': 'E',
|
|
|
'phi': r'\phi',
|
|
|
'log_ez': r'$\rm{log}\/ E_z$',
|
|
|
'log_er': r'$\rm{log}\/ E_r$',
|
|
|
'log_eabs': r'$\rm{log}\/ E$',
|
|
|
'log_electrons': 'n_e',
|
|
|
'charge': r'Charge density',
|
|
|
'los_impact': r'$I(y,z)$'
|
|
|
#'charge': r'\rho - \sigma'
|
|
|
}
|
|
|
|
|
|
def pretty_var(var):
|
|
|
return PRETTY_VARS.get(var, var)
|
|
|
|
|
|
def flex_load(fname, path=''):
|
|
|
try:
|
|
|
# 1. Try with `filename'
|
|
|
return load(os.path.join(path, fname))
|
|
|
except IOError:
|
|
|
try:
|
|
|
# 2. If not, try with `filename.gz'
|
|
|
return load(os.path.join(path, fname + '.gz'))
|
|
|
except IOError:
|
|
|
# 2. If not, extract it from an all.<grid>.tgz '
|
|
|
tar = tarfile.open(os.path.join(path, tar_name(fname)), 'r:gz')
|
|
|
tar.extract(os.path.basename(fname))
|
|
|
r = load(fname)
|
|
|
os.remove(fname)
|
|
|
return r
|
|
|
|
|
|
def default_loader(var, grid, path=''):
|
|
|
return flex_load(input_name(var, grid), path=path)
|
|
|
|
|
|
def charge_loader(var, grid, path=''):
|
|
|
return (default_loader('ions', grid, path=path)
|
|
|
- default_loader('electrons', grid, path=path))
|
|
|
|
|
|
def d_charge_loader(var, grid, path=''):
|
|
|
return (default_loader('d_ions', grid, path=path)
|
|
|
- default_loader('d_electrons', grid, path=path))
|
|
|
|
|
|
def current_loader(var, grid, path=''):
|
|
|
eabs = default_loader('eabs', grid, path=path)
|
|
|
electrons = default_loader('electrons', grid, path=path)
|
|
|
|
|
|
return eabs * electrons
|
|
|
|
|
|
def log_loader(var, grid, path=''):
|
|
|
loader = special_vars.get(var[4:], default_loader);
|
|
|
return log10(abs(loader(var[4:], grid, path=path)) + LOG_EPS)
|
|
|
|
|
|
def minus_loader(var, grid, path=''):
|
|
|
real_var = var[1:]
|
|
|
loader = special_vars.get(real_var, default_loader);
|
|
|
return -loader(real_var, grid, path=path)
|
|
|
|
|
|
def townsend(efield, efield0):
|
|
|
return efield * exp(-efield0 / abs(efield))
|
|
|
|
|
|
def impact_loader(var, grid, path=''):
|
|
|
eabs = default_loader('eabs', grid, path=path)
|
|
|
electrons = default_loader('electrons', grid, path=path)
|
|
|
|
|
|
return townsend(eabs, 1.0) * electrons
|
|
|
|
|
|
def los_loader(var, grid, path=''):
|
|
|
# Only used here:
|
|
|
from losight import losight
|
|
|
real_var = var[var.index('_') + 1:]
|
|
|
#pre_loader = special_vars.get(real_var, default_loader)
|
|
|
|
|
|
#lr = flex_load(input_name('r', grid), path=path)
|
|
|
#lz = flex_load(input_name('z', grid), path=path)
|
|
|
#m = pre_loader(real_var, grid, path=path)
|
|
|
lr, lz, m = load_fine(real_var, grid, path, uptolevel=2)
|
|
|
|
|
|
return lr, lz, losight(lr, lz, m)
|
|
|
|
|
|
special_vars = {'ccharge': charge_loader,
|
|
|
'd_charge': d_charge_loader,
|
|
|
'log_electrons': log_loader,
|
|
|
'log_ions': log_loader, 'log_ccharge': log_loader,
|
|
|
'log_charge': log_loader,
|
|
|
'log_phi': log_loader,
|
|
|
'log_er': log_loader,
|
|
|
'_er': minus_loader,
|
|
|
'_ez': minus_loader,
|
|
|
'_charge': minus_loader,
|
|
|
'log_ez': log_loader,
|
|
|
'log_eabs': log_loader,
|
|
|
'log_error': log_loader,
|
|
|
'log_d_electrons': log_loader,
|
|
|
'log_d_ions': log_loader,
|
|
|
'current': current_loader,
|
|
|
'impact': impact_loader,
|
|
|
'los_impact': los_loader,
|
|
|
'log_los_impact': log_loader
|
|
|
}
|
|
|
|
|
|
def input_name(var, grid):
|
|
|
return "%s.%s.tsv" % (var, grid)
|
|
|
|
|
|
def output_name(var, grid, ext='png', what=""):
|
|
|
return "%s%s.%s.%s" % (var, what, grid, ext)
|
|
|
|
|
|
RE_FNAME = re.compile(r".*\.(\w\d{3}).*\.tsv")
|
|
|
def tar_name(fname):
|
|
|
m = RE_FNAME.match(fname)
|
|
|
if not m:
|
|
|
raise ValueError("File-name format %s not recognized." % fname)
|
|
|
|
|
|
gid = m.group(1)
|
|
|
return "all.%s.tgz" % gid
|
|
|
|
|
|
def old_tar_name(fname):
|
|
|
try:
|
|
|
for i in range(fname.index(".C") + 2, len(fname)):
|
|
|
if not fname[i].isdigit():
|
|
|
return 'all.%s.tgz' % fname[fname.index(".C") + 1:i]
|
|
|
|
|
|
except ValueError:
|
|
|
indx = fname.index(".P")
|
|
|
return 'all.%s.tgz' % fname[indx + 1:indx + 5]
|
|
|
|
|
|
return fname
|
|
|
|
|
|
def get_subgrids(gid, path=output_dir):
|
|
|
try:
|
|
|
# The '.' is a dirty hack to avoid changing tar_name, which expects
|
|
|
# a filename.
|
|
|
tar = tarfile.open(os.path.join(path, tar_name('.' + gid + '.tsv')))
|
|
|
gridnames = tar.getnames()
|
|
|
tar.close()
|
|
|
except (IOError, ValueError):
|
|
|
gridnames = [os.path.basename(g) for g in
|
|
|
glob.glob(os.path.join(path, "z.%s*.tsv" % gid))]
|
|
|
|
|
|
grids = [g for z, g, tsv in [nam.split('.') for nam in gridnames]
|
|
|
if z == 'z']
|
|
|
ret = [g for g in grids if g[0:len(gid)] == gid and g != gid]
|
|
|
|
|
|
return ret
|
|
|
|
|
|
def all_grids(pattern="C???", path=output_dir):
|
|
|
i = glob.glob(os.path.join(path, "all.%s.tgz" % pattern))
|
|
|
if not i:
|
|
|
i = glob.glob(os.path.join(path, "z.%s.tsv" % pattern))
|
|
|
|
|
|
i.sort()
|
|
|
for fn in i:
|
|
|
yield fn.split('.')[-2]
|
|
|
|
|
|
def last_grid(pattern="C???", path=output_dir):
|
|
|
i = glob.glob(os.path.join(path, "all.%s.tgz" % pattern))
|
|
|
i.sort()
|
|
|
try:
|
|
|
yield i[-1].split('.')[-2]
|
|
|
except IndexError:
|
|
|
# There are no grids
|
|
|
pass
|
|
|
|
|
|
LOAD_VAR_MEMOIZE = {}
|
|
|
|
|
|
def clear_load_var_memoize():
|
|
|
global LOAD_VAR_MEMOIZE
|
|
|
del LOAD_VAR_MEMOIZE
|
|
|
LOAD_VAR_MEMOIZE = {}
|
|
|
|
|
|
_default_rescaling = {}
|
|
|
|
|
|
def set_def_rescaling(rescaling):
|
|
|
global _default_rescaling
|
|
|
|
|
|
if isinstance(rescaling, str):
|
|
|
mainmod = sys.modules[__name__]
|
|
|
rescaling = getattr(mainmod, rescaling)
|
|
|
|
|
|
if rescaling is None:
|
|
|
rescaling = {}
|
|
|
|
|
|
_default_rescaling = rescaling
|
|
|
|
|
|
|
|
|
def load_var(var, grid, path='', ntheta=None, opposite=False,
|
|
|
rescaling=None, zcross=None):
|
|
|
""" Loads a variable read from a grid inside a matrix.
|
|
|
Returns a tuple formed by lr, lz, m, where lr and lz are 1d vectors
|
|
|
representing the gridpoints in r and z. """
|
|
|
|
|
|
global _default_rescaling
|
|
|
|
|
|
try:
|
|
|
id, var = var.split(':')
|
|
|
path = os.path.join(BASE_PATH, id)
|
|
|
except ValueError:
|
|
|
pass
|
|
|
|
|
|
if rescaling is None:
|
|
|
rescaling = _default_rescaling
|
|
|
|
|
|
k = (path, var, grid, rescaling.get(var, 1.0), ntheta, zcross)
|
|
|
if k in LOAD_VAR_MEMOIZE:
|
|
|
return LOAD_VAR_MEMOIZE[k]
|
|
|
|
|
|
loader = special_vars.get(var, default_loader);
|
|
|
|
|
|
# Load matrix
|
|
|
m = loader(var, grid, path=path)
|
|
|
|
|
|
try:
|
|
|
ltheta = flex_load(input_name('theta', grid), path=path)
|
|
|
except:
|
|
|
ltheta = array([0])
|
|
|
|
|
|
# Some loaders already provide lr, lz and reshaped m
|
|
|
if isinstance(m, tuple):
|
|
|
lr, lz, m = m
|
|
|
else:
|
|
|
lr = flex_load(input_name('r', grid), path=path)
|
|
|
lz = flex_load(input_name('z', grid), path=path)
|
|
|
|
|
|
# Reshape
|
|
|
m = reshape(m, (ltheta.shape[0], lr.shape[0], lz.shape[0]))
|
|
|
|
|
|
|
|
|
# Remove the theta dimension
|
|
|
if ntheta is None:
|
|
|
ntheta = min(ltheta.shape[0] - 1, 1)
|
|
|
|
|
|
if zcross is None:
|
|
|
mr = m[ntheta, :, :]
|
|
|
else:
|
|
|
iz = (zcross - lz[0]) / (lz[1] - lz[0])
|
|
|
mr = m[:, :, iz]
|
|
|
|
|
|
lz = numpy.outer(numpy.cos(ltheta), lr)
|
|
|
lr = numpy.outer(numpy.sin(ltheta), lr)
|
|
|
|
|
|
|
|
|
if opposite:
|
|
|
max_theta = ltheta.shape[0] - 1
|
|
|
#ntheta2 = (ntheta + max_theta / 2 - 1) % max_theta
|
|
|
ntheta2 = (ntheta + max_theta / 2) % max_theta + 2
|
|
|
|
|
|
mr = concatenate((flipud(m[ntheta2, :, :]), mr), axis=0)
|
|
|
lr = concatenate((-lr[::-1], lr))
|
|
|
|
|
|
|
|
|
try:
|
|
|
f = rescaling['%s_f' % var]
|
|
|
mr = f(lz) * mr
|
|
|
except KeyError:
|
|
|
if not var.startswith('log_'):
|
|
|
mr = rescaling.get(var, 1.0) * mr
|
|
|
else:
|
|
|
mr = log10(rescaling.get(var[4:], 1.0)) + mr
|
|
|
|
|
|
lr = rescaling.get('r', 1.0) * lr + rescaling.get('r0', 0.0)
|
|
|
lz = rescaling.get('z', 1.0) * lz + rescaling.get('z0', 0.0)
|
|
|
|
|
|
if lz[0] > lz[-1]:
|
|
|
mr = numpy.fliplr(mr)
|
|
|
lz = numpy.flipud(lz)
|
|
|
|
|
|
LOAD_VAR_MEMOIZE[k] = (lr, lz, mr)
|
|
|
return lr, lz, mr
|
|
|
|
|
|
|
|
|
def load_axis(var, grid, path=output_dir, col=0, **kwargs):
|
|
|
""" Reads the values along the minimum r of var"""
|
|
|
lr, lz, m = load_var(var, grid, path=path, **kwargs)
|
|
|
return lz, m[col,:]
|
|
|
|
|
|
def save_var(var, grid, lr, lz, m):
|
|
|
save(input_name('r', grid), lr)
|
|
|
save(input_name('z', grid), lz)
|
|
|
save(input_name(var, grid), m)
|
|
|
|
|
|
|
|
|
# For potential plots, we use a lot more contours
|
|
|
NCONTOURS = {'phi': 100, 'charge': 20}
|
|
|
|
|
|
def add_to_dict(thedict, key):
|
|
|
def doit(f):
|
|
|
thedict[key] = f
|
|
|
return f
|
|
|
|
|
|
return doit
|
|
|
|
|
|
def mirror_var(lr, m):
|
|
|
m = concatenate((flipud(m), m), axis=0)
|
|
|
lr = concatenate((-lr[::-1], lr))
|
|
|
|
|
|
return lr, m
|
|
|
|
|
|
def mirrored_var(lr, m):
|
|
|
m = flipud(m)
|
|
|
lr = -lr[::-1]
|
|
|
|
|
|
return lr, m
|
|
|
|
|
|
prev_contours = None
|
|
|
|
|
|
def plot_var(var, grid, mode='pcolor', aspect=True, mirror=False,
|
|
|
setaxis=True, finish=True, var_range=None, path='',
|
|
|
vector=None, ntheta=None, rescaling={}, units={},
|
|
|
fontsize=18, rlabel='r', zlabel='z', clabel=None,
|
|
|
rrange=None, zrange=None,
|
|
|
saffman=None, cmap=None, color='black', linewidth=None,
|
|
|
opposite=False, normranges=False, subgrids=[],
|
|
|
touchaxis=True, frame=False, superimpose=None, ncontours=None,
|
|
|
field=None, zcross=None, fine=False, mirrored=None,
|
|
|
subframe=False):
|
|
|
|
|
|
if not fine:
|
|
|
lr, lz, m = load_var (var, grid, path=path, ntheta=ntheta,
|
|
|
opposite=opposite, rescaling=rescaling,
|
|
|
zcross=zcross)
|
|
|
else:
|
|
|
def _loader(_var, _grid, path='.'):
|
|
|
return load_var (_var, _grid, path=path, ntheta=ntheta,
|
|
|
opposite=opposite, rescaling=rescaling,
|
|
|
zcross=zcross)
|
|
|
lr, lz, m = load_fine(var, grid, path=path, loader=_loader)
|
|
|
|
|
|
|
|
|
if mirror:
|
|
|
lr, m = mirror_var(lr, m)
|
|
|
|
|
|
if mirrored == -1:
|
|
|
lr, m = mirrored_var(lr, m)
|
|
|
|
|
|
if zcross is None:
|
|
|
[z, r] = meshgrid(lz, lr)
|
|
|
else:
|
|
|
z, r = lz, lr
|
|
|
|
|
|
if field is not None and var == 'phi':
|
|
|
m = m - field * z
|
|
|
|
|
|
modes = {}
|
|
|
@add_to_dict(modes, 'pcolor')
|
|
|
def plot_pcolor():
|
|
|
if var_range is not None:
|
|
|
vmin, vmax = var_range
|
|
|
norm = normalize(vmin=vmin, vmax=vmax)
|
|
|
else:
|
|
|
vmin, vmax = None, None
|
|
|
norm = normalize()
|
|
|
|
|
|
if hasattr(cmap, 'center'):
|
|
|
# A dynamic colormap
|
|
|
if vmin is None or vmax is None:
|
|
|
vmin = min(m.flat)
|
|
|
vmax = max(m.flat)
|
|
|
|
|
|
cmap.center = -vmin / (vmax - vmin)
|
|
|
|
|
|
pcolor(r, z, m, shading='flat', vmin=vmin, vmax=vmax,
|
|
|
norm=norm, cmap=cmap)
|
|
|
|
|
|
@add_to_dict(modes, 'contour')
|
|
|
def plot_contour():
|
|
|
global prev_contours
|
|
|
ncont =ncontours or NCONTOURS.get(var, 1)
|
|
|
|
|
|
if not prev_contours:
|
|
|
prev_contours = contour(r, z, m, ncont,
|
|
|
colors=color, linewidths=linewidth)
|
|
|
else:
|
|
|
contour(r, z, m, prev_contours[0], colors='black')
|
|
|
|
|
|
@add_to_dict(modes, 'neg_shape')
|
|
|
def plot_neg_shape():
|
|
|
themin = min(m.flat)
|
|
|
contour(r, z, m, [themin / 2], colors=color, linewidth=linewidth)
|
|
|
|
|
|
@add_to_dict(modes, 'neg_filled')
|
|
|
def plot_neg_filled():
|
|
|
themin = min(m.flat)
|
|
|
contourf(r, z, m, [themin / 2, themin], cmap=cmap)
|
|
|
|
|
|
@add_to_dict(modes, 'pos_shape')
|
|
|
def plot_pos_shape():
|
|
|
themax = max(m.flat)
|
|
|
contour(r, z, m, [themax / 2], colors=color, linewidth=linewidth)
|
|
|
|
|
|
@add_to_dict(modes, 'pos_filled')
|
|
|
def plot_pos_filled():
|
|
|
themax = max(m.flat)
|
|
|
contourf(r, z, m, [themax / 2, themax], cmap=cmap)
|
|
|
|
|
|
@add_to_dict(modes, 'range')
|
|
|
def output_range():
|
|
|
themax, themin = max(m.flat), min(m.flat)
|
|
|
print "%g:%g" % (themin, themax)
|
|
|
|
|
|
modes[mode]()
|
|
|
|
|
|
if frame:
|
|
|
hlines([lz[0], lz[-1]], lr[0], lr[-1], "w", linewidth=1.5)
|
|
|
vlines([lr[0], lr[-1]], lz[0], lz[-1], "w", linewidth=1.5)
|
|
|
|
|
|
|
|
|
if saffman is not None:
|
|
|
plot_saffman(color='blue', *saffman)
|
|
|
|
|
|
if superimpose is not None:
|
|
|
try:
|
|
|
fname, shift, factor = superimpose
|
|
|
except ValueError:
|
|
|
fname, shift, factor = superimpose, 0.0, 1.0
|
|
|
|
|
|
plot_xyfile(fname, shift=shift, factor=factor, color='blue')
|
|
|
|
|
|
if aspect:
|
|
|
axis('scaled')
|
|
|
|
|
|
if setaxis and zcross is None:
|
|
|
axis([0, 0.75, 0, 1.0])
|
|
|
xlim(lr[0], lr[-1])
|
|
|
ylim(lz[0], lz[-1])
|
|
|
|
|
|
if rrange:
|
|
|
if normranges:
|
|
|
factor = rescaling.get('r', 1.0)
|
|
|
rrange = [factor * rl for rl in rrange]
|
|
|
|
|
|
xlim(*rrange)
|
|
|
|
|
|
if zrange:
|
|
|
if normranges:
|
|
|
factor = rescaling.get('z', 1.0)
|
|
|
zrange = [factor * zl for zl in zrange]
|
|
|
|
|
|
ylim(*zrange)
|
|
|
|
|
|
#@print_ret
|
|
|
def get_label(x):
|
|
|
theunits = units.get(x, '')
|
|
|
if theunits:
|
|
|
return (r'$%s$ [$\mathdefault{%s}$]'
|
|
|
% (pretty_var(x), theunits))
|
|
|
else:
|
|
|
return r'%s' % pretty_var(x)
|
|
|
|
|
|
|
|
|
if touchaxis:
|
|
|
xlabel(get_label(rlabel), fontsize=fontsize)
|
|
|
ylabel(get_label(zlabel), fontsize=fontsize)
|
|
|
|
|
|
|
|
|
xmin, xmax = None, None
|
|
|
|
|
|
ax = gca()
|
|
|
ax.xaxis.set_major_formatter(FancyFormatter())
|
|
|
ax.yaxis.set_major_formatter(FancyFormatter())
|
|
|
|
|
|
|
|
|
|
|
|
if mode == 'pcolor' and touchaxis:
|
|
|
|
|
|
#cax = colorbar(format="%.3g", cax=ncax)
|
|
|
|
|
|
# This was for the Saffman-Taylor paper:
|
|
|
#cax = colorbar(format="%.3g", fraction=0.20, pad=0.00, aspect=12)
|
|
|
|
|
|
format = FancyFormatter()
|
|
|
if 'log' in var:
|
|
|
format = r'$\mathdefault{10^{%d}}$'
|
|
|
|
|
|
#format = 'log' in var and '$10^{%d}$' or "%.3g"
|
|
|
cax = colorbar(format=format)
|
|
|
#for t in cax.get_yticklabels():
|
|
|
# t.set_fontsize(fontsize - 2)
|
|
|
#gcf().sca(ax)
|
|
|
#ipshell()
|
|
|
xmin, xmax = cax.get_clim()
|
|
|
|
|
|
#trans = blend_xy_sep_transform(cax.transAxes, cax.transData)
|
|
|
#xpos = 2.5
|
|
|
#xpos = 4.0
|
|
|
#xpos = -0.5
|
|
|
#cax.text(xpos, 0.5 * (xmin + xmax),
|
|
|
# get_label(var), ha='right', va='center', transform=trans,
|
|
|
# rotation='vertical', fontsize=fontsize)
|
|
|
if clabel is None:
|
|
|
clabel = get_label(var)
|
|
|
|
|
|
cax.set_label(clabel, fontsize=fontsize)
|
|
|
|
|
|
if vector is not None:
|
|
|
if ':' in vector:
|
|
|
grid, vector = vector.split(':')
|
|
|
|
|
|
d = {'field': False, 'current': True}
|
|
|
plot_some_vect(grid, path=path, mirror=mirror, current=d[vector],
|
|
|
rescaling=rescaling)
|
|
|
|
|
|
|
|
|
for subgrid in subgrids:
|
|
|
if mirror or mirrored is not None:
|
|
|
mirrs = (-1, 1)
|
|
|
else:
|
|
|
mirrs = (1,)
|
|
|
|
|
|
for mirr_ in mirrs:
|
|
|
plot_var(var, subgrid, mode=mode, setaxis=False, finish=False,
|
|
|
path=path, units=units, opposite=False, cmap=cmap,
|
|
|
mirror=False, frame=subframe, var_range=[xmin, xmax],
|
|
|
rescaling=rescaling, mirrored=mirr_, rrange=rrange,
|
|
|
zrange=zrange, aspect=False, touchaxis=False,
|
|
|
subframe=subframe)
|
|
|
|
|
|
|
|
|
if finish:
|
|
|
show()
|
|
|
|
|
|
def plot_saffman(width, C, color='white'):
|
|
|
eps = 0.05
|
|
|
A = width / math.pi
|
|
|
x = arange(-width / 2 + eps, width / 2 - eps, 0.01)
|
|
|
y = A * log(cos(x / A))+ float(C)
|
|
|
|
|
|
#lam = 0.5
|
|
|
#L = 2 * width
|
|
|
#y = (L * ( 1 - lam) / (2 * math.pi)
|
|
|
# * log (0.5 * (1 + cos(2 * math.pi * x / (lam * L))))) + float(C)
|
|
|
|
|
|
plot(x, y, linewidth=1.5, color=color)
|
|
|
|
|
|
def plot_xyfile(fname, shift=0, factor=1.0, color='white'):
|
|
|
data = load(fname)
|
|
|
|
|
|
plot(factor * data[:, 0], factor * data[:, 1] + shift,
|
|
|
linewidth=1.5, color=color)
|
|
|
|
|
|
|
|
|
def plot_some_vect(grid, path='', dr=1, dz=1,
|
|
|
rescaling={}, mirror=False, current=False):
|
|
|
""" Plots either the electric field (if current == False) or the
|
|
|
current vector field (if current == True)"""
|
|
|
|
|
|
lr, lz, er = load_var('er', grid, path=path, rescaling=rescaling)
|
|
|
void, void, ez = load_var('ez', grid, path=path, rescaling=rescaling)
|
|
|
|
|
|
stepr, stepz = int(dr / (lr[1] - lr[0])), int(dr / (lz[1] - lz[0])),
|
|
|
|
|
|
lr = lr[stepr / 2::stepr]
|
|
|
lz = lz[::stepz]
|
|
|
|
|
|
er = er[stepr / 2::stepr, 0::stepz]
|
|
|
ez = ez[stepr / 2::stepr, 0::stepz]
|
|
|
|
|
|
if mirror:
|
|
|
er = concatenate((-flipud(er), er), axis=0)
|
|
|
ez = concatenate((flipud(ez), ez), axis=0)
|
|
|
|
|
|
lr = concatenate((-lr[::-1], lr))
|
|
|
|
|
|
if current:
|
|
|
void, void, electrons = load_var('electrons', grid, path)
|
|
|
electrons = electrons[stepr / 2::stepr, 0::stepz]
|
|
|
if mirror:
|
|
|
electrons = concatenate((flipud(electrons), electrons), axis=0)
|
|
|
else:
|
|
|
electrons = 1.0
|
|
|
|
|
|
mr, mz = electrons * er, electrons * ez
|
|
|
|
|
|
eunit = 1.5e-3 * dr
|
|
|
|
|
|
plot_vect(lr, lz, mr, mz, eunit=eunit)
|
|
|
|
|
|
def plot_vect(lr, lz, mr, mz, eunit=1.0, threshold=0.0125):
|
|
|
EPS = 1e-15
|
|
|
|
|
|
mr = mr + EPS
|
|
|
mz = mz + EPS
|
|
|
|
|
|
N = sqrt(mr**2 + mz**2)
|
|
|
if threshold is not None:
|
|
|
mr[N > threshold]=0
|
|
|
mz[N > threshold]=0
|
|
|
|
|
|
Nmax = max(N.flat) * 0.00075
|
|
|
|
|
|
|
|
|
[z, r] = meshgrid(lz, lr)
|
|
|
#quiver(r, z, mr, mz, eunit * Nmax, pivot='middle',
|
|
|
# width=1.0 * Nmax, headwidth=500, color="white")
|
|
|
|
|
|
quiver(r, z, mr, mz, scale=Nmax / eunit, pivot='middle', color='w')
|
|
|
|
|
|
|
|
|
def draw_arrow( x, y, dx, dy, width=1, color = 'k' ):
|
|
|
ax = gca()
|
|
|
a = Arrow(x, y, dx, dy, width)
|
|
|
a.set_edgecolor(color)
|
|
|
a.set_facecolor(color)
|
|
|
ax.add_patch(a)
|
|
|
draw_if_interactive()
|
|
|
return a
|
|
|
|
|
|
def plot_axis(var, grid, format=None, xlim=None, col=0, rescaling={}):
|
|
|
if format is not None and '@@' in format:
|
|
|
fmt, linewidth = format.split('@@')
|
|
|
linewidth = float(linewidth)
|
|
|
else:
|
|
|
fmt = format
|
|
|
linewidth = 1.5
|
|
|
|
|
|
lz, m = load_axis (var, grid, col=col, rescaling=rescaling)
|
|
|
if format:
|
|
|
plot (lz, m, fmt, linewidth=linewidth)
|
|
|
else:
|
|
|
plot (lz, m, linewidth=1.5)
|
|
|
|
|
|
|
|
|
def save_fig(var, grid, ext, what="", ofile=None, **kwargs):
|
|
|
if ofile is None:
|
|
|
ofile = output_name(var, grid, ext, what=what)
|
|
|
|
|
|
savefig(ofile, **kwargs)
|
|
|
return ofile
|
|
|
|
|
|
def print_ret(f):
|
|
|
def _f(*args, **kwargs):
|
|
|
ret = f(*args, **kwargs)
|
|
|
print "%s -> %s" % (f.__name__, repr(ret))
|
|
|
return ret
|
|
|
return _f
|
|
|
|
|
|
|
|
|
|
|
|
class FancyFormatter(Formatter):
|
|
|
def __call__(self, x, pos=None):
|
|
|
formatted = "%.3g" % x
|
|
|
if 'e' in formatted:
|
|
|
m, e = formatted.split('e')
|
|
|
formatted = (r'$\mathdefault{%s}\cdot \mathdefault{10^{%d}}$'
|
|
|
% (m, int(e)))
|
|
|
|
|
|
return formatted
|