File: pyradler.cc

package info (click to toggle)
wsclean 3.6-3
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 16,296 kB
  • sloc: cpp: 129,246; python: 22,066; sh: 360; ansic: 230; makefile: 185
file content (251 lines) | stat: -rw-r--r-- 11,825 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
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
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
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
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
// SPDX-License-Identifier: LGPL-3.0-only

#include "radler.h"

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/iostream.h>

#include <aocommon/image.h>
#include <aocommon/logger.h>

#include "work_table.h"
#include "utils/load_image_accessor.h"
#include "utils/load_and_store_image_accessor.h"

namespace py = pybind11;

void init_radler(py::module& m) {
  py::class_<radler::Radler>(m, "Radler",
                             "Radio Astronomical Deconvolution Library core "
                             "class for controlling and excuting the "
                             "deconvolution of radio astronmical images.")
      .def(py::init([](const radler::Settings& settings,
                       radler::WorkTable& work_table, double beam_size) {
             if (settings.thread_count == 0) {
               throw std::runtime_error("Number of threads should be > 0.");
             }
             return std::make_unique<radler::Radler>(
                 settings,
                 std::make_unique<radler::WorkTable>(std::move(work_table)),
                 beam_size);
           }),
           R"pbdoc(
        Constructs a Radler object from a radler.Settings object and a
        radler.WorkTable object.
        Residual and model images in the WorkTable entries are updated in-place
        during Radler.perform() calls.
        Note that the provided WorkTable is moved to the Radler object, hence
        becomes invalid for the Python client after this call.

        Internally, views on the provided images are used. Hence, keep the
        images (numpy arrays) that are attached to the WorkTable alive during
        the lifetime of the instantiated Radler object.

        Parameters
        ----------
        settings: radler.Settings
            Settings object.
        work_table: radler.WorkTable
            Table containing the work instructions for radler.
            Since this WorkTable is moved to the Radler object, it becomes
            invalid for the Python client after this call.
        beam_size: double
            Beam size in [rad]. The beam size is typically calculated from the
            longest baseline, and used as initial value when fitting the
            (accurate) beam.
        )pbdoc",
           py::arg("settings"), py::arg("work_table"), py::arg("beam_size"))
      .def(
          py::init([](const radler::Settings& settings, py::array_t<float>& psf,
                      py::array_t<float>& residual, py::array_t<float>& model,
                      double beam_size, size_t n_deconvolution_groups,
                      py::array_t<double>& frequencies,
                      py::array_t<double>& weights,
                      aocommon::PolarizationEnum polarization) {
            if (settings.thread_count == 0) {
              throw std::runtime_error("Number of threads should be > 0.");
            }

            if (psf.ndim() != residual.ndim() || psf.ndim() != model.ndim()) {
              throw std::runtime_error(
                  "Provided arrays should have equal dimension count.");
            }

            for (pybind11::ssize_t d = 0; d < psf.ndim(); ++d) {
              if (residual.shape(d) != psf.shape(d) ||
                  model.shape(d) != psf.shape(d)) {
                throw std::runtime_error(
                    "Provided arrays should have equal shape.");
              }
            }

            if (psf.ndim() != 2 && psf.ndim() != 3) {
              throw std::runtime_error(
                  "Provided arrays should have 2 or 3 dimensions.");
            }

            const pybind11::ssize_t height = psf.shape(psf.ndim() - 2);
            const pybind11::ssize_t width = psf.shape(psf.ndim() - 1);
            const pybind11::ssize_t n_images =
                (psf.ndim() == 2) ? 1 : psf.shape(0);

            if (settings.spectral_fitting.mode !=
                    schaapcommon::fitters::SpectralFittingMode::kNoFitting &&
                frequencies.size() == 0) {
              throw std::runtime_error(
                  "Frequencies are required when spectral fitting is enabled.");
            }

            if (frequencies.size() > 0 &&
                (frequencies.ndim() != 2 || frequencies.shape(0) != n_images ||
                 frequencies.shape(1) != 2)) {
              throw std::runtime_error(
                  "Provided frequencies should have shape (n_images, 2).");
            }

            if (weights.size() > 0 &&
                (weights.ndim() != 1 || weights.shape(0) != n_images)) {
              throw std::runtime_error(
                  "Provided weights should have shape (n_images).");
            }

            // Create a WorkTable with an entry for each image.
            auto table = std::make_unique<radler::WorkTable>(
                std::vector<radler::PsfOffset>{}, n_images,
                n_deconvolution_groups);
            for (pybind11::ssize_t i = 0; i < n_images; ++i) {
              // This loop supports both 2-D arrays with a single image,
              // and 3-D arrays with multiple images.
              float* psf_data =
                  (i == 0) ? psf.mutable_data() : psf.mutable_data(i, 0, 0);
              float* residual_data = (i == 0) ? residual.mutable_data()
                                              : residual.mutable_data(i, 0, 0);
              float* model_data =
                  (i == 0) ? model.mutable_data() : model.mutable_data(i, 0, 0);

              aocommon::Image psf_image(psf_data, width, height);
              aocommon::Image residual_image(residual_data, width, height);
              aocommon::Image model_image(model_data, width, height);

              auto entry = std::make_unique<radler::WorkTableEntry>();
              if (frequencies.size() > 0) {
                auto unchecked = frequencies.unchecked<2>();
                entry->band_start_frequency = unchecked(i, 0);
                entry->band_end_frequency = unchecked(i, 1);
              }
              entry->polarization = polarization;
              entry->original_channel_index = i;
              entry->image_weight =
                  (weights.size() > 0) ? weights.unchecked<1>()(i) : 1.0;
              entry->psf_accessors.emplace_back(
                  std::make_unique<radler::utils::LoadOnlyImageAccessor>(
                      psf_image));
              entry->residual_accessor =
                  std::make_unique<radler::utils::LoadAndStoreImageAccessor>(
                      residual_image);
              entry->model_accessor =
                  std::make_unique<radler::utils::LoadAndStoreImageAccessor>(
                      model_image);
              table->AddEntry(std::move(entry));
            }
            return std::make_unique<radler::Radler>(settings, std::move(table),
                                                    beam_size);
          }),
          R"pbdoc(
        Constructor expecting a radler::Settings object along with PSF,
        residual, and model images as a numpy array of dtype=np.float32.
        The numpy arrays can be 2-D arrays containing a single image or 3-D
        arrays that contain images for different frequency bands.
        All arrays should have equal shapes.

        The residual and model images are updated in-place in Radler.perform() calls.

        Internally, views on the provided images are used. Hence, keep the images
        (numpy arrays) alive during the lifetime of the instantiated Radler object.

        Parameters
        ----------
        settings: radler.Settings
            Settings object
        psf: np.array
            2-D or 3-D numpy array of the PSF of type np.float32.
            Its lifetime should exceed the lifetime of the Radler object.
        residual: np.array
            2-D or 3-D numpy array of the residual image of type np.float32.
            Its lifetime should exceed the lifetime of the Radler object.
        model: np.array
            2-D or 3-D numpy array of the model image of type np.float32.
            Its lifetime should exceed the lifetime of the Radler object.
        beam_size: double
            Beam size in [rad]. The beam size is typically calculated from the longest
            baseline, and used as initial value when fitting the (accurate) beam.
        n_deconvolution_groups: int, optional
            The number of deconvolution groups. If it is less than the number of
            images, Radler joinedly deconvolves images by averaging them before
            deconvolution and interpolating them after deconvolution.
            If the value is zero, or larger than the number of images,
            Radler deconvolves all images separately.
        frequencies: np.array, optional
            2-D array with the start and end frequencies for each image.
            You may only omit this argument when spectral fitting is disabled in
            the supplied settings.
        weights: np.array, optional
            1-D array with the relative weight of each image. Radler uses these
            weights when joinedly deconvolving images.
            If omitted, all weight values become 1.0.
        polarization: radler.Polarization, optional
            Polarization of the input images. Default rd.Polarization.stokes_i.
       )pbdoc",
          // noconvert() is necessary for PSF, residual and model, since
          // (lists of) images can be large. It is not necessary for
          // frequencies and weights, since those lists are small.
          py::arg("settings"), py::arg("psf").noconvert(),
          py::arg("residual").noconvert(), py::arg("model").noconvert(),
          py::arg("beam_size"), py::arg("n_deconvolution_groups") = 0,
          py::arg("frequencies") = py::array_t<double>(),
          py::arg("weights") = py::array_t<double>(),
          py::arg("polarization") = aocommon::PolarizationEnum::StokesI)
      // Perform needs a lambda-expression, as the boolean input is an
      // immutable type.
      .def(
          "perform",
          [](radler::Radler& self, size_t major_iteration_number) {
            py::scoped_ostream_redirect stream(
                std::cout, py::module_::import("sys").attr("stdout"));
            // The scoped_ostream_redirect may grab the Global Interpreter
            // Lock (GIL) from a different thread, which causes hanging
            // if the GIL is not released beforehand
            py::gil_scoped_release release;
            bool reached_major_threshold = false;
            self.Perform(reached_major_threshold, major_iteration_number);
            return reached_major_threshold;
          },
          R"pbdoc(
        Execute deconvolution minor loop.

        Parameters
        ----------
        major_iteration_number: int
            How many major iterations (calls to @c Perform()) were performed
            so far.

        Returns
        -------
        bool:
            Indicates whether another major iteration should be run. If
            @c True, the caller should do a new prediction-gridding iteration
            to calculate a new residual image, after which the @c perform()
            function should be called again. If @c False, the algorithm is
            finished and the caller can do its last prediction-gridding round.
       )pbdoc",
          py::arg("major_iteration_number"))
      .def_property_readonly("iteration_number",
                             &radler::Radler::IterationNumber, R"pbdoc(
        Return minor loop iteration number of the underlying DeconvolutionAlgorithm.
       )pbdoc")
      .def_property_readonly("component_list",
                             &radler::Radler::GetComponentList, R"pbdoc(
        Return the component list (only stored when specified in settings).
       )pbdoc");
}