cosmotool/src/fourier/euclidian.hpp

#ifndef __COSMOTOOL_FOURIER_EUCLIDIAN_HPP
#define __COSMOTOOL_FOURIER_EUCLIDIAN_HPP

#include <cmath>
#include <boost/function.hpp>
#include <vector>
#include <boost/shared_ptr.hpp>
#include <gsl/gsl_randist.h>
#include "base_types.hpp"
#include "fft/fftw_calls.hpp"
#include "../algo.hpp"

namespace CosmoTool
{
  template<typename T>
  class EuclidianOperator
  {
  public:
    typedef boost::function1<T, T> Function;
    
    Function base, op;
    T operator()(T k) {
      return op(base(k));
    }
  };

  template<typename T>
  class EuclidianSpectrum_1D: public SpectrumFunction<T>
  {
  public:
    typedef boost::function1<T, T> Function;
  protected:
    Function f;

    static T msqrt(T a) { return std::sqrt(a); }
  public:
    typedef typename SpectrumFunction<T>::FourierMapType FourierMapType;
    typedef typename SpectrumFunction<T>::SpectrumFunctionPtr SpectrumFunctionPtr;
    typedef boost::shared_ptr<FourierMapType> ptr_map;

    EuclidianSpectrum_1D(Function P)
      : f(P)
    {
    }

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

    SpectrumFunctionPtr copy() const {
      return SpectrumFunctionPtr(new EuclidianSpectrum_1D(f));
    }

    void sqrt() {
      EuclidianOperator<T> o;
      o.base = f;
      o.op = &EuclidianSpectrum_1D<T>::msqrt;
      f = (Function(o));
    }

    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 EuclidianFourierMapBase: public FourierMap<T>
  {
  public:
    typedef std::vector<int> DimArray;
  private:
    boost::shared_ptr<T> m_data;
    DimArray m_dims;
    long m_size;
  public:
    
    EuclidianFourierMapBase(boost::shared_ptr<T> indata, const DimArray& indims)
    {
      m_data = indata;
      m_dims = indims;
      m_size = 1;
      for (int i = 0; i < m_dims.size(); i++)
	m_size *= m_dims[i];
    }

    virtual ~EuclidianFourierMapBase()
    {
    }

    const DimArray& getDims() const { return m_dims; }

    virtual const T *data() const { return m_data.get(); }
    virtual T *data() { return m_data.get(); }
    virtual long size() const  { return m_size; } 
    
    virtual FourierMap<T> *copy() const
    {
      FourierMap<T> *m = this->mimick();
      m->eigen() = this->eigen();
      return m;
    }

  };

  template<typename T>
  class EuclidianFourierMapReal: public EuclidianFourierMapBase<T>
  {
  public:
    typedef typename EuclidianFourierMapBase<T>::DimArray DimArray;

    EuclidianFourierMapReal(boost::shared_ptr<T> indata, const DimArray& indims)
      : EuclidianFourierMapBase<T>(indata, indims) 
    {}

    virtual FourierMap<T> *mimick() const
    {
      return new EuclidianFourierMapReal<T>(
		      boost::shared_ptr<T>((T *)fftw_malloc(sizeof(T)*this->size()), 
					   std::ptr_fun(fftw_free)),
		      this->getDims());
    }

    virtual T dot_product(const FourierMap<T>& other) const
      throw(std::bad_cast)
    {
      const EuclidianFourierMapReal<T>& m2 = dynamic_cast<const EuclidianFourierMapReal<T>&>(other);
      if (this->size() != m2.size())
        throw std::bad_cast();

      return (this->eigen()*m2.eigen()).sum();
    }
  };

  template<typename T>
  class EuclidianFourierMapComplex: public EuclidianFourierMapBase<std::complex<T> >
  {
  protected:
    typedef boost::shared_ptr<std::complex<T> > ptr_t;
    std::vector<double> delta_k;
    int m_dim0;
    bool even0, alleven;
    long plane_size;
  public:
    typedef typename EuclidianFourierMapBase<std::complex<T> >::DimArray DimArray;

    EuclidianFourierMapComplex(ptr_t indata,
			       int dim0,
			       const DimArray& indims, 
			       const std::vector<double>& dk)
      : EuclidianFourierMapBase<std::complex<T> >(indata, indims), delta_k(dk), m_dim0(dim0), even0((dim0 % 2)==0)
    {
      assert(dk.size() == indims.size()); 
      plane_size = 1;
      alleven = true;
      for (int q = 1; q < indims.size(); q++)
	{
	  plane_size *= indims[q];
	  alleven = alleven && ((indims[q]%2)==0);
	}
    }

    virtual FourierMap<std::complex<T> > *mimick() const
    {
      return 
	new EuclidianFourierMapComplex<T>(
		   ptr_t((std::complex<T> *)
			 fftw_malloc(sizeof(std::complex<T>)*this->size()), 
			 std::ptr_fun(fftw_free)),
		   m_dim0,
		   this->getDims(),
		   this->delta_k);
    }

    template<typename Array>
    double get_K(const Array& ik)
    {
      const DimArray& dims = this->getDims();
      assert(ik.size() == dims.size());
      double k2 = 0;
      k2 += CosmoTool::square(ik[0]*delta_k[0]);

      for (int q = 1; q < ik.size(); q++)
	{
	  int dk = ik[q];

	  if (dk > dims[q]/2)
	    dk = dk - dims[q];
	  
	  k2 += CosmoTool::square(delta_k[q]*dk);
	}
      return std::sqrt(k2);
    }

    double get_K(long p)
    {
      const DimArray& dims = this->getDims();
      DimArray d(delta_k.size());
      for (int q = 0; q < d.size(); q++)
       {
         d[q] = p%dims[q];
         p = (p-d[q])/dims[q];
       }
      return get_K(d);
    }

    bool allDimensionsEven() const { return alleven; }
    bool firstDimensionEven() const { return even0; }

    virtual std::complex<T> dot_product(const FourierMap<std::complex<T> >& other) const
      throw(std::bad_cast)
    {
      const EuclidianFourierMapComplex<T>& m2 = dynamic_cast<const EuclidianFourierMapComplex<T>&>(other);
      if (this->size() != m2.size())
        throw std::bad_cast();
      
      const std::complex<T> *d1 = this->data();
      const std::complex<T> *d2 = m2.data();
      const DimArray& dims = this->getDims();
      int N0 = dims[0] + (even0 ? 0 : 1);
      std::complex<T> result = 0;

      for (long q0 = 1; q0 < N0-1; q0++)
	{
	  for (long p = 0; p < plane_size; p++)
	    {
	      long idx = q0+dims[0]*p;
	      assert(idx < this->size());
	      result += 2*(conj(d1[idx]) * d2[idx]).real();
	    }
	}
      if (even0)
	{
	  for (long p = 0; p < plane_size; p++)
	    {
	      long q0 = N0*p, q1 = (p+1)*N0-1;
	      result += conj(d1[q0]) * d2[q0];
	      result += conj(d1[q1]) * d2[q1];
	    }
	}
      return result;
    }

  };
  
  template<typename T>
  class EuclidianFourierTransform: public FourierTransform<T>
  {
  public:
    typedef typename EuclidianFourierMapBase<T>::DimArray DimArray;
  private:
    typedef FFTW_Calls<T> calls;
    EuclidianFourierMapReal<T> *realMap;
    EuclidianFourierMapComplex<T> *fourierMap;
    typename calls::plan_type m_analysis, m_synthesis;
    double volume;
    long N, Nc;
    DimArray m_dims, m_dims_hc;
    std::vector<double> m_L;
  public:
    EuclidianFourierTransform(const DimArray& dims, const std::vector<double>& L)
    {
      assert(L.size() == dims.size());
      std::vector<double> dk(L.size());
      std::vector<int> swapped_dims(dims.size());
      
      m_dims = dims;
      m_dims_hc = dims;
      m_dims_hc[0] = dims[0]/2+1;
      m_L = L;

      N = 1;
      Nc = 1;
      volume = 1;
      for (int i = 0; i < dims.size(); i++)
	{
	  N *= dims[i];
          Nc *= m_dims_hc[i];
	  volume *= L[i];
	  dk[i] = 2*M_PI/L[i];
          swapped_dims[dims.size()-1-i] = dims[i];
	}

      realMap = new EuclidianFourierMapReal<T>(
		   boost::shared_ptr<T>(calls::alloc_real(N),
					std::ptr_fun(calls::free)),
		   m_dims);
      fourierMap = new EuclidianFourierMapComplex<T>(
		   boost::shared_ptr<std::complex<T> >((std::complex<T>*)calls::alloc_complex(Nc),
						       std::ptr_fun(calls::free)), 
		   dims[0], m_dims_hc, dk);
      m_analysis = calls::plan_dft_r2c(dims.size(), &swapped_dims[0],
				       realMap->data(), (typename calls::complex_type *)fourierMap->data(),
				       FFTW_DESTROY_INPUT|FFTW_MEASURE);
      m_synthesis = calls::plan_dft_c2r(dims.size(), &swapped_dims[0],
					(typename calls::complex_type *)fourierMap->data(), realMap->data(),
					FFTW_DESTROY_INPUT|FFTW_MEASURE);
    }
    
    virtual ~EuclidianFourierTransform()
    {
      delete realMap;
      delete fourierMap;
      calls::destroy_plan(m_synthesis);
      calls::destroy_plan(m_analysis);
    }

    void synthesis()
    {
      calls::execute(m_synthesis);
      realMap->scale(1/volume);
    }
    
    void analysis()
    {
      calls::execute(m_analysis);
      fourierMap->scale(volume/N);
    }
    
    void synthesis_conjugate()
    {
      calls::execute(m_analysis);
      fourierMap->scale(1/volume);
    }
    
    void analysis_conjugate()
    {
      calls::execute(m_synthesis);
      realMap->scale(volume/N);
    }

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

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

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

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

    FourierTransform<T> *mimick() const
    {
      return new EuclidianFourierTransform(m_dims, m_L);
    }
  };

  template<typename T>
  class EuclidianFourierTransform_1d: public EuclidianFourierTransform<T>
  {
  private:
    template<typename T2>
    static std::vector<T2> make_1d_vector(T2 a)
    {
      T2 arr[2] = { a};
      return std::vector<T2>(&arr[0],&arr[1]);
    }
  public:
    EuclidianFourierTransform_1d(int Nx, double Lx)
      : EuclidianFourierTransform<T>(make_1d_vector<int>(Nx), make_1d_vector<double>(Lx))
    {
    }

    virtual ~EuclidianFourierTransform_1d() {}

  };

  template<typename T>
  class EuclidianFourierTransform_2d: public EuclidianFourierTransform<T>
  {
  private:
    template<typename T2>
    static std::vector<T2> make_2d_vector(T2 a, T2 b)
    {
      T2 arr[2] = { a, b};
      return std::vector<T2>(&arr[0], &arr[2]);
    }
  public:
    EuclidianFourierTransform_2d(int Nx, int Ny, double Lx, double Ly)
      : EuclidianFourierTransform<T>(make_2d_vector<int>(Nx, Ny), make_2d_vector<double>(Lx, Ly))
    {
    }

    virtual ~EuclidianFourierTransform_2d() {}

  };

  template<typename T>
  class EuclidianFourierTransform_3d: public EuclidianFourierTransform<T>
  {
  private:
    template<typename T2>
    static std::vector<T2> make_3d_vector(T2 a, T2 b, T2 c)
    {
      T2 arr[2] = { a, b, c};
      return std::vector<T2>(&arr[0], &arr[3]);
    }

  public:
    EuclidianFourierTransform_3d(int Nx, int Ny, int Nz, double Lx, double Ly, double Lz)
      : EuclidianFourierTransform<T>(make_3d_vector<int>(Nx, Ny, Nz), make_3d_vector<double>(Lx, Ly, Lz))
    {
    }

    virtual ~EuclidianFourierTransform_3d() {}
  };


  template<typename T>
  void EuclidianSpectrum_1D<T>::newRandomFourier(gsl_rng *rng, FourierMapType& out_map) const
  {
    typedef EuclidianFourierMapComplex<T> MapT;
    typedef typename EuclidianSpectrum_1D<T>::ptr_map ptr_map;
    typedef typename MapT::DimArray DimArray;

    MapT& rand_map = dynamic_cast<MapT&>(out_map);

    std::complex<T> *d = rand_map.data();
    long idx;
    const DimArray& dims = rand_map.getDims();
    long plane_size;
    bool alleven = rand_map.allDimensionsEven();

    for (long p = 1; p < rand_map.size(); p++)
      {
	double A_k = std::sqrt(0.5*f(rand_map.get_K(p)));
	d[p] = std::complex<T>(gsl_ran_gaussian(rng, A_k),
			       gsl_ran_gaussian(rng, A_k));
      }
    // Generate the mean value
    d[0] = std::complex<T>(gsl_ran_gaussian(rng, std::sqrt(f(0))), 0);

    if (!rand_map.firstDimensionEven())
      return;

    // Correct the Nyquist plane
    idx = dims[0]-1; // Stick to the last element of the first dimension
    d[idx] = std::complex<T>(d[idx].real() + d[idx].imag(), 0);
    // 1D is special case
    if (dims.size() == 1)
      return;
    
    plane_size = 1;
    for (int q = 1; q < dims.size(); q++)
      {
	plane_size *= dims[q];
      }

    for (long p = 1; p < plane_size/2; p++)
      {
	long q = (p+1)*dims[0]-1;
	long q2 = (plane_size-p+1)*dims[0]-1;
	assert(q < plane_size*dims[0]);
	assert(q2 < plane_size*dims[0]);
	d[q] = conj(d[q2]);
      }

    if (alleven)
      {
	long q = 0;
	for (int i = dims.size()-1; i >= 1; i--)
	  q = dims[i]*q + dims[i]/2;
	q += dims[0]-1;
	d[q] = std::complex<T>(d[q].real()+d[q].imag(),0);
      }
  }

  template<typename T>
  void EuclidianSpectrum_1D<T>::mul(FourierMapType& m) const
  {
    EuclidianFourierMapComplex<T>& m_c = dynamic_cast<EuclidianFourierMapComplex<T>&>(m);
    std::complex<T> *d = m.data();

    for (long p = 0; p < m_c.size(); p++)
      d[p] *= f(m_c.get_K(p));
  }

  template<typename T>
  void EuclidianSpectrum_1D<T>::mul_sqrt(FourierMapType& m) const
  {
    EuclidianFourierMapComplex<T>& m_c = dynamic_cast<EuclidianFourierMapComplex<T>&>(m);
    std::complex<T> *d = m.data();

    for (long p = 0; p < m_c.size(); p++)
      d[p] *= std::sqrt(f(m_c.get_K(p)));
  }

  template<typename T>
  void EuclidianSpectrum_1D<T>::mul_inv(FourierMapType& m) const
  {
    EuclidianFourierMapComplex<T>& m_c = dynamic_cast<EuclidianFourierMapComplex<T>&>(m);
    std::complex<T> *d = m.data();

    for (long p = 0; p < m_c.size(); p++)
     {
        T A = f(m_c.get_K(p));
        if (A==0)
          d[p] = 0;
        else
          d[p] /= A;
     }
  }

  template<typename T>
  void EuclidianSpectrum_1D<T>::mul_inv_sqrt(FourierMapType& m) const
  {
    EuclidianFourierMapComplex<T>& m_c = dynamic_cast<EuclidianFourierMapComplex<T>&>(m);
    std::complex<T> *d = m.data();

    for (long p = 0; p < m_c.size(); p++)
      {
        T A = std::sqrt(f(m_c.get_K(p)));
        if (A == 0)
          d[p] = 0;
        else
          d[p] /= A;
      }
  }

};

#endif