* Fix test

dsa2000_cal/dsa2000_cal/adapter/utils.py

@@ -213,6 +213,7 @@ def detect_mixed_corrs(corrs: List[str]):
     Returns:
         True if mixed corrs are present, False otherwise.
     """
+    corrs = [CASA_CORR_TYPES[corr] if isinstance(corr, int) else corr for corr in corrs]
     linear_present = is_linear_present(corrs)
     circular_present = is_circular_present(corrs)
     stokes_present = is_stokes_present(corrs)

dsa2000_cal/dsa2000_cal/visibility_model/source_models/celestial/tests/demo_tests.py

@@ -12,7 +12,7 @@
 from dsa2000_cal.common.ellipse_utils import Gaussian
 from dsa2000_cal.common.quantity_utils import quantity_to_jnp
 from dsa2000_cal.common.types import complex_type
-from dsa2000_cal.common.wgridder import dirty2vis, image_to_vis
+from dsa2000_cal.common.wgridder import image_to_vis
 from dsa2000_cal.delay_models.far_field import VisibilityCoords
 from dsa2000_cal.visibility_model.source_models.celestial.fits_source_model import FITSSourceModel
 from dsa2000_cal.visibility_model.source_models.celestial.gaussian_source_model import GaussianModelData, \
@@ -44,7 +44,6 @@ def test_plot_ateam_sources(source):
     fits_sources.plot()
 
 
-
 def test_gaussian_correctness_order_1():
     major_fwhm_arcsec = 4. * 60
     minor_fwhm_arcsec = 2. * 60

dsa2000_cal/dsa2000_cal/visibility_model/source_models/celestial/tests/test_gaussian_source_model.py

@@ -8,7 +8,7 @@
 
 from dsa2000_cal.common.corr_translation import linear_to_stokes
 from dsa2000_cal.common.types import mp_policy
-from dsa2000_cal.common.wgridder import dirty2vis, vis2dirty, image_to_vis, vis_to_image
+from dsa2000_cal.common.wgridder import image_to_vis, vis_to_image
 from dsa2000_cal.delay_models.far_field import VisibilityCoords
 from dsa2000_cal.visibility_model.source_models.celestial.gaussian_source_model import \
     GaussianPredict, GaussianModelData, GaussianSourceModel

dsa2000_cal/dsa2000_cal/visibility_model/source_models/celestial/tests/test_point_source_model.py

@@ -7,7 +7,7 @@
 
 from dsa2000_cal.common.corr_translation import linear_to_stokes
 from dsa2000_cal.common.types import mp_policy
-from dsa2000_cal.common.wgridder import dirty2vis, vis2dirty, image_to_vis, vis_to_image
+from dsa2000_cal.common.wgridder import image_to_vis, vis_to_image
 from dsa2000_cal.delay_models.far_field import VisibilityCoords
 from dsa2000_cal.visibility_model.source_models.celestial.point_source_model import PointPredict, PointModelData
 
@@ -24,7 +24,7 @@ def build_mock_point_model_data(di_gains: bool, chan: int, source: int, time: in
     model_data = PointModelData(
         freqs=mp_policy.cast_to_freq(freqs),
         image=mp_policy.cast_to_image(jnp.ones((source, chan, 2, 2))),
-        gains=mp_policy.cast_to_gain(jnp.ones(gain_shape)),
+        gains=mp_policy.cast_to_gain(jnp.ones(gain_shape, jnp.complex64)),
         lmn=mp_policy.cast_to_angle(lmn)
     )
     print(model_data)
@@ -214,7 +214,7 @@ def pixel_to_lmn(xi, yi):
     image = np.zeros((2, num_freqs, 2, 2))  # [source, chan, 2, 2]
     image[:, :, 0, 0] = 0.5
     image[:, :, 1, 1] = 0.5
-    gains = np.zeros((2, num_ants, num_freqs, 2, 2))  # [[source,] time, ant, chan, 2, 2]
+    gains = np.ones((2, num_ants, num_freqs, 2, 2), jnp.complex64)  # [[source,] time, ant, chan, 2, 2]
     gains[..., 0, 0] = 1.
     gains[..., 1, 1] = 1.
     lmn = np.stack([lmn1, lmn2], axis=0)  # [source, 3]
@@ -230,7 +230,7 @@ def pixel_to_lmn(xi, yi):
             time_obs=mp_policy.cast_to_time(np.zeros(num_rows)),
             antenna_1=mp_policy.cast_to_index(antenna_1),
             antenna_2=mp_policy.cast_to_index(antenna_2),
-            time_idx=mp_policy.cast_to_index(np.zeros(num_rows))
+            time_idx=mp_policy.cast_to_index(np.zeros(num_rows), quiet=True)
         ))  # [row, chan, 2, 2]
     vis_point_predict_stokes = jax.vmap(jax.vmap(linear_to_stokes))(vis_point_predict_linear)[:, :, 0, 0]
     print(vis_point_predict_stokes)

dsa2000_cal/dsa2000_cal/visibility_model/tests/demo_tests.py

@@ -9,11 +9,10 @@
 
 from dsa2000_cal.assets.content_registry import fill_registries
 from dsa2000_cal.assets.registries import array_registry, source_model_registry, rfi_model_registry
-
 from dsa2000_cal.common.ellipse_utils import Gaussian
 from dsa2000_cal.common.quantity_utils import quantity_to_jnp
 from dsa2000_cal.common.types import complex_type
-from dsa2000_cal.common.wgridder import vis2dirty, vis_to_image
+from dsa2000_cal.common.wgridder import vis_to_image
 from dsa2000_cal.delay_models.far_field import VisibilityCoords, FarFieldDelayEngine
 from dsa2000_cal.delay_models.near_field import NearFieldDelayEngine
 from dsa2000_cal.geodesics.geodesic_model import GeodesicModel