sht.cc 10.5 KB
Newer Older
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
/*
 *  This file is part of pyHealpix.
 *
 *  pyHealpix 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.
 *
 *  pyHealpix 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 pyHealpix; if not, write to the Free Software
 *  Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 *
 *  For more information about HEALPix, see http://healpix.sourceforge.net
 */

/*
 *  pyHealpix is being developed at the Max-Planck-Institut fuer Astrophysik
 */

/*
 *  Copyright (C) 2017-2020 Max-Planck-Society
 *  Author: Martin Reinecke
 */

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <vector>
#include <complex>

35
36
37
38
39
40
41
42
43
#include "mr_util/sharp/sharp.h"
#include "mr_util/sharp/sharp_geomhelpers.h"
#include "mr_util/sharp/sharp_almhelpers.h"
#include "mr_util/infra/string_utils.h"
#include "mr_util/infra/error_handling.h"
#include "mr_util/infra/mav.h"
#include "mr_util/math/fft.h"
#include "mr_util/math/constants.h"
#include "mr_util/bindings/pybind_utils.h"
44

45
46
namespace mr {

Martin Reinecke's avatar
Martin Reinecke committed
47
namespace detail_pymodule_sht {
48

49
50
51
52
using namespace std;

namespace py = pybind11;

53
using a_d = py::array_t<double>;
54
55
56
57
58
59
60
61
62
63
using a_d_c = py::array_t<double, py::array::c_style | py::array::forcecast>;
using a_c_c = py::array_t<complex<double>,
  py::array::c_style | py::array::forcecast>;

template<typename T> class py_sharpjob
  {
  private:
    unique_ptr<sharp_geom_info> ginfo;
    unique_ptr<sharp_alm_info> ainfo;
    int64_t lmax_, mmax_, npix_;
64
    int nthreads;
65
66

  public:
67
    py_sharpjob () : lmax_(0), mmax_(0), npix_(0), nthreads(1) {}
68
69
70
71
72
73
74

    string repr() const
      {
      return "<sharpjob_d: lmax=" + dataToString(lmax_) +
        ", mmax=" + dataToString(mmax_) + ", npix=", dataToString(npix_) +".>";
      }

75
76
    void set_nthreads(int64_t nthreads_)
      { nthreads = int(nthreads_); }
77
    void set_gauss_geometry(int64_t nrings, int64_t nphi)
78
79
80
81
82
      {
      MR_assert((nrings>0)&&(nphi>0),"bad grid dimensions");
      npix_=nrings*nphi;
      ginfo = sharp_make_gauss_geom_info (nrings, nphi, 0., 1, nphi);
      }
83
    void set_healpix_geometry(int64_t nside)
84
85
86
87
88
      {
      MR_assert(nside>0,"bad Nside value");
      npix_=12*nside*nside;
      ginfo = sharp_make_healpix_geom_info (nside, 1);
      }
89
    void set_fejer1_geometry(int64_t nrings, int64_t nphi)
90
91
92
93
      {
      MR_assert(nrings>0,"bad nrings value");
      MR_assert(nphi>0,"bad nphi value");
      npix_=nrings*nphi;
94
      ginfo = sharp_make_fejer1_geom_info (nrings, nphi, 0., 1, nphi);
95
      }
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
    void set_fejer2_geometry(int64_t nrings, int64_t nphi)
      {
      MR_assert(nrings>0,"bad nrings value");
      MR_assert(nphi>0,"bad nphi value");
      npix_=nrings*nphi;
      ginfo = sharp_make_fejer2_geom_info (nrings, nphi, 0., 1, nphi);
      }
    void set_cc_geometry(int64_t nrings, int64_t nphi)
      {
      MR_assert(nrings>0,"bad nrings value");
      MR_assert(nphi>0,"bad nphi value");
      npix_=nrings*nphi;
      ginfo = sharp_make_cc_geom_info (nrings, nphi, 0., 1, nphi);
      }
    void set_dh_geometry(int64_t nrings, int64_t nphi)
111
112
113
114
115
116
      {
      MR_assert(nrings>1,"bad nrings value");
      MR_assert(nphi>0,"bad nphi value");
      npix_=nrings*nphi;
      ginfo = sharp_make_dh_geom_info (nrings, nphi, 0., 1, nphi);
      }
117
118
119
120
121
122
123
    void set_mw_geometry(int64_t nrings, int64_t nphi)
      {
      MR_assert(nrings>0,"bad nrings value");
      MR_assert(nphi>0,"bad nphi value");
      npix_=nrings*nphi;
      ginfo = sharp_make_mw_geom_info (nrings, nphi, 0., 1, nphi);
      }
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
    void set_triangular_alm_info (int64_t lmax, int64_t mmax)
      {
      MR_assert(mmax>=0,"negative mmax");
      MR_assert(mmax<=lmax,"mmax must not be larger than lmax");
      lmax_=lmax; mmax_=mmax;
      ainfo = sharp_make_triangular_alm_info(lmax,mmax,1);
      }

    int64_t n_alm() const
      { return ((mmax_+1)*(mmax_+2))/2 + (mmax_+1)*(lmax_-mmax_); }

    a_d_c alm2map (const a_c_c &alm) const
      {
      MR_assert(npix_>0,"no map geometry specified");
      MR_assert (alm.size()==n_alm(),
        "incorrect size of a_lm array");
      a_d_c map(npix_);
      auto mr=map.mutable_unchecked<1>();
      auto ar=alm.unchecked<1>();
143
      sharp_alm2map(&ar[0], &mr[0], *ginfo, *ainfo, 0, nthreads);
144
145
146
147
148
149
150
151
152
      return map;
      }
    a_c_c alm2map_adjoint (const a_d_c &map) const
      {
      MR_assert(npix_>0,"no map geometry specified");
      MR_assert (map.size()==npix_,"incorrect size of map array");
      a_c_c alm(n_alm());
      auto mr=map.unchecked<1>();
      auto ar=alm.mutable_unchecked<1>();
153
      sharp_map2alm(&ar[0], &mr[0], *ginfo, *ainfo, 0, nthreads);
154
155
156
157
158
159
160
161
162
      return alm;
      }
    a_c_c map2alm (const a_d_c &map) const
      {
      MR_assert(npix_>0,"no map geometry specified");
      MR_assert (map.size()==npix_,"incorrect size of map array");
      a_c_c alm(n_alm());
      auto mr=map.unchecked<1>();
      auto ar=alm.mutable_unchecked<1>();
163
      sharp_map2alm(&ar[0], &mr[0], *ginfo, *ainfo, SHARP_USE_WEIGHTS, nthreads);
164
165
166
167
168
169
170
171
172
173
      return alm;
      }
    a_d_c alm2map_spin (const a_c_c &alm, int64_t spin) const
      {
      MR_assert(npix_>0,"no map geometry specified");
      auto ar=alm.unchecked<2>();
      MR_assert((ar.shape(0)==2)&&(ar.shape(1)==n_alm()),
        "incorrect size of a_lm array");
      a_d_c map(vector<size_t>{2,size_t(npix_)});
      auto mr=map.mutable_unchecked<2>();
174
      sharp_alm2map_spin(spin, &ar(0,0), &ar(1,0), &mr(0,0), &mr(1,0), *ginfo, *ainfo, 0, nthreads);
175
176
177
178
179
180
181
182
183
184
      return map;
      }
    a_c_c map2alm_spin (const a_d_c &map, int64_t spin) const
      {
      MR_assert(npix_>0,"no map geometry specified");
      auto mr=map.unchecked<2>();
      MR_assert ((mr.shape(0)==2)&&(mr.shape(1)==npix_),
        "incorrect size of map array");
      a_c_c alm(vector<size_t>{2,size_t(n_alm())});
      auto ar=alm.mutable_unchecked<2>();
185
      sharp_map2alm_spin(spin, &ar(0,0), &ar(1,0), &mr(0,0), &mr(1,0), *ginfo, *ainfo, SHARP_USE_WEIGHTS, nthreads);
186
187
188
189
      return alm;
      }
  };

Martin Reinecke's avatar
Martin Reinecke committed
190
const char *sht_DS = R"""(
191
192
193
Python interface for some of the libsharp functionality

Error conditions are reported by raising exceptions.
Martin Reinecke's avatar
Martin Reinecke committed
194
)""";
195

Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
196
197
void upsample_to_cc(const mav<double,2> &in, bool has_np, bool has_sp,
  mav<double,2> &out)
198
199
200
201
202
203
204
205
206
207
208
209
210
211
  {
  size_t ntheta_in = in.shape(0),
         ntheta_out = out.shape(0),
         nphi = in.shape(1);
  MR_assert(out.shape(1)==nphi, "phi dimensions must be equal");
  MR_assert((nphi&1)==0, "nphi must be even");
  size_t nrings_in = 2*ntheta_in-has_np-has_sp;
  size_t nrings_out = 2*ntheta_out-2;
  MR_assert(nrings_out>=nrings_in, "number of rings must increase");
  constexpr size_t delta=128;
  for (size_t js=0; js<nphi; js+=delta)
    {
    size_t je = min(js+delta, nphi);
    mav<double,2> tmp({nrings_out,je-js});
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
212
    fmav<double> ftmp(tmp);
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
213
    mav<double,2> tmp2(tmp.vdata(),{nrings_in, je-js}, true);
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
214
    fmav<double> ftmp2(tmp2);
215
216
217
    // enhance to "double sphere"
    if (has_np)
      for (size_t j=js; j<je; ++j)
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
218
        tmp2.v(0,j-js) = in(0,j);
219
220
    if (has_sp)
      for (size_t j=js; j<je; ++j)
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
221
        tmp2.v(ntheta_in-1,j-js) = in(ntheta_in-1,j);
222
223
224
225
    for (size_t i=has_np, i2=nrings_in-1; i+has_sp<ntheta_in; ++i,--i2)
      for (size_t j=js,j2=js+nphi/2; j<je; ++j,++j2)
        {
        if (j2>=nphi) j2-=nphi;
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
226
227
        tmp2.v(i,j-js) = in(i,j);
        tmp2.v(i2,j-js) = in(i,j2);
228
229
        }
    // FFT in theta direction
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
230
    r2r_fftpack(ftmp2,ftmp2,{0},true,true,1./nrings_in,0);
231
232
233
234
235
236
237
238
239
240
    if (!has_np)  // shift
      {
      double ang = -pi/nrings_in;
      for (size_t i=1; i<ntheta_in; ++i)
        {
        complex<double> rot(cos(i*ang),sin(i*ang));
        for (size_t j=js; j<je; ++j)
          {
          complex<double> ctmp(tmp2(2*i-1,j-js),tmp2(2*i,j-js));
          ctmp *= rot;
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
241
242
          tmp2.v(2*i-1,j-js) = ctmp.real();
          tmp2.v(2*i  ,j-js) = ctmp.imag();
243
244
245
246
247
248
          }
        }
      }
    // zero-padding
    for (size_t i=nrings_in; i<nrings_out; ++i)
      for (size_t j=js; j<je; ++j)
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
249
        tmp.v(i,j-js) = 0;
250
    // FFT back
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
251
    r2r_fftpack(ftmp,ftmp,{0},false,false,1.,0);
252
253
254
    // copy to output map
    for (size_t i=0; i<ntheta_out; ++i)
      for (size_t j=js; j<je; ++j)
Martin Reinecke's avatar
stage 2  
Martin Reinecke committed
255
        out.v(i,j) = tmp(i,j-js);
256
257
258
    }
  }

Martin Reinecke's avatar
Martin Reinecke committed
259
260
py::array py_upsample_to_cc(const py::array &in, size_t nrings_out, bool has_np,
  bool has_sp, py::object &out_)
261
  {
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
262
  auto in2 = to_mav<double,2>(in);
Martin Reinecke's avatar
Martin Reinecke committed
263
  auto out = get_optional_Pyarr<double>(out_, {nrings_out,size_t(in.shape(1))});
Martin Reinecke's avatar
stage 1  
Martin Reinecke committed
264
265
  auto out2 = to_mav<double,2>(out,true);
    MR_assert(out2.writable(),"x1");
266
  upsample_to_cc(in2, has_np, has_sp, out2);
Martin Reinecke's avatar
Martin Reinecke committed
267
  return move(out);
268
269
  }

Martin Reinecke's avatar
Martin Reinecke committed
270
void add_sht(py::module &msup)
271
272
  {
  using namespace pybind11::literals;
Martin Reinecke's avatar
Martin Reinecke committed
273
  auto m = msup.def_submodule("sht");
Martin Reinecke's avatar
Martin Reinecke committed
274
  m.doc() = sht_DS;
275

Martin Reinecke's avatar
Martin Reinecke committed
276
  py::class_<py_sharpjob<double>> (m, "sharpjob_d", py::module_local())
277
    .def(py::init<>())
278
    .def("set_nthreads", &py_sharpjob<double>::set_nthreads, "nthreads"_a)
279
    .def("set_gauss_geometry", &py_sharpjob<double>::set_gauss_geometry,
280
      "nrings"_a,"nphi"_a)
281
    .def("set_healpix_geometry", &py_sharpjob<double>::set_healpix_geometry,
282
      "nside"_a)
283
284
285
286
287
288
289
    .def("set_fejer1_geometry", &py_sharpjob<double>::set_fejer1_geometry,
      "nrings"_a, "nphi"_a)
    .def("set_fejer2_geometry", &py_sharpjob<double>::set_fejer2_geometry,
      "nrings"_a, "nphi"_a)
    .def("set_cc_geometry", &py_sharpjob<double>::set_cc_geometry,
      "nrings"_a, "nphi"_a)
    .def("set_dh_geometry", &py_sharpjob<double>::set_dh_geometry,
290
      "nrings"_a, "nphi"_a)
291
    .def("set_mw_geometry", &py_sharpjob<double>::set_mw_geometry,
292
      "nrings"_a, "nphi"_a)
293
294
295
296
297
298
299
300
    .def("set_triangular_alm_info",
      &py_sharpjob<double>::set_triangular_alm_info, "lmax"_a, "mmax"_a)
    .def("n_alm", &py_sharpjob<double>::n_alm)
    .def("alm2map", &py_sharpjob<double>::alm2map,"alm"_a)
    .def("alm2map_adjoint", &py_sharpjob<double>::alm2map_adjoint,"map"_a)
    .def("map2alm", &py_sharpjob<double>::map2alm,"map"_a)
    .def("alm2map_spin", &py_sharpjob<double>::alm2map_spin,"alm"_a,"spin"_a)
    .def("map2alm_spin", &py_sharpjob<double>::map2alm_spin,"map"_a,"spin"_a)
301
    .def("__repr__", &py_sharpjob<double>::repr);
Martin Reinecke's avatar
Martin Reinecke committed
302
303
  m.def("upsample_to_cc",&py_upsample_to_cc, "in"_a, "nrings_out"_a,
    "has_np"_a, "has_sp"_a, "out"_a=py::none());
304
  }
305
306
307

}

Martin Reinecke's avatar
Martin Reinecke committed
308
using detail_pymodule_sht::add_sht;
309
310
311

}