healpix.cc 12.7 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
/*
 *  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
 */

/*
Martin Reinecke's avatar
Martin Reinecke committed
26
 *  Copyright (C) 2017-2020 Max-Planck-Society
27
28
29
30
31
32
33
34
35
 *  Author: Martin Reinecke
 */

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

Martin Reinecke's avatar
Martin Reinecke committed
36
37
38
39
40
#include "mr_util/healpix/healpix_base.h"
#include "mr_util/math/constants.h"
#include "mr_util/infra/string_utils.h"
#include "mr_util/math/geom_utils.h"
#include "mr_util/bindings/pybind_utils.h"
41

42
43
namespace mr {

Martin Reinecke's avatar
Martin Reinecke committed
44
namespace detail_pymodule_healpix {
45

46
47
48
49
using namespace std;

namespace py = pybind11;

Martin Reinecke's avatar
Martin Reinecke committed
50
51
using shape_t = fmav_info::shape_t;

Martin Reinecke's avatar
Martin Reinecke committed
52
53
template<size_t nd1, size_t nd2> shape_t repl_dim(const shape_t &s,
  const array<size_t,nd1> &si, const array<size_t,nd2> &so)
54
  {
Martin Reinecke's avatar
Martin Reinecke committed
55
56
57
58
59
60
61
62
63
  MR_assert(s.size()+nd1,"too few input array dimensions");
  for (size_t i=0; i<nd1; ++i)
    MR_assert(si[i]==s[s.size()-nd1+i], "input dimension mismatch");
  shape_t snew(s.size()-nd1+nd2);
  for (size_t i=0; i<s.size()-nd1; ++i)
    snew[i]=s[i];
  for (size_t i=0; i<nd2; ++i)
    snew[i+s.size()-nd1] = so[i];
  return snew;
64
  }
Martin Reinecke's avatar
Martin Reinecke committed
65
66
67

template<typename T1, typename T2, size_t nd1, size_t nd2, typename Func>
  py::array doStuff(const py::array &ain, const array<size_t,nd1> &a1, const array<size_t,nd2> &a2, Func func)
68
  {
Martin Reinecke's avatar
Martin Reinecke committed
69
70
71
72
73
74
75
76
77
78
79
  auto in = to_fmav<T1>(ain);
  auto oshp = repl_dim(in.shape(), a1, a2);
  auto aout = make_Pyarr<T2>(oshp);
  auto out = to_fmav<T2>(aout,true);
  MavIter<T1,nd1+1> iin(in);
  MavIter<T2,nd2+1> iout(out);
  while (!iin.done())
    {
    func(iin, iout);
    iin.inc();iout.inc();
    }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
80
  return move(aout);
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
  }

class Pyhpbase
  {
  public:
    Healpix_Base2 base;

    Pyhpbase (int64_t nside, const string &scheme)
      : base (nside, RING, SET_NSIDE)
      {
      MR_assert((scheme=="RING")||(scheme=="NEST"),
        "unknown ordering scheme");
      if (scheme=="NEST")
        base.SetNside(nside,NEST);
      }
    string repr() const
      {
      return "<Healpix Base: Nside=" + dataToString(base.Nside()) +
        ", Scheme=" + ((base.Scheme()==RING) ? "RING" : "NEST") +".>";
      }

Martin Reinecke's avatar
Martin Reinecke committed
102
    py::array pix2ang (const py::array &pix) const
103
      {
Martin Reinecke's avatar
Martin Reinecke committed
104
      return doStuff<int64_t, double, 0, 1>(pix, {}, {2}, [this](const MavIter<int64_t,1> &iin, MavIter<double,2> &iout)
105
        {
Martin Reinecke's avatar
Martin Reinecke committed
106
        for (size_t i=0; i<iin.shape(0); ++i)
107
108
          {
          pointing ptg=base.pix2ang(iin(i));
Martin Reinecke's avatar
Martin Reinecke committed
109
          iout.v(i,0) = ptg.theta; iout.v(i,1) = ptg.phi;
110
          }
Martin Reinecke's avatar
Martin Reinecke committed
111
        });
112
      }
Martin Reinecke's avatar
Martin Reinecke committed
113
    py::array ang2pix (const py::array &ang) const
114
      {
Martin Reinecke's avatar
Martin Reinecke committed
115
      return doStuff<double, int64_t, 1, 0>(ang, {2}, {}, [this](const MavIter<double,2> &iin, MavIter<int64_t,1> &iout)
116
        {
Martin Reinecke's avatar
Martin Reinecke committed
117
118
        for (size_t i=0; i<iout.shape(0); ++i)
          iout.v(i)=base.ang2pix(pointing(iin(i,0),iin(i,1)));
Martin Reinecke's avatar
Martin Reinecke committed
119
        });
120
      }
Martin Reinecke's avatar
Martin Reinecke committed
121
    py::array pix2vec (const py::array &pix) const
122
      {
Martin Reinecke's avatar
Martin Reinecke committed
123
      return doStuff<int64_t, double, 0, 1>(pix, {}, {3}, [this](const MavIter<int64_t,1> &iin, MavIter<double,2> &iout)
124
        {
Martin Reinecke's avatar
Martin Reinecke committed
125
        for (size_t i=0; i<iin.shape(0); ++i)
126
127
          {
          vec3 v=base.pix2vec(iin(i));
Martin Reinecke's avatar
Martin Reinecke committed
128
          iout.v(i,0)=v.x; iout.v(i,1)=v.y; iout.v(i,2)=v.z;
129
          }
Martin Reinecke's avatar
Martin Reinecke committed
130
        });
131
      }
Martin Reinecke's avatar
Martin Reinecke committed
132
    py::array vec2pix (const py::array &vec) const
133
      {
Martin Reinecke's avatar
Martin Reinecke committed
134
      return doStuff<double, int64_t, 1, 0>(vec, {3}, {}, [this](const MavIter<double,2> &iin, MavIter<int64_t,1> &iout)
135
        {
Martin Reinecke's avatar
Martin Reinecke committed
136
137
        for (size_t i=0; i<iout.shape(0); ++i)
          iout.v(i)=base.vec2pix(vec3(iin(i,0),iin(i,1),iin(i,2)));
Martin Reinecke's avatar
Martin Reinecke committed
138
        });
139
      }
Martin Reinecke's avatar
Martin Reinecke committed
140
    py::array pix2xyf (const py::array &pix) const
141
      {
Martin Reinecke's avatar
Martin Reinecke committed
142
      return doStuff<int64_t, int64_t, 0, 1>(pix, {}, {3}, [this](const MavIter<int64_t,1> &iin, MavIter<int64_t,2> &iout)
143
        {
Martin Reinecke's avatar
Martin Reinecke committed
144
        for (size_t i=0; i<iin.shape(0); ++i)
145
146
147
          {
          int x,y,f;
          base.pix2xyf(iin(i),x,y,f);
Martin Reinecke's avatar
Martin Reinecke committed
148
          iout.v(i,0)=x; iout.v(i,1)=y; iout.v(i,2)=f;
149
          }
Martin Reinecke's avatar
Martin Reinecke committed
150
        });
151
      }
Martin Reinecke's avatar
Martin Reinecke committed
152
    py::array xyf2pix (const py::array &xyf) const
153
      {
Martin Reinecke's avatar
Martin Reinecke committed
154
      return doStuff<int64_t, int64_t, 1, 0>(xyf, {3}, {}, [this](const MavIter<int64_t,2> &iin, MavIter<int64_t,1> &iout)
155
        {
Martin Reinecke's avatar
Martin Reinecke committed
156
157
        for (size_t i=0; i<iout.shape(0); ++i)
          iout.v(i)=base.xyf2pix(iin(i,0),iin(i,1),iin(i,2));
Martin Reinecke's avatar
Martin Reinecke committed
158
        });
159
      }
Martin Reinecke's avatar
Martin Reinecke committed
160
    py::array neighbors (const py::array &pix) const
161
      {
Martin Reinecke's avatar
Martin Reinecke committed
162
      return doStuff<int64_t, int64_t, 0, 1>(pix, {}, {8}, [this](const MavIter<int64_t,1> &iin, MavIter<int64_t,2> &iout)
163
        {
Martin Reinecke's avatar
Martin Reinecke committed
164
        for (size_t i=0; i<iin.shape(0); ++i)
165
166
167
          {
          array<int64_t,8> res;
          base.neighbors(iin(i),res);
Martin Reinecke's avatar
Martin Reinecke committed
168
          for (size_t j=0; j<8; ++j) iout.v(i,j)=res[j];
169
          }
Martin Reinecke's avatar
Martin Reinecke committed
170
        });
171
      }
Martin Reinecke's avatar
Martin Reinecke committed
172
    py::array ring2nest (const py::array &ring) const
173
      {
Martin Reinecke's avatar
Martin Reinecke committed
174
      return doStuff<int64_t, int64_t, 0, 0>(ring, {}, {}, [this](const MavIter<int64_t,1> &iin, MavIter<int64_t,1> &iout)
175
        {
Martin Reinecke's avatar
Martin Reinecke committed
176
177
        for (size_t i=0; i<iin.shape(0); ++i)
          iout.v(i)=base.ring2nest(iin(i));
Martin Reinecke's avatar
Martin Reinecke committed
178
        });
179
      }
Martin Reinecke's avatar
Martin Reinecke committed
180
    py::array nest2ring (const py::array &nest) const
181
      {
Martin Reinecke's avatar
Martin Reinecke committed
182
      return doStuff<int64_t, int64_t, 0, 0>(nest, {}, {}, [this](const MavIter<int64_t,1> &iin, MavIter<int64_t,1> &iout)
183
        {
Martin Reinecke's avatar
Martin Reinecke committed
184
185
        for (size_t i=0; i<iin.shape(0); ++i)
          iout.v(i)=base.nest2ring(iin(i));
Martin Reinecke's avatar
Martin Reinecke committed
186
        });
187
      }
Martin Reinecke's avatar
Martin Reinecke committed
188
    py::array query_disc(const py::array &ptg, double radius) const
189
190
191
192
      {
      MR_assert((ptg.ndim()==1)&&(ptg.shape(0)==2),
        "ptg must be a 1D array with 2 values");
      rangeset<int64_t> pixset;
Martin Reinecke's avatar
Martin Reinecke committed
193
194
195
      auto ptg2 = to_mav<double,1>(ptg);
      base.query_disc(pointing(ptg2(0),ptg2(1)), radius, pixset);
      auto res = make_Pyarr<int64_t>(shape_t({pixset.nranges(),2}));
196
197
198
199
200
201
      auto oref=res.mutable_unchecked<2>();
      for (size_t i=0; i<pixset.nranges(); ++i)
        {
        oref(i,0)=pixset.ivbegin(i);
        oref(i,1)=pixset.ivend(i);
        }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
202
      return move(res);
203
204
205
      }
  };

Martin Reinecke's avatar
Martin Reinecke committed
206
py::array ang2vec (const py::array &ang)
207
  {
Martin Reinecke's avatar
Martin Reinecke committed
208
  return doStuff<double, double, 1, 1>(ang, {2}, {3}, [](const MavIter<double,2> &iin, MavIter<double,2> &iout)
209
    {
Martin Reinecke's avatar
Martin Reinecke committed
210
    for (size_t i=0; i<iin.shape(0); ++i)
211
212
      {
      vec3 v (pointing(iin(i,0),iin(i,1)));
Martin Reinecke's avatar
Martin Reinecke committed
213
      iout.v(i,0)=v.x; iout.v(i,1)=v.y; iout.v(i,2)=v.z;
214
      }
Martin Reinecke's avatar
Martin Reinecke committed
215
    });
216
  }
Martin Reinecke's avatar
Martin Reinecke committed
217
py::array vec2ang (const py::array &vec)
218
  {
Martin Reinecke's avatar
Martin Reinecke committed
219
  return doStuff<double, double, 1, 1>(vec, {3}, {2}, [](const MavIter<double,2> &iin, MavIter<double,2> &iout)
220
    {
Martin Reinecke's avatar
Martin Reinecke committed
221
    for (size_t i=0; i<iin.shape(0); ++i)
222
223
      {
      pointing ptg (vec3(iin(i,0),iin(i,1),iin(i,2)));
Martin Reinecke's avatar
Martin Reinecke committed
224
      iout.v(i,0)=ptg.theta; iout.v(i,1)=ptg.phi;
225
      }
Martin Reinecke's avatar
Martin Reinecke committed
226
    });
227
  }
Martin Reinecke's avatar
Martin Reinecke committed
228
py::array local_v_angle (const py::array &v1, const py::array &v2)
229
  {
Martin Reinecke's avatar
Martin Reinecke committed
230
231
232
  auto v12 = to_fmav<double>(v1);
  auto v22 = to_fmav<double>(v2);
  MR_assert(v12.shape()==v22.shape());
Martin Reinecke's avatar
Martin Reinecke committed
233
  auto angle = make_Pyarr<double>(repl_dim<1,0>(v12.shape(),{3},{}));
Martin Reinecke's avatar
Martin Reinecke committed
234
235
236
  auto angle2 = to_fmav<double>(angle,true);
  MavIter<double,2> ii1(v12), ii2(v22);
  MavIter<double,1> iout(angle2);
237
238
  while (!iout.done())
    {
Martin Reinecke's avatar
Martin Reinecke committed
239
240
241
242
    for (size_t i=0; i<iout.shape(0); ++i)
      iout.v(i)=v_angle(vec3(ii1(i,0),ii1(i,1),ii1(i,2)),
                        vec3(ii2(i,0),ii2(i,1),ii2(i,2)));
    ii1.inc();ii2.inc();iout.inc();
243
    }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
244
  return move(angle);
245
246
  }

Martin Reinecke's avatar
Martin Reinecke committed
247
const char *healpix_DS = R"""(
248
249
250
251
252
253
254
255
256
257
Python interface for some of the HEALPix C++ functionality

All angles are interpreted as radians.
The theta coordinate is measured as co-latitude, ranging from 0 (North Pole)
to pi (South Pole).

All 3-vectors returned by the functions are normalized.
However, 3-vectors provided as input to the functions need not be normalized.

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

Martin Reinecke's avatar
Martin Reinecke committed
260
const char *order_DS = R"""(
261
262
Returns the ORDER parameter of the pixelisation.
If Nside is a power of 2, this is log_2(Nside), otherwise it is -1.
Martin Reinecke's avatar
Martin Reinecke committed
263
)""";
264

Martin Reinecke's avatar
Martin Reinecke committed
265
const char *nside_DS = R"""(
266
Returns the Nside parameter of the pixelisation.
Martin Reinecke's avatar
Martin Reinecke committed
267
)""";
268

Martin Reinecke's avatar
Martin Reinecke committed
269
const char *npix_DS = R"""(
270
Returns the total number of pixels of the pixelisation.
Martin Reinecke's avatar
Martin Reinecke committed
271
)""";
272

Martin Reinecke's avatar
Martin Reinecke committed
273
const char *scheme_DS = R"""(
274
275
Returns a string representation of the pixelisation's ordering scheme
("RING" or "NEST").
Martin Reinecke's avatar
Martin Reinecke committed
276
)""";
277

Martin Reinecke's avatar
Martin Reinecke committed
278
const char *pix_area_DS = R"""(
279
Returns the area (in steradian) of a single pixel.
Martin Reinecke's avatar
Martin Reinecke committed
280
)""";
281

Martin Reinecke's avatar
Martin Reinecke committed
282
const char *max_pixrad_DS = R"""(
283
284
Returns the maximum angular distance (in radian) between a pixel center
and its corners for this pixelisation.
Martin Reinecke's avatar
Martin Reinecke committed
285
)""";
286

Martin Reinecke's avatar
Martin Reinecke committed
287
const char *pix2ang_DS = R"""(
288
289
290
Returns a (co-latitude, longitude) tuple for each value in pix.
The result array has the same shape as pix, with an added last dimension
of size 2.
Martin Reinecke's avatar
Martin Reinecke committed
291
)""";
292

Martin Reinecke's avatar
Martin Reinecke committed
293
const char *ang2pix_DS = R"""(
294
295
296
Returns the index of the containing pixel for every (co-latitude, longitude)
tuple in ang. ang must have a last dimension of size 2; the result array
has the same shape as ang, except that ang's last dimension is removed.
Martin Reinecke's avatar
Martin Reinecke committed
297
)""";
298

Martin Reinecke's avatar
Martin Reinecke committed
299
const char *pix2vec_DS = R"""(
300
301
302
Returns a normalized 3-vector for each value in pix.
The result array has the same shape as pix, with an added last dimension
of size 3.
Martin Reinecke's avatar
Martin Reinecke committed
303
)""";
304

Martin Reinecke's avatar
Martin Reinecke committed
305
const char *vec2pix_DS = R"""(
306
307
308
Returns the index of the containing pixel for every 3-vector in vec.
vec must have a last dimension of size 3; the result array has the same shape as
vec, except that vec's last dimension is removed.
Martin Reinecke's avatar
Martin Reinecke committed
309
)""";
310

Martin Reinecke's avatar
Martin Reinecke committed
311
const char *ring2nest_DS = R"""(
312
313
Returns the pixel index in NEST scheme for every entry of ring.
The result array has the same shape as ring.
Martin Reinecke's avatar
Martin Reinecke committed
314
)""";
315

Martin Reinecke's avatar
Martin Reinecke committed
316
const char *nest2ring_DS = R"""(
317
318
Returns the pixel index in RING scheme for every entry of nest.
The result array has the same shape as nest.
Martin Reinecke's avatar
Martin Reinecke committed
319
)""";
320

Martin Reinecke's avatar
Martin Reinecke committed
321
const char *query_disc_DS = R"""(
322
323
324
325
Returns a range set of all pixels whose centers fall within "radius" of "ptg".
"ptg" must be a single (co-latitude, longitude) tuple. The result is a 2D array
with last dimension 2; the pixels lying inside the disc are
[res[0,0] .. res[0,1]); [res[1,0] .. res[1,1]) etc.
Martin Reinecke's avatar
Martin Reinecke committed
326
)""";
327

Martin Reinecke's avatar
Martin Reinecke committed
328
const char *ang2vec_DS = R"""(
329
330
331
Returns a normalized 3-vector for every (co-latitude, longitude)
tuple in ang. ang must have a last dimension of size 2; the result array
has the same shape as ang, except that its last dimension is 3 instead of 2.
Martin Reinecke's avatar
Martin Reinecke committed
332
)""";
333

Martin Reinecke's avatar
Martin Reinecke committed
334
const char *vec2ang_DS = R"""(
335
336
337
Returns a (co-latitude, longitude) tuple for every 3-vector in vec.
vec must have a last dimension of size 3; the result array has the same shape as
vec, except that its last dimension is 2 instead of 3.
Martin Reinecke's avatar
Martin Reinecke committed
338
)""";
339

Martin Reinecke's avatar
Martin Reinecke committed
340
const char *v_angle_DS = R"""(
341
342
343
344
Returns the angles between the 3-vectors in v1 and v2. The input arrays must
have identical shapes. The result array has the same shape as v1 or v2, except
that their last dimension is removed.
The employed algorithm is highly accurate, even for angles close to 0 or pi.
Martin Reinecke's avatar
Martin Reinecke committed
345
)""";
346

Martin Reinecke's avatar
Martin Reinecke committed
347
void add_healpix(py::module &msup)
348
349
  {
  using namespace pybind11::literals;
Martin Reinecke's avatar
Martin Reinecke committed
350
  auto m = msup.def_submodule("healpix");
Martin Reinecke's avatar
Martin Reinecke committed
351
  m.doc() = healpix_DS;
352

Martin Reinecke's avatar
Martin Reinecke committed
353
  py::class_<Pyhpbase> (m, "Healpix_Base", py::module_local())
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
    .def(py::init<int,const string &>(),"nside"_a,"scheme"_a)
    .def("order", [](Pyhpbase &self)
      { return self.base.Order(); }, order_DS)
    .def("nside", [](Pyhpbase &self)
      { return self.base.Nside(); }, nside_DS)
    .def("npix", [](Pyhpbase &self)
      { return self.base.Npix(); }, npix_DS)
    .def("scheme", [](Pyhpbase &self)
      { return self.base.Scheme(); }, scheme_DS)
    .def("pix_area", [](Pyhpbase &self)
      { return 4*pi/self.base.Npix(); }, pix_area_DS)
    .def("max_pixrad", [](Pyhpbase &self)
      { return self.base.max_pixrad(); }, max_pixrad_DS)
    .def("pix2ang", &Pyhpbase::pix2ang, pix2ang_DS, "pix"_a)
    .def("ang2pix", &Pyhpbase::ang2pix, ang2pix_DS, "ang"_a)
    .def("pix2vec", &Pyhpbase::pix2vec, pix2vec_DS, "pix"_a)
    .def("vec2pix", &Pyhpbase::vec2pix, vec2pix_DS, "vec"_a)
    .def("pix2xyf", &Pyhpbase::pix2xyf, "pix"_a)
    .def("xyf2pix", &Pyhpbase::xyf2pix, "xyf"_a)
    .def("neighbors", &Pyhpbase::neighbors,"pix"_a)
    .def("ring2nest", &Pyhpbase::ring2nest, ring2nest_DS, "ring"_a)
    .def("nest2ring", &Pyhpbase::nest2ring, nest2ring_DS, "nest"_a)
    .def("query_disc", &Pyhpbase::query_disc, query_disc_DS, "ptg"_a,"radius"_a)
    .def("__repr__", &Pyhpbase::repr)
    ;

  m.def("ang2vec",&ang2vec, ang2vec_DS, "ang"_a);
  m.def("vec2ang",&vec2ang, vec2ang_DS, "vec"_a);
  m.def("v_angle",&local_v_angle, v_angle_DS, "v1"_a, "v2"_a);
  }
384
385
386

}

Martin Reinecke's avatar
Martin Reinecke committed
387
using detail_pymodule_healpix::add_healpix;
388
389

}