merge in schroedinger

.gitignore

@@ -53,3 +53,4 @@ src/CMakeCache.txt
 src/fastLPT
 src/input_powerspec.txt
 src/Makefile
+.DS_Store

example.conf

@@ -11,7 +11,7 @@ BCClattice   = no
 NumThreads   = 4
 
 [output]
-fname_hdf5   = output.hdf5
+fname_hdf5   = output_sch.hdf5
 fbase_analysis = output
 #format       = gadget2
 #filename     = ics_gadget.dat
@@ -33,3 +33,7 @@ Omega_L      = 0.698
 H0           = 70.3
 sigma_8      = 0.811
 nspec        = 0.961
+
+[sch]
+hbar        = 100.0
+dt          = 5.0

include/grid_fft.hh

@@ -88,6 +88,24 @@ public:
             data_[i] = 0.0;
     }
 
+    void copy_from( const Grid_FFT<data_t>& g ){
+        // make sure the two fields are in the same space
+        if( g.space_ != this->space_ ){
+            if( this->space_ == kspace_id ) this->FourierTransformBackward(false);
+            else this->FourierTransformForward(false);
+        }
+
+        // make sure the two fields have the same dimensions
+        assert( this->n_[0] == g.n_[0] );
+        assert( this->n_[1] == g.n_[1] );
+        assert( this->n_[2] == g.n_[2] );
+
+        // now we can copy all the data over
+        #pragma omp parallel for
+        for (size_t i = 0; i < ntot_; ++i)
+            data_[i] = g.data_[i];
+    }
+
     data_t& operator[]( size_t i ){
         return data_[i];
     }
@@ -356,6 +374,7 @@ public:
         }
     }
 
+
     template <typename functional, typename grid1_t, typename grid2_t>
     void assign_function_of_grids_r(const functional &f, const grid1_t &g1, const grid2_t &g2)
     {
@@ -408,6 +427,27 @@ public:
         }
     }
 
+    template <typename functional, typename grid_t>
+    void assign_function_of_grids_k(const functional &f, const grid_t &g)
+    {
+        assert(g.size(0) == size(0) && g.size(1) == size(1)); // && g.size(2) == size(2) );
+
+    #pragma omp parallel for
+        for (size_t i = 0; i < sizes_[0]; ++i)
+        {
+            for (size_t j = 0; j < sizes_[1]; ++j)
+            {
+                for (size_t k = 0; k < sizes_[2]; ++k)
+                {
+                    auto &elem = this->kelem(i, j, k);
+                    const auto &elemg = g.kelem(i, j, k);
+
+                    elem = f(elemg);
+                }
+            }
+        }
+    }
+
     template <typename functional>
     void apply_function_k_dep(const functional &f)
     {
@@ -508,4 +548,5 @@ public:
             }
         }
     }
+
 };

src/ic_generator.cc

@@ -74,6 +74,8 @@ int Run( ConfigFile& the_config )
     Grid_FFT<real_t> A3x({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
     Grid_FFT<real_t> A3y({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
     Grid_FFT<real_t> A3z({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
+    Grid_FFT<ccomplex_t> psi({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
+    Grid_FFT<real_t> rho({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
     //... array [.] access to components of A3:
     std::array< Grid_FFT<real_t>*,3 > A3({&A3x,&A3y,&A3z});
 
@@ -158,6 +160,92 @@ int Run( ConfigFile& the_config )
         csoca::ilog << std::setw(20) << std::setfill(' ') << std::right << "took " << get_wtime()-wtime << "s" << std::endl;
     }
 
+    //======================================================================
+    // initialise psi = exp(i Phi(1)/hbar)
+    //======================================================================
+    real_t hbar= the_config.GetValueSafe<real_t>("sch", "hbar", 0.1);
+    real_t dt = the_config.GetValueSafe<real_t>("sch", "dt", 0.5);
+    phi.FourierTransformBackward();
+    psi.assign_function_of_grids_r([&]( real_t p ){return std::exp(ccomplex_t(0.0,1.0)/hbar*p);}, phi );
+
+    //======================================================================
+    // evolve wave-function (one drift step) psi = psi *exp(-i hbar *k^2 dt / 2)
+    //======================================================================
+    psi.FourierTransformForward();
+    psi.apply_function_k_dep([hbar,dt]( auto epsi, auto k ){
+        auto k2 = k.norm_squared();
+        return epsi * std::exp( - ccomplex_t(0.0,0.5)*hbar* k2 * dt);
+    });
+    psi.FourierTransformBackward();
+
+    //======================================================================
+    // compute rho
+    //======================================================================
+    // psi.assign_function_of_grids_r([&]( auto p ){return p.norm_squared();}, psi );
+    rho.assign_function_of_grids_r([&]( auto p ){
+        auto pp = std::real(p)*std::real(p) + std::imag(p)*std::imag(p);
+        return pp;
+    }, psi);
+
+
+    //======================================================================
+    // compute  v
+    //======================================================================
+    Grid_FFT<real_t>
+        vx({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen}),
+        vy({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen}),
+        vz({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
+
+    for( int idim=0; idim<3; ++idim )
+    {
+        Grid_FFT<ccomplex_t> grad_psi({ngrid, ngrid, ngrid}, {boxlen, boxlen, boxlen});
+        grad_psi.copy_from(psi);
+        grad_psi.FourierTransformForward();
+        grad_psi.apply_function_k_dep([&](auto x, auto k) {
+            return x * ccomplex_t(0.0,k[idim]);
+        });
+        grad_psi.FourierTransformBackward();
+
+        if( idim==0 )
+        {
+            vx.assign_function_of_grids_r([&](auto ppsi, auto pgrad_psi, auto prho) {
+                return std::real((std::conj(ppsi) * pgrad_psi - ppsi * std::conj(pgrad_psi)) / ccomplex_t(0.0, 2.0 / hbar)/prho);
+            }, psi, grad_psi, rho);
+        }
+        else if( idim==1 )
+        {
+            vy.assign_function_of_grids_r([&](auto ppsi, auto pgrad_psi, auto prho) {
+                return std::real((std::conj(ppsi) * pgrad_psi - ppsi * std::conj(pgrad_psi)) / ccomplex_t(0.0, 2.0 / hbar)/prho);
+            }, psi, grad_psi, rho);
+        }
+        else if (idim == 2)
+        {
+            vz.assign_function_of_grids_r([&](auto ppsi, auto pgrad_psi, auto prho) {
+                return std::real((std::conj(ppsi) * pgrad_psi - ppsi * std::conj(pgrad_psi)) / ccomplex_t(0.0, 2.0 / hbar)/prho);
+            }, psi, grad_psi, rho);
+        }
+    }
+
+    // std::cout << "RHO  " << rho.relem(1,2,3) << std::endl;
+    // std::cout << "VX  " << vx.relem(1,2,3) << std::endl;
+    // std::cout << "VY  " << vy.relem(1,2,3) << std::endl;
+    // std::cout << "VZ  " << vz.relem(1,2,3) << std::endl;
+    const std::string fname_sch_hdf5 = the_config.GetValueSafe<std::string>("output", "fname_hdf5", "output_sch.hdf5");
+
+    #if defined(USE_MPI)
+            if( CONFIG::MPI_task_rank == 0 )
+                unlink(fname_sch_hdf5.c_str());
+            MPI_Barrier( MPI_COMM_WORLD );
+    #else
+            unlink(fname_sch_hdf5.c_str());
+    #endif
+
+    rho.Write_to_HDF5(fname_sch_hdf5, "rho");
+    vx.Write_to_HDF5(fname_sch_hdf5, "vx");
+    vy.Write_to_HDF5(fname_sch_hdf5, "vy");
+    vz.Write_to_HDF5(fname_sch_hdf5, "vz");
+
+
     //======================================================================
     //... compute 3LPT displacement potential
     //======================================================================
@@ -230,7 +318,6 @@ int Run( ConfigFile& the_config )
 
     csoca::ilog << "-----------------------------------------------------------------------------" << std::endl;
 
-    // gadget2_output_interface gof( the_config );
 
     ///////////////////////////////////////////////////////////////////////
     // we store the densities here if we compute them