account for frequency dimension in nearest gridder

Author: mkolopanis

src/py21cmwedge/uvgridder.py

@@ -265,10 +265,10 @@ def uv_weights(self, u, v, spatial_function="triangle"):
         #     weights = 1. - (np.abs(uv - grid)/np.diff(grid)[0])**2
         #     weights = np.exp( - (uv - grid)**2/(2*np.diff(grid)[0]**2))
         #     weights = np.exp( - abs(uv - grid)/(np.diff(grid)[0]))
-        _range = np.arange(self.uv_size) - (self.uv_size - 1) / 2.0
-        _range *= self.uv_delta
         match spatial_function.casefold():
             case "triangle":
+                _range = np.arange(self.uv_size) - (self.uv_size - 1) / 2.0
+                _range *= self.uv_delta
                 x, y = np.meshgrid(_range, _range)
                 x.shape += (1,)
                 y.shape += (1,)
@@ -280,11 +280,22 @@ def uv_weights(self, u, v, spatial_function="triangle"):
                 weights /= np.sum(weights, axis=(0, 1))
                 weights = np.transpose(weights, [2, 0, 1])
             case "nearest":
-                u_index = np.argmin(np.abs(u - _range))
-                v_index = np.argmin(np.abs(v - _range))
+                _range = np.arange(self.uv_size) - (self.uv_size - 1) / 2.0
+                _range *= self.uv_delta
+                x, y = _range, _range
+                x.shape += (1,)
+                y.shape += (1,)
+                x = u - x
+                y = v - y
+                u_index = np.argmin(np.abs(x), axis=0)
+                v_index = np.argmin(np.abs(y), axis=0)
+
+                print(f"{u_index.shape=:}, {v_index.shape=:}")
                 print(f"{u_index=:}, {v_index=:}")
-                weights = np.zeros((1, self.uv_size, self.uv_size), dtype=complex)
-                weights[0, u_index, v_index] = 1.0
+                weights = np.zeros(
+                    (x.shape[-1], self.uv_size, self.uv_size), dtype=complex
+                )
+                weights[:, u_index, v_index] = 1.0
             case _:
                 raise ValueError(
                     f"Unknown value for 'spatial_function': {spatial_function}"

tests/test_uvgridder.py

@@ -326,6 +326,20 @@ def test_grid_uv_deltas():
     np.testing.assert_allclose(test_obj.uvf_cube[0, 33, 18], 10)
 
 
+def test_grid_uv_deltas_multi_freq():
+    """Test grid_uv sets up complex type."""
+    test_obj = UVGridder()
+    test_obj.uv_delta = 0.5
+    test_obj.set_freqs([50e6, 150e6])
+    test_uvw = np.zeros((3, 10)) + np.array([[14.6], [0], [0]])
+    test_obj.set_uvw_array(test_uvw)
+    test_obj.uvw_to_dict()
+    test_obj.grid_uvw(convolve_beam=False, spatial_function="nearest")
+
+    np.testing.assert_allclose(test_obj.uvf_cube[1, 33, 18], 10)
+    np.testing.assert_allclose(test_obj.uvf_cube[0, 23, 18], 10)
+
+
 def test_grid_uv_error():
     """Test grid_uv sets up complex type."""
     test_obj = UVGridder()