nifty_gridder.cc

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

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

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

#include "gridder_cxx.h"

using namespace std;
using namespace gridder;

namespace py = pybind11;

namespace {

auto None = py::none();

template<typename T>
  using pyarr = py::array_t<T, 0>;

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

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

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

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

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

template<size_t ndim, typename T> mav<T,ndim> make_mav(pyarr<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return mav<T, ndim>(in.mutable_data(),dims,str);
  }
template<size_t ndim, typename T> const_mav<T,ndim> make_const_mav(const pyarr<T> &in)
  {
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
    {
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
    }
  return const_mav<T, ndim>(in.data(),dims,str);
  }

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    pyarr<T> apply_taper(const pyarr<T> &img, bool divide) const
      {
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto img2 = make_const_mav<2>(img);
      auto res2 = make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::apply_taper(img2, res2, divide);
      }
      return res;
      }
    pyarr<T> grid2dirty(const pyarr<T> &grid) const
      {
      auto res = makeArray<T>({nx_dirty, ny_dirty});
      auto grid2=make_const_mav<2>(grid);
      auto res2=make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::grid2dirty(grid2,res2);
      }
      return res;
      }
    pyarr<complex<T>> grid2dirty_c(const pyarr<complex<T>> &grid) const
      {
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto grid2=make_const_mav<2>(grid);
      auto res2=make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::grid2dirty_c(grid2,res2);
      }
      return res;
      }

    pyarr<T> dirty2grid(const pyarr<T> &dirty) const
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto grid = makeArray<T>({nu, nv});
      auto grid2=make_mav<2>(grid);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::dirty2grid(dirty2, grid2);
      }
      return grid;
      }
    pyarr<complex<T>> dirty2grid_c(const pyarr<complex<T>> &dirty) const
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto grid = makeArray<complex<T>>({nu, nv});
      auto grid2=make_mav<2>(grid);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::dirty2grid_c(dirty2, grid2);
      }
      return grid;
      }
    pyarr<complex<T>> apply_wscreen(const pyarr<complex<T>> &dirty, double w, bool adjoint) const
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto res2 = make_mav<2>(res);
      {
      py::gil_scoped_release release;
      GridderConfig<T>::apply_wscreen(dirty2, res2, w, adjoint);
      }
      return res;
      }
  };


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

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

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
template<typename T> pyarr<complex<T>> Pyvis2grid_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto vis2 = make_const_mav<1>(vis_);
  size_t nvis = vis2.shape(0);
  auto idx2 = make_const_mav<1>(idx_);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nvis});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = provideCArray<complex<T>>(grid_in, {gconf.Nu(), gconf.Nv()});
  auto grid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  vis2grid_c<T>(baselines, gconf, idx2, vis2, grid, wgt2);
  }
  return res;
  }

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

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

Returns
=======
np.array((nu,nv), dtype=np.float64):
    the gridded visibilities (made real by making use of Hermitian symmetry)
)""";
template<typename T> pyarr<T> Pyvis2grid(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
  {
  auto tmp=Pyvis2grid_c(baselines, gconf, idx_, vis_, None, wgt_);
  auto grd=provideCArray<T>(grid_in,{gconf.Nu(), gconf.Nv()});
  gridder::detail::complex2hartley(make_const_mav<2>(tmp), make_mav<2>(grd), gconf.Nthreads());
  return grd;
  }

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

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

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
template<typename T> pyarr<complex<T>> Pyms2grid_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  auto ms2 = make_const_mav<2>(ms_);
  auto idx2 = make_const_mav<1>(idx_);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nrows, nchan});
  auto wgt2 = make_const_mav<2>(wgt);
  auto res = provideCArray<complex<T>>(grid_in, {gconf.Nu(), gconf.Nv()});
  auto grid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  ms2grid_c<T>(baselines, gconf, idx2, ms2, grid, wgt2);
  }
  return res;
  }

template<typename T> pyarr<T> Pyms2grid(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_,
  py::object &grid_in, const py::object &wgt_)
  {
  auto tmp = Pyms2grid_c(baselines, gconf, idx_, ms_, None, wgt_);
  auto res_ = provideCArray<T>(grid_in, {gconf.Nu(), gconf.Nv()});
  auto res = make_mav<2>(res_);
  gridder::detail::complex2hartley(make_const_mav<2>(tmp), res, gconf.Nthreads());
  return res_;
  }

template<typename T> pyarr<complex<T>> Pygrid2vis_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
  const py::object &wgt_)
  {
  auto grid2 = make_const_mav<2>(grid_);
  auto idx2 = make_const_mav<1>(idx_);
  size_t nvis = idx2.shape(0);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nvis});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = makeArray<complex<T>>({nvis});
  auto vis = make_mav<1>(res);
  {
  py::gil_scoped_release release;
  grid2vis_c<T>(baselines, gconf, idx2, grid2, vis, wgt2);
  }
  return res;
  }

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

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

Returns
=======
np.array((nvis,), dtype=np.complex)
    The degridded visibility data
)""";
template<typename T> pyarr<complex<T>> Pygrid2vis(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<T> &grid_, const py::object &wgt_)
  {
  auto tmp=makeArray<complex<T>>({gconf.Nu(), gconf.Nv()});
  gridder::detail::hartley2complex(make_const_mav<2>(grid_),make_mav<2>(tmp), gconf.Nthreads());
  return Pygrid2vis_c(baselines, gconf, idx_, tmp, wgt_);
  }

template<typename T> pyarr<complex<T>> Pygrid2ms_c(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
  py::object &ms_in, const py::object &wgt_)
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  auto grid2 = make_const_mav<2>(grid_);
  auto idx2 = make_const_mav<1>(idx_);
  pyarr<T> wgt = providePotentialArray<T>(wgt_, {nrows, nchan});
  auto wgt2 = make_const_mav<2>(wgt);
  auto res = provideCArray<complex<T>>(ms_in, {nrows, nchan});
  auto ms = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  grid2ms_c<T>(baselines, gconf, idx2, grid2, ms, wgt2);
  }
  return res;
  }

template<typename T> pyarr<complex<T>> Pygrid2ms(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<T> &grid_, py::object &ms_in, const py::object &wgt_)
  {
  auto grid2_ = makeArray<complex<T>>({gconf.Nu(), gconf.Nv()});
  auto grid2 = make_mav<2>(grid2_);
  gridder::detail::hartley2complex(make_const_mav<2>(grid_), grid2, gconf.Nthreads());
  return Pygrid2ms_c(baselines, gconf, idx_, grid2_, ms_in, wgt_);
  }

template<typename T> pyarr<complex<T>> Pyapply_holo(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &grid_,
  const py::object &wgt_)
  {
  auto idx = make_const_mav<1>(idx_);
  auto grid = make_const_mav<2>(grid_);
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {idx.shape(0)});
  auto wgt=make_const_mav<1>(wgt2);

  auto res = makeArray<complex<T>>({grid.shape(0),grid.shape(1)});
  auto ogrid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  apply_holo(baselines, gconf, idx, grid, ogrid, wgt);
  }
  return res;
  }

template<typename T> pyarr<T> Pyget_correlations(
  const PyBaselines<T> &baselines, const PyGridderConfig<T> &gconf,
  const pyarr<uint32_t> &idx_, int du, int dv, const py::object &wgt_)
  {
  auto idx = make_const_mav<1>(idx_);
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, {idx.shape(0)});
  auto wgt=make_const_mav<1>(wgt2);

  auto res = makeArray<T>({gconf.Nu(),gconf.Nv()});
  auto ogrid = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  get_correlations(baselines, gconf, idx, du, dv, ogrid, wgt);
  }
  return res;
  }

constexpr auto getIndices_DS = R"""(
Selects a subset of entries from a `Baselines` object.

Parameters
==========
baselines: Baselines
    the Baselines object
gconf: GridderConf
    the GridderConf object to be used with the returned indices.
    (used to optimize the ordering of the indices)
flags: np.array((nrows, nchannels), dtype=np.bool)
    "True" indicates that the value should not be used
chbegin: int
    first channel to use (-1: start with the first available channel)
chend: int
    one-past last channel to use (-1: one past the last available channel)
wmin: float
    only select entries with w>=wmin
wmax: float
    only select entries with w<wmax

Returns
=======
np.array((nvis,), dtype=np.uint32)
    the compressed indices for all entries which match the selected criteria
    and are not flagged.
)""";
template<typename T> pyarr<uint32_t> PygetIndices(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<bool> &flags_, int chbegin,
  int chend, T wmin, T wmax)
  {
  size_t nidx;
  auto flags = make_const_mav<2>(flags_);
  {
  py::gil_scoped_release release;
  nidx = getIdxSize(baselines, flags, chbegin, chend, wmin, wmax, gconf.Nthreads());
  }
  auto res = makeArray<uint32_t>({nidx});
  auto res2 = make_mav<1>(res);
  {
  py::gil_scoped_release release;
  fillIdx(baselines, gconf, flags, chbegin, chend, wmin, wmax, res2);
  }
  return res;
  }

template<typename T> pyarr<T> Pyvis2dirty_wstack(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &vis_)
  {
  auto nx_dirty=gconf.Nxdirty();
  auto ny_dirty=gconf.Nydirty();
  auto idx2=make_const_mav<1>(idx_);
  auto vis2=make_const_mav<1>(vis_);
  auto dirty = makeArray<T>({nx_dirty, ny_dirty});
  auto dirty2=make_mav<2>(dirty);
  {
  py::gil_scoped_release release;
  vis2dirty_wstack<T>(baselines, gconf, idx2, vis2, dirty2);
  }
  return dirty;
  }

template<typename T> pyarr<complex<T>> Pydirty2vis_wstack(const PyBaselines<T> &baselines,
  const PyGridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<T> &dirty_)
  {
  auto idx2=make_const_mav<1>(idx_);
  auto nvis = idx2.shape(0);
  auto vis = makeArray<complex<T>>({nvis});
  auto vis2=make_mav<1>(vis);
  auto dirty2=make_const_mav<2>(dirty_);
  {
  py::gil_scoped_release release;
  vis2.fill(0);
  dirty2vis_wstack<T>(baselines, gconf, idx2, dirty2, vis2);
  }
  return vis;
  }

template<typename T> pyarr<T> Pygridding(const pyarr<T> &uvw,
  const pyarr<T> &freq, const pyarr<complex<T>> &ms, const py::object &wgt_,
  size_t npix_x, size_t npix_y, T pixsize_x, T pixsize_y, double epsilon,
  size_t nthreads, size_t verbosity)
  {
  auto uvw2 = make_const_mav<2>(uvw);
  auto freq2 = make_const_mav<1>(freq);
  auto ms2 = make_const_mav<2>(ms);
  auto wgt=providePotentialArray<T>(wgt_,{ms2.shape(0),ms2.shape(1)});
  auto wgt2 = make_const_mav<2>(wgt);
  auto dirty=makeArray<T>({npix_x,npix_y});
  auto dirty2 =make_mav<2>(dirty);
  gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,dirty2,verbosity);
  return dirty;
  }

template<typename T> pyarr<complex<T>> Pydegridding(const pyarr<T> &uvw,
  const pyarr<T> &freq, const pyarr<T> &dirty, const py::object &wgt_,
  T pixsize_x, T pixsize_y, double epsilon, size_t nthreads, size_t verbosity)
  {
  auto uvw2 = make_const_mav<2>(uvw);
  auto freq2 = make_const_mav<1>(freq);
  auto dirty2 = make_const_mav<2>(dirty);
  auto wgt=providePotentialArray<T>(wgt_,{uvw2.shape(0),freq2.shape(0)});
  auto wgt2 = make_const_mav<2>(wgt);
  auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
  auto ms2 =make_mav<2>(ms);
  degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,ms2,verbosity);
  return ms;
  }

template<typename T> pyarr<T> Pyfull_gridding(const pyarr<T> &uvw,
  const pyarr<T> &freq, const pyarr<complex<T>> &ms, const py::object &wgt_,
  size_t npix_x, size_t npix_y, T pixsize_x, T pixsize_y, double epsilon,
  size_t nthreads, size_t verbosity)
  {
  auto uvw2 = make_const_mav<2>(uvw);
  auto freq2 = make_const_mav<1>(freq);
  auto ms2 = make_const_mav<2>(ms);
  auto wgt=providePotentialArray<T>(wgt_,{ms2.shape(0),ms2.shape(1)});
  auto wgt2 = make_const_mav<2>(wgt);
  auto dirty=makeArray<T>({npix_x,npix_y});
  auto dirty2 =make_mav<2>(dirty);
  full_gridding(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,dirty2,verbosity);
  return dirty;
  }

template<typename T> pyarr<complex<T>> Pyfull_degridding(const pyarr<T> &uvw,
  const pyarr<T> &freq, const pyarr<T> &dirty, const py::object &wgt_,
  T pixsize_x, T pixsize_y, double epsilon, size_t nthreads, size_t verbosity)
  {
  auto uvw2 = make_const_mav<2>(uvw);
  auto freq2 = make_const_mav<1>(freq);
  auto dirty2 = make_const_mav<2>(dirty);
  auto wgt=providePotentialArray<T>(wgt_,{uvw2.shape(0),freq2.shape(0)});
  auto wgt2 = make_const_mav<2>(wgt);
  auto ms=makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
  auto ms2 =make_mav<2>(ms);
  full_degridding(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,epsilon,nthreads,ms2,verbosity);
  return ms;
  }

} // unnamed namespace

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

  py::class_<PyBaselines<double>> (m, "Baselines", PyBaselines_DS)
    .def(py::init<const pyarr<double> &, const pyarr<double> &>(),
      "coord"_a, "freq"_a)
    .def ("Nrows",&PyBaselines<double>::Nrows)
    .def ("Nchannels",&PyBaselines<double>::Nchannels)
    .def ("ms2vis",&PyBaselines<double>::ms2vis<complex<double>>,
      PyBaselines<double>::ms2vis_DS, "ms"_a, "idx"_a, "nthreads"_a=1)
    .def ("effectiveuvw",&PyBaselines<double>::effectiveuvw, "idx"_a)
    .def ("vis2ms",&PyBaselines<double>::vis2ms<complex<double>>,
      PyBaselines<double>::vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=None, "nthreads"_a=1);
  py::class_<PyGridderConfig<double>> (m, "GridderConfig", GridderConfig_DS)
    .def(py::init<size_t, size_t, double, double, double, size_t>(),"nxdirty"_a,
      "nydirty"_a, "epsilon"_a, "pixsize_x"_a, "pixsize_y"_a, "nthreads"_a=1)
    .def("Nxdirty", &PyGridderConfig<double>::Nxdirty)
    .def("Nydirty", &PyGridderConfig<double>::Nydirty)
    .def("Epsilon", &PyGridderConfig<double>::Epsilon)
    .def("Pixsize_x", &PyGridderConfig<double>::Pixsize_x)
    .def("Pixsize_y", &PyGridderConfig<double>::Pixsize_y)
    .def("Nu", &PyGridderConfig<double>::Nu)
    .def("Nv", &PyGridderConfig<double>::Nv)
    .def("Supp", &PyGridderConfig<double>::Supp)
    .def("apply_taper", &PyGridderConfig<double>::apply_taper, apply_taper_DS,
      "img"_a, "divide"_a=false)
     .def("grid2dirty", &PyGridderConfig<double>::grid2dirty,
        grid2dirty_DS, "grid"_a)
    .def("grid2dirty_c", &PyGridderConfig<double>::grid2dirty_c, "grid"_a)
    .def("dirty2grid", &PyGridderConfig<double>::dirty2grid,
       dirty2grid_DS, "dirty"_a)
    .def("dirty2grid_c", &PyGridderConfig<double>::dirty2grid_c, "dirty"_a)
    .def("apply_wscreen", &PyGridderConfig<double>::apply_wscreen,
      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a)

    // pickle support
    .def(py::pickle(
        // __getstate__
        [](const PyGridderConfig<double> & gc) {
          // Encode object state in tuple
          return py::make_tuple(gc.Nxdirty(), gc.Nydirty(), gc.Epsilon(),
                                gc.Pixsize_x(), gc.Pixsize_y(), gc.Nthreads());
        },
        // __setstate__
        [](py::tuple t) {
          myassert(t.size()==6,"Invalid state");

          // Reconstruct from tuple
          return PyGridderConfig<double>(t[0].cast<size_t>(), t[1].cast<size_t>(),
                                       t[2].cast<double>(), t[3].cast<double>(),
                                       t[4].cast<double>(), t[5].cast<size_t>());

        }));
  m.def("getIndices", PygetIndices<double>, getIndices_DS, "baselines"_a,
    "gconf"_a, "flags"_a, "chbegin"_a=-1, "chend"_a=-1,
    "wmin"_a=-1e30, "wmax"_a=1e30);
   m.def("vis2grid",&Pyvis2grid<double>, vis2grid_DS, "baselines"_a, "gconf"_a,
     "idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
  m.def("ms2grid",&Pyms2grid<double>, "baselines"_a, "gconf"_a, "idx"_a, "ms"_a,
    "grid_in"_a=None, "wgt"_a=None);
  m.def("grid2vis",&Pygrid2vis<double>, grid2vis_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "grid"_a, "wgt"_a=None);
  m.def("grid2ms",&Pygrid2ms<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "ms_in"_a=None, "wgt"_a=None);
  m.def("vis2grid_c",&Pyvis2grid_c<double>, vis2grid_c_DS, "baselines"_a,
    "gconf"_a, "idx"_a, "vis"_a, "grid_in"_a=None, "wgt"_a=None);
  m.def("ms2grid_c",&Pyms2grid_c<double>, ms2grid_c_DS, "baselines"_a, "gconf"_a,
    "idx"_a, "ms"_a, "grid_in"_a=None, "wgt"_a=None);
  m.def("grid2vis_c",&Pygrid2vis_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "wgt"_a=None);
  m.def("grid2ms_c",&Pygrid2ms_c<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "ms_in"_a=None, "wgt"_a=None);
  m.def("apply_holo",&Pyapply_holo<double>, "baselines"_a, "gconf"_a, "idx"_a,
    "grid"_a, "wgt"_a=None);
  m.def("get_correlations", &Pyget_correlations<double>, "baselines"_a, "gconf"_a,
    "idx"_a, "du"_a, "dv"_a, "wgt"_a=None);
  m.def("vis2dirty_wstack",&Pyvis2dirty_wstack<double>, "baselines"_a, "gconf"_a,
    "idx"_a, "vis"_a);
  m.def("dirty2vis_wstack",&Pydirty2vis_wstack<double>, "baselines"_a, "gconf"_a,
    "idx"_a, "dirty"_a);

  m.def("full_gridding",&Pyfull_gridding<double>,"uvw"_a,"freq"_a,"ms"_a,
    "wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
    "nthreads"_a=1, "verbosity"_a=0);
  m.def("full_gridding_f",&Pyfull_gridding<float>,"uvw"_a,"freq"_a,"ms"_a,
    "wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
    "nthreads"_a=1, "verbosity"_a=0);
  m.def("full_degridding",&Pyfull_degridding<double>,"uvw"_a,"freq"_a,"dirty"_a,
    "wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a,"nthreads"_a=1,
    "verbosity"_a=0);
  m.def("full_degridding_f",&Pyfull_degridding<float>,"uvw"_a,"freq"_a,"dirty"_a,
    "wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "nthreads"_a=1,
    "verbosity"_a=0);
  m.def("gridding",&Pygridding<double>,"uvw"_a,"freq"_a,"ms"_a,
    "wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
    "nthreads"_a=1, "verbosity"_a=0);
  m.def("gridding_f",&Pygridding<float>,"uvw"_a,"freq"_a,"ms"_a,
    "wgt"_a=None,"npix_x"_a,"npix_y"_a,"pixsize_x"_a,"pixsize_y"_a,"epsilon"_a,
    "nthreads"_a=1, "verbosity"_a=0);
  m.def("degridding",&Pydegridding<double>,"uvw"_a,"freq"_a,"dirty"_a,
    "wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a,"nthreads"_a=1,
    "verbosity"_a=0);
  m.def("degridding_f",&Pydegridding<float>,"uvw"_a,"freq"_a,"dirty"_a,
    "wgt"_a=None,"pixsize_x"_a,"pixsize_y"_a, "epsilon"_a, "nthreads"_a=1,
    "verbosity"_a=0);
  }