nifty_gridder.cc 39.8 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1 2 3 4 5 6 7 8 9 10 11 12 13 14
/*
 *  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
Martin Reinecke's avatar
Martin Reinecke committed
15
 *  along with nifty_gridder; if not, write to the Free Software
Martin Reinecke's avatar
Martin Reinecke committed
16 17 18
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

Martin Reinecke's avatar
Martin Reinecke committed
19 20 21
/* Copyright (C) 2019 Max-Planck-Society
   Author: Martin Reinecke */

Martin Reinecke's avatar
import  
Martin Reinecke committed
22 23
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
Martin Reinecke's avatar
Martin Reinecke committed
24

Martin Reinecke's avatar
Martin Reinecke committed
25
#include "gridder_cxx.h"
Martin Reinecke's avatar
import  
Martin Reinecke committed
26 27

using namespace std;
Martin Reinecke's avatar
Martin Reinecke committed
28
using namespace gridder;
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
29
using gridder::detail::idx_t;
Martin Reinecke's avatar
import  
Martin Reinecke committed
30 31 32 33
namespace py = pybind11;

namespace {

Martin Reinecke's avatar
Martin Reinecke committed
34 35
auto None = py::none();

36
template<typename T>
Martin Reinecke's avatar
updates  
Martin Reinecke committed
37
  using pyarr = py::array_t<T, 0>;
Martin Reinecke's avatar
import  
Martin Reinecke committed
38

39 40 41 42 43 44 45 46 47 48
template<typename T> bool isPytype(const py::array &arr)
  {
  auto t1=arr.dtype();
  auto t2=pybind11::dtype::of<T>();
  auto k1=t1.kind();
  auto k2=t2.kind();
  auto s1=t1.itemsize();
  auto s2=t2.itemsize();
  return (k1==k2)&&(s1==s2);
  }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
49 50 51 52
template<typename T> pyarr<T> getPyarr(const py::array &arr, const string &name)
  {
  auto t1=arr.dtype();
  auto t2=pybind11::dtype::of<T>();
53 54 55 56 57 58 59
  auto k1=t1.kind();
  auto k2=t2.kind();
  auto s1=t1.itemsize();
  auto s2=t2.itemsize();
  myassert((k1==k2)&&(s1==s2),
    "type mismatch for array '", name, "': expected '", k2, s2,
    "', but got '", k1, s1, "'.");
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
60 61 62
  return arr.cast<pyarr<T>>();
  }

Martin Reinecke's avatar
Martin Reinecke committed
63 64
template<typename T> pyarr<T> makeArray(const vector<size_t> &shape)
  { return pyarr<T>(shape); }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
65

Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
66
void checkArray(const py::array &arr, const string &aname,
Martin Reinecke's avatar
merge  
Martin Reinecke committed
67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83
  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");
      }
  }

Martin Reinecke's avatar
Martin Reinecke committed
84
template<typename T> pyarr<T> provideArray(const py::object &in,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
85
  const string &name, const vector<size_t> &shape)
Martin Reinecke's avatar
merge  
Martin Reinecke committed
86
  {
87
  if (in.is_none())
Martin Reinecke's avatar
merge  
Martin Reinecke committed
88 89 90 91 92 93 94 95
    {
    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_;
    }
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
96 97
  auto tmp_ = getPyarr<T>(in.cast<py::array>(), name);
  checkArray(tmp_, name, shape);
Martin Reinecke's avatar
merge  
Martin Reinecke committed
98 99 100
  return tmp_;
  }

Martin Reinecke's avatar
Martin Reinecke committed
101
template<typename T> pyarr<T> providePotentialArray(const py::object &in,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
102
  const string &name, const vector<size_t> &shape)
Martin Reinecke's avatar
Martin Reinecke committed
103 104 105
  {
  if (in.is_none())
    return makeArray<T>(vector<size_t>(shape.size(), 0));
Martin Reinecke's avatar
Martin Reinecke committed
106
  return getPyarr<T>(in.cast<py::array>(), name);
Martin Reinecke's avatar
merge  
Martin Reinecke committed
107 108
  }

Martin Reinecke's avatar
Martin Reinecke committed
109
template<size_t ndim, typename T> mav<T,ndim> make_mav(pyarr<T> &in)
Martin Reinecke's avatar
Martin Reinecke committed
110
  {
Martin Reinecke's avatar
Martin Reinecke committed
111 112 113 114
  myassert(ndim==in.ndim(), "dimension mismatch");
  array<size_t,ndim> dims;
  array<ptrdiff_t,ndim> str;
  for (size_t i=0; i<ndim; ++i)
Martin Reinecke's avatar
Martin Reinecke committed
115
    {
Martin Reinecke's avatar
Martin Reinecke committed
116 117 118
    dims[i]=in.shape(i);
    str[i]=in.strides(i)/sizeof(T);
    myassert(str[i]*ptrdiff_t(sizeof(T))==in.strides(i), "weird strides");
Martin Reinecke's avatar
Martin Reinecke committed
119
    }
Martin Reinecke's avatar
Martin Reinecke committed
120
  return mav<T, ndim>(in.mutable_data(),dims,str);
Martin Reinecke's avatar
Martin Reinecke committed
121
  }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
122
template<size_t ndim, typename T> const_mav<T,ndim> make_const_mav(const pyarr<T> &in)
Martin Reinecke's avatar
Martin Reinecke committed
123 124 125 126 127 128 129 130 131 132
  {
  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");
    }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
133
  return const_mav<T, ndim>(in.data(),dims,str);
Martin Reinecke's avatar
Martin Reinecke committed
134
  }
Martin Reinecke's avatar
Martin Reinecke committed
135

Martin Reinecke's avatar
Martin Reinecke committed
136
constexpr auto PyBaselines_DS = R"""(
137 138 139 140
Class storing UVW coordinates and channel information.

Parameters
==========
Martin Reinecke's avatar
Martin Reinecke committed
141
coord: np.array((nrows, 3), dtype=np.float64)
142
    u, v and w coordinates for each row
Martin Reinecke's avatar
Martin Reinecke committed
143
freq: np.array((nchannels,), dtype=np.float64)
144 145
    frequency for each individual channel (in Hz)
)""";
146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180

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

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

Martin Reinecke's avatar
step 1  
Martin Reinecke committed
181
class PyBaselines: public Baselines
Martin Reinecke's avatar
Martin Reinecke committed
182 183
  {
  public:
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
184 185 186
    using Baselines::Baselines;
    template<typename T> PyBaselines(const pyarr<T> &coord, const pyarr<T> &freq)
      : Baselines(make_const_mav<2>(coord), make_const_mav<1>(freq))
Martin Reinecke's avatar
Martin Reinecke committed
187
      {}
Martin Reinecke's avatar
Martin Reinecke committed
188

Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
189
    template<typename T> pyarr<T> effectiveuvw(const pyarr<idx_t> &idx_) const
Martin Reinecke's avatar
Martin Reinecke committed
190
      {
191
      size_t nvis = size_t(idx_.shape(0));
Martin Reinecke's avatar
Martin Reinecke committed
192
      auto idx=make_const_mav<1>(idx_);
193
      auto res_=makeArray<T>({nvis, 3});
Martin Reinecke's avatar
Martin Reinecke committed
194
      auto res=make_mav<2>(res_);
Martin Reinecke's avatar
Martin Reinecke committed
195 196
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
197
      Baselines::effectiveUVW(idx,res);
Martin Reinecke's avatar
Martin Reinecke committed
198
      }
199
      return res_;
Martin Reinecke's avatar
Martin Reinecke committed
200
      }
201

Martin Reinecke's avatar
step 1  
Martin Reinecke committed
202
    template<typename T> pyarr<T> ms2vis(const pyarr<T> &ms_,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
203
      const pyarr<idx_t> &idx_, size_t nthreads) const
Martin Reinecke's avatar
updates  
Martin Reinecke committed
204
      {
Martin Reinecke's avatar
Martin Reinecke committed
205
      auto idx=make_const_mav<1>(idx_);
Martin Reinecke's avatar
Martin Reinecke committed
206
      size_t nvis = size_t(idx.shape(0));
Martin Reinecke's avatar
Martin Reinecke committed
207
      auto ms = make_const_mav<2>(ms_);
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
208
      auto res=makeArray<T>({nvis});
Martin Reinecke's avatar
Martin Reinecke committed
209
      auto vis = make_mav<1>(res);
Martin Reinecke's avatar
Martin Reinecke committed
210 211
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
212
      Baselines::ms2vis(ms, idx, vis, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
213
      }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
214 215 216
      return res;
      }

Martin Reinecke's avatar
step 1  
Martin Reinecke committed
217
    template<typename T> pyarr<T> vis2ms(const pyarr<T> &vis_,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
218
      const pyarr<idx_t> &idx_, py::object &ms_in, size_t nthreads) const
Martin Reinecke's avatar
updates  
Martin Reinecke committed
219
      {
Martin Reinecke's avatar
Martin Reinecke committed
220 221
      auto vis=make_const_mav<1>(vis_);
      auto idx=make_const_mav<1>(idx_);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
222
      auto res = provideArray<T>(ms_in, "ms_in", {nrows, nchan});
Martin Reinecke's avatar
Martin Reinecke committed
223
      auto ms = make_mav<2>(res);
Martin Reinecke's avatar
Martin Reinecke committed
224 225
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
226
      Baselines::vis2ms(vis, idx, ms, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
227
      }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
228 229
      return res;
      }
Martin Reinecke's avatar
Martin Reinecke committed
230 231
  };

232 233 234 235 236
constexpr auto grid2dirty_DS = R"""(
Converts from UV grid to dirty image (FFT, cropping, correction)

Parameters
==========
Martin Reinecke's avatar
Martin Reinecke committed
237
grid: np.array((nu, nv), dtype=np.float32 or np.float64)
238 239 240 241
    gridded UV data

Returns
=======
Martin Reinecke's avatar
Martin Reinecke committed
242
nd.array((nxdirty, nydirty), same dtype as `grid`)
243 244 245 246 247 248 249 250
    the dirty image
)""";

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

Parameters
==========
Martin Reinecke's avatar
Martin Reinecke committed
251
dirty: nd.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
252 253 254 255
    the dirty image

Returns
=======
Martin Reinecke's avatar
Martin Reinecke committed
256
np.array((nu, nv), same dtype as `dirty`)
257 258 259
    gridded UV data
)""";

Martin Reinecke's avatar
Martin Reinecke committed
260 261 262 263 264
constexpr auto apply_taper_DS = R"""(
Applies the taper (or its inverse) to an image

Parameters
==========
Martin Reinecke's avatar
Martin Reinecke committed
265
img: nd.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
Martin Reinecke's avatar
Martin Reinecke committed
266 267 268 269 270 271
    the image
divide: bool
    if True, the routine dividex by the taper, otherwise it multiplies by it

Returns
=======
Martin Reinecke's avatar
Martin Reinecke committed
272
np.array((nxdirty, nydirty), same dtype as `img`)
Martin Reinecke's avatar
Martin Reinecke committed
273 274 275
    the image with the taper applied
)""";

Martin Reinecke's avatar
Martin Reinecke committed
276 277 278 279 280
constexpr auto apply_wscreen_DS = R"""(
Applies the w screen to an image

Parameters
==========
Martin Reinecke's avatar
Martin Reinecke committed
281
dirty: nd.array((nxdirty, nydirty), dtype=np.complex64 or np.complex128)
Martin Reinecke's avatar
Martin Reinecke committed
282 283 284 285 286
    the image
w : float
    the w value to use
adjoint: bool
    if True, apply the complex conjugate of the w screen
Martin Reinecke's avatar
typo  
Martin Reinecke committed
287
divide_by_n: bool, default=True
Martin Reinecke's avatar
Martin Reinecke committed
288
    if True, divide ny n.
Martin Reinecke's avatar
Martin Reinecke committed
289 290 291

Returns
=======
Martin Reinecke's avatar
Martin Reinecke committed
292
np.array((nxdirty, nydirty), same dtype as 'dirty')
Martin Reinecke's avatar
Martin Reinecke committed
293 294 295
    the image with the w screen applied
)""";

296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311
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)
Martin Reinecke's avatar
Martin Reinecke committed
312 313
nthreads: int
    the number of threads to use for all calculations involving this object.
314
)""";
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
315
class PyGridderConfig: public GridderConfig
Martin Reinecke's avatar
Martin Reinecke committed
316 317
  {
  public:
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
318
    using GridderConfig::GridderConfig;
319 320 321 322
    PyGridderConfig(size_t nxdirty, size_t nydirty, size_t nu_, size_t nv_,
      double epsilon, double pixsize_x, double pixsize_y, size_t nthreads_)
      : GridderConfig(nxdirty, nydirty, nu_, nv_, epsilon, pixsize_x, pixsize_y, nthreads_) {}

Martin Reinecke's avatar
Martin Reinecke committed
323
    PyGridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
324 325
      double pixsize_x, double pixsize_y, size_t nthreads_)
      : GridderConfig(nxdirty, nydirty, epsilon, pixsize_x, pixsize_y, nthreads_) {}
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
326

Martin Reinecke's avatar
Martin Reinecke committed
327
    template<typename T> pyarr<T> apply_taper2(const pyarr<T> &img, bool divide) const
Martin Reinecke's avatar
Martin Reinecke committed
328
      {
Martin Reinecke's avatar
merge  
Martin Reinecke committed
329
      auto res = makeArray<T>({nx_dirty, ny_dirty});
Martin Reinecke's avatar
Martin Reinecke committed
330 331
      auto img2 = make_const_mav<2>(img);
      auto res2 = make_mav<2>(res);
Martin Reinecke's avatar
Martin Reinecke committed
332 333
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
334
      GridderConfig::apply_taper(img2, res2, divide);
Martin Reinecke's avatar
Martin Reinecke committed
335
      }
Martin Reinecke's avatar
Martin Reinecke committed
336 337
      return res;
      }
Martin Reinecke's avatar
Martin Reinecke committed
338 339
    py::array apply_taper(const py::array &img, bool divide) const
      {
Martin Reinecke's avatar
fix  
Martin Reinecke committed
340
      if (isPytype<float>(img))
Martin Reinecke's avatar
Martin Reinecke committed
341
        return apply_taper2<float>(img, divide);
Martin Reinecke's avatar
fix  
Martin Reinecke committed
342
      if (isPytype<double>(img))
Martin Reinecke's avatar
Martin Reinecke committed
343 344 345 346
        return apply_taper2<double>(img, divide);
      myfail("type matching failed: 'img' has neither type 'f4' nor 'f8'");
      }
    template<typename T> pyarr<T> grid2dirty2(const pyarr<T> &grid) const
347 348 349 350 351 352
      {
      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;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
353
      GridderConfig::grid2dirty(grid2,res2);
354 355 356
      }
      return res;
      }
Martin Reinecke's avatar
Martin Reinecke committed
357 358 359 360 361 362 363 364 365 366
    py::array grid2dirty(const py::array &grid) const
      {
      if (isPytype<float>(grid))
        return grid2dirty2<float>(grid);
      if (isPytype<double>(grid))
        return grid2dirty2<double>(grid);
      myfail("type matching failed: 'grid' has neither type 'f4' nor 'f8'");
      }
    template<typename T> pyarr<complex<T>> grid2dirty_c2
      (const pyarr<complex<T>> &grid) const
Martin Reinecke's avatar
Martin Reinecke committed
367
      {
Martin Reinecke's avatar
merge  
Martin Reinecke committed
368
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
Martin Reinecke's avatar
Martin Reinecke committed
369 370
      auto grid2=make_const_mav<2>(grid);
      auto res2=make_mav<2>(res);
Martin Reinecke's avatar
Martin Reinecke committed
371 372
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
373
      GridderConfig::grid2dirty_c(grid2,res2);
Martin Reinecke's avatar
Martin Reinecke committed
374
      }
375 376
      return res;
      }
Martin Reinecke's avatar
Martin Reinecke committed
377 378 379 380 381 382 383 384
    py::array grid2dirty_c(const py::array &grid) const
      {
      if (isPytype<complex<float>>(grid))
        return grid2dirty_c2<float>(grid);
      if (isPytype<complex<double>>(grid))
        return grid2dirty_c2<double>(grid);
      myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
      }
385

Martin Reinecke's avatar
Martin Reinecke committed
386
    template<typename T> pyarr<T> dirty2grid2(const pyarr<T> &dirty) const
387 388 389 390 391 392
      {
      auto dirty2 = make_const_mav<2>(dirty);
      auto grid = makeArray<T>({nu, nv});
      auto grid2=make_mav<2>(grid);
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
393
      GridderConfig::dirty2grid(dirty2, grid2);
394 395 396
      }
      return grid;
      }
Martin Reinecke's avatar
Martin Reinecke committed
397 398 399 400 401 402 403 404 405
    py::array dirty2grid(const py::array &dirty) const
      {
      if (isPytype<float>(dirty))
        return dirty2grid2<float>(dirty);
      if (isPytype<double>(dirty))
        return dirty2grid2<double>(dirty);
      myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
      }
    template<typename T> pyarr<complex<T>> dirty2grid_c2(const pyarr<complex<T>> &dirty) const
406
      {
Martin Reinecke's avatar
updates  
Martin Reinecke committed
407 408 409
      auto dirty2 = make_const_mav<2>(dirty);
      auto grid = makeArray<complex<T>>({nu, nv});
      auto grid2=make_mav<2>(grid);
Martin Reinecke's avatar
Martin Reinecke committed
410 411
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
412
      GridderConfig::dirty2grid_c(dirty2, grid2);
Martin Reinecke's avatar
Martin Reinecke committed
413
      }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
414
      return grid;
415
      }
Martin Reinecke's avatar
Martin Reinecke committed
416 417 418 419 420 421 422 423 424
    py::array dirty2grid_c(const py::array &dirty) const
      {
      if (isPytype<complex<float>>(dirty))
        return dirty2grid_c2<float>(dirty);
      if (isPytype<complex<double>>(dirty))
        return dirty2grid_c2<double>(dirty);
      myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
      }
    template<typename T> pyarr<complex<T>> apply_wscreen2
Martin Reinecke's avatar
Martin Reinecke committed
425 426
      (const pyarr<complex<T>> &dirty, double w, bool adjoint,
       bool divide_by_n) const
Martin Reinecke's avatar
test1  
Martin Reinecke committed
427
      {
Martin Reinecke's avatar
updates  
Martin Reinecke committed
428 429 430
      auto dirty2 = make_const_mav<2>(dirty);
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
      auto res2 = make_mav<2>(res);
Martin Reinecke's avatar
test1  
Martin Reinecke committed
431 432
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
433
      GridderConfig::apply_wscreen(dirty2, res2, w, adjoint, divide_by_n);
434
      }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
435
      return res;
436
      }
Martin Reinecke's avatar
Martin Reinecke committed
437 438
    py::array apply_wscreen(const py::array &dirty, double w, bool adjoint,
      bool divide_by_n) const
Martin Reinecke's avatar
Martin Reinecke committed
439 440
      {
      if (isPytype<complex<float>>(dirty))
Martin Reinecke's avatar
Martin Reinecke committed
441
        return apply_wscreen2<float>(dirty, w, adjoint, divide_by_n);
Martin Reinecke's avatar
Martin Reinecke committed
442
      if (isPytype<complex<double>>(dirty))
Martin Reinecke's avatar
Martin Reinecke committed
443
        return apply_wscreen2<double>(dirty, w, adjoint, divide_by_n);
Martin Reinecke's avatar
Martin Reinecke committed
444 445
      myfail("type matching failed: 'dirty' has neither type 'c8' nor 'c16'");
      }
Martin Reinecke's avatar
Martin Reinecke committed
446 447
  };

Martin Reinecke's avatar
import  
Martin Reinecke committed
448

449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464
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.
Martin Reinecke's avatar
Martin Reinecke committed
465 466
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.
467 468 469 470 471 472

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
Martin Reinecke's avatar
updates  
Martin Reinecke committed
473
template<typename T> pyarr<complex<T>> Pyvis2grid_c(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
474
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
475
  const pyarr<idx_t> &idx_, const pyarr<complex<T>> &vis_,
Martin Reinecke's avatar
Martin Reinecke committed
476
  py::object &grid_in, const py::object &wgt_)
477
  {
Martin Reinecke's avatar
updates  
Martin Reinecke committed
478 479 480
  auto vis2 = make_const_mav<1>(vis_);
  size_t nvis = vis2.shape(0);
  auto idx2 = make_const_mav<1>(idx_);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
481
  pyarr<T> wgt = providePotentialArray<T>(wgt_, "wgt", {nvis});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
482
  auto wgt2 = make_const_mav<1>(wgt);
Martin Reinecke's avatar
Martin Reinecke committed
483
  auto res = provideArray<complex<T>>(grid_in, "grid_in", {gconf.Nu(), gconf.Nv()});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
484
  auto grid = make_mav<2>(res);
Martin Reinecke's avatar
Martin Reinecke committed
485 486
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
updates  
Martin Reinecke committed
487
  vis2grid_c<T>(baselines, gconf, idx2, vis2, grid, wgt2);
Martin Reinecke's avatar
Martin Reinecke committed
488
  }
489 490 491
  return res;
  }

492 493 494 495 496 497 498 499 500 501 502 503 504 505
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
Martin Reinecke's avatar
Martin Reinecke committed
506 507
grid_in: np.array((nu,nv), dtype=np.float64), optional
    If present, the result is added to this array.
Martin Reinecke's avatar
Martin Reinecke committed
508 509
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.
510 511 512 513 514 515

Returns
=======
np.array((nu,nv), dtype=np.float64):
    the gridded visibilities (made real by making use of Hermitian symmetry)
)""";
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
516
template<typename T> pyarr<T> Pyvis2grid(const PyBaselines &baselines,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
517
  const PyGridderConfig &gconf, const pyarr<idx_t> &idx_,
518 519 520
  const pyarr<complex<T>> &vis_, py::object &grid_in, const py::object &wgt_)
  {
  auto tmp=Pyvis2grid_c(baselines, gconf, idx_, vis_, None, wgt_);
Martin Reinecke's avatar
Martin Reinecke committed
521
  auto grd=provideArray<T>(grid_in, "grid_in", {gconf.Nu(), gconf.Nv()});
Martin Reinecke's avatar
Martin Reinecke committed
522 523
  {
  py::gil_scoped_release release;
524
  gridder::detail::complex2hartley(make_const_mav<2>(tmp), make_mav<2>(grd), gconf.Nthreads());
Martin Reinecke's avatar
Martin Reinecke committed
525
  }
526 527
  return grd;
  }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
528

529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544
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.
Martin Reinecke's avatar
Martin Reinecke committed
545 546
wgt: np.array((nrows, nchannels), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.
547 548 549 550 551 552

Returns
=======
np.array((nu,nv), dtype=np.complex128):
    the gridded visibilities
)""";
Martin Reinecke's avatar
updates  
Martin Reinecke committed
553
template<typename T> pyarr<complex<T>> Pyms2grid_c(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
554
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
555
  const pyarr<idx_t> &idx_, const pyarr<complex<T>> &ms_,
Martin Reinecke's avatar
Martin Reinecke committed
556
  py::object &grid_in, const py::object &wgt_)
Martin Reinecke's avatar
merge  
Martin Reinecke committed
557 558 559
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
Martin Reinecke's avatar
updates  
Martin Reinecke committed
560 561
  auto ms2 = make_const_mav<2>(ms_);
  auto idx2 = make_const_mav<1>(idx_);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
562
  pyarr<T> wgt = providePotentialArray<T>(wgt_, "wgt", {nrows, nchan});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
563
  auto wgt2 = make_const_mav<2>(wgt);
Martin Reinecke's avatar
Martin Reinecke committed
564
  auto res = provideArray<complex<T>>(grid_in, "grid_in", {gconf.Nu(), gconf.Nv()});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
565
  auto grid = make_mav<2>(res);
Martin Reinecke's avatar
Martin Reinecke committed
566 567
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
updates  
Martin Reinecke committed
568
  ms2grid_c<T>(baselines, gconf, idx2, ms2, grid, wgt2);
Martin Reinecke's avatar
Martin Reinecke committed
569
  }
Martin Reinecke's avatar
merge  
Martin Reinecke committed
570 571 572
  return res;
  }

573
template<typename T> pyarr<T> Pyms2grid(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
574
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
575
  const pyarr<idx_t> &idx_, const pyarr<complex<T>> &ms_,
576 577 578
  py::object &grid_in, const py::object &wgt_)
  {
  auto tmp = Pyms2grid_c(baselines, gconf, idx_, ms_, None, wgt_);
Martin Reinecke's avatar
Martin Reinecke committed
579
  auto res_ = provideArray<T>(grid_in, "grid_in", {gconf.Nu(), gconf.Nv()});
580
  auto res = make_mav<2>(res_);
Martin Reinecke's avatar
Martin Reinecke committed
581 582
  {
  py::gil_scoped_release release;
583
  gridder::detail::complex2hartley(make_const_mav<2>(tmp), res, gconf.Nthreads());
Martin Reinecke's avatar
Martin Reinecke committed
584
  }
585 586 587
  return res_;
  }

Martin Reinecke's avatar
updates  
Martin Reinecke committed
588
template<typename T> pyarr<complex<T>> Pygrid2vis_c(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
589
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
590
  const pyarr<idx_t> &idx_, const pyarr<complex<T>> &grid_,
Martin Reinecke's avatar
Martin Reinecke committed
591
  const py::object &wgt_)
592
  {
Martin Reinecke's avatar
updates  
Martin Reinecke committed
593 594 595
  auto grid2 = make_const_mav<2>(grid_);
  auto idx2 = make_const_mav<1>(idx_);
  size_t nvis = idx2.shape(0);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
596
  pyarr<T> wgt = providePotentialArray<T>(wgt_, "wgt", {nvis});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
597
  auto wgt2 = make_const_mav<1>(wgt);
Martin Reinecke's avatar
merge  
Martin Reinecke committed
598
  auto res = makeArray<complex<T>>({nvis});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
599
  auto vis = make_mav<1>(res);
Martin Reinecke's avatar
Martin Reinecke committed
600 601
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
602
  vis.fill(0);
Martin Reinecke's avatar
updates  
Martin Reinecke committed
603
  grid2vis_c<T>(baselines, gconf, idx2, grid2, vis, wgt2);
Martin Reinecke's avatar
Martin Reinecke committed
604
  }
605 606 607
  return res;
  }

608 609 610 611 612 613 614 615 616 617 618 619 620 621
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)
Martin Reinecke's avatar
Martin Reinecke committed
622 623
wgt: np.array((nvis,), dtype=np.float64), optional
    If present, visibilities are multiplied by the corresponding entries.
624 625 626 627 628 629

Returns
=======
np.array((nvis,), dtype=np.complex)
    The degridded visibility data
)""";
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
630
template<typename T> pyarr<complex<T>> Pygrid2vis(const PyBaselines &baselines,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
631
  const PyGridderConfig &gconf, const pyarr<idx_t> &idx_,
632 633 634 635 636 637 638
  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_);
  }

Martin Reinecke's avatar
updates  
Martin Reinecke committed
639
template<typename T> pyarr<complex<T>> Pygrid2ms_c(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
640
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
641
  const pyarr<idx_t> &idx_, const pyarr<complex<T>> &grid_,
Martin Reinecke's avatar
updates  
Martin Reinecke committed
642 643 644 645 646 647
  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_);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
648
  pyarr<T> wgt = providePotentialArray<T>(wgt_, "wgt", {nrows, nchan});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
649
  auto wgt2 = make_const_mav<2>(wgt);
Martin Reinecke's avatar
Martin Reinecke committed
650
  auto res = provideArray<complex<T>>(ms_in, "ms_in", {nrows, nchan});
Martin Reinecke's avatar
updates  
Martin Reinecke committed
651 652 653 654 655 656 657
  auto ms = make_mav<2>(res);
  {
  py::gil_scoped_release release;
  grid2ms_c<T>(baselines, gconf, idx2, grid2, ms, wgt2);
  }
  return res;
  }
Martin Reinecke's avatar
merge  
Martin Reinecke committed
658

Martin Reinecke's avatar
step 1  
Martin Reinecke committed
659
template<typename T> pyarr<complex<T>> Pygrid2ms(const PyBaselines &baselines,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
660
  const PyGridderConfig &gconf, const pyarr<idx_t> &idx_,
661 662 663 664 665 666 667 668
  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_);
  }

Martin Reinecke's avatar
Martin Reinecke committed
669
template<typename T> pyarr<complex<T>> apply_holo2(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
670
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
Martin Reinecke committed
671
  const py::array &idx_, const py::array &grid_, const py::object &wgt_)
672
  {
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
673
  auto idx = getPyarr<idx_t>(idx_, "idx");
Martin Reinecke's avatar
Martin Reinecke committed
674 675 676 677 678 679 680
  auto idx2 = make_const_mav<1>(idx);
  auto grid = getPyarr<complex<T>>(grid_, "grid");
  auto grid2 = make_const_mav<2>(grid);
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
  auto wgt2 = make_const_mav<1>(wgt);
  auto res = makeArray<complex<T>>({grid2.shape(0),grid2.shape(1)});
  auto res2 = make_mav<2>(res);
681 682
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
683
  apply_holo(baselines, gconf, idx2, grid2, res2, wgt2);
684 685 686
  }
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
687 688 689 690 691 692 693 694 695 696
py::array Pyapply_holo(
  const PyBaselines &baselines, const PyGridderConfig &gconf,
  const py::array &idx, const py::array &grid, const py::object &wgt)
  {
  if (isPytype<complex<float>>(grid))
    return apply_holo2<float>(baselines, gconf, idx, grid, wgt);
  if (isPytype<complex<double>>(grid))
    return apply_holo2<double>(baselines, gconf, idx, grid, wgt);
  myfail("type matching failed: 'grid' has neither type 'c8' nor 'c16'");
  }
697

Martin Reinecke's avatar
Martin Reinecke committed
698
template<typename T> pyarr<T> Pyget_correlations(
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
699
  const PyBaselines &baselines, const PyGridderConfig &gconf,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
700
  const pyarr<idx_t> &idx_, int du, int dv, const py::object &wgt_)
701
  {
Martin Reinecke's avatar
Martin Reinecke committed
702
  auto idx = make_const_mav<1>(idx_);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
703
  pyarr<T> wgt2 = providePotentialArray<T>(wgt_, "wgt", {idx.shape(0)});
Martin Reinecke's avatar
Martin Reinecke committed
704 705 706 707
  auto wgt=make_const_mav<1>(wgt2);

  auto res = makeArray<T>({gconf.Nu(),gconf.Nv()});
  auto ogrid = make_mav<2>(res);
708 709
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
710
  get_correlations(baselines, gconf, idx, du, dv, ogrid, wgt);
711 712 713 714
  }
  return res;
  }

715 716 717 718 719 720 721 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739 740 741
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.
)""";
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
742
pyarr<idx_t> PygetIndices(const PyBaselines &baselines,
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
743
  const PyGridderConfig &gconf, const pyarr<bool> &flags_, int chbegin,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
744
  int chend, double wmin, double wmax)
Martin Reinecke's avatar
updates  
Martin Reinecke committed
745
  {
746 747
  size_t nidx;
  auto flags = make_const_mav<2>(flags_);
Martin Reinecke's avatar
Martin Reinecke committed
748 749
  {
  py::gil_scoped_release release;
750
  nidx = getIdxSize(baselines, flags, chbegin, chend, wmin, wmax, gconf.Nthreads());
Martin Reinecke's avatar
Martin Reinecke committed
751
  }
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
752
  auto res = makeArray<idx_t>({nidx});
753
  auto res2 = make_mav<1>(res);
Martin Reinecke's avatar
Martin Reinecke committed
754 755
  {
  py::gil_scoped_release release;
756
  fillIdx(baselines, gconf, flags, chbegin, chend, wmin, wmax, res2);
Martin Reinecke's avatar
Martin Reinecke committed
757
  }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
758 759 760
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
761 762
template<typename T> pyarr<T> vis2dirty2(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const py::array &idx_,
763 764
  const py::array &vis_, const py::object &wgt_, bool do_wstacking,
  size_t verbosity)
Martin Reinecke's avatar
Martin Reinecke committed
765
  {
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
766
  auto idx = getPyarr<idx_t>(idx_, "idx");
Martin Reinecke's avatar
Martin Reinecke committed
767 768 769 770 771
  auto idx2 = make_const_mav<1>(idx);
  auto dirty = makeArray<T>({gconf.Nxdirty(), gconf.Nydirty()});
  auto dirty2 = make_mav<2>(dirty);
  auto vis = getPyarr<complex<T>>(vis_, "vis");
  auto vis2 = make_const_mav<1>(vis);
772 773
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
  auto wgt2 = make_const_mav<1>(wgt);
Martin Reinecke's avatar
Martin Reinecke committed
774 775
  {
  py::gil_scoped_release release;
776 777
  vis2dirty<T>(baselines, gconf, idx2, vis2, wgt2, dirty2, do_wstacking,
    verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
778
  }
Martin Reinecke's avatar
updates  
Martin Reinecke committed
779
  return dirty;
Martin Reinecke's avatar
Martin Reinecke committed
780
  }
Martin Reinecke's avatar
Martin Reinecke committed
781 782 783 784 785 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802

constexpr auto vis2dirty_DS = R"""(
Converts an array of visibilities to a dirty image.

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 provided visibilities
vis: np.array(nvis,), dtype=np.complex64 or np.complex128)
    The input visibilities
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nvis,), dtype=float with same precision as `vis`, optional
    If present, visibilities are multiplied by the corresponding entries.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.
803 804 805 806
verbosity: int
    0: no output
    1: some output
    2: detailed output
Martin Reinecke's avatar
Martin Reinecke committed
807 808 809 810 811 812 813

Returns
=======
np.array((nxdirty, nydirty), dtype=float of same precision as `vis`.)
    The dirty image
)""";

Martin Reinecke's avatar
Martin Reinecke committed
814 815
py::array Pyvis2dirty(const PyBaselines &baselines,
  const PyGridderConfig &gconf, const py::array &idx,
816 817
  const py::array &vis, const py::object &wgt, bool do_wstacking,
  size_t verbosity)
Martin Reinecke's avatar
Martin Reinecke committed
818 819
  {
  if (isPytype<complex<float>>(vis))
820 821
    return vis2dirty2<float>(baselines, gconf, idx, vis, wgt, do_wstacking,
      verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
822
  if (isPytype<complex<double>>(vis))
823 824
    return vis2dirty2<double>(baselines, gconf, idx, vis, wgt, do_wstacking,
      verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
825 826
  myfail("type matching failed: 'vis' has neither type 'c8' nor 'c16'");
  }
Martin Reinecke's avatar
Martin Reinecke committed
827

Martin Reinecke's avatar
Martin Reinecke committed
828
template<typename T> pyarr<complex<T>> dirty2vis2(const PyBaselines &baselines,
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
829
  const PyGridderConfig &gconf, const pyarr<idx_t> &idx_,
830 831
  const pyarr<T> &dirty_, const py::object &wgt_, bool do_wstacking,
  size_t verbosity)
832
  {
Martin Reinecke's avatar
bug fix  
Martin Reinecke committed
833
  auto idx = getPyarr<idx_t>(idx_, "idx");
Martin Reinecke's avatar
Martin Reinecke committed
834 835 836
  auto idx2 = make_const_mav<1>(idx);
  auto dirty = getPyarr<T>(dirty_, "dirty");
  auto dirty2 = make_const_mav<2>(dirty_);
837 838
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {idx2.shape(0)});
  auto wgt2 = make_const_mav<1>(wgt);
Martin Reinecke's avatar
Martin Reinecke committed
839 840
  auto vis = makeArray<complex<T>>({idx2.shape(0)});
  auto vis2 = make_mav<1>(vis);
841 842
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
843
  vis2.fill(0);
844 845
  dirty2vis<T>(baselines, gconf, idx2, dirty2, wgt2, vis2, do_wstacking,
    verbosity);
846 847 848
  }
  return vis;
  }
Martin Reinecke's avatar
Martin Reinecke committed
849 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867 868 869 870

constexpr auto dirty2vis_DS = R"""(
Converts a dirty image into a 1D array of visibilities.

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 visibilities to be computed
dirty: np.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
    dirty image
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nvis,), same dtype as `dirty`, optional
    If present, visibilities are multiplied by the corresponding entries.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.
871 872 873 874
verbosity: int
    0: no output
    1: some output
    2: detailed output
Martin Reinecke's avatar
Martin Reinecke committed
875 876 877 878 879 880

Returns
=======
np.array((nvis,), dtype=complex of same precision as `dirty`.)
    The visibility data
)""";
Martin Reinecke's avatar
Martin Reinecke committed
881
py::array Pydirty2vis(const PyBaselines &baselines,
882
  const PyGridderConfig &gconf, const py::array &idx, const py::array &dirty,
883
  const py::object &wgt, bool do_wstacking, size_t verbosity)
Martin Reinecke's avatar
Martin Reinecke committed
884 885
  {
  if (isPytype<float>(dirty))
886 887
    return dirty2vis2<float>(baselines, gconf, idx, dirty, wgt, do_wstacking,
      verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
888
  if (isPytype<double>(dirty))
889 890
    return dirty2vis2<double>(baselines, gconf, idx, dirty, wgt, do_wstacking,
      verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
891 892
  myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
  }
893

Martin Reinecke's avatar
Martin Reinecke committed
894
template<typename T> py::array ms2dirty_general2(const py::array &uvw_,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
895
  const py::array &freq_, const py::array &ms_, const py::object &wgt_,
Martin Reinecke's avatar
Martin Reinecke committed
896
  size_t npix_x, size_t npix_y, double pixsize_x, double pixsize_y, size_t nu, size_t nv, double epsilon,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
897
  bool do_wstacking, size_t nthreads, size_t verbosity)
898
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
899
  auto uvw = getPyarr<double>(uvw_, "uvw");
900
  auto uvw2 = make_const_mav<2>(uvw);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
901
  auto freq = getPyarr<double>(freq_, "freq");
902
  auto freq2 = make_const_mav<1>(freq);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
903
  auto ms = getPyarr<complex<T>>(ms_, "ms");
904
  auto ms2 = make_const_mav<2>(ms);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
905
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {ms2.shape(0),ms2.shape(1)});
906
  auto wgt2 = make_const_mav<2>(wgt);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
907 908
  auto dirty = makeArray<T>({npix_x,npix_y});
  auto dirty2 = make_mav<2>(dirty);
Martin Reinecke's avatar
Martin Reinecke committed
909 910
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
911
  ms2dirty_general(uvw2,freq2,ms2,wgt2,pixsize_x,pixsize_y, nu, nv, epsilon,do_wstacking,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
912
    nthreads,dirty2,verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
913
  }
914 915
  return dirty;
  }
Martin Reinecke's avatar
Martin Reinecke committed
916 917 918 919 920 921 922 923 924 925 926 927 928
py::array Pyms2dirty_general(const py::array &uvw,
  const py::array &freq, const py::array &ms, const py::object &wgt,
  size_t npix_x, size_t npix_y, double pixsize_x, double pixsize_y, size_t nu, size_t nv, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  if (isPytype<complex<float>>(ms))
    return ms2dirty_general2<float>(uvw, freq, ms, wgt, npix_x, npix_y,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
  if (isPytype<complex<double>>(ms))
    return ms2dirty_general2<double>(uvw, freq, ms, wgt, npix_x, npix_y,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
  myfail("type matching failed: 'ms' has neither type 'c8' nor 'c16'");
  }
Martin Reinecke's avatar
Martin Reinecke committed
929 930 931 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949 950 951 952 953 954 955 956 957 958 959 960 961 962 963 964 965
constexpr auto ms2dirty_DS = R"""(
Converts an MS object to dirty image.

Parameters
==========
uvw: np.array((nrows, 3), dtype=np.float64)
    UVW coordinates from the measurement set
freq: np.array((nchan,), dtype=np.float64)
    channel frequencies
ms: np.array((nrows, nchan,), dtype=np.complex64 or np.complex128)
    the input measurement set data.
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nrows, nchan), float with same precision as `ms`), optional
    If present, its values are multiplied to the output
npix_x, npix_y: int
    dimensions of the dirty image
pixsize_x, pixsize_y: float
    angular pixel size (in radians) of the dirty image
epsilon: float
    accuracy at which the computation should be done. Must be larger than 2e-13.
    If `ms` has type np.complex64, it must be larger than 1e-5.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.
nthreads: int
    number of threads to use for the calculation
verbosity: int
    0: no output
    1: some output
    2: detailed output

Returns
=======
np.array((nxdirty, nydirty), dtype=float of same precision as `ms`)
    the dirty image
)""";
Martin Reinecke's avatar
Martin Reinecke committed
966
py::array Pyms2dirty(const py::array &uvw,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
967
  const py::array &freq, const py::array &ms, const py::object &wgt,
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
968
  size_t npix_x, size_t npix_y, double pixsize_x, double pixsize_y, double epsilon,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
969
  bool do_wstacking, size_t nthreads, size_t verbosity)
Martin Reinecke's avatar
Martin Reinecke committed
970
  {
Martin Reinecke's avatar
Martin Reinecke committed
971 972 973
  return Pyms2dirty_general(uvw, freq, ms, wgt, npix_x, npix_y,
    pixsize_x, pixsize_y, 2*npix_x, 2*npix_y, epsilon, do_wstacking, nthreads,
    verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
974 975
  }

Martin Reinecke's avatar
Martin Reinecke committed
976
template<typename T> py::array dirty2ms_general2(const py::array &uvw_,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
977
  const py::array &freq_, const py::array &dirty_, const py::object &wgt_,
Martin Reinecke's avatar
Martin Reinecke committed
978
  double pixsize_x, double pixsize_y, size_t nu, size_t nv, double epsilon,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
979
  bool do_wstacking, size_t nthreads, size_t verbosity)
Martin Reinecke's avatar
Martin Reinecke committed
980
  {
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
981
  auto uvw = getPyarr<double>(uvw_, "uvw");
Martin Reinecke's avatar
Martin Reinecke committed
982
  auto uvw2 = make_const_mav<2>(uvw);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
983
  auto freq = getPyarr<double>(freq_, "freq");
Martin Reinecke's avatar
Martin Reinecke committed
984
  auto freq2 = make_const_mav<1>(freq);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
985
  auto dirty = getPyarr<T>(dirty_, "dirty");
Martin Reinecke's avatar
Martin Reinecke committed
986
  auto dirty2 = make_const_mav<2>(dirty);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
987
  auto wgt = providePotentialArray<T>(wgt_, "wgt", {uvw2.shape(0),freq2.shape(0)});
Martin Reinecke's avatar
Martin Reinecke committed
988
  auto wgt2 = make_const_mav<2>(wgt);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
989 990
  auto ms = makeArray<complex<T>>({uvw2.shape(0),freq2.shape(0)});
  auto ms2 = make_mav<2>(ms);
Martin Reinecke's avatar
Martin Reinecke committed
991 992
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
993
  dirty2ms_general(uvw2,freq2,dirty2,wgt2,pixsize_x,pixsize_y,nu, nv,epsilon,do_wstacking,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
994
    nthreads,ms2,verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
995
  }
Martin Reinecke's avatar
Martin Reinecke committed
996 997
  return ms;
  }
Martin Reinecke's avatar
Martin Reinecke committed
998 999 1000 1001 1002 1003 1004 1005 1006 1007 1008 1009 1010
py::array Pydirty2ms_general(const py::array &uvw,
  const py::array &freq, const py::array &dirty, const py::object &wgt,
  double pixsize_x, double pixsize_y, size_t nu, size_t nv, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  if (isPytype<float>(dirty))
    return dirty2ms_general2<float>(uvw, freq, dirty, wgt,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
  if (isPytype<double>(dirty))
    return dirty2ms_general2<double>(uvw, freq, dirty, wgt,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
  myfail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
  }
Martin Reinecke's avatar
Martin Reinecke committed
1011 1012 1013 1014 1015 1016 1017 1018 1019 1020 1021 1022 1023 1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037 1038 1039 1040 1041 1042 1043 1044 1045
constexpr auto dirty2ms_DS = R"""(
Converts a dirty image to an MS object.

Parameters
==========
uvw: np.array((nrows, 3), dtype=np.float64)
    UVW coordinates from the measurement set
freq: np.array((nchan,), dtype=np.float64)
    channel frequencies
dirty: np.array((nxdirty, nydirty), dtype=np.float32 or np.float64)
    dirty image
    Its data type determines the precision in which the calculation is carried
    out.
wgt: np.array((nrows, nchan), same dtype as `dirty`), optional
    If present, its values are multiplied to the output
pixsize_x, pixsize_y: float
    angular pixel size (in radians) of the dirty image
epsilon: float
    accuracy at which the computation should be done. Must be larger than 2e-13.
    If `dirty` has type np.float32, it must be larger than 1e-5.
do_wstacking: bool
    if True, the full improved w-stacking algorithm is carried out, otherwise
    the w values are assumed to be zero.
nthreads: int
    number of threads to use for the calculation
verbosity: int
    0: no output
    1: some output
    2: detailed output

Returns
=======
np.array((nrows, nchan,), dtype=complex of same precision as `dirty`)
    the measurement set data.
)""";
Martin Reinecke's avatar
Martin Reinecke committed
1046
py::array Pydirty2ms(const py::array &uvw,
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
1047 1048 1049 1050
  const py::array &freq, const py::array &dirty, const py::object &wgt,
  double pixsize_x, double pixsize_y, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
Martin Reinecke's avatar
Martin Reinecke committed
1051 1052 1053
  return Pydirty2ms_general(uvw, freq, dirty, wgt, pixsize_x, pixsize_y,
    2*dirty.shape(0), 2*dirty.shape(1), epsilon, do_wstacking, nthreads,
    verbosity);
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
1054 1055
  }

Martin Reinecke's avatar
import  
Martin Reinecke committed
1056 1057 1058 1059
} // unnamed namespace

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

Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1062
  py::class_<PyBaselines> (m, "Baselines", PyBaselines_DS)
Martin Reinecke's avatar
Martin Reinecke committed
1063
    .def(py::init<const pyarr<double> &, const pyarr<double> &>(),
1064
      "coord"_a, "freq"_a)
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1065 1066 1067
    .def ("Nrows",&PyBaselines::Nrows)
    .def ("Nchannels",&PyBaselines::Nchannels)
    .def ("ms2vis",&PyBaselines::ms2vis<complex<double>>,
1068
      ms2vis_DS, "ms"_a, "idx"_a, "nthreads"_a=1)
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1069 1070
    .def ("effectiveuvw",&PyBaselines::effectiveuvw<double>, "idx"_a)
    .def ("vis2ms",&PyBaselines::vis2ms<complex<double>>,
1071
      vis2ms_DS, "vis"_a, "idx"_a, "ms_in"_a=None, "nthreads"_a=1);
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1072
  py::class_<PyGridderConfig> (m, "GridderConfig", GridderConfig_DS)
Martin Reinecke's avatar
Martin Reinecke committed
1073 1074
    .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)
1075 1076
    .def(py::init<size_t, size_t, size_t, size_t, double, double, double, size_t>(),"nxdirty"_a,
      "nydirty"_a, "nu"_a, "nv"_a, "epsilon"_a, "pixsize_x"_a, "pixsize_y"_a, "nthreads"_a=1)
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1077 1078 1079 1080 1081 1082 1083 1084
    .def("Nxdirty", &PyGridderConfig::Nxdirty)
    .def("Nydirty", &PyGridderConfig::Nydirty)
    .def("Epsilon", &PyGridderConfig::Epsilon)
    .def("Pixsize_x", &PyGridderConfig::Pixsize_x)
    .def("Pixsize_y", &PyGridderConfig::Pixsize_y)
    .def("Nu", &PyGridderConfig::Nu)
    .def("Nv", &PyGridderConfig::Nv)
    .def("Supp", &PyGridderConfig::Supp)
Martin Reinecke's avatar
Martin Reinecke committed
1085
    .def("apply_taper", &PyGridderConfig::apply_taper, apply_taper_DS,
Martin Reinecke's avatar
Martin Reinecke committed
1086
      "img"_a, "divide"_a=false)
1087 1088
    .def("grid2dirty", &PyGridderConfig::grid2dirty,
      grid2dirty_DS, "grid"_a)
Martin Reinecke's avatar
Martin Reinecke committed
1089 1090
    .def("grid2dirty_c", &PyGridderConfig::grid2dirty_c, "grid"_a)
    .def("dirty2grid", &PyGridderConfig::dirty2grid,
1091
       dirty2grid_DS, "dirty"_a)
Martin Reinecke's avatar
Martin Reinecke committed
1092 1093
    .def("dirty2grid_c", &PyGridderConfig::dirty2grid_c, "dirty"_a)
    .def("apply_wscreen", &PyGridderConfig::apply_wscreen,
Martin Reinecke's avatar
Martin Reinecke committed
1094
      apply_wscreen_DS, "dirty"_a, "w"_a, "adjoint"_a, "divide_by_n"_a=true)
1095 1096 1097 1098

    // pickle support
    .def(py::pickle(
        // __getstate__
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1099
        [](const PyGridderConfig & gc) {
1100
          // Encode object state in tuple
1101 1102
          return py::make_tuple(gc.Nxdirty(), gc.Nydirty(), gc.Nu(), gc.Nv(),
                                gc.Epsilon(),
Martin Reinecke's avatar
Martin Reinecke committed
1103
                                gc.Pixsize_x(), gc.Pixsize_y(), gc.Nthreads());
1104 1105 1106
        },
        // __setstate__
        [](py::tuple t) {
1107
          myassert(t.size()==8,"Invalid state");
1108 1109

          // Reconstruct from tuple
Martin Reinecke's avatar
step 1  
Martin Reinecke committed
1110
          return PyGridderConfig(t[0].cast<size_t>(), t[1].cast<size_t>(),
1111 1112 1113
                                 t[2].cast<size_t>(), t[3].cast<size_t>(),
                                 t[4].cast<double>(), t[5].cast<double>(),
                                 t[6].cast<double>(), t[7].cast<size_t>());
1114 1115

        }));
Martin Reinecke's avatar
cleanup  
Martin Reinecke committed
1116
  m.def("getIndices", PygetIndices, getIndices_DS, "baselines"_a,