Commit d5391c22 authored by Martin Reinecke's avatar Martin Reinecke
Browse files

add interpolator module (by no means complete or working)

parent 124971a9
Pipeline #69869 passed with stages
in 8 minutes and 58 seconds
/*! \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
/*
* 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);
}
../src/mr_util/
\ No newline at end of file
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']
)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment