interpol_ng.cc 8.01 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/*
 *  Copyright (C) 2020 Max-Planck-Society
 *  Author: Martin Reinecke
 */

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <vector>
#include <complex>
#include "mr_util/math/constants.h"
#include "mr_util/math/gl_integrator.h"
#include "mr_util/math/es_kernel.h"
#include "mr_util/infra/mav.h"
#include "mr_util/sharp/sharp.h"
#include "mr_util/sharp/sharp_almhelpers.h"
#include "mr_util/sharp/sharp_geomhelpers.h"
#include "alm.h"
#include "mr_util/math/fft.h"
#include "mr_util/bindings/pybind_utils.h"
Martin Reinecke's avatar
fix    
Martin Reinecke committed
20

21
22
23
24
25
26
27
using namespace std;
using namespace mr;

namespace py = pybind11;

namespace {

Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
28
29
constexpr double ofmin=1.5;

30
31
32
template<typename T> class Interpolator
  {
  protected:
Martin Reinecke's avatar
Martin Reinecke committed
33
    size_t lmax, kmax, nphi0, ntheta0, nphi, ntheta;
Martin Reinecke's avatar
fix    
Martin Reinecke committed
34
35
    double ofactor;
    size_t supp;
36
37
38
39
40
41
    ES_Kernel kernel;
    mav<T,3> cube; // the data cube (theta, phi, 2*mbeam+1[, IQU])

    void correct(mav<T,2> &arr)
      {
      mav<T,2> tmp({nphi,nphi});
Martin Reinecke's avatar
Martin Reinecke committed
42
      auto tmp0=tmp.template subarray<2>({0,0},{nphi0, nphi0});
Martin Reinecke's avatar
Martin Reinecke committed
43
44
45
46
      fmav<T> ftmp0(tmp0);
      for (size_t i=0; i<ntheta0; ++i)
        for (size_t j=0; j<nphi0; ++j)
          tmp0.v(i,j) = arr(i,j);
47
      // extend to second half
Martin Reinecke's avatar
Martin Reinecke committed
48
49
      for (size_t i=1, i2=2*ntheta0-3; i+1<ntheta0; ++i,--i2)
        for (size_t j=0,j2=nphi0/2; j<nphi0; ++j,++j2)
50
          {
Martin Reinecke's avatar
Martin Reinecke committed
51
52
          if (j2>=nphi0) j2-=nphi0;
          tmp0.v(i2,j) = arr(i,j2);
53
          }
Martin Reinecke's avatar
Martin Reinecke committed
54
      // FFT to frequency domain on minimal grid
Martin Reinecke's avatar
Martin Reinecke committed
55
      r2r_fftpack(ftmp0,ftmp0,{0,1},true,true,1./(nphi0*nphi0),0);
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
56
      auto fct = kernel.correction_factors(nphi, nphi0/2+1, 0);
Martin Reinecke's avatar
Martin Reinecke committed
57
58
59
      for (size_t i=0; i<nphi0; ++i)
        for (size_t j=0; j<nphi0; ++j)
          tmp0.v(i,j) *= fct[(i+1)/2] * fct[(j+1)/2];
Martin Reinecke's avatar
Martin Reinecke committed
60
      auto tmp1=tmp.template subarray<2>({0,0},{nphi, nphi0});
Martin Reinecke's avatar
Martin Reinecke committed
61
      fmav<T> ftmp1(tmp1);
Martin Reinecke's avatar
Martin Reinecke committed
62
      // zero-padded FFT in theta direction
Martin Reinecke's avatar
Martin Reinecke committed
63
      r2r_fftpack(ftmp1,ftmp1,{0},false,false,1.,0);
Martin Reinecke's avatar
Martin Reinecke committed
64
      auto tmp2=tmp.template subarray<2>({0,0},{ntheta, nphi});
Martin Reinecke's avatar
Martin Reinecke committed
65
66
      fmav<T> ftmp2(tmp2);
      fmav<T> farr(arr);
Martin Reinecke's avatar
Martin Reinecke committed
67
      // zero-padded FFT in phi direction
Martin Reinecke's avatar
Martin Reinecke committed
68
      r2r_fftpack(ftmp2,farr,{1},false,false,1.,0);
69
70
71
72
73
      }

  public:
    Interpolator(const Alm<complex<T>> &slmT, const Alm<complex<T>> &blmT,
      double epsilon)
Martin Reinecke's avatar
fix    
Martin Reinecke committed
74
      : lmax(slmT.Lmax()),
75
        kmax(blmT.Mmax()),
Martin Reinecke's avatar
Martin Reinecke committed
76
77
        nphi0(2*good_size_real(lmax+1)),
        ntheta0(nphi0/2+1),
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
78
        nphi(2*good_size_real(size_t((2*lmax+1)*ofmin/2.))),
Martin Reinecke's avatar
Martin Reinecke committed
79
80
        ntheta(nphi/2+1),
        ofactor(double(nphi)/(2*lmax+1)),
Martin Reinecke's avatar
fix    
Martin Reinecke committed
81
82
        supp(ES_Kernel::get_supp(epsilon, ofactor)),
        kernel(supp, ofactor, 1),
83
84
85
86
87
88
        cube({ntheta+2*supp, nphi+2*supp, 2*kmax+1})
      {
      MR_assert((supp<=ntheta) && (supp<=nphi), "support too large!");
      MR_assert(slmT.Mmax()==lmax, "Sky lmax must be equal to Sky mmax");
      MR_assert(blmT.Lmax()==lmax, "Sky and beam lmax must be equal");
      Alm<complex<T>> a1(lmax, lmax), a2(lmax,lmax);
Martin Reinecke's avatar
Martin Reinecke committed
89
      auto ginfo = sharp_make_cc_geom_info(ntheta0,nphi0,0,cube.stride(1),cube.stride(0));
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
      auto ainfo = sharp_make_triangular_alm_info(lmax,lmax,1);

      vector<double>lnorm(lmax+1);
      for (size_t i=0; i<=lmax; ++i)
        lnorm[i]=sqrt(4*pi/(2*i+1.));

      for (size_t k=0; k<=kmax; ++k)
        {
        double spinsign = (k==0) ? 1. : -1.;
        for (size_t m=0; m<=lmax; ++m)
          {
          T mfac=T((m&1) ? -1.:1.);
          for (size_t l=m; l<=lmax; ++l)
            {
            if (l<k)
              a1(l,m)=a2(l,m)=0.;
            else
              {
              complex<T> v1=slmT(l,m)*blmT(l,k),
                         v2=conj(slmT(l,m))*blmT(l,k)*mfac;
              a1(l,m) = (v1+conj(v2)*mfac)*T(0.5*spinsign*lnorm[l]);
              if (k>0)
                {
                complex<T> tmp = (v1-conj(v2)*mfac)*T(-spinsign*0.5*lnorm[l]);
                a2(l,m) = complex<T>(-tmp.imag(), tmp.real());
                }
              }
            }
          }
        size_t kidx1 = (k==0) ? 0 : 2*k-1,
               kidx2 = (k==0) ? 0 : 2*k;
        auto quadrant=k%4;
        if (quadrant&1)
          swap(kidx1, kidx2);
Martin Reinecke's avatar
Martin Reinecke committed
124
125
        auto m1 = cube.template subarray<2>({supp,supp,kidx1},{ntheta,nphi,0});
        auto m2 = cube.template subarray<2>({supp,supp,kidx2},{ntheta,nphi,0});
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
        if (k==0)
          sharp_alm2map(a1.Alms().data(), m1.vdata(), *ginfo, *ainfo, 0, nullptr, nullptr);
        else
          sharp_alm2map_spin(k, a1.Alms().data(), a2.Alms().data(), m1.vdata(), m2.vdata(), *ginfo, *ainfo, 0, nullptr, nullptr);
        correct(m1);
        if (k!=0) correct(m2);

        if ((quadrant==1)||(quadrant==2)) m1.apply([](T &v){v=-v;});
        if ((k>0) &&((quadrant==0)||(quadrant==1))) m2.apply([](T &v){v=-v;});
        }
      // fill border regions
      for (size_t i=0; i<supp; ++i)
        for (size_t j=0, j2=nphi/2; j<nphi; ++j,++j2)
          for (size_t k=0; k<cube.shape(2); ++k)
            {
            if (j2>=nphi) j2-=nphi;
            cube.v(supp-1-i,j2+supp,k) = cube(supp+1+i,j+supp,k);
            cube.v(supp+ntheta+i,j2+supp,k) = cube(supp+ntheta-2-i, j+supp,k);
            }
      for (size_t i=0; i<ntheta+2*supp; ++i)
        for (size_t j=0; j<supp; ++j)
          for (size_t k=0; k<cube.shape(2); ++k)
            {
            cube.v(i,j,k) = cube(i,j+nphi,k);
            cube.v(i,j+nphi+supp,k) = cube(i,j+supp,k);
            }
      }

    void interpolx (const mav<T,2> &ptg, mav<T,1> &res) const
      {
      MR_assert(ptg.shape(0)==res.shape(0), "dimension mismatch");
      MR_assert(ptg.shape(1)==3, "second dimension must have length 3");
      vector<T> wt(supp), wp(supp);
      vector<T> psiarr(2*kmax+1);
Martin Reinecke's avatar
Martin Reinecke committed
160
161
162
      double delta = 2./supp;
      double xdtheta = (ntheta-1)/pi,
             xdphi = nphi/(2*pi);
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
163
164
165
166
167
168
169
170
171
172
173
174
175
176
#if 1
      constexpr size_t cellsize=16;
      size_t nct = ntheta/cellsize+1,
             ncp = nphi/cellsize+1;
      vector<vector<size_t>> mapper(nct*ncp);
      for (size_t i=0; i<ptg.shape(0); ++i)
        {
        size_t itheta=size_t(ptg(i,0)/pi*nct),
               iphi=size_t(ptg(i,1)/(2*pi)*ncp);
        mapper[itheta*ncp+iphi].push_back(i);
        }
      for (const auto &vec: mapper)
      for (auto i:vec)
#else
177
      for (size_t i=0; i<ptg.shape(0); ++i)
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
178
#endif
179
        {
Martin Reinecke's avatar
Martin Reinecke committed
180
        double f0=0.5*supp+ptg(i,0)*xdtheta;
181
182
        size_t i0 = size_t(f0+1.);
        for (size_t t=0; t<supp; ++t)
Martin Reinecke's avatar
Martin Reinecke committed
183
184
          wt[t] = kernel((t+i0-f0)*delta - 1);
        double f1=0.5*supp+ptg(i,1)*xdphi;
185
186
        size_t i1 = size_t(f1+1.);
        for (size_t t=0; t<supp; ++t)
Martin Reinecke's avatar
Martin Reinecke committed
187
          wp[t] = kernel((t+i1-f1)*delta - 1);
188
189
        double val=0;
        psiarr[0]=1.;
Martin Reinecke's avatar
Martin Reinecke committed
190
        double cpsi=cos(ptg(i,2)), spsi=sin(ptg(i,2));
191
192
193
194
        double cnpsi=cpsi, snpsi=spsi;
        for (size_t l=1; l<=kmax; ++l)
          {
          psiarr[2*l-1]=cnpsi;
Martin Reinecke's avatar
Martin Reinecke committed
195
          psiarr[2*l]=-snpsi;
196
197
198
199
200
201
202
          const double tmp = snpsi*cpsi + cnpsi*spsi;
          cnpsi=cnpsi*cpsi - snpsi*spsi;
          snpsi=tmp;
          }
        for (size_t j=0; j<supp; ++j)
          for (size_t k=0; k<supp; ++k)
            for (size_t l=0; l<2*kmax+1; ++l)
Martin Reinecke's avatar
fix    
Martin Reinecke committed
203
              val += cube(i0+j,i1+k,l)*wt[j]*wp[k]*psiarr[l];
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
        res.v(i) = val;
        }
      }
  };

template<typename T> class PyInterpolator: public Interpolator<T>
  {
  public:
    PyInterpolator(const py::array &slmT, const py::array &blmT, int64_t lmax, int64_t kmax, double epsilon)
      : Interpolator<T>(Alm<complex<T>>(to_mav<complex<T>,1>(slmT), lmax, lmax),
                        Alm<complex<T>>(to_mav<complex<T>,1>(blmT), lmax, kmax),
                        epsilon) {}
    using Interpolator<T>::interpolx;
    py::array interpol(const py::array &ptg)
      {
      auto ptg2 = to_mav<T,2>(ptg);
      auto res = make_Pyarr<double>({ptg2.shape(0)});
      auto res2 = to_mav<double,1>(res,true);
      interpolx(ptg2, res2);
      return res;
      }
  };

} // unnamed namespace

PYBIND11_MODULE(interpol_ng, m)
  {
  using namespace pybind11::literals;

  py::class_<PyInterpolator<double>> (m, "PyInterpolator")
    .def(py::init<const py::array &, const py::array &, int64_t, int64_t, double>(),
      "sky"_a, "beam"_a, "lmax"_a, "kmax"_a, "epsilon"_a)
    .def ("interpol", &PyInterpolator<double>::interpol, "ptg"_a);
  }