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_fridder; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <iostream>
#include <algorithm>
#include "pocketfft_hdronly.h"

#ifdef __GNUC__
#define RESTRICT __restrict__
#define NOINLINE __attribute__ ((noinline))
#else
#define RESTRICT
#endif

using namespace std;

namespace py = pybind11;

namespace {

//
// basic utilities
//

void myassert(bool cond, const char *msg)
  {
  if (cond) return;
  throw runtime_error(msg);
  }

/*! Returns the largest integer \a n that fulfills \a 2^n<=arg. */
template<typename I> inline int ilog2 (I arg)
  {
#ifdef __GNUC__
  if (arg==0) return 0;
  if (sizeof(I)==sizeof(int))
    return 8*sizeof(int)-1-__builtin_clz(arg);
  if (sizeof(I)==sizeof(long))
    return 8*sizeof(long)-1-__builtin_clzl(arg);
  if (sizeof(I)==sizeof(long long))
    return 8*sizeof(long long)-1-__builtin_clzll(arg);
#endif
  int res=0;
  while (arg > 0xFFFF) { res+=16; arg>>=16; }
  if (arg > 0x00FF) { res|=8; arg>>=8; }
  if (arg > 0x000F) { res|=4; arg>>=4; }
  if (arg > 0x0003) { res|=2; arg>>=2; }
  if (arg > 0x0001) { res|=1; }
  return res;
  }

/*! Returns the number of bits needed to represent \a arg different values.
    \a arg must be >=1. */
template<typename I> inline int bits_needed (I arg)
  {
  myassert(arg>=1, "argument must be >=1");
  if (arg==1) return 0;
  return ilog2(arg-1)+1;
  }

/*! Returns the remainder of the division \a v1/v2.
    The result is non-negative.
    \a v1 can be positive or negative; \a v2 must be positive. */
template<typename T> inline T fmodulo (T v1, T v2)
  {
  if (v1>=0)
    return (v1<v2) ? v1 : fmod(v1,v2);
  T tmp=fmod(v1,v2)+v2;
  return (tmp==v2) ? T(0) : tmp;
  }

//
// Utilities for converting between Cartesian coordinates and Peano index
//

static const uint16_t utab[] = {
#define Z(a) 0x##a##0, 0x##a##1, 0x##a##4, 0x##a##5
#define Y(a) Z(a##0), Z(a##1), Z(a##4), Z(a##5)
#define X(a) Y(a##0), Y(a##1), Y(a##4), Y(a##5)
X(0),X(1),X(4),X(5)
#undef X
#undef Y
#undef Z
};

uint32_t coord2morton2D_32 (uint32_t x, uint32_t y)
  {
  typedef uint32_t I;
  return  (I)(utab[x&0xff])     | ((I)(utab[(x>>8)&0xff])<<16)
       | ((I)(utab[y&0xff])<<1) | ((I)(utab[(y>>8)&0xff])<<17);
  }

static const uint8_t m2p2D_1[4][4] = {
{ 4, 1, 11, 2},{0,15, 5, 6},{10,9,3,12},{14,7,13,8}};
static uint8_t m2p2D_3[4][64];
static const uint8_t p2m2D_1[4][4] = {
{ 4, 1, 3, 10},{0,6,7,13},{15,9,8,2},{11,14,12,5}};
static uint8_t p2m2D_3[4][64];
static int peano2d_done=0;

static void init_peano2d (void)
  {
  peano2d_done=1;

  for (int d=0; d<4; ++d)
    for (uint32_t p=0; p<64; ++p)
      {
      unsigned rot = d;
      uint32_t v = p<<26;
      uint32_t res = 0;
      for (int i=0; i<3; ++i)
        {
        unsigned tab=m2p2D_1[rot][v>>30];
        v<<=2;
        res = (res<<2) | (tab&0x3);
        rot = tab>>2;
        }
      m2p2D_3[d][p]=res|(rot<<6);
      }
  for (int d=0; d<4; ++d)
    for (uint32_t p=0; p<64; ++p)
      {
      unsigned rot = d;
      uint32_t v = p<<26;
      uint32_t res = 0;
      for (int i=0; i<3; ++i)
        {
        unsigned tab=p2m2D_1[rot][v>>30];
        v<<=2;
        res = (res<<2) | (tab&0x3);
        rot = tab>>2;
        }
      p2m2D_3[d][p]=res|(rot<<6);
      }
  }

uint32_t morton2peano2D_32(uint32_t v, int bits)
  {
  if (!peano2d_done) init_peano2d();
  unsigned rot = 0;
  uint32_t res = 0;
  v<<=32-(bits<<1);
  int i=0;
  for (; i<bits-2; i+=3)
    {
    unsigned tab=m2p2D_3[rot][v>>26];
    v<<=6;
    res = (res<<6) | (tab&0x3fu);
    rot = tab>>6;
    }
  for (; i<bits; ++i)
    {
    unsigned tab=m2p2D_1[rot][v>>30];
    v<<=2;
    res = (res<<2) | (tab&0x3);
    rot = tab>>2;
    }
  return res;
  }

//
// Utilities for indirect sorting (argsort)
//

template<typename It, typename Comp> class IdxComp__
  {
  private:
    It begin;
    Comp comp;
  public:
    IdxComp__ (It begin_, Comp comp_): begin(begin_), comp(comp_) {}
    bool operator() (std::size_t a, std::size_t b) const
      { return comp(*(begin+a),*(begin+b)); }
  };
/*! Performs an indirect sort on the supplied iterator range and returns in
    \a idx a \a vector containing the indices of the smallest, second smallest,
    third smallest, etc. element, according to \a comp. */
template<typename It, typename T2, typename Comp>
  inline void buildIndex (It begin, It end, std::vector<T2> &idx, Comp comp)
  {
  using namespace std;
  T2 num=end-begin;
  idx.resize(num);
  for (T2 i=0; i<num; ++i) idx[i] = i;
  sort (idx.begin(),idx.end(),IdxComp__<It,Comp>(begin,comp));
  }

/*! Performs an indirect sort on the supplied iterator range and returns in
    \a idx a \a vector containing the indices of the smallest, second smallest,
    third smallest, etc. element. */
template<typename It, typename T2> inline void buildIndex (It begin, It end,
  std::vector<T2> &idx)
  {
  using namespace std;
  typedef typename iterator_traits<It>::value_type T;
  buildIndex(begin,end,idx,less<T>());
  }

//
// Utilities for Gauss-Legendre quadrature
//

static inline double one_minus_x2 (double x)
  { return (fabs(x)>0.1) ? (1.+x)*(1.-x) : 1.-x*x; }

void legendre_prep(int n, vector<double> &x, vector<double> &w)
  {
  constexpr double pi = 3.141592653589793238462643383279502884197;
  constexpr double eps = 3e-14;
  int m = (n+1)>>1;
  x.resize(m);
  w.resize(m);

  double t0 = 1 - (1-1./n) / (8.*n*n);
  double t1 = 1./(4.*n+2.);

#pragma omp parallel
{
  int i;
#pragma omp for schedule(dynamic,100)
  for (i=1; i<=m; ++i)
    {
    double x0 = cos(pi * ((i<<2)-1) * t1) * t0;

    int dobreak=0;
    int j=0;
    double dpdx;
    while(1)
      {
      double P_1 = 1.0;
      double P0 = x0;
      double dx, x1;

      for (int k=2; k<=n; k++)
        {
        double P_2 = P_1;
        P_1 = P0;
//        P0 = ((2*k-1)*x0*P_1-(k-1)*P_2)/k;
        P0 = x0*P_1 + (k-1.)/k * (x0*P_1-P_2);
        }

      dpdx = (P_1 - x0*P0) * n / one_minus_x2(x0);

      /* Newton step */
      x1 = x0 - P0/dpdx;
      dx = x0-x1;
      x0 = x1;
      if (dobreak) break;

      if (abs(dx)<=eps) dobreak=1;
      if (++j>=100) throw runtime_error("convergence problem");
      }

    x[m-i] = x0;
    w[m-i] = 2. / (one_minus_x2(x0) * dpdx * dpdx);
    }
} // end of parallel region
  }

//
// Start of real gridder functionality
//

template<typename T>
  using pyarr = py::array_t<T>;
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); }

size_t get_w(double epsilon)
  {
  static const vector<double> maxmaperr { 1e8, 0.32, 0.021, 6.2e-4,
    1.08e-5, 1.25e-7, 8.25e-10, 5.70e-12, 1.22e-13, 2.48e-15, 4.82e-17,
    6.74e-19, 5.41e-21, 4.41e-23, 7.88e-25, 3.9e-26 };

  double epssq = epsilon*epsilon;

  for (size_t i=1; i<maxmaperr.size(); ++i)
    if (epssq>maxmaperr[i]) return i;
  throw runtime_error("requested epsilon too small - minimum is 2e-13");
  }

template<typename T> pyarr_c<T> complex2hartley
  (const pyarr_c<complex<T>> &grid_)
  {
  myassert(grid_.ndim()==2, "grid array must be 2D");
  size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
  auto grid = grid_.data();

  auto res = makearray<T>({nu,nv});
  auto grid2 = res.mutable_data();
  for (size_t u=0; u<nu; ++u)
    {
    size_t xu = (u==0) ? 0 : nu-u;
    for (size_t v=0; v<nv; ++v)
      {
      size_t xv = (v==0) ? 0 : nv-v;
      size_t i1 = u*nv+v;
      size_t i2 = xu*nv+xv;
      grid2[i1] = T(0.5)*(grid[i1].real()+grid[i1].imag()+
                          grid[i2].real()-grid[i2].imag());
      }
    }
  return res;
  }

template<typename T> pyarr_c<complex<T>> hartley2complex
  (const pyarr_c<T> &grid_)
  {
  myassert(grid_.ndim()==2, "grid array must be 2D");
  size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
  auto grid = grid_.data();

  auto res=makearray<complex<T>>({nu, nv});
  auto grid2 = res.mutable_data();
  for (size_t u=0; u<nu; ++u)
    {
    size_t xu = (u==0) ? 0 : nu-u;
    for (size_t v=0; v<nv; ++v)
      {
      size_t xv = (v==0) ? 0 : nv-v;
      size_t i1 = u*nv+v;
      size_t i2 = xu*nv+xv;
      T v1 = T(0.5)*grid[i1];
      T v2 = T(0.5)*grid[i2];
      grid2[i1] = complex<T>(v1+v2, v1-v2);
      }
    }
  return res;
  }

/* Compute correction factors for the ES gridding kernel
   This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
vector<double> correction_factors (size_t n, size_t nval, size_t w)
  {
  constexpr double pi = 3.141592653589793238462643383279502884197;
  auto beta = 2.3*w;
  auto p = int(1.5*w+2);
  double alpha = pi*w/n;
  vector<double> x, wgt;
  legendre_prep(2*p,x,wgt);
  auto psi = x;
  for (auto &v:psi)
    v = exp(beta*(sqrt(1-v*v)-1.));
  vector<double> res(nval);
  for (size_t k=0; k<nval; ++k)
    {
    double tmp=0;
    for (int i=0; i<p; ++i)
      tmp += wgt[i]*psi[i]*cos(alpha*k*x[i]);
    res[k] = 1./(w*tmp);
    }
  return res;
  }

template<typename T> struct UV
  {
  T u, v;
  UV () {}
  UV (T u_, T v_) : u(u_), v(v_) {}
  UV operator* (T fct) const
    { return UV(u*fct, v*fct); }
  };
template<typename T> struct UVW
  {
  T u, v, w;
  UVW () {}
  UVW (T u_, T v_, T w_) : u(u_), v(v_), w(w_) {}
  UVW operator* (T fct) const
    { return UVW(u*fct, v*fct, w*fct); }
  };

template<typename T> class Baselines
  {
  private:
    vector<UVW<T>> coord;
    vector<T> scaling;
    size_t nrows, nchan;
    size_t channelbits, channelmask;

  public:
    Baselines(const pyarr_c<T> &coord_, const pyarr_c<T> &scaling_)
      {
      myassert(coord_.ndim()==2, "coord array must be 2D");
      myassert(coord_.shape(1)==3, "coord.shape[1] must be 3");
      myassert(scaling_.ndim()==1, "scaling array must be 1D");
      nrows = coord_.shape(0)/scaling_.shape(0);
      myassert(nrows*size_t(scaling_.shape(0))==size_t(coord_.shape(0)),
        "bad array dimensions");
      nchan = scaling_.shape(0);
      scaling.resize(nchan);
      for (size_t i=0; i<nchan; ++i)
        scaling[i] = scaling_.data()[i];
      coord.resize(nrows);
      auto cood = coord_.data();
      for (size_t i=0; i<coord.size(); ++i)
        coord[i] = UVW<T>(cood[3*i], cood[3*i+1], cood[3*i+2]);
      channelbits = bits_needed(nchan);
      channelmask = (size_t(1)<<channelbits)-1;
      auto rowbits = bits_needed(nrows);
      myassert(rowbits+channelbits<=8*sizeof(uint32_t), "Ti too small");
      }
    vector<uint32_t> getIndices(int chbegin, int chend, T wmin,
      T wmax) const
      {
      if (chbegin<0) chbegin=0;
      if (chend<0) chend=nchan;
      vector<uint32_t> odata;
      for (size_t i=0; i<nrows; ++i)
        for (int j=chbegin; j<chend; ++j)
          {
          auto w = coord[i].w*scaling[j];
          if ((w>=wmin) && (w<wmax))
            odata.push_back((i<<channelbits)+j);
          }
      return odata;
      }

    UVW<T> effectiveCoord(uint32_t index) const
      { return coord[index>>channelbits]*scaling[index&channelmask]; }
    size_t Nrows() const { return nrows; }
    size_t Nchannels() const { return nchan; }
    size_t irow(uint32_t index) const
      { return index>>channelbits; }
    size_t ichannel(uint32_t index) const
      { return index&channelmask; }
    size_t offset(uint32_t index) const
      { return (index>>channelbits)*nchan + (index&channelmask); }

    pyarr_c<complex<T>> ms2vis(const pyarr_c<complex<T>> &ms_,
      const pyarr_c<uint32_t> &idx_) const
      {
      myassert(idx_.ndim()==1, "idx array must be 1D");
      myassert(ms_.ndim()==2, "ms array must be 2D");
      myassert(size_t(ms_.shape(0))==nrows, "baselines 1st dim mismatch");
      myassert(size_t(ms_.shape(1))==nchan, "baselines 2nd dim mismatch");
      size_t nvis = size_t(idx_.shape(0));
      auto idx = idx_.data();
      auto ms = ms_.data();

      auto res=makearray<complex<T>>({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_,
      const pyarr_c<uint32_t> &idx_) const
      {
      myassert(idx_.ndim()==1, "idx array must be 1D");
      myassert(vis_.ndim()==1, "vis array must be 1D");
      size_t nvis = size_t(vis_.shape(0));
      myassert(int(nvis)==idx_.shape(0), "idx/vis size mismatch");
      auto idx = idx_.data();
      auto vis = vis_.data();

      auto res = makearray<complex<T>>({nrows, nchan});
      auto ms = res.mutable_data();
      for (size_t i=0; i<nrows*nchan; ++i)
        ms[i] = complex<T>(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
      {
      myassert(idx_.ndim()==1, "idx array must be 1D");
      myassert(vis_.ndim()==1, "vis array must be 1D");
      myassert(ms_.ndim()==2, "ms array must be 2D");
      myassert(size_t(ms_.shape(0))==nrows, "ms: bad 1st dimension");
      myassert(size_t(ms_.shape(1))==nchan, "ms: bad 2nd dimension");
      size_t nvis = size_t(vis_.shape(0));
      myassert(int(nvis)==idx_.shape(0), "idx/vis size mismatch");
      auto idx = idx_.data();
      auto vis = vis_.data();
      auto ms = ms_.mutable_data();

      for (size_t i=0; i<nvis; ++i)
        ms[offset(idx[i])] += vis[i];
      return ms_;
      }
  };

class GridderConfig
  {
  private:
    size_t nx_dirty, ny_dirty;
    double ucorr, vcorr;
    size_t w, nsafe, nu, nv;
    size_t peano_level;
    vector<double> cfu, cfv;

  public:
    GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
      double urange, double vrange)
      : nx_dirty(nxdirty), ny_dirty(nydirty),
        ucorr(1./urange), vcorr(1./vrange),
        w(get_w(epsilon)), nsafe((w+1)/2),
        nu(max(2*nsafe,2*nx_dirty)), nv(max(2*nsafe,2*ny_dirty)),
        peano_level(bits_needed(max(nu, nv))),
        cfu(nx_dirty), cfv(ny_dirty)
      {
      myassert((nx_dirty&1)==0, "nx_dirty must be even");
      myassert((ny_dirty&1)==0, "ny_dirty must be even");
      myassert(epsilon>0, "epsilon must be positive");
      myassert(urange>0, "urange must be positive");
      myassert(vrange>0, "vrange must be positive");

      auto tmp = correction_factors(nu, nx_dirty/2+1, w);
      cfu[nx_dirty/2]=tmp[0];
      cfu[0]=tmp[nx_dirty/2];
      for (size_t i=1; i<nx_dirty/2; ++i)
        cfu[nx_dirty/2-i] = cfu[nx_dirty/2+i] = tmp[i];
      tmp = correction_factors(nv, ny_dirty/2+1, w);
      cfv[ny_dirty/2]=tmp[0];
      cfv[0]=tmp[ny_dirty/2];
      for (size_t i=1; i<ny_dirty/2; ++i)
        cfv[ny_dirty/2-i] = cfv[ny_dirty/2+i] = tmp[i];
      }
    size_t Nu() const { return nu; }
    size_t Nv() const { return nv; }
    size_t W() const { return w; }
    size_t coord2peano(const UVW<double> &coord) const
      {
      double u=fmodulo(coord.u*ucorr, 1.)*nu,
             v=fmodulo(coord.v*vcorr, 1.)*nv;
      auto iu = min(nu-1, size_t(u));
      auto iv = min(nv-1, size_t(v));
      return morton2peano2D_32(coord2morton2D_32(iu,iv),peano_level);
      }
    pyarr_c<double> grid2dirty(const pyarr_c<double> &grid) const
      {
      auto tmp = makearray<double>({nu, nv});
      auto ptmp = tmp.mutable_data();
      pocketfft::r2r_hartley({nu, nv},
        {grid.strides(0), grid.strides(1)},
        {tmp.strides(0), tmp.strides(1)}, {0,1},
        grid.data(),
        ptmp, 1.);
      for(size_t i=1; i<(nu+1)/2; ++i)
        for(size_t j=1; j<(nv+1)/2; ++j)
           {
           double a = ptmp[i*nv+j];
           double b = ptmp[(nu-i)*nv+j];
           double c = ptmp[i*nv+nv-j];
           double d = ptmp[(nu-i)*nv+nv-j];
           ptmp[i*nv+j] = 0.5*(a+b+c-d);
           ptmp[(nu-i)*nv+j] = 0.5*(a+b+d-c);
           ptmp[i*nv+nv-j] = 0.5*(a+c+d-b);
           ptmp[(nu-i)*nv+nv-j] = 0.5*(b+c+d-a);
           }
      auto res = makearray<double>({nx_dirty, ny_dirty});
      auto pout = res.mutable_data();
      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<double> dirty2grid(const pyarr_c<double> &dirty) const
      {
      auto pdirty = dirty.data();
      auto tmp = makearray<double>({nu, nv});
      auto ptmp = tmp.mutable_data();
      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::r2r_hartley({nu, nv},
        {tmp.strides(0), tmp.strides(1)},
        {tmp.strides(0), tmp.strides(1)}, {0,1},
        tmp.data(), tmp.mutable_data(), 1.);
      for(size_t i=1; i<(nu+1)/2; ++i)
        for(size_t j=1; j<(nv+1)/2; ++j)
           {
           double a = ptmp[i*nv+j];
           double b = ptmp[(nu-i)*nv+j];
           double c = ptmp[i*nv+nv-j];
           double d = ptmp[(nu-i)*nv+nv-j];
           ptmp[i*nv+j] = 0.5*(a+b+c-d);
           ptmp[(nu-i)*nv+j] = 0.5*(a+b+d-c);
           ptmp[i*nv+nv-j] = 0.5*(a+c+d-b);
           ptmp[(nu-i)*nv+nv-j] = 0.5*(b+c+d-a);
           }
      return tmp;
      }
  };

class Helper
  {
  protected:
    int nu, nv;
  public:
    int w;
    double beta;
  protected:
    int nsafe, su;
  public:
    int sv;

    vector<double> kernel;
    int iu0, iv0; // start index of the current visibility
    int bu0, bv0; // start index of the current buffer

    void NOINLINE update(double u_in, double v_in)
      {
      auto u = fmodulo(u_in, 1.)*nu;
      iu0 = int(u-w*0.5 + 1 + nu) - nu;
      if (iu0+w>nu+nsafe) iu0 = nu+nsafe-w;
      auto v = fmodulo(v_in, 1.)*nv;
      iv0 = int(v-w*0.5 + 1 + nv) - nv;
      if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
      double xw=2./w;
      for (int i=0; i<w; ++i)
        {
        kernel[i  ] = xw*(iu0+i-u);
        kernel[i+w] = xw*(iv0+i-v);
        }
      for (auto &k : kernel)
        k = exp(beta*sqrt(1.-k*k));
      }

    bool need_to_move() const
      { return (iu0<bu0) || (iv0<bv0) || (iu0+w>bu0+su) || (iv0+w>bv0+sv); }

    void update_position()
      {
      bu0=max(-nsafe, min(nu+nsafe-su, iu0+nsafe-su/2));
      bv0=max(-nsafe, min(nv+nsafe-sv, iv0+nsafe-sv/2));
      }

  protected:
    Helper(int nu_, int nv_, int w_)
      : nu(nu_), nv(nv_), w(w_), beta(2.3*w), nsafe((w+1)/2),
        su(min(max(2*nsafe,80), nu)), sv(min(max(2*nsafe,80), nv)),
        kernel(2*w),
        bu0(-1000000), bv0(-1000000)
      {
      if (min(nu,nv)<2*nsafe) throw runtime_error("grid dimensions too small");
      }
  };

class WriteHelper: public Helper
  {
  private:
    vector<complex<double>> data;
    complex<double> *grid;

    void dump()
      {
      if (bu0<-nsafe) return; // nothing written into buffer yet
#pragma omp critical
{
      int idxu = (bu0+nu)%nu;
      int idxv0 = (bv0+nv)%nv;
      for (int iu=0; iu<su; ++iu)
        {
        int idxv = idxv0;
        for (int iv=0; iv<sv; ++iv)
          {
          grid[idxu*nv + idxv] += data[iu*sv + iv];
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
}
      }

  public:
    complex<double> *p0;
    WriteHelper(int nu_, int nv_, int w, complex<double> *grid_)
      : Helper(nu_, nv_, w), data(su*sv, 0.), grid(grid_) {}
    ~WriteHelper() { dump(); }

    void prep_write(double u_in, double v_in)
      {
      update(u_in, v_in);
      if (need_to_move())
        {
        dump();
        update_position();
        fill(data.begin(), data.end(), 0.);
        }
      p0 = data.data() + sv*(iu0-bu0) + iv0-bv0;
      }
  };

class ReadHelper: public Helper
  {
  private:
    vector<complex<double>> data;
    const complex<double> *grid;

    void load()
      {
      int idxu = (bu0+nu)%nu;
      int idxv0 = (bv0+nv)%nv;
      for (int iu=0; iu<su; ++iu)
        {
        int idxv = idxv0;
        for (int iv=0; iv<sv; ++iv)
          {
          data[iu*sv + iv] = grid[idxu*nv + idxv];
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
      }

  public:
    const complex<double> *p0;
    ReadHelper(int nu_, int nv_, int w_, const complex<double> *grid_)
      : Helper(nu_, nv_, w_), data(su*sv,0.), grid(grid_), p0(nullptr) {}

    void prep_read(double u_in, double v_in)
      {
      update(u_in, v_in);
      if (need_to_move())
        {
        update_position();
        load();
        }
      p0 = data.data() + sv*(iu0-bu0) + iv0-bv0;
      }
  };

pyarr_c<double> vis2grid(const Baselines<double> &baselines, const GridderConfig &gconf,
  const pyarr_c<uint32_t> &idx_, const pyarr_c<complex<double>> &vis_)
  {
  myassert(idx_.ndim()==1, "idx array must be 1D");
  myassert(vis_.ndim()==1, "vis must be 1D");
  auto vis=vis_.data();
  myassert(vis_.shape(0)==idx_.shape(0), "bad vis dimension");
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.data();

  size_t nu=gconf.Nu(), nv=gconf.Nv();
  auto res = makearray<complex<double>>({nu, nv});
  auto grid = res.mutable_data();
  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;

#pragma omp parallel
{
  WriteHelper hlp(nu, nv, gconf.W(), grid);
  double emb = exp(-2*hlp.beta);
  const double * RESTRICT ku = hlp.kernel.data();
  const double * RESTRICT kv = hlp.kernel.data()+hlp.w;

  // Loop over sampling points
#pragma omp for schedule(dynamic,10000)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<double> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep_write(coord.u, coord.v);
    auto * RESTRICT ptr = hlp.p0;
    int w = hlp.w;
    auto v(vis[ipart]*emb);
    for (int cu=0; cu<w; ++cu)
      {
      complex<double> tmp(v*ku[cu]);
      for (int cv=0; cv<w; ++cv)
        ptr[cv] += tmp*kv[cv];
      ptr+=hlp.sv;
      }
    }
} // end of parallel region
  return complex2hartley(res);
  }
pyarr_c<complex<double>> grid2vis(const Baselines<double> &baselines,
  const GridderConfig &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<double> &grid0_)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  myassert(idx_.ndim()==1, "idx array must be 1D");
  auto grid_ = hartley2complex(grid0_);
  auto grid = grid_.data();
  myassert(grid_.ndim()==2, "data must be 2D");
  myassert(grid_.shape(0)==int(nu), "bad grid dimension");
  myassert(grid_.shape(1)==int(nv), "bad grid dimension");
  size_t nvis = size_t(idx_.shape(0));
  auto idx = idx_.data();

  auto res = makearray<complex<double>>({nvis});
  auto vis = res.mutable_data();

  // Loop over sampling points
#pragma omp parallel
{
  ReadHelper hlp(nu, nv, gconf.W(), grid);
  double emb = exp(-2*hlp.beta);
  const double * RESTRICT ku = hlp.kernel.data();
  const double * RESTRICT kv = hlp.kernel.data()+hlp.w;

#pragma omp for schedule(dynamic,10000)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<double> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep_read(coord.u, coord.v);
    complex<double> r = 0.;
    auto * RESTRICT ptr = hlp.p0;
    int w = hlp.w;
    for (int cu=0; cu<w; ++cu)
      {
      complex<double> tmp(0.);
      for (int cv=0; cv<w; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += hlp.sv;
      }
    vis[ipart] = r*emb;
    }
}
  return res;
  }

pyarr_c<uint32_t> getIndices(const Baselines<double> &baselines,
  const GridderConfig &gconf, int chbegin, int chend, double wmin, double wmax)
  {
  auto idx = baselines.getIndices(chbegin, chend, wmin, wmax);
  vector<size_t> peano(idx.size());
  for (size_t i=0; i<peano.size(); ++i)
    peano[i] = gconf.coord2peano(baselines.effectiveCoord(idx[i]));
  vector<size_t> newind;
  buildIndex(peano.begin(), peano.end(), newind);
  peano=vector<size_t>(); // deallocate
  auto res = makearray<uint32_t>({idx.size()});
  auto iout = res.mutable_data();
  for (size_t i=0; i<idx.size(); ++i)
    iout[i] = idx[newind[i]];
  return res;
  }

const char *Baselines_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
scaling: np.array((nchannels,), dtype=np.float)
    scaling factor for u, v, w for each individual channel
)""";

const char *BL_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
)""";

const char *BL_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

Returns
=======
np.array((nrows, nchannels), dtype=np.complex)
    the measurement set's visibility data (0 where not covered by idx)
)""";

const char *BL_add_vis_to_ms_DS = R"""(
Adds a set of visibilities to an existing MS.

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: np.array((nrows, nchannels), dtype=np.complex)
    a MS

Returns
=======
np.array((nrows, nchannels), dtype=np.complex)
    the updated MS
)""";

} // unnamed namespace

PYBIND11_MODULE(nifty_gridder, m)
  {
  using namespace pybind11::literals;

  py::class_<Baselines<double>> (m, "Baselines", Baselines_DS)
    .def(py::init<pyarr_c<double>, pyarr_c<double>>(), "coord"_a, "scaling"_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,
      "vis"_a, "idx"_a, "ms"_a.noconvert());
  py::class_<GridderConfig> (m, "GridderConfig")
    .def(py::init<size_t, size_t, double, double, double>(),"nxdirty"_a,
      "nydirty"_a, "epsilon"_a, "urange"_a, "vrange"_a)
    .def("Nu", &GridderConfig::Nu)
    .def("Nv", &GridderConfig::Nv)
    .def("grid2dirty", &GridderConfig::grid2dirty, "grid"_a)
    .def("dirty2grid", &GridderConfig::dirty2grid, "dirty"_a);
  m.def("getIndices", getIndices, "baselines"_a, "gconf"_a, "chbegin"_a=-1,
    "chend"_a=-1, "wmin"_a=-1e30, "wmax"_a=1e30);
  m.def("vis2grid",&vis2grid, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a);
  m.def("grid2vis",&grid2vis, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a);
  }