source: Sophya/trunk/Cosmo/RadioBeam/specpk.cc@ 3787

//  Classes to compute 3D power spectrum 
// R. Ansari - Nov 2008, May 2010 

#include "specpk.h"
#include "randr48.h"      

//------------------------------------
// Class SpectralShape 
// -----------------------------------

double Pnu1(double nu) 
{
  return ( sqrt(sqrt(nu)) / ((nu+1.0)/0.2) * 
           (1+0.2*cos(2*M_PI*(nu-2.)*0.15)*exp(-nu/50.)) );
}

double Pnu2(double nu) 
{
  if (nu < 1.e-9) return 0.;
  return ((1.-exp(-nu/0.5))/nu*(1+0.25*cos(2*M_PI*nu*0.1)*exp(-nu/20.)) );
}


double Pnu3(double nu) 
{
  return ( log(nu/100.+1)*(1+sin(2*M_PI*nu/300))*exp(-nu/4000) );
}


double Pnu4(double nu) 
{
  double x = (nu-0.5)/0.05;
  double rc = 2*exp(-x*x);
  x = (nu-3.1)/0.27;
  rc += exp(-x*x);
  x = (nu-7.6)/1.4;
  rc += 0.5*exp(-x*x);
  return ( rc+2.*exp(-x*x) );
}

//--------------------------------------------------
// -- SpectralShape class : test P(k) class
//--------------------------------------------------
// Constructor
SpectralShape::SpectralShape(int typ)
{
  typ_=typ;
}

// Return the spectral power for a given wave number wk 
double SpectralShape::operator() (double wk)
{
  wk/=DeuxPI;
  switch (typ_) 
    {
    case 1:
      return Pnu1(wk);
      break;
    case 2:
      return Pnu2(wk);
      break;
    case 3:
      return Pnu3(wk);
      break;
    case 4:
      return Pnu4(wk);
      break;
    default :
      {
  // global shape
      double csp = pow( (2*sin(sqrt(sqrt(wk/7.)))),2.);
      if (csp < 0.) return 0.;
      
      // Adding some pics
      double picpos[5] = {75.,150.,225.,300.,375.,};
      
      for(int k=0; k<5; k++) {
        double x0 = picpos[k];
        if ( (wk > x0-25.) && (wk < x0+25.) ) {
          double x = (wk-x0);
          csp *= (1.+0.5*exp(-(x*x)/(2.*5*5)));
          break;
        }
      }
      return csp;
      }
      break;
    }
}
// Return a vector representing the power spectrum (for checking) 
Histo SpectralShape::GetPk(int n)
{
  if (n<16) n = 256;
  Histo h(0.,1024.*DeuxPI,n);
  for(int k=0; k<h.NBins(); k++)   h(k) = Value((k+0.5)*h.BinWidth());
  return h;     
}

//--------------------------------------------------
// -- Four2DResponse class : test P(k) class

//---------------------------------------------------------------
// -- Four3DPk class :  3D fourier amplitudes and power spectrum 
//---------------------------------------------------------------
// Constructeur avec Tableau des coeff. de Fourier en argument
Four3DPk::Four3DPk(TArray< complex<TF> > & fourcoedd, RandomGeneratorInterface& rg)
  : rg_(rg), fourAmp(fourcoedd)
{
  SetPrtLevel();
  SetCellSize();
}
// Constructor
Four3DPk::Four3DPk(RandomGeneratorInterface& rg, sa_size_t szx, sa_size_t szy, sa_size_t szz)
  : rg_(rg), fourAmp(szx, szy, szz) 
{
  SetPrtLevel();
  SetCellSize();
}


// Generate mass field Fourier Coefficient
void Four3DPk::ComputeFourierAmp(SpectralShape& pk)
{
  // We generate a random gaussian real field  
  // fourAmp represent 3-D fourier transform of a real input array. 
  // The second half of the array along Y and Z contain negative frequencies
  //  double fnorm = 1./sqrt(2.*fourAmp.Size()); 
  double fnorm = 1.; 
  double kxx, kyy, kzz;
  // sa_size_t is large integer type  
  for(sa_size_t kz=0; kz<fourAmp.SizeZ(); kz++) {
    kzz =  (kz>fourAmp.SizeZ()/2) ? (double)(fourAmp.SizeZ()-kz)*dkz_ : (double)kz*dkz_; 
    for(sa_size_t ky=0; ky<fourAmp.SizeY(); ky++) {
      kyy =  (ky>fourAmp.SizeY()/2) ? (double)(fourAmp.SizeY()-ky)*dky_ : (double)ky*dky_; 
      for(sa_size_t kx=0; kx<fourAmp.SizeX(); kx++) {
        double kxx=(double)kx*dkx_;
        double wk = sqrt(kxx*kxx+kyy*kyy+kzz*kzz);
        double amp = sqrt(pk(wk)*fnorm/2.);      
        fourAmp(kx, ky, kz) = complex<TF>(rg_.Gaussian(amp), rg_.Gaussian(amp));   // renormalize fourier coeff usin 
      }
    }
  }
  if (prtlev_>0)
    cout << " Four3DPk::ComputeFourierAmp() done ..." << endl;
}

// Generate mass field Fourier Coefficient
void Four3DPk::ComputeNoiseFourierAmp(Four2DResponse& resp, bool crmask)
{
  TMatrix<r_4> mask(fourAmp.SizeY(), fourAmp.SizeX());
  // fourAmp represent 3-D fourier transform of a real input array. 
  // The second half of the array along Y and Z contain negative frequencies
  double kxx, kyy, kzz, rep, amp;
  // sa_size_t is large integer type  
  for(sa_size_t kz=0; kz<fourAmp.SizeZ(); kz++) {
    kzz =  (kz>fourAmp.SizeZ()/2) ? -(double)(fourAmp.SizeZ()-kz)*dkz_ : (double)kz*dkz_; 
    for(sa_size_t ky=0; ky<fourAmp.SizeY(); ky++) {
      kyy =  (ky>fourAmp.SizeY()/2) ? -(double)(fourAmp.SizeY()-ky)*dky_ : (double)ky*dky_; 
      for(sa_size_t kx=0; kx<fourAmp.SizeX(); kx++) {
        kxx=(double)kx*dkx_;
        rep = resp(kxx, kyy);
        if (crmask&&(kz==0))  mask(ky,kx)=((rep<1.e-8)?9.e9:(1./rep));
        if (rep<1.e-8)  fourAmp(kx, ky, kz) = complex<TF>(9.e9,0.);
        else {
          amp = 1./sqrt(rep)/sqrt(2.);
          fourAmp(kx, ky, kz) = complex<TF>(rg_.Gaussian(amp), rg_.Gaussian(amp));   
        }
      }
    }
  }
  if (prtlev_>1)  fourAmp.Show();
  if (crmask) {
    POutPersist po("mask.ppf");
    po << mask;
  }
  if (prtlev_>0)
    cout << " Four3DPk::ComputeNoiseFourierAmp() done ..." << endl;
}

// Compute mass field from its Fourier Coefficient
TArray<TF>  Four3DPk::ComputeMassDens()
{
  TArray<TF> massdens;
// Backward fourier transform of the fourierAmp array   
  FFTWServer ffts(true);                     
  ffts.setNormalize(true); 
  ffts.FFTBackward(fourAmp, massdens, true);
  //  cout << " Four3DPk::ComputeMassDens() done NbNeg=" << npbz << " / NPix=" <<  massDens.Size() << endl;
  cout << " Four3DPk::ComputeMassDens() done NPix=" <<  massdens.Size() << endl;
  return massdens;
}

// Compute power spectrum as a function of wave number k 
// cells with amp^2=re^2+im^2>s2cut are ignored
// Output : power spectrum (profile histogram)
HProf Four3DPk::ComputePk(double s2cut, int nbin, double kmin, double kmax)
{
  // The second half of the array along Y (matrix rows) contain
  // negative frequencies
  //  int nbh = sqrt(fourAmp.SizeX()*fourAmp.SizeX()+fourAmp.SizeY()*fourAmp.SizeY()/4.+fourAmp.SizeZ()*fourAmp.SizeY()/4.);
  // The profile histogram will contain the mean value of FFT amplitude
  // as a function of wave-number k = sqrt((double)(kx*kx+ky*ky))
  //  if (nbin < 1) nbin = nbh/2;
  if ((kmax<0.)||(kmax<kmin)) {
    kmin=0.;
    double maxx=fourAmp.SizeX()*dkx_;
    double maxy=fourAmp.SizeY()*dky_/2;
    double maxz=fourAmp.SizeZ()*dkz_/2;
    kmax=sqrt(maxx*maxx+maxy*maxy+maxz*maxz);
  }
  if (nbin<2) nbin=128; 
  HProf hp(kmin, kmax, nbin);
  hp.SetErrOpt(false);
  ComputePkCumul(hp, s2cut);
  return hp;
}

// Compute power spectrum as a function of wave number k 
// Cumul dans hp - cells with amp^2=re^2+im^2>s2cut are ignored
void Four3DPk::ComputePkCumul(HProf& hp, double s2cut)
{

  // fourAmp represent 3-D fourier transform of a real input array. 
  // The second half of the array along Y and Z contain negative frequencies
  double kxx, kyy, kzz;
  // sa_size_t is large integer type  
  // We ignore 0th term in all frequency directions ...
  for(sa_size_t kz=1; kz<fourAmp.SizeZ(); kz++) {
    kzz =  (kz > fourAmp.SizeZ()/2) ? (double)(fourAmp.SizeZ()-kz)*dkz_ : (double)kz*dkz_; 
    for(sa_size_t ky=1; ky<fourAmp.SizeY(); ky++) {
      kyy =  (ky > fourAmp.SizeY()/2) ? (double)(fourAmp.SizeY()-ky)*dky_ : (double)ky*dky_; 
      for(sa_size_t kx=1; kx<fourAmp.SizeX(); kx++) {  // ignore the 0th coefficient (constant term)
        double kxx=(double)kx*dkx_;
        complex<TF> za = fourAmp(kx, ky, kz);
        if (za.real()>8.e9) continue;
        double wk = sqrt(kxx*kxx+kyy*kyy+kzz*kzz);
        double amp2 = za.real()*za.real()+za.imag()*za.imag();
        if ((s2cut>1.e-9)&&(amp2>s2cut))  continue;
        hp.Add(wk, amp2);
      }
    }
  }
  return;
}



//-----------------------------------------------------
// -- MassDist2D class :  2D mass distribution 
//-----------------------------------------------------
// Constructor
MassDist2D::MassDist2D(GenericFunc& pk, int size, double meandens) 
: pkSpec(pk) , sizeA((size>16)?size:16) ,  massDens(sizeA, sizeA), 
  meanRho(meandens) , fg_fourAmp(false) , fg_massDens(false)
{
}

// To the computation job
void MassDist2D::Compute()
{
  ComputeFourierAmp();
  ComputeMassDens();    
}

// Generate mass field Fourier Coefficient
void MassDist2D::ComputeFourierAmp()
{
  if (fg_fourAmp) return; // job already done
  // We generate a random gaussian real field  
  double sigma = 1.;
// The following line fills the array by gaussian random numbers  
//--Replaced--  massDens = RandomSequence(RandomSequence::Gaussian, 0., sigma);
// Can be replaced by 
  DR48RandGen rg;
  for(sa_size_t ir=0; ir<massDens.NRows(); ir++) {
        for(sa_size_t jc=0; jc<massDens.NCols(); jc++) {
      massDens(ir, jc) = rg.Gaussian(sigma); 
        }
  }
// --- End of random filling

  // Compute fourier transform of the random gaussian field -> white noise 
  FFTWServer ffts(true);                     
  ffts.setNormalize(true); 
  ffts.FFTForward(massDens, fourAmp);
    
  // fourAmp represent 2-D fourier transform of a real input array. 
  // The second half of the array along Y (matrix rows) contain
  // negative frequencies
//  double fnorm = 1./sqrt(2.*fourAmp.Size()); 
// PUT smaller value for fnorm and check number of zeros
  double fnorm = 1.; 
  // sa_size_t is large integer type  
  for(sa_size_t ky=0; ky<fourAmp.NRows(); ky++) {
    double kyy = ky;
    if (ky > fourAmp.NRows()/2) kyy = fourAmp.NRows()-ky;  // negative frequencies 
    for(sa_size_t kx=0; kx<fourAmp.NCols(); kx++) {
      double wk = sqrt((double)(kx*kx+kyy*kyy));
      double amp = pkSpec(wk)*fnorm;      
      fourAmp(ky, kx) *= amp;   // renormalize fourier coeff using 
    }
  }
  fg_fourAmp = true;
  cout << " MassDist2D::ComputeFourierAmp() done ..." << endl;
}

// Compute mass field from its Fourier Coefficient
void MassDist2D::ComputeMassDens()
{
  if (fg_massDens) return; // job already done
  if (!fg_fourAmp) ComputeFourierAmp();   // Check fourier amp generation

// Backward fourier transform of the fourierAmp array   
  FFTWServer ffts(true);                     
  ffts.setNormalize(true); 
  ffts.FFTBackward(fourAmp, massDens, true);
// We consider that massDens represents delta rho/rho 
// rho = (delta rho/rho + 1) * MeanDensity 
  massDens += 1.;
// We remove negative values 
  sa_size_t npbz = 0;
  for (sa_size_t i=0; i<massDens.NRows(); i++) 
    for (sa_size_t j=0; j<massDens.NCols(); j++) 
      if (massDens(i,j) < 0.) { npbz++; massDens(i,j) = 0.; }
  massDens *= meanRho;
  cout << " MassDist2D::ComputeMassDens() done NbNeg=" << npbz << " / NPix=" <<  massDens.Size() << endl;
}

// Compute power spectrum as a function of wave number k 
// Output : power spectrum (profile histogram)
HProf MassDist2D::ReconstructPk(int nbin)
{
  // The second half of the array along Y (matrix rows) contain
  // negative frequencies
  int nbh = sqrt(2.0)*fourAmp.NCols();
  // The profile histogram will contain the mean value of FFT amplitude
  // as a function of wave-number k = sqrt((double)(kx*kx+ky*ky))
  if (nbin < 1) nbin = nbh/2;
  HProf hp(-0.5, nbh-0.5, nbin);
  hp.SetErrOpt(false);

  for(int ky=0; ky<fourAmp.NRows(); ky++) {
    double kyy = ky;
    if (ky > fourAmp.NRows()/2)  kyy = fourAmp.NRows()-ky;  // negative frequencies
    for(int kx=0; kx<fourAmp.NCols(); kx++) {
      double wk = sqrt((double)(kx*kx+kyy*kyy));
      complex<r_8> za = fourAmp(ky, kx);
      double amp = sqrt(za.real()*za.real()+za.imag()*za.imag());
      hp.Add(wk, amp);
    }
  }
  return hp;
}