Beam plotting tweaks

CHANGELOG.md

@@ -4,6 +4,8 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
+- A colormap option to the UVBeam and AnalyticBeam `plot` methods, along with
+better default colormaps for phase.
 - A NotImplementedError to `mwa_corr_fits.py` that is thrown when trying to read
 fringe-stopped data.
 

src/pyuvdata/analytic_beam.py

@@ -646,6 +646,7 @@ def plot(
         beam_type: str,
         freq: float,
         complex_type: str = "real",
+        colormap: str | None = None,
         logcolor: bool | None = None,
         plt_kwargs: dict | None = None,
         norm_kwargs: dict | None = None,
@@ -666,6 +667,11 @@ def plot(
             What to plot for complex beams, options are: [real, imag, abs, phase].
             Defaults to "real" for complex beams. Ignored for real beams
             (i.e. power beams, same feed).
+        colormap : str, optional
+            Matplotlib colormap to use. Defaults to "twlight" if complex_type="phase"
+            and logcolor=False, otherwise it defaults to "viridis" if the data to be
+            plotted is positive definite (e.g. if complex_type="abs") and "PRGn"
+            otherwise.
         logcolor : bool, optional
             Option to use log scaling for the color. Defaults to True for power
             beams and False for E-field beams. Results in using
@@ -690,6 +696,7 @@ def plot(
             beam_type=beam_type,
             freq=freq,
             complex_type=complex_type,
+            colormap=colormap,
             logcolor=logcolor,
             plt_kwargs=plt_kwargs,
             norm_kwargs=norm_kwargs,