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..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