Made tau_ a Periodic_function

src/density/density.cpp

@@ -78,15 +78,20 @@ Density::Density(Simulation_context& ctx__)
         RTE_THROW("Simulation_context is not initialized");
     }
 
+    using pf  = Periodic_function<double>;
     using spf = Smooth_periodic_function<double>;
 
     /*  allocate charge density and magnetization on a coarse grid */
     for (int i = 0; i < ctx_.num_mag_dims() + 1; i++) {
         rho_mag_coarse_[i] = std::make_unique<spf>(ctx_.spfft_coarse<double>(), ctx_.gvec_coarse_fft_sptr());
     }
 
-    /// (WIP)TODO: allocate tau kinetic energy density => probably should add a check on wether used or not
-    tau_ = std::make_unique<spf>(ctx_.spfft_coarse<double>(), ctx_.gvec_fft_sptr());
+    if (ctx_.full_potential()) {
+        tau_ = std::make_unique<pf>(
+                ctx_, [&](int ia) { return lmax_t(ctx_.lmax_pot()); }, &ctx_.unit_cell().spl_num_atoms());
+    } else {
+        tau_ = std::make_unique<pf>(ctx_);
+    }
     tau_coarse_ = std::make_unique<spf>(ctx_.spfft_coarse<double>(), ctx_.gvec_coarse_fft_sptr());
 
     /* core density of the pseudopotential method */
@@ -718,16 +723,14 @@ Density::add_k_point_contribution_rg(K_point<T>* kp__, std::array<wf::Wave_funct
     density_rg.zero();
 
     ///(WIP)TODO: test if tau needed, only 1 dimension for now
-    mdarray<T, 2> tau_rg({nr, ctx_.num_mag_dims() + 1}, get_memory_pool(memory_t::host),
-                          mdarray_label("tau_rg"));
+    mdarray<T, 2> tau_rg({nr, ctx_.num_mag_dims() + 1}, get_memory_pool(memory_t::host), mdarray_label("tau_rg"));
     tau_rg.zero();
 
     /// create a PW buffer to hold derivative of Phi (WIP)TODO: assume CPU for now
     mdarray<std::complex<T>, 1> buf_pw({kp__->gkvec_fft_sptr()->count()}, get_memory_pool(ctx_.host_memory_t()));
-    //if (ctx_.processing_unit() == device_t::GPU) {
-    //    buf_pw.allocate(get_memory_pool(memory_t::device));
-    //}
-
+    // if (ctx_.processing_unit() == device_t::GPU) {
+    //     buf_pw.allocate(get_memory_pool(memory_t::device));
+    // }
 
     if (fft.processing_unit() == SPFFT_PU_GPU) {
         density_rg.allocate(get_memory_pool(memory_t::device)).zero(memory_t::device);
@@ -753,10 +756,10 @@ Density::add_k_point_contribution_rg(K_point<T>* kp__, std::array<wf::Wave_funct
                 ///(WIP)TODO: the above calculates |phi^2| from the G+k representation. We can use
                 ///           the same function, if we pass grad phi as an argument, and a factor 0.5
 
-                ///TODO: this assumes CPU, would need to add extra if for GPU
-                ///      will also need to copy back from GPU after that
-                /// Note: it seems that tau is correct
-                /// Can probably use to_gradient() here (see xc.cpp)
+                /// TODO: this assumes CPU, would need to add extra if for GPU
+                ///       will also need to copy back from GPU after that
+                ///  Note: it seems that tau is correct
+                ///  Can probably use to_gradient() here (see xc.cpp)
                 for (int x : {0, 1, 2}) {
                     #pragma omp parallel for
                     for (int igloc = 0; igloc < kp__->gkvec_fft_sptr()->count(); igloc++) {
@@ -765,8 +768,8 @@ Density::add_k_point_contribution_rg(K_point<T>* kp__, std::array<wf::Wave_funct
                         buf_pw[igloc] = wf_fft__[ispn].pw_coeffs(igloc, wf::band_index(i)) * static_cast<T>(gvc[x]);
                     }
                     add_k_point_contribution_rg_collinear(kp__->spfft_transform(), ispn, w,
-                                                          reinterpret_cast<T const*>(buf_pw.at(wf_mem)), 
-                                                          nr, ctx_.gamma_point(), tau_rg);
+                                                          reinterpret_cast<T const*>(buf_pw.at(wf_mem)), nr,
+                                                          ctx_.gamma_point(), tau_rg);
                 }
             }
         } // ispn
@@ -823,10 +826,10 @@ Density::add_k_point_contribution_rg(K_point<T>* kp__, std::array<wf::Wave_funct
                 rho_mag_coarse_[0]->value(ir) += density_rg(ir, 0); // rho
 
                 ///(WIP)TODO: add the 0.5 factor for kinetic density
-                //TODO: check if tau is correct in the first place:
-                //      - if replace G by 1, should get back the normal density
-                //      - should be able to compare to DFTK, if rho is also the same
-                tau_coarse_->value(ir) += 0.5*tau_rg(ir, 0);
+                // TODO: check if tau is correct in the first place:
+                //       - if replace G by 1, should get back the normal density
+                //       - should be able to compare to DFTK, if rho is also the same
+                tau_coarse_->value(ir) += 0.5 * tau_rg(ir, 0);
             }
         }
     }
@@ -1241,9 +1244,9 @@ Density::generate(K_point_set const& ks__, bool symmetrize__, bool add_core__, b
         symmetrize_field4d(*this);
         /// (WIP)TODO: make a test on tau or not. For now, assume spin-restricted
         std::vector<Smooth_periodic_function<double>*> tau_vec;
-        tau_vec.push_back(tau_.get()); 
-        symmetrize_pw_function(ctx_.unit_cell().symmetry(), ctx_.remap_gvec(), ctx_.sym_phase_factors(), 
-                               ctx_.num_mag_dims(), tau_vec);    
+        tau_vec.push_back(&tau_->rg());
+        symmetrize_pw_function(ctx_.unit_cell().symmetry(), ctx_.remap_gvec(), ctx_.sym_phase_factors(),
+                               ctx_.num_mag_dims(), tau_vec);
 
         if (ctx_.electronic_structure_method() == electronic_structure_method_t::pseudopotential) {
             std::unique_ptr<density_matrix_t> dm_ref;
@@ -1335,7 +1338,7 @@ Density::generate(K_point_set const& ks__, bool symmetrize__, bool add_core__, b
     if (transform_to_rg__) {
         this->fft_transform(1);
         /// (WIP)TODO: test on tau
-        tau_->fft_transform(1);
+        tau_->rg().fft_transform(1);
     }
 }
 
@@ -1499,7 +1502,7 @@ Density::generate_valence(K_point_set const& ks__)
     tau_coarse_->fft_transform(-1);
     /* map to fine G-vector grid */
     for (int igloc = 0; igloc < ctx_.gvec_coarse().count(); igloc++) {
-        tau_->f_pw_local(ctx_.gvec().gvec_base_mapping(igloc)) = tau_coarse_->f_pw_local(igloc);
+        tau_->rg().f_pw_local(ctx_.gvec().gvec_base_mapping(igloc)) = tau_coarse_->f_pw_local(igloc);
     }
 
     if (!ctx_.full_potential()) {
@@ -1848,7 +1851,7 @@ Density::generate_valence_mt()
                 /// (WIP)TODO: Calculate here (or around here) the muffin-tin part of tau
                 /// According to PhysRevB.91.075101, the kinetic energy density can be expressed as:
                 /// tau = 0.5*[ lapl(rho) - sum_jk lapl(psi_jk)*psi_jk - sum_jk psi_jk*lapl(psi_jk)]
-                /// The first bit, lapl(rho) is straight forward: rho is already calculated as a 
+                /// The first bit, lapl(rho) is straight forward: rho is already calculated as a
                 /// spectral function, and the laplacian is implemented for it.
                 /// For the second and third term, we need to use the relation:
                 /// lapl(u(r)*Y_lm) = [u''(r) + 2u'(r)/r - l(l+1)u(r)/r^2]*Y_lm, where only the radial
@@ -1858,7 +1861,6 @@ Density::generate_valence_mt()
                 /// the l quantum number with atom_type_.indexr(idxrf).am.l(), and the u''(r) and u'(r)
                 /// derivatives with a spline. Then, we contract with the density matrix the same way
                 /// it is done for the density
-
             }
         }
         for (int j = 0; j < ctx_.num_mag_dims() + 1; j++) {

src/density/density.hpp

@@ -227,7 +227,9 @@ class Density : public Field4D
     std::array<std::unique_ptr<Smooth_periodic_function<double>>, 4> rho_mag_coarse_;
 
     /// (WIP)TODO: kinetic energy density. For now, assume spin restricted system
-    std::unique_ptr<Smooth_periodic_function<double>> tau_;
+    /// Store tau_ as a Periodic function to have both PW and muffin-tin. In the
+    /// future, might want to store tau_ as a Field4D
+    std::unique_ptr<Periodic_function<double>> tau_;
     std::unique_ptr<Smooth_periodic_function<double>> tau_coarse_;
 
     /// Pointer to pseudo core charge density

src/dft/dft_ground_state.cpp

@@ -162,8 +162,8 @@ DFT_ground_state::check_scf_density()
                             << "Eold: " << gs0["energy"]["total"].get<double>()
                             << " Enew: " << gs1["energy"]["total"].get<double>() << std::endl;
 
-        std::vector<std::string> labels({"total", "vha", "vxc", "vtau", "exc", "bxc", "veff", "eval_sum", "kin", "ewald",
-                                         "vloc", "scf_correction", "entropy_sum"});
+        std::vector<std::string> labels({"total", "vha", "vxc", "vtau", "exc", "bxc", "veff", "eval_sum", "kin",
+                                         "ewald", "vloc", "scf_correction", "entropy_sum"});
 
         for (auto e : labels) {
             RTE_OUT(ctx_.out()) << "energy component: " << e << ", diff: "

src/dft/energy.cpp

@@ -157,8 +157,8 @@ ks_energy(Simulation_context const& ctx, std::map<std::string, double> const& en
         }
 
         case electronic_structure_method_t::pseudopotential: {
-            tot_en = energies.at("valence_eval_sum") - energies.at("vxc") - energies.at("bxc") -
-                     energies.at("vtau") - energies.at("PAW_one_elec");
+            tot_en = energies.at("valence_eval_sum") - energies.at("vxc") - energies.at("bxc") - energies.at("vtau") -
+                     energies.at("PAW_one_elec");
             tot_en += -0.5 * energies.at("vha") + energies.at("exc") + energies.at("PAW_total_energy") +
                       energies.at("ewald");
             if (ctx.hubbard_correction()) {
@@ -243,8 +243,8 @@ double
 one_electron_energy(Density const& density, Potential const& potential)
 {
     return energy_vha(potential) + energy_vxc(density, potential) + energy_bxc(density, potential) +
-           energy_vtau(density, potential) +
-           potential.PAW_one_elec_energy(density) + one_electron_energy_hubbard(density, potential);
+           energy_vtau(density, potential) + potential.PAW_one_elec_energy(density) +
+           one_electron_energy_hubbard(density, potential);
 }
 
 double

src/hamiltonian/hamiltonian.hpp

@@ -575,18 +575,18 @@ Hamiltonian0<T>::operator()(K_point<T>& kp__) const
  *  \f$ \kappa = \{\ell, \lambda\} \f$. We will focus on the product of two gradients and use
  *  the following relation
  *  \f[
- *    \nabla \varphi_{\xi}^{*} \nabla \varphi_{\xi'} = 
+ *    \nabla \varphi_{\xi}^{*} \nabla \varphi_{\xi'} =
  *    \frac{1}{2} \left( \nabla^2 (\varphi_{\xi}^{*} \varphi_{\xi'}) - \varphi_{\xi}^{*} \nabla^2 \varphi_{\xi'} -
  *     \varphi_{\xi'} \nabla^2 \varphi_{\xi}^{*} \right) = \sum_{L}p_{L}^{ij}(r) R_{L}(\hat{\bf r})
  *  \f]
  *  Laplacian in spherical coordinates acts on spherical function in the following way:
  *  \f[
- *  \nabla^2 \left( f Y_{L} \right) = \left( \frac{d^2 f}{dr^2} + \frac{2}{r} \frac{df}{dr} - 
+ *  \nabla^2 \left( f Y_{L} \right) = \left( \frac{d^2 f}{dr^2} + \frac{2}{r} \frac{df}{dr} -
  *    \frac{\ell(\ell + 1)}{r^2} f(r) \right) Y_{L}(\hat {\bf r}) = \tilde f(r)Y_{L}(\hat {\bf r})
  *  \f]
  *  So the task is to find radial functions \f$ p_{L}^{ij}(r) \f$ of the product of gradients of spherical functions.
  *  We are aiming at expansion in real spherical harmonics because later we are going to integrate it with the
- *  tau-potential which is also expanded in real harmonics. 
+ *  tau-potential which is also expanded in real harmonics.
  *
  *  First term:
  *  \f[
@@ -614,7 +614,7 @@ Hamiltonian0<T>::operator()(K_point<T>& kp__) const
  *  Third term:
  *  \f[
  *    \varphi_{\xi'} \nabla^2 \varphi_{\xi}^{*} =
- *     f_{\kappa_{\xi'}}(r) Y_{L_{\xi'}} \tilde f_{\kappa_{\xi}}(r) Y_{L_{\xi}}^{*} = 
+ *     f_{\kappa_{\xi'}}(r) Y_{L_{\xi'}} \tilde f_{\kappa_{\xi}}(r) Y_{L_{\xi}}^{*} =
  *     \sum_{L}  p_{2, L}^{\kappa \kappa'}(r) R_{L}
  *  \f]
  *  where
@@ -625,15 +625,15 @@ Hamiltonian0<T>::operator()(K_point<T>& kp__) const
  *  \f]
  *  We can now combine three terms:
  *  \f[
- *    p_{L}^{\kappa,\kappa'}(r) = \frac{1}{2} \langle Y_{L_{\xi}} | R_{L} | Y_{L_{\xi'}} \rangle \left( 
- *      \widetilde{f_{\kappa_{\xi}} f_{\kappa_{\xi'}}} -  f_{\kappa_{\xi}} \tilde f_{\kappa_{\xi'}} - 
+ *    p_{L}^{\kappa,\kappa'}(r) = \frac{1}{2} \langle Y_{L_{\xi}} | R_{L} | Y_{L_{\xi'}} \rangle \left(
+ *      \widetilde{f_{\kappa_{\xi}} f_{\kappa_{\xi'}}} -  f_{\kappa_{\xi}} \tilde f_{\kappa_{\xi'}} -
  *      \tilde f_{\kappa_{\xi}} f_{\kappa_{\xi'}} \right)
  *  \f]
  *  Now we can compute matrix elements of \f$ V_{\tau}({\bf r}) \f$ in muffin-tins:
  *  \f[
- *    \langle \nabla \varphi_{\xi}^{*} | \sum_{L} V_{L}^{\tau}(r) R_{L} | \nabla \varphi_{\xi'} \rangle = 
+ *    \langle \nabla \varphi_{\xi}^{*} | \sum_{L} V_{L}^{\tau}(r) R_{L} | \nabla \varphi_{\xi'} \rangle =
  *      \sum_{L} \frac{1}{2} \langle Y_{L_{\xi}} | R_{L} | Y_{L_{\xi'}} \rangle
- *      \int \left( \widetilde{f_{\kappa_{\xi}} f_{\kappa_{\xi'}}} -  f_{\kappa_{\xi}} \tilde f_{\kappa_{\xi'}} - 
+ *      \int \left( \widetilde{f_{\kappa_{\xi}} f_{\kappa_{\xi'}}} -  f_{\kappa_{\xi}} \tilde f_{\kappa_{\xi'}} -
  *      \tilde f_{\kappa_{\xi}} f_{\kappa_{\xi'}} \right) V^{\tau}_{L}(r) r^2 dr
  *  \f]
  *  Remember that radial integrals are computed for each atom in the muffin-tin. Atomic index \f$ \alpha \f$ is ommited

src/hamiltonian/local_operator.cpp

@@ -131,8 +131,9 @@ Local_operator<T>::Local_operator(Simulation_context const& ctx__, fft::spfft_tr
             #pragma omp parallel for schedule(static)
             for (int igloc = 0; igloc < gvec_coarse_p_->gvec().count(); igloc++) {
                 /* map from fine to coarse set of G-vectors */
-                veff_vec_[6]->f_pw_local(igloc) = 0.5*potential__->tau_potential().rg().f_pw_local(
-                        potential__->tau_potential().rg().gvec().gvec_base_mapping(igloc));
+                veff_vec_[6]->f_pw_local(igloc) =
+                        0.5 * potential__->tau_potential().rg().f_pw_local(
+                                      potential__->tau_potential().rg().gvec().gvec_base_mapping(igloc));
             }
             /* transform to real space */
             veff_vec_[6]->fft_transform(1);
@@ -344,9 +345,9 @@ Local_operator<T>::apply_h(fft::spfft_transform_type<T>& spfftk__, std::shared_p
 
     /// (WIP)TODO: transform the gradient of the wave function to real space, for a given coord
     /// Can probably use to_gradient() here (see xc.cpp)
-    auto gradphi_to_r = [&](wf::spin_index ispn, wf::band_index i, int x){
+    auto gradphi_to_r = [&](wf::spin_index ispn, wf::band_index i, int x) {
         auto phi_mem = phi_fft[ispn.get()].on_device() ? memory_t::device : memory_t::host;
-        ///TODO: this assumes CPU only
+        /// TODO: this assumes CPU only
         #pragma omp parallel for
         for (int ig = 0; ig < ngv_fft; ig++) {
             buf_pw_[ig] = phi_fft[ispn].pw_coeffs(ig, i) * static_cast<T>(gkvec_cart_(ig, x));
@@ -360,8 +361,8 @@ Local_operator<T>::apply_h(fft::spfft_transform_type<T>& spfftk__, std::shared_p
     };
 
     /// (WIP)TODO: compute the divergence of vphi(G) in place, for a given direction
-    auto div_vphi_G = [&](int x){
-        ///TODO: this assumes CPU only
+    auto div_vphi_G = [&](int x) {
+        /// TODO: this assumes CPU only
         #pragma omp parallel for
         for (int ig = 0; ig < ngv_fft; ig++) {
             vphi_[ig] *= static_cast<T>(gkvec_cart_(ig, x));
@@ -372,7 +373,7 @@ Local_operator<T>::apply_h(fft::spfft_transform_type<T>& spfftk__, std::shared_p
        spin block (ispn_block) is used as a bit mask:
         - first bit: spin component which is updated
         - second bit: add or not kinetic energy term */
-    auto add_to_hphi = [&](int ispn_block, wf::band_index i, int add_ekin=1) {
+    auto add_to_hphi = [&](int ispn_block, wf::band_index i, int add_ekin = 1) {
         /* index of spin component */
         int ispn = ispn_block & 1;
         /* add kinetic energy if this is a diagonal block */

src/potential/potential.cpp

@@ -325,11 +325,10 @@ Potential::generate(Density const& density__, bool use_symmetry__, bool transfor
         ///            final density might be correct, but not energy
         /// Note: seems to do nothing at all, probably that sinve tau is symmetrize, Vtau is too
         std::vector<Smooth_periodic_function<double>*> tau_vec;
-        tau_vec.push_back(&tau_potential_->rg()); 
-        symmetrize_pw_function(ctx_.unit_cell().symmetry(), ctx_.remap_gvec(), ctx_.sym_phase_factors(), 
+        tau_vec.push_back(&tau_potential_->rg());
+        symmetrize_pw_function(ctx_.unit_cell().symmetry(), ctx_.remap_gvec(), ctx_.sym_phase_factors(),
                                ctx_.num_mag_dims(), tau_vec);
 
-
         if (transform_to_rg__) {
             /* transform potential to real space after symmetrization */
             this->fft_transform(1);

src/potential/potential.hpp

@@ -875,7 +875,7 @@ class Potential : public Field4D
     auto
     energy_vtau(Density const& density__) const
     {
-        return inner(density__.tau(), tau_potential().rg());
+        return inner(density__.tau(), tau_potential());
     }
 
     /// Integral of \f$ \rho({\bf r}) \epsilon^{XC}({\bf r}) \f$.

src/potential/xc.cpp

@@ -31,7 +31,7 @@ Potential::xc_rg_nonmagnetic(Density const& density__, bool use_lapl__)
 
     bool is_gga = is_gradient_correction();
     bool is_tau = is_meta_tau();
-    
+
     int num_points = ctx_.spfft<double>().local_slice_size();
 
     Smooth_periodic_function<double> rho(ctx_.spfft<double>(), gvp);
@@ -55,7 +55,7 @@ Potential::xc_rg_nonmagnetic(Density const& density__, bool use_lapl__)
 
         rhomin        = std::min(rhomin, d);
         rho.value(ir) = std::max(d, 0.0);
-    } 
+    }
     mpi::Communicator(ctx_.spfft<double>().communicator()).allreduce<double, mpi::op_t::min>(&rhomin, 1);
     /* even a small negative density is a sign of something bing wrong; don't remove this check */
     if (rhomin < 0.0 && ctx_.comm().rank() == 0) {
@@ -101,19 +101,19 @@ Potential::xc_rg_nonmagnetic(Density const& density__, bool use_lapl__)
     for (int ir = 0; ir < num_points; ir++) {
 
         // int ir = spl_np[irloc];
-        double t = grad_rho_grad_rho.value(ir)/(8.0*rho.value(ir));
+        double t = grad_rho_grad_rho.value(ir) / (8.0 * rho.value(ir));
 
-        //tau.value(ir) = std::max(density__.tau().value(ir), t);
-        tau.value(ir) = std::max(density__.tau().value(ir), 0.0);
+        // tau.value(ir) = std::max(density__.tau().value(ir), t);
+        tau.value(ir) = std::max(density__.tau().rg().value(ir), 0.0);
     }
 
     mdarray<double, 1> exc({num_points}, mdarray_label("exc_tmp"));
     mdarray<double, 1> vxc({num_points}, mdarray_label("vxc_tmp"));
 
     /// (WIP)TODO: meta-GGA quantities. For now, pass laplacian as null_ptr
     mdarray<double, 1> vtau({num_points}, mdarray_label("vtau_tmp"));
-    double* lapl = nullptr; 
-    double* vlapl = nullptr; 
+    double* lapl  = nullptr;
+    double* vlapl = nullptr;
 
     /* loop over XC functionals */
     for (auto& ixc : xc_func_) {
@@ -153,11 +153,10 @@ Potential::xc_rg_nonmagnetic(Density const& density__, bool use_lapl__)
                     /// (WIP):TODO get the tau potential. Ignore the laplacian for now
                     if (ixc.is_mgga()) {
                         ixc.get_meta(spl_t.local_size(), &rho.value(spl_t.global_offset()),
-                                     &grad_rho_grad_rho.value(spl_t.global_offset()),
-                                     lapl, &tau.value(spl_t.global_offset()),
-                                     vxc.at(memory_t::host, spl_t.global_offset()), 
-                                     &vsigma.value(spl_t.global_offset()),
-                                     vlapl, vtau.at(memory_t::host, spl_t.global_offset()),
+                                     &grad_rho_grad_rho.value(spl_t.global_offset()), lapl,
+                                     &tau.value(spl_t.global_offset()), vxc.at(memory_t::host, spl_t.global_offset()),
+                                     &vsigma.value(spl_t.global_offset()), vlapl,
+                                     vtau.at(memory_t::host, spl_t.global_offset()),
                                      exc.at(memory_t::host, spl_t.global_offset()));
                     }
                 } // omp parallel region