Source code for spacepy.pybats.rim

#!/usr/bin/env python
'''
Classes, functions, and methods for reading, writing, and plotting output
from the Ridley Ionosphere Model (RIM) and the similar legacy code, 
Ridley Serial.

Copyright 2010 Los Alamos National Security, LLC.

.. rubric:: Classes

.. autosummary::
    :toctree:
    :template: clean_class.rst

    Iono
    OvalDebugFile

.. rubric:: Functions

.. autosummary::
    :toctree:

    fix_format
    get_iono_cb
    tex_label
'''

import os
import numpy as np
from spacepy import deprecated
import spacepy.plot.apionly
from spacepy.plot import set_target
from spacepy.pybats import PbData, dmarray

[docs]@deprecated('0.2.2', 'Functionality integrated into `rim.Iono`.') def fix_format(filename, finalize=True): ''' Some 2D output files for RIM/RidleySerial have a broken format: values for the same lat-lon entry are split across two lines. This function detects and fixes such files in-place. The original file will not be overwritten if kwarg *finalize* is set to **False**. Parameters ========== filename : str The file to repair. Other Parameters ================ finalize : bool If **True**, overwrite the original file. If **False**, perform a dry run for testing purposes. Defaults to **True**. ''' # Slurp in entire file. This is necessary to detect the number of f. f = open(filename, 'r') raw = f.readlines() f.close() i = raw.index('NUMERICAL VALUES\n') nvars = int(raw[i+1].split()[0]) ntheta= int(raw[i+2].split()[0]) nphi = int(raw[i+3].split()[0]) # A broken file will have more than 4X lines than nTheta*nPhi. # This is because nLines = header + 2*nTheta*nPhi in a two-hemisphere file. if not len(raw) > 4*ntheta*nphi: return out = open('temp_iono_fix.txt', 'w') iNorth = raw.index('BEGIN NORTHERN HEMISPHERE\n') iSouth = raw.index('BEGIN SOUTHERN HEMISPHERE\n') # Rewrite header: out.writelines(raw[0:iNorth+1]) # Rewrite northern hemisphere. for l1, l2 in zip(raw[iNorth+1:iSouth:2], raw[iNorth+2:iSouth:2]): out.write(l1.strip()+l2) # Rewrite southern hemisphere. out.write(raw[iSouth]) for l1, l2 in zip(raw[iSouth+1::2], raw[iSouth+2::2]): out.write(l1.strip()+l2) out.close() # Move files. if finalize: os.remove(filename) os.rename(out.name, filename) return True
[docs]def get_iono_cb(ct_name='bwr'): ''' Several custom colorbars used by RIM and AMIE have become standard when visualizing data from these models. These are 'blue_white_red' and 'white_red', used for data that have positive and negative values and for data that have only positive values, respectively. This function builds and returns these colorbars when called with the initials of the color table name as the only argument. Other Parameters ================ ct_name : str Select the color table. Can be 'bwr' for blue-white-red or 'wr' for white-red. Defaults to 'bwr'. Examples ======== >>> bwr_map = get_iono_cb('bwr') >>> wr_map = get_iono_cb('wr') ''' from matplotlib.colors import LinearSegmentedColormap as lsc if ct_name=='bwr': table = { 'red': [(0.,0.,.0),(.34,0.,0.),(.5,1.,1.),(1.,1.,1.)], 'green':[(0.,0.,0.),(.35,1.,1.),(.66,1.,1.),(1.,0.,0.)], 'blue' :[(0.,1.,1.),(.5,1.,1.),(.66,0.,0.),(.85,0.,0.),(1.,.1,.1)] } cmap = lsc('blue_white_red',table) elif ct_name=='wr': table = { 'red': [(0.,1.,1.),(1.,1.,1.)], 'green':[(0.,1.,1.),(1.,0.,0.)], 'blue' :[(0.,1.,1.),(1.,.0,.0)] } cmap = lsc('white_red',table) return cmap
[docs]def tex_label(varname): ''' Many variable names used in the Ridley Ionosphere Model look much better in LaTeX format with their proper Greek letters. This function takes a variable name, and if it is recognized, returns a properly formatted string that uses MatPlotLib's MathText functionality to display the proper characters. If it is not recognized, the varname is returned. Parameters ========== varname : string The variable to convert to a LaTeX label. Examples ======== >>>tex_label('n_phi') '\\Phi_{Ionosphere}' >>>tex_label('Not Recognized') 'Not Recognized' ''' if varname[:2]=='n_' or varname[:2]=='s_': varname=varname[2:] known = { 'phi': r'$\Phi_{Ionosphere}$', 'sigmah':r'$\sigma_{Hall}$', 'sigmap':r'$\sigma_{Peder}$', 'jr':r'$J_{radial}$', 'mA/m^2':r'$\mu A/m^{2}$', 'W/m2':r'$W/m^2$', 'eV':'$eV$' } if varname in known: label = known[varname] else: label = varname return label
[docs]class Iono(PbData): ''' A class for handling 2D output from the Ridley Ionosphere Model. Instantiate an object as follows: >>> iono = rim.Iono('filename.idl') ...where filename.idl is the name of a RIM 2D output file. ''' def __init__(self, infile, *args, **kwargs): super(Iono, self).__init__(*args, **kwargs) self.attrs['file']=infile self.readascii() def readascii(self): ''' Read an ascii ".idl" output file and load the contents into the object. ''' from sys import version_info import re import datetime as dt from numpy import zeros, reshape import gzip # slurp entire file. if self.attrs['file'][-3:]=='.gz': try: infile = gzip.open(self.attrs['file'],'rt') except ValueError: #Python2.7 (Windows) compatibility infile = gzip.open(self.attrs['file']) else: infile = open(self.attrs['file'], 'r') raw = infile.readlines() infile.close() # Parse header title = raw[raw.index('TITLE\n')+1] self.attrs['title'] = title[title.index('"')+1:title.rindex('"')] i = raw.index('NUMERICAL VALUES\n') self.attrs['nvars'] = int(raw[i+1].split()[0]) self.attrs['ntheta']= int(raw[i+2].split()[0]) self.attrs['nphi'] = int(raw[i+3].split()[0]) # Convenience: nphi, ntheta = self.attrs['nphi'], self.attrs['ntheta'] i = raw.index('TIME\n') self.attrs['time'] = dt.datetime( int(raw[i+1].split()[0]), #year int(raw[i+2].split()[0]), #month int(raw[i+3].split()[0]), #day int(raw[i+4].split()[0]), #hour int(raw[i+5].split()[0]), #min int(raw[i+6].split()[0]), #sec int(raw[i+7].split()[0])*1000 #microsec ) i = raw.index('SIMULATION\n') self.attrs['iter'] = int(raw[i+1].split()[0]) self.attrs['simtime'] = float(raw[i+2].split()[0]) i = raw.index('DIPOLE TILT\n') self.tilt = zeros(2) self.tilt[0] = float(raw[i+1].split()[0]) self.tilt[1] = float(raw[i+2].split()[0]) i = raw.index('VARIABLE LIST\n') namevar = [] units = {} for j in range(i+1,i+self.attrs['nvars']+1): match = re.match(r'\s*\d+\s+([\w\s\W]+)\[([\w\s\W]+)\]',raw[j]) if match: name = (match.group(1).strip()).lower() namevar.append(name) units[name] = match.group(2).strip() else: raise ValueError('Could not parse %s' % raw[j]) ### Read all data ### # Create data arrays nPts = self.attrs['ntheta']*self.attrs['nphi'] for key in namevar: self['n_'+key] = dmarray(zeros(nPts), {'units':units[key]}) self['s_'+key] = dmarray(zeros(nPts), {'units':units[key]}) i = raw.index('BEGIN NORTHERN HEMISPHERE\n')+1 # Some compilers insert line breaks automatically when fortran format # string is not adequately specified. Let's see if that's the # case here: how many lines does it take to cover all variables? nvars, nvarline, nwrap = len(namevar), 0, 0 while nvarline<nvars: nvarline += len(raw[i+nwrap].split()) nwrap += 1 # Fill data arrays: for j in range(nPts): # Build list of variables; accounting for line-wrapping: parts = [] iLine = i + j*nwrap for iwrap in range(nwrap): parts += raw[iLine+iwrap].split() for k in range(self.attrs['nvars']): self['n_'+namevar[k]][j] = parts[k] i = raw.index('BEGIN SOUTHERN HEMISPHERE\n')+1 for j in range(nPts): parts = [] iLine = i + j*nwrap for iwrap in range(nwrap): parts += raw[iLine+iwrap].split() for k in range(self.attrs['nvars']): self['s_'+namevar[k]][j] = parts[k] # Create 2-D arrays. for key in namevar: nkey, skey = 'n_'+key, 's_'+key self[nkey] = reshape(self[nkey], (ntheta, nphi), 'F') self[skey] = reshape(self[skey], (ntheta, nphi), 'F') # Some extra grid info: self.dlon = self['n_psi' ][0,3]-self['n_psi' ][0,2] self.dlat = self['n_theta'][3,0]-self['n_theta'][2,0] def calc_I(self): ''' Integrate radial current to get the following values: I: The total radial current over one hemisphere. Iup: The total upward current over one hemisphere. Idown: The total downward current over one hemisphere. Values are stored in the object with the prefix 'n_' or 's_' to indicate the hemisphere and may be accessed via self['n_I'], etc. Values are stored in units of mega Amps. Parameters ========== Returns ======= True Examples ======== >>> a = rim.Iono('spacepy/tests/data/pybats_test/it000321_104510_000.idl.gz') >>> a.calc_I() >>> print(a['n_Iup']) ''' # Calculate some physically meaningful values/units units = 1E-6*1E-6 # micro amps to amps, amps to MegaAmps R = (6371.0+110.0)*1000.0 # Radius of Earth + iono altitude dTheta = np.pi*self.dlat/180. dPhi = np.pi*self.dlon/180. # -----NORTHERN HEMISPHERE----- # Get relevant values: colat = self['n_theta']*np.pi/180. integrand = self['n_jr']*np.sin(colat)*dTheta*dPhi # Get locations of "up" and "down" loc_up = self['n_jr']>0 loc_do = self['n_jr']<0 self['n_I'] = units*R**2 * np.sum(integrand) self['n_Iup'] = units*R**2 * np.sum(integrand[loc_up]) self['n_Idown'] = units*R**2 * np.sum(integrand[loc_do]) # -----SOUTHERN HEMISPHERE----- # Get relevant values: colat = self['s_theta']*np.pi/180. integrand = self['s_jr']*np.sin(colat)*dTheta*dPhi # Get locations of "up" and "down" loc_up = self['s_jr']>0 loc_do = self['s_jr']<0 self['s_I'] = units*R**2 * np.sum(integrand) self['s_Iup'] = units*R**2 * np.sum(integrand[loc_up]) self['s_Idown'] = units*R**2 * np.sum(integrand[loc_do]) def add_cont(self, var, target=None, n=24, maxz=False, lines=True, cmap=False, add_cbar=False, label=None, loc=111, xticksize=12, yticksize=12, max_colat=40, **kwargs): ''' Create a polar contour of variable *var*. Plot will be either drawn on a new matplotlib figure and axes, or you can specify a plot target using the *target* kwarg. Parameters ========== var : str The name of the variable to plot. Returns ======= fig : matplotlib figure object ax : matplotlib axes object cnt : matplotlib contour object cb : matplotlib colorbar object Other Parameters ================ target : Figure or Axes If None (default), a new figure is generated from scratch. If a matplotlib Figure object, a new axis is created to fill that figure. If a matplotlib Axes object, the plot is placed into that axis. loc : int Use to specify the subplot placement of the axis (e.g. loc=212, etc.) Used if target is a Figure or None. Default 111 (single plot). n : int Set number of levels. Should be a multiple of 3 for best match between filled and traced contours. Default is 21. lines : bool Add unfilled black solid/dashed contours to plot for additional contrast. Default is **True**. maxz : real Set the max/min value for the color bar. Default is set by data. max_colat : real Set the co-latitude range of the plot, in degrees. Defaults to 40. cmap : str Set the colormap. Default is to autoselect using classic IE maps. Can be 'bwr', 'wr', or any name of a matplotlib colar map. add_cbar : bool Add colorbar to plot. Default is **False** which will not add one to the plot. label : str or None Label to place at top of plot. Defaults to variable name. xticksize : int Size of longitude markers in text points. Defaults to 12. yticksize : int Size of latitude markers in text points. Defaults to 12. Extra keywords are passed to the contourf routine. ''' # Get only what we need to decrease runtime. from math import pi from numpy import linspace from matplotlib.colors import Normalize from matplotlib.ticker import MaxNLocator, MultipleLocator from matplotlib.pyplot import clabel, colorbar fig, ax = set_target(target, polar=True, loc=loc) hemi = var[:2] # user defined variables may not have hemisphere marking. # Assume nothern hemi in those cases. if hemi!='n_' and hemi!='s_': hemi = 'n_' # Set levels and ticks: if label==None: label=tex_label(var) lt_labels = ['06', label, '18', '00'] xticks = [ 0, pi/2, pi, 3*pi/2] lct = MultipleLocator(10) minz = self[var].min() if minz < 0.0: if not maxz: maxz = max([abs(minz),self[var].max()]) crange = Normalize(vmin=-1.*maxz, vmax=maxz) levs = linspace(-1.*maxz, maxz, n) else: if not maxz: maxz = self[var].max() crange = Normalize(vmin=0., vmax=maxz) levs = linspace(0., maxz, n) # Get color map if not given: if not cmap: if self[var].min() >= 0.0: cmap=get_iono_cb('wr') else: cmap=get_iono_cb('bwr') # Set the latitude based on hemisphere: theta = self[hemi+'theta'] if 's_' in hemi: theta = 180-self[hemi+'theta'] # Create contour: cnt1 = ax.contourf(self[hemi+'psi']*pi/180.0+pi/2., theta, np.array(self[var]), levs, norm=crange, cmap=cmap, **kwargs) # Set xtick label size, increase font of top label. labels = ax.get_xticklabels() for l in labels: l.set_size(xticksize) labels[1].set_size(xticksize*1.25) if lines: nk = int(round(n/3.0)) cnt2 = ax.contour(self[hemi+'psi']*pi/180.0+pi/2., theta, np.array(self[var]), nk, colors='k') #clabel(cnt2,fmt='%3i',fontsize=10) if add_cbar: cbarticks = MaxNLocator(7) cbar = colorbar(cnt1, ticks=cbarticks, shrink=0.75, pad=0.08, ax=ax) cbar.set_label(tex_label(self[var].attrs['units'])) else: cbar=False ax.set_xticks(xticks) ax.set_xticklabels(lt_labels) ax.yaxis.set_major_locator(lct) ax.set_ylim([0,max_colat]) # Use text function to manually add pretty ticks. ax.set_yticklabels('') # old ticks off. opts = {'size':yticksize, 'rotation':-45, 'ha':'center', 'va':'center'} for theta in [80.,70.,60.]: txt = '{:02.0f}'.format(theta)+r'$^{\circ}$' ax.text(pi/4., 90.-theta, txt, color='w', weight='heavy', **opts) ax.text(pi/4., 90.-theta, txt, color='k', weight='light', **opts) return fig, ax, cnt1, cbar
[docs]class OvalDebugFile(PbData): ''' The auroral oval calculations in RIM may spit out special debug files that are extremely useful. This class handles reading and plotting the data contained within those files. ''' def __init__(self, infile, *args, **kwargs): super(OvalDebugFile, self).__init__(*args, **kwargs) self.attrs['file']=infile self._readascii() def _readascii(self): ''' Loads data & populates object upon instantiation. ''' import datetime as dt # Slurp in lines: f = open(self.attrs['file']) lines = f.readlines() f.close() # Parse header to get longitude in radians: l = lines.pop(0) self['lon'] = np.array(l.split('=')[-1].split(), dtype=float)* np.pi/180. l = lines.pop(0) # Some helper vars: nLons = self['lon'].size nLine = len(lines) # Create container arrays: self['time'] = np.zeros(nLine, dtype=object) self['oval'] = np.zeros( (nLine, nLons) ) # Parse rest of file: for j,l in enumerate(lines): self['time'][j] = dt.datetime.strptime(l[:19], '%Y %m %d %H %M %S') self['oval'][j,:] = l.split()[7:] def get_oval(self, time, interp=True): ''' Given a datetime, *time*, interpolate the oval position to that time. If *time* is outside the range of self['time'], an exception will be raised. Linear interpolation is used unless kwarg *interp* is False. In that case, the value of the oval nearest to *time* will be used without interpolation. Parameters ========== time : datetime.datetime The time at which the oval is requested. Returns ======= oval : numpy.ndarray Oval location at time *time* in radians. Zero corresponds to local noon. Other Parameters ================ interp : bool Set interpolation behavior. If **True**, linear interpolation is used. If **False**, the oval location nearest to *time* is used without interpolation. Examples ======== >>> ''' from matplotlib.dates import date2num # If *time* is outside the bounds of self['time'], raise exception # to avoid extrapolation. if time<self['time'][0] or time>self['time'][-1]: raise ValueError('Given time outside object range ' + 'and requires extrapolation') # Turn datetimes into numbers. time = date2num(time) ovaltime = date2num(self['time']) # Start by obtaining the indices of the time array that bound index = np.arange(self['time'].size) i1 = index[time>=ovaltime][-1] # Value before time i2 = index[time<=ovaltime][ 0] # Value after time # Check for exact matches: dT1 = np.abs(ovaltime[i1]-time) dT2 = np.abs(ovaltime[i2]-time) if dT1==0: return self['oval'][i1] if dT2==0: return self['oval'][i2] # If no interpolation, just send back nearest neighbor: if not interp: if dT1<dT2: return self['oval'][i1] else: return self['oval'][i2] # If interpolation, get slope and y-intercept vectors: dT = ovaltime[i2]-ovaltime[i1] m = (self['oval'][i2]-self['oval'][i1])/dT b = self['oval'][i2] - m*ovaltime[i2] # Interpolate and return: return m*time + b def add_oval_line(self, time, *args, **kwargs): ''' Adds the location of the auroral oval at time *time* as a line plot on to *target*, which must be a Matplotlib figure/axes. If *target* not given, a new axes object is created. If *target* is not an existing axes object, a new axes object is created and customized to be an ionosphere polar plot. Extra arguments/kwargs are sent to :meth:`matplotlib.pyplot.plot`. Parameters ========== time : integer or datetime Sets the time at which to plot the line. If an integer is used, the integer is used to index the internal array naively. If a datetime object is given and is within the bounds of self['time'], the oval location will be interpolated to *time*. Control of the interpolation happens via the *interp* keyword. Returns ======= fig : matplotlib figure object ax : matplotlib axes object line : matplotlib line object Other Parameters ================ target : Figure or Axes If None (default), a new figure is generated from scratch. If a matplotlib Figure object, a new axis is created to fill that figure. If a matplotlib Axes object, the plot is placed into that axis. loc : int Use to specify the subplot placement of the axis (e.g. loc=212, etc.) Used if target is a Figure or None. Default 111 (single plot). interp : bool Control the behavior of time interpolation when the *time* argument is given as a datetime object. Defaults to **True**, meaning that oval location is interpolated in time. Examples ======== >>> ''' import datetime as dt # Handle kwargs not being handed to plot: target=None; loc=111; interp=True if 'target' in kwargs: target=kwargs.pop('target') if 'interp' in kwargs: interp=kwargs.pop('interp') if 'loc' in kwargs: loc=kwargs.pop('loc') # Set plot targets: fig, ax = set_target(target, polar=True, loc=loc) # Get oval interpolated in time: if type(time) == dt.datetime: oval = self.get_oval(time, interp=interp) elif type(time) == int: oval = self['oval'][time] else: raise TypeError('Unrecognized type for *time*:'+type(time)) # Plot oval: line = ax.plot(self['lon']+np.pi/2., oval, *args, **kwargs) return fig, ax, line