nifty_gridder.cc 35 KB
Newer Older
Martin Reinecke's avatar
Martin Reinecke committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
/*
 *  This file is part of nifty_gridder.
 *
 *  nifty_gridder is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  nifty_gridder is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with nifty_fridder; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

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
24
25
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <iostream>
#include <algorithm>
Martin Reinecke's avatar
Martin Reinecke committed
26
27

#define POCKETFFT_OPENMP
Martin Reinecke's avatar
updates    
Martin Reinecke committed
28
#include "pocketfft_hdronly.h"
Martin Reinecke's avatar
import  
Martin Reinecke committed
29

Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
30
31
#ifdef __GNUC__
#define RESTRICT __restrict__
Martin Reinecke's avatar
Martin Reinecke committed
32
#define NOINLINE __attribute__ ((noinline))
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
33
34
35
36
#else
#define RESTRICT
#endif

Martin Reinecke's avatar
import  
Martin Reinecke committed
37
38
39
40
41
42
using namespace std;

namespace py = pybind11;

namespace {

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
43
44
45
46
//
// basic utilities
//

47
48
49
50
51
52
void myassert(bool cond, const char *msg)
  {
  if (cond) return;
  throw runtime_error(msg);
  }

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
53
54
55
56
57
58
59
60
61
62
63
/*! Returns the remainder of the division \a v1/v2.
    The result is non-negative.
    \a v1 can be positive or negative; \a v2 must be positive. */
template<typename T> inline T fmodulo (T v1, T v2)
  {
  if (v1>=0)
    return (v1<v2) ? v1 : fmod(v1,v2);
  T tmp=fmod(v1,v2)+v2;
  return (tmp==v2) ? T(0) : tmp;
  }

Martin Reinecke's avatar
Martin Reinecke committed
64
65
66
67
//
// Utilities for Gauss-Legendre quadrature
//

Martin Reinecke's avatar
Martin Reinecke committed
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
static inline double one_minus_x2 (double x)
  { return (fabs(x)>0.1) ? (1.+x)*(1.-x) : 1.-x*x; }

void legendre_prep(int n, vector<double> &x, vector<double> &w)
  {
  constexpr double pi = 3.141592653589793238462643383279502884197;
  constexpr double eps = 3e-14;
  int m = (n+1)>>1;
  x.resize(m);
  w.resize(m);

  double t0 = 1 - (1-1./n) / (8.*n*n);
  double t1 = 1./(4.*n+2.);

#pragma omp parallel
{
  int i;
#pragma omp for schedule(dynamic,100)
  for (i=1; i<=m; ++i)
    {
    double x0 = cos(pi * ((i<<2)-1) * t1) * t0;

    int dobreak=0;
    int j=0;
    double dpdx;
    while(1)
      {
      double P_1 = 1.0;
      double P0 = x0;
      double dx, x1;

      for (int k=2; k<=n; k++)
        {
        double P_2 = P_1;
        P_1 = P0;
//        P0 = ((2*k-1)*x0*P_1-(k-1)*P_2)/k;
        P0 = x0*P_1 + (k-1.)/k * (x0*P_1-P_2);
        }

      dpdx = (P_1 - x0*P0) * n / one_minus_x2(x0);

      /* Newton step */
      x1 = x0 - P0/dpdx;
      dx = x0-x1;
      x0 = x1;
      if (dobreak) break;

      if (abs(dx)<=eps) dobreak=1;
      if (++j>=100) throw runtime_error("convergence problem");
      }

    x[m-i] = x0;
    w[m-i] = 2. / (one_minus_x2(x0) * dpdx * dpdx);
    }
} // end of parallel region
  }

Martin Reinecke's avatar
Martin Reinecke committed
125
126
127
128
//
// Start of real gridder functionality
//

129
130
template<typename T>
  using pyarr = py::array_t<T>;
131
// The "_c" suffix here stands for "C memory order, contiguous"
132
133
template<typename T>
  using pyarr_c = py::array_t<T, py::array::c_style | py::array::forcecast>;
Martin Reinecke's avatar
import  
Martin Reinecke committed
134

Martin Reinecke's avatar
merge    
Martin Reinecke committed
135
template<typename T> pyarr_c<T> makeArray(const vector<size_t> &shape)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
136
137
  { return pyarr_c<T>(shape); }

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
138
size_t get_w(double epsilon)
Martin Reinecke's avatar
Martin Reinecke committed
139
140
141
142
143
144
145
146
147
148
149
  {
  static const vector<double> maxmaperr { 1e8, 0.32, 0.021, 6.2e-4,
    1.08e-5, 1.25e-7, 8.25e-10, 5.70e-12, 1.22e-13, 2.48e-15, 4.82e-17,
    6.74e-19, 5.41e-21, 4.41e-23, 7.88e-25, 3.9e-26 };

  double epssq = epsilon*epsilon;

  for (size_t i=1; i<maxmaperr.size(); ++i)
    if (epssq>maxmaperr[i]) return i;
  throw runtime_error("requested epsilon too small - minimum is 2e-13");
  }
Martin Reinecke's avatar
Martin Reinecke committed
150

Martin Reinecke's avatar
merge    
Martin Reinecke committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
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");
      }
  }

169
template<typename T> pyarr<T> provideArray(py::object &in,
Martin Reinecke's avatar
merge    
Martin Reinecke committed
170
171
172
173
174
175
176
177
178
179
180
  const vector<size_t> &shape)
  {
  if (in.is(py::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_;
    }
181
  auto tmp_ = in.cast<pyarr<T>>();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
182
183
184
185
  checkArray(tmp_, "temporary", shape);
  return tmp_;
  }

Martin Reinecke's avatar
Martin Reinecke committed
186
187
template<typename T> pyarr_c<T> complex2hartley
  (const pyarr_c<complex<T>> &grid_)
Martin Reinecke's avatar
Martin Reinecke committed
188
  {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
189
  checkArray(grid_, "grid", {0,0});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
190
  size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
Martin Reinecke's avatar
Martin Reinecke committed
191
192
  auto grid = grid_.data();

Martin Reinecke's avatar
merge    
Martin Reinecke committed
193
  auto res = makeArray<T>({nu,nv});
Martin Reinecke's avatar
Martin Reinecke committed
194
  auto grid2 = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
195
196
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
197
#pragma omp parallel for
Martin Reinecke's avatar
updates    
Martin Reinecke committed
198
  for (size_t u=0; u<nu; ++u)
Martin Reinecke's avatar
Martin Reinecke committed
199
    {
Martin Reinecke's avatar
updates    
Martin Reinecke committed
200
201
    size_t xu = (u==0) ? 0 : nu-u;
    for (size_t v=0; v<nv; ++v)
Martin Reinecke's avatar
Martin Reinecke committed
202
      {
Martin Reinecke's avatar
updates    
Martin Reinecke committed
203
204
205
      size_t xv = (v==0) ? 0 : nv-v;
      size_t i1 = u*nv+v;
      size_t i2 = xu*nv+xv;
Martin Reinecke's avatar
Martin Reinecke committed
206
207
      grid2[i1] = T(0.5)*(grid[i1].real()+grid[i1].imag()+
                          grid[i2].real()-grid[i2].imag());
Martin Reinecke's avatar
Martin Reinecke committed
208
209
      }
    }
Martin Reinecke's avatar
Martin Reinecke committed
210
  }
Martin Reinecke's avatar
Martin Reinecke committed
211
212
213
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
214
215
template<typename T> pyarr_c<complex<T>> hartley2complex
  (const pyarr_c<T> &grid_)
Martin Reinecke's avatar
Martin Reinecke committed
216
  {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
217
  checkArray(grid_, "grid", {0, 0});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
218
  size_t nu = size_t(grid_.shape(0)), nv = size_t(grid_.shape(1));
Martin Reinecke's avatar
Martin Reinecke committed
219
220
  auto grid = grid_.data();

Martin Reinecke's avatar
merge    
Martin Reinecke committed
221
  auto res=makeArray<complex<T>>({nu, nv});
Martin Reinecke's avatar
Martin Reinecke committed
222
  auto grid2 = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
223
224
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
225
#pragma omp parallel for
Martin Reinecke's avatar
updates    
Martin Reinecke committed
226
  for (size_t u=0; u<nu; ++u)
Martin Reinecke's avatar
Martin Reinecke committed
227
    {
Martin Reinecke's avatar
updates    
Martin Reinecke committed
228
229
    size_t xu = (u==0) ? 0 : nu-u;
    for (size_t v=0; v<nv; ++v)
Martin Reinecke's avatar
Martin Reinecke committed
230
      {
Martin Reinecke's avatar
updates    
Martin Reinecke committed
231
232
233
      size_t xv = (v==0) ? 0 : nv-v;
      size_t i1 = u*nv+v;
      size_t i2 = xu*nv+xv;
Martin Reinecke's avatar
Martin Reinecke committed
234
235
236
      T v1 = T(0.5)*grid[i1];
      T v2 = T(0.5)*grid[i2];
      grid2[i1] = complex<T>(v1+v2, v1-v2);
Martin Reinecke's avatar
Martin Reinecke committed
237
238
      }
    }
Martin Reinecke's avatar
Martin Reinecke committed
239
  }
Martin Reinecke's avatar
Martin Reinecke committed
240
241
242
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
243
244
245
template<typename T> void hartley2_2D(const pyarr_c<T> &in, pyarr_c<T> &out)
  {
  size_t nu=in.shape(0), nv=in.shape(1);
Martin Reinecke's avatar
Martin Reinecke committed
246
247
248
  pocketfft::stride_t s_i{in.strides(0), in.strides(1)},
                      s_o{out.strides(0), out.strides(1)};
  auto d_i = in.data();
Martin Reinecke's avatar
Martin Reinecke committed
249
  auto ptmp = out.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
250
251
252
  {
  py::gil_scoped_release release;
  pocketfft::r2r_hartley({nu, nv}, s_i, s_o, {0,1}, d_i, ptmp, T(1), 0);
Martin Reinecke's avatar
Martin Reinecke committed
253
#pragma omp parallel for
Martin Reinecke's avatar
Martin Reinecke committed
254
255
256
257
258
259
260
261
262
263
264
265
266
  for(size_t i=1; i<(nu+1)/2; ++i)
    for(size_t j=1; j<(nv+1)/2; ++j)
       {
       T a = ptmp[i*nv+j];
       T b = ptmp[(nu-i)*nv+j];
       T c = ptmp[i*nv+nv-j];
       T d = ptmp[(nu-i)*nv+nv-j];
       ptmp[i*nv+j] = T(0.5)*(a+b+c-d);
       ptmp[(nu-i)*nv+j] = T(0.5)*(a+b+d-c);
       ptmp[i*nv+nv-j] = T(0.5)*(a+c+d-b);
       ptmp[(nu-i)*nv+nv-j] = T(0.5)*(b+c+d-a);
       }
  }
Martin Reinecke's avatar
Martin Reinecke committed
267
  }
Martin Reinecke's avatar
Martin Reinecke committed
268

Martin Reinecke's avatar
Martin Reinecke committed
269
270
/* Compute correction factors for the ES gridding kernel
   This implementation follows eqs. (3.8) to (3.10) of Barnett et al. 2018 */
Martin Reinecke's avatar
updates    
Martin Reinecke committed
271
vector<double> correction_factors (size_t n, size_t nval, size_t w)
Martin Reinecke's avatar
Martin Reinecke committed
272
273
274
275
276
277
278
279
280
281
  {
  constexpr double pi = 3.141592653589793238462643383279502884197;
  auto beta = 2.3*w;
  auto p = int(1.5*w+2);
  double alpha = pi*w/n;
  vector<double> x, wgt;
  legendre_prep(2*p,x,wgt);
  auto psi = x;
  for (auto &v:psi)
    v = exp(beta*(sqrt(1-v*v)-1.));
Martin Reinecke's avatar
updates    
Martin Reinecke committed
282
  vector<double> res(nval);
Martin Reinecke's avatar
Martin Reinecke committed
283
#pragma omp parallel for schedule(static)
Martin Reinecke's avatar
Martin Reinecke committed
284
285
286
287
288
  for (size_t k=0; k<nval; ++k)
    {
    double tmp=0;
    for (int i=0; i<p; ++i)
      tmp += wgt[i]*psi[i]*cos(alpha*k*x[i]);
Martin Reinecke's avatar
updates    
Martin Reinecke committed
289
    res[k] = 1./(w*tmp);
Martin Reinecke's avatar
Martin Reinecke committed
290
291
292
293
    }
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
294
template<typename T> struct UVW
Martin Reinecke's avatar
updates    
Martin Reinecke committed
295
  {
Martin Reinecke's avatar
Martin Reinecke committed
296
  T u, v, w;
Martin Reinecke's avatar
updates    
Martin Reinecke committed
297
  UVW () {}
Martin Reinecke's avatar
Martin Reinecke committed
298
299
  UVW (T u_, T v_, T w_) : u(u_), v(v_), w(w_) {}
  UVW operator* (T fct) const
Martin Reinecke's avatar
updates    
Martin Reinecke committed
300
301
    { return UVW(u*fct, v*fct, w*fct); }
  };
Martin Reinecke's avatar
Martin Reinecke committed
302

Martin Reinecke's avatar
Martin Reinecke committed
303
template<typename T> class Baselines
Martin Reinecke's avatar
Martin Reinecke committed
304
305
  {
  private:
Martin Reinecke's avatar
Martin Reinecke committed
306
    vector<UVW<T>> coord;
307
    vector<T> f_over_c;
Martin Reinecke's avatar
updates    
Martin Reinecke committed
308
    size_t nrows, nchan;
Martin Reinecke's avatar
Martin Reinecke committed
309
310

  public:
311
    Baselines(const pyarr<T> &coord_, const pyarr<T> &freq_)
Martin Reinecke's avatar
Martin Reinecke committed
312
      {
313
      constexpr double speedOfLight = 299792458.;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
314
      checkArray(coord_, "coord", {0, 3});
315
      checkArray(freq_, "freq", {0});
Martin Reinecke's avatar
Martin Reinecke committed
316
      nrows = coord_.shape(0);
317
      nchan = freq_.shape(0);
Martin Reinecke's avatar
Martin Reinecke committed
318
      myassert(nrows*nchan<(size_t(1)<<32), "too many entries in MS");
319
320
      auto freq = freq_.template unchecked<1>();
      auto cood = coord_.template unchecked<2>();
Martin Reinecke's avatar
Martin Reinecke committed
321
322
      {
      py::gil_scoped_release release;
323
      f_over_c.resize(nchan);
Martin Reinecke's avatar
updates    
Martin Reinecke committed
324
      for (size_t i=0; i<nchan; ++i)
325
        f_over_c[i] = freq(i)/speedOfLight;
Martin Reinecke's avatar
Martin Reinecke committed
326
327
      coord.resize(nrows);
      for (size_t i=0; i<coord.size(); ++i)
328
        coord[i] = UVW<T>(cood(i,0), cood(i,1), cood(i,2));
Martin Reinecke's avatar
Martin Reinecke committed
329
      }
Martin Reinecke's avatar
Martin Reinecke committed
330
331
      }

Martin Reinecke's avatar
Martin Reinecke committed
332
    UVW<T> effectiveCoord(uint32_t index) const
333
334
335
      {
      size_t irow = index/nchan;
      size_t ichan = index-nchan*irow;
336
      return coord[irow]*f_over_c[ichan];
337
338
      }
    UVW<T> effectiveCoord(size_t irow, size_t ichan) const
339
      { return coord[irow]*f_over_c[ichan]; }
Martin Reinecke's avatar
Martin Reinecke committed
340
    size_t Nrows() const { return nrows; }
Martin Reinecke's avatar
updates    
Martin Reinecke committed
341
342
    size_t Nchannels() const { return nchan; }

343
    template<typename T2> pyarr_c<T2> ms2vis(const pyarr<T2> &ms_,
Martin Reinecke's avatar
updates    
Martin Reinecke committed
344
345
      const pyarr_c<uint32_t> &idx_) const
      {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
346
347
      checkArray(idx_, "idx", {0});
      checkArray(ms_, "ms", {nrows, nchan});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
348
      size_t nvis = size_t(idx_.shape(0));
349
350
      auto idx = idx_.template unchecked<1>();
      auto ms = ms_.template unchecked<2>();
Martin Reinecke's avatar
updates    
Martin Reinecke committed
351

Martin Reinecke's avatar
merge    
Martin Reinecke committed
352
      auto res=makeArray<T2>({nvis});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
353
      auto vis = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
354
355
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
356
#pragma omp parallel for
Martin Reinecke's avatar
updates    
Martin Reinecke committed
357
      for (size_t i=0; i<nvis; ++i)
358
359
360
361
362
363
        {
        auto t = idx(i);
        auto row = t/nchan;
        auto chan = t-row*nchan;
        vis[i] = ms(row, chan);
        }
Martin Reinecke's avatar
Martin Reinecke committed
364
      }
Martin Reinecke's avatar
updates    
Martin Reinecke committed
365
366
367
      return res;
      }

368
369
    template<typename T2> pyarr_c<T2> vis2ms(const pyarr<T2> &vis_,
      const pyarr<uint32_t> &idx_, py::object &ms_in) const
Martin Reinecke's avatar
updates    
Martin Reinecke committed
370
      {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
371
      checkArray(vis_, "vis", {0});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
372
      size_t nvis = size_t(vis_.shape(0));
Martin Reinecke's avatar
merge    
Martin Reinecke committed
373
      checkArray(idx_, "idx", {nvis});
374
375
      auto idx = idx_.template unchecked<1>();
      auto vis = vis_.template unchecked<1>();
Martin Reinecke's avatar
updates    
Martin Reinecke committed
376

Martin Reinecke's avatar
merge    
Martin Reinecke committed
377
      auto res = provideArray<T2>(ms_in, {nrows, nchan});
378
      auto ms = res.template mutable_unchecked<2>();
Martin Reinecke's avatar
Martin Reinecke committed
379
380
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
381
#pragma omp parallel for
Martin Reinecke's avatar
updates    
Martin Reinecke committed
382
      for (size_t i=0; i<nvis; ++i)
383
384
385
386
387
388
        {
        auto t = idx(i);
        auto row = t/nchan;
        auto chan = t-row*nchan;
        ms(row, chan) = vis(i);
        }
Martin Reinecke's avatar
Martin Reinecke committed
389
      }
Martin Reinecke's avatar
updates    
Martin Reinecke committed
390
391
      return res;
      }
Martin Reinecke's avatar
Martin Reinecke committed
392
393
  };

Martin Reinecke's avatar
Martin Reinecke committed
394
395
constexpr int logsquare=4;

Martin Reinecke's avatar
Martin Reinecke committed
396
template<typename T> class GridderConfig
Martin Reinecke's avatar
Martin Reinecke committed
397
398
399
  {
  private:
    size_t nx_dirty, ny_dirty;
Martin Reinecke's avatar
Martin Reinecke committed
400
    double eps, psx, psy;
Martin Reinecke's avatar
Martin Reinecke committed
401
    size_t w, nsafe, nu, nv;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
402
    T beta;
Martin Reinecke's avatar
Martin Reinecke committed
403
    vector<T> cfu, cfv;
Martin Reinecke's avatar
Martin Reinecke committed
404
405
406

  public:
    GridderConfig(size_t nxdirty, size_t nydirty, double epsilon,
Martin Reinecke's avatar
Martin Reinecke committed
407
      double pixsize_x, double pixsize_y)
Martin Reinecke's avatar
Martin Reinecke committed
408
409
      : nx_dirty(nxdirty), ny_dirty(nydirty), eps(epsilon),
        psx(pixsize_x), psy(pixsize_y),
Martin Reinecke's avatar
Martin Reinecke committed
410
411
        w(get_w(epsilon)), nsafe((w+1)/2),
        nu(max(2*nsafe,2*nx_dirty)), nv(max(2*nsafe,2*ny_dirty)),
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
412
        beta(2.3*w),
Martin Reinecke's avatar
updates    
Martin Reinecke committed
413
        cfu(nx_dirty), cfv(ny_dirty)
Martin Reinecke's avatar
Martin Reinecke committed
414
      {
Martin Reinecke's avatar
Martin Reinecke committed
415
416
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
417
418
419
      myassert((nx_dirty&1)==0, "nx_dirty must be even");
      myassert((ny_dirty&1)==0, "ny_dirty must be even");
      myassert(epsilon>0, "epsilon must be positive");
Martin Reinecke's avatar
Martin Reinecke committed
420
421
      myassert(pixsize_x>0, "pixsize_x must be positive");
      myassert(pixsize_y>0, "pixsize_y must be positive");
Martin Reinecke's avatar
updates    
Martin Reinecke committed
422
423
424
425
426
427
428
429
430
431
432

      auto tmp = correction_factors(nu, nx_dirty/2+1, w);
      cfu[nx_dirty/2]=tmp[0];
      cfu[0]=tmp[nx_dirty/2];
      for (size_t i=1; i<nx_dirty/2; ++i)
        cfu[nx_dirty/2-i] = cfu[nx_dirty/2+i] = tmp[i];
      tmp = correction_factors(nv, ny_dirty/2+1, w);
      cfv[ny_dirty/2]=tmp[0];
      cfv[0]=tmp[ny_dirty/2];
      for (size_t i=1; i<ny_dirty/2; ++i)
        cfv[ny_dirty/2-i] = cfv[ny_dirty/2+i] = tmp[i];
Martin Reinecke's avatar
Martin Reinecke committed
433
      }
Martin Reinecke's avatar
Martin Reinecke committed
434
      }
Martin Reinecke's avatar
Martin Reinecke committed
435
436
437
438
439
    size_t Nxdirty() const { return nx_dirty; }
    size_t Nydirty() const { return ny_dirty; }
    double Epsilon() const { return eps; }
    double Pixsize_x() const { return psx; }
    double Pixsize_y() const { return psy; }
Martin Reinecke's avatar
Martin Reinecke committed
440
441
442
    size_t Nu() const { return nu; }
    size_t Nv() const { return nv; }
    size_t W() const { return w; }
443
    size_t Nsafe() const { return nsafe; }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
444
    T Beta() const { return beta; }
Martin Reinecke's avatar
Martin Reinecke committed
445
    pyarr_c<T> grid2dirty(const pyarr_c<T> &grid) const
Martin Reinecke's avatar
Martin Reinecke committed
446
      {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
447
448
      checkArray(grid, "grid", {nu, nv});
      auto tmp = makeArray<T>({nu, nv});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
449
      auto ptmp = tmp.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
450
      hartley2_2D<T>(grid, tmp);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
451
      auto res = makeArray<T>({nx_dirty, ny_dirty});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
452
      auto pout = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
453
454
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
updates    
Martin Reinecke committed
455
456
457
458
459
460
461
462
463
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          pout[ny_dirty*i + j] = ptmp[nv*i2+j2]*cfu[i]*cfv[j];
          }
Martin Reinecke's avatar
Martin Reinecke committed
464
      }
Martin Reinecke's avatar
Martin Reinecke committed
465
466
      return res;
      }
467
468
    pyarr_c<complex<T>> grid2dirty_c(const pyarr_c<complex<T>> &grid) const
      {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
469
470
      checkArray(grid, "grid", {nu, nv});
      auto tmp = makeArray<complex<T>>({nu, nv});
471
472
473
474
      auto ptmp = tmp.mutable_data();
      pocketfft::c2c({nu,nv},{grid.strides(0),grid.strides(1)},
        {tmp.strides(0), tmp.strides(1)}, {0,1}, pocketfft::BACKWARD,
        grid.data(), tmp.mutable_data(), T(1), 0);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
475
      auto res = makeArray<complex<T>>({nx_dirty, ny_dirty});
476
      auto pout = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
477
478
      {
      py::gil_scoped_release release;
479
480
481
482
483
484
485
486
487
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          pout[ny_dirty*i + j] = ptmp[nv*i2+j2]*cfu[i]*cfv[j];
          }
Martin Reinecke's avatar
Martin Reinecke committed
488
      }
489
490
      return res;
      }
Martin Reinecke's avatar
Martin Reinecke committed
491
    pyarr_c<T> dirty2grid(const pyarr_c<T> &dirty) const
Martin Reinecke's avatar
updates    
Martin Reinecke committed
492
      {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
493
      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
494
      auto pdirty = dirty.data();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
495
      auto tmp = makeArray<T>({nu, nv});
Martin Reinecke's avatar
updates    
Martin Reinecke committed
496
      auto ptmp = tmp.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
497
498
      {
      py::gil_scoped_release release;
Martin Reinecke's avatar
updates    
Martin Reinecke committed
499
500
501
502
503
504
505
506
507
508
509
      for (size_t i=0; i<nu*nv; ++i)
        ptmp[i] = 0.;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          ptmp[nv*i2+j2] = pdirty[ny_dirty*i + j]*cfu[i]*cfv[j];
          }
Martin Reinecke's avatar
Martin Reinecke committed
510
      }
Martin Reinecke's avatar
Martin Reinecke committed
511
      hartley2_2D<T>(tmp, tmp);
Martin Reinecke's avatar
updates    
Martin Reinecke committed
512
513
      return tmp;
      }
514
515
    pyarr_c<complex<T>> dirty2grid_c(const pyarr_c<complex<T>> &dirty) const
      {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
516
      checkArray(dirty, "dirty", {nx_dirty, ny_dirty});
517
      auto pdirty = dirty.data();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
518
      auto tmp = makeArray<complex<T>>({nu, nv});
519
      auto ptmp = tmp.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
520
521
522
      pocketfft::stride_t strides{tmp.strides(0),tmp.strides(1)};
      {
      py::gil_scoped_release release;
523
524
525
526
527
528
529
530
531
532
533
      for (size_t i=0; i<nu*nv; ++i)
        ptmp[i] = 0.;
      for (size_t i=0; i<nx_dirty; ++i)
        for (size_t j=0; j<ny_dirty; ++j)
          {
          size_t i2 = nu-nx_dirty/2+i;
          if (i2>=nu) i2-=nu;
          size_t j2 = nv-ny_dirty/2+j;
          if (j2>=nv) j2-=nv;
          ptmp[nv*i2+j2] = pdirty[ny_dirty*i + j]*cfu[i]*cfv[j];
          }
Martin Reinecke's avatar
Martin Reinecke committed
534
535
536
      pocketfft::c2c({nu,nv}, strides, strides, {0,1}, pocketfft::FORWARD,
        ptmp, ptmp, T(1), 0);
      }
537
538
      return tmp;
      }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
539
540
    inline void getpix(T u_in, T v_in, T &u, T &v, int &iu0, int &iv0) const
      {
Martin Reinecke's avatar
Martin Reinecke committed
541
      u=fmodulo(u_in*psx, T(1))*nu,
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
542
543
      iu0 = int(u-w*0.5 + 1 + nu) - nu;
      if (iu0+w>nu+nsafe) iu0 = nu+nsafe-w;
Martin Reinecke's avatar
Martin Reinecke committed
544
      v=fmodulo(v_in*psy, T(1))*nv;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
545
546
547
      iv0 = int(v-w*0.5 + 1 + nv) - nv;
      if (iv0+w>nv+nsafe) iv0 = nv+nsafe-w;
      }
Martin Reinecke's avatar
Martin Reinecke committed
548
549
  };

Martin Reinecke's avatar
Martin Reinecke committed
550
template<typename T> class Helper
Martin Reinecke's avatar
import  
Martin Reinecke committed
551
  {
Martin Reinecke's avatar
Martin Reinecke committed
552
  private:
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
553
    const GridderConfig<T> &gconf;
Martin Reinecke's avatar
Martin Reinecke committed
554
555
    int nu, nv, nsafe, w;
    T beta;
556
557
558
    const complex<T> *grid_r;
    complex<T> *grid_w;
    int su, sv;
Martin Reinecke's avatar
Martin Reinecke committed
559
560
561
    int iu0, iv0; // start index of the current visibility
    int bu0, bv0; // start index of the current buffer

562
    vector<complex<T>> rbuf, wbuf;
Martin Reinecke's avatar
import  
Martin Reinecke committed
563

Martin Reinecke's avatar
Martin Reinecke committed
564
    void dump() const
Martin Reinecke's avatar
import  
Martin Reinecke committed
565
      {
Martin Reinecke's avatar
Martin Reinecke committed
566
      if (bu0<-nsafe) return; // nothing written into buffer yet
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
567

Martin Reinecke's avatar
merge    
Martin Reinecke committed
568
#pragma omp critical (gridder_writing_to_grid)
Martin Reinecke's avatar
import  
Martin Reinecke committed
569
{
Martin Reinecke's avatar
Martin Reinecke committed
570
571
572
      int idxu = (bu0+nu)%nu;
      int idxv0 = (bv0+nv)%nv;
      for (int iu=0; iu<su; ++iu)
Martin Reinecke's avatar
import  
Martin Reinecke committed
573
574
        {
        int idxv = idxv0;
Martin Reinecke's avatar
Martin Reinecke committed
575
        for (int iv=0; iv<sv; ++iv)
Martin Reinecke's avatar
import  
Martin Reinecke committed
576
          {
577
          grid_w[idxu*nv + idxv] += wbuf[iu*sv + iv];
Martin Reinecke's avatar
import  
Martin Reinecke committed
578
579
580
581
582
583
584
585
586
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
}
      }

    void load()
      {
Martin Reinecke's avatar
Martin Reinecke committed
587
588
589
      int idxu = (bu0+nu)%nu;
      int idxv0 = (bv0+nv)%nv;
      for (int iu=0; iu<su; ++iu)
Martin Reinecke's avatar
import  
Martin Reinecke committed
590
591
        {
        int idxv = idxv0;
Martin Reinecke's avatar
Martin Reinecke committed
592
        for (int iv=0; iv<sv; ++iv)
Martin Reinecke's avatar
import  
Martin Reinecke committed
593
          {
594
          rbuf[iu*sv + iv] = grid_r[idxu*nv + idxv];
Martin Reinecke's avatar
import  
Martin Reinecke committed
595
596
597
598
599
600
601
          if (++idxv>=nv) idxv=0;
          }
        if (++idxu>=nu) idxu=0;
        }
      }

  public:
602
603
    const complex<T> *p0r;
    complex<T> *p0w;
Martin Reinecke's avatar
Martin Reinecke committed
604
    vector<T> kernel;
Martin Reinecke's avatar
import  
Martin Reinecke committed
605

606
607
608
609
    Helper(const GridderConfig<T> &gconf_, const complex<T> *grid_r_,
      complex<T> *grid_w_)
      : gconf(gconf_), nu(gconf.Nu()), nv(gconf.Nv()), nsafe(gconf.Nsafe()),
        w(gconf.W()), beta(gconf.Beta()), grid_r(grid_r_), grid_w(grid_w_),
Martin Reinecke's avatar
Martin Reinecke committed
610
        su(2*nsafe+(1<<logsquare)), sv(2*nsafe+(1<<logsquare)),
611
612
613
614
        bu0(-1000000), bv0(-1000000),
        rbuf(su*sv*(grid_r!=nullptr),T(0)),
        wbuf(su*sv*(grid_w!=nullptr),T(0)),
        kernel(2*w)
Martin Reinecke's avatar
Martin Reinecke committed
615
      {}
616
617
618
    ~Helper() { if (grid_w) dump(); }

    int lineJump() const { return sv; }
Martin Reinecke's avatar
Martin Reinecke committed
619
620

    void prep(T u_in, T v_in)
Martin Reinecke's avatar
import  
Martin Reinecke committed
621
      {
Martin Reinecke's avatar
Martin Reinecke committed
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
      T u, v;
      gconf.getpix(u_in, v_in, u, v, iu0, iv0);
      T xw=T(2)/w;
      auto x0 = xw*(iu0-u);
      auto y0 = xw*(iv0-v);
      for (int i=0; i<w; ++i)
        {
        auto x = x0+i*xw;
        kernel[i  ] = beta*sqrt(T(1)-x*x);
        auto y = y0+i*xw;
        kernel[i+w] = beta*sqrt(T(1)-y*y);
        }
      for (auto &k : kernel)
        k = exp(k);

      if ((iu0<bu0) || (iv0<bv0) || (iu0+w>bu0+su) || (iv0+w>bv0+sv))
Martin Reinecke's avatar
import  
Martin Reinecke committed
638
        {
639
        if (grid_w) { dump(); fill(wbuf.begin(), wbuf.end(), T(0)); }
Martin Reinecke's avatar
Martin Reinecke committed
640
641
        bu0=((((iu0+nsafe)>>logsquare)<<logsquare))-nsafe;
        bv0=((((iv0+nsafe)>>logsquare)<<logsquare))-nsafe;
642
        if (grid_r) load();
Martin Reinecke's avatar
import  
Martin Reinecke committed
643
        }
644
645
      p0r = rbuf.data() + sv*(iu0-bu0) + iv0-bv0;
      p0w = wbuf.data() + sv*(iu0-bu0) + iv0-bv0;
Martin Reinecke's avatar
import  
Martin Reinecke committed
646
647
648
      }
  };

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
649
650
template<typename T> pyarr_c<complex<T>> vis2grid_c(
  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
651
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &vis_)
652
  {
Martin Reinecke's avatar
merge    
Martin Reinecke committed
653
654
655
  checkArray(vis_, "vis", {0});
  size_t nvis = size_t(vis_.shape(0));
  checkArray(idx_, "idx", {nvis});
656
657
  auto vis=vis_.template unchecked<1>();
  auto idx = idx_.template unchecked<1>();
658
659

  size_t nu=gconf.Nu(), nv=gconf.Nv();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
660
  auto res = makeArray<complex<T>>({nu, nv});
661
  auto grid = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
662
663
  {
  py::gil_scoped_release release;
664
  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
665
666
  T beta = gconf.Beta();
  size_t w = gconf.W();
667
668
669

#pragma omp parallel
{
670
  Helper<T> hlp(gconf, nullptr, grid);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
671
  T emb = exp(-2*beta);
672
  int jump = hlp.lineJump();
673
  const T * RESTRICT ku = hlp.kernel.data();
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
674
  const T * RESTRICT kv = hlp.kernel.data()+w;
675
676
677
678
679

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
680
    UVW<T> coord = baselines.effectiveCoord(idx(ipart));
Martin Reinecke's avatar
Martin Reinecke committed
681
    hlp.prep(coord.u, coord.v);
682
    auto * RESTRICT ptr = hlp.p0w;
683
    auto v(vis(ipart)*emb);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
684
    for (size_t cu=0; cu<w; ++cu)
685
686
      {
      complex<T> tmp(v*ku[cu]);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
687
      for (size_t cv=0; cv<w; ++cv)
688
        ptr[cv] += tmp*kv[cv];
689
      ptr+=jump;
690
691
692
      }
    }
} // end of parallel region
Martin Reinecke's avatar
Martin Reinecke committed
693
  }
694
695
696
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
697
template<typename T> pyarr_c<T> vis2grid(const Baselines<T> &baselines,
698
699
  const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &vis_)
Martin Reinecke's avatar
Martin Reinecke committed
700
  { return complex2hartley(vis2grid_c(baselines, gconf, idx_, vis_)); }
Martin Reinecke's avatar
updates    
Martin Reinecke committed
701

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
702
703
template<typename T> pyarr_c<complex<T>> ms2grid_c(
  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
704
  const pyarr<uint32_t> &idx_, const pyarr<complex<T>> &ms_)
Martin Reinecke's avatar
merge    
Martin Reinecke committed
705
706
707
708
709
710
  {
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
  checkArray(ms_, "ms", {nrows, nchan});
  checkArray(idx_, "idx", {0});
  size_t nvis = size_t(idx_.shape(0));
711
712
  auto ms = ms_.template unchecked<2>();
  auto idx = idx_.template unchecked<1>();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
713
714
715
716

  size_t nu=gconf.Nu(), nv=gconf.Nv();
  auto res = makeArray<complex<T>>({nu, nv});
  auto grid = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
717
718
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
719
720
721
722
723
724
  for (size_t i=0; i<nu*nv; ++i) grid[i] = 0.;
  T beta = gconf.Beta();
  size_t w = gconf.W();

#pragma omp parallel
{
725
  Helper<T> hlp(gconf, nullptr, grid);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
726
  T emb = exp(-2*beta);
727
  int jump = hlp.lineJump();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
728
729
730
731
732
733
734
  const T * RESTRICT ku = hlp.kernel.data();
  const T * RESTRICT kv = hlp.kernel.data()+w;

  // Loop over sampling points
#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
735
736
737
738
    auto tidx = idx(ipart);
    auto row = tidx/nchan;
    auto chan = tidx-row*nchan;
    UVW<T> coord = baselines.effectiveCoord(tidx);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
739
    hlp.prep(coord.u, coord.v);
740
    auto * RESTRICT ptr = hlp.p0w;
741
    auto v(ms(row,chan)*emb);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
742
743
744
745
746
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(v*ku[cu]);
      for (size_t cv=0; cv<w; ++cv)
        ptr[cv] += tmp*kv[cv];
747
      ptr+=jump;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
748
749
750
      }
    }
} // end of parallel region
Martin Reinecke's avatar
Martin Reinecke committed
751
  }
Martin Reinecke's avatar
merge    
Martin Reinecke committed
752
753
754
755
  return res;
  }

template<typename T> pyarr_c<T> ms2grid(const Baselines<T> &baselines,
756
757
  const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
  const pyarr<complex<T>> &ms_)
Martin Reinecke's avatar
merge    
Martin Reinecke committed
758
759
  { return complex2hartley(ms2grid_c(baselines, gconf, idx_, ms_)); }

760
761
template<typename T> pyarr_c<complex<T>> grid2vis_c(
  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
762
  const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_)
763
764
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
765
  checkArray(idx_, "idx", {0});
766
  auto grid = grid_.data();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
767
  checkArray(grid_, "grid", {nu, nv});
768
  size_t nvis = size_t(idx_.shape(0));
769
  auto idx = idx_.template unchecked<1>();
770

Martin Reinecke's avatar
merge    
Martin Reinecke committed
771
  auto res = makeArray<complex<T>>({nvis});
772
  auto vis = res.mutable_data();
Martin Reinecke's avatar
Martin Reinecke committed
773
774
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
775
776
  T beta = gconf.Beta();
  size_t w = gconf.W();
777
778
779
780

  // Loop over sampling points
#pragma omp parallel
{
781
  Helper<T> hlp(gconf, grid, nullptr);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
782
  T emb = exp(-2*beta);
783
  int jump = hlp.lineJump();
784
  const T * RESTRICT ku = hlp.kernel.data();
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
785
  const T * RESTRICT kv = hlp.kernel.data()+w;
786
787
788
789

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
790
    UVW<T> coord = baselines.effectiveCoord(idx(ipart));
Martin Reinecke's avatar
Martin Reinecke committed
791
    hlp.prep(coord.u, coord.v);
792
    complex<T> r = 0;
793
    const auto * RESTRICT ptr = hlp.p0r;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
794
    for (size_t cu=0; cu<w; ++cu)
795
796
      {
      complex<T> tmp(0);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
797
      for (size_t cv=0; cv<w; ++cv)
798
799
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
800
      ptr += jump;
801
802
803
804
      }
    vis[ipart] = r*emb;
    }
}
Martin Reinecke's avatar
Martin Reinecke committed
805
  }
806
807
808
  return res;
  }

Martin Reinecke's avatar
Martin Reinecke committed
809
template<typename T> pyarr_c<complex<T>> grid2vis(const Baselines<T> &baselines,
810
  const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
Martin Reinecke's avatar
Martin Reinecke committed
811
812
813
  const pyarr_c<T> &grid_)
  { return grid2vis_c(baselines, gconf, idx_, hartley2complex(grid_)); }

Martin Reinecke's avatar
merge    
Martin Reinecke committed
814
template<typename T> pyarr_c<complex<T>> grid2ms_c(const Baselines<T> &baselines,
815
  const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
Martin Reinecke's avatar
merge    
Martin Reinecke committed
816
817
  const pyarr_c<complex<T>> &grid_, py::object &ms_in)
  {
818
819
  auto nrows = baselines.Nrows();
  auto nchan = baselines.Nchannels();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
820
821
822
823
824
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
825
  auto idx = idx_.template unchecked<1>();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
826

827
828
  auto res = provideArray<complex<T>>(ms_in, {nrows, nchan});
  auto ms = res.template mutable_unchecked<2>();
Martin Reinecke's avatar
Martin Reinecke committed
829
830
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
831
832
833
834
835
836
  T beta = gconf.Beta();
  size_t w = gconf.W();

  // Loop over sampling points
#pragma omp parallel
{
837
  Helper<T> hlp(gconf, grid, nullptr);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
838
  T emb = exp(-2*beta);
839
  int jump = hlp.lineJump();
Martin Reinecke's avatar
merge    
Martin Reinecke committed
840
841
842
843
844
845
  const T * RESTRICT ku = hlp.kernel.data();
  const T * RESTRICT kv = hlp.kernel.data()+w;

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
846
847
848
849
    auto tidx = idx(ipart);
    auto row = tidx/nchan;
    auto chan = tidx-row*nchan;
    UVW<T> coord = baselines.effectiveCoord(tidx);
Martin Reinecke's avatar
merge    
Martin Reinecke committed
850
851
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
852
    const auto * RESTRICT ptr = hlp.p0r;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
853
854
855
856
857
858
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<w; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
859
      ptr += jump;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
860
      }
861
    ms(row,chan) += r*emb;
Martin Reinecke's avatar
merge    
Martin Reinecke committed
862
863
    }
}
Martin Reinecke's avatar
Martin Reinecke committed
864
  }
Martin Reinecke's avatar
merge    
Martin Reinecke committed
865
866
867
868
  return res;
  }

template<typename T> pyarr_c<complex<T>> grid2ms(const Baselines<T> &baselines,
869
  const GridderConfig<T> &gconf, const pyarr<uint32_t> &idx_,
Martin Reinecke's avatar
merge    
Martin Reinecke committed
870
871
872
  const pyarr_c<T> &grid_, py::object &ms_in)
  { return grid2ms_c(baselines, gconf, idx_, hartley2complex(grid_), ms_in); }

873
874
template<typename T> pyarr_c<complex<T>> apply_holo(
  const Baselines<T> &baselines, const GridderConfig<T> &gconf,
875
  const pyarr<uint32_t> &idx_, const pyarr_c<complex<T>> &grid_)
876
877
878
879
880
881
  {
  size_t nu=gconf.Nu(), nv=gconf.Nv();
  checkArray(idx_, "idx", {0});
  auto grid = grid_.data();
  checkArray(grid_, "grid", {nu, nv});
  size_t nvis = size_t(idx_.shape(0));
882
  auto idx = idx_.template unchecked<1>();
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902

  auto res = makeArray<complex<T>>({nu, nv});
  auto ogrid = res.mutable_data();
  {
  py::gil_scoped_release release;
  T beta = gconf.Beta();
  size_t w = gconf.W();

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

#pragma omp for schedule(guided,100)
  for (size_t ipart=0; ipart<nvis; ++ipart)
    {
903
    UVW<T> coord = baselines.effectiveCoord(idx(ipart));
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
    hlp.prep(coord.u, coord.v);
    complex<T> r = 0;
    const auto * RESTRICT ptr = hlp.p0r;
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(0);
      for (size_t cv=0; cv<w; ++cv)
        tmp += ptr[cv] * kv[cv];
      r += tmp*ku[cu];
      ptr += jump;
      }
    r*=emb*emb;
    auto * RESTRICT wptr = hlp.p0w;
    for (size_t cu=0; cu<w; ++cu)
      {
      complex<T> tmp(r*ku[cu]);
      for (size_t cv=0; cv<w; ++cv)
        wptr[cv] += tmp*kv[cv];
      wptr+=jump;
      }
    }
}
  }
  return res;
  }
Martin Reinecke's avatar
Martin Reinecke committed
929
template<typename T> pyarr_c<uint32_t> getIndices(const Baselines<T> &baselines,
930
931
  const GridderConfig<T> &gconf, const pyarr_c<bool> &flags_, int chbegin,
  int chend, T wmin, T wmax)
Martin Reinecke's avatar
updates    
Martin Reinecke committed
932
  {
933
934
  size_t nrow=baselines.Nrows(),
         nchan=baselines.Nchannels(),
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
935
         nsafe=gconf.Nsafe();
936
937
  if (chbegin<0) chbegin=0;
  if (chend<0) chend=nchan;
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
938
939
940
  myassert(chend>chbegin, "empty channel range selected");
  myassert(chend<=int(nchan), "chend too large");
  myassert(wmax>wmin, "empty w range selected");
Martin Reinecke's avatar
merge    
Martin Reinecke committed
941
  checkArray(flags_, "flags", {nrow, nchan});
942
943
944
945
  auto flags = flags_.data();
  constexpr int side=1<<logsquare;
  size_t nbu = (gconf.Nu()+1+side-1) >> logsquare,
         nbv = (gconf.Nv()+1+side-1) >> logsquare;
Martin Reinecke's avatar
Martin Reinecke committed
946
947
  vector<uint32_t> acc(nbu*nbv+1, 0);
  vector<uint32_t> tmp(nrow*(chend-chbegin));
Martin Reinecke's avatar
Martin Reinecke committed
948
949
  {
  py::gil_scoped_release release;
Martin Reinecke's avatar
Martin Reinecke committed
950
  for (size_t irow=0, idx=0; irow<nrow; ++irow)
Martin Reinecke's avatar
bug fix    
Martin Reinecke committed
951
    for (int ichan=chbegin; ichan<chend; ++ichan)
Martin Reinecke's avatar
Martin Reinecke committed
952
      if (!flags[irow*nchan+ichan])
953
        {
954
        auto uvw = baselines.effectiveCoord(irow, ichan);
955
956
        if ((uvw.w>=wmin) && (uvw.w<wmax))
          {
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
957
958
959
960
961
          T u, v;
          int iu0, iv0;
          gconf.getpix(uvw.u, uvw.v, u, v, iu0, iv0);
          iu0 = (iu0+nsafe)>>logsquare;
          iv0 = (iv0+nsafe)>>logsquare;
Martin Reinecke's avatar
Martin Reinecke committed
962
          ++acc[nbv*iu0 + iv0 + 1];
Martin Reinecke's avatar
bug fix    
Martin Reinecke committed
963
          tmp[idx++] = nbv*iu0 + iv0;
964
965
          }
        }
Martin Reinecke's avatar
Martin Reinecke committed
966
967