diff --git a/.coveragerc b/.coveragerc index 2e903ab..2dc3a04 100644 --- a/.coveragerc +++ b/.coveragerc @@ -1,5 +1,6 @@ [run] include = fftvis/* +omit = fftvis/logutils.py branch = True [report] diff --git a/fftvis/simulate.py b/fftvis/simulate.py index 738ee5a..238fdbd 100644 --- a/fftvis/simulate.py +++ b/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] diff --git a/fftvis/tests/test_compare_matvis.py b/fftvis/tests/test_compare_matvis.py index d83be35..907cbe1 100644 --- a/fftvis/tests/test_compare_matvis.py +++ b/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) diff --git a/fftvis/tests/test_utils.py b/fftvis/tests/test_utils.py index 0f0ee84..f2d6c0e 100644 --- a/fftvis/tests/test_utils.py +++ b/fftvis/tests/test_utils.py @@ -1,6 +1,48 @@ import pytest +import numpy as np +from fftvis import utils +def test_get_plane_to_xy_rotation_matrix(): + """ + """ + # Rotate the array to the xy-plane + x = np.linspace(0, 100, 100) + y = np.linspace(0, 100, 100) + z = 0.125 * x + 0.5 * y + antvecs = np.array([x, y, z]).T -def test_get_pos_reds(): - """ """ - pass + rotation_matrix = utils.get_plane_to_xy_rotation_matrix(antvecs) + rm_antvecs = np.dot(rotation_matrix.T, antvecs.T) + + # Check that all elements of the z-axis are close to zero + np.testing.assert_allclose(rm_antvecs[-1], 0, atol=1e-12) + + # Check that the lengths of the vectors are preserved + np.testing.assert_allclose( + np.linalg.norm(antvecs, axis=1), + np.linalg.norm(rm_antvecs, axis=0), + atol=1e-12 + ) + +def test_get_plane_to_xy_rotation_matrix_errors(): + """ + """ + # Rotate the array to the xy-plane + x = np.linspace(0, 100, 100) + y = np.linspace(0, 100, 100) + z = 0.125 * x + 0.5 * y + antvecs = np.array([x, y, z]).T + + # Check that method is robust to errors + rng = np.random.default_rng(42) + random_antvecs = antvecs.copy() + random_antvecs[:, -1] += rng.standard_normal(100) + + # Rotate the array to the xy-plane + rotation_matrix = utils.get_plane_to_xy_rotation_matrix( + random_antvecs + ) + rm_antvecs = np.dot(rotation_matrix.T, antvecs.T) + + # Check that all elements of the z-axis within 5-sigma of zero + np.testing.assert_array_less(np.abs(rm_antvecs[-1]), 5) \ No newline at end of file diff --git a/fftvis/utils.py b/fftvis/utils.py index 2bcd650..4c4b9d4 100644 --- a/fftvis/utils.py +++ b/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 \ No newline at end of file