cosmotool/src/fourier/healpix.hpp

#ifndef __COSMOTOOL_FOURIER_HEALPIX_HPP
#define __COSMOTOOL_FOURIER_HEALPIX_HPP

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <cmath>
#include <vector>
#include <boost/bind.hpp>
#include <boost/shared_ptr.hpp>
#include <exception>
#include "base_types.hpp"
#include <sharp_lowlevel.h>
#include <sharp_geomhelpers.h>
#include <sharp_almhelpers.h>

namespace CosmoTool
{
  template<typename T>
  class HealpixSpectrum: public SpectrumFunction<T>
  {
  protected:
    std::vector<T> cls;
  public:
    typedef typename SpectrumFunction<T>::FourierMapType FourierMapType;
    typedef boost::shared_ptr<FourierMapType> ptr_map;


    HealpixSpectrum(long Lmax)
     : cls(Lmax+1) {}

    T *data() { return &cls[0]; }
    const T *data() const { return &cls[0]; }
    long size() const { return cls.size(); }

    void newRandomFourier(gsl_rng *rng, FourierMapType& like_map) const; 

    void mul(FourierMapType& m) const;
    void mul_sqrt(FourierMapType& m) const;
    void mul_inv(FourierMapType& m) const;
    void mul_inv_sqrt(FourierMapType& m) const;
  };

  template<typename T>
  class HealpixFourierMap: public FourierMap<T>
  {
  private:
    T *m_data;
    long Npix, m_Nside;
    Eigen::aligned_allocator<T> alloc;
  public:
    HealpixFourierMap(long nSide)
      : Npix(12*nSide*nSide), m_Nside(nSide)
    {
      m_data = alloc.allocate(Npix);
    }

    virtual ~HealpixFourierMap()
    {
      alloc.deallocate(m_data, Npix);
    }

    long Nside() const { return m_Nside; }
    virtual const T* data() const { return m_data; }
    virtual T *data() { return m_data; }
    virtual long size() const { return Npix; }

    virtual T dot_product(const FourierMap<T>& other) const
      throw(std::bad_cast)
    {
      typedef typename FourierMap<T>::MapType MapType;

      const HealpixFourierMap<T>& mfm = dynamic_cast<const HealpixFourierMap<T>&>(other);
      if (Npix != mfm.size())
        throw std::bad_cast();

      MapType m1(m_data, Npix);
      MapType m2(mfm.m_data, mfm.Npix);
      
      return (m1*m2).sum();
    }
    
    virtual FourierMap<T> *mimick() const
    {
      return new HealpixFourierMap<T>(m_Nside);
    }
  };


  template<typename T>
  class HealpixFourierALM: public FourierMap<std::complex<T> >
  {
  private:
    std::complex<T> *alms;
    long m_size;
    long Lmax_, Mmax_, TVal_;
    Eigen::aligned_allocator<std::complex<T> > alloc;
  public:
    typedef unsigned long LType;

    LType Lmax() const { return Lmax_; }
    LType Mmax() const { return Mmax_; }

    LType Num_Alms() const
    {
      return ((Mmax_+1)*(Mmax_+2))/2 + (Mmax_+1)*(Lmax_-Mmax_);
    }

    LType index_l0(LType m) const
    {
      return ((m*(TVal_-m))/2);
    }

    LType index(LType l, LType m) const
    {
      return index_l0(m) + l;
    }
    
    HealpixFourierALM(LType lmax, LType mmax)
      : Lmax_(lmax), Mmax_(mmax), TVal_(2*lmax+1)
    {
      m_size = Num_Alms();
      alms = alloc.allocate(m_size);
    }

    virtual ~HealpixFourierALM()
    {
      alloc.deallocate(alms, m_size);
    }

    virtual const std::complex<T>* data() const { return alms; }
    virtual std::complex<T> * data() { return alms;} 
    virtual long size() const { return m_size; }

    virtual FourierMap<std::complex<T> > *mimick() const
    {
      return new HealpixFourierALM<T>(Lmax_, Mmax_);
    }

    virtual std::complex<T> dot_product(const FourierMap<std::complex<T> >& other) const
      throw(std::bad_cast)
    {
      const HealpixFourierALM<T>& mfm = dynamic_cast<const HealpixFourierALM<T>&>(other);
      typedef typename FourierMap<std::complex<T> >::MapType MapType;
      std::complex<T> S;

      if (m_size != mfm.m_size)
        throw std::bad_cast();

      MapType m1(alms, m_size);
      MapType m2(mfm.alms, mfm.m_size);
      
      S = (m1.block(0,0,1,Lmax_+1).conjugate() * m2.block(0,0,1,Lmax_+1)).sum();
      S += std::complex<T>(2,0)*(m1.block(0,1+Lmax_,1,m_size-1-Lmax_).conjugate() * m2.block(0,1+Lmax_,1,m_size-1-Lmax_)).sum();
      return S;
    }
  };

  template<typename T> struct HealpixJobHelper__ {};

  template<> struct HealpixJobHelper__<double>
    { enum {val=1}; };

  template<> struct HealpixJobHelper__<float>
    { enum {val=0}; };


  template<typename T>
  class HealpixFourierTransform: public FourierTransform<T>
  {
  private:
    sharp_alm_info *ainfo;
    sharp_geom_info *ginfo;
    HealpixFourierMap<T> realMap;
    HealpixFourierALM<T> fourierMap;
    int m_iterate;
  public:
    HealpixFourierTransform(long nSide, long Lmax, long Mmax, int iterate = 0)
      : realMap(nSide), fourierMap(Lmax, Mmax), ainfo(0), ginfo(0), m_iterate(iterate)
    {
      sharp_make_healpix_geom_info (nSide, 1, &ginfo);
      sharp_make_triangular_alm_info (Lmax, Mmax, 1, &ainfo);
    }

    virtual ~HealpixFourierTransform()
    {
      sharp_destroy_geom_info(ginfo);
      sharp_destroy_alm_info(ainfo);
    }

    virtual const FourierMap<std::complex<T> >& fourierSpace() const { return fourierMap; }

    virtual FourierMap<std::complex<T> >& fourierSpace() { return fourierMap; }

    virtual const FourierMap<T>& realSpace() const { return realMap; }

    virtual FourierMap<T>& realSpace() { return realMap; }
    
    virtual FourierTransform<T> *mimick() const 
    {
      return new HealpixFourierTransform<T>(realMap.Nside(), fourierMap.Lmax(), fourierMap.Mmax());
    }

    virtual void analysis()
    {
      void *aptr=reinterpret_cast<void *>(fourierMap.data()), *mptr=reinterpret_cast<void *>(realMap.data());

      sharp_execute (SHARP_MAP2ALM, 0, 0, &aptr, &mptr, ginfo, ainfo, 1,
        HealpixJobHelper__<T>::val,0,0,0);
      for (int i = 0; i < m_iterate; i++)
        {
          HealpixFourierMap<T> tmp_map(realMap.Nside());          
          void *tmp_ptr=reinterpret_cast<void *>(tmp_map.data());
          typename HealpixFourierMap<T>::MapType m0 = tmp_map.eigen();
          typename HealpixFourierMap<T>::MapType m1 = realMap.eigen();

          sharp_execute (SHARP_ALM2MAP, 0, 0, &aptr, &tmp_ptr, ginfo, ainfo, 1,
            HealpixJobHelper__<T>::val,0,0,0);
          m0 = m1 - m0;
          sharp_execute (SHARP_MAP2ALM, 0, 1, &aptr, &tmp_ptr, ginfo, ainfo, 1,
            HealpixJobHelper__<T>::val,0,0,0);
        }
    }

    virtual void synthesis()
    {
      void *aptr=reinterpret_cast<void *>(fourierMap.data()), *mptr=reinterpret_cast<void *>(realMap.data());

      sharp_execute (SHARP_ALM2MAP, 0, 0, &aptr, &mptr, ginfo, ainfo, 1,
        HealpixJobHelper__<T>::val,0,0,0);
    }

    virtual void analysis_conjugate()
    {
      synthesis();
      realMap.scale(4*M_PI/realMap.size());
    }

    virtual void synthesis_conjugate()
    {
      analysis();
      fourierMap.scale(realMap.size()/(4*M_PI));
    }
    
  };

  template<typename T>
  void HealpixSpectrum<T>::newRandomFourier(gsl_rng *rng, FourierMapType& out_map) const
  { 
    HealpixFourierALM<T>& alms = dynamic_cast<HealpixFourierALM<T>&>(out_map);
    long lmaxGen = std::min(cls.size()-1, alms.Lmax());
    std::complex<T> *new_data = alms.data();

    for (long l = 0; l <= lmaxGen; l++)
      {
        double Al = std::sqrt(cls[l]);

        new_data[alms.index(l,0)] = gsl_ran_gaussian(rng, Al);
        Al *= M_SQRT1_2;
        for (long m = 1; m <= std::min(l,alms.Mmax()); m++)
          {
            std::complex<T>& c = new_data[alms.index(l,m)];
            c.real() = gsl_ran_gaussian(rng, Al); 
            c.imag() = gsl_ran_gaussian(rng, Al); 
          }
      }
  } 

 template<typename T>
 void HealpixSpectrum<T>::mul(FourierMapType& like_map) const
  {
    HealpixFourierALM<T>& alms = dynamic_cast<HealpixFourierALM<T>&>(like_map);
    std::complex<T> *data = alms.data();

    for (long l = 0; l <= alms.Lmax(); l++)
      {
        double Al = cls[l];

        for (long m = 0; m <= std::min(l,alms.Mmax()); m++)
          {
            data[alms.index(l,m)] *= Al;
          }
      }
  }

 template<typename T>
 void HealpixSpectrum<T>::mul_sqrt(FourierMapType& like_map) const
  {
    HealpixFourierALM<T>& alms = dynamic_cast<const HealpixFourierALM<T>&>(like_map);
    std::complex<T> *data = alms.data();

    for (long l = 0; l <= alms.Lmax(); l++)
      {
        double Al = std::sqrt(cls[l]);

        for (long m = 0; m <= std::min(l,alms.Mmax()); m++)
          {
            data[alms.index(l,m)] *= Al;
          }
      }
  }

 template<typename T>
 void HealpixSpectrum<T>::mul_inv(FourierMapType& like_map) const
  {
    HealpixFourierALM<T>& alms = dynamic_cast<HealpixFourierALM<T>&>(like_map);
    std::complex<T> *data = alms.data();

    for (long l = 0; l <= alms.Lmax(); l++)
      {
        double Al = (cls[l] <= 0) ? 0 : (1/cls[l]);

        for (long m = 0; m <= std::min(l,alms.Mmax()); m++)
          {
            data[alms.index(l,m)] *= Al;
          }
      }
  }

 template<typename T>
 void HealpixSpectrum<T>::mul_inv_sqrt(FourierMapType& like_map) const
  {
    HealpixFourierALM<T>& alms = dynamic_cast<HealpixFourierALM<T>&>(like_map);
    std::complex<T> *data = alms.data();

    for (long l = 0; l <= alms.Lmax(); l++)
      {
        double Al = (cls[l] <= 0) ? 0 : std::sqrt(1/cls[l]);

        for (long m = 0; m <= std::min(l,alms.Mmax()); m++)
          {
            data[alms.index(l,m)] *= Al;
          }
      }
  }

};

#endif