In [9]:
import numpy
import matplotlib.pyplot as plt

import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter

In [102]:
data = numpy.load('/g/data/r87/dbi599/basins.npy')

In [103]:
data.shape


Out[103]:
(144, 58, 130, 360)

In [104]:
lats = numpy.array([-64.5, -63.5, -62.5, -61.5, -60.5, -59.5, -58.5, -57.5, -56.5,
       -55.5, -54.5, -53.5, -52.5, -51.5, -50.5, -49.5, -48.5, -47.5,
       -46.5, -45.5, -44.5, -43.5, -42.5, -41.5, -40.5, -39.5, -38.5,
       -37.5, -36.5, -35.5, -34.5, -33.5, -32.5, -31.5, -30.5, -29.5,
       -28.5, -27.5, -26.5, -25.5, -24.5, -23.5, -22.5, -21.5, -20.5,
       -19.5, -18.5, -17.5, -16.5, -15.5, -14.5, -13.5, -12.5, -11.5,
       -10.5,  -9.5,  -8.5,  -7.5,  -6.5,  -5.5,  -4.5,  -3.5,  -2.5,
        -1.5,  -0.5,   0.5,   1.5,   2.5,   3.5,   4.5,   5.5,   6.5,
         7.5,   8.5,   9.5,  10.5,  11.5,  12.5,  13.5,  14.5,  15.5,
        16.5,  17.5,  18.5,  19.5,  20.5,  21.5,  22.5,  23.5,  24.5,
        25.5,  26.5,  27.5,  28.5,  29.5,  30.5,  31.5,  32.5,  33.5,
        34.5,  35.5,  36.5,  37.5,  38.5,  39.5,  40.5,  41.5,  42.5,
        43.5,  44.5,  45.5,  46.5,  47.5,  48.5,  49.5,  50.5,  51.5,
        52.5,  53.5,  54.5,  55.5,  56.5,  57.5,  58.5,  59.5,  60.5,
        61.5,  62.5,  63.5,  64.5])

In [105]:
lons = numpy.array([  20.5,   21.5,   22.5,   23.5,   24.5,   25.5,   26.5,   27.5,
         28.5,   29.5,   30.5,   31.5,   32.5,   33.5,   34.5,   35.5,
         36.5,   37.5,   38.5,   39.5,   40.5,   41.5,   42.5,   43.5,
         44.5,   45.5,   46.5,   47.5,   48.5,   49.5,   50.5,   51.5,
         52.5,   53.5,   54.5,   55.5,   56.5,   57.5,   58.5,   59.5,
         60.5,   61.5,   62.5,   63.5,   64.5,   65.5,   66.5,   67.5,
         68.5,   69.5,   70.5,   71.5,   72.5,   73.5,   74.5,   75.5,
         76.5,   77.5,   78.5,   79.5,   80.5,   81.5,   82.5,   83.5,
         84.5,   85.5,   86.5,   87.5,   88.5,   89.5,   90.5,   91.5,
         92.5,   93.5,   94.5,   95.5,   96.5,   97.5,   98.5,   99.5,
        100.5,  101.5,  102.5,  103.5,  104.5,  105.5,  106.5,  107.5,
        108.5,  109.5,  110.5,  111.5,  112.5,  113.5,  114.5,  115.5,
        116.5,  117.5,  118.5,  119.5,  120.5,  121.5,  122.5,  123.5,
        124.5,  125.5,  126.5,  127.5,  128.5,  129.5,  130.5,  131.5,
        132.5,  133.5,  134.5,  135.5,  136.5,  137.5,  138.5,  139.5,
        140.5,  141.5,  142.5,  143.5,  144.5,  145.5,  146.5,  147.5,
        148.5,  149.5,  150.5,  151.5,  152.5,  153.5,  154.5,  155.5,
        156.5,  157.5,  158.5,  159.5,  160.5,  161.5,  162.5,  163.5,
        164.5,  165.5,  166.5,  167.5,  168.5,  169.5,  170.5,  171.5,
        172.5,  173.5,  174.5,  175.5,  176.5,  177.5,  178.5,  179.5,
        180.5,  181.5,  182.5,  183.5,  184.5,  185.5,  186.5,  187.5,
        188.5,  189.5,  190.5,  191.5,  192.5,  193.5,  194.5,  195.5,
        196.5,  197.5,  198.5,  199.5,  200.5,  201.5,  202.5,  203.5,
        204.5,  205.5,  206.5,  207.5,  208.5,  209.5,  210.5,  211.5,
        212.5,  213.5,  214.5,  215.5,  216.5,  217.5,  218.5,  219.5,
        220.5,  221.5,  222.5,  223.5,  224.5,  225.5,  226.5,  227.5,
        228.5,  229.5,  230.5,  231.5,  232.5,  233.5,  234.5,  235.5,
        236.5,  237.5,  238.5,  239.5,  240.5,  241.5,  242.5,  243.5,
        244.5,  245.5,  246.5,  247.5,  248.5,  249.5,  250.5,  251.5,
        252.5,  253.5,  254.5,  255.5,  256.5,  257.5,  258.5,  259.5,
        260.5,  261.5,  262.5,  263.5,  264.5,  265.5,  266.5,  267.5,
        268.5,  269.5,  270.5,  271.5,  272.5,  273.5,  274.5,  275.5,
        276.5,  277.5,  278.5,  279.5,  280.5,  281.5,  282.5,  283.5,
        284.5,  285.5,  286.5,  287.5,  288.5,  289.5,  290.5,  291.5,
        292.5,  293.5,  294.5,  295.5,  296.5,  297.5,  298.5,  299.5,
        300.5,  301.5,  302.5,  303.5,  304.5,  305.5,  306.5,  307.5,
        308.5,  309.5,  310.5,  311.5,  312.5,  313.5,  314.5,  315.5,
        316.5,  317.5,  318.5,  319.5,  320.5,  321.5,  322.5,  323.5,
        324.5,  325.5,  326.5,  327.5,  328.5,  329.5,  330.5,  331.5,
        332.5,  333.5,  334.5,  335.5,  336.5,  337.5,  338.5,  339.5,
        340.5,  341.5,  342.5,  343.5,  344.5,  345.5,  346.5,  347.5,
        348.5,  349.5,  350.5,  351.5,  352.5,  353.5,  354.5,  355.5,
        356.5,  357.5,  358.5,  359.5,  360.5,  361.5,  362.5,  363.5,
        364.5,  365.5,  366.5,  367.5,  368.5,  369.5,  370.5,  371.5,
        372.5,  373.5,  374.5,  375.5,  376.5,  377.5,  378.5,  379.5])

In [6]:
%matplotlib inline

In [106]:
plt.imshow(data[10,10,:,:])


Out[106]:
<matplotlib.image.AxesImage at 0x7fc15dcd3d90>

In [20]:
ax = plt.subplot(111, projection=ccrs.PlateCarree())
cf = ax.contourf(lons, lats, data[5,5,:,:])#, transform=ccrs.PlateCarree())

ax.coastlines()

#    ax.set_yticks([-80, -60, -40, -20, 0, 20, 40, 60, 80], crs=ccrs.PlateCarree())
#    ax.set_xticks([0, 60, 120, 180, 240, 300, 360], crs=ccrs.PlateCarree())
#    lon_formatter = LongitudeFormatter(zero_direction_label=True)
#    lat_formatter = LatitudeFormatter()
#    ax.xaxis.set_major_formatter(lon_formatter)
#    ax.yaxis.set_major_formatter(lat_formatter)
#    ax.set_xlabel('Longitude', fontsize='small')
#    ax.set_ylabel('Latitude', fontsize='small')


Out[20]:
<cartopy.mpl.feature_artist.FeatureArtist at 0x7fc15e313b50>

In [94]:
def single2list(item, numpy_array=False):
    """Check if item is a list, then convert if not."""
    
    if type(item) == list or type(item) == tuple or type(item) == numpy.ndarray:
        output = item 
    elif type(item) == str:
        output = [item,]
    else:
        try:
            test = len(item)
        except TypeError:
            output = [item,]

    if numpy_array and not isinstance(output, numpy.ndarray):
        return numpy.array(output)
    else:
        return output


def adjust_lon_range(lons, radians=True, start=0.0):
    """Express longitude values in a 360 degree (or 2*pi radians) interval.
    Args:
      lons (list/tuple): Longitude axis values (monotonically increasing)
      radians (bool): Specify whether input data are in radians (True) or
        degrees (False). Output will be the same units.
      start (float, optional): Start value for the output interval (add 360 degrees or 2*pi
        radians to get the end point)
    
    """
    
    lons = single2list(lons, numpy_array=True)    
    
    interval360 = 2.0*numpy.pi if radians else 360.0
    end = start + interval360    
    
    less_than_start = numpy.ones([len(lons),])
    while numpy.sum(less_than_start) != 0:
        lons = numpy.where(lons < start, lons + interval360, lons)
        less_than_start = lons < start
    
    more_than_end = numpy.ones([len(lons),])
    while numpy.sum(more_than_end) != 0:
        lons = numpy.where(lons >= end, lons - interval360, lons)
        more_than_end = lons >= end

    return lons


def broadcast_array(array, axis_index, shape):
    """Broadcast a one dimensional array to a target shape.
    
    Args:
      array (numpy.ndarray): One dimensional array
      axis_index (int): Postion in the target shape that the array
        corresponds to (e.g. if array corresponds to lat in (time, depth
        lat, lon) array then index = 2
      shape (tuple): shape to broadcast to
    
    """

    dim = axis_index - 1
    while dim >= 0:
        array = array[numpy.newaxis, ...]
        array = numpy.repeat(array, shape[dim], axis=0)
        dim = dim - 1
    
    dim = axis_index + 1
    while dim < len(shape):    
        array = array[..., numpy.newaxis]
        array = numpy.repeat(array, shape[dim], axis=-1)
        dim = dim + 1

    return array


def create_basin_array(lat_axis, lon_axis):
    """Create a basin file.

    For similarity with the CMIP5 basin file, in the output:
      Atlantic Ocean = 2
      Pacific Ocean = 3
      Indian Ocean = 5

    """

    shape = (len(lat_axis), len(lon_axis))
    
    pacific_bounds = [147, 294]
    indian_bounds = [23, 147]

    lon_axis = adjust_lon_range(lon_axis, radians=False)
    
    lat_array = broadcast_array(lat_axis, 0, shape)
    lon_array = broadcast_array(lon_axis, 1, shape)

    basin_array = numpy.ones(shape) * 2
    basin_array = numpy.where((lon_array >= pacific_bounds[0]) & (lon_array <= pacific_bounds[1]), 3, basin_array)
    basin_array = numpy.where((lon_array >= indian_bounds[0]) & (lon_array <= indian_bounds[1]), 5, basin_array)

    basin_array = numpy.where((basin_array == 3) & (lon_array >= 279) & (lat_array >= 10), 2, basin_array)
    basin_array = numpy.where((basin_array == 5) & (lon_array >= 121) & (lat_array >= 0), 3, basin_array)

    return basin_array, lat_array, lon_array

In [95]:
conventional_lats = numpy.arange(-65,66,1)
conventional_lons = numpy.arange(0,360,1)

In [99]:
basin_array, lat_array, lon_array = create_basin_array(lats, lons)

In [100]:
ax = plt.subplot(111, projection=ccrs.PlateCarree())
cf = ax.contourf(lons, lats, basin_array)#, transform=ccrs.PlateCarree())

ax.coastlines()
cbar = plt.colorbar(cf)



In [101]:
plt.imshow(basin_array)


Out[101]:
<matplotlib.image.AxesImage at 0x7fc15db7cf50>

In [ ]: