add interpolator module (by no means complete or working)

interpol_ng/alm.h

+/*! \file alm.h
+ *  Class for storing spherical harmonic coefficients.
+ *
+ *  Copyright (C) 2003-2020 Max-Planck-Society
+ *  \author Martin Reinecke
+ */
+
+#ifndef MRUTIL_ALM_H
+#define MRUTIL_ALM_H
+
+#include "mr_util/infra/mav.h"
+#include "mr_util/infra/error_handling.h"
+
+namespace mr {
+
+/*! Base class for calculating the storage layout of spherical harmonic
+    coefficients. */
+class Alm_Base
+  {
+  protected:
+    size_t lmax, mmax, tval;
+
+  public:
+    /*! Returns the total number of coefficients for maximum quantum numbers
+        \a l and \a m. */
+    static size_t Num_Alms (size_t l, size_t m)
+      {
+      MR_assert(m<=l,"mmax must not be larger than lmax");
+      return ((m+1)*(m+2))/2 + (m+1)*(l-m);
+      }
+
+    /*! Constructs an Alm_Base object with given \a lmax and \a mmax. */
+    Alm_Base (size_t lmax_=0, size_t mmax_=0)
+      : lmax(lmax_), mmax(mmax_), tval(2*lmax+1) {}
+
+    /*! Returns the maximum \a l */
+    size_t Lmax() const { return lmax; }
+    /*! Returns the maximum \a m */
+    size_t Mmax() const { return mmax; }
+
+    /*! Returns an array index for a given m, from which the index of a_lm
+        can be obtained by adding l. */
+    size_t index_l0 (size_t m) const
+      { return ((m*(tval-m))>>1); }
+
+    /*! Returns the array index of the specified coefficient. */
+    size_t index (size_t l, size_t m) const
+      { return index_l0(m) + l; }
+
+    /*! Returns \a true, if both objects have the same \a lmax and \a mmax,
+        else  \a false. */
+    bool conformable (const Alm_Base &other) const
+      { return ((lmax==other.lmax) && (mmax==other.mmax)); }
+  };
+
+/*! Class for storing spherical harmonic coefficients. */
+template<typename T> class Alm: public Alm_Base
+  {
+  private:
+    mav<T,1> alm;
+
+  public:
+    /*! Constructs an Alm object with given \a lmax and \a mmax. */
+    Alm (mav<T,1> &data, size_t lmax_, size_t mmax_)
+      : Alm_Base(lmax_, mmax_), alm(data)
+      { MR_assert(alm.size()==Num_Alms(lmax, mmax), "bad array size"); }
+    Alm (const mav<T,1> &data, size_t lmax_, size_t mmax_)
+      : Alm_Base(lmax_, mmax_), alm(data)
+      { MR_assert(alm.size()==Num_Alms(lmax, mmax), "bad array size"); }
+    Alm (size_t lmax_=0, size_t mmax_=0)
+      : Alm_Base(lmax_,mmax_), alm ({Num_Alms(lmax,mmax)}) {}
+
+    /*! Sets all coefficients to zero. */
+    void SetToZero ()
+      { alm.fill(0); }
+
+    /*! Multiplies all coefficients by \a factor. */
+    template<typename T2> void Scale (const T2 &factor)
+      { for (size_t m=0; m<alm.size(); ++m) alm.v(m)*=factor; }
+    /*! \a a(l,m) *= \a factor[l] for all \a l,m. */
+    template<typename T2> void ScaleL (const mav<T2,1> &factor)
+      {
+      MR_assert(factor.size()>size_t(lmax),
+        "alm.ScaleL: factor array too short");
+      for (size_t m=0; m<=mmax; ++m)
+        for (size_t l=m; l<=lmax; ++l)
+          operator()(l,m)*=factor(l);
+      }
+    /*! \a a(l,m) *= \a factor[m] for all \a l,m. */
+    template<typename T2> void ScaleM (const mav<T2,1> &factor)
+      {
+      MR_assert(factor.size()>size_t(mmax),
+        "alm.ScaleM: factor array too short");
+      for (size_t m=0; m<=mmax; ++m)
+        for (size_t l=m; l<=lmax; ++l)
+          operator()(l,m)*=factor(m);
+      }
+    /*! Adds \a num to a_00. */
+    template<typename T2> void Add (const T2 &num)
+      { alm.v(0)+=num; }
+
+    /*! Returns a reference to the specified coefficient. */
+    T &operator() (size_t l, size_t m)
+      { return alm.v(index(l,m)); }
+    /*! Returns a constant reference to the specified coefficient. */
+    const T &operator() (size_t l, size_t m) const
+      { return alm(index(l,m)); }
+
+    /*! Returns a pointer for a given m, from which the address of a_lm
+        can be obtained by adding l. */
+    T *mstart (size_t m)
+      { return &alm.v(index_l0(m)); }
+    /*! Returns a pointer for a given m, from which the address of a_lm
+        can be obtained by adding l. */
+    const T *mstart (size_t m) const
+      { return &alm(index_l0(m)); }
+
+    /*! Returns a constant reference to the a_lm data. */
+    const mav<T,1> &Alms() const { return alm; }
+
+    ptrdiff_t stride() const
+      { return alm.stride(0); }
+
+    /*! Adds all coefficients from \a other to the own coefficients. */
+    void Add (const Alm &other)
+      {
+      MR_assert (conformable(other), "A_lm are not conformable");
+      for (size_t m=0; m<alm.size(); ++m)
+        alm.v(m) += other.alm(m);
+      }
+  };
+
+}
+
+#endif

interpol_ng/interpol_ng.cc

+/*
+ *  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"
+#include <iostream>
+using namespace std;
+using namespace mr;
+
+namespace py = pybind11;
+
+namespace {
+
+template<typename T> class Interpolator
+  {
+  protected:
+    size_t supp;
+    size_t lmax, kmax, ppring, ntheta, nphi;
+    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});
+      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
+      for (size_t i=1, i2=2*ntheta-3; i+1<ntheta; ++i,--i2)
+        for (size_t j=0,j2=nphi/2; j<nphi; ++j,++j2)
+          {
+          if (j2>=nphi) j2-=nphi;
+          tmp.v(i2,j) = arr(i,j2);
+          }
+      // FFT
+      r2r_fftpack(ftmp,ftmp,{0,1},true,true,1.,0);
+      ES_Kernel kernel(supp,nphi/double(2*lmax+1),0);
+      auto fct = kernel.correction_factors(nphi, nphi/2+2, 0);
+      for (size_t i=0; i<nphi; ++i)
+        for (size_t j=0; j<nphi; ++j)
+          tmp.v(i,j) *= fct[(i+1)/2] * fct[(j+1)/2];
+      r2r_fftpack(ftmp,ftmp,{0,1},false,false,1./(nphi*nphi),0);
+      for (size_t i=0; i<ntheta; ++i)
+        for (size_t j=0; j<nphi; ++j)
+          arr.v(i,j) = tmp(i,j);
+      }
+
+  public:
+    Interpolator(const Alm<complex<T>> &slmT, const Alm<complex<T>> &blmT,
+      double epsilon)
+      : supp(ES_Kernel::get_supp(epsilon)),
+        lmax(slmT.Lmax()),
+        kmax(blmT.Mmax()),
+        ppring(2*good_size_real(2*lmax+1)),
+        ntheta(ppring/2+1),
+        nphi(ppring),
+        kernel(supp, 2, 1),
+        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);
+      auto ginfo = sharp_make_cc_geom_info(ntheta,nphi,0,cube.stride(1),cube.stride(0));
+      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);
+        mav<T,2> m1(&cube.v(supp,supp,kidx1),{ntheta,nphi},{cube.stride(0),cube.stride(1)},true);
+        mav<T,2> m2(&cube.v(supp,supp,kidx2),{ntheta,nphi},{cube.stride(0),cube.stride(1)},true);
+        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);
+      for (size_t i=0; i<ptg.shape(0); ++i)
+        {
+        double theta=ptg(i,0);
+        double phi=ptg(i,1);
+        double psi=ptg(i,2);
+        double f0=supp+theta*((ntheta-1)/pi)-0.5*supp;
+        size_t i0 = size_t(f0+1.);
+        for (size_t t=0; t<supp; ++t)
+          wt[t] = kernel(((t+i0)-f0-0.5*supp)/supp);
+        double f1=supp+phi*(nphi/(2*pi))-0.5*supp;
+        size_t i1 = size_t(f1+1.);
+        for (size_t t=0; t<supp; ++t)
+          wp[t] = kernel(((t+i1)-f1-0.5*supp)/supp);
+        double val=0;
+        psiarr[0]=1.;
+        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;
+          }
+        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)
+              val += cube(i0+j,i1+k,0)*wt[j]*wp[k]*psiarr[l];
+        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);
+  }

interpol_ng/mr_util

+../src/mr_util/
\ No newline at end of file

interpol_ng/setup.py

+from setuptools import setup, Extension
+import sys
+
+
+class _deferred_pybind11_include(object):
+    def __init__(self, user=False):
+        self.user = user
+
+    def __str__(self):
+        import pybind11
+        return pybind11.get_include(self.user)
+
+
+include_dirs = ['.', _deferred_pybind11_include(True),
+                _deferred_pybind11_include()]
+extra_compile_args = ['--std=c++17', '-march=native', '-ffast-math', '-O3']
+python_module_link_args = []
+define_macros = []
+
+if sys.platform == 'darwin':
+    import distutils.sysconfig
+    extra_compile_args += ['-mmacosx-version-min=10.9']
+    python_module_link_args += ['-mmacosx-version-min=10.9', '-bundle']
+    vars = distutils.sysconfig.get_config_vars()
+    vars['LDSHARED'] = vars['LDSHARED'].replace('-bundle', '')
+elif sys.platform == 'win32':
+    extra_compile_args = ['/Ox', '/EHsc', '/std:c++17']
+else:
+    extra_compile_args += ['-Wfatal-errors', '-Wfloat-conversion', '-W', '-Wall', '-Wstrict-aliasing=2', '-Wwrite-strings', '-Wredundant-decls', '-Woverloaded-virtual', '-Wcast-qual', '-Wcast-align', '-Wpointer-arith']
+    python_module_link_args += ['-march=native', '-ffast-math', '-Wl,-rpath,$ORIGIN']
+
+# if you don't want debugging info, add "-s" to python_module_link_args
+
+def get_extension_modules():
+    return [Extension('interpol_ng',
+                      language='c++',
+                      sources=['interpol_ng.cc',
+                               'mr_util/infra/threading.cc',
+                               'mr_util/sharp/sharp.cc',
+                               'mr_util/sharp/sharp_core.cc',
+                               'mr_util/sharp/sharp_geomhelpers.cc',
+                               'mr_util/sharp/sharp_almhelpers.cc',
+                               'mr_util/sharp/sharp_ylmgen.cc'],
+                      depends=['setup.py'],
+                      include_dirs=include_dirs,
+                      define_macros=define_macros,
+                      extra_compile_args=extra_compile_args,
+                      extra_link_args=python_module_link_args)]
+
+
+setup(name='interpol_ng',
+      version='0.0.1',
+      description='Python bindings for some libsharp functionality',
+      include_package_data=True,
+      author='Martin Reinecke',
+      author_email='martin@mpa-garching.mpg.de',
+      packages=[],
+      setup_requires=['numpy>=1.15.0', 'pybind11>=2.2.4'],
+      ext_modules=get_extension_modules(),
+      install_requires=['numpy>=1.15.0', 'pybind11>=2.2.4']
+      )