LCAO_deepks.h

Go to the documentation of this file.

#ifndef LCAO_DEEPKS_H
#define LCAO_DEEPKS_H
 
#ifdef __MLALGO
 
#include "deepks_basic.h"
#include "deepks_param.h"
#include "deepks_vdelta.h"
#include "source_base/complexmatrix.h"
#include "source_base/matrix.h"
#include "source_basis/module_ao/parallel_orbitals.h"
#include "source_basis/module_nao/two_center_integrator.h"
#include "source_cell/module_neighbor/sltk_grid_driver.h"
#include "source_lcao/module_hcontainer/hcontainer.h"
 
#include <torch/script.h>
#include <torch/torch.h>
 
// caoyu add 2021-03-29
// wenfei modified 2022-1-5
//
template <typename T>
class LCAO_Deepks
{
 
    //-------------------
    // public variables
    //-------------------
  public:
    double E_delta = 0.0;
    double e_delta_band = 0.0;
 
    std::vector<std::vector<T>> V_delta;
 
    //-------------------
    // private variables
    //-------------------
    //  private:
  public:                      // change to public to reconstuct the code, 2024-07-22 by mohan
    DeePKS_Param deepks_param; // parameters for DeePKS
 
    bool init_pdm = false; // for DeePKS NSCF calculation, set init_pdm to skip the calculation of pdm in SCF iteration
 
    // deep neural network module that provides corrected Hamiltonian term and
    // related derivatives. Used in cal_edelta_gedm.
    torch::jit::script::Module model_deepks;
 
    // saves <phi(0)|alpha(R)> and its derivatives
    // index 0 for itself and index 1-3 for derivatives over x,y,z
    std::vector<hamilt::HContainer<double>*> phialpha;
 
    // density matrix in real space
    hamilt::HContainer<double>* dm_r = nullptr;
 
    // projected density matrix
    // [tot_Inl][2l+1][2l+1], here l is corresponding to inl;
    // [nat][nlm*nlm] for equivariant version
    std::vector<torch::Tensor> pdm;
 
    double** gedm = nullptr; //[tot_Inl][(2l+1)*(2l+1)]
 
    // functions for hr status: 1. get value; 2. set value;
    int get_hr_cal()
    {
        return this->hr_cal;
    }
    void set_hr_cal(bool cal)
    {
        this->hr_cal = cal;
    }
 
    //-------------------
    // LCAO_deepks.cpp
    //-------------------
 
    // This file contains constructor and destructor of the class LCAO_deepks,
    // as well as subroutines for initializing and releasing relevant data structures
 
    // Other than the constructor and the destructor, it contains 3 types of subroutines:
    // 1. subroutines that are related to calculating descriptors:
    //   - init : allocates some arrays
    //   - init_index : records the index (inl)
    // 2. subroutines that are related to V_delta:
    //   - allocate_V_delta : allocates V_delta; if calculating force, it also allocates F_delta
 
  public:
    explicit LCAO_Deepks();
    ~LCAO_Deepks();
 
    void init(const LCAO_Orbitals& orb,
              const int nat,
              const int ntype,
              const int nks,
              const Parallel_Orbitals& pv_in,
              std::vector<int> na,
              std::ofstream& ofs);
 
    void allocate_V_delta(const int nat, const int nks = 1);
 
    void init_DMR(const UnitCell& ucell,
                  const LCAO_Orbitals& orb,
                  const Parallel_Orbitals& pv,
                  const Grid_Driver& GridD);
 
    void dpks_cal_e_delta_band(const std::vector<std::vector<T>>& dm, const int nks);
 
  private:
    // flag of HR status,
    // true : HR should be calculated
    // false : HR has been calculated
    bool hr_cal = true;
 
    // arrange index of descriptor in all atoms
    void init_index(const int ntype,
                    const int nat,
                    std::vector<int> na,
                    const int tot_inl,
                    const LCAO_Orbitals& orb,
                    std::ofstream& ofs);
 
    const Parallel_Orbitals* pv = nullptr;
};
 
#endif
#endif