gridder_cxx.h

#ifndef GRIDDER_CXX_H
#define GRIDDER_CXX_H

/*
 *  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 <iostream>
#include <algorithm>
#include <cstdlib>
#include <cmath>
#include <vector>
#include <array>

#include "pocketfft_hdronly.h"

#if defined(__GNUC__)
#define NOINLINE __attribute__((noinline))
#define ALIGNED(align) __attribute__ ((aligned(align)))
#define RESTRICT __restrict__
#else
#define NOINLINE
#define ALIGNED(align)
#define RESTRICT
#endif

namespace gridder {

namespace detail {

using namespace std;

template<typename T> struct VLEN { static constexpr size_t val=1; };

#if (defined(__AVX512F__))
template<> struct VLEN<float> { static constexpr size_t val=16; };
template<> struct VLEN<double> { static constexpr size_t val=8; };
#elif (defined(__AVX__))
template<> struct VLEN<float> { static constexpr size_t val=8; };
template<> struct VLEN<double> { static constexpr size_t val=4; };
#elif (defined(__SSE2__))
template<> struct VLEN<float> { static constexpr size_t val=4; };
template<> struct VLEN<double> { static constexpr size_t val=2; };
#elif (defined(__VSX__))
template<> struct VLEN<float> { static constexpr size_t val=4; };
template<> struct VLEN<double> { static constexpr size_t val=2; };
#endif

template<typename T, size_t ndim> class mav
  {
  static_assert((ndim>0) && (ndim<3), "only supports 1D and 2D arrays");

  private:
    T *d;
    array<size_t, ndim> shp;
    array<ptrdiff_t, ndim> str;

  public:
    mav(T *d_, const array<size_t,ndim> &shp_,
        const array<ptrdiff_t,ndim> &str_)
      : d(d_), shp(shp_), str(str_) {}
    mav(T *d_, const array<size_t,ndim> &shp_)
      : d(d_), shp(shp_)
      {
      str[ndim-1]=1;
      for (size_t d=2; d<=ndim; ++d)
        str[ndim-d] = str[ndim-d+1]*shp[ndim-d+1];
      }
    operator mav<const T, ndim>() const
      { return mav<const T, ndim>(d,shp,str); }
    T &operator[](size_t i) const
      { return operator()(i); }
    T &operator()(size_t i) const
      {
      static_assert(ndim==1, "ndim must be 1");
      return d[str[0]*i];
      }
    T &operator()(size_t i, size_t j) const
      {
      static_assert(ndim==2, "ndim must be 2");
      return d[str[0]*i + str[1]*j];
      }
    size_t shape(size_t i) const { return shp[i]; }
    const array<size_t,ndim> &shape() const { return shp; }
    size_t size() const
      {
      size_t res=1;
      for (auto v: shp) res*=v;
      return res;
      }
    ptrdiff_t stride(size_t i) const { return str[i]; }
    T *data() const
      { return d; }
    bool contiguous() const
      {
      ptrdiff_t stride=1;
      for (size_t i=0; i<ndim; ++i)
        {
        if (str[ndim-1-i]!=stride) return false;
        stride *= shp[ndim-1-i];
        }
      return true;
      }
    void fill(const T &val) const
      {
      if (ndim==1)
        for (size_t i=0; i<shp[0]; ++i)
          d[str[0]*i]=val;
      else if (ndim==2)
        for (size_t i=0; i<shp[0]; ++i)
          for (size_t j=0; j<shp[1]; ++j)
            d[str[0]*i + str[1]*j] = val;
      }
  };

template<typename T, size_t ndim> using const_mav = mav<const T, ndim>;
template<typename T, size_t ndim> const_mav<T, ndim> cmav (const mav<T, ndim> &mav)
  { return const_mav<T, ndim>(mav.data(), mav.shape()); }
template<typename T, size_t ndim> const_mav<T, ndim> nullmav()
  {
  array<size_t,ndim> shp;
  shp.fill(0);
  return const_mav<T, ndim>(nullptr, shp);
  }
template<typename T, size_t ndim> class tmpStorage
  {
  private:
    vector<T> d;
    mav<T,ndim> mav_;

    static size_t prod(const array<size_t,ndim> &shp)
      {
      size_t res=1;
      for (auto v: shp) res*=v;
      return res;
      }

  public:
    tmpStorage(const array<size_t,ndim> &shp)
      : d(prod(shp)), mav_(d.data(), shp) {}
    mav<T,ndim> &getMav() { return mav_; }
    const_mav<T,ndim> getCmav() { return cmav(mav_); }
    void fill(const T & val)
      { std::fill(d.begin(), d.end(), val); }
  };
//
// basic utilities
//

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

template<size_t ndim> void checkShape
  (const array<size_t, ndim> &shp1, const array<size_t, ndim> &shp2)
  {
  for (size_t i=0; i<ndim; ++i)
    myassert(shp1[i]==shp2[i], "shape mismatch");
  }

/*! 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 = v1+v2;
  if (tmp<0)
    tmp=fmod(v1,v2)+v2;
  return (tmp==v2) ? T(0) : tmp;
  }

//
// 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, size_t nthreads)
  {
  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 num_threads(nthreads)
{
  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;
      myassert(++j<100, "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
//

size_t get_supp(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");
  }

struct RowChan
  {
  size_t row, chan;
  };

template<typename T> void complex2hartley
  (const const_mav<complex<T>, 2> &grid, const mav<T,2> &grid2, size_t nthreads)
  {
  myassert(grid.shape()==grid2.shape(), "shape mismatch");
  size_t nu=grid.shape(0), nv=grid.shape(1);

#pragma omp parallel for num_threads(nthreads)
  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;
      grid2(u,v) += T(0.5)*(grid( u, v).real()+grid( u, v).imag()+
                            grid(xu,xv).real()-grid(xu,xv).imag());
      }
    }
  }

template<typename T> void hartley2complex
  (const const_mav<T,2> &grid, const mav<complex<T>,2> &grid2, size_t nthreads)
  {
  myassert(grid.shape()==grid2.shape(), "shape mismatch");
  size_t nu=grid.shape(0), nv=grid.shape(1);

#pragma omp parallel for num_threads(nthreads)
  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;
      T v1 = T(0.5)*grid( u, v);
      T v2 = T(0.5)*grid(xu,xv);
      grid2(u,v) = std::complex<T>(v1+v2, v1-v2);
      }
    }
  }

template<typename T> void hartley2_2D(const const_mav<T,2> &in, const mav<T,2> &out, size_t nthreads)
  {
  myassert(in.shape()==out.shape(), "shape mismatch");
  size_t nu=in.shape(0), nv=in.shape(1);
  ptrdiff_t sz=ptrdiff_t(sizeof(T));
  pocketfft::stride_t stri{sz*in.stride(0), sz*in.stride(1)};
  pocketfft::stride_t stro{sz*out.stride(0), sz*out.stride(1)};
  auto d_i = in.data();
  auto ptmp = out.data();
  pocketfft::r2r_separable_hartley({nu, nv}, stri, stro, {0,1}, d_i, ptmp, T(1),
    nthreads);
#pragma omp parallel for num_threads(nthreads)
  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);
       }
  }


template<typename T> class EC_Kernel
  {
  protected:
    T beta;

  public:
    EC_Kernel(size_t supp) : beta(2.3*supp) {}
    T operator()(T v) const { return exp(beta*(sqrt(T(1)-v*v)-T(1))); }
  };

template<typename T> class EC_Kernel_with_correction: public EC_Kernel<T>
  {
  protected:
    static constexpr T pi = T(3.141592653589793238462643383279502884197L);
    int p;
    vector<double> x, wgt, psi;
    size_t supp;

  public:
    using EC_Kernel<T>::operator();
    EC_Kernel_with_correction(size_t supp_, size_t nthreads)
      : EC_Kernel<T>(supp_), p(int(1.5*supp_+2)), supp(supp_)
      {
      legendre_prep(2*p,x,wgt,nthreads);
      psi=x;
      for (auto &v:psi)
        v=operator()(v);
      }
    /* Compute correction factors for the ES gridding kernel
       This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
    T corfac(T v) const
      {
      T tmp=0;
      for (int i=0; i<p; ++i)
        tmp += wgt[i]*psi[i]*cos(pi*supp*v*x[i]);
      return T(1./(supp*tmp));
      }
  };
/* 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 supp,
  size_t nthreads)
  {
  EC_Kernel_with_correction<double> kernel(supp, nthreads);
  vector<double> res(nval);
  double xn = 1./n;
#pragma omp parallel for schedule(static) num_threads(nthreads)
  for (size_t k=0; k<nval; ++k)
    res[k] = kernel.corfac(k*xn);
  return res;
  }

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); }
  void Flip() { u=-u; v=-v; w=-w; }
  bool FixW()
    {
    bool flip = w<0;
    if (flip) Flip();
    return flip;
    }
  };

template<typename T> class Baselines
  {
  protected:
    vector<UVW<T>> coord;
    vector<T> f_over_c;
    size_t nrows, nchan;
    size_t shift, mask;

  public:
    Baselines(const const_mav<T,2> &coord_, const const_mav<T,1> &freq)
      {
      constexpr double speedOfLight = 299792458.;
      myassert(coord_.shape(1)==3, "dimension mismatch");
      nrows = coord_.shape(0);
      nchan = freq.shape(0);
      shift=0;
      while((size_t(1)<<shift)<nchan) ++shift;
      mask=(size_t(1)<<shift)-1;
      myassert(nrows*nchan<(size_t(1)<<32), "too many entries in MS");
      f_over_c.resize(nchan);
      for (size_t i=0; i<nchan; ++i)
        {
        myassert(freq[i]>0, "negative channel frequency encountered");
        f_over_c[i] = freq(i)/speedOfLight;
        }
      coord.resize(nrows);
      for (size_t i=0; i<coord.size(); ++i)
        coord[i] = UVW<T>(coord_(i,0), coord_(i,1), coord_(i,2));
      }

    RowChan getRowChan(uint32_t index) const
      { return RowChan{index>>shift, index&mask}; }

    UVW<T> effectiveCoord(const RowChan &rc) const
      { return coord[rc.row]*f_over_c[rc.chan]; }
    UVW<T> effectiveCoord(uint32_t index) const
      { return effectiveCoord(getRowChan(index)); }
    size_t Nrows() const { return nrows; }
    size_t Nchannels() const { return nchan; }
    uint32_t getIdx(size_t irow, size_t ichan) const
      { return ichan+(irow<<shift); }

    void effectiveUVW(const mav<const uint32_t,1> &idx, mav<T,2> &res) const
      {
      size_t nvis = idx.shape(0);
      myassert(res.shape(0)==nvis, "shape mismatch");
      myassert(res.shape(1)==3, "shape mismatch");
      for (size_t i=0; i<nvis; i++)
        {
        auto uvw = effectiveCoord(idx(i));
        res(i,0) = uvw.u;
        res(i,1) = uvw.v;
        res(i,2) = uvw.w;
        }
      }

    template<typename T2> void ms2vis(const mav<const T2,2> &ms,
      const mav<const uint32_t,1> &idx, mav<T2,1> &vis, size_t nthreads) const
      {
      myassert(ms.shape(0)==nrows, "shape mismatch");
      myassert(ms.shape(1)==nchan, "shape mismatch");
      size_t nvis = idx.shape(0);
      myassert(vis.shape(0)==nvis, "shape mismatch");
#pragma omp parallel for num_threads(nthreads)
      for (size_t i=0; i<nvis; ++i)
        {
        auto rc = getRowChan(idx(i));
        vis[i] = ms(rc.row, rc.chan);
        }
      }

    template<typename T2> void vis2ms(const mav<const T2,1> &vis,
      const mav<const uint32_t,1> &idx, mav<T2,2> &ms, size_t nthreads) const
      {
      size_t nvis = vis.shape(0);
      myassert(idx.shape(0)==nvis, "shape mismatch");
      myassert(ms.shape(0)==nrows, "shape mismatch");
      myassert(ms.shape(1)==nchan, "shape mismatch");
#pragma omp parallel for num_threads(nthreads)
      for (size_t i=0; i<nvis; ++i)
        {
        auto rc = getRowChan(idx(i));
        ms(rc.row, rc.chan) += vis(i);
        }
      }
  };

template<typename T> class GridderConfig
  {
  protected:
    size_t nx_dirty, ny_dirty;
    T eps, psx, psy;
    size_t supp, nsafe, nu, nv;
    T beta;
    vector<T> cfu, cfv;
    size_t nthreads;

    complex<T> wscreen(double x, double y, double w, bool adjoint) const
      {
      constexpr double pi = 3.141592653589793238462643383279502884197;
      double tmp = 1-x-y;
      if (tmp<0) return 0.;
      double nm1 = (-x-y)/(sqrt(tmp)+1); // more accurate form of sqrt(1-x-y)
      double n = nm1+1., xn = 1./n;
      double phase = 2*pi*w*nm1;
      if (adjoint) phase *= -1;
      return complex<T>(cos(phase)*xn, sin(phase)*xn);
      }

  public:
    GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
      double pixsize_x, double pixsize_y, size_t nthreads_)
      : nx_dirty(nxdirty), ny_dirty(nydirty), eps(epsilon),
        psx(pixsize_x), psy(pixsize_y),
        supp(get_supp(epsilon)), nsafe((supp+1)/2),
        nu(max(2*nsafe,2*nx_dirty)), nv(max(2*nsafe,2*ny_dirty)),
        beta(2.3*supp),
        cfu(nx_dirty), cfv(ny_dirty), nthreads(nthreads_)
      {
      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(pixsize_x>0, "pixsize_x must be positive");
      myassert(pixsize_y>0, "pixsize_y must be positive");

      auto tmp = correction_factors(nu, nx_dirty/2+1, supp, nthreads);
      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, supp, nthreads);
      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 Nxdirty() const { return nx_dirty; }
    size_t Nydirty() const { return ny_dirty; }
    double Epsilon() const { return eps; }
    double Pixsize_x() const { return psx; }
    double Pixsize_y() const { return psy; }
    size_t Nu() const { return nu; }
    size_t Nv() const { return nv; }
    size_t Supp() const { return supp; }
    size_t Nsafe() const { return nsafe; }
    T Beta() const { return beta; }
    size_t Nthreads() const { return nthreads; }

    void apply_taper(const mav<const T,2> &img, mav<T,2> &img2, bool divide) const
      {
      checkShape(img.shape(), {nx_dirty, ny_dirty});
      checkShape(img2.shape(), {nx_dirty, ny_dirty});
      if (divide)
        for (size_t i=0; i<nx_dirty; ++i)
          for (size_t j=0; j<ny_dirty; ++j)
            img2(i,j) = img(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)
            img2(i,j) = img(i,j)*cfu[i]*cfv[j];
      }

    void grid2dirty(const const_mav<T,2> &grid, const mav<T,2> &dirty) const
      {
      checkShape(grid.shape(), {nu,nv});
      checkShape(dirty.shape(), {nx_dirty,ny_dirty});
      tmpStorage<T,2> tmpdat({nu,nv});
      auto tmav = tmpdat.getMav();
      hartley2_2D<T>(grid, tmav, nthreads);
      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;
          dirty(i,j) = tmav(i2,j2)*cfu[i]*cfv[j];
          }
      }

    void grid2dirty_c(const mav<const complex<T>,2> &grid, mav<complex<T>,2> &dirty) const
      {
      checkShape(grid.shape(), {nu,nv});
      checkShape(dirty.shape(), {nx_dirty,ny_dirty});
      tmpStorage<complex<T>,2> tmpdat({nu,nv});
      auto tmp = tmpdat.getMav();
      constexpr auto sc = ptrdiff_t(sizeof(complex<T>));
      pocketfft::c2c({nu,nv},{grid.stride(0)*sc,grid.stride(1)*sc},
        {tmp.stride(0)*sc, tmp.stride(1)*sc}, {0,1}, pocketfft::BACKWARD,
        grid.data(), tmp.data(), T(1), nthreads);
      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;
          dirty(i,j) = tmp(i2,j2)*cfu[i]*cfv[j];
          }
      }

    void dirty2grid(const const_mav<T,2> &dirty, mav<T,2> &grid) const
      {
      checkShape(dirty.shape(), {nx_dirty, ny_dirty});
      checkShape(grid.shape(), {nu, nv});
      grid.fill(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;
          grid(i2,j2) = dirty(i,j)*cfu[i]*cfv[j];
          }
      hartley2_2D<T>(grid, grid, nthreads);
      }

    void dirty2grid_c(const const_mav<complex<T>,2> &dirty,
      mav<complex<T>,2> &grid) const
      {
      checkShape(dirty.shape(), {nx_dirty, ny_dirty});
      checkShape(grid.shape(), {nu, nv});
      grid.fill(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;
          grid(i2,j2) = dirty(i,j)*cfu[i]*cfv[j];
          }
      constexpr auto sc = ptrdiff_t(sizeof(complex<T>));
      pocketfft::stride_t strides{grid.stride(0)*sc,grid.stride(1)*sc};
      pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
        grid.data(), grid.data(), T(1), nthreads);
      }

    void getpix(T u_in, T v_in, T &u, T &v, int &iu0, int &iv0) const
      {
      u=fmodulo(u_in*psx, T(1))*nu,
      iu0 = int(u-supp*0.5 + 1 + nu) - nu;
      iu0 = min<int>(iu0, (nu+nsafe)-supp);
      v=fmodulo(v_in*psy, T(1))*nv;
      iv0 = int(v-supp*0.5 + 1 + nv) - nv;
      iv0 = min<int>(iv0, (nv+nsafe)-supp);
      }

    void apply_wscreen(const const_mav<complex<T>,2> &dirty,
      mav<complex<T>,2> &dirty2, double w, bool adjoint) const
      {
      checkShape(dirty.shape(), {nx_dirty, ny_dirty});
      checkShape(dirty2.shape(), {nx_dirty, ny_dirty});
      double x0 = -0.5*nx_dirty*psx,
             y0 = -0.5*ny_dirty*psy;
#pragma omp parallel for num_threads(nthreads) 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);
          dirty2(i,j) = dirty(i,j)*ws; // lower left
          size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
          if ((i>0)&&(i<i2))
            {
            dirty2(i2,j) = dirty(i2,j)*ws; // lower right
            if ((j>0)&&(j<j2))
              dirty2(i2,j2) = dirty(i2,j2)*ws; // upper right
            }
          if ((j>0)&&(j<j2))
            dirty2(i,j2) = dirty(i,j2)*ws; // upper left
          }
        }
      }

  };

constexpr int logsquare=4;

template<typename T, typename T2=complex<T>> class Helper
  {
  private:
    const GridderConfig<T> &gconf;
    int nu, nv, nsafe, supp;
    T beta;
    const T2 *grid_r;
    T2 *grid_w;
    int su, sv;
    int iu0, iv0; // start index of the current visibility
    int bu0, bv0; // start index of the current buffer

    vector<T2> rbuf, wbuf;
    bool do_w_gridding;
    T w0, xdw;
    size_t nexp;
    size_t nvecs;

    void dump() const
      {
      if (bu0<-nsafe) return; // nothing written into buffer yet

#pragma omp critical (gridder_writing_to_grid)
{
      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_w[idxu*nv + idxv] += wbuf[iu*sv + iv];
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
}
      }

    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)
          {
          rbuf[iu*sv + iv] = grid_r[idxu*nv + idxv];
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
      }

  public:
    const T2 *p0r;
    T2 *p0w;
    T kernel[64] ALIGNED(64);

    Helper(const GridderConfig<T> &gconf_, const T2 *grid_r_, T2 *grid_w_, T w0_=-1, T dw_=-1)
      : gconf(gconf_), nu(gconf.Nu()), nv(gconf.Nv()), nsafe(gconf.Nsafe()),
        supp(gconf.Supp()), beta(gconf.Beta()), grid_r(grid_r_),
        grid_w(grid_w_), su(2*nsafe+(1<<logsquare)), sv(2*nsafe+(1<<logsquare)),
        bu0(-1000000), bv0(-1000000),
        rbuf(su*sv*(grid_r!=nullptr),T(0)),
        wbuf(su*sv*(grid_w!=nullptr),T(0)),
        do_w_gridding(dw_>0),
        w0(w0_),
        xdw(T(1)/dw_),
        nexp(2*supp + do_w_gridding),
        nvecs(VLEN<T>::val*((nexp+VLEN<T>::val-1)/VLEN<T>::val))
      {}
    ~Helper() { if (grid_w) dump(); }

    int lineJump() const { return sv; }
    T Wfac() const { return kernel[2*supp]; }
    void prep(const UVW<T> &in)
      {
      T u, v;
      gconf.getpix(in.u, in.v, u, v, iu0, iv0);
      T xsupp=T(2)/supp;
      auto x0 = xsupp*(iu0-u);
      auto y0 = xsupp*(iv0-v);
      for (int i=0; i<supp; ++i)
        {
        kernel[i  ] = x0+i*xsupp;
        kernel[i+supp] = y0+i*xsupp;
        }
      if (do_w_gridding)
        kernel[2*supp] = min(T(1), xdw*xsupp*abs(w0-in.w));
      for (size_t i=nexp; i<nvecs; ++i)
        kernel[i]=0;
      for (size_t i=0; i<nvecs; ++i)
        kernel[i] = exp(beta*(sqrt(T(1)-kernel[i]*kernel[i])-T(1)));
      if ((iu0<bu0) || (iv0<bv0) || (iu0+supp>bu0+su) || (iv0+supp>bv0+sv))
        {
        if (grid_w) { dump(); fill(wbuf.begin(), wbuf.end(), T(0)); }
        bu0=((((iu0+nsafe)>>logsquare)<<logsquare))-nsafe;
        bv0=((((iv0+nsafe)>>logsquare)<<logsquare))-nsafe;
        if (grid_r) load();
        }
      p0r = grid_r ? rbuf.data() + sv*(iu0-bu0) + iv0-bv0 : nullptr;
      p0w = grid_w ? wbuf.data() + sv*(iu0-bu0) + iv0-bv0 : nullptr;
      }
  };

template<class T, class T2, class T3> class VisServ
  {
  private:
    const Baselines<T> &baselines;
    const const_mav<uint32_t,1> &idx;
    T2 vis;
    const T3 &wgt;
    size_t nvis;
    bool have_wgt;

  public:
    VisServ(const Baselines<T> &baselines_,
      const const_mav<uint32_t,1> &idx_, T2 vis_, const T3 &wgt_)
      : baselines(baselines_), idx(idx_), vis(vis_), wgt(wgt_)
      {
      nvis = vis.shape(0);
      checkShape(idx.shape(), {nvis});
      have_wgt = wgt.size()!=0;
      if (have_wgt) checkShape(wgt.shape(), {nvis});
      }
    size_t Nvis() const { return nvis; }
    const Baselines<T> &getBaselines() const { return baselines; }
    UVW<T> getCoord(size_t i) const
      { return baselines.effectiveCoord(idx[i]); }
    complex<T> getVis(size_t i) const
      { return have_wgt ? vis(i)*wgt(i) : vis(i); }
    uint32_t getIdx(size_t i) const { return idx[i]; }
    void setVis (size_t i, const complex<T> &v) const
      { vis(i) = have_wgt ? v*wgt(i) : v; }
    void addVis (size_t i, const complex<T> &v) const
      { vis(i) += have_wgt ? v*wgt(i) : v; }
  };
template<class T, class T2, class T3> class MsServ
  {
  private:
    const Baselines<T> &baselines;
    const const_mav<uint32_t,1> &idx;
    T2 ms;
    const T3 &wgt;
    size_t nvis;
    bool have_wgt;

  public:
    MsServ(const Baselines<T> &baselines_,
    const const_mav<uint32_t,1> &idx_, T2 ms_, const T3 &wgt_)
      : baselines(baselines_), idx(idx_), ms(ms_), wgt(wgt_)
      {
      auto nrows = baselines.Nrows();
      auto nchan = baselines.Nchannels();
      nvis = idx.shape(0);
      checkShape(ms.shape(), {nrows, nchan});
      have_wgt = wgt.size()!=0;
      if (have_wgt) checkShape(wgt.shape(), {nrows, nchan});
      }
    MsServ(const MsServ &orig, const const_mav<uint32_t,1> &newidx)
      : MsServ(orig.baselines, newidx, orig.ms, orig.wgt) {}
    size_t Nvis() const { return nvis; }
    const Baselines<T> &getBaselines() const { return baselines; }
    UVW<T> getCoord(size_t i) const
      { return baselines.effectiveCoord(idx(i)); }
    complex<T> getVis(size_t i) const
      {
      auto rc = baselines.getRowChan(idx(i));
      return have_wgt ? ms(rc.row, rc.chan)*wgt(rc.row, rc.chan)
                      : ms(rc.row, rc.chan);
      }
    uint32_t getIdx(size_t i) const { return idx[i]; }
    void setVis (size_t i, const complex<T> &v) const
      {
      auto rc = baselines.getRowChan(idx(i));
      ms(rc.row, rc.chan) = have_wgt ? v*wgt(rc.row, rc.chan) : v;
      }
    void addVis (size_t i, const complex<T> &v) const
      {
      auto rc = baselines.getRowChan(idx(i));
      ms(rc.row, rc.chan) += have_wgt ? v*wgt(rc.row, rc.chan) : v;
      }
  };

template<typename T, typename Serv> void x2grid_c
  (const GridderConfig<T> &gconf, const Serv &srv, mav<complex<T>,2> &grid, T w0=-1, T dw=-1)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkShape(grid.shape(), {nu, nv});
  myassert(grid.contiguous(), "grid is not contiguous");
  size_t supp = gconf.Supp();
  size_t nthreads = gconf.Nthreads();
  bool do_w_gridding=dw>0;

#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, nullptr, grid.data(), w0, dw);
  int jump = hlp.lineJump();
  const T * RESTRICT ku = hlp.kernel;
  const T * RESTRICT kv = hlp.kernel+supp;
  size_t np = srv.Nvis();

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<np; ++ipart)
    {
    UVW<T> coord = srv.getCoord(ipart);
    auto flip = coord.FixW();
    hlp.prep(coord);
    auto * RESTRICT ptr = hlp.p0w;
    auto v(srv.getVis(ipart));
    if (do_w_gridding) v*=hlp.Wfac();
    if (flip) v=conj(v);
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(v*ku[cu]);
      size_t cv=0;
      for (; cv<supp-3; cv+=4)
        {
        ptr[cv  ] += tmp*kv[cv  ];
        ptr[cv+1] += tmp*kv[cv+1];
        ptr[cv+2] += tmp*kv[cv+2];
        ptr[cv+3] += tmp*kv[cv+3];
        }
      for (; cv<supp; ++cv)
        ptr[cv] += tmp*kv[cv];
      ptr+=jump;
      }
    }
} // end of parallel region
  }

template<typename T> void vis2grid_c
  (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
  const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,1> &vis,
  mav<complex<T>,2> &grid, const const_mav<T,1> &wgt, T w0=-1, T dw=-1)
  {
  x2grid_c(gconf, VisServ<T,const_mav<complex<T>,1>,const_mav<T,1>> (baselines, idx, vis, wgt), grid, w0, dw);
  }

template<typename T> void ms2grid_c
  (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
  const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &ms,
  mav<complex<T>,2> &grid, const const_mav<T,2> &wgt)
  {
  x2grid_c(gconf, MsServ<T,const_mav<complex<T>,2>,const_mav<T,2>> (baselines, idx, ms, wgt), grid);
  }

template<typename T, typename Serv> void grid2x_c
  (const GridderConfig<T> &gconf, const const_mav<complex<T>,2> &grid, const Serv &srv,
  T w0=-1, T dw=-1)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkShape(grid.shape(), {nu, nv});
  myassert(grid.contiguous(), "grid is not contiguous");
  size_t supp = gconf.Supp();
  size_t nthreads = gconf.Nthreads();
  bool do_w_gridding=dw>0;

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, grid.data(), nullptr, w0, dw);
  int jump = hlp.lineJump();
  const T * RESTRICT ku = hlp.kernel;
  const T * RESTRICT kv = hlp.kernel+supp;
  size_t np = srv.Nvis();

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<np; ++ipart)
    {
    UVW<T> coord = srv.getCoord(ipart);
    auto flip = coord.FixW();
    hlp.prep(coord);
    complex<T> r = 0;
    const auto * RESTRICT ptr = hlp.p0r;
    for (size_t cu=0; cu<supp; ++cu)
      {
      complex<T> tmp(0);
      size_t cv=0;
      for (; cv<supp-3; cv+=4)
        tmp += ptr[cv  ]*kv[cv  ]
             + ptr[cv+1]*kv[cv+1]
             + ptr[cv+2]*kv[cv+2]
             + ptr[cv+3]*kv[cv+3];
      for (; cv<supp; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += jump;
      }
    if (flip) r=conj(r);
    if (do_w_gridding) r*=hlp.Wfac();
    srv.setVis(ipart, r);
    }
}
  }
template<typename T> void grid2vis_c
  (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
  const const_mav<uint32_t,1> &idx, const const_mav<complex<T>,2> &grid,
  mav<complex<T>,1> &vis, const const_mav<T,1> &wgt, T w0=-1, T dw=-1)
  {
  grid2x_c(gconf, grid, VisServ<T,mav<complex<T>,1>,const_mav<T,1>> (baselines, idx, vis, wgt), w0, dw);
  }

template<typename T> void grid2ms_c
  (const Baselines<T> &baselines, const GridderConfig<T> &gconf,