nifty_gridder.cc 8.24 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_gridder; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

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

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
24
#include "mr_util/bindings/pybind_utils.h"
Martin Reinecke's avatar
Martin Reinecke committed
25
26
27
28
29
#include "gridder_cxx.h"

using namespace std;
using namespace gridder;
using namespace mr;
30

Martin Reinecke's avatar
Martin Reinecke committed
31
32
33
34
35
36
37
38
namespace py = pybind11;

namespace {

auto None = py::none();

template<typename T> py::array ms2dirty_general2(const py::array &uvw_,
  const py::array &freq_, const py::array &ms_, const py::object &wgt_,
39
40
41
  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)
Martin Reinecke's avatar
Martin Reinecke committed
42
  {
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
43
44
45
  auto uvw = to_mav<double,2>(uvw_, false);
  auto freq = to_mav<double,1>(freq_, false);
  auto ms = to_mav<complex<T>,2>(ms_, false);
46
  auto wgt = get_optional_const_Pyarr<T>(wgt_, {ms.shape(0),ms.shape(1)});
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
47
  auto wgt2 = to_mav<T,2>(wgt, false);
48
  auto dirty = make_Pyarr<T>({npix_x,npix_y});
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
49
  auto dirty2 = to_mav<T,2>(dirty, true);
Martin Reinecke's avatar
Martin Reinecke committed
50
51
  {
  py::gil_scoped_release release;
52
53
  ms2dirty_general(uvw,freq,ms,wgt2,pixsize_x,pixsize_y,nu,nv,epsilon,
    do_wstacking,nthreads,dirty2,verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
54
  }
Martin Reinecke's avatar
Martin Reinecke committed
55
  return move(dirty);
Martin Reinecke's avatar
Martin Reinecke committed
56
57
58
  }
py::array Pyms2dirty_general(const py::array &uvw,
  const py::array &freq, const py::array &ms, const py::object &wgt,
59
60
61
  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)
Martin Reinecke's avatar
Martin Reinecke committed
62
  {
63
  if (isPyarr<complex<float>>(ms))
Martin Reinecke's avatar
Martin Reinecke committed
64
65
    return ms2dirty_general2<float>(uvw, freq, ms, wgt, npix_x, npix_y,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
66
  if (isPyarr<complex<double>>(ms))
Martin Reinecke's avatar
Martin Reinecke committed
67
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
    return ms2dirty_general2<double>(uvw, freq, ms, wgt, npix_x, npix_y,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
  MR_fail("type matching failed: 'ms' has neither type 'c8' nor 'c16'");
  }
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
)""";
py::array Pyms2dirty(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, double epsilon,
  bool do_wstacking, size_t nthreads, size_t verbosity)
  {
  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);
  }

template<typename T> py::array dirty2ms_general2(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)
  {
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
123
124
125
  auto uvw = to_mav<double,2>(uvw_, false);
  auto freq = to_mav<double,1>(freq_, false);
  auto dirty = to_mav<T,2>(dirty_, false);
126
  auto wgt = get_optional_const_Pyarr<T>(wgt_, {uvw.shape(0),freq.shape(0)});
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
127
  auto wgt2 = to_mav<T,2>(wgt, false);
128
  auto ms = make_Pyarr<complex<T>>({uvw.shape(0),freq.shape(0)});
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
129
  auto ms2 = to_mav<complex<T>,2>(ms, true);
Martin Reinecke's avatar
Martin Reinecke committed
130
131
  {
  py::gil_scoped_release release;
132
133
  dirty2ms_general(uvw,freq,dirty,wgt2,pixsize_x,pixsize_y,nu,nv,epsilon,
    do_wstacking,nthreads,ms2,verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
134
  }
Martin Reinecke's avatar
Martin Reinecke committed
135
  return move(ms);
Martin Reinecke's avatar
Martin Reinecke committed
136
137
138
139
140
141
  }
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)
  {
142
  if (isPyarr<float>(dirty))
Martin Reinecke's avatar
Martin Reinecke committed
143
144
    return dirty2ms_general2<float>(uvw, freq, dirty, wgt,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
145
  if (isPyarr<double>(dirty))
Martin Reinecke's avatar
Martin Reinecke committed
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
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    return dirty2ms_general2<double>(uvw, freq, dirty, wgt,
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
  MR_fail("type matching failed: 'dirty' has neither type 'f4' nor 'f8'");
  }
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.
)""";
py::array Pydirty2ms(const py::array &uvw,
  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)
  {
  return Pydirty2ms_general(uvw, freq, dirty, wgt, pixsize_x, pixsize_y,
    2*dirty.shape(0), 2*dirty.shape(1), epsilon, do_wstacking, nthreads,
    verbosity);
  }

} // unnamed namespace

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

  m.def("ms2dirty", &Pyms2dirty, ms2dirty_DS, "uvw"_a, "freq"_a, "ms"_a,
    "wgt"_a=None, "npix_x"_a, "npix_y"_a, "pixsize_x"_a, "pixsize_y"_a,
    "epsilon"_a, "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
  m.def("ms2dirty_general", &Pyms2dirty_general, "uvw"_a, "freq"_a, "ms"_a,
    "wgt"_a=None, "npix_x"_a, "npix_y"_a, "pixsize_x"_a, "pixsize_y"_a, "nu"_a, "nv"_a,
    "epsilon"_a, "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
  m.def("dirty2ms", &Pydirty2ms, dirty2ms_DS, "uvw"_a, "freq"_a, "dirty"_a,
    "wgt"_a=None, "pixsize_x"_a, "pixsize_y"_a, "epsilon"_a,
    "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
  m.def("dirty2ms_general", &Pydirty2ms_general, "uvw"_a, "freq"_a, "dirty"_a,
    "wgt"_a=None, "pixsize_x"_a, "pixsize_y"_a, "nu"_a, "nv"_a, "epsilon"_a,
    "do_wstacking"_a=false, "nthreads"_a=1, "verbosity"_a=0);
  }