Proper tilting of input antenna positions when z-component is non-zero

.coveragerc

@@ -1,5 +1,6 @@
 [run]
 include = fftvis/*
+omit = fftvis/logutils.py
 branch = True
 
 [report]

fftvis/simulate.py

@@ -36,6 +36,7 @@ def simulate_vis(
     eps: float = None,
     use_feed: str = "x",
     flat_array_tol: float = 0.0,
+    live_progress: bool = True,
 ):
     """
     Parameters:
@@ -80,6 +81,8 @@ def simulate_vis(
         Tolerance for checking if the array is flat in units of meters. If the
         z-coordinate of all baseline vectors is within this tolerance, the array
         is considered flat and the z-coordinate is set to zero. Default is 0.0.
+    live_progress : bool, default = True
+        Whether to show progress bar during simulation.
 
     Returns:
     -------
@@ -115,6 +118,7 @@ def simulate_vis(
         polarized=polarized,
         eps=eps,
         flat_array_tol=flat_array_tol,
+        live_progress=live_progress,
     )
 
 
@@ -241,14 +245,24 @@ def simulate(
     # Factor of 0.5 accounts for splitting Stokes between polarization channels
     Isky = (0.5 * fluxes).astype(complex_dtype)
 
+    # Flatten antenna positions
+    antkey_to_idx = dict(zip(ants.keys(), range(len(ants))))
+    antvecs = np.array([ants[ant] for ant in ants], dtype=real_dtype)
+
+    # Rotate the array to the xy-plane
+    rotation_matrix = utils.get_plane_to_xy_rotation_matrix(antvecs)
+    rotation_matrix = rotation_matrix.astype(real_dtype)
+    rotated_antvecs = np.dot(rotation_matrix.T, antvecs.T)
+    rotated_ants = {ant: rotated_antvecs[:, antkey_to_idx[ant]] for ant in ants}
+
     # Compute baseline vectors
-    blx, bly, blz = np.array([ants[bl[1]] - ants[bl[0]] for bl in baselines])[
+    blx, bly, blz = np.array([rotated_ants[bl[1]] - rotated_ants[bl[0]] for bl in baselines])[
         :, :
     ].T.astype(real_dtype)
 
     # Check if the array is flat within tolerance
     is_coplanar = np.all(np.less_equal(np.abs(blz), flat_array_tol))
-
+    
     # Generate visibility array
     if expand_vis:
         vis = np.zeros(
@@ -308,6 +322,9 @@ def simulate(
             # Compute azimuth and zenith angles
             az, za = conversions.enu_to_az_za(enu_e=tx, enu_n=ty, orientation="uvbeam")
 
+            # Rotate source coordinates with rotation matrix.
+            tx, ty, tz = np.dot(rotation_matrix.T, [tx, ty, tz])
+
             for fi in range(nfreqs):
                 # Compute uv coordinates
                 u[:], v[:], w[:] = blx * freqs[fi], bly * freqs[fi], blz * freqs[fi]

fftvis/tests/test_compare_matvis.py

@@ -138,6 +138,7 @@ def test_simulate_non_coplanar():
 
     # Set up the antenna positions
     antpos = {k: np.array([k * 10, 0, k]) for k in range(nants)}
+    antpos_flat = {k: np.array([k * 10, 0, 0]) for k in range(nants)}
 
     # Define a Gaussian beam
     beam = AnalyticBeam("gaussian", diameter=14.0)
@@ -151,6 +152,9 @@ def test_simulate_non_coplanar():
     mvis = matvis.simulate_vis(
         antpos, sky_model, ra, dec, freqs, lsts, beams=[beam], precision=2
     )
+    mvis_flat = matvis.simulate_vis(
+        antpos_flat, sky_model, ra, dec, freqs, lsts, beams=[beam], precision=2
+    )
 
     # Use fftvis to simulate visibilities
     fvis = simulate.simulate_vis(
@@ -160,10 +164,5 @@ def test_simulate_non_coplanar():
     # Check that the results are the same
     assert np.allclose(mvis, fvis, atol=1e-5)
 
-    # Use fftvis to simulate visibilities with flat array
-    fvis_flat = simulate.simulate_vis(
-        antpos, sky_model, ra, dec, freqs, lsts, beam, precision=2, eps=1e-10, flat_array_tol=np.inf
-    )
-
     # Check that the results are different
-    assert not np.allclose(fvis, fvis_flat, atol=1e-5)
+    assert not np.allclose(mvis_flat, fvis, atol=1e-5)

fftvis/utils.py

@@ -1,5 +1,5 @@
 import numpy as np
-
+from scipy import linalg
 IDEALIZED_BL_TOL = 1e-8  # bl_error_tol for redcal.get_reds when using antenna positions calculated from reds
 speed_of_light = 299792458.0  # m/s
 
@@ -65,3 +65,51 @@ def get_pos_reds(antpos, decimals=3, include_autos=True):
             reds_list.append(red)
 
     return reds_list
+
+
+def get_plane_to_xy_rotation_matrix(antvecs):
+    """
+    Compute the rotation matrix that projects the antenna positions onto the xy-plane. 
+    This function is used to rotate the antenna positions so that they lie in the xy-plane.
+
+    Parameters:
+    ----------
+        antvecs: np.array
+            Array of antenna positions in the form (Nants, 3).
+
+    Returns:
+    -------
+        rotation_matrix: np.array
+            Rotation matrix that projects the antenna positions onto the xy-plane of shape (3, 3).
+    """
+    # Fit a plane to the antenna positions
+    antx, anty, antz = antvecs.T
+    basis = np.array([antx, anty, np.ones_like(antz)]).T
+    plane, res, rank, s = linalg.lstsq(basis, antz)
+
+    # Project the antenna positions onto the plane
+    slope_x, slope_y, z_offset = plane
+
+    # Plane is already approximately aligned with the xy-axes, 
+    # return identity rotation matrix
+    if np.isclose(slope_x, 0) and np.isclose(slope_y, 0.0):
+        return np.eye(3)
+
+    # Normalize the normal vector
+    normal = np.array([slope_x, slope_y, -1])
+    normal = normal / np.linalg.norm(normal)
+
+    # Compute the rotation axis
+    axis = np.array([slope_y, -slope_x, 0])
+    axis = axis / np.linalg.norm(axis)
+
+    # Compute the rotation angle
+    theta = np.arccos(-normal[2])
+
+    # Compute the rotation matrix using Rodrigues' formula
+    K = np.array([[0, -axis[2], axis[1]],
+                  [axis[2], 0, -axis[0]],
+                  [-axis[1], axis[0], 0]])
+    rotation_matrix = np.eye(3) + np.sin(theta) * K + (1 - np.cos(theta)) * np.dot(K, K)
+
+    return rotation_matrix