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;
    vector<bool> flags;
    size_t nrows, nchan;
    size_t channelbits, channelmask;

  public:
    Baselines(const pyarr_c<T> &coord_, const pyarr_c<T> &scaling_,
       const pyarr_c<bool> &flags_)
      {
      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);
      nchan = scaling_.shape(0);
      myassert(flags_.ndim()==2, "flags array must be 2D");
      myassert(size_t(flags_.shape(0))==nrows, "flags.shape[0] must be nrows");
      myassert(size_t(flags_.shape(1))==nchan, "flags.shape[1] must be chan");
      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]);
      flags.resize(nrows*nchan);
      auto pflags=flags_.data();
      for (size_t i=0; i<nrows*nchan; ++i)
        flags[i] = (pflags[i]!=0);
      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)
          if (!flags[i*nchan + 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); }

    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");
      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<T2>({nvis});
      auto vis = res.mutable_data();
      for (size_t i=0; i<nvis; ++i)
        vis[i] = ms[offset(idx[i])];
      return res;
      }

    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");
      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<T2>({nrows, nchan});
      auto ms = res.mutable_data();
      for (size_t i=0; i<nrows*nchan; ++i)
        ms[i] = T2(0);
      for (size_t i=0; i<nvis; ++i)
        ms[offset(idx[i])] = vis[i];
      return res;
      }

    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");
      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_;
      }
  };

template<typename T> class GridderConfig
  {
  private:
    size_t nx_dirty, ny_dirty;
    T ucorr, vcorr;
    size_t w, nsafe, nu, nv;
    size_t peano_level;
    vector<T> 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; }
    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,
             v=fmodulo(coord.v*vcorr, T(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<T> grid2dirty(const pyarr_c<T> &grid) const
      {
      auto tmp = makearray<T>({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, T(1), 0);
      for(size_t i=1; i<(nu+1)/2; ++i)
        for(size_t j=1; j<(nv+1)/2; ++j)
           {
           T a = ptmp[i*nv+j];
           T b = ptmp[(nu-i)*nv+j];
           T c = ptmp[i*nv+nv-j];
           T d = ptmp[(nu-i)*nv+nv-j];
           ptmp[i*nv+j] = T(0.5)*(a+b+c-d);
           ptmp[(nu-i)*nv+j] = T(0.5)*(a+b+d-c);
           ptmp[i*nv+nv-j] = T(0.5)*(a+c+d-b);
           ptmp[(nu-i)*nv+nv-j] = T(0.5)*(b+c+d-a);
           }
      auto res = makearray<T>({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<T> dirty2grid(const pyarr_c<T> &dirty) const
      {
      auto pdirty = dirty.data();
      auto tmp = makearray<T>({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(), T(1), 0);
      for(size_t i=1; i<(nu+1)/2; ++i)
        for(size_t j=1; j<(nv+1)/2; ++j)
           {
           T a = ptmp[i*nv+j];
           T b = ptmp[(nu-i)*nv+j];
           T c = ptmp[i*nv+nv-j];
           T d = ptmp[(nu-i)*nv+nv-j];
           ptmp[i*nv+j] = T(0.5)*(a+b+c-d);
           ptmp[(nu-i)*nv+j] = T(0.5)*(a+b+d-c);
           ptmp[i*nv+nv-j] = T(0.5)*(a+c+d-b);
           ptmp[(nu-i)*nv+nv-j] = T(0.5)*(b+c+d-a);
           }
      return tmp;
      }
  };

template<typename T> class Helper
  {
  protected:
    int nu, nv;
  public:
    int w;
    T beta;
  protected:
    int nsafe, su;
  public:
    int sv;

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

    void NOINLINE update(T u_in, T v_in)
      {
      auto u = fmodulo(u_in, T(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, T(1))*nv;
      iv0 = int(v-w*0.5 + 1 + nv) - nv;
      if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
      T xw=T(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(T(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");
      }
  };

template<typename T> class WriteHelper: public Helper<T>
  {
  protected:
    using Helper<T>::nu;
    using Helper<T>::nv;
  public:
    using Helper<T>::w;
    using Helper<T>::beta;
  protected:
    using Helper<T>::nsafe;
    using Helper<T>::su;
  public:
    using Helper<T>::sv;
    using Helper<T>::kernel;
    using Helper<T>::iu0;
    using Helper<T>::iv0;
    using Helper<T>::bu0;
    using Helper<T>::bv0;
    using Helper<T>::need_to_move;
    using Helper<T>::update_position;
    using Helper<T>::update;

  private:
    vector<complex<T>> data;
    complex<T> *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<T> *p0;
    WriteHelper(int nu_, int nv_, int w, complex<T> *grid_)
      : Helper<T>(nu_, nv_, w), data(su*sv, T(0)), grid(grid_) {}
    ~WriteHelper() { dump(); }

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

template<typename T> class ReadHelper: public Helper<T>
  {
  protected:
    using Helper<T>::nu;
    using Helper<T>::nv;
  public:
    using Helper<T>::w;
    using Helper<T>::beta;
  protected:
    using Helper<T>::nsafe;
    using Helper<T>::su;
  public:
    using Helper<T>::sv;
    using Helper<T>::kernel;
    using Helper<T>::iu0;
    using Helper<T>::iv0;
    using Helper<T>::bu0;
    using Helper<T>::bv0;
    using Helper<T>::need_to_move;
    using Helper<T>::update_position;
    using Helper<T>::update;

  private:
    vector<complex<T>> data;
    const complex<T> *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<T> *p0;
    ReadHelper(int nu_, int nv_, int w_, const complex<T> *grid_)
      : Helper<T>(nu_, nv_, w_), data(su*sv,T(0)), grid(grid_), p0(nullptr) {}

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

template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<complex<T>> &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<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
{
  WriteHelper<T> hlp(nu, nv, gconf.W(), grid);
  T emb = exp(-2*hlp.beta);
  const T * RESTRICT ku = hlp.kernel.data();
  const T * 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<T> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep_write(coord.u*ucorr, coord.v*vcorr);
    auto * RESTRICT ptr = hlp.p0;
    int w = hlp.w;
    auto v(vis[ipart]*emb);
    for (int cu=0; cu<w; ++cu)
      {
      complex<T> 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);
  }

template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const pyarr_c<uint32_t> &idx_,
  const pyarr_c<T> &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<T>>({nvis});
  auto vis = res.mutable_data();
  T ucorr = gconf.Ucorr(), vcorr=gconf.Vcorr();

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

#pragma omp for schedule(dynamic,10000)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    UVW<T> coord = baselines.effectiveCoord(idx[ipart]);
    hlp.prep_read(coord.u*ucorr, coord.v*vcorr);
    complex<T> r = 0;
    auto * RESTRICT ptr = hlp.p0;
    int w = hlp.w;
    for (int cu=0; cu<w; ++cu)
      {
      complex<T> 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;
  }

template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, int chbegin, int chend, T wmin, T 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
flags: np.array((nrows, nchannels), dtype=np.bool)
    "True" indicates that the value should not be used
)""";

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<const pyarr_c<double> &, const pyarr_c<double> &,
      const pyarr_c<bool> &>(), "coord"_a, "scaling"_a, "flags"_a)
    .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,
      "nydirty"_a, "epsilon"_a, "urange"_a, "vrange"_a)
    .def("Nu", &GridderConfig<double>::Nu)
    .def("Nv", &GridderConfig<double>::Nv)
    .def("grid2dirty", &GridderConfig<double>::grid2dirty, "grid"_a)
    .def("dirty2grid", &GridderConfig<double>::dirty2grid, "dirty"_a);
  m.def("getIndices", getIndices<double>, "baselines"_a, "gconf"_a, "chbegin"_a=-1,
    "chend"_a=-1, "wmin"_a=-1e30, "wmax"_a=1e30);
  m.def("vis2grid",&vis2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a);
  m.def("grid2vis",&grid2vis<double>, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a);

  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<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,
      "nydirty"_a, "epsilon"_a, "urange"_a, "vrange"_a)
    .def("Nu", &GridderConfig<float>::Nu)
    .def("Nv", &GridderConfig<float>::Nv)
    .def("grid2dirty", &GridderConfig<float>::grid2dirty, "grid"_a)
    .def("dirty2grid", &GridderConfig<float>::dirty2grid, "dirty"_a);
  m.def("getIndices_f", getIndices<float>, "baselines"_a, "gconf"_a, "chbegin"_a=-1,
    "chend"_a=-1, "wmin"_a=-1e30, "wmax"_a=1e30);
  m.def("vis2grid_f",&vis2grid<float>, "baselines"_a, "gconf"_a, "idx"_a, "vis"_a);
  m.def("grid2vis_f",&grid2vis<float>, "baselines"_a, "gconf"_a, "idx"_a, "grid"_a);
  }