working, but docs not updated yet

pyinterpol_ng/interpol_ng.h

@@ -450,23 +450,28 @@ template<typename T> class Interpolator
       for (size_t i=0; i<=lmax; ++i)
         lnorm[i]=std::sqrt(4*pi/(2*i+1.));
 
+      for (size_t j=0; j<blm.size(); ++j)
+        slm[j].SetToZero();
+
       for (size_t icomp=0; icomp<ncomp; ++icomp)
         {
-        bool separate = ncomp==blm.size();
+        bool separate = ncomp>1;
         {
-        auto m1 = cube.template subarray<2>({supp,supp,0,separate?icomp:0},{ntheta,nphi,0,0});
+        auto m1 = cube.template subarray<2>({supp,supp,0,icomp},{ntheta,nphi,0,0});
         decorrect(m1,0);
         sharp_alm2map_adjoint(a1.Alms().vdata(), m1.data(), *ginfo, *ainfo, 0, nthreads);
         for (size_t m=0; m<=lmax; ++m)
           for (size_t l=m; l<=lmax; ++l)
-            for (size_t j=0; j<ncomp; ++j)
-              slm[j](l,m)=conj(a1(l,m))*blm[j](l,0).real()*T(lnorm[l]);
+            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]);
         }
-
         for (size_t k=1; k<=kmax; ++k)
           {
-          auto m1 = cube.template subarray<2>({supp,supp,2*k-1,separate?icomp:0},{ntheta,nphi,0,0});
-          auto m2 = cube.template subarray<2>({supp,supp,2*k  ,separate?icomp:0},{ntheta,nphi,0,0});
+          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});
           decorrect(m1,k);
           decorrect(m2,k);
 
@@ -475,12 +480,21 @@ template<typename T> class Interpolator
           for (size_t m=0; m<=lmax; ++m)
             for (size_t l=m; l<=lmax; ++l)
               if (l>=k)
-                for (size_t j=0; j<ncomp; ++j)
+                {
+                if (separate)
                   {
-                  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();
+                  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();
                   }
+                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();
+                    }
+                }
           }
         }
       }

pyinterpol_ng/pyinterpol_ng.cc

@@ -19,30 +19,40 @@ template<typename T> class PyInterpolator: public Interpolator<T>
   protected:
     using Interpolator<T>::lmax;
     using Interpolator<T>::kmax;
+    using Interpolator<T>::ncomp;
     using Interpolator<T>::interpol;
     using Interpolator<T>::deinterpol;
     using Interpolator<T>::getSlm;
 
-vector<Alm<complex<T>>> makevec(const py::array &inp, int64_t lmax, int64_t kmax, bool rw=false)
+vector<Alm<complex<T>>> makevec(const py::array &inp, int64_t lmax, int64_t kmax)
   {
-  auto inp2 = to_mav<complex<T>,2>(inp, rw);
+  auto inp2 = to_mav<complex<T>,2>(inp);
   vector<Alm<complex<T>>> res;
   for (size_t i=0; i<inp2.shape(1); ++i)
     res.push_back(Alm<complex<T>>(inp2.template subarray<1>({0,i},{inp2.shape(0),0}),lmax, kmax));
   return res;
   }
+void makevec_v(py::array &inp, int64_t lmax, int64_t kmax, vector<Alm<complex<T>>> &res)
+  {
+  auto inp2 = to_mav<complex<T>,2>(inp, true);
+  for (size_t i=0; i<inp2.shape(1); ++i)
+    {
+    auto xtmp = inp2.template subarray<1>({0,i},{inp2.shape(0),0});
+    res.emplace_back(xtmp, lmax, kmax);
+    }
+  }
   public:
     PyInterpolator(const py::array &slm, const py::array &blm,
-      int64_t lmax, int64_t kmax, double epsilon, int nthreads=0)
+      bool separate, int64_t lmax, int64_t kmax, double epsilon, int nthreads=0)
       : Interpolator<T>(makevec(slm, lmax, lmax),
                         makevec(blm, lmax, kmax),
-                        false, epsilon, nthreads) {}
-    PyInterpolator(int64_t lmax, int64_t kmax, double epsilon, int nthreads=0)
-      : Interpolator<T>(lmax, kmax, 1, epsilon, nthreads) {}
+                        separate, epsilon, nthreads) {}
+    PyInterpolator(int64_t lmax, int64_t kmax, int64_t ncomp_, double epsilon, int nthreads=0)
+      : Interpolator<T>(lmax, kmax, ncomp_, epsilon, nthreads) {}
     py::array pyinterpol(const py::array &ptg) const
       {
       auto ptg2 = to_mav<T,2>(ptg);
-      auto res = make_Pyarr<double>({ptg2.shape(0),1});
+      auto res = make_Pyarr<double>({ptg2.shape(0),ncomp});
       auto res2 = to_mav<double,2>(res,true);
       interpol(ptg2, res2);
       return res;
@@ -58,7 +68,8 @@ vector<Alm<complex<T>>> makevec(const py::array &inp, int64_t lmax, int64_t kmax
       {
       auto blm = makevec(blm_, lmax, kmax);
       auto res = make_Pyarr<complex<T>>({Alm_Base::Num_Alms(lmax, lmax),blm.size()});
-      auto slm = makevec(res, lmax, lmax, true);
+      vector<Alm<complex<double>>> slm;
+      makevec_v(res, lmax, lmax, slm);
       getSlm(blm, slm);
       return res;
       }
@@ -212,11 +223,11 @@ PYBIND11_MODULE(pyinterpol_ng, m)
   m.doc() = pyinterpol_ng_DS;
 
   py::class_<PyInterpolator<double>> (m, "PyInterpolator", pyinterpolator_DS)
-    .def(py::init<const py::array &, const py::array &, int64_t, int64_t, double, int>(),
-      initnormal_DS, "sky"_a, "beam"_a, "lmax"_a, "kmax"_a, "epsilon"_a,
+    .def(py::init<const py::array &, const py::array &, bool, int64_t, int64_t, double, int>(),
+      initnormal_DS, "sky"_a, "beam"_a, "separate"_a, "lmax"_a, "kmax"_a, "epsilon"_a,
       "nthreads"_a)
-    .def(py::init<int64_t, int64_t, double, int>(), initadjoint_DS,
-      "lmax"_a, "kmax"_a, "epsilon"_a, "nthreads"_a)
+    .def(py::init<int64_t, int64_t, int64_t, double, int>(), initadjoint_DS,
+      "lmax"_a, "kmax"_a, "ncomp"_a, "epsilon"_a, "nthreads"_a)
     .def ("interpol", &PyInterpolator<double>::pyinterpol, interpol_DS, "ptg"_a)
     .def ("deinterpol", &PyInterpolator<double>::pydeinterpol, deinterpol_DS, "ptg"_a, "data"_a)
     .def ("getSlm", &PyInterpolator<double>::pygetSlm, getSlm_DS, "blmT"_a);

pyinterpol_ng/test/test_interpol_ng.py

@@ -51,14 +51,16 @@ def myalmdot(a1,a2,lmax,mmax,spin):
     return np.vdot(compress_alm(a1,lmax),compress_alm(np.conj(a2),lmax))
 
 
-@pmp("lkmax", [(43,43),(2,1),(30,15),(125,2)])
-def test_against_convolution(lkmax):
+@pmp("lkmax", [(13,13),(2,1),(30,15),(35,2)])
+@pmp("ncomp", [1, 3])
+@pmp("separate", [True, False])
+def test_against_convolution(lkmax, ncomp, separate):
     lmax, kmax = lkmax
-    slm = random_alm(lmax, lmax, 1)
-    blm = random_alm(lmax, kmax, 1)
+    slm = random_alm(lmax, lmax, ncomp)
+    blm = random_alm(lmax, kmax, ncomp)
 
-    inter = pyinterpol_ng.PyInterpolator(slm, blm, lmax, kmax, epsilon=1e-8,
-                                         nthreads=2)
+    inter = pyinterpol_ng.PyInterpolator(slm, blm, separate, lmax, kmax,
+                                         epsilon=1e-8, nthreads=2)
     nptg = 50
     ptg = np.zeros((nptg,3))
     ptg[:,0] = np.random.uniform(0, np.pi, nptg)
@@ -67,28 +69,37 @@ def test_against_convolution(lkmax):
 
     res1 = inter.interpol(ptg)
 
-    res2 = np.zeros(nptg)
-    blm2 = np.zeros(nalm(lmax,lmax))+0j
-    blm2[0:blm.shape[0]] = blm[:,0]
-    for i in range(nptg):
-        rbeam=pyinterpol_ng.rotate_alm(blm2, lmax, ptg[i,2],ptg[i,0],ptg[i,1])
-        res2[i] = convolve(slm[:,0], rbeam, lmax).real
-    _assert_close(res1[:,0], res2, 1e-7)
-
-@pmp("lkmax", [(5,2)])#[(43,43),(2,1),(30,15),(125,2)])
-def test_adjointness(lkmax):
+    blm2 = np.zeros((nalm(lmax,lmax), ncomp))+0j
+    blm2[0:blm.shape[0],:] = blm
+    res2 = np.zeros((nptg, ncomp))
+    for c in range(ncomp):
+        for i in range(nptg):
+            rbeam=pyinterpol_ng.rotate_alm(blm2[:,c], lmax, ptg[i,2],ptg[i,0],ptg[i,1])
+            res2[i,c] = convolve(slm[:,c], rbeam, lmax).real
+    if separate:
+        _assert_close(res1, res2, 1e-7)
+    else:
+        _assert_close(res1[:,0], np.sum(res2,axis=1), 1e-7)
+
+@pmp("lkmax", [(13,13),(2,1),(30,15),(35,2)])
+@pmp("ncomp", [1, 3])
+@pmp("separate", [True, False])
+def test_adjointness(lkmax, ncomp, separate):
     lmax, kmax = lkmax
-    slm = random_alm(lmax, lmax, 1)
-    blm = random_alm(lmax, kmax, 1)
+    slm = random_alm(lmax, lmax, ncomp)
+    blm = random_alm(lmax, kmax, ncomp)
     nptg=100000
     ptg=np.random.uniform(0.,1.,nptg*3).reshape(nptg,3)
     ptg[:,0]*=np.pi
     ptg[:,1]*=2*np.pi
     ptg[:,2]*=2*np.pi
-    foo = pyinterpol_ng.PyInterpolator(slm,blm,lmax, kmax, epsilon=1e-6, nthreads=2)
+    foo = pyinterpol_ng.PyInterpolator(slm,blm,separate,lmax, kmax, epsilon=1e-6, nthreads=2)
     inter1=foo.interpol(ptg)
-    fake = np.random.uniform(0.,1., (ptg.shape[0],1))
-    foo2 = pyinterpol_ng.PyInterpolator(lmax, kmax, epsilon=1e-6, nthreads=2)
+    ncomp2 = inter1.shape[1]
+    fake = np.random.uniform(0.,1., (ptg.shape[0],ncomp2))
+    foo2 = pyinterpol_ng.PyInterpolator(lmax, kmax, ncomp2, epsilon=1e-6, nthreads=2)
     foo2.deinterpol(ptg.reshape((-1,3)), fake)
     bla=foo2.getSlm(blm)
-    _assert_close(myalmdot(slm, bla, lmax, lmax, 0), np.vdot(fake,inter1), 1e-12)
+    v1 = np.sum([myalmdot(slm[:,c], bla[:,c], lmax, lmax, 0) for c in range(ncomp)])
+    v2 = np.sum([np.vdot(fake[:,c],inter1[:,c]) for c in range(ncomp2)])
+    _assert_close(v1, v2, 1e-12)

src/mr_util/infra/mav.h

@@ -130,7 +130,7 @@ template<typename T> class membuf
       MR_assert(rw, "array is not writable");
       return const_cast<T *>(d);
       }
-    bool writable() { return rw; }
+    bool writable() const { return rw; }
   };
 
 // "mav" stands for "multidimensional array view"