fixes and extensions

nifty_gridder.cc

@@ -454,7 +454,7 @@ template<typename T> class Baselines
     size_t offset(uint32_t index) const
       { return (index>>channelbits)*nchan + (index&channelmask); }
 
-    pyarr_c<complex<T>> ms2vis(const pyarr_c<complex<T>> &ms_,
+    template<typename T2> pyarr_c<T2> ms2vis(const pyarr_c<T2> &ms_,
       const pyarr_c<uint32_t> &idx_) const
       {
       myassert(idx_.ndim()==1, "idx array must be 1D");
@@ -465,14 +465,14 @@ template<typename T> class Baselines
       auto idx = idx_.data();
       auto ms = ms_.data();
 
-      auto res=makearray<complex<T>>({nvis});
+      auto res=makearray<T2>({nvis});
       auto vis = res.mutable_data();
       for (size_t i=0; i<nvis; ++i)
         vis[i] = ms[offset(idx[i])];
       return res;
       }
 
-    pyarr_c<complex<T>> vis2ms(const pyarr_c<complex<T>> &vis_,
+    template<typename T2> pyarr_c<T2> vis2ms(const pyarr_c<T2> &vis_,
       const pyarr_c<uint32_t> &idx_) const
       {
       myassert(idx_.ndim()==1, "idx array must be 1D");
@@ -482,17 +482,17 @@ template<typename T> class Baselines
       auto idx = idx_.data();
       auto vis = vis_.data();
 
-      auto res = makearray<complex<T>>({nrows, nchan});
+      auto res = makearray<T2>({nrows, nchan});
       auto ms = res.mutable_data();
       for (size_t i=0; i<nrows*nchan; ++i)
-        ms[i] = complex<T>(0);
+        ms[i] = T2(0);
       for (size_t i=0; i<nvis; ++i)
         ms[offset(idx[i])] = vis[i];
       return res;
       }
 
-    pyarr_c<complex<T>> add_vis_to_ms(const pyarr_c<complex<T>> &vis_,
-      const pyarr_c<uint32_t> &idx_, pyarr_c<complex<T>> &ms_) const
+    template<typename T2> pyarr_c<T2> add_vis_to_ms(const pyarr_c<T2> &vis_,
+      const pyarr_c<uint32_t> &idx_, pyarr_c<T2> &ms_) const
       {
       myassert(idx_.ndim()==1, "idx array must be 1D");
       myassert(vis_.ndim()==1, "vis array must be 1D");
@@ -550,6 +550,8 @@ template<typename T> class GridderConfig
     size_t Nu() const { return nu; }
     size_t Nv() const { return nv; }
     size_t W() const { return w; }
+    T Ucorr() const { return ucorr; }
+    T Vcorr() const { return vcorr; }
     size_t coord2peano(const UVW<T> &coord) const
       {
       double u=fmodulo(coord.u*ucorr, T(1))*nu,
@@ -821,6 +823,7 @@ template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
   auto res = makearray<complex<T>>({nu, nv});
   auto grid = res.mutable_data();
   for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
+  T ucorr = gconf.Ucorr(), vcorr=gconf.Vcorr();
 
 #pragma omp parallel
 {
@@ -834,7 +837,7 @@ template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
   for (size_t ipart=0; ipart<nvis; ++ipart)
     {
     UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
-    hlp.prep_write(coord.u, coord.v);
+    hlp.prep_write(coord.u*ucorr, coord.v*vcorr);
     auto * RESTRICT ptr = hlp.p0;
     int w = hlp.w;
     auto v(vis[ipart]*emb);
@@ -866,6 +869,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
 
   auto res = makearray<complex<T>>({nvis});
   auto vis = res.mutable_data();
+  T ucorr = gconf.Ucorr(), vcorr=gconf.Vcorr();
 
   // Loop over sampling points
 #pragma omp parallel
@@ -879,7 +883,7 @@ template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
   for (size_t ipart=0; ipart<nvis; ++ipart)
     {
     UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
-    hlp.prep_read(coord.u, coord.v);
+    hlp.prep_read(coord.u*ucorr, coord.v*vcorr);
     complex<T> r = 0;
     auto * RESTRICT ptr = hlp.p0;
     int w = hlp.w;
@@ -986,9 +990,10 @@ PYBIND11_MODULE(nifty_gridder, m)
   py::class_<Baselines<double>> (m, "Baselines", Baselines_DS)
     .def(py::init<const pyarr_c<double> &, const pyarr_c<double> &,
       const pyarr_c<bool> &>(), "coord"_a, "scaling"_a, "flags"_a)
-    .def ("ms2vis",&Baselines<double>::ms2vis, BL_ms2vis_DS, "ms"_a, "idx"_a)
-    .def ("vis2ms",&Baselines<double>::vis2ms, BL_vis2ms_DS, "vis"_a, "idx"_a)
-    .def ("add_vis_to_ms",&Baselines<double>::add_vis_to_ms, BL_add_vis_to_ms_DS,
+    .def ("ms2vis",&Baselines<double>::ms2vis<complex<double>>, BL_ms2vis_DS, "ms"_a, "idx"_a)
+    .def ("ms2vis_f32",&Baselines<double>::ms2vis<float>, "ms"_a, "idx"_a)
+    .def ("vis2ms",&Baselines<double>::vis2ms<complex<double>>, BL_vis2ms_DS, "vis"_a, "idx"_a)
+    .def ("add_vis_to_ms",&Baselines<double>::add_vis_to_ms<complex<double>>, BL_add_vis_to_ms_DS,
       "vis"_a, "idx"_a, "ms"_a.noconvert());
   py::class_<GridderConfig<double>> (m, "GridderConfig")
     .def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,
@@ -1005,9 +1010,9 @@ PYBIND11_MODULE(nifty_gridder, m)
   py::class_<Baselines<float>> (m, "Baselines_f", Baselines_DS)
     .def(py::init<const pyarr_c<float> &, const pyarr_c<float> &,
       const pyarr_c<bool> &>(), "coord"_a, "scaling"_a, "flags"_a)
-    .def ("ms2vis",&Baselines<float>::ms2vis, BL_ms2vis_DS, "ms"_a, "idx"_a)
-    .def ("vis2ms",&Baselines<float>::vis2ms, BL_vis2ms_DS, "vis"_a, "idx"_a)
-    .def ("add_vis_to_ms",&Baselines<float>::add_vis_to_ms, BL_add_vis_to_ms_DS,
+    .def ("ms2vis",&Baselines<float>::ms2vis<complex<float>>, BL_ms2vis_DS, "ms"_a, "idx"_a)
+    .def ("vis2ms",&Baselines<float>::vis2ms<complex<float>>, BL_vis2ms_DS, "vis"_a, "idx"_a)
+    .def ("add_vis_to_ms",&Baselines<float>::add_vis_to_ms<complex<float>>, BL_add_vis_to_ms_DS,
       "vis"_a, "idx"_a, "ms"_a.noconvert());
   py::class_<GridderConfig<float>> (m, "GridderConfig_f")
     .def(py::init<size_t, size_t, float, float, float>(),"nxdirty"_a,