field.hpp

Go to the documentation of this file.

// Triumvirate: Three-Point Clustering Measurements in LSS
//
// Copyright (C) 2023 Mike S Wang & Naonori S Sugiyama [GPL-3.0-or-later]
//
// This file is part of the Triumvirate program. See the COPYRIGHT
// and LICENCE files at the top-level directory of this distribution
// for details of copyright and licensing.
//
// This program is free software: you can redistribute it and/or modify it
// under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
// See the GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

 
#ifndef TRIUMVIRATE_INCLUDE_FIELD_HPP_INCLUDED_
#define TRIUMVIRATE_INCLUDE_FIELD_HPP_INCLUDED_
 
#if defined(TRV_USE_HIP)
#include <hipfft/hipfftXt.h>
#elif defined(TRV_USE_CUDA)  // !TRV_USE_HIP && TRV_USE_CUDA
#include <cufftXt.h>
#endif                       // TRV_USE_HIP
#include <fftw3.h>
 
#include <chrono>
#include <cmath>
#include <complex>
#include <functional>
#include <vector>
 
#include "arrayops.hpp"
#include "monitor.hpp"
#include "parameters.hpp"
#include "maths.hpp"
#include "dataobjs.hpp"
#include "io.hpp"
#include "particles.hpp"

#ifdef __GNUC__
#define PURE __attribute__((pure))
#else
#define PURE
#endif
 
namespace trvm = trv::maths;
 
namespace trv {
 
// ***********************************************************************
// Mesh field
// ***********************************************************************

class MeshField {
 public:
  trv::ParameterSet params;  
  std::string name;          
  fftw_complex* field;       
  double dr[3];              
  double dk[3];              
  double vol;                
  double vol_cell;           
 
  // ---------------------------------------------------------------------
  // Life cycle
  // ---------------------------------------------------------------------

  explicit MeshField(
    trv::ParameterSet& params,
    bool plan_ini = true,
    const std::string& name = "mesh-field"
  );

  explicit MeshField(
    trv::ParameterSet& params,
    fftw_plan& transform, fftw_plan& inv_transform,
    const std::string& name = "mesh-field"
  );

  ~MeshField();

  void reset_density_field();
 
  // ---------------------------------------------------------------------
  // Operators & reserved methods
  // ---------------------------------------------------------------------

  PURE const fftw_complex& operator[](long long gid);
 
  // ---------------------------------------------------------------------
  // Mesh assignment
  // ---------------------------------------------------------------------

  void assign_weighted_field_to_mesh(
    ParticleCatalogue& particles, fftw_complex* weights
  );
 
  // ---------------------------------------------------------------------
  // Field computations
  // ---------------------------------------------------------------------

  void compute_unweighted_field(ParticleCatalogue& particles);

  void compute_unweighted_field_fluctuations_insitu(
    ParticleCatalogue& particles
  );

  void compute_ylm_wgtd_field(
    ParticleCatalogue& particles_data, ParticleCatalogue& particles_rand,
    LineOfSight* los_data, LineOfSight* los_rand,
    double alpha, int ell, int m
  );

  void compute_ylm_wgtd_field(
    ParticleCatalogue& particles, LineOfSight* los,
    double alpha, int ell, int m
  );

  void compute_ylm_wgtd_quad_field(
    ParticleCatalogue& particles_data, ParticleCatalogue& particles_rand,
    LineOfSight* los_data, LineOfSight* los_rand,
    double alpha, int ell, int m
  );

  void compute_ylm_wgtd_quad_field(
    ParticleCatalogue& particles, LineOfSight* los,
    double alpha, int ell, int m
  );
 
  // ---------------------------------------------------------------------
  // Field transforms
  // ---------------------------------------------------------------------

  void fourier_transform();

  void inv_fourier_transform();
 
  // ---------------------------------------------------------------------
  // Field operations
  // ---------------------------------------------------------------------

  void apply_wide_angle_pow_law_kernel();

  void apply_assignment_compensation();
 
  // ---------------------------------------------------------------------
  // One-point statistics
  // ---------------------------------------------------------------------

  void inv_fourier_transform_ylm_wgtd_field_band_limited(
    MeshField& field_fourier, std::vector< std::complex<double> >& ylm,
    double k_band, double dk_band,
    double& k_eff, int& nmodes
  );

  void inv_fourier_transform_sjl_ylm_wgtd_field(
    MeshField& field_fourier,
    std::vector< std::complex<double> >& ylm,
    trvm::SphericalBesselCalculator& sjl,
    double r
  );
 
  // ---------------------------------------------------------------------
  // Misc
  // ---------------------------------------------------------------------

  double calc_grid_based_powlaw_norm(ParticleCatalogue& particles, int order);
 
 private:
  std::vector<double> window;
  int window_assign_order = -1;

  fftw_complex* field_s = nullptr;

  fftw_plan transform{};
  fftw_plan transform_s{};
  fftw_plan inv_transform{};
#if defined(TRV_USE_HIP) || defined(TRV_USE_CUDA)
  fftHandle transform_gpu{};
#endif  // TRV_USE_HIP || TRV_USE_CUDA

#if defined(TRV_USE_HIP) || defined(TRV_USE_CUDA)
  fft_double_complex* d_field = nullptr;
#endif  // TRV_USE_HIP || TRV_USE_CUDA
#ifdef TRV_USE_CUDA
  cudaLibXtDesc* field_desc = nullptr;
#endif  // TRV_USE_CUDA
 
#ifdef _CUDA_STREAM
  cudaStream_t custream{};
#endif  // _CUDA_STREAM
 
  bool plan_ini = false;  
  bool plan_ext = false;  
 
  friend class FieldStats;
 
  // ---------------------------------------------------------------------
  // Mesh grid properties
  // ---------------------------------------------------------------------

  PURE long long ret_grid_index(int i, int j, int k);

  void shift_grid_indices_fourier(int& i, int& j, int& k);

  void get_grid_pos_vector(int i, int j, int k, double rvec[3]);

  void get_grid_wavevector(int i, int j, int k, double kvec[3]);
 
  // ---------------------------------------------------------------------
  // Mesh assignment
  // ---------------------------------------------------------------------

  void assign_weighted_field_to_mesh_ngp(
    ParticleCatalogue& particles, fftw_complex* weight
  );

  void assign_weighted_field_to_mesh_cic(
    ParticleCatalogue& particles, fftw_complex* weight
  );

  void assign_weighted_field_to_mesh_tsc(
    ParticleCatalogue& particles, fftw_complex* weight
  );

  void assign_weighted_field_to_mesh_pcs(
    ParticleCatalogue& particles, fftw_complex* weight
  );

  double calc_assignment_window_in_fourier(int i, int j, int k, int order = 0);

  void compute_assignment_window_in_fourier(int order = 0);
};
 
 
// ***********************************************************************
// Field statistics
// ***********************************************************************

class FieldStats {
 public:
  std::vector<int> nmodes;  
  std::vector<int> npairs;  
  std::vector<double> k;    
  std::vector<double> r;    
  std::vector< std::complex<double> > sn;
  std::vector< std::complex<double> > pk;
  std::vector< std::complex<double> > xi;
 
  // ---------------------------------------------------------------------
  // Life cycle
  // ---------------------------------------------------------------------

  explicit FieldStats(trv::ParameterSet& params, bool plan_ini = true);

  ~FieldStats();

  void reset_stats();
 
  // ---------------------------------------------------------------------
  // Binned statistics
  // ---------------------------------------------------------------------

  void compute_ylm_wgtd_2pt_stats_in_fourier(
    MeshField& field_a, MeshField& field_b, std::complex<double> shotnoise_amp,
    int ell, int m, trv::Binning& kbinning
  );

  void compute_ylm_wgtd_2pt_stats_in_config(
    MeshField& field_a, MeshField& field_b, std::complex<double> shotnoise_amp,
    int ell, int m, trv::Binning& rbinning
  );

  void compute_uncoupled_shotnoise_for_3pcf(
    MeshField& field_a, MeshField& field_b,
    std::vector< std::complex<double> >& ylm_a,
    std::vector< std::complex<double> >& ylm_b,
    std::complex<double> shotnoise_amp,
    trv::Binning& rbinning
  );

  std::complex<double> compute_uncoupled_shotnoise_for_bispec_per_bin(
    MeshField& field_a, MeshField& field_b,
    std::vector< std::complex<double> >& ylm_a,
    std::vector< std::complex<double> >& ylm_b,
    trvm::SphericalBesselCalculator& sj_a,
    trvm::SphericalBesselCalculator& sj_b,
    std::complex<double> shotnoise_amp,
    double k_a, double k_b
  );

  trv::BinnedVectors record_binned_vectors(
    trv::Binning& binning, const std::string& save_file
  );
 
 private:
  trv::ParameterSet params;  
  double dr[3];              
  double dk[3];              
  double vol;                
  double vol_cell;           

  fftw_plan inv_transform{};
#if defined(TRV_USE_HIP) || defined(TRV_USE_CUDA)
  fftHandle inv_transform_gpu{};
#endif  // TRV_USE_HIP || TRV_USE_CUDA

  fftw_complex* twopt_3d = nullptr;
#if defined(TRV_USE_HIP) || defined(TRV_USE_CUDA)
  fft_double_complex* d_twopt_3d = nullptr;
#endif  // TRV_USE_HIP || TRV_USE_CUDA
#ifdef TRV_USE_CUDA
  cudaLibXtDesc* twopt_3d_desc = nullptr;
#endif  // TRV_USE_CUDA
 
#ifdef _CUDA_STREAM
  cudaStream_t custream{};
#endif  // _CUDA_STREAM

  bool plan_ini = false;

  std::vector<double> alias_sn;

  bool alias_ini = false;
 
  // ---------------------------------------------------------------------
  // Utilities
  // ---------------------------------------------------------------------

  PURE bool if_fields_compatible(MeshField& field_a, MeshField& field_b);

  void resize_stats(int num_bins);

  PURE long long ret_grid_index(int i, int j, int k);

  void shift_grid_indices_fourier(int& i, int& j, int& k);
 
  // ---------------------------------------------------------------------
  // Sampling corrections
  // ---------------------------------------------------------------------

  std::function<double(int, int, int)> ret_calc_shotnoise_aliasing();

  void get_shotnoise_aliasing_sin2(
    int i, int j, int k, double& sx2, double& sy2, double& sz2
  );

  void get_shotnoise_aliasing_sin2_cos2_isoapprox(
    int i, int j, int k, double& s2h, double& c2h
  );

  double calc_shotnoise_aliasing_ngp(int i, int j, int k);

  double calc_shotnoise_aliasing_cic(int i, int j, int k);

  double calc_shotnoise_aliasing_tsc(int i, int j, int k);

  double calc_shotnoise_aliasing_pcs(int i, int j, int k);

  double calc_shotnoise_aliasing_interlaced_isoapprox(
    int i, int j, int k, int order
  );

  void compute_shotnoise_aliasing();
};
 
}  // namespace trv
 
#endif  // !TRIUMVIRATE_INCLUDE_FIELD_HPP_INCLUDED_