def gitm_single_3D_image(plot_type, zkey, gdata, title=None, figname=None,
                         draw=True, aindex=-1, nlat=90, slat=-90, linc=6,
                         earth=False, tlon=90, zmax=None, zmin=None,
                         zcolor=None, data_type="contour", faspect=True,
                         meq=False, terminator=False, m=None, *args, **kwargs):
    '''
    Creates a rectangular or polar map projection plot for a specified latitude
    range.
    Input: plot_type  = key to determine plot type (rectangular, polar)
           zkey       = key for z variable (ie 'Vertical TEC')
           gData      = gitm bin structure
           title      = plot title
           figname    = file name to save figure as (default is none)
           draw       = draw to screen? (default is True)
           aindex     = altitude index (default -1 if it is a 2D parameter)
           nlat       = northern latitude limit (degrees North, default 90)
           slat       = southern latitude limit (degrees North, defalut -90)
           linc       = number of latitude tick incriments (default 6)
           earth      = include continent outlines for Earth (default False)
           tlon       = longitude at the top of a polar dial (degrees east,
                        default 90)
           zmax       = Maximum z range (default None)
           zmin       = Minimum z range (default None)
           zcolor     = Color map for the z variable.  If none, will be chosen
                        based on the z range (default=None)
           data_type  = scatter or contour (default=scatter)
           faspect    = Fix the aspect of Earth if using outlines (default=True)
           meq        = Include the geomagnetic equator?  (default=False)
           terminator = Include the solar terminator? (default=False)
           m          = Handle for earth map (default=None)

    Output: f = figure handle
            m = map handle (or None)
    '''
    # Set the altitude and latitude limits.  For polar plots, latitudes
    # Above and below 90 degrees will cause the routine to fail
    ialt = 0
    if aindex > 0:
        ialt = aindex

    if plot_type.find("polar") >= 0:
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1
    else:
        imin = 0
        imax = gdata.attrs['nLat']

    # If including the solar terminator, extract the UT date
    tdt = None
    if terminator:
        tdt = gdata['time']

    # Format title
    spec_title = "{:} UT".format(gdata['time'])
    if aindex >= 0:
        spec_title = "{:s} slice at {:.2f} km".format(spec_title, 1.0e-3 *
                                                      gdata['Altitude'][0,0,
                                                                        ialt])

    if title:
        title = "{:s}\n{:s}".format(spec_title, title)
    else:
        title = spec_title

    # Output the figure
    fm = p3g.plot_single_3D_image(plot_type,
                                  np.array(gdata['dLat'][:,imin:imax,ialt]),
                                  np.array(gdata['dLon'][:,imin:imax,ialt]),
                                  np.array(gdata[zkey][:,imin:imax,ialt]),
                                  gdata[zkey].attrs['name'],
                                  gdata[zkey].attrs['scale'],
                                  gdata[zkey].attrs['units'], zmax=zmax,
                                  zmin=zmin, zcolor=zcolor, title=title,
                                  figname=figname, draw=draw, nlat=nlat,
                                  slat=slat, linc=linc, tlon=tlon,
                                  data_type=data_type, meq=meq, earth=earth,
                                  m=m, faspect=faspect, term_datetime=tdt)
    return fm
def plot_net_gitm_comp(plot_type, lon_data, lat_data, obs_data, obs_name,
                       obs_scale, obs_units, diff_data, diff_name, diff_scale,
                       diff_units, gitm_key, gitm_alt, gdata, gitm_name,
                       diff_max=None, zmax=None, zmin=None, title=None,
                       color=True, bcolor='#747679', data_coff=False,
                       diff_coff=True, figname=None, draw=True, latlim1=90,
                       latlim2=-90, linc=6, tlon=90, meq=False, earth=False,
                       map_list=[], faspect=True, term_datetime=None,
                       extra_lines=False, *args, **kwargs):
    '''
    Creates three plots of a specified type, one showing the observations, one
    showing the GITM data, and one showing the difference between the two.

    Input: plot_type  = key to determine plot type (rectangular, polar,
                        nsglobal, or snapshot)
           lon_data   = Numpy array with longitude data for matching model-obs
                        points
           lat_data   = Numpy array with latitude data for matching model-obs
                        points
           obs_data   = Numpy array with observational data for matching
                        model-obs points
           obs_name   = Name portion of the observational data label
           obs_scale  = Scale (linear/exponential) for plotting obs. data
           obs_units  = Unit portion of the observational data label
           diff_data  = Numpy array with differences for matching model-obs
                        points
           gitm_key   = Key for the GITM data
           gitm_alt   = Altitude in km to plot the GITM data at.  For a 2D
                        variable like hmF2 or TEC, use 0.0 km.
           gdata      = GitmBin structure with model observations.
           gitm_name  = Name portion of the GITM data label
           diff_max   = Maximum value for the difference (absolute value),
                        if None, will be determined in script (default=None)
           zmin       = minimum value for z variable (default=None)
           zmax       = maximum value for z variable (default=None)
           title      = Plot title (default=None)
           color      = Color (True, default) or black and white (False)?
           bcolor     = Background color (default=)
           data_coff  = Center the data color scale about zero (False, default)?
           diff_coff  = Center the diff color scale about zero (True, default)?
           figname    = Output figure name with a .png suffix (default=None)
           draw       = Draw to screen? (default=True)
           latlim1    = First latitude limit (degrees North, default=90).
                        Purpose varies depending on plot type.  For rectangular,
                        this is the northern latitude limit.  For polar, this
                        is the latitude at the center of the dial.  For
                        snapshot, this is the lower boundary of polar dials.
                        It is not used for nsglobal.
           latlim2    = Second latitude limit (degrees North, default=-90).
                        Purpose varies depending on plot type.  For rectangular,
                        this is the southern latitude limit.  For polar, this
                        is the latitude at the edge of the dial.  This option is
                        not used with the snapshot or nsglobal option.
           linc       = Number of latitude tick incriments (default=6)
           tlon       = Longitude on top of the polar dial (degrees East,
                        default=90)
           meq            = Add a line for the geomagnetic equator?
                            (default=False)
           earth         = Include continent outlines for Earth (default=False)
           map_list      = List of map handles for the specified plot_type
                           (default=empty list)
           faspect       = Keep a true aspect ratio for maps? (default=True)
           term_datetime = Include the solar terminator by shading the night
                           time regions?  If so, include a datetime object
                           with the UT for this map.  Only used if earth=True.
           extra_lines   = Plot a specified lines (good for showing regional
                           boundaries) (default=False).  Provide a list of lists
                           which have the format shown:
                           [x np.array, y np.array, style string (eg 'k-')]
                           where x is in degrees longitude and y is in 
                           degrees latitude

    Output: f = handle to figure
    '''
    rout_name = "plot_net_gitm_comp"

    # Get the desired color bars
    data_color = gpr.choose_contour_map(color, data_coff)
    diff_color = gpr.choose_contour_map(color, diff_coff)

    # Get the altitude index
    ialt = 0
    if gitm_alt > 0.0:
        ialt = gpr.find_alt_index(gdata, 0, 0, alt, units="km")

    # Initialize the z variables, if desired.  GITM and Observational data
    # should share the same scale.
    if(zmin is None):
        obsmin = np.nanmin(obs_data)
        gitmin = np.nanmin(gdata[gitm_key][:,:,ialt])
        zmin = min(obsmin,gitmin)
    if(zmax is None):
        obsmax = np.nanmax(obs_data)
        gitmax = np.nanmax(gdata[gitm_key][:,:,ialt])
        zmax = max(obsmax, gitmax)

    zran = round((zmax-zmin)/6.0)
    if(zran != 0.0):
        zmin = math.floor(float("{:.14f}".format(zmin / zran))) * zran
        zmax = math.ceil(float("{:.14f}".format(zmax / zran))) * zran

    # Set the difference max/min limits, if desired
    if diff_max is None:
        diff_max = max(np.nanmax(diff_data), abs(np.nanmin(diff_data)))

    diff_min = -1.0 * diff_max

    # Initialize the figure, setting the height for a 3 subfigure stack
    fwidth = 6
    fheight = 12
    if(plot_type.find("global") > 0):
        fwidth *= 1.5
    if(plot_type.find("shot") > 0):
        fwidth *= 1.5
        fheight *= 1.5

    f = plt.figure(figsize=(fwidth,fheight))

    # Plot the three datasets using the desired format
    if plot_type.find("shot") > 0:
        if len(map_list) == 3:
            ml = map_list[0]
            mn = map_list[1]
            ms = map_list[2]
        else:
            ml = None
            mn = None
            ms = None
        # Output the observations as a scatter plot
        axl,ml,axn,mn,axs,ms = p3g.plot_snapshot_subfigure(f, 3, 0, lat_data,
                                                           lon_data, obs_data,
                                                           obs_name, obs_scale,
                                                           obs_units, zmax,
                                                           zmin, data_color,
                                                           tlon=tlon,
                                                           blat=latlim1,
                                                           xl=False, yl=False,
                                                           xt=False,
                                                           bcolor=bcolor,
                                                           meq=meq, earth=earth,
                                                           ml=ml, mn=mn, ms=ms,
                                                           faspect=faspect,
                                                           term_datetime=term_datetime)
        # Output the gitm data as a contour after ensuring that the GITM array
        # isn't padded to include unrealistic latitudes
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1 
        p3g.plot_snapshot_subfigure(f, 3, 1,
                                    np.array(gdata['dLat'][:,imin:imax,ialt]),
                                    np.array(gdata['dLon'][:,imin:imax,ialt]),
                                    np.array(gdata[gitm_key][:,imin:imax,ialt]),
                                    gitm_name, gdata[gitm_key].attrs["scale"],
                                    gdata[gitm_key].attrs["units"], zmax, zmin,
                                    data_color, cb=True, cloc="r", tlon=tlon,
                                    blat=latlim1, title=False, xl=False,
                                    xt=False, bcolor=bcolor, meq=meq,
                                    earth=earth, ml=ml, mn=mn, ms=ms,
                                    faspect=faspect, data_type="contour",
                                    term_datetime=term_datetime)
        # Output the differences as a scatter plot
        p3g.plot_snapshot_subfigure(f, 3, 2, lat_data, lon_data, diff_data,
                                    diff_name, diff_scale, diff_units, diff_max,
                                    diff_min, diff_color, tlon=tlon,
                                    blat=latlim1, title=False, yl=False,
                                    bcolor=bcolor, meq=meq, earth=earth, ml=ml,
                                    mn=mn, ms=ms, faspect=faspect,
                                    term_datetime=term_datetime)
        map_list = list([ml, mn, ms])
    elif plot_type.find("nsglobal") >= 0:
        if len(map_list) == 2:
            mn = map_list[0]
            ms = map_list[1]
        else:
            mn = None
            ms = None

        # Check for boundary lines to plot
        eline_north = False
        eline_south = False
        if type(extra_lines) is list:
            if len(extra_lines) >= 1:
                eline_north = extra_lines[0]

                if len(extra_lines) >= 2:
                    eline_south = extra_lines[1]
                else:
                    print "Only one boundary provided, plotting in north"
            else:
                print "No boundaries provided, better to declare as False"

        # Output the observations as a scatter plot
        axn1,mn,axs1,ms = p3g.plot_nsglobal_subfigure(f, 3, 0, lat_data,
                                                      lon_data, obs_data,
                                                      obs_name, obs_scale,
                                                      obs_units, zmax, zmin,
                                                      data_color, title=True, cb=True,
                                                      elat=latlim1, tlon=tlon, rl=False,
                                                      tl=False, bcolor=bcolor,
                                                      earth=earth, mn=mn, ms=ms,
                                                      faspect=faspect, term_datetime=term_datetime,
                                                      extra_line_n=eline_north,
                                                      extra_line_s=eline_south)
        # Output the gitm data as a contour after ensuring that the GITM array
        # isn't padded to include unrealistic latitudes 
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1 

        axn2,mn,axs2,ms = p3g.plot_nsglobal_subfigure(f, 3, 1, np.array(gdata['dLat'][:,imin:imax,ialt]), np.array(gdata['dLon'][:,imin:imax,ialt]), np.array(gdata[gitm_key][:,imin:imax,ialt]), gitm_name, gdata[gitm_key].attrs["scale"], gdata[gitm_key].attrs["units"], zmax, zmin, data_color, title=False, cb=True, elat=latlim1, tlon=tlon, tl=False, bcolor=bcolor, earth=earth, mn=mn, ms=ms, data_type="contour", term_datetime=term_datetime, extra_line_n=eline_north, extra_line_s=eline_south)

        # Output the differences as a scatter plot
        p3g.plot_nsglobal_subfigure(f, 3, 2, lat_data, lon_data, diff_data,
                                    diff_name, diff_scale, diff_units, diff_max,
                                    diff_min, diff_color, title=False, cb=True,
                                    elat=latlim1, tlon=tlon, rl=False, bcolor=bcolor,
                                    earth=earth, mn=mn, ms=ms, faspect=faspect,
                                    term_datetime=term_datetime,
                                    extra_line_n=eline_north,
                                    extra_line_s=eline_south)
        map_list = list([mn, ms])
    elif plot_type.find("rect") >= 0:
        if len(map_list) == 1:
            m = map_list[0]
        else:
            m = None
        # Output the observations as a scatter plot
        ax = f.add_subplot(3,1,1)
        con1, m = p3g.plot_rectangular_3D_global(ax, lat_data, lon_data,
                                                 obs_data, obs_name, obs_scale,
                                                 obs_units, zmin, zmax,
                                                 data_color, nlat=latlim1,
                                                 slat=latlim2, linc=linc,
                                                 cloc="r", xl=False, xt=False,
                                                 yl=False, meq=meq,
                                                 bcolor=bcolor, earth=earth,
                                                 m=m, faspect=faspect,
                                                 term_datetime=term_datetime)
        # Output the gitm data as a contour
        ax = f.add_subplot(3,1,2)
        con2, m = p3g.plot_rectangular_3D_global(ax, np.array(gdata['dLat'][:,:,ialt]), np.array(gdata['dLon'][:,:,ialt]), np.array(gdata[gitm_key][:,:,ialt]),
                                                 gitm_name,
                                                 gdata[gitm_key].attrs["scale"],
                                                 gdata[gitm_key].attrs["units"],
                                                 zmin, zmax, data_color,
                                                 nlat=latlim1, slat=latlim2,
                                                 linc=linc, cb=True, cloc="r",
                                                 xl=False, xt=False,
                                                 bcolor=bcolor, meq=meq,
                                                 earth=earth, m=m,
                                                 faspect=faspect,
                                                 data_type="contour",
                                                 term_datetime=term_datetime)
        # Adjust plot dimensions if necessary
        if not earth:
            con1_dim = list(con1.axes.get_position().bounds)
            con2_dim = list(con2.ax.get_position().bounds)
            con2_dim[2] = con1_dim[2]
            con2.ax.set_position(con2_dim)

        # Output the differences as a scatter plot
        ax = f.add_subplot(3,1,3)
        p3g.plot_rectangular_3D_global(ax, lat_data, lon_data, diff_data,
                                       diff_name, diff_scale, diff_units,
                                       diff_min, diff_max, diff_color,
                                       nlat=latlim1, slat=latlim2, linc=linc,
                                       cloc="r", yl=False, bcolor=bcolor,
                                       meq=meq, earth=earth, m=m,
                                       faspect=faspect,
                                       term_datetime=term_datetime)
        map_list = list([m])
    elif plot_type.find("polar") >= 0:
        if len(map_list) == 1:
            m = map_list[0]
        else:
            m = None

        pf = True
        if earth:
            pf = False
        # Output the observations as a scatter plot
        ax = f.add_subplot(3,1,1, polar=pf)
        con1,m = p3g.plot_polar_3D_global(ax, 3, lat_data, lon_data, obs_data,
                                          obs_name, obs_scale, obs_units, zmin,
                                          zmax, data_color, center_lat=latlim1,
                                          edge_lat=latlim2, linc=linc,
                                          top_lon=tlon, cloc="r", tl=False,
                                          rl=False, bcolor=bcolor, earth=earth,
                                          m=m, faspect=faspect,
                                          term_datetime=term_datetime)
        # Output the gitm data as a contour after ensuring that the GITM
        # array isn't padded to include unrealistic latitudes
        ax = f.add_subplot(3,1,2, polar=pf)
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1 
        con2,m = p3g.plot_polar_3D_global(ax, 3, np.array(gdata['dLat'][:,imin:imax,ialt]), np.array(gdata['dLon'][:,imin:imax,ialt]), np.array(gdata[gitm_key][:,imin:imax,ialt]),
                                          gitm_name,
                                          gdata[gitm_key].attrs["scale"],
                                          gdata[gitm_key].attrs["units"], zmin,
                                          zmax, data_color, center_lat=latlim1,
                                          edge_lat=latlim2, linc=linc,
                                          top_lon=tlon, cb=True, cloc="r",
                                          tl=False, bcolor=bcolor, earth=earth,
                                          m=m, faspect=faspect,
                                          data_type="contour",
                                          term_datetime=term_datetime)

        con1_dim = list(con1.axes.get_position().bounds)
        con2_dim = list(con2.ax.get_position().bounds)
        con2_dim[0] = con2_dim[0] - 0.05
        con2_dim[2] = con1_dim[2]
        con2.ax.set_position(con2_dim)

        # Output the differences as a scatter plot
        ax = f.add_subplot(3,1,3, polar=pf)
        p3g.plot_polar_3D_global(ax, 3, lat_data, lon_data, diff_data,
                                 diff_name, diff_scale, diff_units, diff_min,
                                 diff_max, diff_color, center_lat=latlim1,
                                 edge_lat=latlim2, linc=linc, top_lon=tlon,
                                 cloc="r", rl=False, bcolor=bcolor, earth=earth,
                                 m=m, faspect=faspect,
                                 term_datetime=term_datetime)
        map_list = list([m])
    else:
        print rout_name, "ERROR: uknown plot type [", plot_type, "]"
        return

    if title:
        f.suptitle(title, size="medium")

    # Adjust subplot locations
    if plot_type.find("rect") >= 0 or plot_type.find("polar") >= 0:
        plt.subplots_adjust(left=.15)

    # Draw to screen if desired
    if draw:
        if plt.isinteractive():
            plt.draw() #In interactive mode, you just "draw".
        else:
            # W/o interactive mode, "show" stops the user from typing more 
            # at the terminal until plots are drawn.
            plt.show()

    # Save output file
    if figname is not None:
        plt.savefig(figname)

    return(f, map_list)
def gitm_mult_3D_slices(plot_type, zkey, gdata, aindex, title=None,
                        figname=None, draw=True, nlat=90, slat=-90, linc=6,
                        earth=False, tlon=90, zmax=None, zmin=None, 
                        zcolor=None, data_type="contour", meq=False,
                        faspect=True, terminator=False, *args, **kwargs):
    '''
    Creates a rectangular or polar map projection plot for a specified latitude
    range.
    Input: plot_type  = key to determine plot type (rectangular, polar)
           zkey       = key for z variable (ie 'Vertical TEC')
           gdata      = gitm bin structure
           aindex     = list of altitude indices
           title      = plot title
           figname    = file name to save figure as (default is none)
           draw       = draw to screen? (default is True)
           nlat       = northern latitude limit (degrees North, default 90)
           slat       = southern latitude limit (degrees North, defalut -90)
           linc       = number of latitude tick incriments (default 6)
           earth      = include Earth continent outlines (default False)
           tlon       = longitude on the top, for polar plots (degrees East,
                        default 90)
           zmax       = maximum z range (default None)
           zmin       = minimum z range (default None)
           zcolor     = Color spectrum for z data.  If not specified, will be
                        determined by the z range. (default=None)
           data_type  = Contour or scatter plot? (default=scatter)
           meq        = Add a line for the geomagnetic equator? (default=False)
           earth      = include Earth continent outlines (default False)
           faspect    = Fix aspect ratio if using continents (default=True)
           terminator = Include the solar terminator by shading the night
                        time regions? (default=False)
    '''
    # Set the latitude limits.  For polar plots, latitudes above and below
    # 90 degrees will cause the routine to fail
    if plot_type.find("polar") >= 0:
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1
    else:
        imin = 0
        imax = gdata.attrs['nLat']

    # If including the solar terminator, extract the UT date
    tdt = None
    if terminator:
        tdt = gdata['time']

    # Format title
    spec_title = "{:} UT slice".format(gdata['time'])

    if title:
        title = "{:s}\n{:s}".format(spec_title, title)
    else:
        title = spec_title

    # Output the figure
    f = p3g.plot_mult_3D_slices(plot_type, 2, aindex,
                                np.array(gdata['dLat'][:,imin:imax,:]),
                                np.array(gdata['dLon'][:,imin:imax,:]),
                                np.array(gdata[zkey][:,imin:imax,:]),
                                gdata[zkey].attrs['name'],
                                gdata[zkey].attrs['scale'],
                                gdata[zkey].attrs['units'], zmax=zmax,
                                zmin=zmin, zcolor=zcolor, title=title,
                                figname=figname, draw=draw, nlat=nlat,
                                slat=slat, linc=linc, tlon=tlon,
                                data_type=data_type, meq=meq, earth=earth,
                                faspect=faspect, term_datetime=tdt)
    return f
def gitm_global_3D_snapshot(zkey, gdata, title=None, figname=None, draw=True,
                            aindex=-1, tlon=90, polar_blat=45, rect_blat=45,
                            earth=False, zmax=None, zmin=None,
                            zcolor=None, meq=False, data_type="contour",
                            terminator=False, ml=None, mn=None, ms=None,
                            *args, **kwargs):
    '''
    Creates a map projection plot for the entire globe, seperating the polar
    and central latitude regions.
    Input: zkey       = key for z variable (ie 'Vertical TEC')
           gData      = gitm bin structure
           title      = plot title
           figname    = file name to save figure as (default is none)
           draw       = output a screen image? (default is True)
           aindex     = altitude index (default -1 if it is a 2D parameter)
           tlon       = longitude at the top of the polar dial (degrees East,
                        default 90)
           polar_blat = co-latitude of the lower boundary of the polar dials
                        (default 45)
           rect_blat  = Upper bounding co-latitude of the rectangular map
                        (default 45)
           earth      = include Earth continent outlines (default False)
           zmax       = maximum z limit (default None)
           zmin       = minimum z limit (default None)
           zcolor     = Color scale for z variable.  If not specified, will be
                        determined by the z range (default=None)
           meq        = Add a line for the geomagnetic equator? (default=False)
           data_type  = Type of plot to make scatter/contour (default=contour)
           terminator = Include the solar terminator by shading the night
                        time regions? (default=False)
           ml         = Low latitude map handle (default=None)
           mn         = Northern latitude map handle (default=None)
           ms         = Southern latitude map handle (default=None)

    Output: f  = Figure handle
            ml = Low latitude map handle
            mn = Northern latitude map handle
            ms = Southern latitude map handle
    '''
    # Set the altitude and latitude limits.  For polar plots, latitudes
    # Above and below 90 degrees will cause the routine to fail
    ialt = 0
    if aindex > 0:
        ialt = aindex

    (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
    (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
    imax += 1

    # If including the solar terminator, extract the UT date
    tdt = None
    if terminator:
        tdt = gdata['time']

    # Format title
    spec_title = "{:} UT".format(gdata['time'])
    if aindex >= 0:
        spec_title = "{:s} slice at {:.2f} km".format(spec_title, 1.0e-3 *
                                                      gdata['Altitude'][0,0,
                                                                        ialt])

    if title:
        title = "{:s}\n{:s}".format(spec_title, title)
    else:
        title = spec_title

    # Output figure
    fm = p3g.plot_global_3D_snapshot(np.array(gdata['dLat'][:,imin:imax,ialt]),
                                     np.array(gdata['dLon'][:,imin:imax,ialt]),
                                     np.array(gdata[zkey][:,imin:imax,ialt]),
                                     gdata[zkey].attrs['name'],
                                     gdata[zkey].attrs['scale'],
                                     gdata[zkey].attrs['units'], zmax=zmax,
                                     zmin=zmin, zcolor=zcolor, title=title,
                                     figname=figname, draw=draw, tlon=tlon,
                                     polar_blat=polar_blat, rect_blat=rect_blat,
                                     meq=meq, earth=earth, ml=ml, mn=mn, ms=ms,
                                     data_type=data_type, term_datetime=tdt)
    return fm
def gitm_single_nsglobal_3D_image(zkey, gdata, title=None, figname=None,
                                  draw=True, aindex=-1, plat=90, elat=0, linc=3,
                                  tlon=90, earth=False, zmax=None, zmin=None,
                                  zcolor=None, data_type="contour",
                                  terminator=False, mn=None, ms=None,
                                  *args, **kwargs):
    '''
    Creates a figure with two polar map projections for the northern and 
    southern ends of a specified latitude range.
    Input: zkey       = key for z variable (ie 'Vertical TEC')
           gdata      = gitm bin structure
           title      = plot title
           figname    = file name to save figure as (default is none)
           draw       = draw to screen? (default is True)
           aindex     = altitude index (default -1 if it is a 2D parameter)
           plat       = polar latitude limit (degrees North, default +/-90)
           elat       = equatorial latitude limit (degrees North, defalut 0)
           linc       = number of latitude tick incriments (default 6)
           tlon       = longitude to place on the polar dial top (degrees east,
                        default 90)
           earth      = include Earth continent outlines (default False)
           zmax       = maximum z range (default None)
           zmin       = mininimum z range (default None)
           zcolor     = Color scale for plotting the z data.  If not specified,
                        this will be determined by the z range (default=None)
           data_type  = Type of plot to make scatter/contour (default=scatter)
           terminator = Include the solar terminator by shading the night
                        time regions? (default=False)
           mn         = Northern latitude map handle (default=None)
           ms         = Southern latitude map handle (default=None)

    Output: f  = figure handle
            mn = Northern latitude map handle
            ms = Southern latitude map handle
    '''
    # Set the altitude and latitude limits.  For polar plots, latitudes
    # Above and below 90 degrees will cause the routine to fail
    ialt = 0
    if aindex > 0:
        ialt = aindex

    (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
    (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
    imax += 1

    # If including the solar terminator, extract the UT date
    tdt = None
    if terminator:
        tdt = gdata['time']

    # Format title
    spec_title = "{:} UT".format(gdata['time'])
    if aindex >= 0:
        spec_title = "{:s} slice at {:.2f} km".format(spec_title, 1.0e-3 *
                                                      gdata['Altitude'][0,0,
                                                                        ialt])
    if title:
        title = "{:s}\n{:s}".format(spec_title, title)
    else:
        title = spec_title

    # Output figure
    fm = p3g.plot_single_nsglobal_3D_image(np.array(gdata['dLat'][:,imin:imax,ialt]), np.array(gdata['dLon'][:,imin:imax,ialt]), np.array(gdata[zkey][:,imin:imax,ialt]),
                                           gdata[zkey].attrs['name'],
                                           gdata[zkey].attrs['scale'],
                                           gdata[zkey].attrs['units'],
                                           zmax=zmax, zmin=zmin, zcolor=zcolor,
                                           title=title, figname=figname,
                                           draw=draw, plat=plat, elat=elat,
                                           linc=linc, tlon=tlon, earth=earth,
                                           mn=mn, ms=ms, data_type=data_type,
                                           term_datetime=tdt)

    return fm
#---------------------------------------------
# Determine the limits of each z parameter.  To do this we need to open
# one of the gitm files

for z,zkey in enumerate(zkeys):
    # Find the index for the specified altitude.  This will be the same for
    # all files, longitudes, and latititudes.

    aindex = gpr.find_alt_index(gitmbins[0], 1, 1, zalts[z], zunits[z])
    zindexes.append(aindex)

    # Find the indexes for the specified latitude limits.  The latitude will
    # be the same for all files and longitudes.  The northern limit will be
    # the same for all plot types.

    (nlon,nlat) = gpr.find_lon_lat_index(gitmbins[0], -1.0 , nlimits[z],
                                         "degrees")

    if zplot == "nspolar" or zplot == "snapshot":
        slimit = -1.0 * nlimits[z]
    else:
        slimit = slimits[z]

    (slon,slat) = gpr.find_lon_lat_index(gitmbins[0], 361.0, slimit, "degrees")
    
    # Test to ensure the validity of the indexes

    if nlat == slat:
        slat += 1
    elif nlat < slat:
        temp = nlat
        nlat = slat
示例#7
0
def plot_net_gitm_comp(plot_type,
                       lon_data,
                       lat_data,
                       obs_data,
                       obs_name,
                       obs_scale,
                       obs_units,
                       diff_data,
                       diff_name,
                       diff_scale,
                       diff_units,
                       gitm_key,
                       gitm_alt,
                       gdata,
                       gitm_name,
                       diff_max=None,
                       zmax=None,
                       zmin=None,
                       title=None,
                       color=True,
                       bcolor='#747679',
                       data_coff=False,
                       diff_coff=True,
                       figname=None,
                       draw=True,
                       latlim1=90,
                       latlim2=-90,
                       linc=6,
                       tlon=90,
                       meq=False,
                       earth=False,
                       map_list=[],
                       faspect=True,
                       term_datetime=None,
                       extra_lines=False,
                       *args,
                       **kwargs):
    '''
    Creates three plots of a specified type, one showing the observations, one
    showing the GITM data, and one showing the difference between the two.

    Input: plot_type  = key to determine plot type (rectangular, polar,
                        nsglobal, or snapshot)
           lon_data   = Numpy array with longitude data for matching model-obs
                        points
           lat_data   = Numpy array with latitude data for matching model-obs
                        points
           obs_data   = Numpy array with observational data for matching
                        model-obs points
           obs_name   = Name portion of the observational data label
           obs_scale  = Scale (linear/exponential) for plotting obs. data
           obs_units  = Unit portion of the observational data label
           diff_data  = Numpy array with differences for matching model-obs
                        points
           gitm_key   = Key for the GITM data
           gitm_alt   = Altitude in km to plot the GITM data at.  For a 2D
                        variable like hmF2 or TEC, use 0.0 km.
           gdata      = GitmBin structure with model observations.
           gitm_name  = Name portion of the GITM data label
           diff_max   = Maximum value for the difference (absolute value),
                        if None, will be determined in script (default=None)
           zmin       = minimum value for z variable (default=None)
           zmax       = maximum value for z variable (default=None)
           title      = Plot title (default=None)
           color      = Color (True, default) or black and white (False)?
           bcolor     = Background color (default=)
           data_coff  = Center the data color scale about zero (False, default)?
           diff_coff  = Center the diff color scale about zero (True, default)?
           figname    = Output figure name with a .png suffix (default=None)
           draw       = Draw to screen? (default=True)
           latlim1    = First latitude limit (degrees North, default=90).
                        Purpose varies depending on plot type.  For rectangular,
                        this is the northern latitude limit.  For polar, this
                        is the latitude at the center of the dial.  For
                        snapshot, this is the lower boundary of polar dials.
                        It is not used for nsglobal.
           latlim2    = Second latitude limit (degrees North, default=-90).
                        Purpose varies depending on plot type.  For rectangular,
                        this is the southern latitude limit.  For polar, this
                        is the latitude at the edge of the dial.  This option is
                        not used with the snapshot or nsglobal option.
           linc       = Number of latitude tick incriments (default=6)
           tlon       = Longitude on top of the polar dial (degrees East,
                        default=90)
           meq            = Add a line for the geomagnetic equator?
                            (default=False)
           earth         = Include continent outlines for Earth (default=False)
           map_list      = List of map handles for the specified plot_type
                           (default=empty list)
           faspect       = Keep a true aspect ratio for maps? (default=True)
           term_datetime = Include the solar terminator by shading the night
                           time regions?  If so, include a datetime object
                           with the UT for this map.  Only used if earth=True.
           extra_lines   = Plot a specified lines (good for showing regional
                           boundaries) (default=False).  Provide a list of lists
                           which have the format shown:
                           [x np.array, y np.array, style string (eg 'k-')]
                           where x is in degrees longitude and y is in 
                           degrees latitude

    Output: f = handle to figure
    '''
    rout_name = "plot_net_gitm_comp"

    # Get the desired color bars
    data_color = gpr.choose_contour_map(color, data_coff)
    diff_color = gpr.choose_contour_map(color, diff_coff)

    # Get the altitude index
    ialt = 0
    if gitm_alt > 0.0:
        ialt = gpr.find_alt_index(gdata, 0, 0, alt, units="km")

    # Initialize the z variables, if desired.  GITM and Observational data
    # should share the same scale.
    if (zmin is None):
        obsmin = np.nanmin(obs_data)
        gitmin = np.nanmin(gdata[gitm_key][:, :, ialt])
        zmin = min(obsmin, gitmin)
    if (zmax is None):
        obsmax = np.nanmax(obs_data)
        gitmax = np.nanmax(gdata[gitm_key][:, :, ialt])
        zmax = max(obsmax, gitmax)

    zran = round((zmax - zmin) / 6.0)
    if (zran != 0.0):
        zmin = math.floor(float("{:.14f}".format(zmin / zran))) * zran
        zmax = math.ceil(float("{:.14f}".format(zmax / zran))) * zran

    # Set the difference max/min limits, if desired
    if diff_max is None:
        diff_max = max(np.nanmax(diff_data), abs(np.nanmin(diff_data)))

    diff_min = -1.0 * diff_max

    # Initialize the figure, setting the height for a 3 subfigure stack
    fwidth = 6
    fheight = 12
    if (plot_type.find("global") > 0):
        fwidth *= 1.5
    if (plot_type.find("shot") > 0):
        fwidth *= 1.5
        fheight *= 1.5

    f = plt.figure(figsize=(fwidth, fheight))

    # Plot the three datasets using the desired format
    if plot_type.find("shot") > 0:
        if len(map_list) == 3:
            ml = map_list[0]
            mn = map_list[1]
            ms = map_list[2]
        else:
            ml = None
            mn = None
            ms = None
        # Output the observations as a scatter plot
        axl, ml, axn, mn, axs, ms = p3g.plot_snapshot_subfigure(
            f,
            3,
            0,
            lat_data,
            lon_data,
            obs_data,
            obs_name,
            obs_scale,
            obs_units,
            zmax,
            zmin,
            data_color,
            tlon=tlon,
            blat=latlim1,
            xl=False,
            yl=False,
            xt=False,
            bcolor=bcolor,
            meq=meq,
            earth=earth,
            ml=ml,
            mn=mn,
            ms=ms,
            faspect=faspect,
            term_datetime=term_datetime)
        # Output the gitm data as a contour after ensuring that the GITM array
        # isn't padded to include unrealistic latitudes
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1
        p3g.plot_snapshot_subfigure(f,
                                    3,
                                    1,
                                    np.array(gdata['dLat'][:, imin:imax,
                                                           ialt]),
                                    np.array(gdata['dLon'][:, imin:imax,
                                                           ialt]),
                                    np.array(gdata[gitm_key][:, imin:imax,
                                                             ialt]),
                                    gitm_name,
                                    gdata[gitm_key].attrs["scale"],
                                    gdata[gitm_key].attrs["units"],
                                    zmax,
                                    zmin,
                                    data_color,
                                    cb=True,
                                    cloc="r",
                                    tlon=tlon,
                                    blat=latlim1,
                                    title=False,
                                    xl=False,
                                    xt=False,
                                    bcolor=bcolor,
                                    meq=meq,
                                    earth=earth,
                                    ml=ml,
                                    mn=mn,
                                    ms=ms,
                                    faspect=faspect,
                                    data_type="contour",
                                    term_datetime=term_datetime)
        # Output the differences as a scatter plot
        p3g.plot_snapshot_subfigure(f,
                                    3,
                                    2,
                                    lat_data,
                                    lon_data,
                                    diff_data,
                                    diff_name,
                                    diff_scale,
                                    diff_units,
                                    diff_max,
                                    diff_min,
                                    diff_color,
                                    tlon=tlon,
                                    blat=latlim1,
                                    title=False,
                                    yl=False,
                                    bcolor=bcolor,
                                    meq=meq,
                                    earth=earth,
                                    ml=ml,
                                    mn=mn,
                                    ms=ms,
                                    faspect=faspect,
                                    term_datetime=term_datetime)
        map_list = list([ml, mn, ms])
    elif plot_type.find("nsglobal") >= 0:
        if len(map_list) == 2:
            mn = map_list[0]
            ms = map_list[1]
        else:
            mn = None
            ms = None

        # Check for boundary lines to plot
        eline_north = False
        eline_south = False
        if type(extra_lines) is list:
            if len(extra_lines) >= 1:
                eline_north = extra_lines[0]

                if len(extra_lines) >= 2:
                    eline_south = extra_lines[1]
                else:
                    print "Only one boundary provided, plotting in north"
            else:
                print "No boundaries provided, better to declare as False"

        # Output the observations as a scatter plot
        axn1, mn, axs1, ms = p3g.plot_nsglobal_subfigure(
            f,
            3,
            0,
            lat_data,
            lon_data,
            obs_data,
            obs_name,
            obs_scale,
            obs_units,
            zmax,
            zmin,
            data_color,
            title=True,
            cb=True,
            elat=latlim1,
            tlon=tlon,
            rl=False,
            tl=False,
            bcolor=bcolor,
            earth=earth,
            mn=mn,
            ms=ms,
            faspect=faspect,
            term_datetime=term_datetime,
            extra_line_n=eline_north,
            extra_line_s=eline_south)
        # Output the gitm data as a contour after ensuring that the GITM array
        # isn't padded to include unrealistic latitudes
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1

        axn2, mn, axs2, ms = p3g.plot_nsglobal_subfigure(
            f,
            3,
            1,
            np.array(gdata['dLat'][:, imin:imax, ialt]),
            np.array(gdata['dLon'][:, imin:imax, ialt]),
            np.array(gdata[gitm_key][:, imin:imax, ialt]),
            gitm_name,
            gdata[gitm_key].attrs["scale"],
            gdata[gitm_key].attrs["units"],
            zmax,
            zmin,
            data_color,
            title=False,
            cb=True,
            elat=latlim1,
            tlon=tlon,
            tl=False,
            bcolor=bcolor,
            earth=earth,
            mn=mn,
            ms=ms,
            data_type="contour",
            term_datetime=term_datetime,
            extra_line_n=eline_north,
            extra_line_s=eline_south)

        # Output the differences as a scatter plot
        p3g.plot_nsglobal_subfigure(f,
                                    3,
                                    2,
                                    lat_data,
                                    lon_data,
                                    diff_data,
                                    diff_name,
                                    diff_scale,
                                    diff_units,
                                    diff_max,
                                    diff_min,
                                    diff_color,
                                    title=False,
                                    cb=True,
                                    elat=latlim1,
                                    tlon=tlon,
                                    rl=False,
                                    bcolor=bcolor,
                                    earth=earth,
                                    mn=mn,
                                    ms=ms,
                                    faspect=faspect,
                                    term_datetime=term_datetime,
                                    extra_line_n=eline_north,
                                    extra_line_s=eline_south)
        map_list = list([mn, ms])
    elif plot_type.find("rect") >= 0:
        if len(map_list) == 1:
            m = map_list[0]
        else:
            m = None
        # Output the observations as a scatter plot
        ax = f.add_subplot(3, 1, 1)
        con1, m = p3g.plot_rectangular_3D_global(ax,
                                                 lat_data,
                                                 lon_data,
                                                 obs_data,
                                                 obs_name,
                                                 obs_scale,
                                                 obs_units,
                                                 zmin,
                                                 zmax,
                                                 data_color,
                                                 nlat=latlim1,
                                                 slat=latlim2,
                                                 linc=linc,
                                                 cloc="r",
                                                 xl=False,
                                                 xt=False,
                                                 yl=False,
                                                 meq=meq,
                                                 bcolor=bcolor,
                                                 earth=earth,
                                                 m=m,
                                                 faspect=faspect,
                                                 term_datetime=term_datetime)
        # Output the gitm data as a contour
        ax = f.add_subplot(3, 1, 2)
        con2, m = p3g.plot_rectangular_3D_global(
            ax,
            np.array(gdata['dLat'][:, :, ialt]),
            np.array(gdata['dLon'][:, :, ialt]),
            np.array(gdata[gitm_key][:, :, ialt]),
            gitm_name,
            gdata[gitm_key].attrs["scale"],
            gdata[gitm_key].attrs["units"],
            zmin,
            zmax,
            data_color,
            nlat=latlim1,
            slat=latlim2,
            linc=linc,
            cb=True,
            cloc="r",
            xl=False,
            xt=False,
            bcolor=bcolor,
            meq=meq,
            earth=earth,
            m=m,
            faspect=faspect,
            data_type="contour",
            term_datetime=term_datetime)
        # Adjust plot dimensions if necessary
        if not earth:
            con1_dim = list(con1.axes.get_position().bounds)
            con2_dim = list(con2.ax.get_position().bounds)
            con2_dim[2] = con1_dim[2]
            con2.ax.set_position(con2_dim)

        # Output the differences as a scatter plot
        ax = f.add_subplot(3, 1, 3)
        p3g.plot_rectangular_3D_global(ax,
                                       lat_data,
                                       lon_data,
                                       diff_data,
                                       diff_name,
                                       diff_scale,
                                       diff_units,
                                       diff_min,
                                       diff_max,
                                       diff_color,
                                       nlat=latlim1,
                                       slat=latlim2,
                                       linc=linc,
                                       cloc="r",
                                       yl=False,
                                       bcolor=bcolor,
                                       meq=meq,
                                       earth=earth,
                                       m=m,
                                       faspect=faspect,
                                       term_datetime=term_datetime)
        map_list = list([m])
    elif plot_type.find("polar") >= 0:
        if len(map_list) == 1:
            m = map_list[0]
        else:
            m = None

        pf = True
        if earth:
            pf = False
        # Output the observations as a scatter plot
        ax = f.add_subplot(3, 1, 1, polar=pf)
        con1, m = p3g.plot_polar_3D_global(ax,
                                           3,
                                           lat_data,
                                           lon_data,
                                           obs_data,
                                           obs_name,
                                           obs_scale,
                                           obs_units,
                                           zmin,
                                           zmax,
                                           data_color,
                                           center_lat=latlim1,
                                           edge_lat=latlim2,
                                           linc=linc,
                                           top_lon=tlon,
                                           cloc="r",
                                           tl=False,
                                           rl=False,
                                           bcolor=bcolor,
                                           earth=earth,
                                           m=m,
                                           faspect=faspect,
                                           term_datetime=term_datetime)
        # Output the gitm data as a contour after ensuring that the GITM
        # array isn't padded to include unrealistic latitudes
        ax = f.add_subplot(3, 1, 2, polar=pf)
        (i, imin) = gpr.find_lon_lat_index(gdata, 0.0, -90.0, "degrees")
        (i, imax) = gpr.find_lon_lat_index(gdata, 0.0, 90.0, "degrees")
        imax += 1
        con2, m = p3g.plot_polar_3D_global(ax,
                                           3,
                                           np.array(gdata['dLat'][:, imin:imax,
                                                                  ialt]),
                                           np.array(gdata['dLon'][:, imin:imax,
                                                                  ialt]),
                                           np.array(gdata[gitm_key][:,
                                                                    imin:imax,
                                                                    ialt]),
                                           gitm_name,
                                           gdata[gitm_key].attrs["scale"],
                                           gdata[gitm_key].attrs["units"],
                                           zmin,
                                           zmax,
                                           data_color,
                                           center_lat=latlim1,
                                           edge_lat=latlim2,
                                           linc=linc,
                                           top_lon=tlon,
                                           cb=True,
                                           cloc="r",
                                           tl=False,
                                           bcolor=bcolor,
                                           earth=earth,
                                           m=m,
                                           faspect=faspect,
                                           data_type="contour",
                                           term_datetime=term_datetime)

        con1_dim = list(con1.axes.get_position().bounds)
        con2_dim = list(con2.ax.get_position().bounds)
        con2_dim[0] = con2_dim[0] - 0.05
        con2_dim[2] = con1_dim[2]
        con2.ax.set_position(con2_dim)

        # Output the differences as a scatter plot
        ax = f.add_subplot(3, 1, 3, polar=pf)
        p3g.plot_polar_3D_global(ax,
                                 3,
                                 lat_data,
                                 lon_data,
                                 diff_data,
                                 diff_name,
                                 diff_scale,
                                 diff_units,
                                 diff_min,
                                 diff_max,
                                 diff_color,
                                 center_lat=latlim1,
                                 edge_lat=latlim2,
                                 linc=linc,
                                 top_lon=tlon,
                                 cloc="r",
                                 rl=False,
                                 bcolor=bcolor,
                                 earth=earth,
                                 m=m,
                                 faspect=faspect,
                                 term_datetime=term_datetime)
        map_list = list([m])
    else:
        print rout_name, "ERROR: uknown plot type [", plot_type, "]"
        return

    if title:
        f.suptitle(title, size="medium")

    # Adjust subplot locations
    if plot_type.find("rect") >= 0 or plot_type.find("polar") >= 0:
        plt.subplots_adjust(left=.15)

    # Draw to screen if desired
    if draw:
        if plt.isinteractive():
            plt.draw()  #In interactive mode, you just "draw".
        else:
            # W/o interactive mode, "show" stops the user from typing more
            # at the terminal until plots are drawn.
            plt.show()

    # Save output file
    if figname is not None:
        plt.savefig(figname)

    return (f, map_list)