nifty_gridder.cc 7.82 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
#include "gridder_cxx.h"

27
28
29
30
namespace mr {

namespace detail_nifty_gridder {

Martin Reinecke's avatar
Martin Reinecke committed
31
32
33
using namespace std;
using namespace gridder;
using namespace mr;
34

Martin Reinecke's avatar
Martin Reinecke committed
35
36
37
38
namespace py = pybind11;

auto None = py::none();

39
template<typename T> py::array ms2dirty2(const py::array &uvw_,
Martin Reinecke's avatar
Martin Reinecke committed
40
  const py::array &freq_, const py::array &ms_, const py::object &wgt_,
41
42
43
  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
44
  {
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
45
46
47
  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);
48
  auto wgt = get_optional_const_Pyarr<T>(wgt_, {ms.shape(0),ms.shape(1)});
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
49
  auto wgt2 = to_mav<T,2>(wgt, false);
50
  auto dirty = make_Pyarr<T>({npix_x,npix_y});
Martin Reinecke's avatar
stage 2    
Martin Reinecke committed
51
  auto dirty2 = to_mav<T,2>(dirty, true);
Martin Reinecke's avatar
Martin Reinecke committed
52
53
  {
  py::gil_scoped_release release;
54
  ms2dirty(uvw,freq,ms,wgt2,pixsize_x,pixsize_y,nu,nv,epsilon,
55
    do_wstacking,nthreads,dirty2,verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
56
  }
Martin Reinecke's avatar
Martin Reinecke committed
57
  return move(dirty);
Martin Reinecke's avatar
Martin Reinecke committed
58
  }
59
py::array Pyms2dirty(const py::array &uvw,
Martin Reinecke's avatar
Martin Reinecke committed
60
  const py::array &freq, const py::array &ms, const py::object &wgt,
61
62
63
  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
64
  {
65
  if (isPyarr<complex<float>>(ms))
66
    return ms2dirty2<float>(uvw, freq, ms, wgt, npix_x, npix_y,
Martin Reinecke's avatar
Martin Reinecke committed
67
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
68
  if (isPyarr<complex<double>>(ms))
69
    return ms2dirty2<double>(uvw, freq, ms, wgt, npix_x, npix_y,
Martin Reinecke's avatar
Martin Reinecke committed
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
      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
92
93
94
95
96
97
98
nu, nv: int
    dimensions of the (oversampled) intermediate uv grid
    These values must be >= 1.2*the dimensions of the dirty image; tupical
    oversampling values lie between 1.5 and 2.
    Increasing the oversampling factor decreases the kernel support width
    required for the desired accuracy, so it typically reduces run-time; on the
    other hand, this will increase memory consumption. 
Martin Reinecke's avatar
Martin Reinecke committed
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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
)""";

118
template<typename T> py::array dirty2ms2(const py::array &uvw_,
Martin Reinecke's avatar
Martin Reinecke committed
119
120
121
122
  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
  dirty2ms(uvw,freq,dirty,wgt2,pixsize_x,pixsize_y,nu,nv,epsilon,
133
    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
py::array Pydirty2ms(const py::array &uvw,
Martin Reinecke's avatar
Martin Reinecke committed
138
139
140
141
  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))
143
    return dirty2ms2<float>(uvw, freq, dirty, wgt,
Martin Reinecke's avatar
Martin Reinecke committed
144
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
145
  if (isPyarr<double>(dirty))
146
    return dirty2ms2<double>(uvw, freq, dirty, wgt,
Martin Reinecke's avatar
Martin Reinecke committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
      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
167
168
169
170
171
172
173
nu, nv: int
    dimensions of the (oversampled) intermediate uv grid
    These values must be >= 1.2*the dimensions of the dirty image; tupical
    oversampling values lie between 1.5 and 2.
    Increasing the oversampling factor decreases the kernel support width
    required for the desired accuracy, so it typically reduces run-time; on the
    other hand, this will increase memory consumption. 
Martin Reinecke's avatar
Martin Reinecke committed
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
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.
)""";

193
void add_nifty_gridder(py::module &msup)
Martin Reinecke's avatar
Martin Reinecke committed
194
195
  {
  using namespace pybind11::literals;
Martin Reinecke's avatar
Martin Reinecke committed
196
  auto m = msup.def_submodule("wgridder");
Martin Reinecke's avatar
Martin Reinecke committed
197

198
  m.def("ms2dirty", &Pyms2dirty, "uvw"_a, "freq"_a, "ms"_a,
Martin Reinecke's avatar
Martin Reinecke committed
199
200
    "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);
201
  m.def("dirty2ms", &Pydirty2ms, "uvw"_a, "freq"_a, "dirty"_a,
Martin Reinecke's avatar
Martin Reinecke committed
202
203
204
    "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);
  }
205
206
207
208
209
210

}

using detail_nifty_gridder::add_nifty_gridder;

}