Skip to content

Feature: Add direct gradient label in DeePKS. #7509

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.

Sign up for GitHub

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

Draft
ErjieWu wants to merge 8 commits into
base: develop
Choose a base branch
from
Draft

Feature: Add direct gradient label in DeePKS. #7509

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
fbfd1ce
Add deepks_grad method.
ErjieWu Jun 23, 2026
f0e287d
Merge branch 'deepmodeling:develop' into grad
ErjieWu Jun 23, 2026
bd87ab3
Accelerate cal_vdr_precalc.
ErjieWu Jun 23, 2026
887b324
clang-format change.
ErjieWu Jun 23, 2026
264fc58
Accelerate cal_vdrpre_square.
ErjieWu Jun 23, 2026
1c15df4
Update deepks_grad method and output_hcontainer.
ErjieWu Jul 17, 2026
3cc2c50
Merge branch 'develop' into grad
ErjieWu Jul 17, 2026
20ba86f
Update parameters.yaml
ErjieWu Jul 17, 2026
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
2 changes: 1 addition & 1 deletion docs/advanced/input_files/input-main.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2595,7 +2595,7 @@

- **Type**: Integer
- **Availability**: *Numerical atomic orbital basis*
- **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. NOTICE: To match the unit Normally used in DeePKS, the unit of Hamiltonian in k space is Hartree. However, currently in R space the unit is still Ry.
- **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. All DeePKS energy and Hamiltonian labels are output in Hartree.
- 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.
Expand Down
9 changes: 8 additions & 1 deletion docs/parameters.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2396,14 +2396,21 @@ parameters:
category: DeePKS
type: Integer
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. NOTICE: To match the unit Normally used in DeePKS, the unit of Hamiltonian in k space is Hartree. However, currently in R space the unit is still Ry.
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. All DeePKS energy and Hamiltonian labels are output in Hartree.
* 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_value: "0"
unit: ""
availability: Numerical atomic orbital basis
- name: deepks_grad
category: DeePKS
type: Boolean
description: ""
default_value: "False"
unit: ""
availability: "NAO basis; requires deepks_out_labels >= 1 and deepks_v_delta < 0"
- name: deepks_out_unittest
category: DeePKS
type: Boolean
Expand Down
1 change: 1 addition & 0 deletions source/Makefile.Objects
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -219,6 +219,7 @@ OBJS_DEEPKS=LCAO_deepks.o\
deepks_vdelta.o\
deepks_vdpre.o\
deepks_vdrpre.o\
deepks_grad.o\
deepks_pdm.o\
deepks_phialpha.o\
LCAO_deepks_io.o\
Expand Down
1 change: 1 addition & 0 deletions source/source_io/module_parameter/input_parameter.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,7 @@ struct Input_para
int deepks_bandgap = 0; ///< for bandgap label. QO added 2021-12-15
std::vector<int> deepks_band_range = {-1, 0}; ///< the range of bands to calculate bandgap
int deepks_v_delta = 0; ///< for v_delta label. xuan added
bool deepks_grad = false; ///< output descriptor-gradient label intermediates for DeePKS training
bool deepks_equiv = false; ///< whether to use equivariant version of DeePKS
bool deepks_out_unittest = false; ///< if set to true, prints intermediate quantities that shall
///< be used for making unit test
Expand Down
20 changes: 19 additions & 1 deletion source/source_io/module_parameter/read_input_item_deepks.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -290,7 +290,7 @@ void ReadInput::item_deepks()
" -1 for vdrpre, -2 for phialpha_r and gevdm (can save memory)";
item.category = "DeePKS";
item.type = "Integer";
item.description = R"(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. NOTICE: To match the unit Normally used in DeePKS, the unit of Hamiltonian in k space is Hartree. However, currently in R space the unit is still Ry.
item.description = R"(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. All DeePKS energy and Hamiltonian labels are output in Hartree.
* 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.
Expand All @@ -307,6 +307,24 @@ void ReadInput::item_deepks()
};
this->add_item(item);
}
{
Input_Item item("deepks_grad");
item.annotation = "output descriptor-gradient label intermediates for DeePKS training";
item.category = "DeePKS";
item.type = "Boolean";
item.default_value = "False";
item.unit = "";
item.availability = "NAO basis; requires deepks_out_labels >= 1 and deepks_v_delta < 0";
read_sync_bool(input.deepks_grad);
item.check_value = [](const Input_Item& item, const Parameter& para) {
if (para.input.deepks_grad && para.input.deepks_out_labels == 0)
{
ModuleBase::WARNING_QUIT("ReadInput",
"deepks_grad requires deepks_out_labels 1");
}
};
this->add_item(item);
}
{
Input_Item item("deepks_out_unittest");
item.annotation = "if set 1, prints intermediate quantities that shall "
Expand Down
1 change: 1 addition & 0 deletions source/source_lcao/module_deepks/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ if(ENABLE_MLALGO)
deepks_vdelta.cpp
deepks_vdpre.cpp
deepks_vdrpre.cpp
deepks_grad.cpp
deepks_pdm.cpp
deepks_phialpha.cpp
LCAO_deepks_io.cpp
Expand Down
138 changes: 131 additions & 7 deletions source/source_lcao/module_deepks/LCAO_deepks_interface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
#include "source_lcao/module_deepks/deepks_check.h"
#include "source_lcao/module_deepks/deepks_descriptor.h"
#include "source_lcao/module_deepks/deepks_fpre.h"
#include "source_lcao/module_deepks/deepks_grad.h"
#include "source_lcao/module_deepks/deepks_orbital.h"
#include "source_lcao/module_deepks/deepks_orbpre.h"
#include "source_lcao/module_deepks/deepks_pdm.h"
Expand Down Expand Up @@ -107,8 +108,9 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
const bool output_base
= PARAM.inp.deepks_out_labels == 1
&& is_after_scf; // not output when deepks_out_labels=2 and in electronic step (output true base elsewhere)
const bool output_precalc
= (PARAM.inp.deepks_out_labels == 1) && (PARAM.inp.deepks_scf || PARAM.inp.deepks_out_freq_elec);
const bool output_precalc = (PARAM.inp.deepks_out_labels == 1)
&& (PARAM.inp.deepks_scf || PARAM.inp.deepks_out_freq_elec
|| PARAM.inp.deepks_grad);

//================================================================================
// 1. Update real-space density matrix (DMR) for deepks, projected density matrix (PDM)
Expand Down Expand Up @@ -159,7 +161,13 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
// new gedm is also useful in cal_f_delta, so it should be ld->gedm
if (PARAM.inp.deepks_equiv)
{
DeePKS_domain::cal_edelta_gedm_equiv(nat, deepks_param, descriptor, ld->model_deepks, ld->gedm, E_delta, rank);
DeePKS_domain::cal_edelta_gedm_equiv(nat,
deepks_param,
descriptor,
ld->model_deepks,
ld->gedm,
E_delta,
rank);
}
else // traditional version
{
Expand Down Expand Up @@ -262,6 +270,15 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
}
const std::string file_gradvx = get_filename("gradvx", PARAM.inp.deepks_out_labels, iter);
LCAO_deepks_io::save_tensor2npy<double>(file_gradvx, gvx, rank);
if (PARAM.inp.deepks_grad)
{
torch::Tensor gvx_square;
DeePKS_domain::cal_gvx_square(gvx, gvx_square);
LCAO_deepks_io::save_tensor2npy<double>(PARAM.globalv.global_out_dir
+ "deepks_gradvx_square.npy",
gvx_square,
rank);
}

if (PARAM.inp.deepks_out_unittest)
{
Expand Down Expand Up @@ -306,6 +323,15 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
}
const std::string file_gvepsl = get_filename("gvepsl", PARAM.inp.deepks_out_labels, iter);
LCAO_deepks_io::save_tensor2npy<double>(file_gvepsl, gvepsl, rank);
if (PARAM.inp.deepks_grad)
{
torch::Tensor gvepsl_square;
DeePKS_domain::cal_gvepsl_square(gvepsl, gvepsl_square);
LCAO_deepks_io::save_tensor2npy<double>(PARAM.globalv.global_out_dir
+ "deepks_gvepsl_square.npy",
gvepsl_square,
rank);
}

if (PARAM.inp.deepks_out_unittest)
{
Expand Down Expand Up @@ -541,7 +567,11 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
std::ofstream ofs_hr(file_hrtot, std::ios::out);
ofs_hr << "Matrix Dimension of H(R): " << nbasis << std::endl;
ofs_hr << "Matrix number of H(R): " << hR_tot->size_R_loop() << std::endl;
hamilt::Output_HContainer<TR> out_hr(&hR_serial, ofs_hr, sparse_threshold, precision);
hamilt::Output_HContainer<TR> out_hr(&hR_serial,
ofs_hr,
sparse_threshold,
precision,
LCAO_deepks_io::Ry2Hartree);
out_hr.write(true); // write all the matrices, including empty ones
ofs_hr.close();
}
Expand All @@ -564,7 +594,11 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
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_serial, ofs_hr, sparse_threshold, precision);
hamilt::Output_HContainer<TR> out_hr(&h_deltaR_serial,
ofs_hr,
sparse_threshold,
precision,
LCAO_deepks_io::Ry2Hartree);
out_hr.write(true); // write all the matrices, including empty ones
ofs_hr.close();
}
Expand Down Expand Up @@ -629,7 +663,16 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
int R_size = DeePKS_domain::get_R_size(*h_deltaR);
torch::Tensor overlap_out;
torch::Tensor iRmat;
DeePKS_domain::prepare_phialpha_iRmat(nlocal, R_size, deepks_param, phialpha, ucell, orb, *ParaV, GridD, overlap_out, iRmat);
DeePKS_domain::prepare_phialpha_iRmat(nlocal,
R_size,
deepks_param,
phialpha,
ucell,
orb,
*ParaV,
GridD,
overlap_out,
iRmat);
const std::string file_overlap = PARAM.globalv.global_out_dir + "deepks_phialpha_r.npy";
LCAO_deepks_io::save_tensor2npy<double>(file_overlap, overlap_out, rank);
const std::string file_iRmat = PARAM.globalv.global_out_dir + "deepks_iRmat.npy";
Expand All @@ -641,7 +684,88 @@ void LCAO_Deepks_Interface<TK, TR>::out_deepks_labels(const double& etot,
}

//================================================================================
// 6. Hk
// 6. deepks_grad: descriptor-gradient label intermediates
//================================================================================
if (PARAM.inp.deepks_grad && is_after_scf)
{
{
torch::Tensor dot_phialpha_hamilt; // Energy part of unit is Hartree here
DeePKS_domain::cal_phialpha_hamilt_proj<TR>(nlocal,
nat,
deepks_param,
phialpha,
*p_ham->getHR(),
ucell,
orb,
*ParaV,
GridD,
dot_phialpha_hamilt);

// The model is trained to predict the absolute correction
// H_target - H_base. getHR() already contains the active
// DeePKS correction, so remove its projection before saving
// the current-functional (base) quantity. Target runs have
// deepks_scf=0 and need no subtraction.
if (PARAM.inp.deepks_scf)
{
torch::Tensor dot_phialpha_delta;
DeePKS_domain::cal_phialpha_hamilt_proj<TR>(nlocal,
nat,
deepks_param,
phialpha,
*p_ham->get_V_delta_R(),
ucell,
orb,
*ParaV,
GridD,
dot_phialpha_delta);
dot_phialpha_hamilt -= dot_phialpha_delta;
}
LCAO_deepks_io::save_tensor2npy<TR>(PARAM.globalv.global_out_dir + "deepks_dot_phialpha_hamilt.npy",
dot_phialpha_hamilt,
rank);

if (gevdm.empty())
{
DeePKS_domain::cal_gevdm(nat, deepks_param, pdm, gevdm);
}
// B^T B square matrix: (nat*des_per_atom, nat*des_per_atom)
// R_size: conservatively 2x the phialpha R extent to cover
// all dR2-dR1 differences produced by iterate_ad2.
const int btb_R_size = 2 * DeePKS_domain::get_R_size(*phialpha[0]);
torch::Tensor vdrpre_square;
DeePKS_domain::cal_vdrpre_square(nlocal,
nat,
btb_R_size,
deepks_param,
phialpha,
gevdm,
ucell,
orb,
*ParaV,
GridD,
vdrpre_square);
LCAO_deepks_io::save_tensor2npy<double>(PARAM.globalv.global_out_dir + "deepks_vdrpre_square.npy",
vdrpre_square,
rank);

// Avoid double-writing: the deepks_v_delta==-2 path already
// outputs deepks_gevdm.npy when deepks_out_labels==1.
const bool gevdm_already_output
= (PARAM.inp.deepks_v_delta == -2 && PARAM.inp.deepks_out_labels == 1);
if (!gevdm_already_output)
{
torch::Tensor gevdm_out;
DeePKS_domain::prepare_gevdm(nat, deepks_param, orb, gevdm, gevdm_out);
LCAO_deepks_io::save_tensor2npy<double>(PARAM.globalv.global_out_dir + "deepks_gevdm.npy",
gevdm_out,
rank);
}
}
}

//================================================================================
// 7. Hk
//================================================================================

if (PARAM.inp.deepks_v_delta > 0)
Expand Down
Loading
Loading