Beispiel #1
0
    def add_histplot(self, target=False, loc=111, label='Kyoto $K_{p}$', 
                     time_range=None, **kwargs):
        '''
        Make a quick histogram-style plot of the Kp data.

        Returns
        =======
        fig : matplotlib figure object
        ax  : matplotlib axes 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).

        label : string
            The label applied to the line when a legend is added to the axes.
            Defaults to 'Kyoto $K_{p}$'.

        time_range : tuple of datetimes
            The time range to plot.  Only the first and last values in the
            tuple (or list) are used if the number of elements is greater than
            two.  Defaults to **None**, meaning that the full time range
            available is used.

        Extra keyword arguments are passed to :function:`matplotlib.pyplot.plot`
        to customize the line style.

        Examples
        ========
        >>> import spacepy.pybats.kyoto as kt
        >>> kp = kt.fetch('kp', (1981, 11), (1981, 11)
        >>> kp.add_histplot(lw=2.0, lc='r', label='Example KP')

        '''

        import matplotlib.pyplot as plt

        # Shortcuts for the lazy.
        bstart=self['binstart']
        bstop =self['binstop']
        npts  =self.attrs['npts']

        # Reformulate time to get histogram-type look.
        newtime=np.zeros(npts*24, dtype=object)
        newkp  =np.zeros(npts*24)

        # Set time range.
        if not time_range:
            time_range = newtime
        
        for i in range(npts*8):
            newtime[3*i  ] = bstart[i]
            newtime[3*i+1] = self['time'][i]
            newtime[3*i+2] = bstop[i]
            newkp[3*i:3*i+3] = self['kp'][i], self['kp'][i], self['kp'][i]
    

        fig, ax = set_target(target, figsize=(10,4), loc=loc)
        line=ax.plot(newtime, newkp, label=label, **kwargs)
        applySmartTimeTicks(ax, time_range, dolabel=True)

        return fig, ax
Beispiel #2
0
    def add_slice(self,
                  var,
                  alt,
                  time,
                  nlev=31,
                  zlim=None,
                  target=None,
                  loc=111,
                  title=None,
                  latoffset=1.05,
                  rlim=50.,
                  add_cbar=True,
                  clabel=None,
                  show_pts=False,
                  show_alt=True,
                  dolog=False,
                  lats=[75., 60.],
                  colats=None,
                  figsize=(8.34, 7),
                  *args,
                  **kwargs):
        '''
        Create a plot of variable *var* at altitude *alt*.

        If kwarg **target** is None (default), a new figure is 
        generated from scratch.  If target is a matplotlib Figure
        object, a new axis is created to fill that figure at subplot
        location **loc**.  If **target** is a matplotlib Axes object, 
        the plot is placed into that axis.
        '''

        import matplotlib.pyplot as plt
        from matplotlib.colors import (LogNorm, Normalize)
        from matplotlib.patches import Circle
        from matplotlib.ticker import (LogLocator, LogFormatter,
                                       LogFormatterMathtext, MultipleLocator)

        # Grab the slice of data that we want:
        if type(time) == type(self['time'][0]):
            if time not in self['time']:
                raise ValueError('Time not in object')
            time = np.arange(self.attrs['nTime'])[self['time'] == time][0]

        fig, ax = set_target(target, loc=loc, figsize=figsize)
        ax.set_aspect('equal')

        # Create values in Cartesian plane.
        self._get_cartXY()

        # Grab values from correct time/location.
        x = self._xLat[:, time, alt]
        y = self._yLat[:, time, alt]
        value = self[var][:, time, alt]

        # Get max/min if none given.
        if zlim == None:
            zlim = [0, 0]
            zlim[0] = value.min()
            zlim[1] = value.max()
            if dolog and zlim[0] <= 0:
                zlim[0] = np.min([0.0001, zlim[1] / 1000.0])

        # Create levels and set norm based on dolog.
        if dolog:
            levs = np.power(
                10, np.linspace(np.log10(zlim[0]), np.log10(zlim[1]), nlev))
            z = np.where(value > zlim[0], value, 1.01 * zlim[0])
            norm = LogNorm()
            ticks = LogLocator()
            fmt = LogFormatterMathtext()
        else:
            levs = np.linspace(zlim[0], zlim[1], nlev)
            z = value
            norm = None
            ticks = None
            fmt = None

        # Create contour plot.
        cont = ax.tricontourf(x, y, z, levs, *args, norm=norm, **kwargs)

        # Label altitude.
        if show_alt:
            ax.text(
                0.05,
                0.05,
                r'Alt.={:.1f}$km/s$ ({:.2f}$R_E$)'.format(
                    self['r'][alt], self['r'][alt] / 6371.0 + 1),
                size=12,
                transform=ax.transAxes,
                bbox={
                    'fc': 'white',
                    'ec': 'k'
                })

        if show_pts:
            ax.plot(x, y, '+w')

        # Add cbar if necessary.
        if add_cbar:
            cbar = plt.colorbar(cont, ticks=ticks, format=fmt, pad=0.01)
            if clabel == None:
                clabel = "%s (%s)" % (var, self[var].attrs['units'])
            cbar.set_label(clabel)
        else:
            cbar = None  # Need to return something, even if none.

        # Set title, labels, axis ranges (use defaults where applicable.)
        if title: ax.set_title(title)
        ax.set_yticks([]), ax.set_xticks([])
        ax.set_ylabel(r'Sun $\rightarrow$')
        colat_lim = 90. - rlim
        ax.set_xlim([-1 * colat_lim, colat_lim])
        ax.set_ylim([-1 * colat_lim, colat_lim])

        if colats:
            for colat in colats:
                r = latoffset * colat / np.sqrt(2)
                ax.add_patch(
                    Circle([0, 0], colat, fc='none', ec='k', ls='dashed'))
                ax.text(
                    r,
                    r,
                    '{:.0f}'.format(colat) + r'$^{\circ}$',
                    rotation=315,
                    ha='center',
                    va='center')
        else:
            for lat in lats:
                colat = 90 - lat
                r = latoffset * colat / np.sqrt(2)
                ax.add_patch(
                    Circle([0, 0], colat, fc='none', ec='k', ls='dashed'))
                ax.text(
                    r,
                    r,
                    '{:.0f}'.format(lat) + r'$^{\circ}$',
                    rotation=315,
                    ha='center',
                    va='center')

        return fig, ax, cont, cbar
Beispiel #3
0
    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,
                 **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.
        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.

        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]

        # 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')

        cnt1 = ax.contourf(self[hemi + 'psi'] * pi / 180.0 + pi / 2.,
                           self[hemi + 'theta'],
                           np.array(self[var]),
                           levs,
                           norm=crange,
                           cmap=cmap)
        # 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.,
                              self[hemi + '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)
            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, 40])

        # 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
Beispiel #4
0
    def add_slice(self, var, alt, time, nlev=31, zlim=None, target=None, 
                  loc=111, title=None, latoffset=1.05,
                  rlim=50., add_cbar=True, clabel=None,
                  show_pts=False, show_alt=True, dolog=False, 
                  lats=[75., 60.], colats=None, figsize=(8.34,7),
                  *args, **kwargs):
        '''
        Create a plot of variable *var* at altitude *alt*.

        If kwarg **target** is None (default), a new figure is 
        generated from scratch.  If target is a matplotlib Figure
        object, a new axis is created to fill that figure at subplot
        location **loc**.  If **target** is a matplotlib Axes object, 
        the plot is placed into that axis.
        '''

        import matplotlib.pyplot as plt
        from matplotlib.colors import (LogNorm, Normalize)
        from matplotlib.patches import Circle
        from matplotlib.ticker import (LogLocator, LogFormatter, 
                                       LogFormatterMathtext, MultipleLocator)

        # Grab the slice of data that we want:
        if type(time) == type(self['time'][0]):
            if time not in self['time']:
                raise ValueError('Time not in object')
            time = np.arange(self.attrs['nTime'])[self['time']==time][0]

        fig, ax = set_target(target, loc=loc, figsize=figsize)
        ax.set_aspect('equal')

        # Create values in Cartesian plane.
        self._get_cartXY()

        # Grab values from correct time/location.
        x = self._xLat[:,time, alt]
        y = self._yLat[:,time, alt]
        value = self[var][:, time, alt]

        # Get max/min if none given.
        if zlim==None:
            zlim=[0,0]
            zlim[0]=value.min(); zlim[1]=value.max()
            if dolog and zlim[0]<=0:
                zlim[0] = np.min( [0.0001, zlim[1]/1000.0] )

        # Create levels and set norm based on dolog.
        if dolog:
            levs = np.power(10, np.linspace(np.log10(zlim[0]), 
                                            np.log10(zlim[1]), nlev))
            z=np.where(value>zlim[0], value, 1.01*zlim[0])
            norm=LogNorm()
            ticks=LogLocator()
            fmt=LogFormatterMathtext()
        else:
            levs = np.linspace(zlim[0], zlim[1], nlev)
            z=value
            norm=None
            ticks=None
            fmt=None
        
        # Create contour plot.
        cont=ax.tricontourf(x, y, z, levs, *args, norm=norm, **kwargs)
        
        # Label altitude.
        if show_alt:
            ax.text(0.05, 0.05, r'Alt.={:.1f}$km/s$ ({:.2f}$R_E$)'.format(
                self['r'][alt], self['r'][alt]/6371.0+1), size=12, 
                    transform=ax.transAxes, bbox={'fc':'white', 'ec':'k'})

        if show_pts:
            ax.plot(x, y, '+w')

        # Add cbar if necessary.
        if add_cbar:
            cbar=plt.colorbar(cont, ticks=ticks, format=fmt, pad=0.01)
            if clabel==None: 
                clabel="%s (%s)" % (var, self[var].attrs['units'])
            cbar.set_label(clabel)
        else:
            cbar=None # Need to return something, even if none.
 
        # Set title, labels, axis ranges (use defaults where applicable.)
        if title: ax.set_title(title)
        ax.set_yticks([]), ax.set_xticks([])
        ax.set_ylabel(r'Sun $\rightarrow$')
        colat_lim = 90.-rlim
        ax.set_xlim([-1*colat_lim, colat_lim])
        ax.set_ylim([-1*colat_lim, colat_lim])

        if colats:
            for colat in colats:
                r = latoffset*colat/np.sqrt(2)
                ax.add_patch(Circle([0,0],colat,fc='none',ec='k',ls='dashed'))
                ax.text(r, r, '{:.0f}'.format(colat)+r'$^{\circ}$',
                        rotation=315,ha='center', va='center')
        else:
            for lat in lats:
                colat = 90 - lat
                r = latoffset*colat/np.sqrt(2)
                ax.add_patch(Circle([0,0],colat,fc='none',ec='k',ls='dashed'))
                ax.text(r, r, '{:.0f}'.format(lat)+r'$^{\circ}$', 
                        rotation=315,ha='center', va='center')
                

        return fig, ax, cont, cbar
Beispiel #5
0
    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, **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.
        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.

        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]

        # 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')

        cnt1 = ax.contourf(self[hemi+'psi']*pi/180.0+pi/2., 
                           self[hemi+'theta'], np.array(self[var]),
                           levs, norm=crange, cmap=cmap)
        # 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., 
                              self[hemi+'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)
            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,40])

        # 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