A script for plotting Bedmap2 Antarctic bathymetry and ice topography

create_env_and_plot.bash

#!/usr/bin/env bash

conda create -y -n plot_topo_env -c conda-forge python=3.9 cartopy matplotlib xarray numpy progressbar2 requests cmocean
conda activate plot_topo_env
./plot_bedmap2_topo.py

plot_bedmap2_topo.py

#!/usr/bin/env python
import xarray as xr
import numpy as np
import os
import requests
import zipfile
import progressbar
import cartopy

import matplotlib.pyplot as plt
from matplotlib import colors, path, ticker
from mpl_toolkits.axes_grid1 import make_axes_locatable
from mpl_toolkits.axes_grid1.inset_locator import zoomed_inset_axes, mark_inset


def bedmap2_bin_to_netcdf():
    """
    Download Bedmap2 data to a directory called ``bedmap2`` and convert it to
    NetCDF format (if it isn't already there).
    """
    # make a directory for the output if it doesn't exist
    try:
        os.makedirs('bedmap2')
    except OSError:
        pass

    out_file_name = 'bedmap2/bedmap2.nc'

    if os.path.exists(out_file_name):
        return

    fields = ['bed', 'surface', 'thickness', 'coverage', 'rockmask',
              'grounded_bed_uncertainty', 'icemask_grounded_and_shelves']

    all_exist = True
    for field in fields:
        file_name = f'bedmap2/bedmap2_bin/bedmap2_{field}.flt'
        if not os.path.exists(file_name):
            all_exist = False
            break

    if not all_exist:
        # download
        base_url = 'https://secure.antarctica.ac.uk/data/bedmap2'
        file_names = ['bedmap2_bin.zip']

        download_files(file_names, base_url, 'bedmap2')

        print('Decompressing Bedmap2 data...')
        # unzip
        with zipfile.ZipFile('bedmap2/bedmap2_bin.zip', 'r') as f:
            f.extractall('bedmap2/')
        print('  Done.')

    print('Converting Bedmap2 to NetCDF...')
    ds = xr.Dataset()
    x = np.linspace(-3333000., 3333000., 6667)
    y = x
    ds['x'] = ('x', x)
    ds.x.attrs['units'] = 'meters'
    ds['y'] = ('y', y)
    ds.y.attrs['units'] = 'meters'
    ds.attrs['Grid'] = "Datum = WGS84, earth_radius = 6378137., " \
                       "earth_eccentricity = 0.081819190842621, " \
                       "falseeasting = -3333000., " \
                       "falsenorthing = -3333000., " \
                       "standard_parallel = -71., central_meridien = 0, " \
                       "EPSG=3031"
    ds.attrs['proj'] = "+proj=stere +lat_0=-90 +lat_ts=-71 +lon_0=0 +k=1 " \
                       "+x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
    ds.attrs['proj4'] = "+init=epsg:3031"

    # add Bedmap2 data
    for fieldName in fields:
        file_name = 'bedmap2/bedmap2_bin/bedmap2_{}.flt'.format(fieldName)
        with open(file_name, 'r') as f:
            field = np.fromfile(f, dtype=np.float32).reshape(6667, 6667)
            # flip the y axis
            field = field[::-1, :]
            # switch invalid values to be NaN (as expected by xarray)
            field[field == -9999.] = np.nan
        if fieldName == 'rockmask':
            # rock mask is zero where rock and -9999 (now NaN) elsewhere
            field = np.array(np.isfinite(field), np.float32)
        if fieldName == 'icemask_grounded_and_shelves':
            # split into separate grounded and floating masks
            ds['icemask_grounded'] = \
                (('y', 'x'), np.array(field == 0, np.float32))
            ds['icemask_shelves'] = \
                (('y', 'x'), np.array(field == 1, np.float32))
            ds['open_ocean_mask'] = \
                (('y', 'x'), np.array(np.isnan(field), np.float32))
        else:
            ds[fieldName] = (('y', 'x'), field)

    ds.to_netcdf(out_file_name)
    print('  Done.')


def download_files(file_list, url_base, out_dir):
    """
    Download a list of files from a URL to a directory

    Parameters
    ----------
    file_list : list of str
        A list of file to download

    url_base : str
        The base URL where the files can be found

    out_dir : str
        The directory where the files should be written
    """

    for file_name in file_list:
        out_file_name = os.path.join(out_dir, file_name)
        # out_file_name contains full path, so we need to make the relevant
        # subdirectories if they do not exist already
        directory = os.path.dirname(out_file_name)
        try:
            os.makedirs(directory)
        except OSError:
            pass

        url = f'{url_base}/{file_name}'
        try:
            response = requests.get(url, stream=True)
            encoding = response.headers.get('content-encoding')
            total_size = response.headers.get('content-length')

        except requests.exceptions.RequestException:
            print('  {} could not be reached!'.format(file_name))
            continue

        try:
            response.raise_for_status()
        except requests.exceptions.HTTPError as e:
            print('ERROR while downloading {}:'.format(file_name))
            print(e)
            continue

        if total_size is None or encoding == 'gzip':
            # no content length header, or not related unzipped size
            if not os.path.exists(out_file_name):
                with open(out_file_name, 'wb') as f:
                    print(f'Downloading {file_name}...')
                    try:
                        f.write(response.content)
                    except requests.exceptions.RequestException:
                        print(f'  {file_name} failed!')
                    else:
                        print(f'  {file_name} done.')
        else:
            # we can do the download in chunks and use a progress bar, yay!

            total_size = int(total_size)
            if os.path.exists(out_file_name) and \
                    total_size == os.path.getsize(out_file_name):
                # we already have the file, so just continue
                continue

            file_size = sizeof_fmt(total_size)
            print(f'Downloading {file_name} ({file_size})...')
            widgets = [progressbar.Percentage(), ' ', progressbar.Bar(),
                       ' ', progressbar.ETA()]
            bar = progressbar.ProgressBar(widgets=widgets,
                                          maxval=total_size).start()
            size = 0
            with open(out_file_name, 'wb') as f:
                try:
                    for data in response.iter_content(chunk_size=4096):
                        size += len(data)
                        f.write(data)
                        bar.update(size)
                    bar.finish()
                except requests.exceptions.RequestException:
                    print(f'  {file_name} failed!')
                else:
                    print(f'  {file_name} done.')


# From https://stackoverflow.com/a/1094933/7728169
def sizeof_fmt(num, suffix='B'):
    """
    Covert a number of bytes to a human-readable file size
    """
    for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
        if abs(num) < 1024.0:
            return f'{num:3.1f}{unit}{suffix}'
        num /= 1024.0
    return f'{num:.1f}Yi{suffix}'


def plot_panel(ax, ds, is_inset, labels=None, max_depth=-1000.):
    """
    Make a plot of Bedmap2, possibly with labels

    Parameters
    ----------
    ax : matplotlib.axes.Axes
        The matplotlib axis to plot on
    ds : xarray.Dataset
        The Bedmap2 data set to plot
    is_inset : bool
        Whether this is an inset
    labels : list of dict, optional
        A list of dictionaries used to make labels
    max_depth : float, optional
        The maximum depth of the bathymetry colormap
    """

    # get the x and y arrays
    x = ds.x.values
    y = ds.y.values

    if not is_inset:
        # plot grid lines and labels
        gl = ax.gridlines(crs=cartopy.crs.PlateCarree(), color='k',
                          linestyle=':', zorder=6, draw_labels=True)
        # every 30 degrees in longitude
        gl.xlocator = ticker.FixedLocator(np.arange(-180., 181., 30.))
        # every 10 degrees in latitude
        gl.ylocator = ticker.FixedLocator(np.arange(-80., 81., 10.))
        # a lot of steps so the circles look smooth
        gl.n_steps = 100
        # labels on the right interfere with the colorbar
        gl.right_labels = False
        # labels on the left are covered up by the insets
        gl.left_labels = False
        # format the longitude and latitude labels (degrees E/W or N/S)
        gl.xformatter = cartopy.mpl.gridliner.LONGITUDE_FORMATTER
        gl.yformatter = cartopy.mpl.gridliner.LATITUDE_FORMATTER
        # the font size
        gl.xlabel_style['size'] = 12
        gl.ylabel_style['size'] = 12
        # don't rotate the labels to align with the axis (this usually looks
        # bad)
        gl.rotate_labels = False

    # make a light source to shade based on the ice surface elevation
    ls = colors.LightSource(azdeg=135, altdeg=45)

    # Bedmap2 has invalid values north of 60 S -- make these the max depth.
    # This could be handled differently (e.g. as a masked array) or ignored
    # as long as these parts of the domain aren't plotted.
    bed = (ds.bed.where(ds.bed.notnull(), other=max_depth)).values

    # A colormap for the topography (blue at depth to white at sea level in
    # this case)
    cmap = plt.cm.Blues_r

    # An image with the topography shown in blue colors and with shadowing
    # based on the gradient of the topography
    bed_color = ls.shade(bed, cmap=cmap, vert_exag=0.005, vmin=max_depth,
                         vmax=0., blend_mode='soft')

    # An image showing the ice surface in gray
    ice_shade = 0.8 + 0.2*ls.hillshade(ds.surface.values, vert_exag=.01)
    # Mask out the ice where there is none
    ice_shade = np.ma.masked_array(ice_shade, mask=np.isnan(ice_shade))

    # An array of ones where there are ice shelves and masked out where there
    # are none
    shelves = ds.icemask_shelves.values
    shelves = np.ma.masked_array(shelves, mask=(shelves == 0))

    if not is_inset:
        # plot the topography once without shading to use as the basis for the
        # colorbar
        plot_handle = ax.pcolormesh(x, y, bed, zorder=1, cmap=cmap,
                                    vmin=max_depth, vmax=0)

        # create an axes on the right side of ax. The width of cax will be 5%
        # of ax and the padding between cax and ax will be fixed at 0.05 inch.
        divider = make_axes_locatable(ax)
        cax = divider.append_axes("right", size="5%", pad=0.05,
                                  axes_class=plt.Axes)
        cbar = plt.colorbar(plot_handle, cax=cax)
        cbar.set_label('depth (m)')

    # Plot the bed again, this time as an image with the shading.
    # Flip it over on the y axis (the first axis) because imshow references
    # the image from the bottom, not the top.  In python, an index of ::-1
    # means all indices but with a step of -1 (so flipped).
    ax.imshow(bed_color[::-1, :], zorder=2, extent=[x[0], x[-1], y[0],y[-1]])

    # plot the ice shading on top of the topography
    ax.pcolormesh(x, y, ice_shade, zorder=3, cmap='gray', vmin=0., vmax=1.,
                  rasterized=True, shading='nearest')

    # plot the shelves in gray on top of the topography and ice shading
    ax.pcolormesh(x, y, 0.6*shelves, cmap='gray', vmin=0., vmax=1.,
                  zorder=4, rasterized=True, shading='nearest')

    if is_inset:
        # turn off tick marks if this isn't an inset
        plt.tick_params(
            axis='both',
            which='both',
            bottom=False,
            top=False,
            left=False,
            right=False,
            labelbottom=False,
            labelleft=False)

    if not is_inset:
        # From https://scitools.org.uk/cartopy/docs/v0.15/examples/always_circular_stereo.html
        # Make a circular boundary for the plot so we hide the missing data in
        # Bedmap2.  However, this seems to get rid of the lon/lat annotations.
        theta = np.linspace(0, 2*np.pi, 100)
        center = np.array([0.5, 0.5])
        radius = 0.5
        verts = np.vstack([np.sin(theta), np.cos(theta)]).T
        circle = path.Path(verts * radius + center)

        ax.set_boundary(circle, transform=ax.transAxes)

    for label in labels:
        # add the annotation text, possibly with arrows.
        if 'color' in label.keys():
            color = label['color']
        else:
            color = 'k'
        if 'arrow' in label.keys():
            # this is a hack to subtract 3000 km from the locations I
            # originally defined (using a coordinate that went from 0 to
            # 6000 km, rather than -3333 km to 3333 km in Bedmap2)
            arrow_loc = [loc - 3e6 for loc in label['arrow']]
            text_loc = [loc - 3e6 for loc in label['text']]
            ax.annotate(label['label'], xy=arrow_loc,
                        xytext=text_loc, zorder=5, color=color,
                        ha="center", va="center",
                        arrowprops=dict(arrowstyle="->"), fontsize=10)
        else:
            # this is a hack to subtract 3000 km from the locations I
            # originally defined (using a coordinate that went from 0 to
            # 6000 km, rather than -3333 km to 3333 km in Bedmap2)
            loc = [loc - 3e6 for loc in label['text']]
            ax.text(loc[0], loc[1], label['label'], zorder=5, color=color,
                    ha="center", va="center", fontsize=10)


def plot_bedmap2():
    # open the Bedmap2 data set using xarray
    ds = xr.open_dataset('bedmap2/bedmap2.nc')

    # We reducing the resolution of the data to 4 km because the full
    # resolution os overkill and creates a really big file.  Change 4 to 8 or
    # 16 for faster testing.
    ds = ds.isel(x=slice(0, -1, 8), y=slice(0, -1, 8))

    # These dictionaries define the positions of labels and (optionally) arrows
    # and font colors
    labels = [{'label': 'Ross Sea',
               'text': (2.31e6, 1.27e6),
               'color': 'w'},
              {'label': 'Ross IS',
               'arrow': (2.95e6, 2.02e6),
               'text': (2.98e6, 2.62e6)},
              {'label': 'Mertz IS',
               'arrow': (4.43e6, 0.93e6),
               'text': (4.17e6, 1.65e6)},
              {'label': 'George V\nLand',
               'text': (3.83e6, 1.35e6)},
              {'label': r'${\rm Ad\'elie}$' '\nLand',
               'text': (4.73e6, 1.34e6)},
              {'label': 'Sabrina\nCoast',
               'text': (5.63e6, 1.5e6),
               'color': 'w'},
              {'label': 'Dalton IS',
               'arrow': (5.18e6, 1.66e6),
               'text': (4.5e6, 1.86e6)},
              {'label': 'Totten IS',
               'arrow': (5.29e6, 1.87e6),
               'text': (4.64e6, 2.28e6)},
              {'label': 'Amery IS',
               'arrow': (5.02e6, 3.7e6),
               'text': (4.4e6, 3.44e6)},
              {'label': 'Roi Baudouin IS',
               'arrow': (3.92e6, 4.92e6),
               'text': (3.9e6, 4.48e6)},
              {'label': 'Fimbul IS',
               'arrow': (2.99e6, 5.11e6),
               'text': (3.07e6, 4.75e6)},
              {'label': 'Filchner IS',
               'arrow': (2.33e6, 3.76e6),
               'text': (2.79e6, 3.45e6)},
              {'label': 'Ronne IS',
               'arrow': (1.89e6, 3.65e6),
               'text': (2.21e6, 3.05e6)},
              {'label': 'Weddell Sea',
               'text': (1.48e6, 4.8e6),
               'color': 'w'}]

    # These dictionaries define the positions of labels and (optionally) arrows
    # and font colors as well as the location of the inset itself for the
    # Amundsen Sea inset
    dict_ase = {'labels': [{'label': 'Amundsen\nSea',
                            'text': (1.21e6, 2.37e6)},
                           {'label': 'Abbot IS',
                            'arrow': (1.15e6, 2.73e6),
                            'text': (1.30e6, 2.84e6)},
                           {'label': 'Pine Island IS',
                            'arrow': (1.41e6, 2.69e6),
                            'text': (1.58e6, 2.77e6)},
                           {'label': 'Thwaites IS',
                            'arrow': (1.43e6, 2.55e6),
                            'text': (1.61e6, 2.63e6)},
                           {'label': 'Crosson',
                            'arrow': (1.49e6, 2.43e6),
                            'text': (1.66e6, 2.53e6)},
                           {'label': 'IS',
                            'text': (1.66e6, 2.46e6)},
                           {'label': 'Dotson IS',
                            'arrow': (1.45e6, 2.33e6),
                            'text': (1.64e6, 2.23e6)},
                           {'label': 'Getz',
                            'arrow': (1.59e6, 2.04e6),
                            'text': (1.69e6, 2.14e6)},
                           {'label': 'IS',
                            'text': (1.69e6, 2.07e6)},
                           {'label': 'Dotson\nTrough',
                            'arrow': (1.34e6, 2.21e6),
                            'text': (1.23e6, 1.96e6),
                            'color': 'w'}],
                'bounds': [1.0e6, 1.8e6, 1.7e6, 2.9e6],
                'bbox': (10., 90.),
                'loc': 3,
                'loc1': 2,
                'loc2': 4}

    # These dictionaries define the positions of labels and (optionally) arrows
    # and font colors as well as the location of the inset itself for the
    # Bellingshausen Sea inset
    dict_bse = {'labels': [{'label': 'Bellings.\nSea',
                            'text': (0.93e6, 3.3e6)},
                           {'label': 'Marguerite\nTrough',
                            'arrow': (0.81e6, 3.8e6),
                            'text': (0.91e6, 4.0e6)},
                           {'label': 'Belgica\nTrough',
                            'arrow': (0.98e6, 3.24e6),
                            'text': (1e6, 3.05e6)}],
                'bounds': [0.6e6, 1.3e6, 2.9e6, 3.9e6],
                'bbox': (10., 350.),
                'loc': 3,
                'loc1': 1,
                'loc2': 3}

    # This is the Bedmap2 projection
    projection = cartopy.crs.Stereographic(
        central_latitude=-90., central_longitude=0.0,
        true_scale_latitude=-71.0)

    # Start with a square figure 8 inches x 8 inches with a resolution of
    # 200 dots per inch (matplotlib, get on board with the metric system!)
    plt.figure(figsize=(8, 8), dpi=200)
    # get a set of axes to plot the main figure onto
    ax = plt.subplot(111, projection=projection)

    # make a little extra room for the insets on the left-hand side, keeping
    # the same aspect ratio and centering vertically

    # get the original position
    pos1 = ax.get_position()
    x_offset = 0.05
    x0 = pos1.x0 + x_offset
    width = pos1.width-x_offset
    scale = width/pos1.width
    height = scale*pos1.height
    y0 = pos1.y0 + 0.5*(pos1.height - height)
    pos2 = [x0, y0, width, height]
    # set a new position
    ax.set_position(pos2)

    x = ds.x.values
    y = ds.y.values

    # set the extent of the figure to be the min and max of x and y in the
    # Bedmap2 data.  You could zoom in by setting the extent to something
    # smaller.
    extent = [x[0], x[-1], y[0], y[-1]]
    ax.set_extent(extent, crs=projection)

    # plot the main figure.  We will use the same function to plot the insets
    plot_panel(ax=ax, ds=ds, is_inset=False, labels=labels)

    for dict_inset in [dict_ase, dict_bse]:
        # define a set of axes for the inset.  The location "loc" is an integer
        # that defines a position (top left, bottom left, etc.) as described
        # in the matplotlib documentation
        # https://matplotlib.org/stable/api/_as_gen/mpl_toolkits.axes_grid1.inset_locator.zoomed_inset_axes.html
        # "zoom" is how many times the inset is scaled up from the original
        # size.  "bbox_to_anchor" is the (left, bottom) location (in pixels)
        # of the inset.
        ax_inset = zoomed_inset_axes(ax, zoom=2.5, loc=dict_inset['loc'],
                                     bbox_to_anchor=dict_inset['bbox'])

        # plot the inset
        plot_panel(ax=ax_inset, ds=ds, is_inset=True,
                   labels=dict_inset['labels'])

        # this is a hack to subtract 3000 km from the locations I originally
        # defined (using a coordinate that went from 0 to 6000 km, rather than
        # -3333 km to 3333 km in Bedmap2)
        bounds_inset = [loc - 3e6 for loc in dict_inset['bounds']]

        # set the limits of the inset in Bedmap2 coordinates (meters)
        ax_inset.set_xlim(bounds_inset[0], bounds_inset[1])
        ax_inset.set_ylim(bounds_inset[2], bounds_inset[3])

        # make a box for the inset on the original figure and add lines
        # this box to the inset itself
        mark_inset(ax, ax_inset, loc1=dict_inset['loc1'],
                   loc2=dict_inset['loc2'], fc="none", ec='k', zorder=4)

    # save the result to a high resolution PDF
    plt.savefig('topo_complicated.pdf', dpi=200)

    # uncomment this to also display the result to the screen
    # plt.show()


def main():
    # download Bedmap2 and convert to NetCDF if this hasn't been done already
    bedmap2_bin_to_netcdf()

    # plot the figure
    plot_bedmap2()


if __name__ == '__main__':
    main()