Merge branch 'multicomponent' into 'master'

Multicomponent support

See merge request mtr/cxxbase!4

pyinterpol_ng/demo.py

@@ -10,11 +10,11 @@ def nalm(lmax, mmax):
     return ((mmax+1)*(mmax+2))//2 + (mmax+1)*(lmax-mmax)
 
 
-def random_alm(lmax, mmax):
-    res = np.random.uniform(-1., 1., nalm(lmax, mmax)) \
-     + 1j*np.random.uniform(-1., 1., nalm(lmax, mmax))
+def random_alm(lmax, mmax, ncomp):
+    res = np.random.uniform(-1., 1., (nalm(lmax, mmax), ncomp)) \
+     + 1j*np.random.uniform(-1., 1., (nalm(lmax, mmax), ncomp))
     # make a_lm with m==0 real-valued
-    res[0:lmax+1].imag = 0.
+    res[0:lmax+1,:].imag = 0.
     return res
 
 
@@ -41,19 +41,21 @@ def convolve(alm1, alm2, lmax):
 
 lmax=60
 kmax=13
-
+ncomp=1
+separate=True
+ncomp2 = ncomp if separate else 1
 
 # get random sky a_lm
 # the a_lm arrays follow the same conventions as those in healpy
-slmT = random_alm(lmax, lmax)
+slm = random_alm(lmax, lmax, ncomp)
 
 # build beam a_lm
-blmT = random_alm(lmax, kmax)
+blm = random_alm(lmax, kmax, ncomp)
 
 
 t0=time.time()
-# build interpolator object for slmT and blmT
-foo = pyinterpol_ng.PyInterpolator(slmT,blmT,lmax, kmax, epsilon=1e-6, nthreads=2)
+# build interpolator object for slm and blm
+foo = pyinterpol_ng.PyInterpolator(slm,blm,separate,lmax, kmax, epsilon=1e-4, nthreads=2)
 print("setup time: ",time.time()-t0)
 nth = lmax+1
 nph = 2*lmax+1
@@ -68,22 +70,22 @@ ptg[:,:,1] = (2*np.pi*(0.5+np.arange(nph))/nph).reshape((1,-1))
 ptg[:,:,2] = np.pi*0.2
 t0=time.time()
 # do the actual interpolation
-bar=foo.interpol(ptg.reshape((-1,3))).reshape((nth,nph))
+bar=foo.interpol(ptg.reshape((-1,3))).reshape((nth,nph,ncomp2))
 print("interpolation time: ", time.time()-t0)
 plt.subplot(2,2,1)
-plt.imshow(bar.reshape((nth,nph)))
+plt.imshow(bar[:,:,0])
 bar2 = np.zeros((nth,nph))
-blmTfull = np.zeros(slmT.size)+0j
-blmTfull[0:blmT.size] = blmT
+blmfull = np.zeros(slm.shape)+0j
+blmfull[0:blm.shape[0],:] = blm
 for ith in range(nth):
-    rbeamth=pyinterpol_ng.rotate_alm(blmTfull, lmax, ptg[ith,0,2],ptg[ith,0,0],0)
+    rbeamth=pyinterpol_ng.rotate_alm(blmfull[:,0], lmax, ptg[ith,0,2],ptg[ith,0,0],0)
     for iph in range(nph):
         rbeam=pyinterpol_ng.rotate_alm(rbeamth, lmax, 0, 0, ptg[ith,iph,1])
-        bar2[ith,iph] = convolve(slmT, rbeam, lmax).real
+        bar2[ith,iph] = convolve(slm[:,0], rbeam, lmax).real
 plt.subplot(2,2,2)
 plt.imshow(bar2)
 plt.subplot(2,2,3)
-plt.imshow((bar2-bar.reshape((nth,nph))))
+plt.imshow(bar2-bar[:,:,0])
 plt.show()
 
 
@@ -91,12 +93,18 @@ ptg=np.random.uniform(0.,1.,3*1000000).reshape(1000000,3)
 ptg[:,0]*=np.pi
 ptg[:,1]*=2*np.pi
 ptg[:,2]*=2*np.pi
-foo = pyinterpol_ng.PyInterpolator(slmT,blmT,lmax, kmax, epsilon=1e-6, nthreads=2)
+#foo = pyinterpol_ng.PyInterpolator(slm,blm,separate,lmax, kmax, epsilon=1e-6, nthreads=2)
+t0=time.time()
 bar=foo.interpol(ptg)
-fake = np.random.uniform(0.,1., ptg.shape[0])
-foo2 = pyinterpol_ng.PyInterpolator(lmax, kmax, epsilon=1e-6, nthreads=2)
+print("interpolation time: ", time.time()-t0)
+fake = np.random.uniform(0.,1., (ptg.shape[0],ncomp2))
+foo2 = pyinterpol_ng.PyInterpolator(lmax, kmax, ncomp2, epsilon=1e-4, nthreads=2)
+t0=time.time()
 foo2.deinterpol(ptg.reshape((-1,3)), fake)
-bla=foo2.getSlm(blmT)
-print(myalmdot(slmT, bla, lmax, lmax, 0))
-print(np.vdot(fake,bar))
-print(myalmdot(slmT, bla, lmax, lmax, 0)/np.vdot(fake,bar))
+print("deinterpolation time: ", time.time()-t0)
+t0=time.time()
+bla=foo2.getSlm(blm)
+print("getSlm time: ", time.time()-t0)
+v1 = np.sum([myalmdot(slm[:,i], bla[:,i] , lmax, lmax, 0) for i in range(ncomp)])
+v2 = np.sum([np.vdot(fake[:,i],bar[:,i]) for i in range(ncomp2)])
+print(v1/v2-1.)

pyinterpol_ng/pyinterpol_ng.cc

@@ -14,46 +14,66 @@ namespace py = pybind11;
 
 namespace {
 
+using fptype = double;
+
 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)
+  {
+  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 &slmT, const py::array &blmT,
-      int64_t lmax, int64_t kmax, double epsilon, int nthreads=0)
-      : 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, nthreads) {}
-    PyInterpolator(int64_t lmax, int64_t kmax, double epsilon, int nthreads=0)
-      : Interpolator<T>(lmax, kmax, epsilon, nthreads) {}
+    PyInterpolator(const py::array &slm, const py::array &blm,
+      bool separate, int64_t lmax, int64_t kmax, T epsilon, int nthreads=0)
+      : Interpolator<T>(makevec(slm, lmax, lmax),
+                        makevec(blm, lmax, kmax),
+                        separate, epsilon, nthreads) {}
+    PyInterpolator(int64_t lmax, int64_t kmax, int64_t ncomp_, T 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)});
-      auto res2 = to_mav<double,1>(res,true);
+      auto res = make_Pyarr<T>({ptg2.shape(0),ncomp});
+      auto res2 = to_mav<T,2>(res,true);
       interpol(ptg2, res2);
-      return res;
+      return move(res);
       }
 
     void pydeinterpol(const py::array &ptg, const py::array &data)
       {
       auto ptg2 = to_mav<T,2>(ptg);
-      auto data2 = to_mav<T,1>(data);
+      auto data2 = to_mav<T,2>(data);
       deinterpol(ptg2, data2);
       }
-    py::array pygetSlm(const py::array &blmT_)
+    py::array pygetSlm(const py::array &blm_)
       {
-      auto res = make_Pyarr<complex<T>>({Alm_Base::Num_Alms(lmax, lmax)});
-      Alm<complex<T>> blmT(to_mav<complex<T>,1>(blmT_, false), lmax, kmax);
-      auto slmT_=to_mav<complex<T>,1>(res, true);
-      Alm<complex<T>> slmT(slmT_, lmax, lmax);
-      getSlm(blmT, slmT);
-      return res;
+      auto blm = makevec(blm_, lmax, kmax);
+      auto res = make_Pyarr<complex<T>>({Alm_Base::Num_Alms(lmax, lmax),blm.size()});
+      vector<Alm<complex<T>>> slm;
+      makevec_v(res, lmax, lmax, slm);
+      getSlm(blm, slm);
+      return move(res);
       }
   };
 
@@ -67,7 +87,7 @@ template<typename T> py::array pyrotate_alm(const py::array &alm_, int64_t lmax,
   for (size_t i=0; i<a1.shape(0); ++i) a2.v(i)=a1(i);
   auto tmp = Alm<complex<T>>(a2,lmax,lmax);
   rotate_alm(tmp, psi, theta, phi);
-  return alm;
+  return move(alm);
   }
 #endif
 
@@ -99,14 +119,14 @@ Constructor for interpolation mode
 
 Parameters
 ----------
-sky : numpy.ndarray(numpy.complex)
-    spherical harmonic coefficients of the sky
-beam : numpy.ndarray(numpy.complex)
-    spherical harmonic coefficients of the sky
+sky : numpy.ndarray((nalm_sky, ncomp), dtype=numpy.complex)
+    spherical harmonic coefficients of the sky. ncomp can be 1 or 3.
+beam : numpy.ndarray((nalm_beam, ncomp), dtype=numpy.complex)
+    spherical harmonic coefficients of the beam. ncomp can be 1 or 3
+separate : bool
+    whether contributions of individual components should be added together.
 lmax : int
     maximum l in the coefficient arays
-mmax : int
-    maximum m in the sky coefficients
 kmax : int
     maximum azimuthal moment in the beam coefficients
 epsilon : float
@@ -121,10 +141,11 @@ Parameters
 ----------
 lmax : int
     maximum l in the coefficient arays
-mmax : int
-    maximum m in the sky coefficients
 kmax : int
     maximum azimuthal moment in the beam coefficients
+ncomp : int
+    the number of components which are going to input to `deinterpol`.
+    Can be 1 or 3.
 epsilon : float
     desired accuracy for the interpolation; a typical value is 1e-5
 nthreads : the number of threads to use for computation
@@ -135,7 +156,7 @@ Computes the interpolated values for a given set of angle triplets
 
 Parameters
 ----------
-ptg : numpy.ndarray(numpy.float64) of shape(N,3)
+ptg : numpy.ndarray((N, 3), dtype=numpy.float64)
     theta, phi and psi angles (in radian) for N pointings
     theta must be in the range [0; pi]
     phi must be in the range [0; 2pi]
@@ -143,8 +164,10 @@ ptg : numpy.ndarray(numpy.float64) of shape(N,3)
 
 Returns
 -------
-numpy.array(numpy.float64) of shape (N,)
+numpy.array((N, n2), dtype=numpy.float64)
     the interpolated values
+    n2 is either 1 (if separate=True was used in the constructor) or the
+    second dimension of the input slm and blm arrays (otherwise)
 
 Notes
 -----
@@ -159,13 +182,14 @@ data cube.
 
 Parameters
 ----------
-ptg : numpy.ndarray(numpy.float64) of shape(N,3)
+ptg : numpy.ndarray((N,3), dtype=numpy.float64)
     theta, phi and psi angles (in radian) for N pointings
     theta must be in the range [0; pi]
     phi must be in the range [0; 2pi]
     psi should be in the range [-2pi; 2pi]
-data : numpy.ndarray(numpy.float64) of shape (N,)
+data : numpy.ndarray((N, n2), dtype=numpy.float64)
     the interpolated values
+    n2 must match the `ncomp` value specified in the constructor.
 
 Notes
 -----
@@ -176,18 +200,19 @@ Notes
 
 constexpr const char *getSlm_DS = R"""(
 Returns a set of sky spherical hamonic coefficients resulting from adjoint
-interplation
+interpolation
 
 Parameters
 ----------
-blmT : numpy.array(numpy.complex)
+beam : numpy.array(nalm_beam, nbeam), dtype=numpy.complex)
     spherical harmonic coefficients of the beam with lmax and kmax defined
     in the constructor call
+    nbeam must match the ncomp specified in the constructor, unless ncomp was 1.
 
 Returns
 -------
-numpy.array(numpy.complex):
-    spherical harmonic coefficients of the sky with lmax and mmax defined
+numpy.array(nalm_sky, nbeam), dtype=numpy.complex):
+    spherical harmonic coefficients of the sky with lmax defined
     in the constructor call
 
 Notes
@@ -204,17 +229,17 @@ 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,
+  py::class_<PyInterpolator<fptype>> (m, "PyInterpolator", pyinterpolator_DS)
+    .def(py::init<const py::array &, const py::array &, bool, int64_t, int64_t, fptype, 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 ("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);
+    .def(py::init<int64_t, int64_t, int64_t, fptype, int>(), initadjoint_DS,
+      "lmax"_a, "kmax"_a, "ncomp"_a, "epsilon"_a, "nthreads"_a)
+    .def ("interpol", &PyInterpolator<fptype>::pyinterpol, interpol_DS, "ptg"_a)
+    .def ("deinterpol", &PyInterpolator<fptype>::pydeinterpol, deinterpol_DS, "ptg"_a, "data"_a)
+    .def ("getSlm", &PyInterpolator<fptype>::pygetSlm, getSlm_DS, "beam"_a);
 #if 1
-  m.def("rotate_alm", &pyrotate_alm<double>, "alm"_a, "lmax"_a, "psi"_a, "theta"_a,
+  m.def("rotate_alm", &pyrotate_alm<fptype>, "alm"_a, "lmax"_a, "psi"_a, "theta"_a,
     "phi"_a);
 #endif
   }

pyinterpol_ng/test/test_interpol_ng.py

@@ -22,11 +22,11 @@ def nalm(lmax, mmax):
     return ((mmax+1)*(mmax+2))//2 + (mmax+1)*(lmax-mmax)
 
 
-def random_alm(lmax, mmax):
-    res = np.random.uniform(-1., 1., nalm(lmax, mmax)) \
-     + 1j*np.random.uniform(-1., 1., nalm(lmax, mmax))
+def random_alm(lmax, mmax, ncomp):
+    res = np.random.uniform(-1., 1., (nalm(lmax, mmax), ncomp)) \
+     + 1j*np.random.uniform(-1., 1., (nalm(lmax, mmax), ncomp))
     # make a_lm with m==0 real-valued
-    res[0:lmax+1].imag = 0.
+    res[0:lmax+1,:].imag = 0.
     return res
 
 
@@ -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
-    slmT = random_alm(lmax, lmax)
-    blmT = random_alm(lmax, kmax)
+    slm = random_alm(lmax, lmax, ncomp)
+    blm = random_alm(lmax, kmax, ncomp)
 
-    inter = pyinterpol_ng.PyInterpolator(slmT, blmT, 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)
-    blmT2 = np.zeros(nalm(lmax,lmax))+0j
-    blmT2[0:blmT.shape[0]] = blmT
-    for i in range(nptg):
-        rbeam=pyinterpol_ng.rotate_alm(blmT2, lmax, ptg[i,2],ptg[i,0],ptg[i,1])
-        res2[i] = convolve(slmT, rbeam, lmax).real
-    _assert_close(res1, res2, 1e-7)
-
-@pmp("lkmax", [(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
-    slmT = random_alm(lmax, lmax)
-    blmT = random_alm(lmax, kmax)
+    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(slmT,blmT,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])
-    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(blmT)
-    _assert_close(myalmdot(slmT, bla, lmax, lmax, 0), np.vdot(fake,inter1), 1e-12)
+    bla=foo2.getSlm(blm)
+    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"
@@ -224,7 +224,7 @@ template<size_t ndim> class mav_info
 
 template<typename T, size_t ndim> class mav: public mav_info<ndim>, public membuf<T>
   {
-  static_assert((ndim>0) && (ndim<4), "only supports 1D, 2D, and 3D arrays");
+//  static_assert((ndim>0) && (ndim<4), "only supports 1D, 2D, and 3D arrays");
 
   protected:
     using typename mav_info<ndim>::shape_t;