Skip to content

Instantly share code, notes, and snippets.

@xylar

xylar/create_env_and_plot.bash

Created September 15, 2021 13:28
Show Gist options
  • Save xylar/1fc958ad8a5cd0dbabc5b13caae8ad3e to your computer and use it in GitHub Desktop.
Save xylar/1fc958ad8a5cd0dbabc5b13caae8ad3e to your computer and use it in GitHub Desktop.
Download ZIP
A script for plotting Bedmap2 Antarctic bathymetry and ice topography
Raw
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
Raw
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()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment