diff --git a/gpu4pyscf/lib/pbc/nr_eval_gto.cu b/gpu4pyscf/lib/pbc/nr_eval_gto.cu index 5eecda901..9c9c90cac 100644 --- a/gpu4pyscf/lib/pbc/nr_eval_gto.cu +++ b/gpu4pyscf/lib/pbc/nr_eval_gto.cu @@ -151,8 +151,8 @@ void _cart_deriv1_strain_tensor( } } -template __device__ -void _eval_cart_deriv1_strain_tensor( +template __device__ +void _eval_deriv1_strain_tensor( double *out, double *img_coords, double *env, double xi, double yi, double zi, double rrcutoff, int *bas, int nimgs, int nao, int ngrids) @@ -198,12 +198,203 @@ void _eval_cart_deriv1_strain_tensor( ao, gx, gy, gz, a2, rx, ry, rz, Rx, Ry, Rz, n0); } } - for (int n = 0; n < min(n_cart-n0, 2); ++n) { - for (int x = 0; x < 9; ++x) { - for (int s = 0; s < 4; ++s) { - out[(x*4+s)*naog+(n0+n)*ngrids] = ao[n*36+s*9+x]; + if constexpr (cart) { + for (int n = 0; n < min(n_cart-n0, 2); ++n) { + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+(n0+n)*ngrids] = ao[n*36+s*9+x]; + } } } + } else { + switch (ANG) { + case 0: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog] = ao[s*9+x]; + } + } + break; + case 1: + switch (n0) { + case 0: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+0*ngrids] = ao[0*36+s*9+x]; + out[(x*4+s)*naog+1*ngrids] = ao[1*36+s*9+x]; + } + } + break; + case 2: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+2*ngrids] = ao[0*36+s*9+x]; + } + } + break; + } + break; + case 2: + switch (n0) { + case 0: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+2*ngrids] = -0.315391565252520002 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+4*ngrids] = 0.546274215296039535 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+0*ngrids] = 1.092548430592079070 * ao[1*36+s*9+x]; + } + } + break; + case 2: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+3*ngrids] = 1.092548430592079070 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+2*ngrids]+= -0.315391565252520002 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+4*ngrids]+= -0.546274215296039535 * ao[1*36+s*9+x]; + } + } + break; + case 4: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+1*ngrids] = 1.092548430592079070 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+2*ngrids]+= 0.630783130505040012 * ao[1*36+s*9+x]; + } + } + break; + } + break; + case 3: + switch (n0) { + case 0: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+4*ngrids] = -0.457045799464465739 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+6*ngrids] = 0.590043589926643510 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+0*ngrids] = 1.770130769779930531 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+2*ngrids] = -0.457045799464465739 * ao[1*36+s*9+x]; + } + } + break; + case 2: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+3*ngrids] = -1.119528997770346170 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+5*ngrids] = 1.445305721320277020 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+4*ngrids]+= -0.457045799464465739 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+6*ngrids]+= -1.770130769779930530 * ao[1*36+s*9+x]; + } + } + break; + case 4: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+1*ngrids] = 2.890611442640554055 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+4*ngrids]+= 1.828183197857862944 * ao[1*36+s*9+x]; + } + } + break; + case 6: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+0*ngrids]+= -0.590043589926643510 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+2*ngrids]+= -0.457045799464465739 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+3*ngrids]+= -1.119528997770346170 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+5*ngrids]+= -1.445305721320277020 * ao[1*36+s*9+x]; + } + } + break; + case 8: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+2*ngrids]+= 1.828183197857862944 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+3*ngrids]+= 0.746352665180230782 * ao[1*36+s*9+x]; + } + } + break; + } + break; + case 4: + switch (n0) { + case 0: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+4*ngrids] = 0.317356640745612911 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+6*ngrids] = -0.473087347878780009 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+8*ngrids] = 0.625835735449176134 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+0*ngrids] = 2.503342941796704538 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+2*ngrids] = -0.946174695757560014 * ao[1*36+s*9+x]; + } + } + break; + case 2: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+5*ngrids] = -2.007139630671867500 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+7*ngrids] = 1.770130769779930531 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+4*ngrids]+= 0.634713281491225822 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+8*ngrids]+= -3.755014412695056800 * ao[1*36+s*9+x]; + } + } + break; + case 4: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+1*ngrids] = 5.310392309339791593 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+3*ngrids] = -2.007139630671867500 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+4*ngrids]+= -2.538853125964903290 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+6*ngrids]+= 2.838524087272680054 * ao[1*36+s*9+x]; + } + } + break; + case 6: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+0*ngrids]+= -2.503342941796704538 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+2*ngrids]+= -0.946174695757560014 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+5*ngrids]+= -2.007139630671867500 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+7*ngrids]+= -5.310392309339791590 * ao[1*36+s*9+x]; + } + } + break; + case 8: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+2*ngrids]+= 5.677048174545360108 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+5*ngrids]+= 2.676186174229156671 * ao[1*36+s*9+x]; + } + } + break; + case 10: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+4*ngrids]+= 0.317356640745612911 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+6*ngrids]+= 0.473087347878780009 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+8*ngrids]+= 0.625835735449176134 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+1*ngrids]+= -1.770130769779930530 * ao[1*36+s*9+x]; + out[(x*4+s)*naog+3*ngrids]+= -2.007139630671867500 * ao[1*36+s*9+x]; + } + } + break; + case 12: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+4*ngrids]+= -2.538853125964903290 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+6*ngrids]+= -2.838524087272680054 * ao[0*36+s*9+x]; + out[(x*4+s)*naog+3*ngrids]+= 2.676186174229156671 * ao[1*36+s*9+x]; + } + } + break; + case 14: + for (int x = 0; x < 9; ++x) { + for (int s = 0; s < 4; ++s) { + out[(x*4+s)*naog+4*ngrids]+= 0.846284375321634430 * ao[0*36+s*9+x]; + } + } + break; + } + break; + } } } } @@ -1158,8 +1349,9 @@ static void _sph_ip2_kernel(double *out, PBCIntEnvVars envs, double *grids, } } +template __global__ -static void _cart_deriv0_strain_tensor_kernel( +static void _deriv0_strain_tensor_kernel( double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { @@ -1375,83 +1567,186 @@ static void _cart_deriv0_strain_tensor_kernel( } } } } - if (li < 4) { - int *ao_loc = envs.ao_loc; + const int *ao_loc = envs.ao_loc; + out += (size_t)ao_loc[bas_id] * ngrids + grid_id; + const size_t naog = nao * ngrids; + if constexpr (cart) { int nf = (li + 1) * (li + 2) / 2; - out += (size_t)ao_loc[bas_id] * ngrids + grid_id; - size_t naog = nao * ngrids; - for (int n = 0; n < 10; ++n) { - if (n >= nf) break; - for (int x = 0; x < 9; ++x) { - out[x*naog+n*ngrids] = ao[n*9+x]; + if (li < 4) { + for (int n = 0; n < 10; ++n) { + if (n >= nf) break; + for (int x = 0; x < 9; ++x) { + out[x*naog+n*ngrids] = ao[n*9+x]; + } + } + } else { + for (int n = 0; n < 15; ++n) { + if (n >= nf) break; + for (int x = 0; x < 6; ++x) { + out[x*naog+n*ngrids] = ao[n*6+x]; + } + } + out += 6 * naog; // To process zx, zy, zz + + for (int n = 0; n < 45; ++n) { + ao[n] = 0; + } + for (int img = 0; img < nimgs; ++img) { + double ce = 0; + double ce_2a = 0; + double Rx = img_coords[img*3+0] + cell0_Rx; + double Ry = img_coords[img*3+1] + cell0_Ry; + double Rz = img_coords[img*3+2] + cell0_Rz; + double rx = xi - Rx; + double ry = yi - Ry; + double rz = zi - Rz; + double rr = rx * rx + ry * ry + rz * rz; + if (rr > rrcutoff) continue; + for (int ip = 0; ip < nprim; ++ip) { + double ai = expi[ip]; + double c_exp = ci[ip] * exp(-ai * rr); + ce += c_exp; + ce_2a -= c_exp * ai * 2; + } + if (fabs(ce) < 1e-18) continue; + double az = ce_2a * rz; + gto[0 ] = az * rx * rx * rx * rx; + gto[1 ] = az * rx * rx * rx * ry; + gto[2 ] = (az * rz + ce) * rx * rx * rx; + gto[3 ] = az * rx * rx * ry * ry; + gto[4 ] = (az * rz + ce) * rx * rx * ry; + gto[5 ] = (az * rz + 2 * ce) * rx * rx * rz; + gto[6 ] = az * rx * ry * ry * ry; + gto[7 ] = (az * rz + ce) * rx * ry * ry; + gto[8 ] = (az * rz + 2 * ce) * rx * ry * rz; + gto[9 ] = (az * rz + 3 * ce) * rx * rz * rz; + gto[10] = az * ry * ry * ry * ry; + gto[11] = (az * rz + ce) * ry * ry * ry; + gto[12] = (az * rz + 2 * ce) * ry * ry * rz; + gto[13] = (az * rz + 3 * ce) * ry * rz * rz; + gto[14] = (az * rz + 4 * ce) * rz * rz * rz; + for (int n = 0; n < 15; n++) { + ao[0+3*n] -= gto[n] * Rx; + ao[1+3*n] -= gto[n] * Ry; + ao[2+3*n] -= gto[n] * Rz; + } + } + for (int n = 0; n < 15; ++n) { + for (int x = 0; x < 3; ++x) { + out[x*naog+n*ngrids] = ao[n*3+x]; + } } } } else { - int *ao_loc = envs.ao_loc; - int nf = (li + 1) * (li + 2) / 2; - out += (size_t)ao_loc[bas_id] * ngrids + grid_id; - size_t naog = nao * ngrids; - for (int n = 0; n < 15; ++n) { - if (n >= nf) break; - for (int x = 0; x < 6; ++x) { - out[x*naog+n*ngrids] = ao[n*6+x]; + switch (li) { + case 0: + for (int n = 0; n < 9; ++n) { + out[n * naog] = ao[n]; } - } - out += 6 * naog; // To process zx, zy, zz + break; + case 1: + for (int n = 0; n < 9; ++n) { + out[n*naog+0*ngrids] = ao[0*9+n]; + out[n*naog+1*ngrids] = ao[1*9+n]; + out[n*naog+2*ngrids] = ao[2*9+n]; + } + break; + case 2: + for (int n = 0; n < 9; ++n) { + out[n*naog+0*ngrids] = 1.092548430592079070 * ao[1*9+n]; + out[n*naog+1*ngrids] = 1.092548430592079070 * ao[4*9+n]; + out[n*naog+2*ngrids] = 0.630783130505040012 * ao[5*9+n] - 0.315391565252520002 * (ao[0*9+n] + ao[3*9+n]); + out[n*naog+3*ngrids] = 1.092548430592079070 * ao[2*9+n]; + out[n*naog+4*ngrids] = 0.546274215296039535 * (ao[0*9+n] - ao[3*9+n]); + } + break; + case 3: + for (int n = 0; n < 9; ++n) { + out[n*naog+0*ngrids] = 1.770130769779930531 * ao[1*9+n] - 0.590043589926643510 * ao[6*9+n]; + out[n*naog+1*ngrids] = 2.890611442640554055 * ao[4*9+n]; + out[n*naog+2*ngrids] = 1.828183197857862944 * ao[8*9+n] - 0.457045799464465739 * (ao[1*9+n] + ao[6*9+n]); + out[n*naog+3*ngrids] = 0.746352665180230782 * ao[9*9+n] - 1.119528997770346170 * (ao[2*9+n] + ao[7*9+n]); + out[n*naog+4*ngrids] = 1.828183197857862944 * ao[5*9+n] - 0.457045799464465739 * (ao[0*9+n] + ao[3*9+n]); + out[n*naog+5*ngrids] = 1.445305721320277020 * (ao[2*9+n] - ao[7*9+n]); + out[n*naog+6*ngrids] = 0.590043589926643510 * ao[0*9+n] - 1.770130769779930530 * ao[3*9+n]; + } + break; + case 4: + for (int n = 0; n < 6; ++n) { + out[n*naog+0*ngrids] = 2.503342941796704538 * (ao[1*6+n] - ao[6*6+n]) ; + out[n*naog+1*ngrids] = 5.310392309339791593 * ao[4*6+n] - 1.770130769779930530 * ao[11*6+n]; + out[n*naog+2*ngrids] = 5.677048174545360108 * ao[8*6+n] - 0.946174695757560014 * (ao[1*6+n] + ao[6*6+n]); + out[n*naog+3*ngrids] = 2.676186174229156671 * ao[13*6+n]- 2.007139630671867500 * (ao[4*6+n] + ao[11*6+n]); + out[n*naog+4*ngrids] = 0.317356640745612911 * (ao[0*6+n] + ao[10*6+n]) + 0.634713281491225822 * ao[3*6+n] - 2.538853125964903290 * (ao[5*6+n] + ao[12*6+n]) + 0.846284375321634430 * ao[14*6+n]; + out[n*naog+5*ngrids] = 2.676186174229156671 * ao[9*6+n] - 2.007139630671867500 * (ao[2*6+n] + ao[7*6+n]); + out[n*naog+6*ngrids] = 2.838524087272680054 * (ao[5*6+n] - ao[12*6+n]) + 0.473087347878780009 * (ao[10*6+n]- ao[0*6+n]); + out[n*naog+7*ngrids] = 1.770130769779930531 * ao[2*6+n] - 5.310392309339791590 * ao[7*6+n]; + out[n*naog+8*ngrids] = 0.625835735449176134 * (ao[0*6+n] + ao[10*6+n]) - 3.755014412695056800 * ao[3*6+n]; + } + out += 6 * naog; // To process zx, zy, zz - for (int n = 0; n < 45; ++n) { - ao[n] = 0; - } - for (int img = 0; img < nimgs; ++img) { - double ce = 0; - double ce_2a = 0; - double Rx = img_coords[img*3+0] + cell0_Rx; - double Ry = img_coords[img*3+1] + cell0_Ry; - double Rz = img_coords[img*3+2] + cell0_Rz; - double rx = xi - Rx; - double ry = yi - Ry; - double rz = zi - Rz; - double rr = rx * rx + ry * ry + rz * rz; - if (rr > rrcutoff) continue; - for (int ip = 0; ip < nprim; ++ip) { - double ai = expi[ip]; - double c_exp = ci[ip] * exp(-ai * rr); - ce += c_exp; - ce_2a -= c_exp * ai * 2; + for (int n = 0; n < 45; ++n) { + ao[n] = 0; } - if (fabs(ce) < 1e-18) continue; - double az = ce_2a * rz; - gto[0 ] = az * rx * rx * rx * rx; - gto[1 ] = az * rx * rx * rx * ry; - gto[2 ] = (az * rz + ce) * rx * rx * rx; - gto[3 ] = az * rx * rx * ry * ry; - gto[4 ] = (az * rz + ce) * rx * rx * ry; - gto[5 ] = (az * rz + 2 * ce) * rx * rx * rz; - gto[6 ] = az * rx * ry * ry * ry; - gto[7 ] = (az * rz + ce) * rx * ry * ry; - gto[8 ] = (az * rz + 2 * ce) * rx * ry * rz; - gto[9 ] = (az * rz + 3 * ce) * rx * rz * rz; - gto[10] = az * ry * ry * ry * ry; - gto[11] = (az * rz + ce) * ry * ry * ry; - gto[12] = (az * rz + 2 * ce) * ry * ry * rz; - gto[13] = (az * rz + 3 * ce) * ry * rz * rz; - gto[14] = (az * rz + 4 * ce) * rz * rz * rz; - for (int n = 0; n < 15; n++) { - ao[0+3*n] -= gto[n] * Rx; - ao[1+3*n] -= gto[n] * Ry; - ao[2+3*n] -= gto[n] * Rz; + for (int img = 0; img < nimgs; ++img) { + double ce = 0; + double ce_2a = 0; + double Rx = img_coords[img*3+0] + cell0_Rx; + double Ry = img_coords[img*3+1] + cell0_Ry; + double Rz = img_coords[img*3+2] + cell0_Rz; + double rx = xi - Rx; + double ry = yi - Ry; + double rz = zi - Rz; + double rr = rx * rx + ry * ry + rz * rz; + if (rr > rrcutoff) continue; + for (int ip = 0; ip < nprim; ++ip) { + double ai = expi[ip]; + double c_exp = ci[ip] * exp(-ai * rr); + ce += c_exp; + ce_2a -= c_exp * ai * 2; + } + if (fabs(ce) < 1e-18) continue; + double az = ce_2a * rz; + gto[0 ] = az * rx * rx * rx * rx; + gto[1 ] = az * rx * rx * rx * ry; + gto[2 ] = (az * rz + ce) * rx * rx * rx; + gto[3 ] = az * rx * rx * ry * ry; + gto[4 ] = (az * rz + ce) * rx * rx * ry; + gto[5 ] = (az * rz + 2 * ce) * rx * rx * rz; + gto[6 ] = az * rx * ry * ry * ry; + gto[7 ] = (az * rz + ce) * rx * ry * ry; + gto[8 ] = (az * rz + 2 * ce) * rx * ry * rz; + gto[9 ] = (az * rz + 3 * ce) * rx * rz * rz; + gto[10] = az * ry * ry * ry * ry; + gto[11] = (az * rz + ce) * ry * ry * ry; + gto[12] = (az * rz + 2 * ce) * ry * ry * rz; + gto[13] = (az * rz + 3 * ce) * ry * rz * rz; + gto[14] = (az * rz + 4 * ce) * rz * rz * rz; + for (int n = 0; n < 15; n++) { + ao[0+3*n] -= gto[n] * Rx; + ao[1+3*n] -= gto[n] * Ry; + ao[2+3*n] -= gto[n] * Rz; + } } - } - for (int n = 0; n < 15; ++n) { - for (int x = 0; x < 3; ++x) { - out[x*naog+n*ngrids] = ao[n*3+x]; + for (int n = 0; n < 3; ++n) { + out[n*naog+0*ngrids] = 2.503342941796704538 * (ao[1*3+n] - ao[6*3+n]) ; + out[n*naog+1*ngrids] = 5.310392309339791593 * ao[4*3+n] - 1.770130769779930530 * ao[11*3+n]; + out[n*naog+2*ngrids] = 5.677048174545360108 * ao[8*3+n] - 0.946174695757560014 * (ao[1*3+n] + ao[6*3+n]); + out[n*naog+3*ngrids] = 2.676186174229156671 * ao[13*3+n]- 2.007139630671867500 * (ao[4*3+n] + ao[11*3+n]); + out[n*naog+4*ngrids] = 0.317356640745612911 * (ao[0*3+n] + ao[10*3+n]) + 0.634713281491225822 * ao[3*3+n] - 2.538853125964903290 * (ao[5*3+n] + ao[12*3+n]) + 0.846284375321634430 * ao[14*3+n]; + out[n*naog+5*ngrids] = 2.676186174229156671 * ao[9*3+n] - 2.007139630671867500 * (ao[2*3+n] + ao[7*3+n]); + out[n*naog+6*ngrids] = 2.838524087272680054 * (ao[5*3+n] - ao[12*3+n]) + 0.473087347878780009 * (ao[10*3+n]- ao[0*3+n]); + out[n*naog+7*ngrids] = 1.770130769779930531 * ao[2*3+n] - 5.310392309339791590 * ao[7*3+n]; + out[n*naog+8*ngrids] = 0.625835735449176134 * (ao[0*3+n] + ao[10*3+n]) - 3.755014412695056800 * ao[3*3+n]; } + break; } } } +template __global__ -static void _cart_deriv1_strain_tensor_kernel( +static void _deriv1_strain_tensor_kernel( double *out, PBCIntEnvVars envs, double *grids, size_t ngrids, int nao, double *rcut) { @@ -1477,19 +1772,19 @@ static void _cart_deriv1_strain_tensor_kernel( int nimgs = envs.nimgs; switch (li) { - case 0: _eval_cart_deriv1_strain_tensor<0>(out, img_coords, env, + case 0: _eval_deriv1_strain_tensor<0, cart>(out, img_coords, env, xi, yi, zi, rrcutoff, bas, nimgs, nao, ngrids); break; - case 1: _eval_cart_deriv1_strain_tensor<1>(out, img_coords, env, + case 1: _eval_deriv1_strain_tensor<1, cart>(out, img_coords, env, xi, yi, zi, rrcutoff, bas, nimgs, nao, ngrids); break; - case 2: _eval_cart_deriv1_strain_tensor<2>(out, img_coords, env, + case 2: _eval_deriv1_strain_tensor<2, cart>(out, img_coords, env, xi, yi, zi, rrcutoff, bas, nimgs, nao, ngrids); break; - case 3: _eval_cart_deriv1_strain_tensor<3>(out, img_coords, env, + case 3: _eval_deriv1_strain_tensor<3, cart>(out, img_coords, env, xi, yi, zi, rrcutoff, bas, nimgs, nao, ngrids); break; - case 4: _eval_cart_deriv1_strain_tensor<4>(out, img_coords, env, + case 4: _eval_deriv1_strain_tensor<4, cart>(out, img_coords, env, xi, yi, zi, rrcutoff, bas, nimgs, nao, ngrids); break; } @@ -1543,19 +1838,23 @@ int PBCeval_gto_strain_tensor(double *out, PBCIntEnvVars *envs, double *grids, int ngrids, int nao, int nbas, int deriv, int cart, double *rcut) { - if (!cart) { - fprintf(stderr, "PBCeval_gto_strain_tensor does not support spherical GTOs\n"); - return 1; - } constexpr int ngrids_per_block = THREADS; int threads = ngrids_per_block; dim3 blocks((ngrids+ngrids_per_block-1)/ngrids_per_block, nbas); switch (deriv) { case 0: - _cart_deriv0_strain_tensor_kernel<<>>(out, *envs, grids, ngrids, nao, rcut); + if (cart) { + _deriv0_strain_tensor_kernel <<>>(out, *envs, grids, ngrids, nao, rcut); + } else { + _deriv0_strain_tensor_kernel <<>>(out, *envs, grids, ngrids, nao, rcut); + } break; case 1: - _cart_deriv1_strain_tensor_kernel<<>>(out, *envs, grids, ngrids, nao, rcut); + if (cart) { + _deriv1_strain_tensor_kernel <<>>(out, *envs, grids, ngrids, nao, rcut); + } else { + _deriv1_strain_tensor_kernel <<>>(out, *envs, grids, ngrids, nao, rcut); + } break; default: fprintf(stderr, "PBCeval_gto_strain_tensor deriv = %d not supported\n", deriv); diff --git a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py index 8c3c2230f..b4a05f18f 100644 --- a/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py +++ b/gpu4pyscf/pbc/df/tests/test_pbc_hcore_derivatives.py @@ -221,7 +221,7 @@ def get_energy(cell): assert np.abs(test_energy - ref_energy) < 1e-9 assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 3e-8 - # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 3e-7 # TODO: Support Becke grid stress tensor + assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-7 dm = mf.make_rdm1() kmesh = np.array([1,1,1]) @@ -336,7 +336,7 @@ def get_energy(cell): assert np.abs(test_energy - ref_energy) < 1e-9 assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 5e-7 - # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 5e-6 # TODO: Support Becke grid stress tensor + assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-6 dm = mf.make_rdm1() kpts = cell.make_kpts(kmesh) @@ -458,7 +458,7 @@ def get_energy(cell): assert np.abs(test_energy - ref_energy) < 1e-9 assert np.max(np.abs(test_derivatives[:-3, :] - ref_derivatives[:-3, :])) < 1e-6 - # assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 1e-6 # TODO: Support Becke grid stress tensor + assert np.max(np.abs(test_derivatives[-3:, :] - ref_derivatives[-3:, :])) < 2e-6 if __name__ == '__main__': print("Full Tests for PBC GDF Hcore gradient and stress tensor") diff --git a/gpu4pyscf/pbc/dft/gen_grid.py b/gpu4pyscf/pbc/dft/gen_grid.py index 31e990669..c42a64b90 100644 --- a/gpu4pyscf/pbc/dft/gen_grid.py +++ b/gpu4pyscf/pbc/dft/gen_grid.py @@ -246,7 +246,10 @@ def get_becke_weight_derivative(grids, natm, grid_range = None): dweight_dA_unitcell = cp.zeros([natm, 3, ngrids]) cp.add.at(dweight_dA_unitcell, grids_supatm_to_atm_idx, dweight_dA_supercell) - return dweight_dA_unitcell + weight_stress = cp.einsum("Axg,Ay->xyg", dweight_dA_supercell, grids_supatm_coords) + + dweight_gradient_stress = cp.vstack((dweight_dA_unitcell, weight_stress)) + return dweight_gradient_stress class UniformGrids(lib.StreamObject): '''Uniform Grid class.''' diff --git a/gpu4pyscf/pbc/dft/tests/test_pbc_grids.py b/gpu4pyscf/pbc/dft/tests/test_pbc_grids.py index 3824de845..00c783135 100644 --- a/gpu4pyscf/pbc/dft/tests/test_pbc_grids.py +++ b/gpu4pyscf/pbc/dft/tests/test_pbc_grids.py @@ -19,7 +19,7 @@ from pyscf.pbc.dft import gen_grid as gen_grid_cpu from gpu4pyscf.pbc.dft import gen_grid from pyscf.pbc.dft import rks as rks_cpu -from gpu4pyscf.pbc.dft import rks +from gpu4pyscf.pbc.dft import rks, uks from pyscf.pbc.dft import krks as krks_cpu from gpu4pyscf.pbc.dft import krks, kuks from gpu4pyscf.pbc.dft.gen_grid import get_becke_weight_derivative @@ -27,9 +27,12 @@ from gpu4pyscf.pbc.grad.kuks import get_vxc_full_response as unrestricted_get_vxc_full_response from gpu4pyscf.pbc.grad.kuks import get_vxc as unrestricted_get_vxc from gpu4pyscf.dft.tests.test_grids import find_matching_index_between_two_grids +from gpu4pyscf.pbc.grad.rhf import _finite_diff_cells -def numerical_gradient_exc_becke(cell, xc, kpts, auxbasis, atom_grid, dm, unrestricted=False): +def numerical_gradient_exc_becke(cell, xc, kmesh, auxbasis, atom_grid, dm, unrestricted=False, dx = 1e-4): + assert np.array(kmesh).shape == (3,) def get_energy(cell): + kpts = cell.make_kpts(kmesh) if unrestricted: mf = kuks.KUKS(cell, xc=xc, kpts=kpts) else: @@ -46,8 +49,7 @@ def get_energy(cell): n, exc, vxc = mf._numint.nr_rks(cell, mf.grids, mf.xc, dm, 0, hermi=1, kpts=kpts, kpts_band=None) return exc - numerical_gradient = np.zeros((cell.natm, 3)) - dx = 1e-4 + numerical_gradient = np.zeros((cell.natm + 3, 3)) cell_copy = cell.copy() for i_atom in range(cell.natm): for i_xyz in range(3): @@ -68,6 +70,14 @@ def get_energy(cell): translation_invariance = np.sum(numerical_gradient, axis=0) assert np.max(np.abs(translation_invariance)) < 1e-8, "Bad numerical gradient" + for i_xyz in range(3): + for j_xyz in range(3): + cell_p, cell_m = _finite_diff_cells(cell, i_xyz, j_xyz, disp = dx) + E_p = get_energy(cell_p) + E_m = get_energy(cell_m) + + numerical_gradient[cell.natm + i_xyz, j_xyz] = (E_p - E_m) / (2 * dx) + # np.set_printoptions(precision=16, suppress=True, linewidth=np.inf) # print(repr(numerical_gradient)) return numerical_gradient @@ -191,7 +201,7 @@ def test_becke_weight_derivative(self): test_dw_truncated = get_becke_weight_derivative(grids, cell.natm, truncation_range) dx = 1e-5 - reference_dw = cp.empty([cell.natm, 3, grids.coords.shape[0]]) + reference_dw = cp.empty([cell.natm + 3, 3, grids.coords.shape[0]]) cell_copy = cell.copy() for i_atom in range(cell.natm): for i_xyz in range(3): @@ -213,33 +223,60 @@ def test_becke_weight_derivative(self): reference_dw[i_atom, i_xyz, :] = (w_p - w_m) / (2 * dx) + for i_xyz in range(3): + for j_xyz in range(3): + cell_p, cell_m = _finite_diff_cells(cell, i_xyz, j_xyz, disp = dx) + grids.reset(cell_p) + grids.build() + w_p = grids.weights.copy() + + grids.reset(cell_m) + grids.build() + w_m = grids.weights.copy() + + reference_dw[cell.natm + i_xyz, j_xyz] = (w_p - w_m) / (2 * dx) + reference_dw_truncated = reference_dw[:, :, truncation_range[0] : truncation_range[1]] - assert cp.max(cp.abs(test_dw - reference_dw)) < 2e-9 - assert cp.max(cp.abs(test_dw_truncated - reference_dw_truncated)) < 2e-9 + assert cp.max(cp.abs(test_dw[:-3] - reference_dw[:-3])) < 2e-9 + assert cp.max(cp.abs(test_dw_truncated[:-3] - reference_dw_truncated[:-3])) < 2e-9 + + assert cp.max(cp.abs(test_dw[-3:] - reference_dw[-3:])) < 5e-9 + assert cp.max(cp.abs(test_dw_truncated[-3:] - reference_dw_truncated[-3:])) < 5e-9 def test_xc_gradient_lda_with_response(self): cell = pyscf.M( - a = np.eye(3) * 3.5668, - atom = ''' - C 0. 0. 0. - C 0.8917 0.8917 0.8917 - C 1.7834 1.7834 0. - C 2.6751 2.6751 0.8917 - C 1.7834 0. 1.7834 - C 2.6751 0.8917 2.6751 - C 0. 1.7834 1.7834 - C 0.8917 2.6751 2.6751 - ''', - basis = 'def2-svp', + a = np.array([ + [15.9069652593, 0, 0], + [0, 15.9069652593, 0], + [0, 6, 15.9069652593], + ]), + atom = """ + O 15.43509000 9.59549000 8.94968000 + H 15.05724000 9.21878000 9.73314000 + H 0.51550474 9.33856000 9.01857000 + He 2.51550474 9.33856000 9.01857000 + """, + basis = {'default': 'def2-svp', 'He': """ + #BASIS SET: (2s) -> [1s] + He S + 0.2432879285E+01 0.4301284983E+00 + 0.4330512863E+00 0.6789135305E+00 + He F + 1.0 1.0 + He G + 1.2 1.0 + END + """}, verbose = 0, ) kpts = np.array([[0,0,0]]) - mf = krks.KRKS(cell, xc="LDA0", kpts=kpts).density_fit(auxbasis='def2-universal-jkfit') + mf = rks.RKS(cell, xc="LDA0").density_fit(auxbasis='def2-universal-jkfit') mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (50,194) mf.conv_tol = 1e-10 + mf.with_df.linear_dep_threshold = 1e-10 # The default in pyscf==2.8.0 is 1e-9, in pyscf==2.14.0 is 1e-10 mf.kernel() @@ -248,21 +285,9 @@ def test_xc_gradient_lda_with_response(self): dm = dm[None,:,:] test_gradient = get_vxc_full_response(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) - # ref_gradient = numerical_gradient_exc_becke(cell, "LDA0", kpts, 'def2-universal-jkfit', (50,194), dm) - ref_gradient = np.array([ - [ 0.0000000094857455, 0.0000000095390362, 0.0000000095212727], - [-0.0002930449305438, -0.0002930448950167, -0.0002930448950167], - [-0.0000000263611355, -0.0000000263611355, 0.0000000985167503], - [ 0.0002930621789687, 0.0002930622144959, -0.0002931344589285], - [-0.0000000261479727, 0.000000098570041 , -0.0000000263966626], - [ 0.0002930622322594, -0.0002931345122192, 0.0002930622144959], - [ 0.0000000986766224, -0.0000000262012634, -0.0000000262012634], - [-0.0002931345122192, 0.0002930622144959, 0.0002930622144959], - ]) + ref_gradient = numerical_gradient_exc_becke(cell, "LDA0", [1,1,1], 'def2-universal-jkfit', (50,194), dm) - # It can match down to 1e-9, if the finite difference is computed using the same dm from SCF. - # However if we save the finite difference result, it suffers from the numerical instability of dm, and the a 3e-7 error is observed. - assert np.max(np.abs(test_gradient - ref_gradient)) < 5e-7 + assert np.max(np.abs(test_gradient - ref_gradient)) < 2e-8 def test_xc_gradient_gga_with_response(self): cell = pyscf.M( @@ -275,7 +300,8 @@ def test_xc_gradient_gga_with_response(self): verbose = 0, ) - kpts = cell.make_kpts((1,2,3)) + kmesh = (1,2,3) + kpts = cell.make_kpts(kmesh) mf = krks.KRKS(cell, xc="HSE06", kpts=kpts).density_fit(auxbasis='def2-universal-jkfit') mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (50,194) @@ -289,13 +315,13 @@ def test_xc_gradient_gga_with_response(self): test_gradient = get_vxc_full_response(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) # dm is not very stable, and numerical gradient is super fast - ref_gradient = numerical_gradient_exc_becke(cell, "HSE06", kpts, 'def2-universal-jkfit', (50,194), dm) + ref_gradient = numerical_gradient_exc_becke(cell, "HSE06", kmesh, 'def2-universal-jkfit', (50,194), dm) assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-7 def test_xc_gradient_gga_without_response(self): cell = pyscf.M( - a = np.eye(3) * 3.5668, + a = np.eye(3) * 3.6668, atom = ''' C 0. 0. 0. C 0.8917 0.9017 0.8917 @@ -323,19 +349,22 @@ def test_xc_gradient_gga_without_response(self): dm = dm[None,:,:] test_gradient = get_vxc(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) - # ref_gradient = numerical_gradient_exc_becke(cell, "PBE", kpts, 'def2-universal-jkfit', (99,590), dm) + # ref_gradient = numerical_gradient_exc_becke(cell, "PBE", [1,1,1], 'def2-universal-jkfit', (99,590), dm) ref_gradient = np.array([ - [ 0.000429164934701 , -0.0022409300370896, 0.0004291650412824], - [-0.0000602672400873, -0.0009193305317012, -0.0000602672400873], - [ 0.0004409704956743, -0.0019386623151263, -0.0004303180389797], - [ 0.0000603361982598, 0.0037129129282221, -0.0000604878991339], - [-0.0004055164026795, -0.0022406373290096, -0.0004055164382066], - [ 0.0000604298833196, 0.0024061113990115, 0.0000604298477924], - [-0.0004303180745069, -0.0019386622795992, 0.0004409704601471], - [-0.0000604880057153, 0.0037129129637492, 0.0000603362693141], + [ 0.0057493621596905, 0.0040870682127547, 0.0057493626570704], + [-0.0347991650784252, -0.0320027493927455, -0.0347991649363166], + [ 0.0246147672555708, 0.0208271039525698, 0.0145351780389547], + [-0.0077134854237215, -0.0061064405798561, -0.014724848114156 ], + [ 0.0216872688696412, 0.0122791700718494, 0.0216872685854241], + [-0.0093490771035931, -0.0138048160280846, -0.0093490768904303], + [ 0.0145351779679004, 0.0208271048052211, 0.0246147670779351], + [-0.014724848647063 , -0.0061064410417089, -0.0077134864895356], + [ 2.679998687966645 , 0.204392930385211 , 0.179311419188366 ], + [ 0.204415101734412 , 2.6794902670701504, 0.204415101734412 ], + [ 0.1793114174120092, 0.2043929328365834, 2.6799986901693273], ]) - assert np.max(np.abs(test_gradient - ref_gradient)) < 3e-4 + assert np.max(np.abs(test_gradient - ref_gradient)) < 6e-4 def test_xc_gradient_mgga_with_response(self): cell = pyscf.M( @@ -347,7 +376,8 @@ def test_xc_gradient_mgga_with_response(self): verbose = 0, ) - kpts = cell.make_kpts((1,1,3)) + kmesh = (1,1,3) + kpts = cell.make_kpts(kmesh) mf = krks.KRKS(cell, xc="r2scan", kpts=kpts).density_fit(auxbasis='def2-universal-jkfit') mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (50,194) @@ -361,9 +391,10 @@ def test_xc_gradient_mgga_with_response(self): test_gradient = get_vxc_full_response(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) # dm is not very stable, and numerical gradient is super fast - ref_gradient = numerical_gradient_exc_becke(cell, "r2scan", kpts, 'def2-universal-jkfit', (50,194), dm) + ref_gradient = numerical_gradient_exc_becke(cell, "r2scan", kmesh, 'def2-universal-jkfit', (50,194), dm) - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-9 + assert np.max(np.abs(test_gradient[:-3] - ref_gradient[:-3])) < 1e-9 + assert np.max(np.abs(test_gradient[-3:] - ref_gradient[-3:])) < 3e-8 def test_xc_gradient_mgga_without_response(self): cell = pyscf.M( @@ -376,7 +407,8 @@ def test_xc_gradient_mgga_without_response(self): verbose = 0, ) - kpts = cell.make_kpts((1,2,3)) + kmesh = (1,2,3) + kpts = cell.make_kpts(kmesh) mf = krks.KRKS(cell, xc="r2scan0", kpts=kpts).density_fit(auxbasis='def2-universal-jkfit') mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (120,590) @@ -390,9 +422,10 @@ def test_xc_gradient_mgga_without_response(self): test_gradient = get_vxc(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) # dm is not very stable, and numerical gradient is super fast - ref_gradient = numerical_gradient_exc_becke(cell, "r2scan0", kpts, 'def2-universal-jkfit', (120,590), dm) + ref_gradient = numerical_gradient_exc_becke(cell, "r2scan0", kmesh, 'def2-universal-jkfit', (120,590), dm) - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-4 + assert np.max(np.abs(test_gradient[:-3] - ref_gradient[:-3])) < 1e-4 + assert np.max(np.abs(test_gradient[-3:] - ref_gradient[-3:])) < 2e-3 def test_xc_gradient_unrestricted_no_k_without_response(self): cell = pyscf.M( @@ -405,7 +438,7 @@ def test_xc_gradient_unrestricted_no_k_without_response(self): ) kpts = cell.make_kpts((1,1,1)) - mf = kuks.KUKS(cell, xc="HSE06", kpts=kpts).density_fit(auxbasis='def2-universal-jkfit') + mf = uks.UKS(cell, xc="HSE06").density_fit(auxbasis='def2-universal-jkfit') mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (99,590) mf.conv_tol = 1e-10 @@ -417,13 +450,17 @@ def test_xc_gradient_unrestricted_no_k_without_response(self): dm = dm[:,None,:,:] test_gradient = unrestricted_get_vxc(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) - # ref_gradient = numerical_gradient_exc_becke(cell, "HSE06", kpts, 'def2-universal-jkfit', (99,590), dm, unrestricted=True) + # ref_gradient = numerical_gradient_exc_becke(cell, "HSE06", (1,1,1), 'def2-universal-jkfit', (99,590), dm, unrestricted=True, dx=1e-5) ref_gradient = np.array([ - [ 0.0000210273753964, -0.0175452356021566, 0.0000210258033206], - [-0.0000210273665147, 0.0175452356021566, -0.0000210257944389], + [ 0.0000210217621088, -0.0175452376183216, 0.000021022206198 ], + [-0.0000210217621088, 0.0175452377071394, -0.000021022206198 ], + [-0.5849730015938803, -0.0000565129276708, -0.0047818454085302], + [-0.0000557453638805, -0.5846170649803639, -0.0000557508705867], + [-0.0047818451420767, -0.0000565183455592, -0.5849730222884375], ]) - assert np.max(np.abs(test_gradient - ref_gradient)) < 2e-4 + assert np.max(np.abs(test_gradient[:-3] - ref_gradient[:-3])) < 2e-4 + assert np.max(np.abs(test_gradient[-3:] - ref_gradient[-3:])) < 5e-4 def test_xc_gradient_unrestricted_no_k_with_response(self): cell = pyscf.M( @@ -448,6 +485,7 @@ def test_xc_gradient_unrestricted_no_k_with_response(self): mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (50,194) mf.conv_tol = 1e-10 + mf.with_df.linear_dep_threshold = 1e-10 # The default in pyscf==2.8.0 is 1e-9, in pyscf==2.14.0 is 1e-10 mf.kernel() @@ -456,16 +494,19 @@ def test_xc_gradient_unrestricted_no_k_with_response(self): dm = dm[:,None,:,:] test_gradient = unrestricted_get_vxc_full_response(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) - # ref_gradient = numerical_gradient_exc_becke(cell, "r2scan", kpts, 'def2-universal-jkfit', (50,194), dm, unrestricted=True) + # ref_gradient = numerical_gradient_exc_becke(cell, "r2scan", (1,1,1), 'def2-universal-jkfit', (50,194), dm, unrestricted=True) ref_gradient = np.array([ - [ 0.0000003321254383, 0.0000003318678665, 0.0000003322675468], - [-0.0003303815532263, -0.0003303816153988, -0.0003303810025557], - [-0.0000001970335006, -0.0000001969180374, 0.0000003920330727], - [ 0.0003302472517674, 0.0003302459639087, -0.0003304415585603], - [-0.000000197513117 , 0.0000003920863634, -0.0000001963496032], - [ 0.0003302466300426, -0.0003304415407968, 0.0003302464079979], - [ 0.0000003911893032, -0.0000001970335006, -0.0000001977085162], - [-0.0003304409634808, 0.0003302473672306, 0.0003302460438448], + [ 0.000000331477068 , 0.0000003322497832, 0.0000003327649267], + [-0.0003304407325544, -0.0003304415763239, -0.0003304411144711], + [-0.0000001969979735, -0.0000001975841712, 0.0000003919886638], + [ 0.0003303068574212, 0.0003303051521186, -0.0003305007467702], + [-0.0000001974953534, 0.0000003921751812, -0.0000001968025742], + [ 0.0003303066264948, -0.0003305007645338, 0.000330305631735 ], + [ 0.0000003910827218, -0.0000001970779095, -0.0000001967759289], + [-0.0003305006668342, 0.0003303075324368, 0.0003303051876458], + [10.485175710082117 , -0.0000000322586402, -0.0000000312638804], + [-0.0000000314681614, 10.48517571377694 , -0.0000000326494387], + [-0.000000033377745 , -0.0000000317967874, 10.485175704362248 ], ]) # It can match down to 1e-9, if the finite difference is computed using the same dm from SCF. @@ -483,16 +524,13 @@ def test_xc_gradient_unrestricted_k_with_response(self): verbose = 0, ) - kpts = cell.make_kpts((1,3,1)) + kmesh = (1,4,1) + kpts = cell.make_kpts(kmesh) mf = kuks.KUKS(cell, xc="lda", kpts=kpts) mf.grids = gen_grid.BeckeGrids(cell) mf.grids.atom_grid = (40,194) mf.conv_tol = 1e-10 - # TODO: This is a hack to avoid a OOM issue in get_hcore() function (using multigrid_v3 internally) - hcore = mf.to_cpu().get_hcore() - mf.get_hcore = lambda: cp.asarray(hcore) - mf.kernel() dm = mf.make_rdm1() @@ -501,9 +539,10 @@ def test_xc_gradient_unrestricted_k_with_response(self): test_gradient = unrestricted_get_vxc_full_response(mf._numint, cell, mf.grids, mf.xc, dm, kpts, hermi=1) # dm is not very stable, and numerical gradient is super fast - ref_gradient = numerical_gradient_exc_becke(cell, "lda", kpts, None, (40,194), dm, unrestricted=True) + ref_gradient = numerical_gradient_exc_becke(cell, "lda", kmesh, None, (40,194), dm, unrestricted=True, dx=1e-5) - assert np.max(np.abs(test_gradient - ref_gradient)) < 1e-9 + assert np.max(np.abs(test_gradient[:-3] - ref_gradient[:-3])) < 1e-9 + assert np.max(np.abs(test_gradient[-3:] - ref_gradient[-3:])) < 4e-9 if __name__ == '__main__': print("Full Tests for PBC Becke grids") diff --git a/gpu4pyscf/pbc/grad/krks.py b/gpu4pyscf/pbc/grad/krks.py index 835ff77ef..c512d1209 100644 --- a/gpu4pyscf/pbc/grad/krks.py +++ b/gpu4pyscf/pbc/grad/krks.py @@ -27,44 +27,115 @@ from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.pbc.dft import multigrid, multigrid_v3, BeckeGrids from gpu4pyscf.pbc.dft.gen_grid import get_becke_weight_derivative +from gpu4pyscf.pbc.dft.numint import _GTOvalOpt +from gpu4pyscf.pbc.grad.krks_stress import _eval_ao_strain_derivatives __all__ = ['Gradients'] +XX, XY, XZ = 4, 5, 6 +YX, YY, YZ = 5, 7, 8 +ZX, ZY, ZZ = 6, 8, 9 + +def get_d2mu_dr2(ao_ks): + assert ao_ks.ndim == 4 + nkpts = ao_ks.shape[0] + ngrids = ao_ks.shape[2] + nao = ao_ks.shape[3] + + d2mu_dr2 = cp.empty([nkpts, 3, 3, ngrids, nao], dtype = ao_ks.dtype) + d2mu_dr2[:,0,0,:,:] = ao_ks[:,XX,:,:] + d2mu_dr2[:,0,1,:,:] = ao_ks[:,XY,:,:] + d2mu_dr2[:,1,0,:,:] = ao_ks[:,XY,:,:] + d2mu_dr2[:,0,2,:,:] = ao_ks[:,XZ,:,:] + d2mu_dr2[:,2,0,:,:] = ao_ks[:,XZ,:,:] + d2mu_dr2[:,1,1,:,:] = ao_ks[:,YY,:,:] + d2mu_dr2[:,1,2,:,:] = ao_ks[:,YZ,:,:] + d2mu_dr2[:,2,1,:,:] = ao_ks[:,YZ,:,:] + d2mu_dr2[:,2,2,:,:] = ao_ks[:,ZZ,:,:] + return d2mu_dr2 + def get_vxc(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): '''derivatives of the Exc per cell''' assert dm_kpts.ndim == 3 xctype = ni._xc_type(xc_code) nao = cell.nao nkpts = len(kpts) - vmat = cp.zeros((nkpts,3,nao,nao), dtype=dm_kpts.dtype) + if xctype == 'LDA': + ao_deriv = 0 + elif xctype == 'GGA': ao_deriv = 1 - for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts, - sort_grids=True): + elif xctype == 'MGGA': + ao_deriv = 1 + else: + raise NotImplementedError(f"Unrecognized xctype = {xctype}") + eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=ao_deriv) + + vmat = cp.zeros((nkpts,3,nao,nao), dtype=dm_kpts.dtype) + de_stress_rho = cp.zeros((3,3)) + exc_sum = 0 + + if xctype == 'LDA': + for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True): rho = ni.eval_rho(cell, ao_ks[:,0], dm_kpts, xctype=xctype, hermi=hermi) - vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[1] + exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2] wv = weight * vxc[0] aow = cp.einsum('kpi,p->kpi', ao_ks[:,0], wv) for kn in range(nkpts): vmat[kn] += _d1_dot_(ao_ks[kn,1:4], aow[kn]) + del aow, vxc + + exc_sum += cp.sum(weight * (rho * exc)) + + del rho, exc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain = ao_ks_strain[:,:,:,0] + ao_ks_strain += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T) + dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj()) + drho_stress = contract('kxypg,kpg->xyg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('xyg,g->xy', drho_stress, wv).real + + del ao_ks_strain, dm_nu, drho_stress, wv + elif xctype == 'GGA': - ao_deriv = 2 - for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts, - sort_grids=True): + for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True): rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi) - vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[1] + exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2] wv = weight * vxc wv[0] *= .5 for kn in range(nkpts): vmat[kn] += _gga_grad_sum_(ao_ks[kn], wv) + del vxc + + exc_sum += cp.sum(weight * (rho[0] * exc)) + + del rho, exc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, coords.T) + del d2ao + + wv[0] *= 2 + dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv).real + del dmu_stress, dm_nu + dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real + del dnu_stress, dm_dmu + + del ao_ks_strain, wv + elif xctype == 'MGGA': - ao_deriv = 2 - for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts, - sort_grids=True): + for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True): rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi) - vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[1] + exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=0)[:2] wv = weight * vxc wv[0] *= .5 wv[4] *= .5 # for the factor 1/2 in tau @@ -72,15 +143,47 @@ def get_vxc(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): vmat[kn] += _gga_grad_sum_(ao_ks[kn], wv[:4]) vmat[kn] += _tau_grad_dot_(ao_ks[kn], wv[4]) + del vxc + + exc_sum += cp.sum(weight * (rho[0] * exc)) + + del rho, exc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, coords.T) + del d2ao + + wv[0] *= 2 + dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv[0:4]).real + del dmu_stress, dm_nu + dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real + del dnu_stress, dm_dmu + dm_dnu = contract('kpq,kdgq->kdpg', dm_kpts, ao_ks[:,1:4].conj()) + tau_stress = contract('kxydpg,kdpg->xyg', ao_ks_strain[:,:,:,1:4], dm_dnu) + de_stress_rho += 2 * contract('xyg,g->xy', tau_stress, wv[4]).real + del tau_stress, dm_dnu + + del ao_ks_strain, wv + elif xctype == 'HF': pass elif xctype == 'NLC': raise NotImplementedError("NLC") else: - raise NotImplementedError(xc_code) + raise NotImplementedError(f"Unrecognized xctype = {xctype}") + + de_stress_weight = exc_sum * cp.eye(3) - exc = krhf_grad.contract_h1e_dm(cell, vmat, dm_kpts, hermi=1) - exc *= -1.0 / nkpts + exc = np.zeros((cell.natm + 3, 3)) + exc[:-3] = -krhf_grad.contract_h1e_dm(cell, vmat, dm_kpts, hermi=1) + exc[:-3] *= 1.0 / nkpts + exc[-3:] = (de_stress_rho / nkpts + de_stress_weight).get() return exc def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): @@ -103,8 +206,9 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): ao_deriv = 1 else: raise NotImplementedError(f"Unrecognized xctype = {xctype}") + eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=ao_deriv) - de_grid_response_weight = cp.zeros((natm, 3), dtype=cp.float64) + de_grid_response_weight = cp.zeros((natm + 3, 3), dtype=cp.float64) g1 = 0 for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts): g0, g1 = g1, g1 + weight.size @@ -121,6 +225,7 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response = cp.zeros((nkpts,3,nao,nao), dtype=dm_kpts.dtype) de_grid_response_rho = cp.zeros((natm, 3), dtype=dm_kpts.dtype) + de_stress_rho = cp.zeros((3,3), dtype=cp.float64) g1 = 0 for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts): @@ -140,7 +245,20 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response[kn] += vtmp_k de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_k, dm_kpts[kn]) * 2 del vtmp_k - del wv, aow, rho, vxc + + del aow, rho, vxc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain = ao_ks_strain[:,:,:,0] + associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]] + ao_ks_strain += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T) + del associated_supatm_coords + + dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj()) + drho_stress = contract('kxypg,kpg->xyg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('xyg,g->xy', drho_stress, wv).real + + del ao_ks_strain, dm_nu, drho_stress, wv elif xctype == 'GGA': rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi) @@ -153,7 +271,26 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response[kn] += vtmp_k de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_k, dm_kpts[kn]) * 2 del vtmp_k - del wv, rho, vxc + del rho, vxc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]] + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, associated_supatm_coords.T) + del d2ao, associated_supatm_coords + + wv[0] *= 2 + dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv).real + del dmu_stress, dm_nu + dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real + del dnu_stress, dm_dmu + + del ao_ks_strain, wv elif xctype == 'MGGA': rho = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts, xctype=xctype, hermi=hermi) @@ -167,14 +304,39 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response[kn] += vtmp_k de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_k, dm_kpts[kn]) * 2 del vtmp_k - del wv, rho, vxc + del rho, vxc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]] + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, associated_supatm_coords.T) + del d2ao, associated_supatm_coords + + wv[0] *= 2 + dm_nu = contract('kpq,kgq->kpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,kpg->xydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('xydg,dg->xy', dmu_stress, wv[0:4]).real + del dmu_stress, dm_nu + dm_dmu = contract('kpq,kdgp->kdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,kdqg->xydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('xydg,dg->xy', dnu_stress, wv[1:4]).real + del dnu_stress, dm_dmu + dm_dnu = contract('kpq,kdgq->kdpg', dm_kpts, ao_ks[:,1:4].conj()) + tau_stress = contract('kxydpg,kdpg->xyg', ao_ks_strain[:,:,:,1:4], dm_dnu) + de_stress_rho += 2 * contract('xyg,g->xy', tau_stress, wv[4]).real + del tau_stress, dm_dnu + + del ao_ks_strain, wv else: raise NotImplementedError(f"Unrecognized xctype = {xctype}") assert g1 == ngrids - exc = de_grid_response_rho.get().real - exc -= krhf_grad.contract_h1e_dm(cell, dvmat_orbital_response, dm_kpts, hermi=1) + exc = np.zeros((cell.natm + 3, 3), dtype=np.float64) + exc[:-3] = de_grid_response_rho.get().real + exc[:-3] -= krhf_grad.contract_h1e_dm(cell, dvmat_orbital_response, dm_kpts, hermi=1) + exc[-3:] = de_stress_rho.get() exc *= 1.0 / nkpts exc += de_grid_response_weight.get() return exc @@ -233,12 +395,7 @@ def energy_ee(self, dm, kpts): fn = get_vxc_full_response else: fn = get_vxc - de[:-3] = fn(ni, cell, grids, xc, dm, kpts) - if isinstance(grids, BeckeGrids): - de[-3:] = np.nan - else: - de[-3:] = multigrid_v3.MultiGridNumInt(cell).energy_strain_gradient( - xc, dm, kpts=kpts, spin=0, with_j=False, with_nuc=False) + de = fn(ni, cell, grids, xc, dm, kpts) t0 = log.timer_debug1('vxc', *t0) if j_factor != 0 or k_sr != 0 or k_lr != 0: diff --git a/gpu4pyscf/pbc/grad/krks_stress.py b/gpu4pyscf/pbc/grad/krks_stress.py index 25cfc4124..c94b43c70 100644 --- a/gpu4pyscf/pbc/grad/krks_stress.py +++ b/gpu4pyscf/pbc/grad/krks_stress.py @@ -100,13 +100,6 @@ def get_vxc(ks_grad, cell, dm_kpts, kpts, with_j=False, with_nuc=False): assert kpts.ndim == 2 assert dm_kpts.ndim == 3 - if not cell.cart: - c2s = asarray(cell.cart2sph_coeff()) - dm_kpts = sandwich_dot(dm_kpts, c2s.T) - # Ensure all AOs are evaluated in the Cartesian GTOs as ao_ks strain - # derivatives currently supports Cartesian format only - cell = cell.copy() - cell.cart = True nkpts, nao = dm_kpts.shape[:2] assert nkpts == len(kpts) @@ -303,8 +296,8 @@ def _eval_ao_strain_derivatives(cell, coords, kpts=None, deriv=0, out=None, coords = cp.asarray(coords.T, order='C') bvk_ncells = opt.bvk_ncells comp = (deriv+1)*(deriv+2)*(deriv+3)//6 - nao = cell.nao_nr(cart=True) - cart = 1 + nao = cell.nao_nr() + cart = cell.cart out = cp.empty((3, 3, comp, bvk_ncells, nao, ngrids)) drv = libpbc.PBCeval_gto_strain_tensor diff --git a/gpu4pyscf/pbc/grad/kuks.py b/gpu4pyscf/pbc/grad/kuks.py index 53be3f39e..3e2d0268a 100644 --- a/gpu4pyscf/pbc/grad/kuks.py +++ b/gpu4pyscf/pbc/grad/kuks.py @@ -28,6 +28,9 @@ from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.pbc.dft import multigrid, multigrid_v3, BeckeGrids from gpu4pyscf.pbc.dft.gen_grid import get_becke_weight_derivative +from gpu4pyscf.pbc.dft.numint import _GTOvalOpt +from gpu4pyscf.pbc.grad.krks_stress import _eval_ao_strain_derivatives +from gpu4pyscf.pbc.grad.krks import get_d2mu_dr2 __all__ = ['Gradients'] @@ -36,15 +39,28 @@ def get_vxc(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): xctype = ni._xc_type(xc_code) nao = cell.nao nkpts = len(kpts) - vmat = cp.zeros((2,nkpts,3,nao,nao), dtype=dm_kpts.dtype) + if xctype == 'LDA': + ao_deriv = 0 + elif xctype == 'GGA': + ao_deriv = 1 + elif xctype == 'MGGA': ao_deriv = 1 - for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts, - sort_grids=True): + else: + raise NotImplementedError(f"Unrecognized xctype = {xctype}") + eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=ao_deriv) + + vmat = cp.zeros((2,nkpts,3,nao,nao), dtype=dm_kpts.dtype) + de_stress_rho = cp.zeros((3,3)) + exc_sum = 0 + + if xctype == 'LDA': + for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True): rho_a = ni.eval_rho(cell, ao_ks[:,0], dm_kpts[0], xctype=xctype, hermi=hermi) rho_b = ni.eval_rho(cell, ao_ks[:,0], dm_kpts[1], xctype=xctype, hermi=hermi) rho = cp.stack([rho_a, rho_b], axis=0) - vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[1] + del rho_a, rho_b + exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[:2] wv = weight * vxc[:,0] aowa = cp.einsum('xpi,p->xpi', ao_ks[:,0], wv[0]) aowb = cp.einsum('xpi,p->xpi', ao_ks[:,0], wv[1]) @@ -52,28 +68,65 @@ def get_vxc(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): vmat[0,kn] += krks_grad._d1_dot_(ao_ks[kn,1:4], aowa[kn]) vmat[1,kn] += krks_grad._d1_dot_(ao_ks[kn,1:4], aowb[kn]) + del aowa, aowb, vxc + + exc_sum += cp.sum(weight * ((rho[0] + rho[1]) * exc)) + + del rho, exc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain = ao_ks_strain[:,:,:,0] + ao_ks_strain += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T) + dm_nu = contract('ukpq,kgq->ukpg', dm_kpts, ao_ks[:,0].conj()) + drho_stress = contract('kxypg,ukpg->uxyg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('uxyg,ug->xy', drho_stress, wv).real + + del ao_ks_strain, dm_nu, drho_stress, wv + elif xctype == 'GGA': - ao_deriv = 2 - for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts, - sort_grids=True): + for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True): rho_a = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts[0], xctype=xctype, hermi=hermi) rho_b = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts[1], xctype=xctype, hermi=hermi) rho = cp.stack([rho_a, rho_b], axis=0) - vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[1] + del rho_a, rho_b + exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[:2] wv = weight * vxc wv[:,0] *= .5 for kn in range(nkpts): vmat[0,kn] += krks_grad._gga_grad_sum_(ao_ks[kn], wv[0]) vmat[1,kn] += krks_grad._gga_grad_sum_(ao_ks[kn], wv[1]) + del vxc + + exc_sum += cp.sum(weight * ((rho[0,0] + rho[1,0]) * exc)) + + del rho, exc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, coords.T) + del d2ao + + wv[:,0] *= 2 + dm_nu = contract('ukpq,kgq->ukpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,ukpg->uxydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dmu_stress, wv).real + del dmu_stress, dm_nu + dm_dmu = contract('ukpq,kdgp->ukdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,ukdqg->uxydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dnu_stress, wv[:,1:4]).real + del dnu_stress, dm_dmu + + del ao_ks_strain, wv + elif xctype == 'MGGA': - ao_deriv = 2 - for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts, - sort_grids=True): + for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts, sort_grids=True): rho_a = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts[0], xctype=xctype, hermi=hermi) rho_b = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts[1], xctype=xctype, hermi=hermi) rho = cp.stack([rho_a, rho_b], axis=0) - vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[1] + del rho_a, rho_b + exc, vxc = ni.eval_xc_eff(xc_code, rho, deriv=1, xctype=xctype, spin=1)[:2] wv = weight * vxc wv[:,0] *= .5 wv[:,4] *= .5 # for the factor 1/2 in tau @@ -83,15 +136,47 @@ def get_vxc(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): vmat[0,kn] += krks_grad._tau_grad_dot_(ao_ks[kn], wv[0,4]) vmat[1,kn] += krks_grad._tau_grad_dot_(ao_ks[kn], wv[1,4]) + del vxc + + exc_sum += cp.sum(weight * ((rho[0,0] + rho[1,0]) * exc)) + + del rho, exc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, coords.T) + del d2ao + + wv[:,0] *= 2 + dm_nu = contract('ukpq,kgq->ukpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,ukpg->uxydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dmu_stress, wv[:,0:4]).real + del dmu_stress, dm_nu + dm_dmu = contract('ukpq,kdgp->ukdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,ukdqg->uxydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dnu_stress, wv[:,1:4]).real + del dnu_stress, dm_dmu + dm_dnu = contract('ukpq,kdgq->ukdpg', dm_kpts, ao_ks[:,1:4].conj()) + tau_stress = contract('kxydpg,ukdpg->uxyg', ao_ks_strain[:,:,:,1:4], dm_dnu) + de_stress_rho += 2 * contract('uxyg,ug->xy', tau_stress, wv[:,4]).real + del tau_stress, dm_dnu + + del ao_ks_strain, wv + elif xctype == 'HF': pass elif xctype == 'NLC': raise NotImplementedError("NLC") else: - raise NotImplementedError(xc_code) + raise NotImplementedError(f"Unrecognized xctype = {xctype}") - exc = krhf_grad.contract_h1e_dm(cell, vmat, dm_kpts, hermi=1) - exc *= -1.0 / nkpts + de_stress_weight = exc_sum * cp.eye(3) + + exc = np.zeros((cell.natm + 3, 3)) + exc[:-3] = -krhf_grad.contract_h1e_dm(cell, vmat, dm_kpts, hermi=1) + exc[:-3] *= 1.0 / nkpts + exc[-3:] = (de_stress_rho / nkpts + de_stress_weight).get() return exc def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): @@ -115,8 +200,9 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): ao_deriv = 1 else: raise NotImplementedError(f"Unrecognized xctype = {xctype}") + eval_gto_opt = _GTOvalOpt(cell, kpts, deriv=ao_deriv) - de_grid_response_weight = cp.zeros((natm, 3), dtype=cp.float64) + de_grid_response_weight = cp.zeros((natm + 3, 3), dtype=cp.float64) g1 = 0 for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv, kpts): g0, g1 = g1, g1 + weight.size @@ -137,6 +223,7 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response = cp.zeros((2,nkpts,3,nao,nao), dtype=dm_kpts.dtype) de_grid_response_rho = cp.zeros((natm, 3), dtype=dm_kpts.dtype) + de_stress_rho = cp.zeros((3,3), dtype=cp.float64) g1 = 0 for ao_ks, weight, coords in ni.block_loop(cell, grids, ao_deriv + 1, kpts): @@ -162,7 +249,19 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response[1,kn] += vtmp_b de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_b, dm_kpts[1,kn]) * 2 del vtmp_a, vtmp_b - del wv, rho, aowa, aowb, vxc + del rho, aowa, aowb, vxc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + ao_ks_strain = ao_ks_strain[:,:,:,0] + associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]] + ao_ks_strain += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T) + del associated_supatm_coords + + dm_nu = contract('ukpq,kgq->ukpg', dm_kpts, ao_ks[:,0].conj()) + drho_stress = contract('kxypg,ukpg->uxyg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('uxyg,ug->xy', drho_stress, wv).real + + del ao_ks_strain, dm_nu, drho_stress, wv elif xctype == 'GGA': rho_a = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts[0], xctype=xctype, hermi=hermi) @@ -180,7 +279,26 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response[1,kn] += vtmp_b de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_b, dm_kpts[1,kn]) * 2 del vtmp_a, vtmp_b - del wv, rho, vxc + del rho, vxc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]] + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, associated_supatm_coords.T) + del d2ao, associated_supatm_coords + + wv[:,0] *= 2 + dm_nu = contract('ukpq,kgq->ukpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,ukpg->uxydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dmu_stress, wv).real + del dmu_stress, dm_nu + dm_dmu = contract('ukpq,kdgp->ukdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,ukdqg->uxydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dnu_stress, wv[:,1:4]).real + del dnu_stress, dm_dmu + + del ao_ks_strain, wv elif xctype == 'MGGA': rho_a = ni.eval_rho(cell, ao_ks[:,:4], dm_kpts[0], xctype=xctype, hermi=hermi) @@ -199,14 +317,39 @@ def get_vxc_full_response(ni, cell, grids, xc_code, dm_kpts, kpts, hermi=1): dvmat_orbital_response[1,kn] += vtmp_b de_grid_response_rho[i_atom] += cp.einsum('xij,ji->x', vtmp_b, dm_kpts[1,kn]) * 2 del vtmp_a, vtmp_b - del wv, rho, vxc + del rho, vxc + + ao_ks_strain = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=ao_deriv, opt=eval_gto_opt) + associated_supatm_coords = grids.supatm_coords[grids.supatm_idx[g0:g1]] + ao_ks_strain[:,:,:,0] += contract('kxgp,yg->kxypg', ao_ks[:,1:4], associated_supatm_coords.T) + d2ao = get_d2mu_dr2(ao_ks) + ao_ks_strain[:,:,:,1:4] += contract('kxdgp,yg->kxydpg', d2ao, associated_supatm_coords.T) + del d2ao, associated_supatm_coords + + wv[:,0] *= 2 + dm_nu = contract('ukpq,kgq->ukpg', dm_kpts, ao_ks[:,0].conj()) + dmu_stress = contract('kxydpg,ukpg->uxydg', ao_ks_strain, dm_nu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dmu_stress, wv[:,0:4]).real + del dmu_stress, dm_nu + dm_dmu = contract('ukpq,kdgp->ukdqg', dm_kpts, ao_ks[:,1:4]) + dnu_stress = contract('kxyqg,ukdqg->uxydg', ao_ks_strain[:,:,:,0].conj(), dm_dmu) + de_stress_rho += 2 * contract('uxydg,udg->xy', dnu_stress, wv[:,1:4]).real + del dnu_stress, dm_dmu + dm_dnu = contract('ukpq,kdgq->ukdpg', dm_kpts, ao_ks[:,1:4].conj()) + tau_stress = contract('kxydpg,ukdpg->uxyg', ao_ks_strain[:,:,:,1:4], dm_dnu) + de_stress_rho += 2 * contract('uxyg,ug->xy', tau_stress, wv[:,4]).real + del tau_stress, dm_dnu + + del ao_ks_strain, wv else: raise NotImplementedError(f"Unrecognized xctype = {xctype}") assert g1 == ngrids - exc = de_grid_response_rho.get().real - exc -= krhf_grad.contract_h1e_dm(cell, dvmat_orbital_response, dm_kpts, hermi=1) + exc = np.zeros((cell.natm + 3, 3), dtype=np.float64) + exc[:-3] = de_grid_response_rho.get().real + exc[:-3] -= krhf_grad.contract_h1e_dm(cell, dvmat_orbital_response, dm_kpts, hermi=1) + exc[-3:] = de_stress_rho.get() exc *= 1.0 / nkpts exc += de_grid_response_weight.get() return exc @@ -252,12 +395,7 @@ def energy_ee(self, dm, kpts): else: fn = get_vxc cell = self.cell - de[:-3] = fn(ni, cell, grids, xc, dm, kpts) - if isinstance(grids, BeckeGrids): - de[-3:] = np.nan - else: - de[-3:] = multigrid_v3.MultiGridNumInt(cell).energy_strain_gradient( - xc, dm, kpts=kpts, spin=1, with_j=False, with_nuc=False) + de = fn(ni, cell, grids, xc, dm, kpts) t0 = log.timer_debug1('vxc', *t0) if j_factor != 0 or k_sr != 0 or k_lr != 0: diff --git a/gpu4pyscf/pbc/grad/rks.py b/gpu4pyscf/pbc/grad/rks.py index 70521667e..45a7b49a9 100644 --- a/gpu4pyscf/pbc/grad/rks.py +++ b/gpu4pyscf/pbc/grad/rks.py @@ -63,12 +63,7 @@ def energy_ee(self, dm): else: fn = get_vxc cell = self.cell - de[:-3] = fn(ni, cell, grids, xc, dm[None], np.zeros((1, 3))) - if isinstance(grids, BeckeGrids): - de[-3:] = np.nan - else: - de[-3:] = multigrid_v3.MultiGridNumInt(cell).energy_strain_gradient( - xc, dm, spin=0, with_j=False, with_nuc=False) + de = fn(ni, cell, grids, xc, dm[None], np.zeros((1, 3))) if j_factor != 0 or k_sr != 0 or k_lr != 0: de += rhf._get_ejk_derivatives(mf, dm, None, j_factor, omega, k_lr, k_sr) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py index 34c5c65fe..e4de8db50 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_krks_stress.py @@ -66,14 +66,40 @@ def _check_vs_finite_diff(dat, mf_scanner): assert abs(dat[i,j] - (e1-e2)/2/disp/vol) < 1e-6 class KnownValues(unittest.TestCase): - def test_eval_ao_kpts(self): + def test_eval_ao_cart_kpts(self): a = np.eye(3) * 5 np.random.seed(5) a += np.random.rand(3, 3) - .5 cell = gto.M(atom='He 1 1 1; He 2 1.5 2.4', basis=[[0, [.5, 1]], [1, [1.5, 1], [.5, 1]], - [2, [.8, 1]]], a=a, unit='Bohr', cart=True) + [2, [.8, 1]], + [4, [.8, 1]]], a=a, unit='Bohr', cart=True) + kmesh = [3, 1, 1] + kpts = cell.make_kpts(kmesh) + coords = np.random.rand(10, 3) + ao_value = _eval_ao_strain_derivatives(cell, coords, kpts) + ao_value = ao_value.get().transpose(0,1,2,3,5,4) + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1.precision = 1e-10 + cell2.precision = 1e-10 + ao1 = dft.numint.eval_ao_kpts(cell1, coords, cell1.make_kpts(kmesh)) + ao2 = dft.numint.eval_ao_kpts(cell2, coords, cell2.make_kpts(kmesh)) + assert abs(ao_value[0][i,j,0] - (ao1[0] - ao2[0]) / 2e-5).max() < 1e-9 + assert abs(ao_value[1][i,j,0] - (ao1[1] - ao2[1]) / 2e-5).max() < 1e-9 + assert abs(ao_value[2][i,j,0] - (ao1[2] - ao2[2]) / 2e-5).max() < 1e-9 + + def test_eval_ao_sph_kpts(self): + a = np.eye(3) * 5 + np.random.seed(5) + a += np.random.rand(3, 3) - .5 + cell = gto.M(atom='He 1 1 1; He 2 1.5 2.4', + basis=[[0, [.5, 1]], + [1, [1.5, 1], [.5, 1]], + [2, [.8, 1]], + [4, [.8, 1]]], a=a, unit='Bohr') + assert not cell.cart kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) coords = np.random.rand(10, 3) @@ -96,7 +122,8 @@ def test_eval_ao_deriv1_cart_kpts(self): cell = gto.M(atom='He 1 1 1; He 2 1.5 2.4', basis=[[0, [.5, 1]], [1, [1.5, 1], [.5, 1]], - [2, [.8, 1]]], a=a, unit='Bohr', cart=True) + [2, [.8, 1]], + [4, [.8, 1]]], a=a, unit='Bohr', cart=True) kmesh = [3, 1, 1] kpts = cell.make_kpts(kmesh) coords = np.random.rand(10, 3) @@ -112,6 +139,31 @@ def test_eval_ao_deriv1_cart_kpts(self): assert abs(ao_value[1][i,j] - (ao1[1] - ao2[1]) / 2e-5).max() < 1e-9 assert abs(ao_value[2][i,j] - (ao1[2] - ao2[2]) / 2e-5).max() < 1e-9 + def test_eval_ao_deriv1_sph_kpts(self): + a = np.eye(3) * 5 + np.random.seed(5) + a += np.random.rand(3, 3) - .5 + cell = gto.M(atom='He 1 1 1; He 2 1.5 2.4', + basis=[[0, [.5, 1]], + [1, [1.5, 1], [.5, 1]], + [2, [.8, 1]], + [4, [.8, 1]]], a=a, unit='Bohr') + assert not cell.cart + kmesh = [3, 1, 1] + kpts = cell.make_kpts(kmesh) + coords = np.random.rand(10, 3) + ao_value = _eval_ao_strain_derivatives(cell, coords, kpts, deriv=1) + ao_value = ao_value.get().transpose(0,1,2,3,5,4) + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + cell1.precision = 1e-10 + cell2.precision = 1e-10 + ao1 = dft.numint.eval_ao_kpts(cell1, coords, cell1.make_kpts(kmesh), deriv=1) + ao2 = dft.numint.eval_ao_kpts(cell2, coords, cell2.make_kpts(kmesh), deriv=1) + assert abs(ao_value[0][i,j] - (ao1[0] - ao2[0]) / 2e-5).max() < 2e-9 + assert abs(ao_value[1][i,j] - (ao1[1] - ao2[1]) / 2e-5).max() < 2e-9 + assert abs(ao_value[2][i,j] - (ao1[2] - ao2[2]) / 2e-5).max() < 2e-9 + def test_get_vxc_lda(self): a = np.eye(3) * 5 np.random.seed(5) diff --git a/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py b/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py index fb72d43f2..7173d5225 100644 --- a/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py +++ b/gpu4pyscf/pbc/grad/tests/test_pbc_rks_stress.py @@ -94,6 +94,27 @@ def test_eval_ao_cart(self): ao2 = ni.eval_ao(cell2, coords) assert abs(ao_value[i,j,0] - (ao1 - ao2) / 2e-5).max() < 5e-9 + def test_eval_ao_sph(self): + a = np.eye(3) * 5 + np.random.seed(5) + a += np.random.rand(3, 3) - .5 + cell = gto.M(atom='He 1 1 1; He 2 1.5 2.4', + basis=[[0, [.5, 1]], + [1, [1.5, 1], [.5, 1]], + [2, [.8, 1]], + [3, [.7, 1]], + [4, [.6, 1]]], a=a, unit='Bohr') + assert not cell.cart + coords = np.random.rand(10, 3) + ao_value = _eval_ao_strain_derivatives(cell, coords) + ao_value = ao_value.get().transpose(0,1,2,3,5,4)[0] + ni = NumInt() + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-5) + ao1 = ni.eval_ao(cell1, coords) + ao2 = ni.eval_ao(cell2, coords) + assert abs(ao_value[i,j,0] - (ao1 - ao2) / 2e-5).max() < 5e-9 + def test_eval_ao_deriv1_cart(self): a = np.eye(3) * 5 np.random.seed(5) @@ -114,6 +135,27 @@ def test_eval_ao_deriv1_cart(self): ao2 = ni.eval_ao(cell2, coords, deriv=1) assert abs(ao_value[i,j] - (ao1 - ao2) / 2e-4).max() < 1e-7 + def test_eval_ao_deriv1_sph(self): + a = np.eye(3) * 5 + np.random.seed(5) + a += np.random.rand(3, 3) - .5 + cell = gto.M(atom='He 1 1 1; He 2 1.5 2.4', + basis=[[0, [.5, 1]], + [1, [1.5, 1], [.5, 1]], + [2, [.8, 1]], + [3, [.7, 1]], + [4, [.6, 1]]], a=a, unit='Bohr') + assert not cell.cart + coords = np.random.rand(10, 3) + ao_value = _eval_ao_strain_derivatives(cell, coords, deriv=1) + ao_value = ao_value.get().transpose(0,1,2,3,5,4)[0] + ni = NumInt() + for (i, j) in [(0, 0), (0, 1), (0, 2), (2, 0), (2, 2)]: + cell1, cell2 = _finite_diff_cells(cell, i, j, disp=1e-4) + ao1 = ni.eval_ao(cell1, coords, deriv=1) + ao2 = ni.eval_ao(cell2, coords, deriv=1) + assert abs(ao_value[i,j] - (ao1 - ao2) / 2e-4).max() < 1e-7 + def test_get_vxc_lda(self): a = np.eye(3) * 5 np.random.seed(5) diff --git a/gpu4pyscf/pbc/grad/uks.py b/gpu4pyscf/pbc/grad/uks.py index 296271b7d..b3d311e03 100644 --- a/gpu4pyscf/pbc/grad/uks.py +++ b/gpu4pyscf/pbc/grad/uks.py @@ -60,13 +60,7 @@ def energy_ee(self, dm): fn = get_vxc cell = self.cell # UKS densities use (spin, k-point, AO, AO) ordering. - de[:-3] = fn(ni, cell, grids, xc, dm[:,None], np.zeros((1, 3))) - if isinstance(grids, BeckeGrids): - # Atomic grid response is supported; Becke-grid stress is not. - de[-3:] = np.nan - else: - de[-3:] = multigrid_v3.MultiGridNumInt(cell).energy_strain_gradient( - xc, dm, spin=1, with_j=False, with_nuc=False) + de = fn(ni, cell, grids, xc, dm[:,None], np.zeros((1, 3))) if j_factor != 0 or k_sr != 0 or k_lr != 0: de += rhf._get_ejk_derivatives(mf, dm, None, j_factor, omega, k_lr, k_sr)