nifty_gridder.cc

/*
 *  This file is part of nifty_gridder.
 *
 *  nifty_gridder is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  nifty_gridder is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with nifty_gridder; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

/* Copyright (C) 2019 Max-Planck-Society
   Author: Martin Reinecke */

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>

#include "gridder_cxx.h"

using namespace std;

namespace py = pybind11;

namespace {

auto None = py::none();

//
// basic utilities
//

//
// Start of real gridder functionality
//

template<typename T>
  using pyarr = py::array_t<T>;
// The "_c" suffix here stands for "C memory order, contiguous"
template<typename T>
  using pyarr_c = py::array_t<T, py::array::c_style | py::array::forcecast>;

template<typename T> pyarr_c<T> makeArray(const vector<size_t> &shape)
  { return pyarr_c<T>(shape); }

void checkArray(const py::array &arr, const char *aname,
  const vector<size_t> &shape)
  {
  if (size_t(arr.ndim())!=shape.size())
    {
    cerr << "Array '" << aname << "' has " << arr.ndim() << " dimensions; "
            "expected " << shape.size() << endl;
    throw runtime_error("bad dimensionality");
    }
  for (size_t i=0; i<shape.size(); ++i)
    if ((shape[i]!=0) && (size_t(arr.shape(i))!=shape[i]))
      {
      cerr << "Dimension " << i << " of array '" << aname << "' has size "
           << arr.shape(i) << "; expected " << shape[i] << endl;
      throw runtime_error("bad array size");
      }
  }

template<typename T> pyarr<T> provideArray(const py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is_none())
    {
    auto tmp_ = makeArray<T>(shape);
    size_t sz = size_t(tmp_.size());
    auto tmp = tmp_.mutable_data();
    for (size_t i=0; i<sz; ++i)
      tmp[i] = T(0);
    return tmp_;
    }
  auto tmp_ = in.cast<pyarr<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<typename T> pyarr<T> providePotentialArray(const py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is_none())
    return makeArray<T>(vector<size_t>(shape.size(), 0));
  auto tmp_ = in.cast<pyarr<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<typename T> pyarr_c<T> provideCArray(py::object &in,
  const vector<size_t> &shape)
  {
  if (in.is_none())
    {
    auto tmp_ = makeArray<T>(shape);
    size_t sz = size_t(tmp_.size());
    auto tmp = tmp_.mutable_data();
    for (size_t i=0; i<sz; ++i)
      tmp[i] = T(0);
    return tmp_;
    }
  auto tmp_ = in.cast<pyarr_c<T>>();
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

template<size_t ndim, typename T> cmav<T,ndim> make_cmav(pyarr_c<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  for (size_t i=0; i<ndim; ++i) dims[i]=in.shape(i);
  return cmav<T, ndim>(in.mutable_data(),dims);
  }
template<size_t ndim, typename T> smav<T,ndim> make_smav(pyarr<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return smav<T, ndim>(in.mutable_data(),dims,str);
  }
template<size_t ndim, typename T> smav<T,ndim> make_smav(pyarr_c<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return smav<T, ndim>(in.mutable_data(),dims,str);
  }
template<size_t ndim, typename T> cmav<const T,ndim> make_const_cmav(const pyarr_c<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  for (size_t i=0; i<ndim; ++i) dims[i]=in.shape(i);
  return cmav<const T, ndim>(in.data(),dims);
  }
template<size_t ndim, typename T> smav<const T,ndim> make_const_smav(const pyarr<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return smav<const T, ndim>(in.data(),dims,str);
  }
template<size_t ndim, typename T> smav<const T,ndim> make_const_smav(const pyarr_c<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return smav<const T, ndim>(in.data(),dims,str);
  }
template<typename T> pyarr_c<T> complex2hartley
  (const pyarr_c<complex<T>> &grid_, py::object &grid_in)
  {
  auto grid = make_const_smav<2>(grid_);
  size_t nu=grid.shape(0), nv=grid.shape(1);

  auto res = provideCArray<T>(grid_in, {nu, nv});
  auto grid2 = make_smav<2>(res);
  {
  py::gil_scoped_release release;
  complex2hartley(grid, grid2);
  }
  return res;
  }

template<typename T> pyarr_c<complex<T>> hartley2complex
  (const pyarr_c<T> &grid_)
  {
  auto grid = make_const_smav<2>(grid_);
  size_t nu=grid.shape(0), nv=grid.shape(1);

  auto res = makeArray<complex<T>>({nu, nv});
  auto grid2 = make_smav<2>(res);
  {
  py::gil_scoped_release release;
  hartley2complex(grid, grid2);
  }
  return res;
  }

template<typename T> void hartley2_2D(const pyarr_c<T> &in, pyarr_c<T> &out)
  {
  auto grid = make_const_smav<2>(in);
  auto grid2 = make_smav<2>(out);
  py::gil_scoped_release release;
  hartley2_2D(grid, grid2);
  }

constexpr auto PyBaselines_DS = R"""(
Class storing UVW coordinates and channel information.

Parameters
==========
coord: np.array((nrows, 3), dtype=np.float)
    u, v and w coordinates for each row
freq: np.array((nchannels,), dtype=np.float)
    frequency for each individual channel (in Hz)
)""";
template<typename T> class PyBaselines: public Baselines<T>
  {
  protected:
    using Baselines<T>::coord;
    using Baselines<T>::f_over_c;
    using Baselines<T>::nrows;
    using Baselines<T>::nchan;

  public:
    using Baselines<T>::Baselines;
    PyBaselines(const pyarr<T> &coord, const pyarr<T> &freq)
      : Baselines<T>(make_const_smav<2>(coord), make_const_smav<1>(freq))
      {}

    using Baselines<T>::effectiveCoord;
    using Baselines<T>::Nrows;
    using Baselines<T>::Nchannels;

    static constexpr auto ms2vis_DS = R"""(
    Extracts visibility data from a measurement for the provided indices.

    Parameters
    ==========
    ms: np.array((nrows, nchannels), dtype=np.complex)
        the measurement set's visibility data
    idx: np.array((nvis,), dtype=np.uint32)
        the indices to be extracted

    Returns
    =======
    np.array((nvis,), dtype=np.complex)
        The visibility data for the index array
    )""";

   // using Baselines<T>::effectiveUVW;
    pyarr_c<T> effectiveuvw(const pyarr<uint32_t> &idx_) const
      {
      size_t nvis = size_t(idx_.shape(0));
      auto idx=make_const_smav<1>(idx_);
      auto res_=makeArray<T>({nvis, 3});
      auto res=make_smav<2>(res_);
      {
      py::gil_scoped_release release;
      Baselines<T>::effectiveUVW(idx,res);
      }
      return res_;
      }

    template<typename T2> pyarr_c<T2> ms2vis(const pyarr<T2> &ms_,
      const pyarr_c<uint32_t> &idx_) const
      {
      auto idx=make_const_smav<1>(idx_);
      size_t nvis = size_t(idx.shape(0));
      auto ms = make_const_smav<2>(ms_);
      auto res=makeArray<T2>({nvis});
      auto vis = make_smav<1>(res);
      {
      py::gil_scoped_release release;
      Baselines<T>::ms2vis(ms, idx, vis);
      }
      return res;
      }

    static constexpr auto vis2ms_DS = R"""(
    Produces a new MS with the provided visibilities set.

    Parameters
    ==========
    vis: np.array((nvis,), dtype=np.complex)
        The visibility data for the index array
    idx: np.array((nvis,), dtype=np.uint32)
        the indices to be inserted
    ms_in: np.array((nrows, nchannels), dtype=np.complex), optional
        input measurement set to which the visibilities are added.

    Returns
    =======
    np.array((nrows, nchannels), dtype=np.complex)
        the measurement set's visibility data (0 where not covered by idx)
    )""";
    template<typename T2> pyarr_c<T2> vis2ms(const pyarr<T2> &vis_,
      const pyarr<uint32_t> &idx_, py::object &ms_in) const
      {
      auto vis=make_const_smav<1>(vis_);
      auto idx=make_const_smav<1>(idx_);
      auto res = provideArray<T2>(ms_in, {nrows, nchan});
      auto ms = make_smav<2>(res);
      {
      py::gil_scoped_release release;
      Baselines<T>::vis2ms(vis, idx, ms);
      }
      return res;
      }
  };

constexpr auto grid2dirty_DS = R"""(
Converts from UV grid to dirty image (FFT, cropping, correction)

Parameters
==========
grid: np.array((nu, nv), dtype=np.float64)
    gridded UV data

Returns
=======
nd.array((nxdirty, nydirty), dtype=np.float64)
    the dirty image
)""";

constexpr auto dirty2grid_DS = R"""(
Converts from a dirty image to a UV grid (correction, padding, FFT)

Parameters
==========
dirty: nd.array((nxdirty, nydirty), dtype=np.float64)
    the dirty image

Returns
=======
np.array((nu, nv), dtype=np.float64)
    gridded UV data
)""";

constexpr auto apply_taper_DS = R"""(
Applies the taper (or its inverse) to an image

Parameters
==========
img: nd.array((nxdirty, nydirty), dtype=np.float64)
    the image
divide: bool
    if True, the routine dividex by the taper, otherwise it multiplies by it

Returns
=======
np.array((nxdirty, nydirty), dtype=np.float64)
    the image with the taper applied
)""";

constexpr auto apply_wscreen_DS = R"""(
Applies the w screen to an image

Parameters
==========
dirty: nd.array((nxdirty, nydirty), dtype=np.complex128)
    the image
w : float
    the w value to use
adjoint: bool
    if True, apply the complex conjugate of the w screen

Returns
=======
np.array((nxdirty, nydirty), dtype=np.complex128)
    the image with the w screen applied
)""";

constexpr auto GridderConfig_DS = R"""(
Class storing information related to the gridding/degridding process.

Parameters
==========
nxdirty: int
    x resolution of the dirty image; must be even
nydirty: int
    y resolution of the dirty image; must be even
epsilon: float
    required accuracy for the gridding/degridding step
    Must be >= 2e-13.
pixsize_x: float
    Pixel size in x direction (radians)
pixsize_y: float
    Pixel size in y direction (radians)
)""";
template<typename T> class PyGridderConfig: public GridderConfig<T>
  {
  protected:
    using GridderConfig<T>::nx_dirty;
    using GridderConfig<T>::ny_dirty;
    using GridderConfig<T>::eps;
    using GridderConfig<T>::psx;
    using GridderConfig<T>::psy;
    using GridderConfig<T>::supp;
    using GridderConfig<T>::nsafe;
    using GridderConfig<T>::nu;
    using GridderConfig<T>::nv;
    using GridderConfig<T>::beta;
    using GridderConfig<T>::cfu;
    using GridderConfig<T>::cfv;
    using GridderConfig<T>::wscreen;

  public:
    using GridderConfig<T>::GridderConfig;
    PyGridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
      double pixsize_x, double pixsize_y)
      : GridderConfig<T>(nxdirty, nydirty, epsilon, pixsize_x, pixsize_y) {}
    using GridderConfig<T>::Nxdirty;
    using GridderConfig<T>::Nydirty;
    using GridderConfig<T>::Epsilon;
    using GridderConfig<T>::Pixsize_x;
    using GridderConfig<T>::Pixsize_y;
    using GridderConfig<T>::Nu;
    using GridderConfig<T>::Nv;
    using GridderConfig<T>::Supp;
    using GridderConfig<T>::Nsafe;
    using GridderConfig<T>::Beta;

    pyarr_c<T> grid2dirty(const pyarr_c<T> &grid) const
      {
      checkArray(grid, "grid", {nu, nv});
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto grid2=make_const_smav<2>(grid);
      auto res2=make_smav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::grid2dirty(grid2,res2);
      }
      return res;
      }
    pyarr_c<T> apply_taper(const pyarr_c<T> &img, bool divide) const
      {
      checkArray(img, "img", {nx_dirty, ny_dirty});
      auto pin = img.data();
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto pout = res.mutable_data();
      {
      py::gil_scoped_release release;
      if (divide)
        for (size_t i=0; i<nx_dirty; ++i)
          for (size_t j=0; j<ny_dirty; ++j)
            pout[ny_dirty*i + j] = pin[ny_dirty*i + j]/(cfu[i]*cfv[j]);
      else
        for (size_t i=0; i<nx_dirty; ++i)
          for (size_t j=0; j<ny_dirty; ++j)
            pout[ny_dirty*i + j] = pin[ny_dirty*i + j]*cfu[i]*cfv[j];
      }
      return res;
      }
    pyarr_c<complex<T>> grid2dirty_c(const pyarr_c<complex<T>> &grid) const
      {
      checkArray(grid, "grid", {nu, nv});
      auto tmp = makeArray<complex<T>>({nu, nv});
      auto ptmp = tmp.mutable_data();
      pocketfft::c2c({nu,nv},{grid.strides(0),grid.strides(1)},
        {tmp.strides(0), tmp.strides(1)}, {0,1}, pocketfft::BACKWARD,
        grid.data(), tmp.mutable_data(), T(1), nthreads);
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto pout = res.mutable_data();
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          pout[ny_dirty*i + j] = ptmp[nv*i2+j2]*cfu[i]*cfv[j];
          }
      }
      return res;
      }

    pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
      {
      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
      auto pdirty = dirty.data();
      auto tmp = makeArray<T>({nu, nv});
      auto ptmp = tmp.mutable_data();
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nu*nv; ++i)
        ptmp[i] = 0.;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          ptmp[nv*i2+j2] = pdirty[ny_dirty*i + j]*cfu[i]*cfv[j];
          }
      }
      hartley2_2D<T>(tmp, tmp);
      return tmp;
      }
    pyarr_c<complex<T>> dirty2grid_c(const pyarr_c<complex<T>> &dirty) const
      {
      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
      auto pdirty = dirty.data();
      auto tmp = makeArray<complex<T>>({nu, nv});
      auto ptmp = tmp.mutable_data();
      pocketfft::stride_t strides{tmp.strides(0),tmp.strides(1)};
      {
      py::gil_scoped_release release;
      for (size_t i=0; i<nu*nv; ++i)
        ptmp[i] = 0.;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          ptmp[nv*i2+j2] = pdirty[ny_dirty*i + j]*cfu[i]*cfv[j];
          }
      pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
        ptmp, ptmp, T(1), nthreads);
      }
      return tmp;
      }
    pyarr_c<complex<T>> apply_wscreen(const pyarr_c<complex<T>> &dirty_, double w, bool adjoint) const
      {
      checkArray(dirty_, "dirty", {nx_dirty, ny_dirty});
      auto dirty = dirty_.data();
      auto res_ = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto res = res_.mutable_data();
      double x0 = -0.5*nx_dirty*psx,
             y0 = -0.5*ny_dirty*psy;
      {
      py::gil_scoped_release release;
#pragma omp parallel num_threads(nthreads)
{
#pragma omp for schedule(static)
      for (size_t i=0; i<=nx_dirty/2; ++i)
        {
        double fx = x0+i*psx;
        fx *= fx;
        for (size_t j=0; j<=ny_dirty/2; ++j)
          {
          double fy = y0+j*psy;
          auto ws = wscreen(fx, fy*fy, w, adjoint);
          res[ny_dirty*i+j] = dirty[ny_dirty*i+j]*ws; // lower left
          size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
          if ((i>0)&&(i<i2))
            {
            res[ny_dirty*i2+j] = dirty[ny_dirty*i2+j]*ws; // lower right
            if ((j>0)&&(j<j2))
              res[ny_dirty*i2+j2] = dirty[ny_dirty*i2+j2]*ws; // upper right
            }
          if ((j>0)&&(j<j2))
            res[ny_dirty*i+j2] = dirty[ny_dirty*i+j2]*ws; // upper left
          }
        }
}
      }
      return res_;
      }
  };


constexpr auto vis2grid_c_DS = R"""(
Grids visibilities onto a UV grid

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used
    (used to optimize the ordering of the indices)
idx: np.array((nvis,), dtype=np.uint32)
    the indices for the entries to be gridded
vis: np.array((nvis,), dtype=np.complex)
    The visibility data for the index array
grid_in: np.array((nu,nv), dtype=np.complex128), optional
    If present, the result is added to this array.
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
template<typename T> pyarr_c<complex<T>> vis2grid_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_,
  py::object &grid_in, const py::object &wgt_)
  {
  checkArray(vis_, "vis", {0});
  size_t nvis = size_t(vis_.shape(0));
  checkArray(idx_, "idx", {nvis});
  auto vis=vis_.template unchecked<1>();
  auto idx = idx_.template unchecked<1>();
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<1>();

  size_t nu=gconf.Nu(), nv=gconf.Nv();
  auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
  auto grid = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();

#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, nullptr, grid);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx(ipart));
    hlp.prep(coord.u, coord.v);
    auto * ptr = hlp.p0w;
    auto v(vis(ipart)*emb);
    if (have_wgt)
      v*=wgt(ipart);
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(v*ku[cu]);
      for (size_t cv=0; cv<supp; ++cv)
        ptr[cv] += tmp*kv[cv];
      ptr+=jump;
      }
    }
} // end of parallel region
  }
  return res;
  }

constexpr auto vis2grid_DS = R"""(
Grids visibilities onto a UV grid

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used
    (used to optimize the ordering of the indices)
idx: np.array((nvis,), dtype=np.uint32)
    the indices for the entries to be gridded
vis: np.array((nvis,), dtype=np.complex)
    The visibility data for the index array
grid_in: np.array((nu,nv), dtype=np.float64), optional
    If present, the result is added to this array.
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nu,nv), dtype=np.float64):
    the gridded visibilities (made real by making use of Hermitian symmetry)
)""";
template<typename T> pyarr_c<T> vis2grid(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
  { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_, None, wgt_), grid_in); }

constexpr auto ms2grid_c_DS = R"""(
Grids measurement set data onto a UV grid

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used
    (used to optimize the ordering of the indices)
idx: np.array((nvis,), dtype=np.uint32)
    the indices for the entries to be gridded
ms: np.array((nrows, nchannels), dtype=np.complex128)
    the measurement set.
grid_in: np.array((nu,nv), dtype=np.complex128), optional
    If present, the result is added to this array.
wgt: np.array((nrows, nchannels), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
template<typename T> pyarr_c<complex<T>> ms2grid_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  checkArray(ms_, "ms", {nrows, nchan});
  checkArray(idx_, "idx", {0});
  size_t nvis = size_t(idx_.shape(0));
  auto ms = ms_.template unchecked<2>();
  auto idx = idx_.template unchecked<1>();
  auto wgt2 = providePotentialArray<T>(wgt_, {nrows, nchan});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<2>();

  size_t nu=gconf.Nu(), nv=gconf.Nv();
  auto res = provideCArray<complex<T>>(grid_in, {nu, nv});
  auto grid = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();

#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, nullptr, grid);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    auto tidx = idx(ipart);
    auto row = tidx/nchan;
    auto chan = tidx-row*nchan;
    UVW<T> coord = baselines.effectiveCoord(tidx);
    hlp.prep(coord.u, coord.v);
    auto * ptr = hlp.p0w;
    auto v(ms(row,chan)*emb);
    if (have_wgt)
      v*=wgt(row, chan);
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(v*ku[cu]);
      for (size_t cv=0; cv<supp; ++cv)
        ptr[cv] += tmp*kv[cv];
      ptr+=jump;
      }
    }
} // end of parallel region
  }
  return res;
  }

template<typename T> pyarr_c<T> ms2grid(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &ms_, py::object &grid_in, const py::object &wgt_)
  { return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_, None, wgt_), grid_in); }

template<typename T> pyarr_c<complex<T>> grid2vis_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_,
  const py::object &wgt_)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.template unchecked<1>();
  auto wgt2 = providePotentialArray<T>(wgt_, {nvis});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<1>();

  auto res = makeArray<complex<T>>({nvis});
  auto vis = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, grid, nullptr);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx(ipart));
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    const auto * ptr = hlp.p0r;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<supp; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += jump;
      }
    if (have_wgt) r*=wgt[ipart];
    vis[ipart] = r*emb;
    }
}
  }
  return res;
  }

constexpr auto grid2vis_DS = R"""(
Degrids visibilities from a UV grid

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used
    (used to optimize the ordering of the indices)
idx: np.array((nvis,), dtype=np.uint32)
    the indices for the entries to be degridded
grid: np.array((nu,nv), dtype=np.float64):
    the gridded visibilities (made real by making use of Hermitian symmetry)
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.

Returns
=======
np.array((nvis,), dtype=np.complex)
    The degridded visibility data
)""";
template<typename T> pyarr_c<complex<T>> grid2vis(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr_c<T> &grid_, const py::object &wgt_)
  { return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_), wgt_); }

template<typename T> pyarr_c<complex<T>> grid2ms_c(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr_c<complex<T>> &grid_, py::object &ms_in, const py::object &wgt_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.template unchecked<1>();
  auto wgt2 = providePotentialArray<T>(wgt_, {nrows, nchan});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<2>();

  auto res = provideArray<complex<T>>(ms_in, {nrows, nchan});
  auto ms = res.template mutable_unchecked<2>();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, grid, nullptr);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    auto tidx = idx(ipart);
    auto row = tidx/nchan;
    auto chan = tidx-row*nchan;
    UVW<T> coord = baselines.effectiveCoord(tidx);
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    const auto * ptr = hlp.p0r;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<supp; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += jump;
      }
    r*=emb;
    if (have_wgt)
      r*=wgt(row, chan);
    ms(row,chan) += r*emb;
    }
}
  }
  return res;
  }

template<typename T> pyarr_c<complex<T>> grid2ms(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr_c<T> &grid_, py::object &ms_in, const py::object &wgt_)
  { return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in, wgt_); }

template<typename T> pyarr_c<complex<T>> apply_holo(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_,
  const py::object &wgt_)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.template unchecked<1>();
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<1>();

  auto res = makeArray<complex<T>>({nu, nv});
  auto ogrid = res.mutable_data();
  {
  py::gil_scoped_release release;

  for (size_t i=0; i<nu*nv; ++i) ogrid[i] = T(0);
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, grid, ogrid);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    auto tidx = idx(ipart);
    UVW<T> coord = baselines.effectiveCoord(tidx);
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    const auto * ptr = hlp.p0r;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<supp; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += jump;
      }
    r*=emb*emb;
    if (have_wgt)
      {
      auto twgt = wgt(ipart);
      r*=twgt*twgt;
      }
    auto * wptr = hlp.p0w;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(r*ku[cu]);
      for (size_t cv=0; cv<supp; ++cv)
        wptr[cv] += tmp*kv[cv];
      wptr += jump;
      }
    }
}
  }
  return res;
  }

template<typename T> pyarr_c<T> get_correlations(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, int du, int dv, const py::object &wgt_)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  size_t supp = gconf.Supp();
  myassert(size_t(abs(du))<supp, "|du| must be smaller than Supp");
  myassert(size_t(abs(dv))<supp, "|dv| must be smaller than Supp");
  checkArray(idx_, "idx", {0});
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.template unchecked<1>();
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {nvis});
  bool have_wgt = wgt2.size()>0;
  auto wgt = wgt2.template unchecked<1>();

  auto res = makeArray<T>({nu, nv});
  auto ogrid = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  for (size_t i=0; i<nu*nv; ++i) ogrid[i] = 0.;

  size_t u0, u1, v0, v1;
  if (du>=0)
    { u0=0; u1=supp-du; }
  else
    { u0=-du; u1=supp; }
  if (dv>=0)
    { v0=0; v1=supp-dv; }
  else
    { v0=-dv; v1=supp; }

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T,T> hlp(gconf, nullptr, ogrid);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+supp;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx(ipart));
    hlp.prep(coord.u, coord.v);