generalizing jones matrix for Nfeeds=1,2

RadioAstronomySoftwareGroup · Jan 28, 2025 · a81053d · a81053d
1 parent c93c507
commit a81053d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 11 deletions.
diff --git a/src/pyuvsim/antenna.py b/src/pyuvsim/antenna.py
@@ -53,7 +53,7 @@ def get_beam_jones(
         """
         Calculate the jones matrix for this antenna in the direction of sources.
 
-        A 2 x 2 x Nsources array of Efield vectors in Az/Alt.
+        A Nfeeds x 2 x Nsources array of Efield vectors in Az/Alt.
 
         Parameters
         ----------
@@ -82,7 +82,7 @@ def get_beam_jones(
         Returns
         -------
         jones_matrix : array_like of float
-            Jones matricies for each source location, shape (2,2, Ncomponents). The
+            Jones matricies for each source location, shape (Nfeeds,2, Ncomponents). The
             first axis is feed, the second axis is vector component on the sky in az/za.
 
         """
@@ -137,17 +137,14 @@ def get_beam_jones(
         interp_data = bi.compute_response(**interp_kwargs)
 
         Ncomponents = source_za.shape[-1]
-
+        Nfeeds = interp_data.shape[1]
         # interp_data has shape:
         #   (Naxes_vec, Nfeeds, 1 (freq),  Ncomponents (source positions))
-        jones_matrix = np.zeros((2, 2, Ncomponents), dtype=complex)
+        jones_matrix = np.zeros((Nfeeds, 2, Ncomponents), dtype=complex)
 
         # first axis is feed, second axis is theta, phi (opposite order of beam!)
-        jones_matrix[0, 0] = interp_data[1, 0, 0, :]
-        jones_matrix[1, 1] = interp_data[0, 1, 0, :]
-        jones_matrix[0, 1] = interp_data[0, 0, 0, :]
-        jones_matrix[1, 0] = interp_data[1, 1, 0, :]
-
+        jones_matrix[:, 0] = interp_data[1, :, 0, :]
+        jones_matrix[:, 1] = interp_data[0, :, 0, :]
         return jones_matrix
 
     def __eq__(self, other):

diff --git a/src/pyuvsim/uvsim.py b/src/pyuvsim/uvsim.py
@@ -391,8 +391,13 @@ def make_visibility(self):
         # Sum over source component axis:
         vij = np.sum(vij, axis=2)
 
-        # Reshape to be [xx, yy, xy, yx]
-        vis_vector = np.asarray([vij[0, 0], vij[1, 1], vij[0, 1], vij[1, 0]])
+        Nfeeds = vij.shape[0]
+        assert (Nfeeds<=2)
+        # Reshape to be [xx] for one feed or [xx, yy, xy, yx] for two feeds
+        if Nfeeds == 1:
+            vis_vector = np.asarray([vij[0,0]])
+        elif Nfeeds == 2:
+            vis_vector = np.asarray([vij[0, 0], vij[1, 1], vij[0, 1], vij[1, 0]])
         return vis_vector