Run the tutorial plotting code in doctests

To prevent bugs in the plotting code.
Also include the plots in the docs.

README.md

@@ -210,9 +210,10 @@ you will need the following packages:
 * pytest >= 6.2
 * pytest-cases >= 3.6.9
 * pytest-cov
-* cython >=0.23  (This is necessary for coverage reporting of cython extensions)
+* cython >=0.23
 * coverage
 * pre-commit
+* matplotlib
 * sphinx
 * pypandoc
 

ci/pyuvdata_tests.yml

@@ -8,6 +8,7 @@ dependencies:
   - docstring_parser>=0.15
   - h5py>=3.1
   - hdf5plugin>=3.1.0
+  - matplotlib  # this is just for the doctests.
   - numpy>=1.20.*
   - pyerfa>=2.0
   - python-casacore>=3.3.1

docs/conftest.py

@@ -17,6 +17,8 @@ def setup_and_teardown_package(tmp_path_factory):
     tmp_path = tmp_path_factory.mktemp("uvdata_tests")
     try:
         os.chdir(tmp_path)
+        if not os.path.exists("Images"):
+            os.makedirs("Images")
         yield
     finally:
         os.chdir(cwd)

docs/uvbeam_tutorial.rst

@@ -288,18 +288,27 @@ a) Selecting a range of Zenith Angles
   >>> import numpy as np
   >>> from pyuvdata import UVBeam
   >>> from pyuvdata.data import DATA_PATH
-  >>> import matplotlib.pyplot as plt # doctest: +SKIP
+  >>> import matplotlib.pyplot as plt
   >>> settings_file = os.path.join(DATA_PATH, 'NicCSTbeams/NicCSTbeams.yaml')
   >>> beam = UVBeam.from_file(settings_file, beam_type='power', use_future_array_shapes=True)
+  >>> # Make a new object with a reduced zenith angle range with the select method
   >>> new_beam = beam.select(axis2_inds=np.arange(0, 20), inplace=False)
 
   >>> # plot zenith angle cut through beams
-  >>> plt.plot(beam.axis2_array, beam.data_array[0, 0, 0, 0, :, 0], # doctest: +SKIP
-  ...         new_beam.axis2_array, new_beam.data_array[0, 0, 0, 0, :, 0], 'r')
-  >>> plt.xscale('log') # doctest: +SKIP
-  >>> plt.xlabel('Zenith Angle (radians)') # doctest: +SKIP
-  >>> plt.ylabel('Power') # doctest: +SKIP
-  >>> plt.show() # doctest: +SKIP
+  >>> fig, ax = plt.subplots(1, 1)
+  >>> _ = ax.plot(np.rad2deg(beam.axis2_array), beam.data_array[0, 0, 0, :, 0], label="original")
+  >>> _ = ax.plot(np.rad2deg(new_beam.axis2_array), new_beam.data_array[0, 0, 0, :, 0], 'r', label="cut down")
+  >>> _ = ax.set_xscale('log')
+  >>> _ = ax.set_yscale('log')
+  >>> _ = ax.set_xlabel('Zenith Angle (degrees)')
+  >>> _ = ax.set_ylabel('Power')
+  >>> _ = fig.legend(loc="upper right", bbox_to_anchor=[0.9,0.88])
+  >>> plt.show()  # doctest: +SKIP
+  >>> plt.savefig("Images/beam_cut.png", bbox_inches='tight')
+  >>> plt.clf()
+
+.. image:: Images/beam_cut.png
+  :width: 600
 
 a) Selecting Feeds or Polarizations
 ***********************************

docs/uvcal_tutorial.rst

@@ -93,7 +93,7 @@ a) Reading a cal fits gain calibration file.
 
   >>> import os
   >>> import numpy as np
-  >>> import matplotlib.pyplot as plt # doctest: +SKIP
+  >>> import matplotlib.pyplot as plt
   >>> from pyuvdata import UVCal
   >>> from pyuvdata.data import DATA_PATH
   >>> filename = os.path.join(DATA_PATH, 'zen.2457698.40355.xx.gain.calfits')
@@ -124,12 +124,18 @@ a) Reading a cal fits gain calibration file.
   (19, 10, 5, 1)
 
   >>> # plot abs of all gains for first time and first jones component.
-  >>> for ant in range(cal.Nants_data): # doctest: +SKIP
-  ...    plt.plot(cal.freq_array.flatten(), np.abs(cal.gain_array[ant, 0, :, 0, 0]))
-  >>> plt.xlabel('Frequency (Hz)') # doctest: +SKIP
-  >>> plt.ylabel('Abs(gains)') # doctest: +SKIP
+  >>> fig, ax = plt.subplots(1, 1)
+  >>> for ant in range(cal.Nants_data):
+  ...    _ = ax.plot(cal.freq_array.flatten(), np.abs(cal.gain_array[ant, :, 0, 0]), label=f"ant {ant}")
+  >>> _ = ax.set_xlabel('Frequency (Hz)')
+  >>> _ = ax.set_ylabel('Abs(gains)')
+  >>> _ = fig.legend(bbox_to_anchor=(1.08, 0.5), loc="outside center right")
   >>> plt.show() # doctest: +SKIP
+  >>> plt.savefig("Images/abs_gains.png", bbox_inches='tight')
+  >>> plt.clf()
 
+.. image:: Images/abs_gains.png
+    :width: 600
 
 b) FHD cal to cal fits
 ***********************

docs/uvdata_tutorial.rst

@@ -800,37 +800,61 @@ Note: there is now support for reading in only part of a uvfits, uvh5 or miriad
 .. code-block:: python
 
   >>> import os
+  >>> from astropy.time import Time
   >>> import numpy as np
-  >>> import matplotlib.pyplot as plt # doctest: +SKIP
+  >>> import matplotlib.pyplot as plt
   >>> from pyuvdata import UVData
   >>> from pyuvdata.data import DATA_PATH
-  >>> filename = os.path.join(DATA_PATH, 'day2_TDEM0003_10s_norx_1src_1spw.uvfits')
+  >>> filename = os.path.join(DATA_PATH, 'day2_TDEM0003_10s_norx_1scan.uvfits')
   >>> uvd = UVData.from_file(filename, use_future_array_shapes=True)
 
   >>> print(uvd.data_array.shape)
-  (1360, 64, 4)
+  (1414, 128, 4)
   >>> print(uvd.Ntimes)
-  15
+  17
   >>> print(uvd.Nfreqs)
-  64
-  >>> bl = uvd.antnums_to_baseline(1, 2)
-  >>> print(bl)
-  67586
-  >>> bl_ind = np.where(uvd.baseline_array == bl)[0]
+  128
+  >>> # get the data for a single baseline and polarization
+  >>> waterfall_data = uvd.get_data((1, 2, uvd.polarization_array[0]))
+  >>> # get the corresponding times for this waterfall
+  >>> waterfall_times = Time(uvd.get_times((1, 2, uvd.polarization_array[0])), format="jd").iso
 
   >>> # Amplitude waterfall for all spectral channels and 0th polarization
-  >>> plt.imshow(np.abs(uvd.data_array[bl_ind, :, 0])) # doctest: +SKIP
-  >>> plt.show() # doctest: +SKIP
+  >>> fig, ax = plt.subplots(1, 1)
+  >>> _ = ax.imshow(np.abs(waterfall_data), interpolation='none')
+  >>> _ = ax.set_yticks([0, waterfall_times.size - 1])
+  >>> _ = ax.set_yticklabels([waterfall_times[0], waterfall_times[1]])
+  >>> freq_tick_inds = np.concatenate((np.arange(0, uvd.Nfreqs, 16), [uvd.Nfreqs-1]))
+  >>> _ = ax.set_xticks(freq_tick_inds)
+  >>> _ = ax.set_xticklabels([f"{val:.3f}" for val in uvd.freq_array[freq_tick_inds]*1e-9])
+  >>> _ = ax.set_xlabel("Frequency (GHz)")
+  >>> fig.show() # doctest: +SKIP
+  >>> plt.savefig("Images/amplitude_waterfall.png", bbox_inches='tight')
+  >>> plt.clf()
+
+.. image:: Images/amplitude_waterfall.png
+    :width: 600
 
-  >>> # If using flexible spectral windows (flex_spw=True), plot the 0th window, 0th pol
-  >>> if uvd.flex_spw: # doctest: +SKIP
-  ...     window_sel = uvd.flex_spw_id_array == 0 # doctest: +SKIP
-  ...     plt.imshow(np.abs(uvd.data_array[bl_ind, 0, window_sel, 0])) # doctest: +SKIP
-  ...     plt.show() # doctest: +SKIP
+.. code-block:: python
 
-  >>> # Update: With new UI features, making waterfalls is easier than ever!
-  >>> plt.imshow(np.abs(uvd.get_data((1, 2, uvd.polarization_array[0])))) # doctest: +SKIP
+  >>> # The plot above has a discontinuity in the frequency axis because this
+  >>> # data set has spectral windows. Let's just plot the 0th one
+  >>> window_sel = uvd.flex_spw_id_array == 0
+  >>> freqs_use = uvd.freq_array[window_sel]
+  >>> fig, ax = plt.subplots(1, 1)
+  >>> _ = plt.imshow(np.abs(waterfall_data[:,  window_sel]))
+  >>> _ = ax.set_yticks([0, waterfall_times.size - 1])
+  >>> _ = ax.set_yticklabels([waterfall_times[0], waterfall_times[1]])
+  >>> freq_tick_inds = np.concatenate((np.arange(0, freqs_use.size, 16), [freqs_use.size-1]))
+  >>> _ = ax.set_xticks(freq_tick_inds)
+  >>> _ = ax.set_xticklabels([f"{val:.3f}" for val in freqs_use[freq_tick_inds]*1e-9])
+  >>> _ = ax.set_xlabel("Frequency (GHz)")
   >>> plt.show() # doctest: +SKIP
+  >>> plt.savefig("Images/amplitude_waterfall_spw0.png", bbox_inches='tight')
+  >>> plt.clf()
+
+.. image:: Images/amplitude_waterfall_spw0.png
+    :width: 600
 
 
 UVData: Location conversions

environment.yaml

@@ -10,6 +10,7 @@ dependencies:
   - docstring_parser>=0.15
   - h5py>=3.1
   - hdf5plugin>=3.1.0
+  - matplotlib
   - numpy>=1.20
   - pip
   - pre-commit

setup.py

@@ -127,7 +127,7 @@ def is_platform_windows():
     "coverage",
     "pre-commit",
 ]
-doc_reqs = ["sphinx", "pypandoc"]
+doc_reqs = ["matplotlib", "sphinx", "pypandoc"]
 
 setup_args = {
     "name": "pyuvdata",