Modified test_response to also check for jax and nifty consistency of radio response

test/test_response.py

@@ -22,9 +22,12 @@ import numpy as np
 import pytest
 
 import resolve as rve
+import resolve.re as jrve
 
 from .common import setup_function, teardown_function
 
+import configparser
+
 pmp = pytest.mark.parametrize
 np.seterr(all="raise")
 
@@ -66,3 +69,39 @@ def test_facet_consistency():
         if res0 is None:
             res0 = res
         ift.extra.assert_allclose(res0, res, rtol=1e-4, atol=1e-4)
+
+def test_jax_response_consistency():
+    '''
+    Ugly code MUST CLEAN UP
+    '''
+    obs = rve.Observation.legacy_load('../../Vishal_RESOLVE/polarization_imaging/A3667_single_freq.npz')
+    obs = obs.to_double_precision()
+    
+    cfg = configparser.ConfigParser()
+    cfg.read("../../Vishal_RESOLVE/polarization_imaging/a3667.cfg")
+
+    diffuse, additional = rve.sky_model_diffuse(cfg["sky"], observations=obs, nthreads=8)
+
+    ex_field = ift.from_random(diffuse.domain)
+    radio_sky = diffuse(ex_field)
+    radio_sky_arr = radio_sky.val
+
+    sky_shape = (float(cfg['sky']['space npix x']), float(cfg['sky']['space npix y']))
+
+    R_old = rve.InterferometryResponse(obs, diffuse.target, do_wgridding=True, epsilon=1e-9, nthreads=8, verbosity=2)
+    sky_domain_dict = dict(npix_x=sky_shape[0],
+                       npix_y=sky_shape[1],
+                       pixsize_x=float(diffuse.target[3].distances[0]),
+                       pixsize_y=float(diffuse.target[3].distances[1]),
+                       pol_labels=['I', 'Q', 'U', 'V'],
+                       times=[0.],
+                       freqs=[0.])
+
+    R_new = jrve.InterferometryResponse(obs, sky_domain_dict, True, 1e-9, nthreads=8, verbosity=2)
+
+    vis_field_old = R_old(diffuse(ex_field)).val
+    vis_field_new = R_new(radio_sky_arr)
+
+    np.allclose(vis_field_old, vis_field_new)
+    
+