Repository: kevin218/eureka
Issue: [Bug]: S3 crash generating plot

### FAQ check

- [X] Yes, I checked the FAQ and my question has not been addressed.

### Instrument

NIRSpec (Stages 1-3)

### What happened?

Stage 3 was running happily, then crashed trying to generate a figure

### Error traceback output

2022-08-26 09:22:13,248 - stpipe - WARNING - /home/dlong/Eureka/src/eureka/S3_data_reduction/nircam.py:187: RankWarning: Polyfit may be poorly conditioned
  bg, mask = background.fitbg(dataim, meta, datamask, meta.bg_y1,

2022-08-26 09:22:13,249 - stpipe - WARNING - /home/dlong/Eureka/src/eureka/S3_data_reduction/nircam.py:187: RankWarning: Polyfit may be poorly conditioned
  bg, mask = background.fitbg(dataim, meta, datamask, meta.bg_y1,

100%|█████████████████████████████████████████████████████████████████████████████| 3051/3051 [00:33<00:00, 90.02it/s]
  Creating figures for background subtraction...
100%|███████████████████████████████████████████████████████████████████████████████████| 5/5 [00:03<00:00,  1.32it/s]
  Recording y position and width for all integrations...
100%|███████████████████████████████████████████████████████████████████████████| 3051/3051 [00:01<00:00, 2091.45it/s]
  Extracting aperture region...
  Performing optimal spectral extraction...
100%|█████████████████████████████████████████████████████████████████████████████| 3051/3051 [13:11<00:00,  3.85it/s]
  Creating figures for optimal spectral extraction
100%|███████████████████████████████████████████████████████████████████████████████████| 5/5 [00:02<00:00,  1.83it/s]
Finished writing to /astro/jtste/doug/TRAPPIST-1c-Eureka/Stage3/S3_2022-08-26_trappist1c_run1/ap6_bg7/S3_trappist1c_ap6_bg7_FluxData_seg0002.h5
Finished writing to /astro/jtste/doug/TRAPPIST-1c-Eureka/Stage3/S3_2022-08-26_trappist1c_run1/ap6_bg7/S3_trappist1c_ap6_bg7_SpecData.h5
Stage 3 MAD = 43543 ppm
Generating figure
Traceback (most recent call last):
  File "/astro/jtste/doug/TRAPPIST-1c-Eureka/run_eureka.py", line 28, in <module>
    s3_spec, s3_meta = s3.reduce(eventlabel, ecf_path=ecf_path,s2_meta=s2_meta)
  File "/home/dlong/Eureka/src/eureka/S3_data_reduction/s3_reduce.py", line 534, in reduce
    plots_s3.lc_nodriftcorr(meta, spec.wave_1d, spec.optspec,
  File "/home/dlong/Eureka/src/eureka/S3_data_reduction/plots_s3.py", line 59, in lc_nodriftcorr
    im1 = ax1.pcolormesh(wave_1d, np.arange(meta.n_int),
  File "/user/dlong/linux/miniconda3/envs/eureka/lib/python3.9/site-packages/matplotlib/__init__.py", line 1412, in inner
    return func(ax, *map(sanitize_sequence, args), **kwargs)
  File "/user/dlong/linux/miniconda3/envs/eureka/lib/python3.9/site-packages/matplotlib/axes/_axes.py", line 6058, in pcolormesh
    X, Y, C, shading = self._pcolorargs('pcolormesh', *args,
  File "/user/dlong/linux/miniconda3/envs/eureka/lib/python3.9/site-packages/matplotlib/axes/_axes.py", line 5533, in _pcolorargs
    raise ValueError(
ValueError: x and y arguments to pcolormesh cannot have non-finite values or be of type numpy.ma.core.MaskedArray with masked values

### What operating system are you using?

RHEL 8.6

### What version of Python are you running?

Python 3.9

### What Python packages do you have installed?

# packages in environment at /user/dlong/linux/miniconda3/envs/eureka:
#
# Name                    Version                   Build  Channel
_libgcc_mutex             0.1                 conda_forge    conda-forge
_openmp_mutex             4.5                       2_gnu    conda-forge
alabaster                 0.7.12                   pypi_0    pypi
asdf                      2.12.0                   pypi_0    pypi
asdf-astropy              0.2.1                    pypi_0    pypi
asdf-coordinates-schemas  0.1.0                    pypi_0    pypi
asdf-standard             1.0.2                    pypi_0    pypi
asdf-transform-schemas    0.2.2                    pypi_0    pypi
asdf-wcs-schemas          0.1.1                    pypi_0    pypi
asteval                   0.9.27                   pypi_0    pypi
astraeus                  0.2                      pypi_0    pypi
astropy                   5.1                      pypi_0    pypi
astropy-healpix           0.6                      pypi_0    pypi
astroquery                0.4.6                    pypi_0    pypi
astroscrappy              1.1.0                    pypi_0    pypi
asttokens                 2.0.7                    pypi_0    pypi
attrs                     22.1.0                   pypi_0    pypi
babel                     2.10.3                   pypi_0    pypi
backcall                  0.2.0                    pypi_0    pypi
batman-package            2.4.9                    pypi_0    pypi
bayesicfitting            3.0.1                    pypi_0    pypi
beautifulsoup4            4.11.1                   pypi_0    pypi
bokeh                     2.4.3                    pypi_0    pypi
bottleneck                1.3.5                    pypi_0    pypi
ca-certificates           2022.6.15            ha878542_0    conda-forge
ccdproc                   2.3.1                    pypi_0    pypi
celerite                  0.4.2                    pypi_0    pypi
certifi                   2022.6.15                pypi_0    pypi
cffi                      1.15.1                   pypi_0    pypi
cftime                    1.6.1                    pypi_0    pypi
charset-normalizer        2.1.0                    pypi_0    pypi
cloudpickle               2.1.0                    pypi_0    pypi
corner                    2.2.1                    pypi_0    pypi
crds                      11.16.7                  pypi_0    pypi
cryptography              37.0.4                   pypi_0    pypi
cycler                    0.11.0                   pypi_0    pypi
cython                    0.29.32                  pypi_0    pypi
dask                      2022.6.0                 pypi_0    pypi
decorator                 5.1.1                    pypi_0    pypi
docutils                  0.19                     pypi_0    pypi
drizzle                   1.13.6                   pypi_0    pypi
dynesty                   1.2.3                    pypi_0    pypi
emcee                     3.1.2                    pypi_0    pypi
eureka                    0.5.dev197+gf1b824ec          pypi_0    pypi
executing                 0.9.1                    pypi_0    pypi
exotic-ld                 2.0.0                    pypi_0    pypi
filelock                  3.7.1                    pypi_0    pypi
fonttools                 4.34.4                   pypi_0    pypi
fsspec                    2022.7.1                 pypi_0    pypi
future                    0.18.2                   pypi_0    pypi
george                    0.4.0                    pypi_0    pypi
gwcs                      0.18.1                   pypi_0    pypi
h5netcdf                  1.0.2                    pypi_0    pypi
h5py                      3.1.0                    pypi_0    pypi
html5lib                  1.1                      pypi_0    pypi
idna                      3.3                      pypi_0    pypi
imageio                   2.21.1                   pypi_0    pypi
imagesize                 1.4.1                    pypi_0    pypi
importlib-metadata        4.12.0                   pypi_0    pypi
iniconfig                 1.1.1                    pypi_0    pypi
ipython                   8.4.0                    pypi_0    pypi
jedi                      0.18.1                   pypi_0    pypi
jeepney                   0.8.0                    pypi_0    pypi
jinja2                    3.1.2                    pypi_0    pypi
jmespath                  1.0.1                    pypi_0    pypi
jsonschema                4.9.1                    pypi_0    pypi
jwst                      1.6.0                    pypi_0    pypi
keyring                   23.8.2                   pypi_0    pypi
kiwisolver                1.4.4                    pypi_0    pypi
ld_impl_linux-64          2.36.1               hea4e1c9_2    conda-forge
libffi                    3.4.2                h7f98852_5    conda-forge
libgcc-ng                 12.1.0              h8d9b700_16    conda-forge
libgomp                   12.1.0              h8d9b700_16    conda-forge
libzlib                   1.2.12               h166bdaf_2    conda-forge
lmfit                     1.0.3                    pypi_0    pypi
locket                    1.0.0                    pypi_0    pypi
lxml                      4.9.1                    pypi_0    pypi
markupsafe                2.1.1                    pypi_0    pypi
matplotlib                3.5.2                    pypi_0    pypi
matplotlib-inline         0.1.3                    pypi_0    pypi
ncurses                   6.3                  h27087fc_1    conda-forge
netcdf4                   1.6.0                    pypi_0    pypi
networkx                  2.8.5                    pypi_0    pypi
numpy                     1.23.1                   pypi_0    pypi
numpydoc                  1.4.0                    pypi_0    pypi
openssl                   3.0.5                h166bdaf_1    conda-forge
packaging                 21.3                     pypi_0    pypi
pandas                    1.4.3                    pypi_0    pypi
parsley                   1.3                      pypi_0    pypi
parso                     0.8.3                    pypi_0    pypi
partd                     1.2.0                    pypi_0    pypi
pexpect                   4.8.0                    pypi_0    pypi
photutils                 1.5.0                    pypi_0    pypi
pickleshare               0.7.5                    pypi_0    pypi
pillow                    9.2.0                    pypi_0    pypi
pip                       22.2.2             pyhd8ed1ab_0    conda-forge
pluggy                    1.0.0                    pypi_0    pypi
poppy                     1.0.3                    pypi_0    pypi
prompt-toolkit            3.0.30                   pypi_0    pypi
psutil                    5.9.1                    pypi_0    pypi
ptyprocess                0.7.0                    pypi_0    pypi
pure-eval                 0.2.2                    pypi_0    pypi
py                        1.11.0                   pypi_0    pypi
pycparser                 2.21                     pypi_0    pypi
pyerfa                    2.0.0.1                  pypi_0    pypi
pygments                  2.12.0                   pypi_0    pypi
pyparsing                 3.0.9                    pypi_0    pypi
pyrsistent                0.18.1                   pypi_0    pypi
pysynphot                 2.0.0                    pypi_0    pypi
pytest                    7.1.2                    pypi_0    pypi
python                    3.9.7           hf930737_3_cpython    conda-forge
python-dateutil           2.8.2                    pypi_0    pypi
python_abi                3.9                      2_cp39    conda-forge
pytz                      2022.1                   pypi_0    pypi
pyvo                      1.3                      pypi_0    pypi
pywavelets                1.3.0                    pypi_0    pypi
pyyaml                    6.0                      pypi_0    pypi
readline                  8.1.2                h0f457ee_0    conda-forge
reproject                 0.8                      pypi_0    pypi
requests                  2.28.1                   pypi_0    pypi
scikit-image              0.19.3                   pypi_0    pypi
scipy                     1.9.0                    pypi_0    pypi
secretstorage             3.3.2                    pypi_0    pypi
semantic-version          2.10.0                   pypi_0    pypi
setuptools                63.4.2           py39hf3d152e_0    conda-forge
setuptools-scm            7.0.5                    pypi_0    pypi
six                       1.16.0                   pypi_0    pypi
snowballstemmer           2.2.0                    pypi_0    pypi
soupsieve                 2.3.2.post1              pypi_0    pypi
spherical-geometry        1.2.22                   pypi_0    pypi
sphinx                    5.1.1                    pypi_0    pypi
sphinxcontrib-applehelp   1.0.2                    pypi_0    pypi
sphinxcontrib-devhelp     1.0.2                    pypi_0    pypi
sphinxcontrib-htmlhelp    2.0.0                    pypi_0    pypi
sphinxcontrib-jsmath      1.0.1                    pypi_0    pypi
sphinxcontrib-qthelp      1.0.3                    pypi_0    pypi
sphinxcontrib-serializinghtml 1.1.5                    pypi_0    pypi
sqlite                    3.39.2               h4ff8645_0    conda-forge
stack-data                0.3.0                    pypi_0    pypi
stcal                     1.0.0                    pypi_0    pypi
stdatamodels              0.4.3                    pypi_0    pypi
stpipe                    0.4.2                    pypi_0    pypi
stsci-image               2.3.5                    pypi_0    pypi
stsci-imagestats          1.6.3                    pypi_0    pypi
stsci-stimage             0.2.5                    pypi_0    pypi
svo-filters               0.4.3                    pypi_0    pypi
tifffile                  2022.8.8                 pypi_0    pypi
tk                        8.6.12               h27826a3_0    conda-forge
tomli                     2.0.1                    pypi_0    pypi
toolz                     0.12.0                   pypi_0    pypi
tornado                   6.2                      pypi_0    pypi
tqdm                      4.64.0                   pypi_0    pypi
traitlets                 5.3.0                    pypi_0    pypi
tweakwcs                  0.7.4                    pypi_0    pypi
typing-extensions         4.3.0                    pypi_0    pypi
tzdata                    2022a                h191b570_0    conda-forge
uncertainties             3.1.7                    pypi_0    pypi
urllib3                   1.26.11                  pypi_0    pypi
wcwidth                   0.2.5                    pypi_0    pypi
webencodings              0.5.1                    pypi_0    pypi
wheel                     0.37.1             pyhd8ed1ab_0    conda-forge
wiimatch                  0.3.1                    pypi_0    pypi
xarray                    2022.6.0                 pypi_0    pypi
xz                        5.2.5                h516909a_1    conda-forge
zipp                      3.8.1                    pypi_0    pypi
zlib                      1.2.12               h166bdaf_2    conda-forge

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct

Changed files:
src/eureka/S3_data_reduction/plots_s3.py, src/eureka/lib/util.py, src/eureka/S3_data_reduction/s3_reduce.py

diff --git a/src/eureka/S3_data_reduction/plots_s3.py b/src/eureka/S3_data_reduction/plots_s3.py
index a9950f45..93086b02 100644
--- a/src/eureka/S3_data_reduction/plots_s3.py
+++ b/src/eureka/S3_data_reduction/plots_s3.py
@@ -3,13 +3,15 @@ import os
 from tqdm import tqdm
 import matplotlib.pyplot as plt
 import scipy.interpolate as spi
-from .source_pos import gauss
-from ..lib import util
-from ..lib.plots import figure_filetype
 import scipy.stats as stats
+from scipy.interpolate import interp1d
 from matplotlib.colors import LogNorm
 from mpl_toolkits import axes_grid1
 
+from .source_pos import gauss
+from ..lib import util
+from ..lib.plots import figure_filetype
+
 
 def lc_nodriftcorr(meta, wave_1d, optspec, optmask=None):
     '''Plot a 2D light curve without drift correction. (Fig 3101+3102)
@@ -32,6 +34,23 @@ def lc_nodriftcorr(meta, wave_1d, optspec, optmask=None):
     '''
     normspec = util.normalize_spectrum(meta, optspec.values,
                                        optmask=optmask.values)
+
+    # Save the wavelength units as the copy below will erase them
+    wave_units = wave_1d.wave_units
+
+    # For plotting purposes, extrapolate NaN wavelengths
+    wave_1d = np.ma.masked_invalid(np.copy(wave_1d))
+    if np.any(wave_1d.mask):
+        masked = np.where(wave_1d.mask)[0]
+        inds = np.arange(len(wave_1d))
+        wave_1d_valid = np.delete(wave_1d, masked)
+        inds_valid = np.delete(inds, masked)
+        # Do a spline extrapolation of third order
+        interp_fn = interp1d(inds_valid, wave_1d_valid, kind='cubic',
+                             fill_value="extrapolate", assume_sorted=True)
+        wave_1d[masked] = interp_fn(masked)
+    wave_1d = wave_1d.data
+
     wmin = np.nanmin(wave_1d)
     wmax = np.nanmax(wave_1d)
     # Don't do min and max because MIRI is backwards
@@ -68,7 +87,7 @@ def lc_nodriftcorr(meta, wave_1d, optspec, optmask=None):
         ax2.set_ylim(0, meta.n_int)
         ax1.set_ylabel('Integration Number')
         ax2.set_ylabel('Integration Number')
-        ax1.set_xlabel(r'Wavelength ($\\mu m$)')
+        ax1.set_xlabel(f'Wavelength ({wave_units})')
         ax2.set_xlabel('Detector Pixel Position')
     else:
         im1 = ax1.pcolormesh(np.arange(meta.n_int), wave_1d,
@@ -81,7 +100,7 @@ def lc_nodriftcorr(meta, wave_1d, optspec, optmask=None):
         ax2.set_ylim(pmin, pmax)
         ax1.set_xlim(0, meta.n_int)
         ax2.set_xlim(0, meta.n_int)
-        ax1.set_ylabel(r'Wavelength ($\\mu m$)')
+        ax1.set_ylabel(f'Wavelength ({wave_units})')
         ax2.set_ylabel('Detector Pixel Position')
         ax1.set_xlabel('Integration Number')
         ax2.set_xlabel('Integration Number')
diff --git a/src/eureka/S3_data_reduction/s3_reduce.py b/src/eureka/S3_data_reduction/s3_reduce.py
index 25c9b9eb..0df27d1b 100644
--- a/src/eureka/S3_data_reduction/s3_reduce.py
+++ b/src/eureka/S3_data_reduction/s3_reduce.py
@@ -302,6 +302,10 @@ def reduce(eventlabel, ecf_path=None, s2_meta=None):
                                        data.wave_2d[meta.src_ypos].values)
                     data['wave_1d'].attrs['wave_units'] = \\
                         data.wave_2d.attrs['wave_units']
+                
+                # Check for bad wavelength pixels (beyond wavelength solution)
+                util.check_nans(data.wave_1d.values, np.ones(meta.subnx), log,
+                                name='wavelength')
 
                 # Convert flux units to electrons
                 # (eg. MJy/sr -> DN -> Electrons)
diff --git a/src/eureka/lib/util.py b/src/eureka/lib/util.py
index b9b80449..c11441eb 100644
--- a/src/eureka/lib/util.py
+++ b/src/eureka/lib/util.py
@@ -159,13 +159,19 @@ def check_nans(data, mask, log, name=''):
         or infs.
     """
     data = np.ma.masked_where(mask == 0, np.copy(data))
-    num_nans = np.sum(np.ma.masked_invalid(data).mask)
+    masked = np.ma.masked_invalid(data).mask
+    inan = np.where(masked)
+    num_nans = np.sum(masked)
     num_pixels = np.size(data)
     perc_nans = 100*num_nans/num_pixels
-    if num_nans > 0:
+    if num_nans > 0 and name == 'wavelength':
+        log.writelog(f"  WARNING: Your {name} array has {num_nans} NaNs, which"
+                     f" are outside of the wavelength solution. You should "
+                     f"consider removing indices {inan} as their data quality "
+                     f"may be poor.")
+    elif num_nans > 0:
         log.writelog(f"  {name} has {num_nans} NaNs/infs, which is "
                      f"{perc_nans:.2f}% of all pixels.")
-        inan = np.where(np.ma.masked_invalid(data).mask)
         mask[inan] = 0
     if perc_nans > 10:
         log.writelog("  WARNING: Your region of interest may be off the edge "