Additional fixes to trending.py

webbpsf/trending.py

@@ -7,6 +7,7 @@
 import astropy.io.fits as fits
 import astropy.time
 import astropy.units as u
+from astroquery.mast import Observations
 import matplotlib
 import matplotlib.pyplot as plt
 from matplotlib.backends.backend_pdf import PdfPages
@@ -96,7 +97,8 @@ def wavefront_time_series_plot(opdtable, start_date=None, end_date=None, ymin=0,
 
     rms_nm = np.asarray(rmses) * 1000
 
-    routine_pids = [1163, 2586, 2724, 2725, 2726]  # commissioning OTE-26 and cycle 1 maintenance
+    routine_pids = [1163, 2586, 2724, 2725, 2726, 4431, # Commissioning OTE-26 and Cycle 1 maintenance
+                    4500, 4501, 4502, 4503, 4504, 4505, 4506, 4507, 4508, 4509, 4510]  # Cycle 2
 
     is_routine = np.asarray([int(v[1:6]) in routine_pids for v in opdtable[where_pre]['visitId']])
 
@@ -136,7 +138,7 @@ def wavefront_time_series_plot(opdtable, start_date=None, end_date=None, ymin=0,
     ax.xaxis.set_minor_locator(matplotlib.dates.DayLocator())
     ax.tick_params('x', length=10)
     for tick in ax.get_xticklabels():
-        tick.set_rotation(45)
+        tick.set_rotation(75)
 
     # label events
     if label_events:
@@ -175,7 +177,7 @@ def wavefront_time_series_plot(opdtable, start_date=None, end_date=None, ymin=0,
 
 
 
-def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, download_opds=True, mark_corrections='lines'):
+def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, pid=None, download_opds=True, mark_corrections='lines'):
     """ Plot histogram and cumulative histogram of WFE over some time range.
 
     Parameters
@@ -188,18 +190,29 @@ def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, do
         threshold to filter the RMS WFE
     download_opds : bool
         toggle downloading of OPDs from MAST
+    pid : int, optional
+        Program ID for which dates of observations are to be overplotted
 
     Returns
     -------
     Nothing, but makes a plot
 
     """
-
+    
     if start_date is None:
         start_date = astropy.time.Time('2022-07-16')
     if end_date is None:
         end_date = astropy.time.Time.now()
 
+
+    # Look up wavefront sensing and mirror move corrections for that month
+    if pid:
+        pid_dates = get_dates_for_pid(pid)
+        if pid_dates is not None:
+            pid_dates = pid_dates[(pid_dates >= start_date) & (pid_dates <= end_date)]
+        else:
+            pid = None
+
     opdtable0 = webbpsf.mast_wss.deduplicate_opd_table(opdtable)
     opdtable1 = webbpsf.mast_wss.filter_opd_table(opdtable0, start_time=start_date, end_time=end_date)
 
@@ -235,14 +248,16 @@ def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, do
     interp_rmses = interp_fn(mjdrange)
 
     # Plot
-    fig, axes = plt.subplots(figsize=(16,9), nrows=2, gridspec_kw = {'hspace':0.3})
+    hspace = 0.3
+    nrows  = 2
+    fig, axes = plt.subplots(figsize=(16,9), nrows=nrows, gridspec_kw = {'hspace':hspace})
 
+    ms = 14 #markersize
 
-    axes[0].plot_date(dates.plot_date, np.asarray(rmses)*1e3, 'o', ls='-', label='Sensing visit')
+    axes[0].plot_date(dates.plot_date, np.asarray(rmses)*1e3, '.', ms=ms, ls='-', label='Sensing visit')
     axes[0].xaxis.set_major_locator(matplotlib.dates.DayLocator(bymonthday=[1, 15]))
     axes[0].xaxis.set_minor_locator(matplotlib.dates.DayLocator(interval=1))
     axes[0].tick_params('x', length=10)
-
 
     if mark_corrections=='lines':
         # Add vertical lines for corrections
@@ -254,18 +269,30 @@ def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, do
                     plot.set_label('Corrections')
                     icorr += 1
     elif mark_corrections=='triangles':
-        yval = (np.asarray(rmses)*1e3).max()*0.95
+        yval = (np.asarray(rmses)*1e3).max() *1.01
         axes[0].scatter(dates[where_post].plot_date, np.ones(np.sum(where_post))*yval, 
                         marker='v', s=100, color='limegreen', label='Corrections')
     elif mark_corrections=='arrows':
         rms_nm =  np.asarray(rmses)*1e3
-        axes[0].scatter(dates[where_post].plot_date, rms_nm[where_post]+1, 
-                        marker='v', s=100, color='limegreen', label='Corrections')
+
+        sub_height = fig.get_figheight() / (nrows+hspace)        
+        plot_size_points = np.asarray([fig.get_figwidth(), sub_height]) * fig.dpi
+        plot_size_data   = [np.diff(axes[0].get_xlim())[0], np.diff(axes[0].get_ylim())[0]]
+
+        yoffset =  (1.2*ms) * plot_size_data[1] / plot_size_points[1]
+        axes[0].scatter(dates[where_post].plot_date, rms_nm[where_post] + yoffset, 
+                        marker='v', s=100, color='limegreen', label='Corrections', zorder=99)
+
+        yoffsets = [0.6 * ms * plot_size_data[0] / plot_size_points[0],
+                    0.6 * ms * plot_size_data[1] / plot_size_points[1]]
+
         for i, idate in enumerate(where_post):
-            plot_offsets = np.array([-1, 2])
             if idate:
-                axes[0].plot(dates[i-1:i+1].plot_date, rms_nm[i-1:i+1]+plot_offsets, color='limegreen', lw=3, ls='-', alpha=0.5)
+                xtmp = dates[i-1:i+1]
+                ytmp = [rms_nm[i-1] - yoffsets[1], rms_nm[i] + yoffsets[1]]
+                axes[0].plot(xtmp.plot_date, ytmp, color='limegreen', lw=2, ls='-')
 
+    if pid: axes[0].set_ylim(0.975*axes[0].get_ylim()[0], 1.025*axes[0].get_ylim()[1])
     axes[0].set_xlabel("Date")
     axes[0].set_ylabel("RMS WFE\n(OTE+NIRCam)", fontweight='bold')
     axes[0].set_title(f"Observatory WFE from {start_date.isot[0:10]} to {end_date.isot[0:10]}",
@@ -274,7 +301,6 @@ def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, do
     if thresh:
         axes[0].axhline(thresh, color='C2', label='Correction threshold', linestyle='dashed')
 
-    axes[0].legend()
     axes[0].tick_params(right=True, which='both', direction = 'in')
 
     nbins=100
@@ -313,6 +339,18 @@ def wfe_histogram_plot(opdtable, start_date=None, end_date=None, thresh=None, do
                     fontweight='bold', fontsize=14)    
 
 
+    # Add vertical lines for dates of PID observations
+    if pid:
+        for i, obs_date in enumerate(pid_dates):
+            axes[0].axvline(obs_date.plot_date, color='darkgrey', ls='--', alpha=0.5, zorder=1)
+            y_star = axes[0].get_ylim()[0] + 0.10*np.diff(axes[0].get_ylim())
+            if i==0:
+                label = "PID {:d} obs. date(s)".format(pid)
+            else:
+                label = None
+            axes[0].scatter(obs_date.plot_date, y_star, marker='*', s=200, color='darkgrey', label=label)
+
+    axes[0].legend()
 
 ##### Wavefront Drifts Plot #####
 
@@ -919,11 +957,54 @@ def check_colnames(opdtable):
     return opdtable
 
 
-def monthly_trending_plot(year, month, verbose=True, instrument='NIRCam', filter='F200W', vmax=200, opdtable=None):
+def get_dates_for_pid(pid, project='jwst'):
+    """ Check the archive for the start dates of each observation of the specified PID
+
+    Parameters
+    -----------
+    pid: int
+         Program ID for which the observation start dates are to be overplotted
+
+    Returns:
+    start_times: list of astropy.time.Time dates
+
+    """
+
+
+    try:
+        obs_table = Observations.query_criteria(project='jwst', proposal_id=pid)
+
+        cond_calib = (obs_table['calib_level']>1)
+        obs_table  = obs_table[cond_calib]
+
+        obs_table['obs_num'] = [x['obs_id'][7:10] if x['calib_level']==2 else x['obs_id'][8:13] for x in obs_table]
+        obs_table['obs_num'] = [x.strip("_") for x in obs_table['obs_num'] ]
+
+        obs_by_num = obs_table.group_by('obs_num')
+
+        start_times = []
+        for key, group in zip(obs_by_num.groups.keys, obs_by_num.groups):
+            start_times.append(group['t_min'].min())
+
+        start_times = astropy.time.Time(start_times, format='mjd')
+        start_times = start_times.to_value('fits')#, subfmt='date')
+        start_times = astropy.time.Time( np.unique(start_times) )
+
+        return start_times
+    except:
+        print("No access to data for PID {:d}".format(pid))
+        return 
+
+
+
+
+def monthly_trending_plot(year, month, verbose=True, instrument='NIRCam', filter='F200W', vmax=200, pid=None, opdtable=None):
     """Make monthly trending plot showing OPDs, mirror moves, RMS WFE, and the resulting PSF EEs
 
     year, month : integers
         Self explanatory
+    pid : int, optional
+        Program ID for which dates of observations are to be overplotted
     verbose : bool
         Print more verbose text output
     vmax : float
@@ -932,11 +1013,20 @@ def monthly_trending_plot(year, month, verbose=True, instrument='NIRCam', filter
         Table of available OPDs, Default None: as returned by retrieve_mast_opd_table()
     """
 
+
     def vprint(*text):
         if verbose: print(*text)
 
-    # Look up wavefront sensing and mirror move corrections for that month
     start_date, end_date = get_month_start_end(year, month)
+
+    # Look up wavefront sensing and mirror move corrections for that month
+    if pid:
+        pid_dates = get_dates_for_pid(pid)
+        if pid_dates is not None:
+            pid_dates = pid_dates[(pid_dates >= start_date) & (pid_dates <= end_date)]
+        else:
+            pid = None
+
     if opdtable is None:
         opdtable = get_opdtable_for_month(year, month)
     else:
@@ -1021,7 +1111,7 @@ def basic_show_image(image, ax, vmax=.3, nanmask=1):
         ax.set_yticks([])
 
     # Make the plots
-    fig = plt.figure(constrained_layout=False, figsize=(16, 9), )
+    fig = plt.figure(constrained_layout=False, figsize=(16, 10), )
 
     subfigs = fig.subfigures(2, 1, hspace=0.02, height_ratios=[2, 2], )
 
@@ -1043,17 +1133,18 @@ def basic_show_image(image, ax, vmax=.3, nanmask=1):
                       color='C1', ls='-', label='Observatory WFE at NIRCam NRCA3')
     axes[0].plot_date(dates_array.plot_date, rms_ote * 1e9,
                       color='C0', ls='-', label='Telescope WFE')
+
     for ax in axes:
         for corr_date in correction_times:
             ax.axvline(corr_date.plot_date, color='darkgreen', ls='--', alpha=0.5)
 
+
     #axes[0].axhline(59, ls=":", color='gray')
     #axes[0].axhline(80, ls=":", color='gray')
     axes[0].fill_between( [start_date.plot_date - 0.5, end_date.plot_date + 0.5],
                           [59,59], [80, 80], color='blue', alpha=0.08, label='Wavefront control target range')
-    axes[0].legend(loc='lower right', fontsize=9)
 
-    axes[0].set_ylim(0, 150)
+    axes[0].set_ylim(0.8*rms_ote.min()*1e9, 1.2*rms_obs.max()*1e9)
     axes[0].set_ylabel("Wavefront Error\n[nm rms]", fontsize=fs, fontweight='bold')
     axes[0].set_xticklabels([])
 
@@ -1068,9 +1159,9 @@ def basic_show_image(image, ax, vmax=.3, nanmask=1):
         ee_measurements[ee_npix] = ees_at_rad  # save for later
 
         median_ee = np.median(ees_at_rad)
-        ee_ax_ylim = np.max([ee_ax_ylim, np.abs(ees_at_rad-median_ee).max()*1.1]) # display tweak: adjust the plot Y scale sensibly to its contents
+        ee_ax_ylim = np.max([ee_ax_ylim, np.abs((ees_at_rad-median_ee)/median_ee).max()*1.1]) # display tweak: adjust the plot Y scale sensibly to its contents
 
-        axes[1].plot_date(dates_array.plot_date, ees_at_rad-median_ee, ls='-', color=color,
+        axes[1].plot_date(dates_array.plot_date, (ees_at_rad - median_ee)/median_ee, ls='-', color=color,
                           label=f"$\Delta$EE within {ee_rad:.2f} arcsec ({ee_npix} pix)")
 
         axes[1].text(0.01, 0.75-i*0.12, f'Median EE within {ee_rad:.2f} arcsec = {median_ee:.3f}', color=color,
@@ -1079,11 +1170,11 @@ def basic_show_image(image, ax, vmax=.3, nanmask=1):
 
     axes[1].fill_between( [start_date.plot_date - 0.5, end_date.plot_date + 0.5], -0.03, 0.03, color='gray', alpha=0.1, label="±3% change (stability requirement)")
     axes[1].set_xlabel("Date", fontsize=fs, fontweight='bold')
-    axes[1].set_ylabel(f"Change in \nEncircled Energy\n{instrument} {filter}", fontsize=fs, fontweight='bold')
+    axes[1].set_ylabel(f"% Change in \nEncircled Energy\n{instrument} {filter}", fontsize=fs, fontweight='bold')
     axes[1].set_ylim(0.5, 1.0)
     axes[1].axhline(0, ls=":", color='gray')
     axes[1].set_ylim(-ee_ax_ylim, ee_ax_ylim)
-    axes[1].legend(loc='upper right', fontsize=9)
+
 
     # Configure Axes for the time series plots
     for ax in axes[0:2]:
@@ -1093,6 +1184,21 @@ def basic_show_image(image, ax, vmax=.3, nanmask=1):
         ax.xaxis.set_major_locator(matplotlib.dates.WeekdayLocator(byweekday=matplotlib.dates.MONDAY,
                                                                    interval=1))
 
+    # Add vertical lines for dates of PID observations
+    if pid:
+        for iax, ax in enumerate(axes):
+            for i, obs_date in enumerate(pid_dates):
+                ax.axvline(obs_date.plot_date, color='darkgrey', ls='--', alpha=0.5, zorder=1)
+                y_star = ax.get_ylim()[0] + 0.20*np.diff(ax.get_ylim())
+                if iax==0 and i==0:
+                    label = "PID {:d} obs. date(s)".format(pid)
+                else:
+                    label = None
+                ax.scatter(obs_date.plot_date, y_star, marker='*', s=200, color='darkgrey', label=label)
+
+    axes[0].legend(loc='best', fontsize=9, ncol=2)
+    axes[1].legend(loc='upper right', fontsize=9)
+
     nanmask = np.zeros_like(apmask) + np.nan
     nanmask[apmask == 1] = 1
     plot_index = -1
@@ -1143,7 +1249,7 @@ def basic_show_image(image, ax, vmax=.3, nanmask=1):
             basic_show_image(delta_opd * 1e6, ax=im_axes[2, plot_index], nanmask=nanmask, vmax=vmax_micron)
             im_axes[2, plot_index].text(20, 20, f"{webbpsf.utils.rms(delta_opd, mask=apmask)*1e9:.1f}", color='yellow', fontsize=fs*0.6)
 
-    for i, l in enumerate(['WF Sensing', "Drifts", 'Corrections']):
+    for i, l in enumerate(['Measured\nWFE', "Drifts", 'Mirror\nCorrections']):
         im_axes[i, 0].yaxis.set_visible(True)
         im_axes[i, 0].set_ylabel(l + "\n\n", fontsize=fs, fontweight='bold')