Exemplo n.º 1
0
def makeImfInput(indata, fname=None, keymap=None):
    """Make an SWMF IMF input file from an input SpaceData"""
    if keymap is None:
        keymap = {
            'DateTime': 'time',
            'Bx': 'bx',
            'By': 'by',
            'Bz_OB': 'bz',
            'V_sw': 'ux',
            'Vy': 'uy',
            'Vz': 'uz',
            'Den_P': 'rho',
            'Plasma_temp': 'temp',
        }
    numpts = indata['DateTime'].shape[0]
    swmfdata = pybats.ImfInput(filename=False, load=False, npoints=numpts)
    for key_o, newkey in keymap.items():
        if newkey == 'ux':
            swmfdata[newkey] = -1 * np.abs(dm.dmcopy(indata[key_o]))
        else:
            swmfdata[newkey] = dm.dmcopy(indata[key_o])
    swmfdata.attrs['coor'] = 'GSM'
    if fname is not None:
        swmfdata.write(fname)
    return swmfdata
Exemplo n.º 2
0
def applyRefractory(process1, period):
    '''Apply a refractory period to an input discrete event time sequence

    All events in the refractory period are removed from the point process.

    Parameters
    ==========
    process1 : iterable
        an iterable of datetimes, or a spacepy.time.Ticktock
    period : datetime.timedelta
        length of refractory period

    Returns
    =======
    keep : iterable
        returns pruned set of datetimes with same type as input
        NOTE: array subclasses will be lost
    '''
    import spacepy.time as spt
    if isinstance(process1, spt.Ticktock):
        #SpacePy Ticktock
        p1 = dm.dmcopy(process1.UTC).tolist()
        tickt = True
    elif isinstance(process1[0], dt.datetime):
        #iterable containing datetimes
        p1 = dm.dmcopy(process1)
        try:
            p1 = p1.tolist()
            wasArr = True
        except:
            wasArr = False
        tickt = False
    else:
        raise NotImplementedError(
            'Input process must be a list/array of datetimes, or a spacepy.time.Ticktock'
        )

    try:
        assert period.seconds
    except AssertionError:
        raise AttributeError('period must be a datetime.timedelta')

    done = len(p1) < 2
    keep, discard = [], []
    while not done:
        t1 = p1[0]
        t2 = t1 + period
        inds = tb.tOverlapHalf([t1, t2], p1[1:])
        for idx in inds:
            discard.append(p1.pop(idx + 1))
        keep.append(p1.pop(0))  # put test element into keep array
        done = len(p1) < 2

    if tickt:
        return spt.Ticktock(keep)
    else:
        if wasArr:
            return np.array(keep)
        else:
            return keep
Exemplo n.º 3
0
def applyRefractory(process1, period):
    '''Apply a refractory period to an input discrete event time sequence

    All events in the refractory period are removed from the point process.

    Parameters
    ==========
    process1 : iterable
        an iterable of datetimes, or a spacepy.time.Ticktock
    period : datetime.timedelta
        length of refractory period

    Returns
    =======
    keep : iterable
        returns pruned set of datetimes with same type as input
        NOTE: array subclasses will be lost
    '''
    import spacepy.time as spt
    if isinstance(process1, spt.Ticktock):
        #SpacePy Ticktock
        p1 = dm.dmcopy(process1.UTC).tolist()
        tickt = True
    elif isinstance(process1[0], dt.datetime):
        #iterable containing datetimes
        p1 = dm.dmcopy(process1)
        try:
            p1 = p1.tolist()
            wasArr = True
        except:
            wasArr = False
        tickt = False
    else:
        raise NotImplementedError('Input process must be a list/array of datetimes, or a spacepy.time.Ticktock')

    try:
        assert period.seconds
    except AssertionError:
        raise AttributeError('period must be a datetime.timedelta')

    done = len(p1)<2
    keep, discard = [], []
    while not done:
        t1 = p1[0]
        t2 = t1 + period
        inds = tb.tOverlapHalf([t1, t2], p1[1:])
        for idx in inds:
            discard.append(p1.pop(idx+1))
        keep.append(p1.pop(0)) # put test element into keep array
        done = len(p1)<2

    if tickt:
        return spt.Ticktock(keep)
    else:
        if wasArr:
            return np.array(keep)
        else:
            return keep
Exemplo n.º 4
0
 def test_dmcopy(self):
     """dmcopy should copy datamodel objects"""
     a = dm.SpaceData()
     a[1] = dm.dmarray([1,2,3], attrs={1:1})
     b = dm.dmcopy(a)
     self.assertFalse(a is b) # they are not the same memory
     np.testing.assert_almost_equal(a[1], b[1])
     self.assertEqual(a[1].attrs, b[1].attrs)
     b = dm.dmcopy(a[1])
     np.testing.assert_almost_equal(a[1], b)
     self.assertEqual(a[1].attrs, b.attrs)
     a = np.arange(10)
     b = dm.dmcopy(a)
     np.testing.assert_almost_equal(a, b)
     a = [1,2,3]
     b = dm.dmcopy(a)
     self.assertEqual(a, b)
Exemplo n.º 5
0
    def plotSpectrogram(self, ecol=0, **kwargs):
        '''
        Plot a spectrogram of the flux along the requested orbit, as a function of Lm and time

        Other Parameters
        ----------------
        zlim : list
            2-element list with upper and lower bounds for color scale
        colorbar_label : string
            text to appear next to colorbar (default is 'Flux' plus the units)
        ylabel : string
            text to label y-axis (default is 'Lm' plus the field model name)
        title : string
            text to appear above spectrogram (default is climatology model name, data type and energy)
        '''
        import spacepy.plot as splot
        if 'Lm' not in self:
            self.getLm()
        sd = dm.SpaceData()
        sd['Lm'] = self['Lm']
        #filter any bad Lm
        goodidx = sd['Lm'] > 1
        sd['Lm'] = sd['Lm'][goodidx]
        Lm_lim = [2.0, 8.0]
        #TODO: allow user-definition of bins in time and Lm
        varname = self.attrs['varname']
        sd['Epoch'] = dm.dmcopy(
            self['Epoch'])[goodidx]  #TODO: assumes 1 pitch angle, generalize
        sd['1D_dataset'] = self[varname][
            goodidx, ecol]  #TODO: assumes 1 pitch angle, generalize
        spec = splot.spectrogram(sd,
                                 variables=['Epoch', 'Lm', '1D_dataset'],
                                 ylim=Lm_lim)
        if 'zlim' not in kwargs:
            zmax = 10**(int(np.log10(max(sd['1D_dataset']))) + 1)
            idx = np.logical_and(sd['1D_dataset'] > 0, sd['Lm'] > Lm_lim[0])
            idx = np.logical_and(idx, sd['Lm'] <= Lm_lim[1])
            zmin = 10**int(np.log10(min(sd['1D_dataset'][idx])))
            kwargs['zlim'] = [zmin, zmax]
        if 'colorbar_label' not in kwargs:
            flux_units = self[varname].attrs['UNITS']
            kwargs['colorbar_label'] = '{0} ['.format(varname) + re.sub(
                '(\^[\d|-]*)+', _grp2mathmode, flux_units) + ']'
        if 'ylabel' not in kwargs:
            kwargs['ylabel'] = 'L$_M$' + ' ' + '[{0}]'.format(
                self['Lm'].attrs['MODEL'])
        if 'title' not in kwargs:
            kwargs['title'] = '{model_type} {varname}: '.format(**self.attrs) + \
                              '{0} {1}'.format(self['Energy'][ecol], self['Energy'].attrs['UNITS'])
        reset_shrink = splot.mpl.mathtext.SHRINK_FACTOR
        splot.mpl.mathtext.SHRINK_FACTOR = 0.85
        splot.mpl.mathtext.GROW_FACTOR = 1 / 0.85
        ax = spec.plot(cmap='plasma', **kwargs)
        splot.mpl.mathtext.SHRINK_FACTOR = reset_shrink
        splot.mpl.mathtext.GROW_FACTOR = 1 / reset_shrink
        return ax
Exemplo n.º 6
0
    def plotSpectrogram(self, ecol=0, **kwargs):
        '''
        Plot a spectrogram of the flux along the requested orbit, as a function of Lm and time

        Other Parameters
        ----------------
        zlim : list
            2-element list with upper and lower bounds for color scale
        colorbar_label : string
            text to appear next to colorbar (default is 'Flux' plus the units)
        ylabel : string
            text to label y-axis (default is 'Lm' plus the field model name)
        title : string
            text to appear above spectrogram (default is climatology model name, data type and energy)
        '''
        import spacepy.plot as splot
        if 'Lm' not in self:
            self.getLm()
        sd = dm.SpaceData()
        sd['Lm'] = self['Lm']
        #filter any bad Lm
        goodidx = sd['Lm']>1
        sd['Lm'] = sd['Lm'][goodidx]
        Lm_lim = [2.0,8.0]
        #TODO: allow user-definition of bins in time and Lm
        varname = self.attrs['varname']
        sd['Epoch'] = dm.dmcopy(self['Epoch'])[goodidx] #TODO: assumes 1 pitch angle, generalize
        sd['1D_dataset'] = self[varname][goodidx,ecol] #TODO: assumes 1 pitch angle, generalize
        spec = splot.spectrogram(sd, variables=['Epoch', 'Lm', '1D_dataset'], ylim=Lm_lim)
        if 'zlim' not in kwargs:
            zmax = 10**(int(np.log10(max(sd['1D_dataset'])))+1)
            idx = np.logical_and(sd['1D_dataset']>0, sd['Lm']>Lm_lim[0])
            idx = np.logical_and(idx, sd['Lm']<=Lm_lim[1])
            zmin = 10**int(np.log10(min(sd['1D_dataset'][idx])))
            kwargs['zlim'] = [zmin, zmax]
        if 'colorbar_label' not in kwargs:
            flux_units = self[varname].attrs['UNITS']
            kwargs['colorbar_label'] = '{0} ['.format(varname) + re.sub('(\^[\d|-]*)+', _grp2mathmode, flux_units) + ']'
        if 'ylabel' not in kwargs:
            kwargs['ylabel'] = 'L$_M$'+' '+'[{0}]'.format(self['Lm'].attrs['MODEL'])
        if 'title' not in kwargs:
            kwargs['title'] = '{model_type} {varname}: '.format(**self.attrs) + \
                              '{0} {1}'.format(self['Energy'][ecol], self['Energy'].attrs['UNITS'])
        reset_shrink = splot.mpl.mathtext.SHRINK_FACTOR
        splot.mpl.mathtext.SHRINK_FACTOR = 0.85
        splot.mpl.mathtext.GROW_FACTOR =  1/0.85
        ax = spec.plot(cmap='plasma', **kwargs)
        splot.mpl.mathtext.SHRINK_FACTOR = reset_shrink
        splot.mpl.mathtext.GROW_FACTOR =  1/reset_shrink
        return ax
Exemplo n.º 7
0
    except KeyError:
        usestyle = lookdict['default']
    try:
        plt.style.use(usestyle)
    except AttributeError: #plt.style.use not available, old matplotlib?
        dum = matplotlib.rc_params_from_file(usestyle)
        styapply = dict()
        #remove None values as these seem to cause issues...
        for key in dum:
            if dum[key] is not None: styapply[key] = dum[key]
        for key in styapply:
            matplotlib.rcParams[key] = styapply[key]
    matplotlib.rcParams['image.cmap'] = cmap

#save current rcParams before applying spacepy style
oldParams = dmcopy(matplotlib.rcParams)
style()

def revert_style():
    import matplotlib
    for key in oldParams:
        matplotlib.rcParams[key] = oldParams[key]

def dual_half_circle(center=(0,0), radius=1.0,
                     sun_direction='right', ax=None, colors=('w','k'),
                     **kwargs):
    """
    Plot two half circles to a plot with the specified face colors and
    rotation. This is normal to use to denote the sun direction in
    magnetospheric science plots.
Exemplo n.º 8
0
def levelPlot(data,
              var=None,
              time=None,
              levels=(3, 5),
              target=None,
              colors=None,
              **kwargs):
    """
    Draw a step-plot with up to 5 levels following a color cycle (e.g. Kp index "stoplight")

    Parameters
    ----------
    data : array-like, or dict-like
        Data for plotting. If dict-like, the key providing an array-like 
        to plot must be given to var keyword argument.

    Other Parameters
    ----------------
    var    : string
        Name of key in dict-like input that contains data
    time   : array-like or string
        Name of key in dict-like that contains time, or arraylike of datetimes
    levels : array-like, up to 5 levels
        Breaks between levels in data that should be shown as distinct colors
    target : figure or axes
        Target axes or figure window
    colors : array-like
        Colors to use for the color sequence (if insufficient colors, will use as a cycle)
    **kwargs : other keywords
        Other keywords to pass to spacepy.toolbox.binHisto

    Returns
    -------
    binned : tuple
        Tuple of the binned data and bins

    Examples
    --------
    >>> import spacepy.plot as splot
    >>> import spacepy.time as spt
    >>> import spacepy.omni as om
    >>> tt = spt.tickrange('2012/09/28','2012/10/2', 3/24.)
    >>> omni = om.get_omni(tt)
    >>> splot.levelPlot(omni, var='Kp', time='UTC', colors=['seagreen', 'orange', 'crimson'])
    """
    #assume dict-like/key-access, before moving to array-like
    if var is not None:
        try:
            usearr = data[var]
        except KeyError:
            raise KeyError('Key "{1}" not present in data'.format(var))
    else:
        #var is None, so make sure we don't have a dict-like
        if not isinstance(data, Mapping):
            usearr = np.asarray(data)
        else:
            raise TypeError(
                'Data appears to be dict-like without a key being given')
    tflag = False
    if time is not None:
        from scipy.stats import mode
        try:
            times = data[time]
        except (KeyError, ValueError, IndexError):
            times = time
        try:
            times = matplotlib.dates.date2num(times)
            tflag = True
        except AttributeError:
            #the x-data are a non-datetime
            times = np.asarray(time)
        #now add the end-point
        stepsize, dum = mode(np.diff(times), axis=None)
        times = np.hstack([times, times[-1] + stepsize])
    else:
        times = np.asarray(range(0, len(usearr) + 1))
    if not colors:
        if len(levels) <= 3:
            #traffic light colours that are distinct to protanopes and deuteranopes
            colors = ['lime', 'yellow', 'crimson', 'saddlebrown']
        else:
            colors = matplotlib.rcParams['axes.color_cycle']
    else:
        try:
            assert len(colors) > len(levels)
        except AssertionError:
            #cycle the given colors, if not enough are given
            colors = list(colors) * int(1 + len(levels) / len(colors))
    if 'alpha' not in kwargs:
        kwargs['alpha'] = 0.75
    if 'legend' not in kwargs:
        legend = False
    else:
        legend = kwargs['legend']
        del kwargs['legend']
    fig, ax = set_target(target)
    subset = np.asarray(dmcopy(usearr))

    def fill_between_steps(ax, x, y1, **kwargs):
        y2 = np.zeros_like(y1)
        stepsxx = x.repeat(2)[1:-1]
        stepsyy = y1.repeat(2)
        y2 = np.zeros_like(stepsyy)
        ax.fill_between(stepsxx, stepsyy, y2, **kwargs)
        if mpl.__version__ < '1.5.0':
            #pre-v1.5.0, need to manually add an artist for the legend
            p = plt.Rectangle((0, 0), 0, 0, **kwargs)
            ax.add_patch(p)

    #below threshold 1
    idx = 0
    inds = usearr > levels[0]
    subset[inds] = np.nan
    kwargs['label'] = u'≤{0}'.format(levels[idx])
    fill_between_steps(ax, times, subset, color=colors[0], zorder=30, **kwargs)
    #for each of the "between" thresholds
    for idx in range(1, len(levels)):
        subset = np.asarray(dmcopy(usearr))
        inds = np.bitwise_or(usearr <= levels[idx - 1], usearr > levels[idx])
        subset[inds] = np.nan
        kwargs['label'] = u'>{0},≤{1}'.format(levels[idx - 1], levels[idx])
        fill_between_steps(ax,
                           times,
                           subset,
                           color=colors[idx],
                           zorder=30 - (idx * 2),
                           **kwargs)
    #last
    idx += 1
    try:
        inds = usearr <= levels[idx - 1]
        subset = np.asarray(dmcopy(usearr))
        subset[inds] = np.nan
        kwargs['label'] = '>{0}'.format(levels[-1])
        fill_between_steps(ax,
                           times,
                           subset,
                           color=colors[idx],
                           zorder=30 - (idx * 2),
                           **kwargs)
    except:
        pass

    #if required, set x axis to times
    if tflag:
        try:
            applySmartTimeTicks(ax, data[time])
        except (IndexError, KeyError):
            #using data array to index, so should just use time
            applySmartTimeTicks(ax, time)
        ax.grid('off',
                which='minor')  #minor grid usually looks bad on these...

    if legend:
        ncols = len(levels) + 1
        if ncols > 3: ncols = ncols // 2
        ax.legend(loc='upper left', ncol=ncols)

    return ax
Exemplo n.º 9
0
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            dum = mpl.rc_params_from_file(usestyle)
            styapply = dict()
            #remove None values as these seem to cause issues...
            for key in dum:
                if dum[key] is not None: styapply[key] = dum[key]
            for key in styapply:
                mpl.rcParams[key] = styapply[key]
    mpl.rcParams['image.cmap'] = cmap


#save current rcParams before applying spacepy style
oldParams = dict()
for key, val in mpl.rcParams.items():
    oldParams[key] = dmcopy(val)
if config['apply_plot_styles']:
    style()


def revert_style():
    '''Revert plot style settings to those in use prior to importing spacepy.plot
    '''
    import warnings
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        for key in oldParams:
            try:
                mpl.rcParams[key] = oldParams[key]
            except ValueError:
                pass
Exemplo n.º 10
0
    def plotSpectrogram(self, ecol=0, pvars=None, **kwargs):
        '''
        Plot a spectrogram of the flux along the requested orbit, as a function of Lm and time

        Other Parameters
        ----------------
        zlim : list
            2-element list with upper and lower bounds for color scale
        colorbar_label : string
            text to appear next to colorbar (default is 'Flux' plus the units)
        ylabel : string
            text to label y-axis (default is 'Lm' plus the field model name)
        title : string
            text to appear above spectrogram (default is climatology model name, data type and energy)
        pvars : list
            list of plotting variable names in order [Epoch-like (X axis), Flux-like (Z axis), Energy (Index var for Flux-like)]
        ylim : list
            2-element list with upper and lower bounds for y axis
        '''
        import spacepy.plot as splot
        if 'Lm' not in self:
            self.getLm()
        sd = dm.SpaceData()

        if pvars is None:
            varname = self.attrs['varname']
            enname = 'Energy'
            epvar = 'Epoch'
        else:
            epvar = pvars[0]
            varname = pvars[1]
            enname = pvars[2]
        if len(self['Lm']) != len(self[epvar]): #Lm needs interpolating to new timebase
            import matplotlib.dates as mpd
            tmptimetarg, tmptimesrc = mpd.date2num(self[epvar]), mpd.date2num(self['Epoch'])
            sd['Lm'] = dm.dmarray(np.interp(tmptimetarg, tmptimesrc, self['Lm'], left=np.nan, right=np.nan))
        else:
            sd['Lm'] = self['Lm']
        #filter any bad Lm
        goodidx = sd['Lm']>1
        sd['Lm'] = sd['Lm'][goodidx]
        if 'ylim' in kwargs:
            Lm_lim = kwargs['ylim']
            del kwargs['ylim']
        else:
            Lm_lim = [2.0,8.0]
        #TODO: allow user-definition of bins in time and Lm
        sd['Epoch'] = dm.dmcopy(self[epvar])[goodidx] #TODO: assumes 1 pitch angle, generalize
        try:
            sd['1D_dataset'] = self[varname][goodidx,ecol] #TODO: assumes 1 pitch angle, generalize
        except IndexError: #1-D
            sd['1D_dataset'] = self[varname][goodidx]
        bins = [[],[]]
        if 'tbins' not in kwargs:
            bins[0] = spt.tickrange(self[epvar][0], self[epvar][-1], 3/24.).UTC
        else:
            bins[0] = kwargs['tbins']
            del kwargs['tbins']
        if 'Lbins' not in kwargs:
            bins[1] = np.arange(Lm_lim[0], Lm_lim[1], 1./4.)
        else:
            bins[1] = kwargs['Lbins']
            del kwargs['Lbins']
        spec = splot.spectrogram(sd, variables=['Epoch', 'Lm', '1D_dataset'], ylim=Lm_lim, bins=bins)

        if 'zlim' not in kwargs:
            zmax = 10 ** (int(np.log10(max(sd['1D_dataset']))) + 1)
            idx = np.logical_and(sd['1D_dataset'] > 0, sd['Lm'] > Lm_lim[0])
            idx = np.logical_and(idx, sd['Lm'] <= Lm_lim[1])
            zmin = 10 ** int(np.log10(min(sd['1D_dataset'][idx])))
            kwargs['zlim'] = [zmin, zmax]
        if 'colorbar_label' not in kwargs:
            flux_units = self[varname].attrs['UNITS']
            kwargs['colorbar_label'] = '{0} ['.format(varname) + re.sub('(\^[\d|-]*)+', _grp2mathmode, flux_units) + ']'
        if 'ylabel' not in kwargs:
            kwargs['ylabel'] = 'L$_M$' + ' ' + '[{0}]'.format(self['Lm'].attrs['MODEL'])
        if 'title' not in kwargs:
            kwargs['title'] = '{model_type} {varname}: '.format(**self.attrs) + \
                              '{0:5.2f} {1}'.format(self[enname][ecol], self[enname].attrs['UNITS'])
        reset_shrink = splot.mpl.mathtext.SHRINK_FACTOR
        splot.mpl.mathtext.SHRINK_FACTOR = 0.85
        splot.mpl.mathtext.GROW_FACTOR = 1 / 0.85
        ax = spec.plot(cmap='plasma', **kwargs)
        splot.mpl.mathtext.SHRINK_FACTOR = reset_shrink
        splot.mpl.mathtext.GROW_FACTOR = 1 / reset_shrink
        return ax
Exemplo n.º 11
0
def get_omni(ticks, dbase='QDhourly', **kwargs):
    '''
    Returns Qin-Denton OMNI values, interpolated to any time-base from a default hourly resolution

    The update function in toolbox retrieves all available hourly Qin-Denton data, 
    and this function accesses that and interpolates to the given times,
    returning the OMNI values as a SpaceData (dict-like) with
    Kp, Dst, dens, velo, Pdyn, ByIMF, BzIMF, G1, G2, G3, etc.
    (see also http://www.dartmouth.edu/~rdenton/magpar/index.html and
    http://www.agu.org/pubs/crossref/2007/2006SW000296.shtml )

    Parameters
    ==========
    ticks : Ticktock class or array-like of datetimes
        time values for desired output

    dbase : str (optional)
        Select data source, options are 'QDhourly', 'OMNI2', 'Mergedhourly'
        Note - Custom data sources can be specified in the spacepy config file
        as described in the module documentation.

    Returns
    =======
    out : spacepy.datamodel.SpaceData
        containing all Qin-Denton values at times given by ticks

    Examples
    ========
    >>> import spacepy.time as spt
    >>> import spacepy.omni as om
    >>> ticks = spt.Ticktock(['2002-02-02T12:00:00', '2002-02-02T12:10:00'], 'ISO')
    >>> d = om.get_omni(ticks)
    >>> d.tree(levels=1)
    +
    |____ByIMF
    |____Bz1
    |____Bz2
    |____Bz3
    |____Bz4
    |____Bz5
    |____Bz6
    |____BzIMF
    |____DOY
    |____Dst
    |____G1
    |____G2
    |____G3
    |____Hr
    |____Kp
    |____Pdyn
    |____Qbits
    |____RDT
    |____UTC
    |____W1
    |____W2
    |____W3
    |____W4
    |____W5
    |____W6
    |____Year
    |____akp3
    |____dens
    |____ticks
    |____velo


    Notes
    =====
    Note about Qbits: If the status variable is 2, the quantity you are using is fairly well
    determined. If it is 1, the value has some connection to measured values, but is not directly
    measured. These values are still better than just using an average value, but not as good
    as those with the status variable equal to 2. If the status variable is 0, the quantity is
    based on average quantities, and the values listed are no better than an average value. The
    lower the status variable, the less confident you should be in the value.

    '''
    dbase_options = {'QDhourly'    : 1,
                     'OMNI2hourly' : 2,
                     'Mergedhourly': 3,
                     'Test'        : -9,
                     }

    if not isinstance(ticks, spt.Ticktock):
        try:
            ticks = spt.Ticktock(ticks, 'UTC')
        except:
            raise TypeError('get_omni: Input times must be a Ticktock object or a list of datetime objects')

    if not dbase in dbase_options:
        from spacepy import config
        if dbase in config:
            #If a dbase is specified that isn't a default, then it MUST be in the spacepy config
            qdpath = os.path.split(os.path.split(config[dbase])[0])[0]
            if not os.path.isdir(qdpath): raise IOError('Specified dbase ({0}) does not have a valid location ({1})'.format(dbase, config[dbase]))
            days = list(set([tt.date() for tt in ticks.UTC]))
            flist = ['']*len(days)
            fnpath, fnformat = os.path.split(config[dbase])
            for idx, day in enumerate(days):
                dp = fnpath.replace('YYYY', '{0}'.format(day.year))
                df = fnformat.replace('YYYY', '{0}'.format(day.year))
                df = df.replace('MM', '{0:02d}'.format(day.month))
                df = df.replace('DD', '{0:02d}'.format(day.day))
                flist[idx] = os.path.join(dp, df)
            if 'convert' in kwargs:
                convdict = kwargs['convert']
            else:
                convdict = True #set to True as default?
            if 'interp' not in kwargs:
                kwargs['interp'] = True
            data = readJSONheadedASCII(sorted(flist), convert=convdict)
            omniout = SpaceData()

            time_var = [var for var in ['DateTime', 'Time', 'Epoch', 'UTC'] if var in data]
            if time_var:
                use_t_var = time_var[0]
            else:
                #no obvious time variable in input files ... can't continue
                raise ValueError('No clear time variable in file')
            
            if kwargs['interp'] is True:    
                data['RDT'] = spt.Ticktock(data[use_t_var]).RDT
                keylist = sorted(data.keys())
                dum = keylist.pop(keylist.index(use_t_var))
                for key in keylist:
                    try:
                        omniout[key] = dmarray(np.interp(ticks.RDT, data['RDT'], data[key], left=np.NaN, right=np.NaN))
                        omniout[key].attrs = dmcopy(data[key].attrs)
                    except:
                        try:
                            omniout[key] = dmfilled([len(ticks.RDT), data[key].shape[1]], fillval=np.NaN, attrs=dmcopy(data[key].attrs))
                            for col in range(data[key].shape[1]):
                                omniout[key][:,col] = np.interp(ticks.RDT, data['RDT'], data[key][:,col], left=np.NaN, right=np.NaN)
                        except ValueError:
                            print('Failed to interpolate {0} to new time base, skipping variable'.format(key))
                        except IndexError:
                            print('Variable {0} appears to be non-record varying, skipping interpolation'.format(key))
                            omniout[key] = data[key]
                omniout['UTC'] = ticks.UTC 
            else:
                #Trim to specified times
                inds = tOverlapHalf([ticks[0].RDT, ticks[-1].RDT], spt.Ticktock(data['DateTime']).RDT)
                for key in data:
                    if len(inds) == len(data[key]):
                        omniout[key] = data[key][inds]
                    else: #is ancillary data
                        omniout[key] = data[key]
                #TODO: convert to same format as OMNI/QD read (or vice versa)
                omniout['UTC'] = omniout[use_t_var]
            return omniout
        else:
            raise IOError('Specified dbase ({0}) must be specified in spacepy.config'.format(dbase))

    def getattrs(hf, key):
        out = {}
        if hasattr(hf[key],'attrs'):
            for kk, value in hf[key].attrs.items():
                try:
                    out[kk] = value
                except:
                    pass
        return out

    def HrFromDT(indt):
        hour = indt.hour
        minute = indt.minute
        second = indt.second
        musecond = indt.microsecond
        return hour+(minute/60.0)+(second/3600.0)+(musecond/3600.0e3)

    import h5py as h5
    fname, QDkeylist, O2keylist = '', [], []
    omnivals = SpaceData()
    dbase_select = dbase_options[dbase]
    if dbase_select in [1, 3, -9]:
        if dbase_select > 0:
            ldb = 'QDhourly'
            fln = omnifln
        else:
            ldb = 'Test'
            fln = testfln
        with h5.File(fln, 'r') as hfile:
            QDkeylist = [kk for kk in hfile if kk not in ['Qbits', 'UTC']]
            st, en = ticks[0].RDT, ticks[-1].RDT
            ##check that requested requested times are within range of data
            enval, stval = omnirange(dbase=ldb)[1], omnirange(dbase=ldb)[0]
            if (ticks.UTC[0]>enval) or (ticks[-1]<stval):
                raise ValueError('Requested dates are outside data range')
            if (ticks.UTC[-1]>enval) or (ticks[0]<stval):
                print('Warning: Some requested dates are outside data range ({0})'.format(ldb))
            inds = tOverlapHalf([st, en], hfile['RDT'], presort=True) #returns an xrange
            inds = indsFromXrange(inds)
            if inds[0] < 1: inds[0] = 1
            sl_op = slice(inds[0]-1, inds[-1]+2)
    
            fname = ','.join([fname,hfile.filename])
            omnivals.attrs = getattrs(hfile, '/')
            for key in QDkeylist:
                omnivals[key] = dmarray(hfile[key][sl_op]) #TODO: add attrs from h5
                omnivals[key].attrs = getattrs(hfile, key)
            for key in hfile['Qbits']:
                omnivals['Qbits<--{0}'.format(key)] = dmarray(hfile['/Qbits/{0}'.format(key)][sl_op])
                omnivals['Qbits<--{0}'.format(key)].attrs = getattrs(hfile, '/Qbits/{0}'.format(key))
                QDkeylist.append('Qbits<--{0}'.format(key))

    if dbase_options[dbase] == 2 or dbase_options[dbase] == 3:
        ldb = 'OMNI2hourly'
        with h5.File(omni2fln) as hfile:
            O2keylist = [kk for kk in hfile if kk not in ['Epoch','RDT']]
            st, en = ticks[0].RDT, ticks[-1].RDT
            ##check that requested requested times are within range of data
            enval, stval = omnirange(dbase=ldb)[1], omnirange(dbase=ldb)[0]
            if (ticks[0].UTC>enval) or (ticks[-1]<stval):
                raise ValueError('Requested dates are outside data range')
            if (ticks[-1].UTC>enval) or (ticks[0]<stval):
                print('Warning: Some requested dates are outside data range ({0})'.format(ldb))
            inds = tOverlapHalf([st, en], hfile['RDT'], presort=True) #returns an xrange
            inds = indsFromXrange(inds)
            if inds[0] < 1: inds[0] = 1
            sl_op = slice(inds[0]-1, inds[-1]+2)
        
            fname = ','.join([fname,hfile.filename])
            omnivals.attrs = getattrs(hfile, '/') #TODO: This overwrites the previous set on Merged load... Fix!
            omnivals['RDT_OMNI'] = dmarray(hfile['RDT'][sl_op])
            for key in O2keylist:
                omnivals[key] = dmarray(hfile[key][sl_op]) #TODO: add attrs from h5
                omnivals[key].attrs = getattrs(hfile, key)

    if dbase_options[dbase] == 3:
        #prune "merged" SpaceData
        sigmas = [key for key in omnivals if 'sigma' in key]
        for sk in sigmas: del omnivals[sk]
        bees = [key for key in omnivals if re.search('B._', key)]
        for bs in bees: del omnivals[bs]
        aves = [key for key in omnivals if ('_ave' in key) or ('ave_' in key)]
        for av in aves: del omnivals[av]

    omniout = SpaceData(attrs=dmcopy(omnivals.attrs))
    omniout.attrs['filename'] = fname[1:]
    ###print('QDkeys: {0}\n\nO2keys: {1}'.format(QDkeylist, O2keylist))
    for key in sorted(omnivals.keys()):
        if key in O2keylist:
            omniout[key] = dmarray(np.interp(ticks.RDT, omnivals['RDT_OMNI'], omnivals[key], left=np.NaN, right=np.NaN))
            #set metadata -- assume this has been set properly in d/l'd file to match ECT-SOC files
            omniout[key].attrs = dmcopy(omnivals[key].attrs)
        elif key in QDkeylist:
            omniout[key] = dmarray(np.interp(ticks.RDT, omnivals['RDT'], omnivals[key], left=np.NaN, right=np.NaN))
            omniout[key].attrs = dmcopy(omnivals[key].attrs)
        if key == 'G3': #then we have all the Gs
            omniout['G'] = dmarray(np.vstack([omniout['G1'], omniout['G2'], omniout['G3']]).T)
            omniout['G'].attrs = dmcopy(omnivals['G1'].attrs)
            for i in range(1,4): del omniout['G{0}'.format(i)]
        if key == 'W6':
            omniout['W'] = dmarray(np.vstack([omniout['W1'], omniout['W2'], omniout['W3'], omniout['W4'], omniout['W5'], omniout['W6']]).T)
            omniout['W'].attrs = dmcopy(omnivals['W1'].attrs)
            for i in range(1,7): del omniout['W{0}'.format(i)]
        if 'Qbits' in key:
            #Qbits are integer vals, higher is better, so floor to get best representation of interpolated val
            omniout[key] = np.floor(omnivals[key]) 
            omniout[key].attrs = dmcopy(omnivals[key].attrs)
            if 'G3' in key: #then we have all the Gs
                omniout['Qbits<--G'] = dmarray(np.vstack([omniout['Qbits<--G1'], omniout['Qbits<--G2'], omniout['Qbits<--G3']]).T)
                for i in range(1,4): del omniout['Qbits<--G{0}'.format(i)]
            if 'W6' in key:
                omniout['Qbits<--W'] = dmarray(np.vstack([omniout['Qbits<--W1'], omniout['Qbits<--W2'], omniout['Qbits<--W3'], omniout['Qbits<--W4'], omniout['Qbits<--W5'], omniout['Qbits<--W6']]).T)
                for i in range(1,7): del omniout['Qbits<--W{0}'.format(i)]

    omniout['ticks'] = ticks
    omniout['UTC'] = ticks.UTC
    omniout['Hr'] = dmarray([HrFromDT(val) for val in omniout['UTC']])
    omniout['Year'] = dmarray([val.year for val in omniout['UTC']])
    omniout = unflatten(omniout)

    return omniout
Exemplo n.º 12
0
    def plotSpectrogram(self, ecol=0, pvars=None, **kwargs):
        '''
        Plot a spectrogram of the flux along the requested orbit, as a function of Lm and time

        Other Parameters
        ----------------
        zlim : list
            2-element list with upper and lower bounds for color scale
        colorbar_label : string
            text to appear next to colorbar (default is 'Flux' plus the units)
        ylabel : string
            text to label y-axis (default is 'Lm' plus the field model name)
        title : string
            text to appear above spectrogram (default is climatology model name, data type and energy)
        pvars : list
            list of plotting variable names in order [Epoch-like (X axis), Flux-like (Z axis), Energy (Index var for Flux-like)]
        ylim : list
            2-element list with upper and lower bounds for y axis
        '''
        import spacepy.plot as splot
        if 'Lm' not in self:
            self.getLm()
        sd = dm.SpaceData()

        if pvars is None:
            varname = self.attrs['varname']
            enname = 'Energy'
            epvar = 'Epoch'
        else:
            epvar = pvars[0]
            varname = pvars[1]
            enname = pvars[2]
        if len(self['Lm']) != len(
                self[epvar]):  #Lm needs interpolating to new timebase
            import matplotlib.dates as mpd
            tmptimetarg, tmptimesrc = mpd.date2num(self[epvar]), mpd.date2num(
                self['Epoch'])
            sd['Lm'] = dm.dmarray(
                np.interp(tmptimetarg,
                          tmptimesrc,
                          self['Lm'],
                          left=np.nan,
                          right=np.nan))
        else:
            sd['Lm'] = self['Lm']
        #filter any bad Lm
        goodidx = sd['Lm'] > 1
        sd['Lm'] = sd['Lm'][goodidx]
        if 'ylim' in kwargs:
            Lm_lim = kwargs['ylim']
            del kwargs['ylim']
        else:
            Lm_lim = [2.0, 8.0]
        #TODO: allow user-definition of bins in time and Lm
        sd['Epoch'] = dm.dmcopy(
            self[epvar])[goodidx]  #TODO: assumes 1 pitch angle, generalize
        try:
            sd['1D_dataset'] = self[varname][
                goodidx, ecol]  #TODO: assumes 1 pitch angle, generalize
        except IndexError:  #1-D
            sd['1D_dataset'] = self[varname][goodidx]
        bins = [[], []]
        if 'tbins' not in kwargs:
            bins[0] = spt.tickrange(self[epvar][0], self[epvar][-1],
                                    3 / 24.).UTC
        else:
            bins[0] = kwargs['tbins']
            del kwargs['tbins']
        if 'Lbins' not in kwargs:
            bins[1] = np.arange(Lm_lim[0], Lm_lim[1], 1. / 4.)
        else:
            bins[1] = kwargs['Lbins']
            del kwargs['Lbins']
        spec = splot.spectrogram(sd,
                                 variables=['Epoch', 'Lm', '1D_dataset'],
                                 ylim=Lm_lim,
                                 bins=bins)

        if 'zlim' not in kwargs:
            zmax = 10**(int(np.log10(max(sd['1D_dataset']))) + 1)
            idx = np.logical_and(sd['1D_dataset'] > 0, sd['Lm'] > Lm_lim[0])
            idx = np.logical_and(idx, sd['Lm'] <= Lm_lim[1])
            zmin = 10**int(np.log10(min(sd['1D_dataset'][idx])))
            kwargs['zlim'] = [zmin, zmax]
        if 'colorbar_label' not in kwargs:
            flux_units = self[varname].attrs['UNITS']
            kwargs['colorbar_label'] = '{0} ['.format(varname) + re.sub(
                '(\^[\d|-]*)+', _grp2mathmode, flux_units) + ']'
        if 'ylabel' not in kwargs:
            kwargs['ylabel'] = 'L$_M$' + ' ' + '[{0}]'.format(
                self['Lm'].attrs['MODEL'])
        if 'title' not in kwargs:
            kwargs['title'] = '{model_type} {varname}: '.format(**self.attrs) + \
                              '{0:5.2f} {1}'.format(self[enname][ecol], self[enname].attrs['UNITS'])
        reset_shrink = splot.mpl.mathtext.SHRINK_FACTOR
        splot.mpl.mathtext.SHRINK_FACTOR = 0.85
        splot.mpl.mathtext.GROW_FACTOR = 1 / 0.85
        ax = spec.plot(cmap='plasma', **kwargs)
        splot.mpl.mathtext.SHRINK_FACTOR = reset_shrink
        splot.mpl.mathtext.GROW_FACTOR = 1 / reset_shrink
        return ax
Exemplo n.º 13
0
def singleRunEwithEnsemble(runname,
                           ensname,
                           searchpatt='1_geoe_*.csv',
                           outdir='E_maps',
                           eVar='Estd'):
    '''
    Optional Parameters
    ------------------

    eVar : string
        Estd or Emag
    '''
    import cartopy.feature as cfea
    import cartopy.feature.nightshade as night
    from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
    useProject = ccrs.PlateCarree  #Projection to use for plotting, e.g., ccrs.AlbersEqualArea
    plotprojection = useProject(central_longitude=-90.0)

    pdata = {
        'plotvar': 'Emag',
        'plotvec': 'E',
        'dataprojection': ccrs.PlateCarree(),
        'plotprojection': plotprojection,
    }
    doQuiver = False  #may want to change this in future, so I'm leaving the code in
    #add features from ArcGIS shapefile?
    pdata['shapes'] = None

    rundir = runname[4:]
    globterm = os.path.join(runname, searchpatt)
    allfiles = sorted(glob.glob(globterm))

    #downselect files to use
    #TODO: remove hardcoded downselect of timerange
    allfiles = [
        fn for fn in allfiles
        if (int(re.search('\d{8}-(\d{6})', fn).groups()[0]) >= 81500) and (
            int(re.search('\d{8}-(\d{6})', fn).groups()[0]) <= 84500)
    ]
    infiles = [
        fn for fn in allfiles if re.search('\d{8}-(\d{6})', fn).groups()[0]
        [-2:] in ['00', '15', '30', '45']
    ]

    #get ensemble member directories
    members = [en for en in glob.glob(ensname) if en != runname]
    print('Using {0} ensemble members'.format(len(members)))
    for mm in members:
        print('{0}'.format(mm))

    #loop over files and plot map of |E|
    for fname in infiles:
        edata = gmdtools.readEcsv(fname)
        #find matching file for each ens. member
        filepart = os.path.split(fname)[-1]
        collect = [edata]
        for mdir in members:
            fname = os.path.join(mdir, filepart)
            try:
                collect.append(gmdtools.readEcsv(fname))
            except:
                print('Hit an issue with {0}'.format(fname))
                continue
        #collect Emag and calculate stadard deviation across ensemble members at each gridpoint
        combinedMag = np.dstack(cc['Emag'] for cc in collect)
        combinedEast = np.dstack(cc['Ee'] for cc in collect)
        combinedNorth = np.dstack(cc['En'] for cc in collect)
        ensmean = dm.dmcopy(edata)
        ensmean['Ee_raw'] = np.stack(cc['Ee_raw']
                                     for cc in collect).mean(axis=0)
        ensmean['En_raw'] = np.stack(cc['En_raw']
                                     for cc in collect).mean(axis=0)
        ensmean['Estd'] = combinedMag.std(axis=-1)
        ensmean['Ee_ensMean'] = combinedEast.mean(axis=-1)
        ensmean['En_ensMean'] = combinedNorth.mean(axis=-1)
        ensmean['Emag'] = np.sqrt(ensmean['Ee_ensMean']**2 +
                                  ensmean['En_ensMean']**2)
        #
        print('Working on timestep {0}. # members = {1}'.format(
            edata.attrs['time'], len(collect)))
        #
        fig = plt.figure(figsize=(10, 5.5))
        #set up first map panel (reference run)
        pdata['plotvar'] = 'Emag'
        sstyle = {'color': 'black'}
        qstyle = {
            'color': 'darkgrey',
            'pivot': 'mid',
            'alpha': 0.6,
            'scale': 1,
            'scale_units': 'xy'
        }
        ax = plt.subplot(2, 1, 1, projection=pdata['plotprojection'])
        ax.coastlines(color='darkgrey')
        ax.add_feature(cfea.BORDERS, linestyle=':', edgecolor='darkgrey')
        ax.add_feature(night.Nightshade(edata.attrs['time'], alpha=0.3))
        lon_formatter = LongitudeFormatter(number_format='.1f',
                                           degree_symbol='',
                                           dateline_direction_label=True)
        lat_formatter = LatitudeFormatter(number_format='.1f',
                                          degree_symbol='')
        ax.xaxis.set_major_formatter(lon_formatter)
        ax.yaxis.set_major_formatter(lat_formatter)
        ax.set_extent([-165.0, 45.0, 30.0, 80.0], crs=ccrs.PlateCarree())
        ax = gmdtools.plotFilledContours(edata, pdata, addTo=ax)
        ax = gmdtools.plotVectors(edata,
                                  pdata,
                                  addTo=ax,
                                  maxVec=2.75,
                                  sstyle=sstyle)
        ax.set_title('{0}'.format(edata.attrs['time'].isoformat()))
        #set up second map panel (use requested (Emag_ens or Estd)
        ax2 = plt.subplot(2, 1, 2, projection=pdata['plotprojection'])
        ax2.coastlines(color='darkgrey')
        ax2.add_feature(cfea.BORDERS, linestyle=':', edgecolor='darkgrey')
        ax2.add_feature(night.Nightshade(edata.attrs['time'], alpha=0.3))
        ax2.xaxis.set_major_formatter(lon_formatter)
        ax2.yaxis.set_major_formatter(lat_formatter)
        ax2.set_extent([-165.0, 45.0, 30.0, 80.0], crs=ccrs.PlateCarree())
        pdata['plotvar'] = eVar
        ax2 = gmdtools.plotFilledContours(ensmean, pdata, addTo=ax2)
        #plot all ensemble members as light grey quivers
        if doQuiver:
            for ense in collect:
                ax2 = gmdtools.plotVectors(ense,
                                           pdata,
                                           addTo=ax2,
                                           quiver=True,
                                           qstyle=qstyle)
            qstyle['alpha'] = 1
            qstyle['color'] = 'black'
        #then plot ensemble mean with streamlines...
        sstyle['color'] = 'darkblue'
        #collect Emag and calculate [mean/standard deviation] across ensemble members at each gridpoint
        ax2 = gmdtools.plotVectors(ensmean,
                                   pdata,
                                   addTo=ax2,
                                   quiver=False,
                                   maxVec=2.75,
                                   qstyle=qstyle,
                                   sstyle=sstyle)
        for aa in [ax, ax2]:
            aa.set_xticks([-150, -120, -90, -60, -30, 0, 30],
                          crs=ccrs.PlateCarree())
            aa.set_yticks([35, 45, 55, 65, 75], crs=ccrs.PlateCarree())

        #now annotate panels with "Reference" and "Ensemble" (put text at Lisbon, Portugal)
        anchor = (-9.195, 38.744)  #Lisbon #(-7.9304, 37.0194) #Faro, Port.
        #ensText = 'Ensemble Mean + $\sigma(|E|)$' if eVar=='Estd' else 'Ensemble Mean'
        ensText = 'Ensemble Mean'
        ax.text(anchor[0],
                anchor[1],
                'Unperturbed',
                verticalalignment='top',
                weight='semibold',
                bbox=dict(facecolor='white', alpha=0.3, edgecolor='None'),
                transform=ccrs.PlateCarree())
        ax2.text(anchor[0],
                 anchor[1],
                 ensText + '\n(N={0})'.format(len(collect)),
                 verticalalignment='top',
                 weight='semibold',
                 color='darkblue',
                 bbox=dict(facecolor='white', alpha=0.3, edgecolor='None'),
                 transform=ccrs.PlateCarree())

        #windows can't handle colons in filenames...
        isotime = edata.attrs['time'].isoformat()
        plt.tight_layout()
        plt.savefig(os.path.join(
            outdir, r'{0}_{1}.png'.format(pdata['plotvar'],
                                          isotime.replace(':', ''))),
                    dpi=300)
        plt.close()
    gmdtools.makeSymlinks(outdir, kind='E')
Exemplo n.º 14
0
                    error_series[run_num,
                                 it * blocksize:it * blocksize + blocksize,
                                 vidx] = savedata[var][bidx:bidx + blocksize]
            elif Ntimes - it * blocksize > 0:
                room = len(error_series[run_num, it * blocksize:, vidx])
                error_series[run_num, it * blocksize:,
                             vidx] = savedata[var][bidx:bidx + room]
            else:
                pass
        #modify SWMF ImfInput and write new file
        outfilename = '.'.join([
            '_'.join([infilename.split('.')[0], '{0:03d}'.format(run_num)]),
            'dat'
        ])
        if generateInputs:
            surrogateIMF = dm.dmcopy(eventIMF)
            for vidx, var in enumerate(varlist):

                surrogateIMF[map_dict[var]] += error_series[run_num, :Ntimes,
                                                            vidx]
            #then write to file
            surrogateIMF.write(outfilename)

    #save error series if req'd
    if saveErrors:
        out = dm.SpaceData()
        out['errors'] = dm.dmarray(error_series)
        out['errors'].attrs['DEPEND_0'] = 'EnsembleNumber'
        out['errors'].attrs['DEPEND_1'] = 'Timestep'
        out['errors'].attrs['DEPEND_2'] = 'Variable'
        out.toHDF5('MBB_errors.h5'.format(var))
Exemplo n.º 15
0
def levelPlot(data, var=None, time=None, levels=(3, 5), target=None, colors=None, **kwargs):
    """
    Draw a step-plot with up to 5 levels following a color cycle (e.g. Kp index "stoplight")

    Parameters
    ----------
    data : array-like, or dict-like
        Data for plotting. If dict-like, the key providing an array-like 
        to plot must be given to var keyword argument.

    Other Parameters
    ----------------
    var    : string
        Name of key in dict-like input that contains data
    time   : array-like or string
        Name of key in dict-like that contains time, or arraylike of datetimes
    levels : array-like, up to 5 levels
        Breaks between levels in data that should be shown as distinct colors
    target : figure or axes
        Target axes or figure window
    colors : array-like
        Colors to use for the color sequence (if insufficient colors, will use as a cycle)
    **kwargs : other keywords
        Other keywords to pass to spacepy.toolbox.binHisto

    Returns
    -------
    binned : tuple
        Tuple of the binned data and bins

    Examples
    --------
    >>> import spacepy.plot as splot
    >>> import spacepy.time as spt
    >>> import spacepy.omni as om
    >>> tt = spt.tickrange('2012/09/28','2012/10/2', 3/24.)
    >>> omni = om.get_omni(tt)
    >>> splot.levelPlot(omni, var='Kp', time='UTC', colors=['seagreen', 'orange', 'crimson'])
    """
    #assume dict-like/key-access, before moving to array-like
    if var is not None:
        try:
            usearr = data[var]
        except KeyError:
            raise KeyError('Key "{1}" not present in data'.format(var))
    else:
        #var is None, so make sure we don't have a dict-like
        import collections
        if not isinstance(data, collections.Mapping):
            usearr = np.asarray(data)
        else:
            raise TypeError('Data appears to be dict-like without a key being given')
    tflag = False
    if time is not None:
        from scipy.stats import mode
        try:
            times = data[time]
        except (KeyError, ValueError, IndexError):
            times = time
        try:
            times = matplotlib.dates.date2num(times)
            tflag = True
        except AttributeError:
            #the x-data are a non-datetime
            times = np.asarray(time)
        #now add the end-point
        stepsize, dum = mode(np.diff(times), axis=None)
        times = np.hstack([times, times[-1]+stepsize])
    else:
        times = np.asarray(range(0, len(usearr)+1))
    if not colors:
        if len(levels)<=3:
            #traffic light colours that are distinct to protanopes and deuteranopes
            colors = ['lime', 'yellow', 'crimson', 'saddlebrown']
        else:
            colors = matplotlib.rcParams['axes.color_cycle']
    else:
        try:
            assert len(colors) > len(levels)
        except AssertionError:
            #cycle the given colors, if not enough are given
            colors = list(colors)*int(1+len(levels)/len(colors))
    if 'alpha' not in kwargs:
        kwargs['alpha']=0.75
    if 'legend' not in kwargs:
        legend = False
    else:
        legend = kwargs['legend']
        del kwargs['legend']
    fig, ax = set_target(target)
    subset = np.asarray(dmcopy(usearr))

    def fill_between_steps(ax, x, y1, **kwargs):
        y2 = np.zeros_like(y1)
        stepsxx = x.repeat(2)[1:-1]
        stepsyy = y1.repeat(2)
        y2 = np.zeros_like(stepsyy)
        ax.fill_between(stepsxx, stepsyy, y2, **kwargs)
        p = plt.Rectangle((0, 0), 0, 0, **kwargs)
        ax.add_patch(p)
    
    #below threshold 1
    idx = 0
    inds = usearr>levels[0]
    subset[inds] = np.nan
    kwargs['label'] = '<{0}'.format(levels[idx])
    fill_between_steps(ax, times, subset, color=colors[0], zorder=30, **kwargs)
    #for each of the "between" thresholds
    for idx in range(1,len(levels)):
        subset = np.asarray(dmcopy(usearr))
        inds = np.bitwise_or(usearr<=levels[idx-1], usearr>levels[idx])
        subset[inds] = np.nan
        kwargs['label'] = '{0}-{1}'.format(levels[idx-1], levels[idx])
        fill_between_steps(ax, times, subset, color=colors[idx], zorder=30-(idx*2), **kwargs)
    #last
    idx += 1
    try:
        inds = usearr<=levels[idx-1]
        subset = np.asarray(dmcopy(usearr))
        subset[inds] = np.nan
        kwargs['label'] = '>{0}'.format(levels[-1])
        fill_between_steps(ax, times, subset, color=colors[idx], zorder=30-(idx*2), **kwargs)
    except:
        pass

    #if required, set x axis to times
    if tflag:
        try:
            applySmartTimeTicks(ax, data[time])
        except (IndexError, KeyError):
            #using data array to index, so should just use time
            applySmartTimeTicks(ax, time)
        ax.grid('off', which='minor') #minor grid usually looks bad on these...
    
    if legend:
        ncols = len(levels)+1
        if ncols > 3: ncols = ncols//2
        ax.legend(loc='upper left', ncol=ncols)

    return ax
Exemplo n.º 16
0
        import warnings
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            dum = mpl.rc_params_from_file(usestyle)
            styapply = dict()
            #remove None values as these seem to cause issues...
            for key in dum:
                if dum[key] is not None: styapply[key] = dum[key]
            for key in styapply:
                mpl.rcParams[key] = styapply[key]
    mpl.rcParams['image.cmap'] = cmap

#save current rcParams before applying spacepy style
oldParams = dict()
for key, val in mpl.rcParams.items():
        oldParams[key] = dmcopy(val)
if config['apply_plot_styles']:
    style()

def revert_style():
    '''Revert plot style settings to those in use prior to importing spacepy.plot
    '''
    import warnings
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        for key in oldParams:
            try:
                mpl.rcParams[key] = oldParams[key]
            except ValueError:
                pass
Exemplo n.º 17
0
def get_omni(ticks, dbase='QDhourly', **kwargs):
    '''
    Returns Qin-Denton OMNI values, interpolated to any time-base from a default hourly resolution

    The update function in toolbox retrieves all available hourly Qin-Denton data, 
    and this function accesses that and interpolates to the given times,
    returning the OMNI values as a SpaceData (dict-like) with
    Kp, Dst, dens, velo, Pdyn, ByIMF, BzIMF, G1, G2, G3, etc.
    (see also http://www.dartmouth.edu/~rdenton/magpar/index.html and
    http://www.agu.org/pubs/crossref/2007/2006SW000296.shtml )

    Parameters
    ==========
    ticks : Ticktock class or array-like of datetimes
        time values for desired output

    dbase : str (optional)
        Select data source, options are 'QDhourly', 'OMNI2hourly', 'Mergedhourly'
        Note - Custom data sources can be specified in the spacepy config file
        as described in the module documentation.

    Returns
    =======
    out : spacepy.datamodel.SpaceData
        containing all Qin-Denton values at times given by ticks

    Examples
    ========
    >>> import spacepy.time as spt
    >>> import spacepy.omni as om
    >>> ticks = spt.Ticktock(['2002-02-02T12:00:00', '2002-02-02T12:10:00'], 'ISO')
    >>> d = om.get_omni(ticks)
    >>> d.tree(levels=1)
    +
    |____ByIMF
    |____Bz1
    |____Bz2
    |____Bz3
    |____Bz4
    |____Bz5
    |____Bz6
    |____BzIMF
    |____DOY
    |____Dst
    |____G1
    |____G2
    |____G3
    |____Hr
    |____Kp
    |____Pdyn
    |____Qbits
    |____RDT
    |____UTC
    |____W1
    |____W2
    |____W3
    |____W4
    |____W5
    |____W6
    |____Year
    |____akp3
    |____dens
    |____ticks
    |____velo


    Notes
    =====
    Note about Qbits: If the status variable is 2, the quantity you are using is fairly well
    determined. If it is 1, the value has some connection to measured values, but is not directly
    measured. These values are still better than just using an average value, but not as good
    as those with the status variable equal to 2. If the status variable is 0, the quantity is
    based on average quantities, and the values listed are no better than an average value. The
    lower the status variable, the less confident you should be in the value.

    '''
    dbase_options = {
        'QDhourly': 1,
        'OMNI2hourly': 2,
        'Mergedhourly': 3,
        'Test': -9,
    }

    if not isinstance(ticks, spt.Ticktock):
        try:
            ticks = spt.Ticktock(ticks, 'UTC')
        except:
            raise TypeError(
                'get_omni: Input times must be a Ticktock object or a list of datetime objects'
            )

    if not dbase in dbase_options:
        from spacepy import config
        if dbase in config:
            #If a dbase is specified that isn't a default, then it MUST be in the spacepy config
            qdpath = os.path.split(os.path.split(config[dbase])[0])[0]
            if not os.path.isdir(qdpath):
                raise IOError(
                    'Specified dbase ({0}) does not have a valid location ({1})'
                    .format(dbase, config[dbase]))
            days = list(set([tt.date() for tt in ticks.UTC]))
            flist = [''] * len(days)
            fnpath, fnformat = os.path.split(config[dbase])
            for idx, day in enumerate(days):
                dp = fnpath.replace('YYYY', '{0}'.format(day.year))
                df = fnformat.replace('YYYY', '{0}'.format(day.year))
                df = df.replace('MM', '{0:02d}'.format(day.month))
                df = df.replace('DD', '{0:02d}'.format(day.day))
                flist[idx] = os.path.join(dp, df)
            if 'convert' in kwargs:
                convdict = kwargs['convert']
            else:
                convdict = True  #set to True as default?
            if 'interp' not in kwargs:
                kwargs['interp'] = True
            data = readJSONheadedASCII(sorted(flist), convert=convdict)
            omniout = SpaceData()

            time_var = [
                var for var in ['DateTime', 'Time', 'Epoch', 'UTC']
                if var in data
            ]
            if time_var:
                use_t_var = time_var[0]
            else:
                #no obvious time variable in input files ... can't continue
                raise ValueError('No clear time variable in file')

            if kwargs['interp'] is True:
                data['RDT'] = spt.Ticktock(data[use_t_var]).RDT
                keylist = sorted(data.keys())
                dum = keylist.pop(keylist.index(use_t_var))
                for key in keylist:
                    try:
                        omniout[key] = dmarray(
                            np.interp(ticks.RDT,
                                      data['RDT'],
                                      data[key],
                                      left=np.NaN,
                                      right=np.NaN))
                        omniout[key].attrs = dmcopy(data[key].attrs)
                    except:
                        try:
                            omniout[key] = dmfilled(
                                [len(ticks.RDT), data[key].shape[1]],
                                fillval=np.NaN,
                                attrs=dmcopy(data[key].attrs))
                            for col in range(data[key].shape[1]):
                                omniout[key][:,
                                             col] = np.interp(ticks.RDT,
                                                              data['RDT'],
                                                              data[key][:,
                                                                        col],
                                                              left=np.NaN,
                                                              right=np.NaN)
                        except ValueError:
                            print(
                                'Failed to interpolate {0} to new time base, skipping variable'
                                .format(key))
                        except IndexError:
                            print(
                                'Variable {0} appears to be non-record varying, skipping interpolation'
                                .format(key))
                            omniout[key] = data[key]
                omniout['UTC'] = ticks.UTC
            else:
                #Trim to specified times
                inds = tOverlapHalf([ticks[0].RDT, ticks[-1].RDT],
                                    spt.Ticktock(data['DateTime']).RDT)
                for key in data:
                    if len(inds) == len(data[key]):
                        omniout[key] = data[key][inds]
                    else:  #is ancillary data
                        omniout[key] = data[key]
                #TODO: convert to same format as OMNI/QD read (or vice versa)
                omniout['UTC'] = omniout[use_t_var]
            return omniout
        else:
            raise IOError(
                'Specified dbase ({0}) must be specified in spacepy.config'.
                format(dbase))

    def getattrs(hf, key):
        out = {}
        if hasattr(hf[key], 'attrs'):
            for kk, value in hf[key].attrs.items():
                try:
                    out[kk] = value
                except:
                    pass
        return out

    def HrFromDT(indt):
        hour = indt.hour
        minute = indt.minute
        second = indt.second
        musecond = indt.microsecond
        return hour + (minute / 60.0) + (second / 3600.0) + (musecond /
                                                             3600.0e3)

    import h5py as h5
    fname, QDkeylist, O2keylist = '', [], []
    omnivals = SpaceData()
    dbase_select = dbase_options[dbase]
    if dbase_select in [1, 3, -9]:
        if dbase_select > 0:
            ldb = 'QDhourly'
            fln = omnifln
        else:
            ldb = 'Test'
            fln = testfln
        with h5.File(fln, 'r') as hfile:
            QDkeylist = [kk for kk in hfile if kk not in ['Qbits', 'UTC']]
            st, en = ticks[0].RDT, ticks[-1].RDT
            ##check that requested requested times are within range of data
            enval, stval = omnirange(dbase=ldb)[1], omnirange(dbase=ldb)[0]
            if (ticks.UTC[0] > enval) or (ticks[-1] < stval):
                raise ValueError('Requested dates are outside data range')
            if (ticks.UTC[-1] > enval) or (ticks[0] < stval):
                print(
                    'Warning: Some requested dates are outside data range ({0})'
                    .format(ldb))
            inds = tOverlapHalf([st, en], hfile['RDT'],
                                presort=True)  #returns an xrange
            inds = indsFromXrange(inds)
            if inds[0] < 1: inds[0] = 1
            sl_op = slice(inds[0] - 1, inds[-1] + 2)

            fname = ','.join([fname, hfile.filename])
            omnivals.attrs = getattrs(hfile, '/')
            for key in QDkeylist:
                omnivals[key] = dmarray(
                    hfile[key][sl_op])  #TODO: add attrs from h5
                omnivals[key].attrs = getattrs(hfile, key)
            for key in hfile['Qbits']:
                omnivals['Qbits<--{0}'.format(key)] = dmarray(
                    hfile['/Qbits/{0}'.format(key)][sl_op])
                omnivals['Qbits<--{0}'.format(key)].attrs = getattrs(
                    hfile, '/Qbits/{0}'.format(key))
                QDkeylist.append('Qbits<--{0}'.format(key))

    if dbase_options[dbase] == 2 or dbase_options[dbase] == 3:
        ldb = 'OMNI2hourly'
        with h5.File(omni2fln) as hfile:
            O2keylist = [kk for kk in hfile if kk not in ['Epoch', 'RDT']]
            st, en = ticks[0].RDT, ticks[-1].RDT
            ##check that requested requested times are within range of data
            enval, stval = omnirange(dbase=ldb)[1], omnirange(dbase=ldb)[0]
            if (ticks[0].UTC > enval) or (ticks[-1] < stval):
                raise ValueError('Requested dates are outside data range')
            if (ticks[-1].UTC > enval) or (ticks[0] < stval):
                print(
                    'Warning: Some requested dates are outside data range ({0})'
                    .format(ldb))
            inds = tOverlapHalf([st, en], hfile['RDT'],
                                presort=True)  #returns an xrange
            inds = indsFromXrange(inds)
            if inds[0] < 1: inds[0] = 1
            sl_op = slice(inds[0] - 1, inds[-1] + 2)

            fname = ','.join([fname, hfile.filename])
            omnivals.attrs = getattrs(
                hfile, '/'
            )  #TODO: This overwrites the previous set on Merged load... Fix!
            omnivals['RDT_OMNI'] = dmarray(hfile['RDT'][sl_op])
            for key in O2keylist:
                omnivals[key] = dmarray(
                    hfile[key][sl_op])  #TODO: add attrs from h5
                omnivals[key].attrs = getattrs(hfile, key)

    if dbase_options[dbase] == 3:
        #prune "merged" SpaceData
        sigmas = [key for key in omnivals if 'sigma' in key]
        for sk in sigmas:
            del omnivals[sk]
        bees = [key for key in omnivals if re.search('B._', key)]
        for bs in bees:
            del omnivals[bs]
        aves = [key for key in omnivals if ('_ave' in key) or ('ave_' in key)]
        for av in aves:
            del omnivals[av]

    omniout = SpaceData(attrs=dmcopy(omnivals.attrs))
    omniout.attrs['filename'] = fname[1:]
    ###print('QDkeys: {0}\n\nO2keys: {1}'.format(QDkeylist, O2keylist))
    for key in sorted(omnivals.keys()):
        if key in O2keylist:
            omniout[key] = dmarray(
                np.interp(ticks.RDT,
                          omnivals['RDT_OMNI'],
                          omnivals[key],
                          left=np.NaN,
                          right=np.NaN))
            #set metadata -- assume this has been set properly in d/l'd file to match ECT-SOC files
            omniout[key].attrs = dmcopy(omnivals[key].attrs)
        elif key in QDkeylist:
            omniout[key] = dmarray(
                np.interp(ticks.RDT,
                          omnivals['RDT'],
                          omnivals[key],
                          left=np.NaN,
                          right=np.NaN))
            omniout[key].attrs = dmcopy(omnivals[key].attrs)
        if key == 'G3':  #then we have all the Gs
            omniout['G'] = dmarray(
                np.vstack([omniout['G1'], omniout['G2'], omniout['G3']]).T)
            omniout['G'].attrs = dmcopy(omnivals['G1'].attrs)
            for i in range(1, 4):
                del omniout['G{0}'.format(i)]
        if key == 'W6':
            omniout['W'] = dmarray(
                np.vstack([
                    omniout['W1'], omniout['W2'], omniout['W3'], omniout['W4'],
                    omniout['W5'], omniout['W6']
                ]).T)
            omniout['W'].attrs = dmcopy(omnivals['W1'].attrs)
            for i in range(1, 7):
                del omniout['W{0}'.format(i)]
        if 'Qbits' in key:
            #Qbits are integer vals, higher is better, so floor to get best representation of interpolated val
            omniout[key] = np.floor(omnivals[key])
            omniout[key].attrs = dmcopy(omnivals[key].attrs)
            if 'G3' in key:  #then we have all the Gs
                omniout['Qbits<--G'] = dmarray(
                    np.vstack([
                        omniout['Qbits<--G1'], omniout['Qbits<--G2'],
                        omniout['Qbits<--G3']
                    ]).T)
                for i in range(1, 4):
                    del omniout['Qbits<--G{0}'.format(i)]
            if 'W6' in key:
                omniout['Qbits<--W'] = dmarray(
                    np.vstack([
                        omniout['Qbits<--W1'], omniout['Qbits<--W2'],
                        omniout['Qbits<--W3'], omniout['Qbits<--W4'],
                        omniout['Qbits<--W5'], omniout['Qbits<--W6']
                    ]).T)
                for i in range(1, 7):
                    del omniout['Qbits<--W{0}'.format(i)]

    omniout['ticks'] = ticks
    omniout['UTC'] = ticks.UTC
    omniout['Hr'] = dmarray([HrFromDT(val) for val in omniout['UTC']])
    omniout['Year'] = dmarray([val.year for val in omniout['UTC']])
    omniout = unflatten(omniout)

    return omniout