From d124ef11913da257af501fd1c134e384d0ad3e54 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Fri, 8 Nov 2024 10:28:03 +0100
Subject: [PATCH 1/8] feat: new convenience methods on BeamInterface

---
 pyproject.toml                 |  1 +
 src/pyuvdata/analytic_beam.py  |  6 +++-
 src/pyuvdata/beam_interface.py | 65 +++++++++++++++++++++++++++++++--
 tests/test_analytic_beam.py    |  6 ++++
 tests/test_beam_interface.py   | 66 ++++++++++++++++++++++++++++++++++
 5 files changed, 140 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index fdb5cea64..a83638e8c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -113,6 +113,7 @@ select = [
 ]
 ignore = [
     "N806", # non-lowercase variable (we use N* for axes lengths)
+    "N802", # non-lowercase function name (we use N* for axes lengths)
     "B028", # no-explicit-stacklevel for warnings
     "SIM108", # prefer ternary opperators. I find them difficult to read.
     "D203", # one-blank-line-before-class. we use two.
diff --git a/src/pyuvdata/analytic_beam.py b/src/pyuvdata/analytic_beam.py
index 14e93ce3c..80a6765b5 100644
--- a/src/pyuvdata/analytic_beam.py
+++ b/src/pyuvdata/analytic_beam.py
@@ -8,7 +8,7 @@
 import dataclasses
 import importlib
 import warnings
-from dataclasses import InitVar, astuple, dataclass, field
+from dataclasses import InitVar, astuple, dataclass, field, replace
 from typing import ClassVar, Literal
 
 import numpy as np
@@ -189,6 +189,10 @@ def Npols(self):  # noqa N802
         """The number of polarizations."""
         return self.polarization_array.size
 
+    def clone(self, **kw):
+        """Create a new instance of the object with updated parameters."""
+        return replace(self, **kw)
+
     @property
     def east_ind(self):
         """The index of the east feed in the feed array."""
diff --git a/src/pyuvdata/beam_interface.py b/src/pyuvdata/beam_interface.py
index 97cf37f5a..dcc172846 100644
--- a/src/pyuvdata/beam_interface.py
+++ b/src/pyuvdata/beam_interface.py
@@ -7,7 +7,7 @@
 
 import copy
 import warnings
-from dataclasses import InitVar, dataclass
+from dataclasses import InitVar, asdict, dataclass, replace
 from typing import Literal
 
 import numpy as np
@@ -32,8 +32,9 @@ class BeamInterface:
 
     Attributes
     ----------
-    beam : pyuvdata.UVBeam or pyuvdata.AnalyticBeam
-        Beam object to use for computations
+    beam : pyuvdata.UVBeam or pyuvdata.AnalyticBeam or BeamInterface
+        Beam object to use for computations. If a BeamInterface is passed, a new
+        view of the same object is created.
     beam_type : str
         The beam type, either "efield" or "power".
     include_cross_pols : bool
@@ -59,6 +60,11 @@ def __post_init__(self, include_cross_pols: bool):
             for the power beam.
 
         """
+        if isinstance(self.beam, BeamInterface):
+            self.beam = self.beam.beam
+            self.__post_init__(include_cross_pols=include_cross_pols)
+            return
+
         if not isinstance(self.beam, UVBeam) and not issubclass(
             type(self.beam), AnalyticBeam
         ):
@@ -83,6 +89,59 @@ def __post_init__(self, include_cross_pols: bool):
                     "specify `beam_type`."
                 )
 
+    @property
+    def Npols(self):
+        """The number of polarizations in the beam."""
+        return self.beam.Npols
+
+    @property
+    def polarization_array(self):
+        """The polarizations defined on the beam."""
+        return self.beam.polarization_array
+
+    @property
+    def feed_array(self):
+        """The feeds for which the beam is defined."""
+        return self.beam.feed_array
+
+    @property
+    def Nfeeds(self):
+        """The number of feeds defined on the beam."""
+        return self.beam.Nfeeds
+
+    def clone(self, **kw):
+        """Return a new instance with updated parameters."""
+        return replace(self, **kw)
+
+    def as_power_beam(
+        self, include_cross_pols: bool = True, allow_beam_mutation: bool = False
+    ):
+        """Return a new interface instance that is in the power-beam mode.
+
+        Parameters
+        ----------
+        include_cross_pols : bool, optional
+            Whether to include cross-pols in the power beam.
+        allow_beam_mutation : bool, optional
+            Whether to allow the underlying beam to be updated in-place.
+        """
+        beam = self.beam if allow_beam_mutation else copy.deepcopy(self.beam)
+
+        # We cannot simply use .clone() here, because we need to be able to pass
+        # include_cross_pols, which can only be passed to the constructor proper.
+        this = asdict(self)
+        this["beam"] = beam
+        this["beam_type"] = "power"
+        this["include_cross_pols"] = include_cross_pols
+        return BeamInterface(**this)
+
+    def with_feeds(self, feeds):
+        """Return a new interface instance with updated feed_array."""
+        if not self._isuvbeam:
+            return self.clone(beam=self.beam.clone(feed_array=feeds))
+        new_beam = self.beam.select(feeds=feeds, inplace=False)
+        return self.clone(beam=new_beam)
+
     @property
     def _isuvbeam(self):
         return isinstance(self.beam, UVBeam)
diff --git a/tests/test_analytic_beam.py b/tests/test_analytic_beam.py
index 054efd423..2a2de3635 100644
--- a/tests/test_analytic_beam.py
+++ b/tests/test_analytic_beam.py
@@ -568,3 +568,9 @@ def test_single_feed():
     beam = GaussianBeam(diameter=14.0, feed_array=["x"], include_cross_pols=True)
     assert beam.feed_array == ["x"]
     assert beam.polarization_array == [-5]
+
+
+def test_clone():
+    beam = GaussianBeam(diameter=14.0, feed_array=["x", "y"])
+    new_beam = beam.clone(feed_array=["x"])
+    assert new_beam.feed_array == ["x"]
diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
index 62773bbf9..3e6070bec 100644
--- a/tests/test_beam_interface.py
+++ b/tests/test_beam_interface.py
@@ -243,3 +243,69 @@ def test_compute_response_errors(param, value):
     else:
         # this shouldn't error
         bi_uvb.compute_response(**compute_kwargs)
+
+
+@pytest.mark.parametrize(
+    "beam_obj",
+    [
+        AiryBeam(diameter=14.0),
+        GaussianBeam(diameter=14.0),
+        AiryBeam(diameter=14.0).to_uvbeam(
+            freq_array=np.array([1e8]), pixel_coordinate_system="healpix", nside=32
+        ),
+    ],
+)
+def test_idempotent_instantiation(beam_obj):
+    beam = BeamInterface(beam_obj)
+    beam2 = BeamInterface(beam)
+    assert beam == beam2
+
+
+def test_properties():
+    beam = AiryBeam(diameter=14.0)
+    intf = BeamInterface(beam)
+    assert beam.Npols == intf.Npols
+    assert beam.Nfeeds == intf.Nfeeds
+    assert np.all(beam.polarization_array == intf.polarization_array)
+    assert np.all(beam.feed_array == intf.feed_array)
+
+
+def test_clone():
+    beam = AiryBeam(diameter=14.0)
+    intf = BeamInterface(beam)
+    intf_clone = intf.clone(beam_type="power")
+    assert intf != intf_clone
+
+
+@pytest.mark.parametrize("uvbeam", [True, False], ids=["uvbeam", "analytic"])
+@pytest.mark.parametrize("allow_mutation", [True, False], ids=["mutate", "nomutate"])
+@pytest.mark.parametrize("include_cross_pols", [True, False], ids=["incx", "nox"])
+def test_astype(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
+    beam = AiryBeam(diameter=14.0)
+    if uvbeam:
+        beam = beam.to_uvbeam(freq_array=np.array([1e8]), nside=32)
+
+    intf = BeamInterface(beam)
+    intf_power = intf.as_power_beam(
+        allow_beam_mutation=allow_mutation, include_cross_pols=include_cross_pols
+    )
+    assert intf_power.beam_type == "power"
+    assert intf_power.Npols == 4 if include_cross_pols else 2
+
+    # Ensure that the original beam is not mutated unless we say it can be.
+    if uvbeam:
+        if allow_mutation:
+            assert intf.beam.beam_type == "power"
+        else:
+            assert intf.beam.beam_type == "efield"
+
+
+@pytest.mark.parametrize("uvbeam", [True, False])
+def test_with_feeds(uvbeam: bool):
+    beam = AiryBeam(diameter=14.0)
+    if uvbeam:
+        beam = beam.to_uvbeam(freq_array=np.array([1e8]), nside=32)
+
+    intf = BeamInterface(beam)
+    intf_feedx = intf.with_feeds(["x"])
+    assert intf_feedx.feed_array == ["x"]

From 2431d74aa517c5882d82db6184532e5b2ca36f46 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Fri, 8 Nov 2024 10:37:20 +0100
Subject: [PATCH 2/8] docs: update changelog

---
 CHANGELOG.md | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1f303896c..0ee1ba648 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@ All notable changes to this project will be documented in this file.
 - Added support for partial read for MWA correlator FITS files.
 - Added `antenna_names`, `time_range`, `lsts` and `lst_range` parameters to
 `UVFlag.select` to match UVData and UVCal select methods.
+- New convenience methods on `BeamInterface` to simplify the handling of analytic vs
+UVBeam objects.
 
 ### Changed
 - Made it possible to *not* return the `interp_basis_vector` array from beam

From ec7c6755a359b2fa2481038d29df9759a028b405 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Sat, 9 Nov 2024 14:42:23 +0100
Subject: [PATCH 3/8] test: add tests of with_feeds and as_power_beam

---
 src/pyuvdata/beam_interface.py | 62 ++++++++++++++++++++++++++++++----
 tests/test_beam_interface.py   | 41 +++++++++++++++++++++-
 2 files changed, 95 insertions(+), 8 deletions(-)

diff --git a/src/pyuvdata/beam_interface.py b/src/pyuvdata/beam_interface.py
index dcc172846..09f90ee40 100644
--- a/src/pyuvdata/beam_interface.py
+++ b/src/pyuvdata/beam_interface.py
@@ -8,12 +8,14 @@
 import copy
 import warnings
 from dataclasses import InitVar, asdict, dataclass, replace
+from itertools import product
 from typing import Literal
 
 import numpy as np
 import numpy.typing as npt
 
 from .analytic_beam import AnalyticBeam
+from .utils import pol as upol
 from .uvbeam import UVBeam
 
 # Other methods we may want to include:
@@ -75,7 +77,7 @@ def __post_init__(self, include_cross_pols: bool):
         if isinstance(self.beam, UVBeam):
             if self.beam_type is None or self.beam_type == self.beam.beam_type:
                 self.beam_type = self.beam.beam_type
-            elif self.beam_type == "power":
+            elif self.beam_type == "power" and self.beam.beam_type != "power":
                 warnings.warn(
                     "Input beam is an efield UVBeam but beam_type is specified as "
                     "'power'. Converting efield beam to power."
@@ -88,11 +90,13 @@ def __post_init__(self, include_cross_pols: bool):
                     "efield beam, either provide an efield UVBeam or do not "
                     "specify `beam_type`."
                 )
+        elif self.beam_type is None:
+            self.beam_type = "efield"
 
     @property
     def Npols(self):
         """The number of polarizations in the beam."""
-        return self.beam.Npols
+        return self.beam.Npols or len(self.polarization_array)
 
     @property
     def polarization_array(self):
@@ -107,7 +111,7 @@ def feed_array(self):
     @property
     def Nfeeds(self):
         """The number of feeds defined on the beam."""
-        return self.beam.Nfeeds
+        return self.beam.Nfeeds or len(self.feed_array)
 
     def clone(self, **kw):
         """Return a new instance with updated parameters."""
@@ -118,6 +122,10 @@ def as_power_beam(
     ):
         """Return a new interface instance that is in the power-beam mode.
 
+        If already in the power-beam mode, this is a no-op. Note that this might be
+        slighty unexpected, because the effect of `include_cross_pols` is not accounted
+        for in this case.
+
         Parameters
         ----------
         include_cross_pols : bool, optional
@@ -125,6 +133,9 @@ def as_power_beam(
         allow_beam_mutation : bool, optional
             Whether to allow the underlying beam to be updated in-place.
         """
+        if self.beam_type == "power":
+            return self
+
         beam = self.beam if allow_beam_mutation else copy.deepcopy(self.beam)
 
         # We cannot simply use .clone() here, because we need to be able to pass
@@ -133,13 +144,50 @@ def as_power_beam(
         this["beam"] = beam
         this["beam_type"] = "power"
         this["include_cross_pols"] = include_cross_pols
-        return BeamInterface(**this)
+        with warnings.catch_warnings():
+            # Don't emit the warning that we're converting to power, because that is
+            # explicitly desired.
+            warnings.simplefilter("ignore", UserWarning)
+            return BeamInterface(**this)
+
+    def with_feeds(self, feeds, *, maintain_ordering: bool = True):
+        """Return a new interface instance with updated feed_array.
 
-    def with_feeds(self, feeds):
-        """Return a new interface instance with updated feed_array."""
+        Parameters
+        ----------
+        feeds : array_like of str
+            The feeds to keep in the beam. Each value should be a string, e.g. 'n', 'x'.
+        maintain_ordering : bool, optional
+            If True, maintain the same polarization ordering as in the beam currently.
+            If False, change ordering to match the input feeds, which are turned into
+            pols (if a power beam) by using product(feeds, feeds).
+        """
         if not self._isuvbeam:
+            if maintain_ordering:
+                feeds = [fd for fd in self.feed_array if fd in feeds]
             return self.clone(beam=self.beam.clone(feed_array=feeds))
-        new_beam = self.beam.select(feeds=feeds, inplace=False)
+        if self.beam_type == "power":
+            # Down-select polarizations based on the feeds input.
+            possible_pols = [f1 + f2 for f1, f2 in product(feeds, feeds)]
+            possible_pol_ints = upol.polstr2num(
+                possible_pols, x_orientation=self.beam.x_orientation
+            )
+            if maintain_ordering:
+                use_pols = [
+                    p for p in self.beam.polarization_array if p in possible_pol_ints
+                ]
+                print("use_pols: ", use_pols)
+            else:
+                use_pols = [
+                    p for p in possible_pol_ints if p in self.beam.polarization_array
+                ]
+
+            new_beam = self.beam.select(polarizations=use_pols, inplace=False)
+        else:
+            if maintain_ordering:
+                feeds = [fd for fd in self.feed_array if fd in feeds]
+
+            new_beam = self.beam.select(feeds=feeds, inplace=False)
         return self.clone(beam=new_beam)
 
     @property
diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
index 3e6070bec..b93272bb2 100644
--- a/tests/test_beam_interface.py
+++ b/tests/test_beam_interface.py
@@ -280,7 +280,7 @@ def test_clone():
 @pytest.mark.parametrize("uvbeam", [True, False], ids=["uvbeam", "analytic"])
 @pytest.mark.parametrize("allow_mutation", [True, False], ids=["mutate", "nomutate"])
 @pytest.mark.parametrize("include_cross_pols", [True, False], ids=["incx", "nox"])
-def test_astype(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
+def test_as_power(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
     beam = AiryBeam(diameter=14.0)
     if uvbeam:
         beam = beam.to_uvbeam(freq_array=np.array([1e8]), nside=32)
@@ -300,6 +300,14 @@ def test_astype(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
             assert intf.beam.beam_type == "efield"
 
 
+def test_as_power_noop():
+    """Ensure that calling as_power_beam on a power beam is a no-op."""
+    beam = AiryBeam(diameter=14.0)
+    intf = BeamInterface(beam, beam_type="power")
+    intf2 = intf.as_power_beam()
+    assert intf is intf2
+
+
 @pytest.mark.parametrize("uvbeam", [True, False])
 def test_with_feeds(uvbeam: bool):
     beam = AiryBeam(diameter=14.0)
@@ -307,5 +315,36 @@ def test_with_feeds(uvbeam: bool):
         beam = beam.to_uvbeam(freq_array=np.array([1e8]), nside=32)
 
     intf = BeamInterface(beam)
+
     intf_feedx = intf.with_feeds(["x"])
     assert intf_feedx.feed_array == ["x"]
+
+
+def test_with_feeds_ordering():
+    beam = AiryBeam(diameter=14.0)
+    intf = BeamInterface(beam)
+
+    intf_feedx = intf.with_feeds(["y", "x"], maintain_ordering=True)
+    assert np.all(intf_feedx.feed_array == ["x", "y"])
+
+    intf_feedyx = intf.with_feeds(["y", "x"], maintain_ordering=False)
+    assert np.all(intf_feedyx.feed_array == ["y", "x"])
+
+
+@pytest.mark.filterwarnings("ignore:Input beam is an efield UVBeam")
+@pytest.mark.filterwarnings("ignore:Selected polarizations are not evenly spaced")
+def test_with_feeds_ordering_power():
+    beam = AiryBeam(diameter=14.0).to_uvbeam(freq_array=np.array([1e8]), nside=16)
+
+    intf = BeamInterface(beam, beam_type="power")
+    print(intf.polarization_array)
+    intf_feedx = intf.with_feeds(["y", "x"], maintain_ordering=True)
+    assert np.all(intf_feedx.polarization_array == [-5, -6, -7, -8])
+
+    intf_feedyx = intf.with_feeds(["y", "x"], maintain_ordering=False)
+    print(
+        utils.pol.polnum2str(
+            intf_feedyx.polarization_array, x_orientation=intf_feedyx.beam.x_orientation
+        )
+    )
+    assert np.all(intf_feedyx.polarization_array == [-6, -8, -7, -5])

From 5b378e6b21ee1b29cd6a2d95b77d75ccaf9147f2 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Mon, 11 Nov 2024 11:15:03 +0100
Subject: [PATCH 4/8] test: don't use healpix coords to go to uvbeam

---
 tests/test_beam_interface.py | 91 +++++++++++++++++-------------------
 1 file changed, 43 insertions(+), 48 deletions(-)

diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
index b93272bb2..483c3cc65 100644
--- a/tests/test_beam_interface.py
+++ b/tests/test_beam_interface.py
@@ -11,11 +11,22 @@
     GaussianBeam,
     ShortDipoleBeam,
     UniformBeam,
+    UVBeam,
     utils,
 )
 from pyuvdata.testing import check_warnings
 
 
+@pytest.fixture(scope="function")
+def airy() -> AiryBeam:
+    return AiryBeam(diameter=14.0)
+
+
+@pytest.fixture()
+def gaussian() -> GaussianBeam:
+    return GaussianBeam(diameter=14.0)
+
+
 @pytest.fixture()
 def xy_grid_coarse():
     nfreqs = 5
@@ -34,6 +45,14 @@ def xy_grid_coarse():
     return az_array, za_array, freqs
 
 
+@pytest.fixture()
+def gaussian_uv(gaussian, az_za_coords) -> UVBeam:
+    az, za = az_za_coords
+    return gaussian.to_uvbeam(
+        axis1_array=az, axis2_array=za, freq_array=np.array([1e8])
+    )
+
+
 @pytest.mark.parametrize(
     ["beam_obj", "kwargs"],
     [
@@ -245,34 +264,23 @@ def test_compute_response_errors(param, value):
         bi_uvb.compute_response(**compute_kwargs)
 
 
-@pytest.mark.parametrize(
-    "beam_obj",
-    [
-        AiryBeam(diameter=14.0),
-        GaussianBeam(diameter=14.0),
-        AiryBeam(diameter=14.0).to_uvbeam(
-            freq_array=np.array([1e8]), pixel_coordinate_system="healpix", nside=32
-        ),
-    ],
-)
-def test_idempotent_instantiation(beam_obj):
-    beam = BeamInterface(beam_obj)
+@pytest.mark.parametrize("beam_obj", ["airy", "gaussian", "gaussian_uv"])
+def test_idempotent_instantiation(beam_obj, request):
+    beam = BeamInterface(request.getfixturevalue(beam_obj))
     beam2 = BeamInterface(beam)
     assert beam == beam2
 
 
-def test_properties():
-    beam = AiryBeam(diameter=14.0)
-    intf = BeamInterface(beam)
-    assert beam.Npols == intf.Npols
-    assert beam.Nfeeds == intf.Nfeeds
-    assert np.all(beam.polarization_array == intf.polarization_array)
-    assert np.all(beam.feed_array == intf.feed_array)
+def test_properties(airy: AiryBeam):
+    intf = BeamInterface(airy)
+    assert airy.Npols == intf.Npols
+    assert airy.Nfeeds == intf.Nfeeds
+    assert np.all(airy.polarization_array == intf.polarization_array)
+    assert np.all(airy.feed_array == intf.feed_array)
 
 
-def test_clone():
-    beam = AiryBeam(diameter=14.0)
-    intf = BeamInterface(beam)
+def test_clone(airy):
+    intf = BeamInterface(airy)
     intf_clone = intf.clone(beam_type="power")
     assert intf != intf_clone
 
@@ -280,11 +288,10 @@ def test_clone():
 @pytest.mark.parametrize("uvbeam", [True, False], ids=["uvbeam", "analytic"])
 @pytest.mark.parametrize("allow_mutation", [True, False], ids=["mutate", "nomutate"])
 @pytest.mark.parametrize("include_cross_pols", [True, False], ids=["incx", "nox"])
-def test_as_power(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
-    beam = AiryBeam(diameter=14.0)
-    if uvbeam:
-        beam = beam.to_uvbeam(freq_array=np.array([1e8]), nside=32)
-
+def test_as_power(
+    uvbeam: bool, allow_mutation: bool, include_cross_pols: bool, gaussian, gaussian_uv
+):
+    beam = gaussian_uv if uvbeam else gaussian
     intf = BeamInterface(beam)
     intf_power = intf.as_power_beam(
         allow_beam_mutation=allow_mutation, include_cross_pols=include_cross_pols
@@ -292,7 +299,6 @@ def test_as_power(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
     assert intf_power.beam_type == "power"
     assert intf_power.Npols == 4 if include_cross_pols else 2
 
-    # Ensure that the original beam is not mutated unless we say it can be.
     if uvbeam:
         if allow_mutation:
             assert intf.beam.beam_type == "power"
@@ -300,19 +306,16 @@ def test_as_power(uvbeam: bool, allow_mutation: bool, include_cross_pols: bool):
             assert intf.beam.beam_type == "efield"
 
 
-def test_as_power_noop():
+def test_as_power_noop(airy):
     """Ensure that calling as_power_beam on a power beam is a no-op."""
-    beam = AiryBeam(diameter=14.0)
-    intf = BeamInterface(beam, beam_type="power")
+    intf = BeamInterface(airy, beam_type="power")
     intf2 = intf.as_power_beam()
     assert intf is intf2
 
 
 @pytest.mark.parametrize("uvbeam", [True, False])
-def test_with_feeds(uvbeam: bool):
-    beam = AiryBeam(diameter=14.0)
-    if uvbeam:
-        beam = beam.to_uvbeam(freq_array=np.array([1e8]), nside=32)
+def test_with_feeds(uvbeam: bool, gaussian, gaussian_uv):
+    beam = gaussian_uv if uvbeam else gaussian
 
     intf = BeamInterface(beam)
 
@@ -320,9 +323,8 @@ def test_with_feeds(uvbeam: bool):
     assert intf_feedx.feed_array == ["x"]
 
 
-def test_with_feeds_ordering():
-    beam = AiryBeam(diameter=14.0)
-    intf = BeamInterface(beam)
+def test_with_feeds_ordering(airy):
+    intf = BeamInterface(airy)
 
     intf_feedx = intf.with_feeds(["y", "x"], maintain_ordering=True)
     assert np.all(intf_feedx.feed_array == ["x", "y"])
@@ -333,18 +335,11 @@ def test_with_feeds_ordering():
 
 @pytest.mark.filterwarnings("ignore:Input beam is an efield UVBeam")
 @pytest.mark.filterwarnings("ignore:Selected polarizations are not evenly spaced")
-def test_with_feeds_ordering_power():
-    beam = AiryBeam(diameter=14.0).to_uvbeam(freq_array=np.array([1e8]), nside=16)
-
-    intf = BeamInterface(beam, beam_type="power")
-    print(intf.polarization_array)
+def test_with_feeds_ordering_power(gaussian_uv):
+    # beam = AiryBeam(diameter=14.0).to_uvbeam(freq_array=np.array([1e8]), nside=16)
+    intf = BeamInterface(gaussian_uv, beam_type="power")
     intf_feedx = intf.with_feeds(["y", "x"], maintain_ordering=True)
     assert np.all(intf_feedx.polarization_array == [-5, -6, -7, -8])
 
     intf_feedyx = intf.with_feeds(["y", "x"], maintain_ordering=False)
-    print(
-        utils.pol.polnum2str(
-            intf_feedyx.polarization_array, x_orientation=intf_feedyx.beam.x_orientation
-        )
-    )
     assert np.all(intf_feedyx.polarization_array == [-6, -8, -7, -5])

From 304c63dbf964a050709a0ab492d541ac12c25d78 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Wed, 13 Nov 2024 13:59:51 +0100
Subject: [PATCH 5/8] fix: updates to Bryna's comments

---
 CHANGELOG.md                   |  3 +--
 pyproject.toml                 |  1 -
 src/pyuvdata/beam_interface.py | 28 ++++++++++++++++++++++------
 3 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0ee1ba648..f31557d8e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,11 +4,10 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
+- New convenience methods on `BeamInterface` to simplify the handling of analytic vs UVBeam objects.
 - Added support for partial read for MWA correlator FITS files.
 - Added `antenna_names`, `time_range`, `lsts` and `lst_range` parameters to
 `UVFlag.select` to match UVData and UVCal select methods.
-- New convenience methods on `BeamInterface` to simplify the handling of analytic vs
-UVBeam objects.
 
 ### Changed
 - Made it possible to *not* return the `interp_basis_vector` array from beam
diff --git a/pyproject.toml b/pyproject.toml
index a83638e8c..fdb5cea64 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -113,7 +113,6 @@ select = [
 ]
 ignore = [
     "N806", # non-lowercase variable (we use N* for axes lengths)
-    "N802", # non-lowercase function name (we use N* for axes lengths)
     "B028", # no-explicit-stacklevel for warnings
     "SIM108", # prefer ternary opperators. I find them difficult to read.
     "D203", # one-blank-line-before-class. we use two.
diff --git a/src/pyuvdata/beam_interface.py b/src/pyuvdata/beam_interface.py
index 09f90ee40..a8e553447 100644
--- a/src/pyuvdata/beam_interface.py
+++ b/src/pyuvdata/beam_interface.py
@@ -77,7 +77,7 @@ def __post_init__(self, include_cross_pols: bool):
         if isinstance(self.beam, UVBeam):
             if self.beam_type is None or self.beam_type == self.beam.beam_type:
                 self.beam_type = self.beam.beam_type
-            elif self.beam_type == "power" and self.beam.beam_type != "power":
+            elif self.beam_type == "power":
                 warnings.warn(
                     "Input beam is an efield UVBeam but beam_type is specified as "
                     "'power'. Converting efield beam to power."
@@ -94,7 +94,7 @@ def __post_init__(self, include_cross_pols: bool):
             self.beam_type = "efield"
 
     @property
-    def Npols(self):
+    def Npols(self):  # noqa N802
         """The number of polarizations in the beam."""
         return self.beam.Npols or len(self.polarization_array)
 
@@ -109,7 +109,7 @@ def feed_array(self):
         return self.beam.feed_array
 
     @property
-    def Nfeeds(self):
+    def Nfeeds(self):  # noqa N802
         """The number of feeds defined on the beam."""
         return self.beam.Nfeeds or len(self.feed_array)
 
@@ -118,7 +118,7 @@ def clone(self, **kw):
         return replace(self, **kw)
 
     def as_power_beam(
-        self, include_cross_pols: bool = True, allow_beam_mutation: bool = False
+        self, include_cross_pols: bool | None = None, allow_beam_mutation: bool = False
     ):
         """Return a new interface instance that is in the power-beam mode.
 
@@ -129,13 +129,30 @@ def as_power_beam(
         Parameters
         ----------
         include_cross_pols : bool, optional
-            Whether to include cross-pols in the power beam.
+            Whether to include cross-pols in the power beam. By default, this is True
+            for E-field beams, and takes the same value as the existing beam if the
+            existing beam is already a power beam.
         allow_beam_mutation : bool, optional
             Whether to allow the underlying beam to be updated in-place.
         """
         if self.beam_type == "power":
+            if include_cross_pols is None:
+                # By default, keep the value of include_cross_pols the same.
+                include_cross_pols = self.Npols > 2
+
+            if self.Npols > 1 and (
+                (include_cross_pols and self.Npols != 4)
+                or (not include_cross_pols and self.Npols != 4)
+            ):
+                warnings.warn(
+                    "as_power_beam does not modify cross pols when the beam is"
+                    "already in power mode!"
+                )
             return self
 
+        if include_cross_pols is None:
+            include_cross_pols = True
+
         beam = self.beam if allow_beam_mutation else copy.deepcopy(self.beam)
 
         # We cannot simply use .clone() here, because we need to be able to pass
@@ -176,7 +193,6 @@ def with_feeds(self, feeds, *, maintain_ordering: bool = True):
                 use_pols = [
                     p for p in self.beam.polarization_array if p in possible_pol_ints
                 ]
-                print("use_pols: ", use_pols)
             else:
                 use_pols = [
                     p for p in possible_pol_ints if p in self.beam.polarization_array

From 12e72dc9393945f1b75a5ad507d7eaca573d1d03 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Wed, 13 Nov 2024 19:17:20 +0100
Subject: [PATCH 6/8] test: cover all  new lines

---
 tests/test_beam_interface.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
index 483c3cc65..9847fb0ee 100644
--- a/tests/test_beam_interface.py
+++ b/tests/test_beam_interface.py
@@ -287,7 +287,7 @@ def test_clone(airy):
 
 @pytest.mark.parametrize("uvbeam", [True, False], ids=["uvbeam", "analytic"])
 @pytest.mark.parametrize("allow_mutation", [True, False], ids=["mutate", "nomutate"])
-@pytest.mark.parametrize("include_cross_pols", [True, False], ids=["incx", "nox"])
+@pytest.mark.parametrize("include_cross_pols", [True, False, None], ids=["incx", "nox"])
 def test_as_power(
     uvbeam: bool, allow_mutation: bool, include_cross_pols: bool, gaussian, gaussian_uv
 ):
@@ -296,6 +296,9 @@ def test_as_power(
     intf_power = intf.as_power_beam(
         allow_beam_mutation=allow_mutation, include_cross_pols=include_cross_pols
     )
+    if include_cross_pols is None:
+        include_cross_pols = True
+
     assert intf_power.beam_type == "power"
     assert intf_power.Npols == 4 if include_cross_pols else 2
 
@@ -312,6 +315,10 @@ def test_as_power_noop(airy):
     intf2 = intf.as_power_beam()
     assert intf is intf2
 
+    with pytest.warns(UserWarning, match="as_power_beam does not modify cross pols"):
+        intf2 = intf.as_power_beam(include_cross_pols=False)
+    assert intf is intf2
+
 
 @pytest.mark.parametrize("uvbeam", [True, False])
 def test_with_feeds(uvbeam: bool, gaussian, gaussian_uv):

From 4d25cf669ed156bda6a759a78bd5b794b2d840ce Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Wed, 13 Nov 2024 19:57:09 +0100
Subject: [PATCH 7/8] test: fix tests that I broke

---
 src/pyuvdata/beam_interface.py | 2 +-
 tests/test_beam_interface.py   | 5 ++++-
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/src/pyuvdata/beam_interface.py b/src/pyuvdata/beam_interface.py
index a8e553447..c5595e915 100644
--- a/src/pyuvdata/beam_interface.py
+++ b/src/pyuvdata/beam_interface.py
@@ -142,7 +142,7 @@ def as_power_beam(
 
             if self.Npols > 1 and (
                 (include_cross_pols and self.Npols != 4)
-                or (not include_cross_pols and self.Npols != 4)
+                or (not include_cross_pols and self.Npols != 2)
             ):
                 warnings.warn(
                     "as_power_beam does not modify cross pols when the beam is"
diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
index 9847fb0ee..e8e73f140 100644
--- a/tests/test_beam_interface.py
+++ b/tests/test_beam_interface.py
@@ -287,7 +287,9 @@ def test_clone(airy):
 
 @pytest.mark.parametrize("uvbeam", [True, False], ids=["uvbeam", "analytic"])
 @pytest.mark.parametrize("allow_mutation", [True, False], ids=["mutate", "nomutate"])
-@pytest.mark.parametrize("include_cross_pols", [True, False, None], ids=["incx", "nox"])
+@pytest.mark.parametrize(
+    "include_cross_pols", [True, False, None], ids=["incx", "nox", "xpolnone"]
+)
 def test_as_power(
     uvbeam: bool, allow_mutation: bool, include_cross_pols: bool, gaussian, gaussian_uv
 ):
@@ -316,6 +318,7 @@ def test_as_power_noop(airy):
     assert intf is intf2
 
     with pytest.warns(UserWarning, match="as_power_beam does not modify cross pols"):
+        print(intf.Npols)
         intf2 = intf.as_power_beam(include_cross_pols=False)
     assert intf is intf2
 

From 5a23a953f5d227d6cca85694d7149f310f9a22f6 Mon Sep 17 00:00:00 2001
From: Steven Murray <steven.g.murray@asu.edu>
Date: Sun, 17 Nov 2024 10:36:30 +0100
Subject: [PATCH 8/8] docs: update docstring with info about what the user
 asked for

---
 src/pyuvdata/beam_interface.py | 5 ++++-
 tests/test_beam_interface.py   | 1 -
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/pyuvdata/beam_interface.py b/src/pyuvdata/beam_interface.py
index c5595e915..0bd0bde9a 100644
--- a/src/pyuvdata/beam_interface.py
+++ b/src/pyuvdata/beam_interface.py
@@ -146,7 +146,10 @@ def as_power_beam(
             ):
                 warnings.warn(
                     "as_power_beam does not modify cross pols when the beam is"
-                    "already in power mode!"
+                    f"already in power mode! You have polarizations: "
+                    f"{self.polarization_array} but asked to "
+                    f"*{'include' if include_cross_pols else 'not include'}* "
+                    "cross-pols."
                 )
             return self
 
diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
index e8e73f140..1c23a88ca 100644
--- a/tests/test_beam_interface.py
+++ b/tests/test_beam_interface.py
@@ -318,7 +318,6 @@ def test_as_power_noop(airy):
     assert intf is intf2
 
     with pytest.warns(UserWarning, match="as_power_beam does not modify cross pols"):
-        print(intf.Npols)
         intf2 = intf.as_power_beam(include_cross_pols=False)
     assert intf is intf2