Compute spectra of bias operators advected from Lagrangian space

FML/COLASolver/src/AnalyzeOutput.h

@@ -476,6 +476,118 @@ void compute_power_spectrum(NBodySimulation<NDIM, T> & sim, double redshift, std
     sim.pofk_every_output.push_back({redshift, pofk_cb_binning});
 }
 
+template <int NDIM, class T>
+void compute_power_spectrum_bias(NBodySimulation<NDIM, T> & sim, double redshift, std::string snapshot_folder) {
+
+    std::stringstream stream;
+    stream << std::fixed << std::setprecision(3) << redshift;
+    std::string redshiftstring = stream.str();
+
+    //=============================================================
+    // Fetch parameters
+    //=============================================================
+    const double simulation_boxsize = sim.simulation_boxsize;
+    const int pofk_nmesh = sim.pofk_nmesh;
+    const std::string pofk_density_assignment_method = sim.pofk_density_assignment_method;
+    const bool pofk_interlacing = sim.pofk_interlacing;
+    const bool pofk_subtract_shotnoise = sim.pofk_subtract_shotnoise;
+    auto & part = sim.part;
+
+    // TODO: New pofk_bias_... input file parameters
+    if (FML::ThisTask == 0) {
+        std::cout << "\n";
+        std::cout << "#=====================================================\n";
+        std::cout << "# Computing power-spectrum of bias operators\n";
+        std::cout << "# pofk_nmesh                     : " << pofk_nmesh << "\n";
+        std::cout << "# pofk_density_assignment_method : " << pofk_density_assignment_method << "\n";
+        std::cout << "# pofk_interlacing               : " << pofk_interlacing << "\n";
+        std::cout << "# pofk_subtract_shotnoise        : " << pofk_subtract_shotnoise << "\n";
+        std::cout << "#=====================================================\n";
+    }
+
+    // TODO: neutrinos accounted for?
+    const auto nleftright = FML::INTERPOLATION::get_extra_slices_needed_for_density_assignment(pofk_density_assignment_method);
+    const int nleft = nleftright.first;
+    const int nright = nleftright.second + (pofk_interlacing ? 1 : 0);
+
+    // grids for 1(x), delta(x), delta^2(x)
+    std::vector<FFTWGrid<NDIM>> grids;
+    grids.reserve(3);
+    const std::vector<std::string> grid_labels = {"1", "d", "d2"}; // 1(x), delta(x), delta^2(x)
+    for (int i = 0; i < 3; i++) {
+        auto & grid = grids.emplace_back(pofk_nmesh, nleft, nright);
+        grid.set_grid_status_real(true);
+    }
+
+    // Save original masses, then restore them below
+    std::vector<double> original_masses(part.get_npart());
+    for (size_t i = 0; i < part.get_npart(); i++) {
+        original_masses[i] = part[i].mass;
+    }
+
+    // Grid 0: 1(q) -> 1(x)
+    for (auto & p : part) {
+        p.mass = 1.0;
+    }
+    FML::INTERPOLATION::particles_to_grid(part.get_particles_ptr(), part.get_npart(), part.get_npart_total(), grids[0], pofk_density_assignment_method);
+
+    // Grid 1: delta(q) -> delta(x)
+    for (auto & p : part) {
+        p.mass = p.delta_q;
+    }
+    FML::INTERPOLATION::particles_to_grid(part.get_particles_ptr(), part.get_npart(), part.get_npart_total(), grids[1], pofk_density_assignment_method);
+
+    // Grid 2: delta^2(q) -> delta^2(x)
+    for (auto & p : part) {
+        p.mass = p.delta2_q;
+    }
+    FML::INTERPOLATION::particles_to_grid(part.get_particles_ptr(), part.get_npart(), part.get_npart_total(), grids[2], pofk_density_assignment_method);
+
+    // Restore original particle masses
+    for (size_t i = 0; i < part.get_npart(); i++) {
+        part[i].mass = original_masses[i];
+    }
+
+    // Compute spectra
+    std::vector<FML::CORRELATIONFUNCTIONS::PowerSpectrumBinning<NDIM>> Ps;
+    std::vector<std::string> P_labels;
+    for (size_t i = 0; i < grids.size(); i++) {
+        double mean = grids[i].mean();
+        grids[i].fill_real_grid([mean](std::array<double, 3> &, FML::GRID::FloatType val) { return val - mean; }); // subtract mean, so mean is zero
+        grids[i].fftw_r2c();
+        FML::INTERPOLATION::deconvolve_window_function_fourier<NDIM>(grids[i], pofk_density_assignment_method);
+        for (size_t j = 0; j <= i; j++) {
+            // grids j <= i have already been transformed and deconvolved
+            auto & P = Ps.emplace_back(pofk_nmesh / 2);
+            if (i == j) {
+                FML::CORRELATIONFUNCTIONS::bin_up_power_spectrum(grids[i], P); // auto-spectra
+            } else {
+                FML::CORRELATIONFUNCTIONS::bin_up_cross_power_spectrum(grids[i], grids[j], P); // cross-spectra
+            }
+            P.scale(simulation_boxsize);
+            P_labels.emplace_back(grid_labels[i] + grid_labels[j]); // e.g. "d1"
+        }
+    }
+
+    // Output to file
+    if (FML::ThisTask == 0) {
+        std::string filename = snapshot_folder + "/pofk_bias_z" + redshiftstring + ".txt";
+        std::ofstream fp(filename.c_str());
+        if (not fp.is_open()) {
+            std::cout << "Warning: Cannot write power-spectrum of bias operators to file, failed to open [" << filename << "]\n";
+        } else {
+            fp << "#" << std::setw(16-1) << "k/(h/Mpc)";
+            for (auto & P_label : P_labels) fp << " " << std::setw(16) << ("P_" + P_label + "/(Mpc/h)^3");
+            fp << std::endl;
+            for (int i = 0; i < Ps[0].n; i++) {
+                fp << std::setw(16) << Ps[0].kbin[i]; // k-values are common
+                for (auto & P : Ps) fp << " "  << std::setw(16) << P.pofk[i];
+                fp << std::endl;
+            }
+        }
+    }
+}
+
 template <int NDIM, class T>
 void compute_fof_halos(NBodySimulation<NDIM, T> & sim, double redshift, std::string snapshot_folder) {
 

FML/COLASolver/src/Main.cpp

@@ -70,13 +70,23 @@ class Particle {
     void set_id(long long int _id) { id = _id; }
 
     //=============================================================
-    // Initial Lagrangian position
+    // Mass (for weighting in when interpolating particles to grid)
+    //=============================================================
+    double mass;
+    double get_mass() const { return mass; }
+
+    //=============================================================
+    // Initial Lagrangian position and functions thereof (delta, ...)
     // Only needed for COLA with scaledependent growth
     // NB: should ideally have same type as [pos] to avoid truncating the
     // precision of pos (these are temporarily swapped by some algorithms)
     //=============================================================
-    double q[NDIM];
+    double q[NDIM]; // position
+    double delta_q; // overdensity
+    double delta2_q; // overdensity squared
     double * get_q() { return q; }
+    double get_delta_q() { return delta_q; }
+    double get_delta2_q() { return delta2_q; }
 
     //=============================================================
     // 1LPT displacement field Psi (needed if you want >= 1LPT COLA)

FML/COLASolver/src/Simulation.h

@@ -273,6 +273,8 @@ class NBodySimulation {
     template <int _NDIM, class _T>
     friend void compute_power_spectrum(NBodySimulation<_NDIM, _T> & sim, double redshift, std::string snapshot_folder);
     template <int _NDIM, class _T>
+    friend void compute_power_spectrum_bias(NBodySimulation<_NDIM, _T> & sim, double redshift, std::string snapshot_folder);
+    template <int _NDIM, class _T>
     friend void
     compute_power_spectrum_multipoles(NBodySimulation<_NDIM, _T> & sim, double redshift, std::string snapshot_folder);
     template <int _NDIM, class _T>
@@ -1677,6 +1679,7 @@ void NBodySimulation<NDIM, T>::analyze_and_output(int ioutput, double redshift)
     if (pofk) {
         timer.StartTiming("Power-spectrum");
         compute_power_spectrum(*this, redshift, snapshot_folder);
+        compute_power_spectrum_bias(*this, redshift, snapshot_folder);
         timer.EndTiming("Power-spectrum");
     }
 

FML/FFTWGrid/FFTWGrid.h

@@ -170,7 +170,7 @@ namespace FML {
             void fill_fourier_grid(const ComplexType val);
 
             /// Fill the main grid from a function specifying the value at a given position
-            void fill_real_grid(std::function<FloatType(std::array<double, N> &)> & func);
+            void fill_real_grid(const std::function<FloatType(std::array<double, N> &, FloatType)> & func);
             /// Fill the main grid from a function specifying the value at a given fourier wave-vector
             void fill_fourier_grid(std::function<ComplexType(std::array<double, N> &)> & func);
 
@@ -213,6 +213,10 @@ namespace FML {
             /// Add to value in grid
             void add_real(const std::array<int, N> & coord, const FloatType value);
 
+            /// Get sum and mean of all (global) real cell values
+            FloatType sum() const;
+            FloatType mean() const;
+
             /// Set value of cell in fourier grid using (local) coordinate in Local_nx x [0,Nmesh)^(Ndim-2) x
             /// [0,Nmesh/2+1)
             void set_fourier(const std::array<int, N> & coord, const ComplexType value);
@@ -600,7 +604,7 @@ namespace FML {
         }
 
         template <int N>
-        void FFTWGrid<N>::fill_real_grid(std::function<FloatType(std::array<double, N> &)> & func) {
+        void FFTWGrid<N>::fill_real_grid(const std::function<FloatType(std::array<double, N> &, FloatType)> & func) {
 #ifdef DEBUG_FFTWGRID
             if (not grid_is_in_real_space) {
                 if (FML::ThisTask == 0)
@@ -615,8 +619,9 @@ namespace FML {
             for (int islice = 0; islice < Local_nx; islice++) {
                 for (auto && real_index : get_real_range(islice, islice + 1)) {
                     auto coord = get_coord_from_index(real_index);
+                    auto value = get_real_from_index(real_index);
                     auto pos = get_real_position(coord);
-                    auto value = func(pos);
+                    value = func(pos, value);
                     set_real_from_index(real_index, value);
                 }
             }
@@ -1017,6 +1022,34 @@ namespace FML {
             get_real_grid()[index] += value;
         }
 
+        template <int N>
+        FloatType FFTWGrid<N>::sum() const {
+#ifdef DEBUG_FFTWGRID
+            if (not grid_is_in_real_space) {
+                if (FML::ThisTask == 0)
+                    std::cout << "Warning: [FFTWGrid::fill_real_grid] The grid status is [Fourierspace]. Label: " +
+                                     name + "\n";
+            }
+#endif
+            FloatType tot = 0.0;
+            #ifdef USE_OMP
+            #pragma omp parallel for reduction(+ : tot)
+            #endif
+            for (int islice = 0; islice < Local_nx; islice++) {
+                for (auto && real_index : get_real_range(islice, islice + 1)) {
+                    tot += get_real_from_index(real_index);
+                }
+            }
+            FML::SumOverTasks(&tot);
+            return tot;
+        }
+
+        template <int N>
+        FloatType FFTWGrid<N>::mean() const {
+            return sum() / std::pow(get_nmesh(), N);
+        }
+
+
         template <int N>
         void FFTWGrid<N>::set_real_from_index(const IndexIntType index, const FloatType value) {
             get_real_grid()[index] = value;

FML/Interpolation/ParticleGridInterpolation.h

@@ -512,7 +512,7 @@ namespace FML {
                 }
                 SumOverTasks(&mean_mass);
                 mean_mass /= double(NumPartTot);
-                norm_fac /= mean_mass;
+                //norm_fac /= mean_mass; // normalization in presence of multiple particle species?
             }
 
             // Loop over all particles and add them to the grid

FML/NBody/NBody.h

@@ -737,7 +737,7 @@ namespace FML {
                 if (FML::PARTICLE::has_get_vel<T>())
                     std::cout << "# Particle has [Velocity] v_code = a^2 dxdt / (H0 Box) ("
                               << sizeof(FML::PARTICLE::GetPos(tmp)[0]) * N << " bytes)\n";
-                if (FML::PARTICLE::has_set_mass<T>())
+                if (FML::PARTICLE::has_get_mass<T>())
                     std::cout << "# Particle has [Mass] (" << sizeof(FML::PARTICLE::GetMass(tmp)) << " bytes)\n";
                 if (FML::PARTICLE::has_set_id<T>())
                     std::cout << "# Particle has [ID] (" << sizeof(FML::PARTICLE::GetID(tmp)) << " bytes)\n";
@@ -756,6 +756,12 @@ namespace FML {
                 if (FML::PARTICLE::has_get_q<T>())
                     std::cout << "# Particle has [Lagrangian position] ("
                               << sizeof(FML::PARTICLE::GetLagrangianPos(tmp)[0]) * N << " bytes)\n";
+                if (FML::PARTICLE::has_get_delta_q<T>())
+                    std::cout << "# Particle has [Lagrangian overdensity] ("
+                              << sizeof(FML::PARTICLE::GetLagrangianDelta(tmp)) << " bytes)\n";
+                if (FML::PARTICLE::has_get_delta2_q<T>())
+                    std::cout << "# Particle has [Lagrangian overdensity squared] ("
+                              << sizeof(FML::PARTICLE::GetLagrangianDelta2(tmp)) << " bytes)\n";
                 std::cout << "# Total size of particle is " << FML::PARTICLE::GetSize(tmp) << " bytes\n";
                 std::cout << "# We will make " << Npart_1D << "^" << N << " particles\n";
                 std::cout << "# Plus a buffer with room for " << (buffer_factor - 1.0) * 100.0 << "%% more particles\n";
@@ -990,6 +996,38 @@ namespace FML {
                 }
             }
 
+            // Set Lagrangian overdensity of the particle if we have that available
+            if constexpr (FML::PARTICLE::has_get_delta_q<T>()) {
+                FFTWGrid<N> delta_real = delta_fourier; // copy
+                delta_real.fftw_c2r(); // to real space
+                std::vector<FML::GRID::FloatType> weights;
+                weights.reserve(part.get_npart());
+                FML::INTERPOLATION::interpolate_grid_to_particle_positions<N, T>(delta_real, part.get_particles_ptr(), part.get_npart(), weights, interpolation_method);
+                for (size_t i = 0; i < part.get_npart(); i++) {
+                    part[i].delta_q = weights[i];
+                }
+                if (FML::ThisTask == 0) {
+                    std::cout << "Storing Lagrangian overdensity delta(q) in particle\n";
+                    std::cout << "Minimum Lagrangian overdensity delta(q): " << *std::min_element(weights.begin(), weights.end()) << std::endl;
+                    std::cout << "Maximum Lagrangian overdensity delta(q): " << *std::max_element(weights.begin(), weights.end()) << std::endl;
+                }
+
+                if constexpr (FML::PARTICLE::has_get_delta2_q<T>()) {
+                    FFTWGrid<N> delta2_real = delta_real;
+                    delta2_real.fill_real_grid([](std::array<FML::GRID::FloatType, 3> &, FML::GRID::FloatType delta) { return delta * delta; }); // delta^2
+                    weights.clear();
+                    FML::INTERPOLATION::interpolate_grid_to_particle_positions<N, T>(delta2_real, part.get_particles_ptr(), part.get_npart(), weights, interpolation_method);
+                    for (size_t i = 0; i < part.get_npart(); i++) {
+                        part[i].delta2_q = weights[i];
+                    }
+                    if (FML::ThisTask == 0) {
+                        std::cout << "Storing Lagrangian overdensity squared delta^2(q) in particle\n";
+                        std::cout << "Minimum Lagrangian overdensity squared delta^2(q): " << *std::min_element(weights.begin(), weights.end()) << std::endl;
+                        std::cout << "Maximum Lagrangian overdensity squared delta^2(q): " << *std::max_element(weights.begin(), weights.end()) << std::endl;
+                    }
+                }
+            }
+
             // Compute and add displacements
             // NB: we must do this in one go as add_displacement changes the position of the particles
             std::array<std::vector<FML::GRID::FloatType>, N> displacements_1LPT;

FML/ParticleTypes/ReflectOnParticleMethods.h

@@ -228,6 +228,8 @@ namespace FML {
         SFINAE_TEST_GET(GetdDdloga_1LPT, get_dDdloga_1LPT)
         SFINAE_TEST_GET(GetdDdloga_2LPT, get_dDdloga_2LPT)
         SFINAE_TEST_GET(GetLagrangianPos, get_q)
+        SFINAE_TEST_GET(GetLagrangianDelta, get_delta_q)
+        SFINAE_TEST_GET(GetLagrangianDelta2, get_delta2_q)
         constexpr double * GetD_1LPT(...) {
             assert_mpi(false, "Trying to get D_1LPT from a particle that has no get_D_1LPT method");
             return nullptr;
@@ -256,6 +258,14 @@ namespace FML {
             assert_mpi(false, "Trying to get the Lagrangian coordinate q from a particle that has no get_q method");
             return nullptr;
         };
+        constexpr double * GetLagrangianDelta(...) {
+            assert_mpi(false, "Trying to get the Lagrangian overdensity from a particle that has no get_delta_q method");
+            return nullptr;
+        };
+        constexpr double * GetLagrangianDelta2(...) {
+            assert_mpi(false, "Trying to get the Lagrangian overdensity squared from a particle that has no get_delta2_q method");
+            return nullptr;
+        };
 
         //=====================================================================
         // Halo finding
@@ -390,7 +400,7 @@ namespace FML {
                     std::cout << "# Particle has [Velocity] (" << sizeof(FML::PARTICLE::GetPos(tmp)[0]) * N
                               << " bytes)\n";
 
-                if constexpr (FML::PARTICLE::has_set_mass<T>())
+                if constexpr (FML::PARTICLE::has_get_mass<T>())
                     std::cout << "# Particle has [Mass] (" << sizeof(FML::PARTICLE::GetMass(tmp)) << " bytes)\n";
 
                 if constexpr (FML::PARTICLE::has_set_id<T>())
@@ -440,6 +450,12 @@ namespace FML {
                 if constexpr (FML::PARTICLE::has_get_q<T>())
                     std::cout << "# Particle has [Lagrangian position] ("
                               << sizeof(FML::PARTICLE::GetLagrangianPos(tmp)[0]) * N << " bytes)\n";
+                if constexpr (FML::PARTICLE::has_get_delta_q<T>())
+                    std::cout << "# Particle has [Lagrangian overdensity] ("
+                              << sizeof(FML::PARTICLE::GetLagrangianDelta(tmp)) << " bytes)\n";
+                if constexpr (FML::PARTICLE::has_get_delta2_q<T>())
+                    std::cout << "# Particle has [Lagrangian overdensity squared] ("
+                              << sizeof(FML::PARTICLE::GetLagrangianDelta2(tmp)) << " bytes)\n";
                 if constexpr (FML::PARTICLE::has_get_D_1LPT<T>() and FML::PARTICLE::has_get_D_2LPT<T>() and
                               FML::PARTICLE::has_get_D_3LPTa<T>() and FML::PARTICLE::has_get_D_3LPTb<T>()) {
                     std::cout << "# Particle compatible with 3LPT COLA\n";