wgridder.cc 7.78 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"
25
#include "python/gridder_cxx.h"
Martin Reinecke's avatar
Martin Reinecke committed
26

27
28
namespace mr {

Martin Reinecke's avatar
Martin Reinecke committed
29
namespace detail_pymodule_wgridder {
30

Martin Reinecke's avatar
Martin Reinecke committed
31
using namespace std;
32

Martin Reinecke's avatar
Martin Reinecke committed
33
34
35
36
namespace py = pybind11;

auto None = py::none();

37
template<typename T> py::array ms2dirty2(const py::array &uvw_,
Martin Reinecke's avatar
Martin Reinecke committed
38
  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
  ms2dirty(uvw,freq,ms,wgt2,pixsize_x,pixsize_y,nu,nv,epsilon,
53
    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
py::array Pyms2dirty(const py::array &uvw,
Martin Reinecke's avatar
Martin Reinecke committed
58
  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))
64
    return ms2dirty2<float>(uvw, freq, ms, wgt, npix_x, npix_y,
Martin Reinecke's avatar
Martin Reinecke committed
65
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
66
  if (isPyarr<complex<double>>(ms))
67
    return ms2dirty2<double>(uvw, freq, ms, wgt, npix_x, npix_y,
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
      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
90
91
92
93
94
95
96
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
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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
)""";

116
template<typename T> py::array dirty2ms2(const py::array &uvw_,
Martin Reinecke's avatar
Martin Reinecke committed
117
118
119
120
  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
121
122
123
  auto uvw = to_mav<double,2>(uvw_, false);
  auto freq = to_mav<double,1>(freq_, false);
  auto dirty = to_mav<T,2>(dirty_, false);
124
  auto wgt = get_optional_const_Pyarr<T>(wgt_, {uvw.shape(0),freq.shape(0)});
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
125
  auto wgt2 = to_mav<T,2>(wgt, false);
126
  auto ms = make_Pyarr<complex<T>>({uvw.shape(0),freq.shape(0)});
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
127
  auto ms2 = to_mav<complex<T>,2>(ms, true);
Martin Reinecke's avatar
Martin Reinecke committed
128
129
  {
  py::gil_scoped_release release;
130
  dirty2ms(uvw,freq,dirty,wgt2,pixsize_x,pixsize_y,nu,nv,epsilon,
131
    do_wstacking,nthreads,ms2,verbosity);
Martin Reinecke's avatar
Martin Reinecke committed
132
  }
Martin Reinecke's avatar
Martin Reinecke committed
133
  return move(ms);
Martin Reinecke's avatar
Martin Reinecke committed
134
  }
135
py::array Pydirty2ms(const py::array &uvw,
Martin Reinecke's avatar
Martin Reinecke committed
136
137
138
139
  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)
  {
140
  if (isPyarr<float>(dirty))
141
    return dirty2ms2<float>(uvw, freq, dirty, wgt,
Martin Reinecke's avatar
Martin Reinecke committed
142
      pixsize_x, pixsize_y, nu, nv, epsilon, do_wstacking, nthreads, verbosity);
143
  if (isPyarr<double>(dirty))
144
    return dirty2ms2<double>(uvw, freq, dirty, wgt,
Martin Reinecke's avatar
Martin Reinecke committed
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
      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
165
166
167
168
169
170
171
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
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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
191
void add_wgridder(py::module &msup)
Martin Reinecke's avatar
Martin Reinecke committed
192
193
  {
  using namespace pybind11::literals;
Martin Reinecke's avatar
Martin Reinecke committed
194
  auto m = msup.def_submodule("wgridder");
Martin Reinecke's avatar
Martin Reinecke committed
195

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

}

Martin Reinecke's avatar
Martin Reinecke committed
206
using detail_pymodule_wgridder::add_wgridder;
207
208

}