interpol_ng.h 19.4 KB
Newer Older
1
2
3
4
5
/*
 *  Copyright (C) 2020 Max-Planck-Society
 *  Author: Martin Reinecke
 */

Martin Reinecke's avatar
Martin Reinecke committed
6
7
8
#ifndef MRUTIL_INTERPOL_NG_H
#define MRUTIL_INTERPOL_NG_H

9
10
#include <vector>
#include <complex>
Martin Reinecke's avatar
Martin Reinecke committed
11
#include <cmath>
12
13
14
15
16
17
18
19
20
21
#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
Martin Reinecke committed
22

Martin Reinecke's avatar
Martin Reinecke committed
23
namespace mr {
24

Martin Reinecke's avatar
Martin Reinecke committed
25
namespace detail_interpol_ng {
26

Martin Reinecke's avatar
Martin Reinecke committed
27
using namespace std;
28

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

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

43
    void correct(mav<T,2> &arr, int spin)
44
      {
45
      double sfct = (spin&1) ? -1 : 1;
46
      mav<T,2> tmp({nphi,nphi});
Martin Reinecke's avatar
Martin Reinecke committed
47
      tmp.apply([](T &v){v=0.;});
Martin Reinecke's avatar
Martin Reinecke committed
48
      auto tmp0=tmp.template subarray<2>({0,0},{nphi0, nphi0});
Martin Reinecke's avatar
Martin Reinecke committed
49
50
51
52
      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);
53
      // extend to second half
Martin Reinecke's avatar
Martin Reinecke committed
54
      for (size_t i=1, i2=nphi0-1; i+1<ntheta0; ++i,--i2)
Martin Reinecke's avatar
Martin Reinecke committed
55
        for (size_t j=0,j2=nphi0/2; j<nphi0; ++j,++j2)
56
          {
Martin Reinecke's avatar
Martin Reinecke committed
57
          if (j2>=nphi0) j2-=nphi0;
Martin Reinecke's avatar
Martin Reinecke committed
58
          tmp0.v(i2,j) = sfct*tmp0(i,j2);
59
          }
Martin Reinecke's avatar
Martin Reinecke committed
60
      // FFT to frequency domain on minimal grid
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
61
62
      r2r_fftpack(ftmp0,ftmp0,{0,1},true,true,1./(nphi0*nphi0),nthreads);
      // correct amplitude at Nyquist frequency
Martin Reinecke's avatar
Martin Reinecke committed
63
64
65
66
67
      for (size_t i=0; i<nphi0; ++i)
        {
        tmp0.v(i,nphi0-1)*=0.5;
        tmp0.v(nphi0-1,i)*=0.5;
        }
Martin Reinecke's avatar
Martin Reinecke committed
68
      auto fct = kernel.correction_factors(nphi, nphi0/2+1, nthreads);
Martin Reinecke's avatar
Martin Reinecke committed
69
70
71
      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
tweaks    
Martin Reinecke committed
72
73
74
75
76
77
78
79
80
      auto tmp1=tmp.template subarray<2>({0,0},{nphi, nphi0});
      fmav<T> ftmp1(tmp1);
      // zero-padded FFT in theta direction
      r2r_fftpack(ftmp1,ftmp1,{0},false,false,1.,nthreads);
      auto tmp2=tmp.template subarray<2>({0,0},{ntheta, nphi});
      fmav<T> ftmp2(tmp2);
      fmav<T> farr(arr);
      // zero-padded FFT in phi direction
      r2r_fftpack(ftmp2,farr,{1},false,false,1.,nthreads);
81
      }
82
83
84
85
86
87
88
89
90
91
    void decorrect(mav<T,2> &arr, int spin)
      {
      double sfct = (spin&1) ? -1 : 1;
      mav<T,2> tmp({nphi,nphi});
      fmav<T> ftmp(tmp);

      for (size_t i=0; i<ntheta; ++i)
        for (size_t j=0; j<nphi; ++j)
          tmp.v(i,j) = arr(i,j);
      // extend to second half
Martin Reinecke's avatar
Martin Reinecke committed
92
      for (size_t i=1, i2=nphi-1; i+1<ntheta; ++i,--i2)
93
94
95
96
97
        for (size_t j=0,j2=nphi/2; j<nphi; ++j,++j2)
          {
          if (j2>=nphi) j2-=nphi;
          tmp.v(i2,j) = sfct*tmp(i,j2);
          }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
98
99
100
101
      r2r_fftpack(ftmp,ftmp,{1},true,true,1.,nthreads);
      auto tmp1=tmp.template subarray<2>({0,0},{nphi, nphi0});
      fmav<T> ftmp1(tmp1);
      r2r_fftpack(ftmp1,ftmp1,{0},true,true,1.,nthreads);
102
103
104
105
106
107
108
      auto fct = kernel.correction_factors(nphi, nphi0/2+1, nthreads);
      auto tmp0=tmp.template subarray<2>({0,0},{nphi0, nphi0});
      fmav<T> ftmp0(tmp0);
      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];
      // FFT to (theta, phi) domain on minimal grid
Martin Reinecke's avatar
tweaks    
Martin Reinecke committed
109
      r2r_fftpack(ftmp0,ftmp0,{0,1},false, false,1./(nphi0*nphi0),nthreads);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
110
111
112
113
114
      for (size_t j=0; j<nphi0; ++j)
        {
        tmp0.v(0,j)*=0.5;
        tmp0.v(ntheta0-1,j)*=0.5;
        }
115
116
117
118
      for (size_t i=0; i<ntheta0; ++i)
        for (size_t j=0; j<nphi0; ++j)
          arr.v(i,j) = tmp0(i,j);
      }
119

Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
    vector<size_t> getIdx(const mav<T,2> &ptg) const
      {
      vector<size_t> idx(ptg.shape(0));
      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=min(nct-1,size_t(ptg(i,0)/pi*nct)),
               iphi=min(ncp-1,size_t(ptg(i,1)/(2*pi)*ncp));
        mapper[itheta*ncp+iphi].push_back(i);
        }
      size_t cnt=0;
      for (const auto &vec: mapper)
        for (auto i:vec)
          idx[cnt++] = i;
      return idx;
      }

140
  public:
141
142
143
    Interpolator(const vector<Alm<complex<T>>> &slm,
                 const vector<Alm<complex<T>>> &blm,
                 bool separate, double epsilon, int nthreads_)
144
      : adjoint(false),
145
146
        lmax(slm.at(0).Lmax()),
        kmax(blm.at(0).Mmax()),
Martin Reinecke's avatar
Martin Reinecke committed
147
148
        nphi0(2*good_size_real(lmax+1)),
        ntheta0(nphi0/2+1),
Martin Reinecke's avatar
Martin Reinecke committed
149
        nphi(std::max<size_t>(20,2*good_size_real(size_t((2*lmax+1)*ofmin/2.)))),
Martin Reinecke's avatar
Martin Reinecke committed
150
        ntheta(nphi/2+1),
Martin Reinecke's avatar
Martin Reinecke committed
151
        nthreads(nthreads_),
Martin Reinecke's avatar
Martin Reinecke committed
152
        ofactor(double(nphi)/(2*lmax+1)),
Martin Reinecke's avatar
fix    
Martin Reinecke committed
153
        supp(ES_Kernel::get_supp(epsilon, ofactor)),
Martin Reinecke's avatar
Martin Reinecke committed
154
        kernel(supp, ofactor, nthreads),
155
156
        ncomp(separate ? slm.size() : 1),
        cube({ntheta+2*supp, nphi+2*supp, 2*kmax+1, ncomp})
157
      {
158
      MR_assert((ncomp==1)||(ncomp==3), "currently only 1 or 3 components allowed");
159
160
161
162
163
164
165
166
167
      MR_assert(slm.size()==blm.size(), "inconsistent slm and blm vectors");
      for (size_t i=0; i<slm.size(); ++i)
        {
        MR_assert(slm[i].Lmax()==lmax, "inconsistent Sky lmax");
        MR_assert(slm[i].Mmax()==lmax, "Sky lmax must be equal to Sky mmax");
        MR_assert(blm[i].Lmax()==lmax, "Sky and beam lmax must be equal");
        MR_assert(blm[i].Mmax()==kmax, "Inconcistent beam mmax");
        }

168
169
      MR_assert((supp<=ntheta) && (supp<=nphi), "support too large!");
      Alm<complex<T>> a1(lmax, lmax), a2(lmax,lmax);
Martin Reinecke's avatar
Martin Reinecke committed
170
      auto ginfo = sharp_make_cc_geom_info(ntheta0,nphi0,0.,cube.stride(1),cube.stride(0));
171
172
173
174
      auto ainfo = sharp_make_triangular_alm_info(lmax,lmax,1);

      vector<double>lnorm(lmax+1);
      for (size_t i=0; i<=lmax; ++i)
Martin Reinecke's avatar
Martin Reinecke committed
175
        lnorm[i]=std::sqrt(4*pi/(2*i+1.));
176

177
      for (size_t icomp=0; icomp<ncomp; ++icomp)
178
179
180
181
        {
        for (size_t m=0; m<=lmax; ++m)
          for (size_t l=m; l<=lmax; ++l)
            {
182
183
            if (separate)
              a1(l,m) = slm[icomp](l,m)*blm[icomp](l,0).real()*T(lnorm[l]);
184
185
            else
              {
186
187
188
              a1(l,m) = 0;
              for (size_t j=0; j<slm.size(); ++j)
                a1(l,m) += slm[j](l,m)*blm[j](l,0).real()*T(lnorm[l]);
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
        auto m1 = cube.template subarray<2>({supp,supp,0,icomp},{ntheta,nphi,0,0});
        sharp_alm2map(a1.Alms().data(), m1.vdata(), *ginfo, *ainfo, 0, nthreads);
        correct(m1,0);

        for (size_t k=1; k<=kmax; ++k)
          {
          for (size_t m=0; m<=lmax; ++m)
            for (size_t l=m; l<=lmax; ++l)
              {
              if (l<k)
                a1(l,m)=a2(l,m)=0.;
              else
                {
                if (separate)
                  {
                  auto tmp = -2.*blm[icomp](l,k)*T(lnorm[l]);
                  a1(l,m) = slm[icomp](l,m)*tmp.real();
                  a2(l,m) = slm[icomp](l,m)*tmp.imag();
                  }
                else
                  {
                  a1(l,m) = a2(l,m) = 0;
                  for (size_t j=0; j<slm.size(); ++j)
                    {
                    auto tmp = -2.*blm[j](l,k)*T(lnorm[l]);
                    a1(l,m) += slm[j](l,m)*tmp.real();
                    a2(l,m) += slm[j](l,m)*tmp.imag();
                    }
                  }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
220
                }
221
              }
Martin Reinecke's avatar
fixes    
Martin Reinecke committed
222
223
224
225
226
227
          auto m1 = cube.template subarray<2>({supp,supp,2*k-1,icomp},{ntheta,nphi,0,0});
          auto m2 = cube.template subarray<2>({supp,supp,2*k  ,icomp},{ntheta,nphi,0,0});
          sharp_alm2map_spin(k, a1.Alms().data(), a2.Alms().data(), m1.vdata(),
            m2.vdata(), *ginfo, *ainfo, 0, nthreads);
          correct(m1,k);
          correct(m2,k);
228
          }
229
        }
230

231
232
233
234
235
      // 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)
            {
236
            double fct = (((k+1)/2)&1) ? -1 : 1;
237
            if (j2>=nphi) j2-=nphi;
238
239
240
241
242
            for (size_t l=0; l<cube.shape(3); ++l)
              {
              cube.v(supp-1-i,j2+supp,k,l) = fct*cube(supp+1+i,j+supp,k,l);
              cube.v(supp+ntheta+i,j2+supp,k,l) = fct*cube(supp+ntheta-2-i,j+supp,k,l);
              }
243
244
245
246
            }
      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)
247
            for (size_t l=0; l<cube.shape(3); ++l)
248
            {
249
250
            cube.v(i,j,k,l) = cube(i,j+nphi,k,l);
            cube.v(i,j+nphi+supp,k,l) = cube(i,j+supp,k,l);
251
252
253
            }
      }

254
    Interpolator(size_t lmax_, size_t kmax_, size_t ncomp_, double epsilon, int nthreads_)
255
256
257
258
259
      : adjoint(true),
        lmax(lmax_),
        kmax(kmax_),
        nphi0(2*good_size_real(lmax+1)),
        ntheta0(nphi0/2+1),
Martin Reinecke's avatar
Martin Reinecke committed
260
        nphi(std::max<size_t>(20,2*good_size_real(size_t((2*lmax+1)*ofmin/2.)))),
261
262
263
264
265
        ntheta(nphi/2+1),
        nthreads(nthreads_),
        ofactor(double(nphi)/(2*lmax+1)),
        supp(ES_Kernel::get_supp(epsilon, ofactor)),
        kernel(supp, ofactor, nthreads),
266
267
        ncomp(ncomp_),
        cube({ntheta+2*supp, nphi+2*supp, 2*kmax+1, ncomp_})
268
      {
269
      MR_assert((ncomp==1)||(ncomp==3), "currently only 1 or 3 components allowed");
270
271
272
273
      MR_assert((supp<=ntheta) && (supp<=nphi), "support too large!");
      cube.apply([](T &v){v=0.;});
      }

274
    void interpol (const mav<T,2> &ptg, mav<T,2> &res) const
275
      {
276
      MR_assert(!adjoint, "cannot be called in adjoint mode");
277
278
      MR_assert(ptg.shape(0)==res.shape(0), "dimension mismatch");
      MR_assert(ptg.shape(1)==3, "second dimension must have length 3");
279
      MR_assert(res.shape(1)==ncomp, "# of components mismatch");
Martin Reinecke's avatar
Martin Reinecke committed
280
281
282
      double delta = 2./supp;
      double xdtheta = (ntheta-1)/pi,
             xdphi = nphi/(2*pi);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
283
      auto idx = getIdx(ptg);
Martin Reinecke's avatar
Martin Reinecke committed
284
      execStatic(idx.size(), nthreads, 0, [&](Scheduler &sched)
285
        {
Martin Reinecke's avatar
Martin Reinecke committed
286
287
288
        vector<T> wt(supp), wp(supp);
        vector<T> psiarr(2*kmax+1);
        while (auto rng=sched.getNext()) for(auto ind=rng.lo; ind<rng.hi; ++ind)
289
          {
Martin Reinecke's avatar
Martin Reinecke committed
290
291
292
293
294
295
296
297
298
299
          size_t i=idx[ind];
          double f0=0.5*supp+ptg(i,0)*xdtheta;
          size_t i0 = size_t(f0+1.);
          for (size_t t=0; t<supp; ++t)
            wt[t] = kernel((t+i0-f0)*delta - 1);
          double f1=0.5*supp+ptg(i,1)*xdphi;
          size_t i1 = size_t(f1+1.);
          for (size_t t=0; t<supp; ++t)
            wp[t] = kernel((t+i1-f1)*delta - 1);
          psiarr[0]=1.;
300
301
          double psi=ptg(i,2);
          double cpsi=cos(psi), spsi=sin(psi);
Martin Reinecke's avatar
Martin Reinecke committed
302
303
304
305
          double cnpsi=cpsi, snpsi=spsi;
          for (size_t l=1; l<=kmax; ++l)
            {
            psiarr[2*l-1]=cnpsi;
306
            psiarr[2*l]=snpsi;
Martin Reinecke's avatar
Martin Reinecke committed
307
308
309
310
            const double tmp = snpsi*cpsi + cnpsi*spsi;
            cnpsi=cnpsi*cpsi - snpsi*spsi;
            snpsi=tmp;
            }
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
          if (ncomp==1)
            {
            double vv=0.;
            for (size_t j=0; j<supp; ++j)
              for (size_t k=0; k<supp; ++k)
                {
                double t0=wt[j]*wp[k];
                for (size_t l=0; l<2*kmax+1; ++l)
                  vv += cube(i0+j,i1+k,l,0)*t0*psiarr[l];
                }
            res.v(i,0) = vv;
            }
          else // ncomp==3
            {
            double v0=0., v1=0., v2=0.;
            for (size_t j=0; j<supp; ++j)
              for (size_t k=0; k<supp; ++k)
                {
                double t0 = wt[j]*wp[k];
                for (size_t l=0; l<2*kmax+1; ++l)
                  {
                  auto tmp = t0*psiarr[l];
                  v0 += cube(i0+j,i1+k,l,0)*tmp;
                  v1 += cube(i0+j,i1+k,l,1)*tmp;
                  v2 += cube(i0+j,i1+k,l,2)*tmp;
                  }
                }
            res.v(i,0) = v0;
            res.v(i,1) = v1;
            res.v(i,2) = v2;
            }
342
          }
Martin Reinecke's avatar
Martin Reinecke committed
343
        });
344
      }
345

346
    void deinterpol (const mav<T,2> &ptg, const mav<T,2> &data)
347
348
349
350
      {
      MR_assert(adjoint, "can only be called in adjoint mode");
      MR_assert(ptg.shape(0)==data.shape(0), "dimension mismatch");
      MR_assert(ptg.shape(1)==3, "second dimension must have length 3");
351
      MR_assert(data.shape(1)==ncomp, "# of components mismatch");
352
353
354
      double delta = 2./supp;
      double xdtheta = (ntheta-1)/pi,
             xdphi = nphi/(2*pi);
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
355
      auto idx = getIdx(ptg);
356
357
358
359
360
361
362

      constexpr size_t cellsize=16;
      size_t nct = ntheta/cellsize+5,
             ncp = nphi/cellsize+5;
      mav<std::mutex,2> locks({nct,ncp});

      execStatic(idx.size(), nthreads, 0, [&](Scheduler &sched)
363
        {
364
        size_t b_theta=99999999999999, b_phi=9999999999999999;
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
        vector<T> wt(supp), wp(supp);
        vector<T> psiarr(2*kmax+1);
        while (auto rng=sched.getNext()) for(auto ind=rng.lo; ind<rng.hi; ++ind)
          {
          size_t i=idx[ind];
          double f0=0.5*supp+ptg(i,0)*xdtheta;
          size_t i0 = size_t(f0+1.);
          for (size_t t=0; t<supp; ++t)
            wt[t] = kernel((t+i0-f0)*delta - 1);
          double f1=0.5*supp+ptg(i,1)*xdphi;
          size_t i1 = size_t(f1+1.);
          for (size_t t=0; t<supp; ++t)
            wp[t] = kernel((t+i1-f1)*delta - 1);
          psiarr[0]=1.;
          double psi=ptg(i,2);
          double cpsi=cos(psi), spsi=sin(psi);
          double cnpsi=cpsi, snpsi=spsi;
          for (size_t l=1; l<=kmax; ++l)
            {
            psiarr[2*l-1]=cnpsi;
            psiarr[2*l]=snpsi;
            const double tmp = snpsi*cpsi + cnpsi*spsi;
            cnpsi=cnpsi*cpsi - snpsi*spsi;
            snpsi=tmp;
            }
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
          size_t b_theta_new = i0/cellsize,
                 b_phi_new = i1/cellsize;
          if ((b_theta_new!=b_theta) || (b_phi_new!=b_phi))
            {
            if (b_theta<locks.shape(0))  // unlock
              {
              locks.v(b_theta,b_phi).unlock();
              locks.v(b_theta,b_phi+1).unlock();
              locks.v(b_theta+1,b_phi).unlock();
              locks.v(b_theta+1,b_phi+1).unlock();
              }
            b_theta = b_theta_new;
            b_phi = b_phi_new;
            locks.v(b_theta,b_phi).lock();
            locks.v(b_theta,b_phi+1).lock();
            locks.v(b_theta+1,b_phi).lock();
            locks.v(b_theta+1,b_phi+1).lock();
            }
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
          if (ncomp==1)
            {
            double val = data(i,0);
            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)
                  cube.v(i0+j,i1+k,l,0) += val*wt[j]*wp[k]*psiarr[l];
            }
          else // ncomp==3
            {
            double v0=data(i,0), v1=data(i,1), v2=data(i,2);
            for (size_t j=0; j<supp; ++j)
              for (size_t k=0; k<supp; ++k)
                {
                double t0 = wt[j]*wp[k];
                for (size_t l=0; l<2*kmax+1; ++l)
                  {
                  double tmp = t0*psiarr[l];
                  cube.v(i0+j,i1+k,l,0) += v0*tmp;
                  cube.v(i0+j,i1+k,l,1) += v1*tmp;
                  cube.v(i0+j,i1+k,l,2) += v2*tmp;
                  }
                }
            }
432
          }
433
434
435
436
437
438
439
        if (b_theta<locks.shape(0))  // unlock
          {
          locks.v(b_theta,b_phi).unlock();
          locks.v(b_theta,b_phi+1).unlock();
          locks.v(b_theta+1,b_phi).unlock();
          locks.v(b_theta+1,b_phi+1).unlock();
          }
440
441
        });
      }
442
    void getSlm (const vector<Alm<complex<T>>> &blm, vector<Alm<complex<T>>> &slm)
443
444
      {
      MR_assert(adjoint, "can only be called in adjoint mode");
445
446
      MR_assert((blm.size()==ncomp) || (ncomp==1), "incorrect number of beam a_lm sets");
      MR_assert((slm.size()==ncomp) || (ncomp==1), "incorrect number of sky a_lm sets");
447
448
449
450
451
      Alm<complex<T>> a1(lmax, lmax), a2(lmax,lmax);
      auto ginfo = sharp_make_cc_geom_info(ntheta0,nphi0,0.,cube.stride(1),cube.stride(0));
      auto ainfo = sharp_make_triangular_alm_info(lmax,lmax,1);

      // move stuff from border regions onto the main grid
452
      for (size_t i=0; i<cube.shape(0); ++i)
453
454
        for (size_t j=0; j<supp; ++j)
          for (size_t k=0; k<cube.shape(2); ++k)
455
456
457
458
459
            for (size_t l=0; l<cube.shape(3); ++l)
              {
              cube.v(i,j+nphi,k,l) += cube(i,j,k,l);
              cube.v(i,j+supp,k,l) += cube(i,j+nphi+supp,k,l);
              }
460
461
462
463
464
465
      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)
            {
            double fct = (((k+1)/2)&1) ? -1 : 1;
            if (j2>=nphi) j2-=nphi;
466
467
468
469
470
            for (size_t l=0; l<cube.shape(3); ++l)
              {
              cube.v(supp+1+i,j+supp,k,l) += fct*cube(supp-1-i,j2+supp,k,l);
              cube.v(supp+ntheta-2-i, j+supp,k,l) += fct*cube(supp+ntheta+i,j2+supp,k,l);
              }
471
            }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
472
473
474
475

      // special treatment for poles
      for (size_t j=0,j2=nphi/2; j<nphi/2; ++j,++j2)
        for (size_t k=0; k<cube.shape(2); ++k)
476
477
478
479
480
481
482
483
484
485
486
          for (size_t l=0; l<cube.shape(3); ++l)
            {
            double fct = (((k+1)/2)&1) ? -1 : 1;
            if (j2>=nphi) j2-=nphi;
            double tval = (cube(supp,j+supp,k,l) + fct*cube(supp,j2+supp,k,l));
            cube.v(supp,j+supp,k,l) = tval;
            cube.v(supp,j2+supp,k,l) = fct*tval;
            tval = (cube(supp+ntheta-1,j+supp,k,l) + fct*cube(supp+ntheta-1,j2+supp,k,l));
            cube.v(supp+ntheta-1,j+supp,k,l) = tval;
            cube.v(supp+ntheta-1,j2+supp,k,l) = fct*tval;
            }
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
487

488
489
      vector<double>lnorm(lmax+1);
      for (size_t i=0; i<=lmax; ++i)
Martin Reinecke's avatar
Martin Reinecke committed
490
        lnorm[i]=std::sqrt(4*pi/(2*i+1.));
491

492
493
494
      for (size_t j=0; j<blm.size(); ++j)
        slm[j].SetToZero();

495
      for (size_t icomp=0; icomp<ncomp; ++icomp)
496
        {
497
        bool separate = ncomp>1;
498
        {
499
        auto m1 = cube.template subarray<2>({supp,supp,0,icomp},{ntheta,nphi,0,0});
500
501
        decorrect(m1,0);
        sharp_alm2map_adjoint(a1.Alms().vdata(), m1.data(), *ginfo, *ainfo, 0, nthreads);
502
503
        for (size_t m=0; m<=lmax; ++m)
          for (size_t l=m; l<=lmax; ++l)
504
505
506
507
508
            if (separate)
              slm[icomp](l,m) += conj(a1(l,m))*blm[icomp](l,0).real()*T(lnorm[l]);
            else
              for (size_t j=0; j<blm.size(); ++j)
                slm[j](l,m) += conj(a1(l,m))*blm[j](l,0).real()*T(lnorm[l]);
509
510
511
        }
        for (size_t k=1; k<=kmax; ++k)
          {
512
513
          auto m1 = cube.template subarray<2>({supp,supp,2*k-1,icomp},{ntheta,nphi,0,0});
          auto m2 = cube.template subarray<2>({supp,supp,2*k  ,icomp},{ntheta,nphi,0,0});
514
515
516
517
518
519
520
521
          decorrect(m1,k);
          decorrect(m2,k);

          sharp_alm2map_spin_adjoint(k, a1.Alms().vdata(), a2.Alms().vdata(), m1.data(),
            m2.data(), *ginfo, *ainfo, 0, nthreads);
          for (size_t m=0; m<=lmax; ++m)
            for (size_t l=m; l<=lmax; ++l)
              if (l>=k)
522
523
                {
                if (separate)
524
                  {
525
526
527
                  auto tmp = -2.*conj(blm[icomp](l,k))*T(lnorm[l]);
                  slm[icomp](l,m) += conj(a1(l,m))*tmp.real();
                  slm[icomp](l,m) -= conj(a2(l,m))*tmp.imag();
528
                  }
529
530
531
532
533
534
535
536
                else
                  for (size_t j=0; j<blm.size(); ++j)
                    {
                    auto tmp = -2.*conj(blm[j](l,k))*T(lnorm[l]);
                    slm[j](l,m) += conj(a1(l,m))*tmp.real();
                    slm[j](l,m) -= conj(a2(l,m))*tmp.imag();
                    }
                }
537
          }
538
539
        }
      }
540
541
  };

Martin Reinecke's avatar
Martin Reinecke committed
542
}
Martin Reinecke's avatar
cleanup    
Martin Reinecke committed
543

Martin Reinecke's avatar
Martin Reinecke committed
544
using detail_interpol_ng::Interpolator;
545

Martin Reinecke's avatar
Martin Reinecke committed
546
}
547

Martin Reinecke's avatar
Martin Reinecke committed
548
#endif