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 "pocketfft_hdronly.h"
#include <array>
#define NOINLINE __attribute__((noinline))

namespace gridder {

namespace detail {

using namespace std;

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)
      { return operator()(i); }
    const 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()
      { return d; }
    const 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> 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_; };
    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=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> class EC_Kernel
  {
  protected:
    size_t supp;
    T beta, emb_;

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

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;
    using EC_Kernel<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))
      {
      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;

  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);
      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
      {
      size_t irow = index/nchan;
      size_t ichan = index-nchan*irow;
      return RowChan{irow, ichan};
      }

    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; }

    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_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});
      vector<complex<T>> tmpdat(nu*nv);
      auto tmp=mav<complex<T>,2>(tmpdat.data(),{nu,nv},{ptrdiff_t(nv),ptrdiff_t(1)});
      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_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});
      for (size_t i=0; i<nu; ++i)
        for (size_t j=0; j<nv; ++j)
          grid(i,j)=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);
      }

    inline 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;
      if (iu0+supp>nu+nsafe) iu0 = nu+nsafe-supp;
      v=fmodulo(v_in*psy, T(1))*nv;
      iv0 = int(v-supp*0.5 + 1 + nv) - nv;
      if (iv0+supp>nv+nsafe) 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;

    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;
    vector<T> kernel;

    Helper(const GridderConfig<T> &gconf_, const T2 *grid_r_, T2 *grid_w_)
      : 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)),
        kernel(2*supp)
      {}
    ~Helper() { if (grid_w) dump(); }

    int lineJump() const { return sv; }

    void prep(T u_in, T v_in)
      {
      T u, v;
      gconf.getpix(u_in, v_in, 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)
        {
        auto x = x0+i*xsupp;
        kernel[i  ] = beta*sqrt(T(1)-x*x);
        auto y = y0+i*xsupp;
        kernel[i+supp] = beta*sqrt(T(1)-y*y);
        }
      for (auto &k : kernel)
        k = exp(k);

      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});
      }
    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)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkShape(grid.shape(), {nu, nv});
  myassert(grid.contiguous(), "grid is not contiguous");
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();
  size_t nthreads = gconf.Nthreads();

#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, nullptr, grid.data());
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+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.u, coord.v);
    auto * ptr = hlp.p0w;
    auto v(srv.getVis(ipart)*emb);
    if (flip) v=conj(v);
    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
  }

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)
  {
  x2grid_c(gconf, VisServ<T,const_mav<complex<T>,1>,const_mav<T,1>> (baselines, idx, vis, wgt), grid);
  }

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)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkShape(grid.shape(), {nu, nv});
  myassert(grid.contiguous(), "grid is not contiguous");
  T beta = gconf.Beta();
  size_t supp = gconf.Supp();
  size_t nthreads = gconf.Nthreads();

  // Loop over sampling points
#pragma omp parallel num_threads(nthreads)
{
  Helper<T> hlp(gconf, grid.data(), nullptr);
  T emb = exp(-2*beta);
  int jump = hlp.lineJump();
  const T * ku = hlp.kernel.data();
  const T * kv = hlp.kernel.data()+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.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 (flip) r=conj(r);
    srv.setVis(ipart, r*emb);
    }
}
  }
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)
  {
  grid2x_c(gconf, grid, VisServ<T,mav<complex<T>,1>,const_mav<T,1>> (baselines, idx, vis, wgt));
  }

template<typename T> void grid2ms_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>,2> &ms, const const_mav<T,2> &wgt)
  {
  grid2x_c(gconf, grid, MsServ<T,mav<complex<T>,2>,const_mav<T,2>> (baselines, idx, ms, wgt));
  }

template<typename T> void apply_holo
  (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>,2> &ogrid, const const_mav<T,1> &wgt)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkShape(grid.shape(), {nu, nv});
  size_t nvis = idx.shape(0);
  size_t nthreads = gconf.Nthreads();
  bool have_wgt = wgt.size()!=0;
  if (have_wgt) checkShape(wgt.shape(), {nvis});
  checkShape(ogrid.shape(), {nu, nv});

  ogrid.fill(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.data(), ogrid.data());
  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));
    coord.FixW();
    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;
      }
    }
}
  }

template<typename T> void get_correlations
  (const Baselines<T> &baselines, const GridderConfig<T> &gconf,
  const const_mav<uint32_t,1> &idx, int du, int dv,
  mav<T,2> &ogrid, const const_mav<T,1> &wgt)
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  size_t nvis = idx.shape(0);
  bool have_wgt = wgt.size()!=0;
  if (have_wgt) checkShape(wgt.shape(), {nvis});
  checkShape(ogrid.shape(), {nu, 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");
  size_t nthreads = gconf.Nthreads();

  T beta = gconf.Beta();
  ogrid.fill(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.data());
  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));
    coord.FixW();
    hlp.prep(coord.u, coord.v);
    auto * wptr = hlp.p0w + u0*jump;
    auto f0 = emb*emb;
    if (have_wgt)
      {
      auto twgt = wgt(ipart);
      f0*=twgt*twgt;
      }
    for (size_t cu=u0; cu<u1; ++cu)
      {
      auto f1=ku[cu]*ku[cu+du]*f0;
      for (size_t cv=v0; cv<v1; ++cv)
        wptr[cv] += f1*kv[cv]*kv[cv+dv];
      wptr += jump;
      }
    }
}
  }


template<typename T, typename T2> void apply_wcorr(const GridderConfig<T> &gconf,
  const mav<T2,2> &dirty, const EC_Kernel_with_correction<T> &kernel, double dw)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  size_t nthreads = gconf.Nthreads();
  auto psx=gconf.Pixsize_x();
  auto psy=gconf.Pixsize_y();
  double x0 = -0.5*nx_dirty*psx,
         y0 = -0.5*ny_dirty*psy;
#pragma omp parallel for schedule(static) num_threads(nthreads)
  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;
      fy*=fy;
      T fct = 0;
      double tmp = 1.-fx-fy;
      if (tmp>=0)
        {
        auto nm1 = (-fx-fy)/(sqrt(1.-fx-fy)+1.); // accurate form of sqrt(1-x-y)
        fct = (nm1<=-1) ? 0. : kernel.corfac(nm1*dw);
        }
      size_t i2 = nx_dirty-i, j2 = ny_dirty-j;
      dirty(i,j)*=fct;
      if ((i>0)&&(i<i2))
        {
        dirty(i2,j)*=fct;
        if ((j>0)&&(j<j2))
          dirty(i2,j2)*=fct;
        }
      if ((j>0)&&(j<j2))
        dirty(i,j2)*=fct;
      }
    }
  }

template<typename T, typename Serv> void x2dirty(
  const GridderConfig<T> &gconf, const Serv &srv, const mav<T,2> &dirty, size_t verbosity=0)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto nu=gconf.Nu();
  auto nv=gconf.Nv();

  if (verbosity>0)
    cout << "Gridding without w-stacking: " << srv.Nvis()
         << " visibilities" << endl;
  if (verbosity>0) cout << "Using " << gconf.Nthreads() << " threads" << endl;

  tmpStorage<complex<T>,2> grid_({nu,nv});
  auto grid=grid_.getMav();
  grid_.fill(0.);
  x2grid_c(gconf, srv, grid);
  tmpStorage<complex<T>,2> cdirty_({nx_dirty,ny_dirty});
  auto cdirty=cdirty_.getMav();
  // FIXME: this could use symmetries to accelerate the FFT
  gconf.grid2dirty_c(grid, cdirty);
  for (size_t i=0; i<nx_dirty; ++i)
    for (size_t j=0; j<ny_dirty; ++j)
      dirty(i,j) = cdirty(i,j).real();
  }

template<typename T, typename Serv> void dirty2x(
  const GridderConfig<T> &gconf,  const const_mav<T,2> &dirty,
  const Serv &srv, size_t verbosity=0)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto nu=gconf.Nu();
  auto nv=gconf.Nv();
  size_t nvis = srv.Nvis();

  if (verbosity>0)
    cout << "Degridding without w-stacking: " << srv.Nvis()
         << " visibilities" << endl;
  if (verbosity>0) cout << "Using " << gconf.Nthreads() << " threads" << endl;

  for (size_t i=0; i<nvis; ++i)
    srv.setVis(i, 0.);

  tmpStorage<complex<T>,2> cdirty_({nx_dirty,ny_dirty});
  auto cdirty=cdirty_.getMav();
  for (size_t i=0; i<nx_dirty; ++i)
    for (size_t j=0; j<ny_dirty; ++j)
      cdirty(i,j) = dirty(i,j);

  tmpStorage<complex<T>,2> grid_({nu,nv});
  auto grid=grid_.getMav();
  // FIXME: this could use symmetries to accelerate the FFT
  gconf.dirty2grid_c(cdirty, grid);
  grid2x_c(gconf, const_mav<complex<T>,2>(grid.data(),{nu,nv}), srv);
  }

template<typename T, typename Serv> void wstack_common(
  const GridderConfig<T> &gconf, const Serv &srv, T &wmin, vector<T> &wval,
  double &dw, size_t &nplanes, vector<size_t> &nvis_plane, vector<int> &minplane, size_t verbosity)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto psx=gconf.Pixsize_x();
  auto psy=gconf.Pixsize_y();
  size_t nvis = srv.Nvis();
  size_t nthreads = gconf.Nthreads();

  // determine w values for every visibility, and min/max w;
  wmin=T(1e38);
  T wmax=T(-1e38);
  wval.resize(nvis);

#pragma omp parallel for num_threads(nthreads) reduction(min:wmin) reduction(max:wmax)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    wval[ipart] = abs(srv.getCoord(ipart).w);
    wmin = min(wmin,wval[ipart]);
    wmax = max(wmax,wval[ipart]);
    }
  if (verbosity>0) cout << "Using " << nthreads << " threads" << endl;
  if (verbosity>0) cout << "W range: " << wmin << " to " << wmax << endl;
  double x0 = -0.5*nx_dirty*psx,
         y0 = -0.5*ny_dirty*psy;
  double nmin = sqrt(max(1.-x0*x0-y0*y0,0.))-1.;
  dw = 0.25/abs(nmin);
  nplanes = max<size_t>(2,((wmax-wmin)/dw+2));
  dw=(wmax-wmin)/(nplanes-1);

  auto w_supp = gconf.Supp();
  wmin -= (0.5*w_supp-1)*dw;
  wmax += (0.5*w_supp-1)*dw;
  nplanes += w_supp-2;
  if (verbosity>0) cout << "Kernel support: " << w_supp << endl;
  if (verbosity>0) cout << "nplanes: " << nplanes << endl;
  nvis_plane.resize(nplanes,0);
  minplane.resize(nvis);
#pragma omp parallel num_threads(nthreads)
{
  vector<size_t> nvploc(nplanes,0);
#pragma omp for
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
    int plane0 = int((wval[ipart]-wmin)/dw+0.5*w_supp)-int(w_supp-1);
    minplane[ipart]=plane0;
    for (int i=max<int>(0,plane0); i<min<int>(nplanes,plane0+w_supp); ++i)
      ++nvploc[i];
    }
#pragma omp critical
  for (size_t i=0; i<nplanes; ++i)
    nvis_plane[i]+=nvploc[i];
}
  }

template<typename T, typename Serv> void x2dirty_wstack(
  const GridderConfig<T> &gconf, const Serv &srv, const mav<T,2> &dirty, size_t verbosity=0)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto nu=gconf.Nu();
  auto nv=gconf.Nv();
  size_t nvis = srv.Nvis();
  auto w_supp = gconf.Supp();
  size_t nthreads = gconf.Nthreads();
  EC_Kernel_with_correction<T> kernel(w_supp, nthreads);
  T wmin;
  vector<T> wval;
  double dw;
  size_t nplanes;
  vector<size_t> nvis_plane;
  vector<int> minplane;
  if (verbosity>0) cout << "Gridding using improved w-stacking" << endl;
  wstack_common(gconf, srv, wmin, wval, dw, nplanes, nvis_plane, minplane, verbosity);

  dirty.fill(0);
  vector<complex<T>> vis_loc_;
  vector<uint32_t> idx_loc_;
  vector<uint32_t> mapper;
  tmpStorage<complex<T>,2> grid_({nu,nv});
  auto grid=grid_.getMav();
  tmpStorage<complex<T>,2> tdirty_({nx_dirty,ny_dirty});
  auto tdirty=tdirty_.getMav();
  for (size_t iw=0; iw<nplanes; ++iw)
    {
    if (nvis_plane[iw]==0) continue;
    if (verbosity>1)
      cout << "Working on plane " << iw << " containing " << nvis_plane[iw]
           << " visibilities" << endl;
    T wcur = wmin+iw*dw;
    size_t cnt=0;
    vis_loc_.resize(nvis_plane[iw]);
    auto vis_loc = mav<complex<T>,1>(vis_loc_.data(), {nvis_plane[iw]});
    idx_loc_.resize(nvis_plane[iw]);
    auto idx_loc = mav<uint32_t,1>(idx_loc_.data(), {nvis_plane[iw]});
    mapper.resize(nvis_plane[iw]);
    for (size_t ipart=0; ipart<nvis; ++ipart)
      if ((int(iw)>=minplane[ipart]) && (iw<minplane[ipart]+w_supp))
        mapper[cnt++] = ipart;

    T tmpval=T(2)/(w_supp*dw);
#pragma omp parallel for num_threads(nthreads)
    for (size_t i=0; i<nvis_plane[iw]; ++i)
      {
      auto ipart = mapper[i];
      T x=min(T(1), tmpval*abs(wcur-wval[ipart]));
      vis_loc[i] = srv.getVis(ipart)*kernel(x);
      idx_loc[i] = srv.getIdx(ipart);
      }
    grid.fill(0);
    vis2grid_c(srv.getBaselines(), gconf, const_mav<uint32_t,1>(idx_loc), const_mav<complex<T>,1>(vis_loc), grid, const_mav<T,1>(nullptr,{0}));
    gconf.grid2dirty_c(grid, tdirty);
    gconf.apply_wscreen(tdirty, tdirty, wcur, false);
#pragma omp parallel for num_threads(nthreads)
    for (size_t i=0; i<nx_dirty; ++i)
      for (size_t j=0; j<ny_dirty; ++j)
        dirty(i,j) += tdirty(i,j).real();
    }
  // correct for w gridding
  apply_wcorr(gconf, dirty, kernel, dw);
  }
template<typename T> void vis2dirty_wstack(const Baselines<T> &baselines,
  const GridderConfig<T> &gconf, const const_mav<uint32_t,1> &idx,
  const const_mav<complex<T>,1> &vis, mav<T,2> &dirty)
  {
  x2dirty_wstack(gconf, VisServ<T,const_mav<complex<T>,1>,const_mav<T,1>>(baselines, idx, vis, const_mav<T,1>(nullptr,{0})), dirty);
  }

template<typename T, typename Serv> void dirty2x_wstack(
  const GridderConfig<T> &gconf,  const const_mav<T,2> &dirty,
  const Serv &srv, size_t verbosity=0)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto nu=gconf.Nu();
  auto nv=gconf.Nv();
  size_t nvis = srv.Nvis();
  auto w_supp = gconf.Supp();
  size_t nthreads = gconf.Nthreads();
  EC_Kernel_with_correction<T> kernel(w_supp, nthreads);
  T wmin;
  vector<T> wval;
  double dw;
  size_t nplanes;
  vector<size_t> nvis_plane;
  vector<int> minplane;
  if (verbosity>0) cout << "Degridding using improved w-stacking" << endl;
  wstack_common(gconf, srv, wmin, wval, dw, nplanes, nvis_plane, minplane, verbosity);

  for (size_t i=0; i<nvis; ++i)
    srv.setVis(i, 0.);