Skip to content

Feature&Doc&Tests: Support for deepks_v_delta=-1&-2 #6245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 14 commits into from
Jun 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 6 additions & 3 deletions docs/advanced/input_files/input-main.md
Original file line number Diff line number Diff line change
Expand Up @@ -2145,7 +2145,7 @@ Warning: this function is not robust enough for the current version. Please try
- **Description**: Print labels and descriptors for DeePKS in OUT.${suffix}. The names of these files start with "deepks".
- 0 : No output.
- 1 : Output intermediate files needed during DeePKS training.
- 2 : Output target labels for label preperation. The label files are named as `deepks_<property>.npy`, where the units and formats are the same as label files `<property>.npy` required for training, except that the first dimension (`nframes`) is excluded. System structrue files are also given in `deepks_atom.npy` and `deepks_box.npy` in the unit of *Bohr*, which means `lattice_constant` should be set to 1 when training.
- 2 : Output target labels for label preperation. The label files are named as `deepks_<property>.npy` or `deepks_<property>.csr`, where the units and formats are the same as label files `<property>.npy` or `<property>.csr` required for training, except that the first dimension (`nframes`) is excluded. System structrue files are also given in `deepks_atom.npy` and `deepks_box.npy` in the unit of *Bohr*, which means `lattice_constant` should be set to 1 when training.
- **Note**: When `deepks_out_labels` equals **1**, the path of a numerical descriptor (an `orb` file) is needed to be specified under the `NUMERICAL_DESCRIPTOR` tag in the `STRU` file. For example:

```text
Expand Down Expand Up @@ -2252,8 +2252,11 @@ Warning: this function is not robust enough for the current version. Please try

- **Type**: int
- **Availability**: numerical atomic orbital basis
- **Description**: Include V_delta label for DeePKS training. When `deepks_out_labels` is true and `deepks_v_delta` > 0, ABACUS will output h_base.npy, v_delta.npy and h_tot.npy(h_tot=h_base+v_delta).
Meanwhile, when `deepks_v_delta` equals 1, ABACUS will also output v_delta_precalc.npy, which is used to calculate V_delta during DeePKS training. However, when the number of atoms grows, the size of v_delta_precalc.npy will be very large. In this case, it's recommended to set `deepks_v_delta` as 2, and ABACUS will output phialpha.npy and grad_evdm.npy but not v_delta_precalc.npy. These two files are small and can be used to calculate v_delta_precalc in the procedure of training DeePKS.
- **Description**: Include V_delta/V_delta_R (Hamiltonian in k/real space) label for DeePKS training. When `deepks_out_labels` is true and `deepks_v_delta` > 0 (k space), ABACUS will output `deepks_hbase.npy`, `deepks_vdelta.npy` and `deepks_htot.npy`(htot=hbase+vdelta). When `deepks_out_labels` is true and `deepks_v_delta` < 0 (real space), ABACUS will output `deepks_hrtot.csr`, `deepks_hrdelta.csr`. Some more files output for different settings.
- `deepks_v_delta` = 1: `deepks_vdpre.npy`, which is used to calculate V_delta during DeePKS training.
- `deepks_v_delta` = 2: `deepks_phialpha.npy` and `deepks_gevdm.npy`, which can be used to calculate `deepks_vdpre.npy`. A recommanded method for memory saving.
- `deepks_v_delta` = -1: `deepks_vdrpre.npy`, which is used to calculate V_delta_R during DeePKS training.
- `deepks_v_delta` = -2: `deepks_phialpha_r.npy` and `deepks_gevdm.npy`, which can be used to calculate `deepks_vdrpre.npy`. A recommanded method for memory saving.
- **Default**: 0

### deepks_out_unittest
Expand Down
107 changes: 68 additions & 39 deletions source/module_hamilt_lcao/module_deepks/LCAO_deepks_interface.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -2,12 +2,12 @@
#include "LCAO_deepks_interface.h"

#include "LCAO_deepks_io.h" // mohan add 2024-07-22
#include "source_base/global_variable.h"
#include "source_base/tool_title.h"
#include "module_elecstate/cal_dm.h"
#include "module_hamilt_lcao/module_hcontainer/hcontainer.h"
#include "module_hamilt_lcao/module_hcontainer/output_hcontainer.h"
#include "module_parameter/parameter.h"
#include "source_base/global_variable.h"
#include "source_base/tool_title.h"

template <typename TK, typename TR>
LCAO_Deepks_Interface<TK, TR>::LCAO_Deepks_Interface(std::shared_ptr<LCAO_Deepks<TK>> ld_in) : ld(ld_in)
Expand Down Expand Up @@ -353,14 +353,15 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
} // end deepks_out_labels == 1
}

// not add deepks_out_labels = 2 for HR yet
// H(R) matrix part, for HR, base will not be calculated since they are HContainer objects
if (PARAM.inp.deepks_v_delta < 0)
{
// set the output
const double sparse_threshold = 1e-10;
const int precision = 8;
const std::string file_hrtot = PARAM.globalv.global_out_dir + "deepks_hrtot.csr";
const std::string file_hrtot
= PARAM.globalv.global_out_dir
+ (PARAM.inp.deepks_out_labels == 1 ? "deepks_hrtot.csr" : "deepks_hamiltonian_r.csr");
hamilt::HContainer<TR>* hR_tot = (p_ham->getHR());

if (rank == 0)
Expand All @@ -369,47 +370,75 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
ofs_hr << "Matrix Dimension of H(R): " << hR_tot->get_nbasis() << std::endl;
ofs_hr << "Matrix number of H(R): " << hR_tot->size_R_loop() << std::endl;
hamilt::Output_HContainer<TR> out_hr(hR_tot, ofs_hr, sparse_threshold, precision);
out_hr.write();
out_hr.write(true); // write all the matrices, including empty ones
ofs_hr.close();
}

if (PARAM.inp.deepks_scf)
{
const std::string file_vdeltar = PARAM.globalv.global_out_dir + "deepks_hrdelta.csr";
hamilt::HContainer<TR>* h_deltaR = p_ham->get_V_delta_R();

if (rank == 0)
if (PARAM.inp.deepks_out_labels == 1)
{
std::ofstream ofs_hr(file_vdeltar, std::ios::out);
ofs_hr << "Matrix Dimension of H_delta(R): " << h_deltaR->get_nbasis() << std::endl;
ofs_hr << "Matrix number of H_delta(R): " << h_deltaR->size_R_loop() << std::endl;
hamilt::Output_HContainer<TR> out_hr(h_deltaR, ofs_hr, sparse_threshold, precision);
out_hr.write();
ofs_hr.close();
}
const std::string file_vdeltar = PARAM.globalv.global_out_dir + "deepks_hrdelta.csr";
hamilt::HContainer<TR>* h_deltaR = p_ham->get_V_delta_R();

torch::Tensor phialpha_r_out;
torch::Tensor R_query;
DeePKS_domain::prepare_phialpha_r(nlocal,
lmaxd,
inlmax,
nat,
phialpha,
ucell,
orb,
*ParaV,
GridD,
phialpha_r_out,
R_query);
const std::string file_phialpha_r = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
const std::string file_R_query = PARAM.globalv.global_out_dir + "deepks_R_query.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_phialpha_r, phialpha_r_out, rank);
LCAO_deepks_io::save_tensor2npy<int>(file_R_query, R_query, rank);

torch::Tensor gevdm_out;
DeePKS_domain::prepare_gevdm(nat, lmaxd, inlmax, orb, gevdm, gevdm_out);
const std::string file_gevdm = PARAM.globalv.global_out_dir + "deepks_gevdm.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_gevdm, gevdm_out, rank);
if (rank == 0)
{
std::ofstream ofs_hr(file_vdeltar, std::ios::out);
ofs_hr << "Matrix Dimension of H_delta(R): " << h_deltaR->get_nbasis() << std::endl;
ofs_hr << "Matrix number of H_delta(R): " << h_deltaR->size_R_loop() << std::endl;
hamilt::Output_HContainer<TR> out_hr(h_deltaR, ofs_hr, sparse_threshold, precision);
out_hr.write(true); // write all the matrices, including empty ones
ofs_hr.close();
}

if (PARAM.inp.deepks_v_delta == -1)
{
int R_size = DeePKS_domain::get_R_size(*h_deltaR);
torch::Tensor vdr_precalc;
DeePKS_domain::cal_vdr_precalc(nlocal,
lmaxd,
inlmax,
nat,
nks,
R_size,
inl2l,
kvec_d,
phialpha,
gevdm,
inl_index,
ucell,
orb,
*ParaV,
GridD,
vdr_precalc);

const std::string file_vdrpre = PARAM.globalv.global_out_dir + "deepks_vdrpre.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_vdrpre, vdr_precalc, rank);
}
else if (PARAM.inp.deepks_v_delta == -2)
{
int R_size = DeePKS_domain::get_R_size(*h_deltaR);
torch::Tensor phialpha_r_out;
DeePKS_domain::prepare_phialpha_r(nlocal,
lmaxd,
inlmax,
nat,
R_size,
phialpha,
ucell,
orb,
*ParaV,
GridD,
phialpha_r_out);
const std::string file_phialpha_r = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_phialpha_r, phialpha_r_out, rank);

torch::Tensor gevdm_out;
DeePKS_domain::prepare_gevdm(nat, lmaxd, inlmax, orb, gevdm, gevdm_out);
const std::string file_gevdm = PARAM.globalv.global_out_dir + "deepks_gevdm.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_gevdm, gevdm_out, rank);
}
}
}
}

Expand Down Expand Up @@ -544,7 +573,7 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
<< " = " << std::setprecision(8) << e_delta_band * ModuleBase::Ry_to_eV << " eV" << std::endl;
ofs_running << " E_delta_NN = " << std::setprecision(8) << E_delta << " Ry"
<< " = " << std::setprecision(8) << E_delta * ModuleBase::Ry_to_eV << " eV" << std::endl;
ofs_running << " -----------------------------------------------" << std::endl;
ofs_running << " -----------------------------------------------" << std::endl;
}
if (PARAM.inp.deepks_out_unittest)
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,9 @@

#ifdef __MLALGO
#include "LCAO_deepks.h"
#include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"
#include "source_base/complexmatrix.h"
#include "source_base/matrix.h"
#include "module_hamilt_lcao/hamilt_lcaodft/hamilt_lcao.h"

#include <memory>

Expand Down
21 changes: 15 additions & 6 deletions source/module_hamilt_lcao/module_deepks/deepks_check.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,8 @@ void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string&
{
return;
}
using T_tensor = typename std::conditional<std::is_same<T, std::complex<double>>::value, c10::complex<double>, T>::type;
using T_tensor =
typename std::conditional<std::is_same<T, std::complex<double>>::value, c10::complex<double>, T>::type;

std::ofstream ofs(filename.c_str());
ofs << std::setprecision(10);
Expand All @@ -21,15 +22,18 @@ void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string&

// stride for each dimension
std::vector<int64_t> strides(ndim, 1);
for (int i = ndim - 2; i >= 0; --i) {
for (int i = ndim - 2; i >= 0; --i)
{
strides[i] = strides[i + 1] * sizes[i + 1];
}

for (int64_t idx = 0; idx < numel; ++idx) {
for (int64_t idx = 0; idx < numel; ++idx)
{
// index to multi-dimensional indices
std::vector<int64_t> indices(ndim);
int64_t tmp = idx;
for (int d = 0; d < ndim; ++d) {
for (int d = 0; d < ndim; ++d)
{
indices[d] = tmp / strides[d];
tmp = tmp % strides[d];
}
Expand All @@ -39,16 +43,21 @@ void DeePKS_domain::check_tensor(const torch::Tensor& tensor, const std::string&
ofs << *tmp_ptr;

// print space or newline
if ( ( (idx+1) % sizes[ndim-1] ) == 0 ) {
if (((idx + 1) % sizes[ndim - 1]) == 0)
{
ofs << std::endl;
} else {
}
else
{
ofs << " ";
}
}

ofs.close();
}



template void DeePKS_domain::check_tensor<int>(const torch::Tensor& tensor, const std::string& filename, const int rank);
template void DeePKS_domain::check_tensor<double>(const torch::Tensor& tensor, const std::string& filename, const int rank);
template void DeePKS_domain::check_tensor<std::complex<double>>(const torch::Tensor& tensor, const std::string& filename, const int rank);
Expand Down
Loading
Loading