Skip to content
Snippets Groups Projects
graphics.py 13.08 KiB
#!/usr/bin/env python
# encoding: utf-8
"""
This file was completely rewritten in 2019 to use cartopy objects (instead of basemap)
to plot data on the Earth.

~ Eva Schiffer, Aug. 2019
"""

import logging

from cartopy import crs
import cartopy.feature as cfeature

import numpy
from numpy import argmax, ma

LOG = logging.getLogger(__name__)

def _cap_extent (my_extent, projection_object, ) :
    """
    Make sure my extent doesn't go beyond the available extent for the projection.
    """

    global_bounds = projection_object.boundary.bounds
    g_extent = (global_bounds[0], global_bounds[2], global_bounds[1], global_bounds[3])

    return [max(my_extent[0], g_extent[0]), min(my_extent[1], g_extent[1]), max(my_extent[2], g_extent[2]), min(my_extent[3], g_extent[3]),]

def clean_lon_and_lat (lon_data_object, lat_data_object) :
    """
    Given longitude and latitude data objects, clean up values that can't be valid lons or lats.
    """

    # clean up the longitudes
    if lon_data_object is not None :
        lon_valid_mask = (lon_data_object.data >= -180.0) & (lon_data_object.data <= 360.0)
        lon_data_object.data[~lon_valid_mask] = lon_data_object.fill_value

    # clean up the latitudes
    if lat_data_object is not None :
        lat_valid_mask = (lat_data_object.data >= -90.0) & (lat_data_object.data <= 90.0)
        lat_data_object.data[~lat_valid_mask] = lat_data_object.fill_value

def make_or_find_good_mask(data_array, good_mask=None,) :
    """
    If we don't have a good mask, then try to find or make one.
    """

    # if there is no good mask, find one
    if good_mask is None :

        # check to see if we can get a mask from a masked array
        if data_array is not None and type(data_array) is ma :
            good_mask = data_array.mask

        # if we still don't have a mask, make a simple one
        if good_mask is None :
            good_mask = numpy.ones(data_array.shape, dtype=numpy.bool,)

    return good_mask

def lon_subtract_360(lon1, lon2) :
    """
    return lon1 - lon2, correctly handling the wraparound case
    inputs can be in [-180, 180) form or [0, 360) form
    output is in the [0, 360) range
    """
    return (lon1 - lon2) % 360.

def lon_add_180(lon1, lon2) :
    """
    return lon1 + lon2, correctly handling the wraparound case
    inputs can be in [-180, 180) form or [0, 360) form
    output is in the [-180, 180) range
    """
    return (lon1 + lon2 + 180.) % 360. - 180.

def calculate_lon_mid(lon_left, lon_right) :
    """
    calculate midpoint between two longitudes
    """
    return lon_add_180(lon_left, (lon_subtract_360(lon_right, lon_left) / 2.))

def calculate_lon_range(longitude_data, good_mask, longitude_data_b=None, good_mask_b=None, ) :
    """
    calculate range of longitudes.
    Can handle data that spans the anti-meridian

    returns the westmost and eastmost longitude that define the range, in that order

    Note: This assumes that the ranges of the main data and b data are mostly overlapping.
    This strategy will fail when the two data sets are disparate in longitude, but if that much
    data is missing even though the base lon/lat data matches Glance is going to have some
    problems drawing your image correctly anyway.

    Note: This assumes that data you pass in has a valid good mask.
    """

    # select only the good data values and flatten out our data for ease of use
    good_data = longitude_data[good_mask].flatten()
    if longitude_data_b is not None :
        numpy.append(good_data, longitude_data_b[good_mask_b].flatten())

    # if we have nothing, just stop here
    if good_data.size == 0 :
        return None, None

    # sort the lons, smallest to largest
    good_data.sort()

    # find the largest gap in lons. OK to subtract lons because values are sorted from -180 to 180
    lonDiffs = good_data[1:] - good_data[:-1]
    i = argmax(lonDiffs)
    lonMax = good_data[i]
    lonMin = good_data[i+1]

    LOG.debug("naive gap detected: lonMin, lonMax, gap = " + str(lonMin) + ", " + str(lonMax) + ", " + str(lonMin - lonMax))

    # check for largest gap spanning anti-meridian
    if lon_subtract_360(good_data[0], good_data[-1]) > lon_subtract_360(lonMin, lonMax):
        lonMin, lonMax = good_data[0], good_data[-1]
        LOG.debug("anti-meridian gap detected: lonMin, lonMax, gap = " + str(lonMin) + ", " + str(lonMax) + ", " + str(lonMin - lonMax))

    return lonMin, lonMax

def calculate_lat_range(latitude_data, good_mask, latitude_data_b=None, good_mask_b=None,) :
    """
    Calculate the range of latitudes.
    If a B data set is given, include the range of the latitude_data and the latitude_data_B.

    returns the min latitude and max latitude that define the range in that order

    Note: This assumes that data you pass in has a valid good mask.
    """

    # select only the good data values and flatten out our data for ease of use
    good_data = latitude_data[good_mask].flatten()
    if latitude_data_b is not None:
        numpy.append(good_data, latitude_data_b[good_mask_b].flatten())

    # if we have nothing, just stop here
    if good_data.size == 0:
        return None, None

    return numpy.min(good_data), numpy.max(good_data)

def get_extents(longitude_data,        latitude_data,
                lon_good_mask=None,    lat_good_mask=None,
                longitude_data_b=None, latitude_data_b=None,
                lon_good_mask_b=None,  lat_good_mask_b=None,) :
    """
    Calculate the extents that will include all this data.

    returns extents in the format [lon westmost, lon eastmost, lat northmost, lat southmost]

    Note: The b data is entirely optional, but masks must be included
    for the data provided, eithere as explicit numpy arrays or by making
    the original lat/lon masked arrays instead of simple numpy arrays.
    """

    # if we don't have latitude, assume the longitude is a generator
    if latitude_data is None:
        # then assume (lon,lat) pairs in a generator
        perimeter = list(longitude_data)
        longitude_data = numpy.array([x[0] for x in perimeter])
        latitude_data = numpy.array([x[1] for x in perimeter])

    # if we have the optional b data, also check to see if it has a generator
    if longitude_data_b is not None and latitude_data_b is None:
        # again, assume this is a generator
        perimeter_b = list(longitude_data_b)
        longitude_data_b = numpy.array([x[0] for x in perimeter_b])
        latitude_data_b = numpy.array([x[1] for x in perimeter_b])

    # make sure we aren't processing lon or lat values that are way outside the acceptable range
    lat_valid_mask   = (latitude_data    >=  -90.0) & (latitude_data    <=  90.0)
    lon_valid_mask   = (longitude_data   >= -180.0) & (longitude_data   <= 360.0)
    lat_valid_mask_b = (latitude_data_b  >=  -90.0) & (latitude_data_b  <=  90.0) if latitude_data_b  is not None else None
    lon_valid_mask_b = (longitude_data_b >= -180.0) & (longitude_data_b <= 360.0) if longitude_data_b is not None else None

    # double check that all our good masks exist
    lat_good_mask   = make_or_find_good_mask(latitude_data,    good_mask=lat_good_mask,)
    lon_good_mask   = make_or_find_good_mask(longitude_data,   good_mask=lon_good_mask,)
    lat_good_mask_b = make_or_find_good_mask(latitude_data_b,  good_mask=lat_good_mask_b,) if latitude_data_b  is not None else None
    lon_good_mask_b = make_or_find_good_mask(longitude_data_b, good_mask=lon_good_mask_b,) if longitude_data_b is not None else None

    # make some combined masks
    a_mask = lon_good_mask & lon_valid_mask & lat_good_mask & lat_valid_mask
    b_mask = lon_good_mask_b & lon_valid_mask_b & lat_good_mask_b & lat_valid_mask_b if latitude_data_b is not None else None

    # get the lon and lat bounds for our extents
    lon_westmost,  lon_eastmost  = calculate_lon_range(longitude_data,                    good_mask=a_mask,
                                                       longitude_data_b=longitude_data_b, good_mask_b=b_mask,)
    lat_northmost, lat_southmost = calculate_lat_range(latitude_data,                     good_mask=a_mask,
                                                       latitude_data_b=latitude_data_b,   good_mask_b=b_mask,)

    return [lon_westmost, lon_eastmost, lat_northmost, lat_southmost,]

def create_cartopy_proj_objects (extents, out_proj4_str=None,) :
    """
    Create cartopy "crs" projection objects for the input data and output projections.

    This is only a first draft of how Glance will probably use projections in the future.
    For now, we will assume that all data comes in with lat/lon and therefore
    corresponds to a PlateCarree crs object.

    The format of the extents is [lon westmost, lon eastmost, lat northmost, lat southmost]
    where these extents define the whole area you wish to show.

    If the out projection string is None we will be using PlateCarree as our output projection.
    TODO, other projections are not currently supported, but should be someday.

    Note: You will need to call get_extents on your own now. And the extents will be in lat/lon,
    so you will need to use the input_proj_obj when you're setting your plot extents later.
    """

    LOG.debug("Using extents: " + str(extents))
    if out_proj4_str is not None :
        LOG.debug("Input proj4 string provided, but this code cannot use it yet: " + out_proj4_str)

    # FUTURE, this is not actually useful here now, but hang on to it for bounds setting later in our display process
    # # pull out the longitude/latitude info
    # lon_left   = extents[0]
    # lon_right  = extents[1]
    # lat_bottom = extents[2]
    # lat_top    = extents[3]
    # lon_mid = calculate_lon_mid(lon_left, lon_right)  # handle lon wraparound case
    # lat_mid = (lat_top + lat_bottom) / 2.

    # for now just use PlateCarree, in
    # FUTURE we will want to pick projections based on our input so we could use x and y when desired
    input_proj_obj = crs.PlateCarree()

    # FUTURE, we will probably want different projections; select based on the lat/lon range?
    # for reference, here's what cartopy supports: https://scitools.org.uk/cartopy/docs/v0.15/crs/projections.html#cartopy-projections
    output_proj_obj = crs.PlateCarree()

    return input_proj_obj, output_proj_obj

def draw_basic_features(in_proj_obj, axes_object, extents,) :
    """Draw the basic outlines of the earth's features.
    """

    # set the extents (Note: we may have problems setting extents while plotting full disks TODO test that)
    axes_object.set_extent(extents, crs=in_proj_obj)

    # draw some features on our map
    line_color = 'black'    # Glance always uses white background and black lines
    # draw some lat/lon gridlines
    axes_object.gridlines(linestyle=':', zorder=2, edgecolor=line_color, )
    # TODO, do we need to add anything to label the gridlines?
    # draw physical and geopolitical stuff
    axes_object.coastlines(resolution='50m', zorder=3, color=line_color, )
    axes_object.add_feature(cfeature.BORDERS, linewidth=0.5, zorder=4, edgecolor=line_color, )  # this makes country borders
    axes_object.add_feature(cfeature.LAKES, linestyle=':', linewidth=0.5, facecolor='none', edgecolor=line_color, zorder=5, )
    # add states and provinces
    # Create a feature for States/Admin 1 regions at 1:50m from Natural Earth
    states_provinces = cfeature.NaturalEarthFeature(
        category='cultural',
        name='admin_1_states_provinces_lines',
        scale='50m',
        facecolor='none')
    axes_object.add_feature(states_provinces, edgecolor=line_color, linewidth=0.25, zorder=6, )
    
    return

def contourf_on_map(lon, lat, data, axes_object, in_projection, levelsToUse=None, **kwargs ) :
    """
    Contour plot the data on a map. Assume the input coordinates are in the in_projection form
    and the output should be plotted in out_projection.

    optional: levelsToUse is a list of numbers representing data ranges that will be used
    """

    if levelsToUse is not None:
        img_temp = axes_object.contourf(lon, lat, data, levelsToUse, transform=in_projection, **kwargs)
    else:
        img_temp = axes_object.contourf(lon, lat, data,              transform=in_projection, **kwargs)

    return img_temp

def quiver_plot_on_map(lon, lat, axes_object, in_projection, uData=None, vData=None, colordata=None, **kwargs) :
    """
    Show a quiver plot of the given vector data at the given longitude and latitude
    """

    # show the quiver plot if there is data

    if (uData is not None) and (vData is not None):
        if colordata is None:
            img_temp = axes_object.quiver(lon, lat, uData, vData,            transform=in_projection, **kwargs)
        else:
            img_temp = axes_object.quiver(lon, lat, uData, vData, colordata, transform=in_projection, **kwargs)

    return img_temp

def plot_on_map(lon, lat, axes_object, in_projection, marker_const='o', color_const='#000000', marker_size=1, **kwargs) :
    """
    Show a simple point plot at the given longitude and latitude.
    """

    # show the plot if there is data

    if (lon.size > 0) and (lat.size > 0) and (lon.size == lat.size) :

        img_temp = axes_object.plot(lon, lat, marker_const, color=color_const, markersize=marker_size, transform=in_projection, **kwargs)

    return img_temp