Renamed JOSS plot script, updated code, added more comments

.gitignore

@@ -6,6 +6,7 @@ __pycache__
 /*.egg-info/
 /examples/jsoc.stanford.edu/
 /examples/downloads/
+/examples/joss_figure.pdf
 /nogit/
 /.tox/
 /doc/build/

examples/create_sharp_plots.py → examples/create_joss_figure.py

@@ -1,41 +1,66 @@
+"""
+This example creates the figure shown in the JOSS paper.
+"""
+from __future__ import absolute_import, division, print_function
 import os.path
 import numpy as np
 import matplotlib.pyplot as plt
 from matplotlib import dates
 from astropy.io import fits
-import pandas
+import example_helpers
 import drms
 
+import pandas
+pandas_version = tuple(map(int, pandas.__version__.split('.')[:2]))
+if pandas_version >= (0, 22):
+    # Since pandas v0.22, we need to explicitely register matplotlib
+    # converters to use pandas.Timestamp objects in plots.
+    pandas.plotting.register_matplotlib_converters()
+
 
 def read_fits_data(fname):
     """Reads FITS data and fixes/ignores any non-standard FITS keywords."""
     hdulist = fits.open(fname)
     hdulist.verify('silentfix+warn')
     return hdulist[1].data
 
+# Print the doc string of this example.
+print(__doc__)
 
-do_download = 1
-save_plots = 1
-show_plots = 0
+# This example requires a registered export email address. You can register
+# JSOC exports at: http://jsoc.stanford.edu/ajax/register_email.html
+#
+# You will be asked for your registered email address during execution of
+# this example. If you don't want to enter it every time you run this script,
+# you can set the environment variable JSOC_EXPORT_EMAIL or the variable
+# below to your registered email address.
+email = ''
 
-export_email = '[email protected]'
+# Should the plots been shown and/or saved?
+save_plots = True
+show_plots = True
+
+# Series, harpnum and segment selection
 series = 'hmi.sharp_720s'
 sharpnum = 4315
 segments = ['magnetogram', 'continuum']
-# segments = ['magnetogram', 'field', 'continuum', 'Dopplergram']
 fname_fmt_str = '{series}.{sharpnum}.{tstr}.{segment}.fits'
 
-kwlist = ['T_REC', 'LON_FWT', 'OBS_VR', 'CROTA2',
-          'CRPIX1', 'CRPIX2', 'CDELT1', 'CDELT2', 'CRVAL1', 'CRVAL2',
-          'AREA_ACR', 'MEANGAM', 'USFLUX', 'ERRVF', 'MEANJZH', 'ERRMIH']
+# Download directory
+out_dir = os.path.join('downloads', 'sharp_joss')
 
-if do_download:
-    c = drms.Client(email=export_email)
-else:
-    c = drms.Client('kis')
+# Create download directory if it does not exist yet.
+if not os.path.exists(out_dir):
+    os.makedirs(out_dir)
+
+# Keywords to be queried
+kwlist = ['T_REC', 'LON_FWT', 'CROTA2', 'AREA_ACR', 'USFLUX', 'ERRVF',
+          'CRPIX1', 'CRPIX2', 'CDELT1', 'CDELT2', 'CRVAL1', 'CRVAL2']
+
+# Create DRMS client, use debug=True to see the query URLs.
+c = drms.Client(verbose=True)
 
 print('Querying metadata...')
-si = c.info(series)
 kw = c.query('%s[%d]' % (series, sharpnum), key=kwlist, rec_index=True)
 t = drms.to_datetime(kw.T_REC)
 
@@ -45,36 +70,43 @@ def read_fits_data(fname):
 t_cm = drms.to_datetime(kw.T_REC[rec_cm])
 print('-> rec_cm:', rec_cm, '@', kw.LON_FWT[rec_cm], 'deg')
 
-# Generate filenames to check if the files are already downloaded.
+# Check if any files were already downloaded.
+fnames = {}
+download_segments = []
 t_cm_str = t_cm.strftime('%Y%m%d_%H%M%S_TAI')
-fnames = {
-    s: fname_fmt_str.format(
+for s in segments:
+    fnames[s] = fname_fmt_str.format(
         series=series, sharpnum=sharpnum, tstr=t_cm_str, segment=s)
-    for s in segments}
-
-# Check if any files were already downloaded and download them if not.
-download_segments = []
-for k, v in fnames.items():
-    if not os.path.exists(v):
-        download_segments.append(k)
-if do_download and download_segments:
+    if not os.path.exists(os.path.join(out_dir, fnames[s])):
+        download_segments.append(s)
+
+# Only download missing files.
+if download_segments:
+    print()
+    if not email:
+        email = example_helpers.get_export_email()
+    if not email or not c.check_email(email):
+        raise RuntimeError('Email address is not valid or not registered.')
     print('Downloading files...')
-    exp_query = '%s{%s}' % (rec_cm, ','.join(download_segments))
-    r = c.export(exp_query)
-    dl = r.download('.', verbose=True)
-
+    export_query = '%s{%s}' % (rec_cm, ','.join(download_segments))
+    r = c.export(export_query, email=email)
+    dl = r.download(out_dir)
+    print()
 
-print('Reading data files...')
-data = {
-    s: read_fits_data(fnames[s])
-    for s in segments}
+print('Reading image files...')
+data_mag = read_fits_data(os.path.join(out_dir, fnames['magnetogram']))
+data_cont = read_fits_data(os.path.join(out_dir, fnames['continuum']))
 
-# Remove observer's LOS velocity
-if 'Dopplergram' in data:
-    data['Dopplergram'] -= k_cm.OBS_VR
+print('Creating plots...')
+plt.rc('figure', figsize=(11, 6.75))
+plt.rc('axes', titlesize='medium')
+plt.rc('axes.formatter', use_mathtext=True)
+plt.rc('mathtext', default='regular')
+plt.rc('legend', fontsize='medium', framealpha=1.0)
+plt.rc('image', origin='lower', interpolation='none', cmap='gray')
 
 # Convert pixel to world coordinates using WCS keywords
-ny, nx = data[segments[0]].shape
+ny, nx = data_mag.shape
 xmin = (1 - k_cm.CRPIX1)*k_cm.CDELT1 + k_cm.CRVAL1
 xmax = (nx - k_cm.CRPIX1)*k_cm.CDELT1 + k_cm.CRVAL1
 ymin = (1 - k_cm.CRPIX2)*k_cm.CDELT2 + k_cm.CRVAL2
@@ -83,8 +115,8 @@ def read_fits_data(fname):
 # We assume a CROTA2 value close to 180 degree, so we can approximate the
 # rotation by inverting the image axes.
 if abs(180 - k_cm.CROTA2) < 0.1:
-    for s in segments:
-        data[s] = data[s][::-1, ::-1]
+    data_mag = data_mag[::-1, ::-1]
+    data_cont = data_cont[::-1, ::-1]
     xmin, xmax = -xmax, -xmin
     ymin, ymax = -ymax, -ymin
 else:
@@ -94,44 +126,23 @@ def read_fits_data(fname):
 extent = (xmin - abs(k_cm.CDELT1)/2, xmax + abs(k_cm.CDELT1)/2,
           ymin - abs(k_cm.CDELT2)/2, ymax + abs(k_cm.CDELT2)/2)
 
-
-print('Creating plots...')
-plt.rc('figure', figsize=(11, 6.75))
-plt.rc('axes', titlesize='medium')
-plt.rc('axes.formatter', use_mathtext=True)
-plt.rc('mathtext', default='regular')
-plt.rc('legend', fontsize='medium', framealpha=1.0)
-plt.rc('image', origin='lower', interpolation='none', cmap='gray')
-pandas.plotting.register_matplotlib_converters()
-
+# Create figure with 2x2 axes, with enough space for colorbars on the right
 fig, ax = plt.subplots(
     2, 2, num=1, clear=True,
     gridspec_kw={'width_ratios': (2, 2.5)})
 ax_meta, ax_img = ax[:, 0], ax[:, 1]
 
-###
-
-axi = ax_meta[1]
-axi.plot(t, kw.AREA_ACR/1e3, '.', ms=2, label='AREA_ACR')
-axi.set_title('LoS area of active pixels')
-axi.set_ylabel(r'$\mu$Hem $\times 1000$')
-
-# axi = ax_meta[1]
-# axi.plot(t, kw.MEANGAM, '.', ms=2, label='MEANGAM')
-# axi.set_title('Mean inclination angle')
-# axi.set_ylabel('degree')
-
+# Create metadata line plots in the left column
 axi = ax_meta[0]
 axi.errorbar(t, kw.USFLUX/1e22, yerr=kw.ERRVF/1e22, fmt='.', ms=2,
              capsize=0, label='USFLUX')
 axi.set_title('Total unsigned flux')
 axi.set_ylabel(r'Mx $\times 10^{\minus 22}$')
 
-# axi = ax_meta[1]
-# axi.errorbar(t, kw.MEANJZH*1e3, yerr=kw.ERRMIH*1e3, fmt='.', ms=2,
-#              capsize=0, label='MEANJZH')
-# axi.set_title('Mean current helicity')
-# axi.set_ylabel(r'G$^2$ m$^{\minus 1}$ $\times 1000$')
+axi = ax_meta[1]
+axi.plot(t, kw.AREA_ACR/1e3, '.', ms=2, label='AREA_ACR')
+axi.set_title('LoS area of active pixels')
+axi.set_ylabel(r'$\mu$Hem $\times 1000$')
 
 ax_meta[0].set_xticklabels([])
 ax_meta[1].xaxis.set_major_locator(dates.AutoDateLocator())
@@ -142,33 +153,18 @@ def read_fits_data(fname):
     axi.axvline(t_cm, ls='--', color='tab:orange')
     axi.legend(loc='upper left', numpoints=1)
 
-###
+# Create image plots in the right column
+axi = ax_img[0]
+axi.set_title('Continuum intensity')
+im = axi.imshow(data_cont/1e3, extent=extent, vmax=61)
+cb = plt.colorbar(im, ax=axi, label=r'$I_{\mathrm{c}}$ [kDN/s]', pad=0.03)
 
 axi = ax_img[1]
-di = data['magnetogram']
 axi.set_title('LoS magnetogram')
-im = axi.imshow(di/1e3, extent=extent, vmin=-1, vmax=1)
+im = axi.imshow(data_mag/1e3, extent=extent, vmin=-1, vmax=1)
 cb = plt.colorbar(im, ax=axi, label=r'$B_{\mathrm{los}}$ [kG]', pad=0.03)
 cb.set_ticks([-1, -0.5, 0, 0.5, 1, 2])
 
-axi = ax_img[0]
-di = data['continuum']
-axi.set_title('Continuum intensity')
-im = axi.imshow(di/1e3, extent=extent, vmax=61)
-cb = plt.colorbar(im, ax=axi, label=r'$I_{\mathrm{c}}$ [kDN/s]', pad=0.03)
-
-# axi = ax_img[0]
-# di = data['field']
-# axi.set_title('Magnetic field strength')
-# im = axi.imshow(di/1e3, extent=extent, vmin=0.05, vmax=2.6)
-# cb = plt.colorbar(im, ax=axi, label=r'$|B|$ [kG]')
-
-# axi = ax_img[1]
-# di = data['Dopplergram']
-# axi.set_title('LoS Doppler velocity')
-# im = axi.imshow(di/1e3, extent=extent, vmin=-1.3, vmax=1.3)
-# cb = plt.colorbar(im, ax=axi, label=r'$v_{\mathrm{los}}$ [km/s]')
-
 for axi in ax_img:
     axi.set_xlim(-130, 141)
     axi.set_ylim(-262, -86)
@@ -179,15 +175,14 @@ def read_fits_data(fname):
 ax_img[0].set_ylabel('Solar Y [arcsec]')
 ax_img[1].set_ylabel('Solar Y [arcsec]')
 
-###
-
+# Make better use of figure space
 fig.tight_layout(pad=1.2, w_pad=2)
 plt.draw()
 
 if save_plots:
-    print('Saving plots...')
-    fig.savefig('sharp.pdf', dpi=200)
-    fig.savefig('sharp.png', dpi=200)
+    print('Saving figure...')
+    fig.savefig('joss_figure.pdf', dpi=200)
 
 if show_plots:
+    print('Showing figure...')
     plt.show()