From 5c7575a31e80a8f619274807a72ae7e8206d344c Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 27 Mar 2025 08:42:24 +0800 Subject: [PATCH 01/31] add the TDA solvent --- gpu4pyscf/solvent/_attach_solvent.py | 8 ++++- gpu4pyscf/solvent/pcm.py | 6 ++++ gpu4pyscf/solvent/tdscf/__init__.py | 13 +++++++ gpu4pyscf/solvent/tdscf/pcm.py | 54 ++++++++++++++++++++++++++++ 4 files changed, 80 insertions(+), 1 deletion(-) create mode 100644 gpu4pyscf/solvent/tdscf/__init__.py create mode 100644 gpu4pyscf/solvent/tdscf/pcm.py diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 3321909fd..b65db3f7f 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -16,7 +16,7 @@ from pyscf import lib from pyscf.lib import logger from gpu4pyscf.lib.cupy_helper import tag_array -from gpu4pyscf import scf +from gpu4pyscf import scf, dft, tdscf def _for_scf(mf, solvent_obj, dm=None): '''Add solvent model to SCF (HF and DFT) method. @@ -125,6 +125,12 @@ def nuc_grad_method(self): grad_method = super().nuc_grad_method() return self.with_solvent.nuc_grad_method(grad_method) + def TDA(self): + # if isinstance(self, dft.rks.RKS): + # tda_method = + tda_method = super().TDA() + return self.with_solvent.TDA(tda_method) + Gradients = nuc_grad_method def Hessian(self): diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 4c177063f..cce326913 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -413,6 +413,12 @@ def nuc_grad_method(self, grad_method): return pcm_grad.make_grad_object(grad_method) else: raise RuntimeError('Only SCF gradient is supported') + + def TDA(self, td): + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + if self.frozen: + raise RuntimeError('Frozen solvent model is not supported') + return pcm_td.make_tda_object(td) def Hessian(self, hess_method): from gpu4pyscf.solvent.hessian import pcm as pcm_hess diff --git a/gpu4pyscf/solvent/tdscf/__init__.py b/gpu4pyscf/solvent/tdscf/__init__.py new file mode 100644 index 000000000..2b3145125 --- /dev/null +++ b/gpu4pyscf/solvent/tdscf/__init__.py @@ -0,0 +1,13 @@ +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py new file mode 100644 index 000000000..89551d7f7 --- /dev/null +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -0,0 +1,54 @@ +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +''' +TD of PCM family solvent model +''' + +import numpy +import cupy +import ctypes +from cupyx import scipy +from pyscf import lib +from pyscf import gto +from pyscf.grad import rhf as rhf_grad +from gpu4pyscf.gto import int3c1e +from gpu4pyscf.solvent.pcm import PI, switch_h, libsolvent +from gpu4pyscf.gto.int3c1e_ip import int1e_grids_ip1, int1e_grids_ip2 +from gpu4pyscf.lib.cupy_helper import contract +from gpu4pyscf.lib import logger +from pyscf import lib as pyscf_lib + +libdft = lib.load_library('libdft') + +def make_tda_object(tda_method): + '''For td_method in vacuum, add td of solvent pcmobj''' + name = (tda_method._scf.with_solvent.__class__.__name__ + + tda_method.__class__.__name__) + return lib.set_class(WithSolventTDA(tda_method), + (WithSolventTDA, tda_method.__class__), name) + +class WithSolventTDA: + from gpu4pyscf.lib.utils import to_gpu, device + + def __init__(self, grad_method): + self.__dict__.update(grad_method.__dict__) + + def undo_solvent(self): + cls = self.__class__ + name_mixin = self.base.with_solvent.__class__.__name__ + obj = lib.view(self, lib.drop_class(cls, WithSolventTDA, name_mixin)) + return obj + + to_cpu = NotImplemented \ No newline at end of file From d2e1d969f964f32dc58e5467340f1d00cbdcb561 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 27 Mar 2025 14:08:57 +0800 Subject: [PATCH 02/31] add the support for tddft using PCM --- gpu4pyscf/solvent/_attach_solvent.py | 31 +++++++++++- gpu4pyscf/solvent/pcm.py | 58 +++++++++++++++++++++-- gpu4pyscf/solvent/tdscf/pcm.py | 29 ++++-------- gpu4pyscf/solvent/tests/test_pcm_tdscf.py | 13 +++++ 4 files changed, 106 insertions(+), 25 deletions(-) create mode 100644 gpu4pyscf/solvent/tests/test_pcm_tdscf.py diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index b65db3f7f..45748b693 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -131,6 +131,24 @@ def TDA(self): tda_method = super().TDA() return self.with_solvent.TDA(tda_method) + def TDDFT(self): + # if isinstance(self, dft.rks.RKS): + # tda_method = + tda_method = super().TDDFT() + return self.with_solvent.TDDFT(tda_method) + + def TDHF(self): + # if isinstance(self, dft.rks.RKS): + # tda_method = + tda_method = super().TDHF() + return self.with_solvent.TDHF(tda_method) + + def CasidaTDDFT(self): + # if isinstance(self, dft.rks.RKS): + # tda_method = + tda_method = super().CasidaTDDFT() + return self.with_solvent.CasidaTDDFT(tda_method) + Gradients = nuc_grad_method def Hessian(self): @@ -145,12 +163,21 @@ def gen_response(self, *args, **kwargs): singlet = singlet or singlet is None def vind_with_solvent(dm1): v = vind(dm1) - if self.with_solvent.equilibrium_solvation: + if self.with_solvent.eps_optical is not None: if is_uhf: v_solvent = self.with_solvent._B_dot_x(dm1) v += v_solvent[0] + v_solvent[1] elif singlet: - v += self.with_solvent._B_dot_x(dm1) + v += self.with_solvent._B + else: + logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') + else: + if self.with_solvent.equilibrium_solvation: + if is_uhf: + v_solvent = self.with_solvent._B_dot_x(dm1) + v += v_solvent[0] + v_solvent[1] + elif singlet: + v += self.with_solvent._B_dot_x(dm1) return v return vind_with_solvent diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index cce326913..861630aa8 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -245,7 +245,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', + 'equilibrium_solvation', 'e', 'v', 'eps_optical' } from gpu4pyscf.lib.utils import to_gpu, device kernel = ddcosmo.DDCOSMO.kernel @@ -265,6 +265,7 @@ def __init__(self, mol): self.lebedev_order = 29 self._intermediates = {} self.eps = 78.3553 + self.eps_optical = None self.max_cycle = 20 self.conv_tol = 1e-7 @@ -418,7 +419,25 @@ def TDA(self, td): from gpu4pyscf.solvent.tdscf import pcm as pcm_td if self.frozen: raise RuntimeError('Frozen solvent model is not supported') - return pcm_td.make_tda_object(td) + return pcm_td.make_tdscf_object(td) + + def TDHF(self, td): + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + if self.frozen: + raise RuntimeError('Frozen solvent model is not supported') + return pcm_td.make_tdscf_object(td) + + def TDDFT(self, td): + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + if self.frozen: + raise RuntimeError('Frozen solvent model is not supported') + return pcm_td.make_tdscf_object(td) + + def CasidaTDDFT(self, td): + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + if self.frozen: + raise RuntimeError('Frozen solvent model is not supported') + return pcm_td.make_tdscf_object(td) def Hessian(self, hess_method): from gpu4pyscf.solvent.hessian import pcm as pcm_hess @@ -445,8 +464,39 @@ def _B_dot_x(self, dms): nao = dms.shape[-1] dms = dms.reshape(-1,nao,nao) - K = self._intermediates['K'] - R = self._intermediates['R'] + if self.eps_optical is not None: + assert not self.equilibrium_solvation + epsilon = self.eps_optical + logger.info(self, 'eps optical = %s', self.eps_optical) + S = self._intermediates['S'] + D = self._intermediates['D'] + A = self._intermediates['A'] + epsilon = self.eps_optical + if self.method.upper() in ['C-PCM', 'CPCM']: + f_epsilon = (epsilon-1.)/epsilon + K = S + R = -f_epsilon * cupy.eye(K.shape[0]) + elif self.method.upper() == 'COSMO': + f_epsilon = (epsilon - 1.0)/(epsilon + 1.0/2.0) + K = S + R = -f_epsilon * cupy.eye(K.shape[0]) + elif self.method.upper() in ['IEF-PCM', 'IEFPCM']: + f_epsilon = (epsilon - 1.0)/(epsilon + 1.0) + DA = D*A + DAS = cupy.dot(DA, S) + K = S - f_epsilon/(2.0*PI) * DAS + R = -f_epsilon * (cupy.eye(K.shape[0]) - 1.0/(2.0*PI)*DA) + elif self.method.upper() == 'SS(V)PE': + f_epsilon = (epsilon - 1.0)/(epsilon + 1.0) + DA = D*A + DAS = cupy.dot(DA, S) + K = S - f_epsilon/(4.0*PI) * (DAS + DAS.T) + R = -f_epsilon * (cupy.eye(K.shape[0]) - 1.0/(2.0*PI)*DA) + else: + raise RuntimeError(f"Unknown implicit solvent model: {self.method}") + else: + K = self._intermediates['K'] + R = self._intermediates['R'] v_grids = -self._get_v(dms) b = cupy.dot(R, v_grids.T) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 89551d7f7..f4a8cabc5 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -16,39 +16,30 @@ TD of PCM family solvent model ''' -import numpy -import cupy -import ctypes -from cupyx import scipy +import cupy as cp from pyscf import lib -from pyscf import gto -from pyscf.grad import rhf as rhf_grad -from gpu4pyscf.gto import int3c1e from gpu4pyscf.solvent.pcm import PI, switch_h, libsolvent from gpu4pyscf.gto.int3c1e_ip import int1e_grids_ip1, int1e_grids_ip2 -from gpu4pyscf.lib.cupy_helper import contract +from gpu4pyscf.lib.cupy_helper import contract, tag_array from gpu4pyscf.lib import logger -from pyscf import lib as pyscf_lib +from gpu4pyscf import scf -libdft = lib.load_library('libdft') -def make_tda_object(tda_method): +def make_tdscf_object(tda_method): '''For td_method in vacuum, add td of solvent pcmobj''' name = (tda_method._scf.with_solvent.__class__.__name__ + tda_method.__class__.__name__) - return lib.set_class(WithSolventTDA(tda_method), - (WithSolventTDA, tda_method.__class__), name) + return lib.set_class(WithSolventTDSCF(tda_method), + (WithSolventTDSCF, tda_method.__class__), name) -class WithSolventTDA: +class WithSolventTDSCF: from gpu4pyscf.lib.utils import to_gpu, device - def __init__(self, grad_method): - self.__dict__.update(grad_method.__dict__) + def __init__(self, tda_method): + self.__dict__.update(tda_method.__dict__) def undo_solvent(self): cls = self.__class__ name_mixin = self.base.with_solvent.__class__.__name__ - obj = lib.view(self, lib.drop_class(cls, WithSolventTDA, name_mixin)) + obj = lib.view(self, lib.drop_class(cls, WithSolventTDSCF, name_mixin)) return obj - - to_cpu = NotImplemented \ No newline at end of file diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py new file mode 100644 index 000000000..3787aed7e --- /dev/null +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py @@ -0,0 +1,13 @@ +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. \ No newline at end of file From 199563d8afafbb9a29e4b1dc6b9f5b881c5d7627 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 27 Mar 2025 17:43:58 +0800 Subject: [PATCH 03/31] fix a small bug --- gpu4pyscf/solvent/tdscf/pcm.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index f4a8cabc5..4bd240878 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -43,3 +43,6 @@ def undo_solvent(self): name_mixin = self.base.with_solvent.__class__.__name__ obj = lib.view(self, lib.drop_class(cls, WithSolventTDSCF, name_mixin)) return obj + + nuc_grad_method = NotImplemented + From c8c88a152d0aa418cb2a284878abf520a0a032f6 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 27 Mar 2025 17:44:12 +0800 Subject: [PATCH 04/31] fix a typo --- gpu4pyscf/solvent/_attach_solvent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 45748b693..0de52dd58 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -168,7 +168,7 @@ def vind_with_solvent(dm1): v_solvent = self.with_solvent._B_dot_x(dm1) v += v_solvent[0] + v_solvent[1] elif singlet: - v += self.with_solvent._B + v += self.with_solvent._B_dot_x(dm1) else: logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') else: From 5f64e6e1ab69ad59bcb635521eb63031e8a33758 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 31 Mar 2025 14:57:16 +0800 Subject: [PATCH 05/31] Add the unit test for PCM tddft --- gpu4pyscf/solvent/_attach_solvent.py | 17 +- gpu4pyscf/solvent/pcm.py | 45 ++++-- gpu4pyscf/solvent/tdscf/pcm.py | 5 + gpu4pyscf/solvent/tests/test_pcm_tdscf.py | 183 +++++++++++++++++++++- gpu4pyscf/tdscf/rhf.py | 2 + gpu4pyscf/tdscf/uhf.py | 3 +- 6 files changed, 229 insertions(+), 26 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 0de52dd58..f7f81716a 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -163,14 +163,17 @@ def gen_response(self, *args, **kwargs): singlet = singlet or singlet is None def vind_with_solvent(dm1): v = vind(dm1) - if self.with_solvent.eps_optical is not None: - if is_uhf: - v_solvent = self.with_solvent._B_dot_x(dm1) - v += v_solvent[0] + v_solvent[1] - elif singlet: - v += self.with_solvent._B_dot_x(dm1) + if self.with_solvent.tdscf: + if not self.with_solvent.equilibrium_solvation: + if is_uhf: + v_solvent = self.with_solvent._B_dot_x(dm1) + v += v_solvent[0] + v_solvent[1] + elif singlet: + v += self.with_solvent._B_dot_x(dm1) + else: + logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') else: - logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') + raise NotImplementedError('Equilibrium solvation TDDFT (SS) is not implemented!') else: if self.with_solvent.equilibrium_solvation: if is_uhf: diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 9cc8c4ba6..53fe4bf28 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -246,7 +246,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'eps_optical' + 'equilibrium_solvation', 'e', 'v', 'eps_optical', 'tdscf' } from gpu4pyscf.lib.utils import to_gpu, device kernel = ddcosmo.DDCOSMO.kernel @@ -266,7 +266,8 @@ def __init__(self, mol): self.lebedev_order = 29 self._intermediates = {} self.eps = 78.3553 - self.eps_optical = None + self.eps_optical = 1.78 + self.tdscf = False self.max_cycle = 20 self.conv_tol = 1e-7 @@ -478,35 +479,42 @@ def _B_dot_x(self, dms): assert not self.equilibrium_solvation epsilon = self.eps_optical logger.info(self, 'eps optical = %s', self.eps_optical) - S = self._intermediates['S'] - D = self._intermediates['D'] - A = self._intermediates['A'] + F, A = get_F_A(self.surface) + D, S = get_D_S(self.surface, with_S = True, with_D = not self.if_method_in_CPCM_category) epsilon = self.eps_optical if self.method.upper() in ['C-PCM', 'CPCM']: f_epsilon = (epsilon-1.)/epsilon K = S - R = -f_epsilon * cupy.eye(K.shape[0]) elif self.method.upper() == 'COSMO': f_epsilon = (epsilon - 1.0)/(epsilon + 1.0/2.0) K = S - R = -f_epsilon * cupy.eye(K.shape[0]) elif self.method.upper() in ['IEF-PCM', 'IEFPCM']: f_epsilon = (epsilon - 1.0)/(epsilon + 1.0) DA = D*A DAS = cupy.dot(DA, S) K = S - f_epsilon/(2.0*PI) * DAS - R = -f_epsilon * (cupy.eye(K.shape[0]) - 1.0/(2.0*PI)*DA) elif self.method.upper() == 'SS(V)PE': f_epsilon = (epsilon - 1.0)/(epsilon + 1.0) DA = D*A DAS = cupy.dot(DA, S) K = S - f_epsilon/(4.0*PI) * (DAS + DAS.T) - R = -f_epsilon * (cupy.eye(K.shape[0]) - 1.0/(2.0*PI)*DA) else: raise RuntimeError(f"Unknown implicit solvent model: {self.method}") - else: - K = self._intermediates['K'] - R = self._intermediates['R'] + + K_LU, K_LU_pivot = lu_factor(K, overwrite_a = True, check_finite = False) + K = None + v_grids = -self._get_v(dms) + + b = self.left_multiply_R(v_grids.T, f_epsilon = f_epsilon) + q = self.left_solve_K(b, K_LU=K_LU, K_LU_pivot=K_LU_pivot).T + + vK_1 = self.left_solve_K(v_grids.T, K_LU=K_LU, K_LU_pivot=K_LU_pivot, K_transpose = True) + qt = self.left_multiply_R(vK_1, f_epsilon = f_epsilon, R_transpose = True).T + q_sym = (q + qt)/2.0 + + vmat = self._get_vmat(q_sym) + return vmat.reshape(out_shape) + v_grids = -self._get_v(dms) b = self.left_multiply_R(v_grids.T) @@ -523,8 +531,9 @@ def _B_dot_x(self, dms): def if_method_in_CPCM_category(self): return self.method.upper() in ['C-PCM', 'CPCM', "COSMO"] - def left_multiply_R(self, right_vector, R_transpose = False): - f_epsilon = self._intermediates['f_epsilon'] + def left_multiply_R(self, right_vector, f_epsilon = None, R_transpose = False): + if f_epsilon is None: + f_epsilon = self._intermediates['f_epsilon'] if self.if_method_in_CPCM_category: # R = -f_epsilon * cupy.eye(K.shape[0]) return -f_epsilon * right_vector @@ -537,9 +546,11 @@ def left_multiply_R(self, right_vector, R_transpose = False): DA = DA.T return -f_epsilon * (right_vector - 1.0/(2.0*PI) * cupy.dot(DA, right_vector)) - def left_solve_K(self, right_vector, K_transpose = False): + def left_solve_K(self, right_vector, K_LU = None, K_LU_pivot = None, K_transpose = False): ''' K^{-1} @ right_vector ''' - K_LU = self._intermediates['K_LU'] - K_LU_pivot = self._intermediates['K_LU_pivot'] + if K_LU is None: + K_LU = self._intermediates['K_LU'] + if K_LU_pivot is None: + K_LU_pivot = self._intermediates['K_LU_pivot'] return lu_solve((K_LU, K_LU_pivot), right_vector, trans = K_transpose, overwrite_b = False, check_finite = False) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 4bd240878..84a791912 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -37,12 +37,17 @@ class WithSolventTDSCF: def __init__(self, tda_method): self.__dict__.update(tda_method.__dict__) + tda_method._scf.with_solvent.tdscf = True def undo_solvent(self): cls = self.__class__ name_mixin = self.base.with_solvent.__class__.__name__ obj = lib.view(self, lib.drop_class(cls, WithSolventTDSCF, name_mixin)) return obj + + def _finalize(self): + super()._finalize() + self._scf.with_solvent.tdscf = False nuc_grad_method = NotImplemented diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py index 3787aed7e..7a66f6656 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py @@ -10,4 +10,185 @@ # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and -# limitations under the License. \ No newline at end of file +# limitations under the License. + +import unittest +import numpy as np +import cupy as cp +from pyscf import lib, gto +from gpu4pyscf.tdscf import rhf, rks +from gpu4pyscf import tdscf + +class KnownValues(unittest.TestCase): + @classmethod + def setUpClass(cls): + mol = gto.Mole() + mol.verbose = 5 + mol.output = '/dev/null' + mol.atom = [ + ['O', (0. , 0., 0.)], + ['H', (0. , -0.757, 0.587)], + ['H', (0. , 0.757, 0.587)], ] + mol.basis = 'def2svp' + cls.mol = mol.build() + + cls.mf = mf = mol.RHF().PCM().to_gpu() + cls.mf.with_solvent.method = 'C-PCM' + cls.mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids + cls.mf.with_solvent.eps = 78 + cls.mf.with_solvent.eps_optical = 1.78 + cls.mf = mf.run(conv_tol=1e-10) + + cls.mfu = mfu = mol.UHF().PCM().to_gpu() + cls.mfu.with_solvent.method = 'C-PCM' + cls.mfu.with_solvent.lebedev_order = 29 # 302 Lebedev grids + cls.mfu.with_solvent.eps = 78 + cls.mfu.with_solvent.eps_optical = 1.78 + cls.mfu = mfu.run(conv_tol=1e-10) + + mf_b3lyp_nodf = mol.RKS().PCM().to_gpu() + mf_b3lyp_nodf.xc = 'b3lyp' + mf_b3lyp_nodf.grids.atom_grid = (99,590) + mf_b3lyp_nodf.with_solvent.method = 'C-PCM' + mf_b3lyp_nodf.with_solvent.lebedev_order = 29 # 302 Lebedev grids + mf_b3lyp_nodf.with_solvent.eps = 78 + mf_b3lyp_nodf.with_solvent.eps_optical = 1.78 + mf_b3lyp_nodf.cphf_grids = mf_b3lyp_nodf.grids + cls.mf_b3lyp_nodf = mf_b3lyp_nodf.run(conv_tol=1e-10) + + mf_b3lyp_nodf_u = mol.UKS().PCM().to_gpu() + mf_b3lyp_nodf_u.xc = 'b3lyp' + mf_b3lyp_nodf_u.grids.atom_grid = (99,590) + mf_b3lyp_nodf_u.with_solvent.method = 'C-PCM' + mf_b3lyp_nodf_u.with_solvent.lebedev_order = 29 # 302 Lebedev grids + mf_b3lyp_nodf_u.with_solvent.eps = 78 + mf_b3lyp_nodf_u.with_solvent.eps_optical = 1.78 + mf_b3lyp_nodf_u.cphf_grids = mf_b3lyp_nodf_u.grids + cls.mf_b3lyp_nodf_u = mf_b3lyp_nodf_u.run(conv_tol=1e-10) + + mf_b3lyp_nodf_iefpcm = mol.RKS().PCM().to_gpu() + mf_b3lyp_nodf_iefpcm.xc = 'b3lyp' + mf_b3lyp_nodf_iefpcm.grids.atom_grid = (99,590) + mf_b3lyp_nodf_iefpcm.with_solvent.method = 'IEF-PCM' + mf_b3lyp_nodf_iefpcm.with_solvent.lebedev_order = 29 # 302 Lebedev grids + mf_b3lyp_nodf_iefpcm.with_solvent.eps = 78 + mf_b3lyp_nodf_iefpcm.with_solvent.eps_optical = 1.78 + mf_b3lyp_nodf_iefpcm.cphf_grids = mf_b3lyp_nodf_iefpcm.grids + cls.mf_b3lyp_nodf_iefpcm = mf_b3lyp_nodf_iefpcm.run(conv_tol=1e-10) + + mf_b3lyp_nodf_iefpcm_u = mol.RKS().PCM().to_gpu() + mf_b3lyp_nodf_iefpcm_u.xc = 'b3lyp' + mf_b3lyp_nodf_iefpcm_u.grids.atom_grid = (99,590) + mf_b3lyp_nodf_iefpcm_u.with_solvent.method = 'IEF-PCM' + mf_b3lyp_nodf_iefpcm_u.with_solvent.lebedev_order = 29 # 302 Lebedev grids + mf_b3lyp_nodf_iefpcm_u.with_solvent.eps = 78 + mf_b3lyp_nodf_iefpcm_u.with_solvent.eps_optical = 1.78 + mf_b3lyp_nodf_iefpcm_u.cphf_grids = mf_b3lyp_nodf_iefpcm_u.grids + cls.mf_b3lyp_nodf_iefpcm_u = mf_b3lyp_nodf_iefpcm.run(conv_tol=1e-10) + + @classmethod + def tearDownClass(cls): + cls.mol.stdout.close() + + def test_hf_CPCM(self): + """ + $rem + JOBTYPE sp + METHOD hf + BASIS def2-svp + CIS_N_ROOTS 5 + CIS_SINGLETS TRUE + CIS_TRIPLETS FALSE + SYMMETRY FALSE + SYM_IGNORE TRUE + ! RPA 2 # whether tddft or tda + BASIS_LIN_DEP_THRESH 12 + SOLVENT_METHOD PCM + $end + + $PCM + Theory CPCM + HeavyPoints 302 + HPoints 302 + $end + + $solvent + dielectric 78 + $end + """ + mf = self.mf + td = mf.TDHF() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-75.61072291, -75.54419399, -75.51949191, -75.45219025, -75.40975027]) + assert np.allclose(es_gound, ref) + + td = mf.TDA() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-75.60864828, -75.54169327, -75.51738767, -75.44915784, -75.40839714]) + assert np.allclose(es_gound, ref) + + def test_b3lyp_CPCM(self): + mf = self.mf_b3lyp_nodf + td = mf.TDDFT() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-76.06898428, -75.99630982, -75.98765186, -75.91045133, -75.84783748]) + assert np.allclose(es_gound, ref) + + td = mf.TDA() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-76.06789176, -75.99609709, -75.98589720, -75.90894600, -75.84699115]) + assert np.allclose(es_gound, ref) + + def test_b3lyp_IEFPCM(self): + mf = self.mf_b3lyp_nodf_iefpcm + td = mf.TDDFT() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-76.06881645, -75.99631929, -75.98713725, -75.91015704, -75.84668800]) + assert np.allclose(es_gound, ref) + + td = mf.TDA() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-76.06773319, -75.99610928, -75.98534912, -75.90861455, -75.84576041]) + assert np.allclose(es_gound, ref) + + def test_unrestricted_hf_CPCM(self): + mf = self.mfu + td = mf.TDHF() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-75.64482315, -75.61072291, -75.57156784, -75.56769949, -75.54419399]) + assert np.allclose(es_gound, ref) + + td = mf.TDA() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-75.64008952, -75.60864828, -75.56268842, -75.56241960, -75.54169327]) + assert np.allclose(es_gound, ref) + + def test_unrestricted_b3lyp_CPCM(self): + mf = self.mf_b3lyp_nodf_u + td = mf.TDDFT() + assert td.device == 'gpu' + es = td.kernel(nstates=5)[0] + es_gound = es + mf.e_tot + ref = np.array([-76.09301571, -76.06898428, -76.01822101, -76.01369024, -75.99630982]) + assert np.allclose(es_gound, ref) + + +if __name__ == "__main__": + print("Full Tests for PCM TDDFT") + unittest.main() \ No newline at end of file diff --git a/gpu4pyscf/tdscf/rhf.py b/gpu4pyscf/tdscf/rhf.py index f04acc72a..0e9353d13 100644 --- a/gpu4pyscf/tdscf/rhf.py +++ b/gpu4pyscf/tdscf/rhf.py @@ -43,6 +43,8 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True): Ref: Chem Phys Lett, 256, 454 ''' + if getattr(mf, 'with_solvent', None): + raise NotImplementedError("PCM is not supported for get_ab") if mo_energy is None: mo_energy = mf.mo_energy if mo_coeff is None: diff --git a/gpu4pyscf/tdscf/uhf.py b/gpu4pyscf/tdscf/uhf.py index 0145381b6..c9281ea64 100644 --- a/gpu4pyscf/tdscf/uhf.py +++ b/gpu4pyscf/tdscf/uhf.py @@ -42,7 +42,8 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None): B has three items: (B_aaaa, B_aabb, B_bbbb). B_bbaa = B_aabb.transpose(2,3,0,1). ''' - + if getattr(mf, 'with_solvent', None): + raise NotImplementedError("PCM is not supported for get_ab") if mo_energy is None: mo_energy = mf.mo_energy if mo_coeff is None: From cfcd4fe5863c51a99cca41e5581a8bf1cca6b745 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 31 Mar 2025 16:59:56 +0800 Subject: [PATCH 06/31] in writting --- gpu4pyscf/solvent/tdscf/pcm.py | 227 ++++++++++++++++++++++++++++++++- 1 file changed, 226 insertions(+), 1 deletion(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 84a791912..2e34ca234 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -23,6 +23,7 @@ from gpu4pyscf.lib.cupy_helper import contract, tag_array from gpu4pyscf.lib import logger from gpu4pyscf import scf +from gpu4pyscf.gto import int3c1e def make_tdscf_object(tda_method): @@ -32,6 +33,15 @@ def make_tdscf_object(tda_method): return lib.set_class(WithSolventTDSCF(tda_method), (WithSolventTDSCF, tda_method.__class__), name) + +def make_tdscf_gradient_object(tda_grad_method): + '''For td_method in vacuum, add td of solvent pcmobj''' + name = (tda_grad_method.base._scf.with_solvent.__class__.__name__ + + tda_grad_method.__class__.__name__) + return lib.set_class(WithSolventTDSCFGradient(tda_grad_method), + (WithSolventTDSCFGradient, tda_grad_method.__class__), name) + + class WithSolventTDSCF: from gpu4pyscf.lib.utils import to_gpu, device @@ -49,5 +59,220 @@ def _finalize(self): super()._finalize() self._scf.with_solvent.tdscf = False - nuc_grad_method = NotImplemented + def nuc_grad_method(self): + grad_method = super().nuc_grad_method() + return make_tdscf_gradient_object(grad_method) + + +def grad_qv(pcmobj, dm, q_sym = None): + ''' + contributions due to integrals + ''' + if not pcmobj._intermediates: + pcmobj.build() + dm_cache = pcmobj._intermediates.get('dm', None) + if dm_cache is not None and cp.linalg.norm(dm_cache - dm) < 1e-10: + pass + else: + pcmobj._get_vind(dm) + mol = pcmobj.mol + log = logger.new_logger(mol, mol.verbose) + t1 = log.init_timer() + gridslice = pcmobj.surface['gslice_by_atom'] + charge_exp = pcmobj.surface['charge_exp'] + grid_coords = pcmobj.surface['grid_coords'] + if q_sym is None: + q_sym = pcmobj._intermediates['q_sym'] + + intopt = int3c1e.VHFOpt(mol) + intopt.build(1e-14, aosym=False) + dvj = int1e_grids_ip1(mol, grid_coords, dm = dm, charges = q_sym, + direct_scf_tol = 1e-14, charge_exponents = charge_exp**2, + intopt=intopt) + dq = int1e_grids_ip2(mol, grid_coords, dm = dm, charges = q_sym, + direct_scf_tol = 1e-14, charge_exponents = charge_exp**2, + intopt=intopt) + + aoslice = mol.aoslice_by_atom() + dvj = 2.0 * cp.asarray([cp.sum(dvj[:,p0:p1], axis=1) for p0,p1 in aoslice[:,2:]]) + dq = cp.asarray([cp.sum(dq[:,p0:p1], axis=1) for p0,p1 in gridslice]) + de = dq + dvj + t1 = log.timer_debug1('grad qv', *t1) + return de.get() + +def grad_solver(pcmobj, dm): + ''' + dE = 0.5*v* d(K^-1 R) *v + q*dv + v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR - dK q) + ''' + mol = pcmobj.mol + log = logger.new_logger(mol, mol.verbose) + t1 = log.init_timer() + if not pcmobj._intermediates: + pcmobj.build() + dm_cache = pcmobj._intermediates.get('dm', None) + if dm_cache is not None and cp.linalg.norm(dm_cache - dm) < 1e-10: + pass + else: + pcmobj._get_vind(dm) + + gridslice = pcmobj.surface['gslice_by_atom'] + v_grids = pcmobj._intermediates['v_grids'] + q = pcmobj._intermediates['q'] + f_epsilon = pcmobj._intermediates['f_epsilon'] + if not pcmobj.if_method_in_CPCM_category: + A = pcmobj._intermediates['A'] + D = pcmobj._intermediates['D'] + S = pcmobj._intermediates['S'] + + vK_1 = pcmobj.left_solve_K(v_grids, K_transpose = True) + + def contract_bra(a, B, c): + ''' i,xij,j->jx ''' + tmp = a.dot(B) + return (tmp*c).T + + def contract_ket(a, B, c): + ''' i,xij,j->ix ''' + tmp = B.dot(c) + return (a*tmp).T + + de = cp.zeros([pcmobj.mol.natm,3]) + if pcmobj.method.upper() in ['C-PCM', 'CPCM', 'COSMO']: + # dR = 0, dK = dS + de_dS = 0.5 * vK_1.reshape(-1, 1) * left_multiply_dS(pcmobj.surface, q, stream=None) + de_dS -= 0.5 * q.reshape(-1, 1) * right_multiply_dS(pcmobj.surface, vK_1, stream=None) + de -= cp.asarray([cp.sum(de_dS[p0:p1], axis=0) for p0,p1 in gridslice]) + + dF, _ = get_dF_dA(pcmobj.surface, with_dA = False) + dSii = get_dSii(pcmobj.surface, dF) + de -= 0.5*contract('i,xij->jx', vK_1*q, dSii) # 0.5*cp.einsum('i,xij,i->jx', vK_1, dSii, q) + + elif pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SMD']: + dF, dA = get_dF_dA(pcmobj.surface) + dSii = get_dSii(pcmobj.surface, dF) + dF = None + + dD, dS = get_dD_dS(pcmobj.surface, with_D=True, with_S=True) + + # dR = f_eps/(2*pi) * (dD*A + D*dA), + # dK = dS - f_eps/(2*pi) * (dD*A*S + D*dA*S + D*A*dS) + fac = f_epsilon/(2.0*PI) + + Av = A*v_grids + de_dR = 0.5*fac * contract_ket(vK_1, dD, Av) + de_dR -= 0.5*fac * contract_bra(vK_1, dD, Av) + de_dR = cp.asarray([cp.sum(de_dR[p0:p1], axis=0) for p0,p1 in gridslice]) + + vK_1_D = vK_1.dot(D) + vK_1_Dv = vK_1_D * v_grids + de_dR += 0.5*fac * contract('j,xjn->nx', vK_1_Dv, dA) + + de_dS0 = 0.5*contract_ket(vK_1, dS, q) + de_dS0 -= 0.5*contract_bra(vK_1, dS, q) + de_dS0 = cp.asarray([cp.sum(de_dS0[p0:p1], axis=0) for p0,p1 in gridslice]) + + vK_1_q = vK_1 * q + de_dS0 += 0.5*contract('i,xin->nx', vK_1_q, dSii) + + vK_1_DA = vK_1_D*A + de_dS1 = 0.5*contract_ket(vK_1_DA, dS, q) + de_dS1 -= 0.5*contract_bra(vK_1_DA, dS, q) + de_dS1 = cp.asarray([cp.sum(de_dS1[p0:p1], axis=0) for p0,p1 in gridslice]) + + vK_1_DAq = vK_1_DA*q + de_dS1 += 0.5*contract('j,xjn->nx', vK_1_DAq, dSii) + + Sq = cp.dot(S,q) + ASq = A*Sq + de_dD = 0.5*contract_ket(vK_1, dD, ASq) + de_dD -= 0.5*contract_bra(vK_1, dD, ASq) + de_dD = cp.asarray([cp.sum(de_dD[p0:p1], axis=0) for p0,p1 in gridslice]) + + de_dA = 0.5*contract('j,xjn->nx', vK_1_D*Sq, dA) # 0.5*cp.einsum('j,xjn,j->nx', vK_1_D, dA, Sq) + + de_dK = de_dS0 - fac * (de_dD + de_dA + de_dS1) + de += de_dR - de_dK + + elif pcmobj.method.upper() in [ 'SS(V)PE' ]: + dF, dA = get_dF_dA(pcmobj.surface) + dSii = get_dSii(pcmobj.surface, dF) + dF = None + + dD, dS = get_dD_dS(pcmobj.surface, with_D=True, with_S=True) + + # dR = f_eps/(2*pi) * (dD*A + D*dA), + # dK = dS - f_eps/(2*pi) * (dD*A*S + D*dA*S + D*A*dS) + fac = f_epsilon/(2.0*PI) + + Av = A*v_grids + de_dR = 0.5*fac * contract_ket(vK_1, dD, Av) + de_dR -= 0.5*fac * contract_bra(vK_1, dD, Av) + de_dR = cp.asarray([cp.sum(de_dR[p0:p1], axis=0) for p0,p1 in gridslice]) + + vK_1_D = vK_1.dot(D) + vK_1_Dv = vK_1_D * v_grids + de_dR += 0.5*fac * contract('j,xjn->nx', vK_1_Dv, dA) + + de_dS0 = 0.5*contract_ket(vK_1, dS, q) + de_dS0 -= 0.5*contract_bra(vK_1, dS, q) + de_dS0 = cp.asarray([cp.sum(de_dS0[p0:p1], axis=0) for p0,p1 in gridslice]) + + vK_1_q = vK_1 * q + de_dS0 += 0.5*contract('i,xin->nx', vK_1_q, dSii) + + vK_1_DA = vK_1_D*A + de_dS1 = 0.5*contract_ket(vK_1_DA, dS, q) + de_dS1 -= 0.5*contract_bra(vK_1_DA, dS, q) + de_dS1 = cp.asarray([cp.sum(de_dS1[p0:p1], axis=0) for p0,p1 in gridslice]) + vK_1_DAq = vK_1_DA*q + de_dS1 += 0.5*contract('j,xjn->nx', vK_1_DAq, dSii) + + DT_q = cp.dot(D.T, q) + ADT_q = A * DT_q + de_dS1_T = 0.5*contract_ket(vK_1, dS, ADT_q) + de_dS1_T -= 0.5*contract_bra(vK_1, dS, ADT_q) + de_dS1_T = cp.asarray([cp.sum(de_dS1_T[p0:p1], axis=0) for p0,p1 in gridslice]) + vK_1_ADT_q = vK_1 * ADT_q + de_dS1_T += 0.5*contract('j,xjn->nx', vK_1_ADT_q, dSii) + + Sq = cp.dot(S,q) + ASq = A*Sq + de_dD = 0.5*contract_ket(vK_1, dD, ASq) + de_dD -= 0.5*contract_bra(vK_1, dD, ASq) + de_dD = cp.asarray([cp.sum(de_dD[p0:p1], axis=0) for p0,p1 in gridslice]) + + vK_1_S = cp.dot(vK_1, S) + vK_1_SA = vK_1_S * A + de_dD_T = 0.5*contract_ket(vK_1_SA, -dD.transpose(0,2,1), q) + de_dD_T -= 0.5*contract_bra(vK_1_SA, -dD.transpose(0,2,1), q) + de_dD_T = cp.asarray([cp.sum(de_dD_T[p0:p1], axis=0) for p0,p1 in gridslice]) + + de_dA = 0.5*contract('j,xjn->nx', vK_1_D*Sq, dA) # 0.5*cp.einsum('j,xjn,j->nx', vK_1_D, dA, Sq) + + de_dA_T = 0.5*contract('j,xjn->nx', vK_1_S*DT_q, dA) + + de_dK = de_dS0 - 0.5 * fac * (de_dD + de_dA + de_dS1 + de_dD_T + de_dA_T + de_dS1_T) + de += de_dR - de_dK + + else: + raise RuntimeError(f"Unknown implicit solvent model: {pcmobj.method}") + t1 = log.timer_debug1('grad solver', *t1) + return de.get() + + +class WithSolventTDSCFGradient: + from gpu4pyscf.lib.utils import to_gpu, device + + def __init__(self, tda_grad_method): + self.__dict__.update(tda_grad_method.__dict__) + tda_grad_method.base._scf.with_solvent.tdscf = True + + def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): + de = super().grad_elec(self, xy, singlet, atmlst, verbose) + + + def _finalize(self): + super()._finalize() + self._scf.with_solvent.tdscf = False From b00c9363d02b1c503223c3babec345a3889d94bf Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 31 Mar 2025 17:22:25 +0800 Subject: [PATCH 07/31] fix a typo --- gpu4pyscf/solvent/pcm.py | 23 ++++++++++++++++++++++- gpu4pyscf/solvent/tests/test_pcm_tdscf.py | 7 ------- 2 files changed, 22 insertions(+), 8 deletions(-) diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 53fe4bf28..9ea7102ab 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -400,6 +400,27 @@ def _get_vind(self, dms): self._intermediates['v_grids'] = v_grids[0] return epcm, vmat[0] + def _get_qsym(self, dms): + if not self._intermediates: + self.build() + nao = dms.shape[-1] + dms = dms.reshape(-1,nao,nao) + if dms.shape[0] == 2: + dms = (dms[0] + dms[1]).reshape(-1,nao,nao) + if not isinstance(dms, cupy.ndarray): + dms = cupy.asarray(dms) + v_grids_e = self._get_v(dms) + v_grids = self.v_grids_n - v_grids_e + + b = self.left_multiply_R(v_grids.T) + q = self.left_solve_K(b).T + + vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) + qt = self.left_multiply_R(vK_1, R_transpose = True).T + q_sym = (q + qt)/2.0 + + return q_sym[0] + def _get_v(self, dms): ''' return electrostatic potential on surface @@ -475,7 +496,7 @@ def _B_dot_x(self, dms): nao = dms.shape[-1] dms = dms.reshape(-1,nao,nao) - if self.eps_optical is not None: + if self.tdscf: assert not self.equilibrium_solvation epsilon = self.eps_optical logger.info(self, 'eps optical = %s', self.eps_optical) diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py index 7a66f6656..823e2ffdf 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py @@ -172,13 +172,6 @@ def test_unrestricted_hf_CPCM(self): ref = np.array([-75.64482315, -75.61072291, -75.57156784, -75.56769949, -75.54419399]) assert np.allclose(es_gound, ref) - td = mf.TDA() - assert td.device == 'gpu' - es = td.kernel(nstates=5)[0] - es_gound = es + mf.e_tot - ref = np.array([-75.64008952, -75.60864828, -75.56268842, -75.56241960, -75.54169327]) - assert np.allclose(es_gound, ref) - def test_unrestricted_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf_u td = mf.TDDFT() From c27e68a49df7c1b40e8b5746fcab8fea73e2da96 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 31 Mar 2025 17:26:04 +0800 Subject: [PATCH 08/31] fix a typo --- gpu4pyscf/solvent/tdscf/pcm.py | 41 +++------------------------------- 1 file changed, 3 insertions(+), 38 deletions(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 2e34ca234..32a6a0839 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -19,6 +19,7 @@ import cupy as cp from pyscf import lib from gpu4pyscf.solvent.pcm import PI, switch_h, libsolvent +from gpu4pyscf.solvent.grad.pcm import left_multiply_dS, right_multiply_dS, get_dF_dA, get_dSii, get_dD_dS from gpu4pyscf.gto.int3c1e_ip import int1e_grids_ip1, int1e_grids_ip2 from gpu4pyscf.lib.cupy_helper import contract, tag_array from gpu4pyscf.lib import logger @@ -62,43 +63,7 @@ def _finalize(self): def nuc_grad_method(self): grad_method = super().nuc_grad_method() return make_tdscf_gradient_object(grad_method) - -def grad_qv(pcmobj, dm, q_sym = None): - ''' - contributions due to integrals - ''' - if not pcmobj._intermediates: - pcmobj.build() - dm_cache = pcmobj._intermediates.get('dm', None) - if dm_cache is not None and cp.linalg.norm(dm_cache - dm) < 1e-10: - pass - else: - pcmobj._get_vind(dm) - mol = pcmobj.mol - log = logger.new_logger(mol, mol.verbose) - t1 = log.init_timer() - gridslice = pcmobj.surface['gslice_by_atom'] - charge_exp = pcmobj.surface['charge_exp'] - grid_coords = pcmobj.surface['grid_coords'] - if q_sym is None: - q_sym = pcmobj._intermediates['q_sym'] - - intopt = int3c1e.VHFOpt(mol) - intopt.build(1e-14, aosym=False) - dvj = int1e_grids_ip1(mol, grid_coords, dm = dm, charges = q_sym, - direct_scf_tol = 1e-14, charge_exponents = charge_exp**2, - intopt=intopt) - dq = int1e_grids_ip2(mol, grid_coords, dm = dm, charges = q_sym, - direct_scf_tol = 1e-14, charge_exponents = charge_exp**2, - intopt=intopt) - - aoslice = mol.aoslice_by_atom() - dvj = 2.0 * cp.asarray([cp.sum(dvj[:,p0:p1], axis=1) for p0,p1 in aoslice[:,2:]]) - dq = cp.asarray([cp.sum(dq[:,p0:p1], axis=1) for p0,p1 in gridslice]) - de = dq + dvj - t1 = log.timer_debug1('grad qv', *t1) - return de.get() def grad_solver(pcmobj, dm): ''' @@ -268,8 +233,8 @@ def __init__(self, tda_grad_method): self.__dict__.update(tda_grad_method.__dict__) tda_grad_method.base._scf.with_solvent.tdscf = True - def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): - de = super().grad_elec(self, xy, singlet, atmlst, verbose) + # def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): + # de = super().grad_elec(self, xy, singlet, atmlst, verbose) def _finalize(self): From 02b8a038b9800133e64161d64a71172e08986025 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 3 Apr 2025 11:48:38 +0800 Subject: [PATCH 09/31] Finish tdgrad with solvent. TODO: 1. specify epsilon_optical --- gpu4pyscf/grad/tdrhf.py | 10 ++++++++ gpu4pyscf/solvent/grad/pcm.py | 2 +- gpu4pyscf/solvent/pcm.py | 26 ++++++++++++++++--- gpu4pyscf/solvent/tdscf/pcm.py | 47 ++++++++++++++++++++++++++-------- 4 files changed, 70 insertions(+), 15 deletions(-) diff --git a/gpu4pyscf/grad/tdrhf.py b/gpu4pyscf/grad/tdrhf.py index 736209137..66f4c5372 100644 --- a/gpu4pyscf/grad/tdrhf.py +++ b/gpu4pyscf/grad/tdrhf.py @@ -60,6 +60,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): dmxmy = reduce(cp.dot, (orbv, xmy, orbo.T)) # (X-Y) in ao basis dmzoo = reduce(cp.dot, (orbo, doo, orbo.T)) # T_{ij}*2 in ao basis dmzoo += reduce(cp.dot, (orbv, dvv, orbv.T)) # T_{ij}*2 + T_{ab}*2 in ao basis + td_grad.dmxpy = dmxpy vj0, vk0 = mf.get_jk(mol, dmzoo, hermi=0) vj1, vk1 = mf.get_jk(mol, dmxpy + dmxpy.T, hermi=0) @@ -79,11 +80,15 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): vj = cp.stack((vj0, vj1, vj2)) vk = cp.stack((vk0, vk1, vk2)) veff0doo = vj[0] * 2 - vk[0] # 2 for alpha and beta + if getattr(mf, 'with_solvent', None): + veff0doo += mf.with_solvent._B_dot_x(dmzoo)* 2.0 wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2 if singlet: veff = vj[1] * 2 - vk[1] else: veff = -vk[1] + if getattr(mf, 'with_solvent', None): + veff += mf.with_solvent._B_dot_x((dmxpy + dmxpy.T)*1.0)*2.0 veff0mop = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff)) wvo -= contract("ki,ai->ak", veff0mop[:nocc, :nocc], xpy) * 2 # 2 for dm + dm.T wvo += contract("ac,ai->ci", veff0mop[nocc:, nocc:], xpy) * 2 @@ -149,6 +154,7 @@ def fvind(x): # For singlet, closed shell ground state s1 = mf_grad.get_ovlp(mol) dmz1doo = z1ao + dmzoo # P + td_grad.dmz1doo = dmz1doo oo0 = reduce(cp.dot, (orbo, orbo.T)) # D if atmlst is None: @@ -222,6 +228,8 @@ class Gradients(rhf_grad.GradientsBase): "state", "atmlst", "de", + "dmz1doo", + "dmxpy" } def __init__(self, td): @@ -234,6 +242,8 @@ def __init__(self, td): self.state = 1 # of which the gradients to be computed. self.atmlst = None self.de = None + self.dmz1doo = None + self.dmxpy = None def dump_flags(self, verbose=None): log = logger.new_logger(self, verbose) diff --git a/gpu4pyscf/solvent/grad/pcm.py b/gpu4pyscf/solvent/grad/pcm.py index 3483fd33c..98aec3a4c 100644 --- a/gpu4pyscf/solvent/grad/pcm.py +++ b/gpu4pyscf/solvent/grad/pcm.py @@ -278,7 +278,7 @@ def grad_qv(pcmobj, dm, q_sym = None): def grad_solver(pcmobj, dm): ''' dE = 0.5*v* d(K^-1 R) *v + q*dv - v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR - dK q) + v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR v - dK q) ''' mol = pcmobj.mol log = logger.new_logger(mol, mol.verbose) diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 9ea7102ab..748c3e2fb 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -400,7 +400,7 @@ def _get_vind(self, dms): self._intermediates['v_grids'] = v_grids[0] return epcm, vmat[0] - def _get_qsym(self, dms): + def _get_qsym(self, dms, with_nuc = False): if not self._intermediates: self.build() nao = dms.shape[-1] @@ -410,7 +410,10 @@ def _get_qsym(self, dms): if not isinstance(dms, cupy.ndarray): dms = cupy.asarray(dms) v_grids_e = self._get_v(dms) - v_grids = self.v_grids_n - v_grids_e + if with_nuc: + v_grids = self.v_grids_n - v_grids_e + else: + v_grids = -1.0 * v_grids_e b = self.left_multiply_R(v_grids.T) q = self.left_solve_K(b).T @@ -419,7 +422,24 @@ def _get_qsym(self, dms): qt = self.left_multiply_R(vK_1, R_transpose = True).T q_sym = (q + qt)/2.0 - return q_sym[0] + return q_sym[0], q[0] + + def _get_vgrids(self, dms, with_nuc = False): + if not self._intermediates: + self.build() + nao = dms.shape[-1] + dms = dms.reshape(-1,nao,nao) + if dms.shape[0] == 2: + dms = (dms[0] + dms[1]).reshape(-1,nao,nao) + if not isinstance(dms, cupy.ndarray): + dms = cupy.asarray(dms) + v_grids_e = self._get_v(dms) + if with_nuc: + v_grids = self.v_grids_n - v_grids_e + else: + v_grids = -1.0 * v_grids_e + + return v_grids[0] def _get_v(self, dms): ''' diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 32a6a0839..f6691e4da 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -19,9 +19,9 @@ import cupy as cp from pyscf import lib from gpu4pyscf.solvent.pcm import PI, switch_h, libsolvent +from gpu4pyscf.solvent.grad.pcm import grad_nuc, grad_qv, grad_solver from gpu4pyscf.solvent.grad.pcm import left_multiply_dS, right_multiply_dS, get_dF_dA, get_dSii, get_dD_dS -from gpu4pyscf.gto.int3c1e_ip import int1e_grids_ip1, int1e_grids_ip2 -from gpu4pyscf.lib.cupy_helper import contract, tag_array +from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.lib import logger from gpu4pyscf import scf from gpu4pyscf.gto import int3c1e @@ -65,10 +65,10 @@ def nuc_grad_method(self): return make_tdscf_gradient_object(grad_method) -def grad_solver(pcmobj, dm): +def grad_solver_td(pcmobj, dm, dm1, with_nuc_v = False, with_nuc_q = False): ''' dE = 0.5*v* d(K^-1 R) *v + q*dv - v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR - dK q) + v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR v - dK q) ''' mol = pcmobj.mol log = logger.new_logger(mol, mol.verbose) @@ -82,15 +82,17 @@ def grad_solver(pcmobj, dm): pcmobj._get_vind(dm) gridslice = pcmobj.surface['gslice_by_atom'] - v_grids = pcmobj._intermediates['v_grids'] - q = pcmobj._intermediates['q'] f_epsilon = pcmobj._intermediates['f_epsilon'] if not pcmobj.if_method_in_CPCM_category: A = pcmobj._intermediates['A'] D = pcmobj._intermediates['D'] S = pcmobj._intermediates['S'] - vK_1 = pcmobj.left_solve_K(v_grids, K_transpose = True) + v_grids = pcmobj._get_vgrids(dm, with_nuc_q) + q = pcmobj._get_qsym(dm, with_nuc_q)[1] + v_grids_1 = pcmobj._get_vgrids(dm1, with_nuc_v) + + vK_1 = pcmobj.left_solve_K(v_grids_1, K_transpose = True) def contract_bra(a, B, c): ''' i,xij,j->jx ''' @@ -233,11 +235,34 @@ def __init__(self, tda_grad_method): self.__dict__.update(tda_grad_method.__dict__) tda_grad_method.base._scf.with_solvent.tdscf = True - # def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): - # de = super().grad_elec(self, xy, singlet, atmlst, verbose) - + def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): + de = super().grad_elec(xy, singlet, atmlst, verbose) + if self.base._scf.with_solvent.tdscf and not self.base._scf.with_solvent.equilibrium_solvation: + dm = self.base._scf.make_rdm1(ao_repr=True) + if dm.ndim == 3: + dm = dm[0] + dm[1] + # TODO: add unrestricted case support + dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T) + dmxpy = 1.0 * (self.dmxpy + self.dmxpy.T) + pcmobj = self.base._scf.with_solvent + de += grad_qv(pcmobj, dm) + de += grad_solver(pcmobj, dm) + de += grad_nuc(pcmobj, dm) + q_sym_dm = pcmobj._get_qsym(dm, with_nuc = True)[0] + qE_sym_dmP = pcmobj._get_qsym(dmP)[0] + qE_sym_dmxpy = pcmobj._get_qsym(dmxpy)[0] + de += grad_qv(pcmobj, dm, q_sym = qE_sym_dmP) + de += grad_nuc(pcmobj, dm, q_sym = qE_sym_dmP.get()) + de += grad_qv(pcmobj, dmP, q_sym = q_sym_dm) + de += grad_solver_td(pcmobj, dm, dmP, with_nuc_q = True) * 2 + de += grad_qv(pcmobj, dmxpy, q_sym = qE_sym_dmxpy) * 2.0 + de += grad_solver_td(pcmobj, dmxpy, dmxpy) * 2 + else: + raise NotImplementedError("State-specific approach is not supported TDDFT gradient") + + return de def _finalize(self): super()._finalize() - self._scf.with_solvent.tdscf = False + self.base._scf.with_solvent.tdscf = False From 0b2deca4a03e7c069af9283b2550c7a5c12fcefb Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Tue, 8 Apr 2025 14:14:20 +0800 Subject: [PATCH 10/31] finish writting the TDDFT gradient. More tests should be done --- gpu4pyscf/grad/tdrhf.py | 2 + gpu4pyscf/grad/tdrks.py | 2 + gpu4pyscf/solvent/_attach_solvent.py | 15 +++--- gpu4pyscf/solvent/grad/pcm.py | 12 +++-- gpu4pyscf/solvent/pcm.py | 43 +--------------- gpu4pyscf/solvent/tdscf/pcm.py | 51 +++++++++---------- gpu4pyscf/solvent/tests/test_pcm_tdscf.py | 31 ++++++----- .../solvent/tests/test_pcm_tdscf_grad.py | 20 ++++++++ 8 files changed, 81 insertions(+), 95 deletions(-) create mode 100644 gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py diff --git a/gpu4pyscf/grad/tdrhf.py b/gpu4pyscf/grad/tdrhf.py index 66f4c5372..31f2c6992 100644 --- a/gpu4pyscf/grad/tdrhf.py +++ b/gpu4pyscf/grad/tdrhf.py @@ -37,6 +37,8 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): TDDFT X and Y amplitudes. If Y is set to 0, this function computes TDA energy gradients. """ + if singlet is None: + singlet = True log = logger.new_logger(td_grad, verbose) time0 = logger.init_timer(td_grad) mol = td_grad.mol diff --git a/gpu4pyscf/grad/tdrks.py b/gpu4pyscf/grad/tdrks.py index 97fc0c7c6..842b15eb8 100644 --- a/gpu4pyscf/grad/tdrks.py +++ b/gpu4pyscf/grad/tdrks.py @@ -43,6 +43,8 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): TDDFT X and Y amplitudes. If Y is set to 0, this function computes TDA energy gradients. """ + if singlet is None: + singlet = True log = logger.new_logger(td_grad, verbose) time0 = logger.init_timer(td_grad) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index f7f81716a..53e39e08a 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -164,16 +164,13 @@ def gen_response(self, *args, **kwargs): def vind_with_solvent(dm1): v = vind(dm1) if self.with_solvent.tdscf: - if not self.with_solvent.equilibrium_solvation: - if is_uhf: - v_solvent = self.with_solvent._B_dot_x(dm1) - v += v_solvent[0] + v_solvent[1] - elif singlet: - v += self.with_solvent._B_dot_x(dm1) - else: - logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') + if is_uhf: + v_solvent = self.with_solvent._B_dot_x(dm1) + v += v_solvent[0] + v_solvent[1] + elif singlet: + v += self.with_solvent._B_dot_x(dm1) else: - raise NotImplementedError('Equilibrium solvation TDDFT (SS) is not implemented!') + logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') else: if self.with_solvent.equilibrium_solvation: if is_uhf: diff --git a/gpu4pyscf/solvent/grad/pcm.py b/gpu4pyscf/solvent/grad/pcm.py index 98aec3a4c..dd1297aab 100644 --- a/gpu4pyscf/solvent/grad/pcm.py +++ b/gpu4pyscf/solvent/grad/pcm.py @@ -275,7 +275,7 @@ def grad_qv(pcmobj, dm, q_sym = None): t1 = log.timer_debug1('grad qv', *t1) return de.get() -def grad_solver(pcmobj, dm): +def grad_solver(pcmobj, dm, v_grids = None, v_grids_1 = None, q = None): ''' dE = 0.5*v* d(K^-1 R) *v + q*dv v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR v - dK q) @@ -292,15 +292,19 @@ def grad_solver(pcmobj, dm): pcmobj._get_vind(dm) gridslice = pcmobj.surface['gslice_by_atom'] - v_grids = pcmobj._intermediates['v_grids'] - q = pcmobj._intermediates['q'] + if v_grids is None: + v_grids = pcmobj._intermediates['v_grids'] + if v_grids_1 is None: + v_grids_1 = pcmobj._intermediates['v_grids'] + if q is None: + q = pcmobj._intermediates['q'] f_epsilon = pcmobj._intermediates['f_epsilon'] if not pcmobj.if_method_in_CPCM_category: A = pcmobj._intermediates['A'] D = pcmobj._intermediates['D'] S = pcmobj._intermediates['S'] - vK_1 = pcmobj.left_solve_K(v_grids, K_transpose = True) + vK_1 = pcmobj.left_solve_K(v_grids_1, K_transpose = True) def contract_bra(a, B, c): ''' i,xij,j->jx ''' diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 748c3e2fb..d633c9d7e 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -246,7 +246,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'eps_optical', 'tdscf' + 'equilibrium_solvation', 'e', 'v', 'tdscf' } from gpu4pyscf.lib.utils import to_gpu, device kernel = ddcosmo.DDCOSMO.kernel @@ -266,7 +266,6 @@ def __init__(self, mol): self.lebedev_order = 29 self._intermediates = {} self.eps = 78.3553 - self.eps_optical = 1.78 self.tdscf = False self.max_cycle = 20 @@ -515,46 +514,6 @@ def _B_dot_x(self, dms): out_shape = dms.shape nao = dms.shape[-1] dms = dms.reshape(-1,nao,nao) - - if self.tdscf: - assert not self.equilibrium_solvation - epsilon = self.eps_optical - logger.info(self, 'eps optical = %s', self.eps_optical) - F, A = get_F_A(self.surface) - D, S = get_D_S(self.surface, with_S = True, with_D = not self.if_method_in_CPCM_category) - epsilon = self.eps_optical - if self.method.upper() in ['C-PCM', 'CPCM']: - f_epsilon = (epsilon-1.)/epsilon - K = S - elif self.method.upper() == 'COSMO': - f_epsilon = (epsilon - 1.0)/(epsilon + 1.0/2.0) - K = S - elif self.method.upper() in ['IEF-PCM', 'IEFPCM']: - f_epsilon = (epsilon - 1.0)/(epsilon + 1.0) - DA = D*A - DAS = cupy.dot(DA, S) - K = S - f_epsilon/(2.0*PI) * DAS - elif self.method.upper() == 'SS(V)PE': - f_epsilon = (epsilon - 1.0)/(epsilon + 1.0) - DA = D*A - DAS = cupy.dot(DA, S) - K = S - f_epsilon/(4.0*PI) * (DAS + DAS.T) - else: - raise RuntimeError(f"Unknown implicit solvent model: {self.method}") - - K_LU, K_LU_pivot = lu_factor(K, overwrite_a = True, check_finite = False) - K = None - v_grids = -self._get_v(dms) - - b = self.left_multiply_R(v_grids.T, f_epsilon = f_epsilon) - q = self.left_solve_K(b, K_LU=K_LU, K_LU_pivot=K_LU_pivot).T - - vK_1 = self.left_solve_K(v_grids.T, K_LU=K_LU, K_LU_pivot=K_LU_pivot, K_transpose = True) - qt = self.left_multiply_R(vK_1, f_epsilon = f_epsilon, R_transpose = True).T - q_sym = (q + qt)/2.0 - - vmat = self._get_vmat(q_sym) - return vmat.reshape(out_shape) v_grids = -self._get_v(dms) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index f6691e4da..bbe8427f2 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -48,7 +48,6 @@ class WithSolventTDSCF: def __init__(self, tda_method): self.__dict__.update(tda_method.__dict__) - tda_method._scf.with_solvent.tdscf = True def undo_solvent(self): cls = self.__class__ @@ -58,7 +57,6 @@ def undo_solvent(self): def _finalize(self): super()._finalize() - self._scf.with_solvent.tdscf = False def nuc_grad_method(self): grad_method = super().nuc_grad_method() @@ -233,36 +231,37 @@ class WithSolventTDSCFGradient: def __init__(self, tda_grad_method): self.__dict__.update(tda_grad_method.__dict__) - tda_grad_method.base._scf.with_solvent.tdscf = True def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): de = super().grad_elec(xy, singlet, atmlst, verbose) - if self.base._scf.with_solvent.tdscf and not self.base._scf.with_solvent.equilibrium_solvation: - dm = self.base._scf.make_rdm1(ao_repr=True) - if dm.ndim == 3: - dm = dm[0] + dm[1] - # TODO: add unrestricted case support - dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T) - dmxpy = 1.0 * (self.dmxpy + self.dmxpy.T) - pcmobj = self.base._scf.with_solvent - de += grad_qv(pcmobj, dm) - de += grad_solver(pcmobj, dm) - de += grad_nuc(pcmobj, dm) - q_sym_dm = pcmobj._get_qsym(dm, with_nuc = True)[0] - qE_sym_dmP = pcmobj._get_qsym(dmP)[0] - qE_sym_dmxpy = pcmobj._get_qsym(dmxpy)[0] - de += grad_qv(pcmobj, dm, q_sym = qE_sym_dmP) - de += grad_nuc(pcmobj, dm, q_sym = qE_sym_dmP.get()) - de += grad_qv(pcmobj, dmP, q_sym = q_sym_dm) - de += grad_solver_td(pcmobj, dm, dmP, with_nuc_q = True) * 2 - de += grad_qv(pcmobj, dmxpy, q_sym = qE_sym_dmxpy) * 2.0 - de += grad_solver_td(pcmobj, dmxpy, dmxpy) * 2 - else: - raise NotImplementedError("State-specific approach is not supported TDDFT gradient") + + assert self.base._scf.with_solvent.equilibrium_solvation + + dm = self.base._scf.make_rdm1(ao_repr=True) + # TODO: add unrestricted case support + dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T) + dmxpy = 1.0 * (self.dmxpy + self.dmxpy.T) + pcmobj = self.base._scf.with_solvent + de += grad_qv(pcmobj, dm) + de += grad_solver(pcmobj, dm) + de += grad_nuc(pcmobj, dm) + + q_sym_dm = pcmobj._get_qsym(dm, with_nuc = True)[0] + qE_sym_dmP = pcmobj._get_qsym(dmP)[0] + qE_sym_dmxpy = pcmobj._get_qsym(dmxpy)[0] + de += grad_qv(pcmobj, dm, q_sym = qE_sym_dmP) + de += grad_nuc(pcmobj, dm, q_sym = qE_sym_dmP.get()) + de += grad_qv(pcmobj, dmP, q_sym = q_sym_dm) + v_grids_1 = pcmobj._get_vgrids(dmP, with_nuc = False) + de += grad_solver(pcmobj, dm, v_grids_1 = v_grids_1) * 2.0 + de += grad_qv(pcmobj, dmxpy, q_sym = qE_sym_dmxpy) * 2.0 + v_grids = pcmobj._get_vgrids(dmxpy, with_nuc = False) + q = pcmobj._get_qsym(dmxpy, with_nuc = False)[1] + v_grids_1 = pcmobj._get_vgrids(dmxpy, with_nuc = False) + de += grad_solver(pcmobj, dmxpy, v_grids=v_grids, v_grids_1=v_grids_1, q=q) * 2.0 return de def _finalize(self): super()._finalize() - self.base._scf.with_solvent.tdscf = False diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py index 823e2ffdf..9246b2cfa 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py @@ -36,14 +36,12 @@ def setUpClass(cls): cls.mf.with_solvent.method = 'C-PCM' cls.mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids cls.mf.with_solvent.eps = 78 - cls.mf.with_solvent.eps_optical = 1.78 cls.mf = mf.run(conv_tol=1e-10) cls.mfu = mfu = mol.UHF().PCM().to_gpu() cls.mfu.with_solvent.method = 'C-PCM' cls.mfu.with_solvent.lebedev_order = 29 # 302 Lebedev grids cls.mfu.with_solvent.eps = 78 - cls.mfu.with_solvent.eps_optical = 1.78 cls.mfu = mfu.run(conv_tol=1e-10) mf_b3lyp_nodf = mol.RKS().PCM().to_gpu() @@ -52,7 +50,6 @@ def setUpClass(cls): mf_b3lyp_nodf.with_solvent.method = 'C-PCM' mf_b3lyp_nodf.with_solvent.lebedev_order = 29 # 302 Lebedev grids mf_b3lyp_nodf.with_solvent.eps = 78 - mf_b3lyp_nodf.with_solvent.eps_optical = 1.78 mf_b3lyp_nodf.cphf_grids = mf_b3lyp_nodf.grids cls.mf_b3lyp_nodf = mf_b3lyp_nodf.run(conv_tol=1e-10) @@ -62,7 +59,6 @@ def setUpClass(cls): mf_b3lyp_nodf_u.with_solvent.method = 'C-PCM' mf_b3lyp_nodf_u.with_solvent.lebedev_order = 29 # 302 Lebedev grids mf_b3lyp_nodf_u.with_solvent.eps = 78 - mf_b3lyp_nodf_u.with_solvent.eps_optical = 1.78 mf_b3lyp_nodf_u.cphf_grids = mf_b3lyp_nodf_u.grids cls.mf_b3lyp_nodf_u = mf_b3lyp_nodf_u.run(conv_tol=1e-10) @@ -72,7 +68,6 @@ def setUpClass(cls): mf_b3lyp_nodf_iefpcm.with_solvent.method = 'IEF-PCM' mf_b3lyp_nodf_iefpcm.with_solvent.lebedev_order = 29 # 302 Lebedev grids mf_b3lyp_nodf_iefpcm.with_solvent.eps = 78 - mf_b3lyp_nodf_iefpcm.with_solvent.eps_optical = 1.78 mf_b3lyp_nodf_iefpcm.cphf_grids = mf_b3lyp_nodf_iefpcm.grids cls.mf_b3lyp_nodf_iefpcm = mf_b3lyp_nodf_iefpcm.run(conv_tol=1e-10) @@ -82,7 +77,6 @@ def setUpClass(cls): mf_b3lyp_nodf_iefpcm_u.with_solvent.method = 'IEF-PCM' mf_b3lyp_nodf_iefpcm_u.with_solvent.lebedev_order = 29 # 302 Lebedev grids mf_b3lyp_nodf_iefpcm_u.with_solvent.eps = 78 - mf_b3lyp_nodf_iefpcm_u.with_solvent.eps_optical = 1.78 mf_b3lyp_nodf_iefpcm_u.cphf_grids = mf_b3lyp_nodf_iefpcm_u.grids cls.mf_b3lyp_nodf_iefpcm_u = mf_b3lyp_nodf_iefpcm.run(conv_tol=1e-10) @@ -118,30 +112,34 @@ def test_hf_CPCM(self): """ mf = self.mf td = mf.TDHF() - assert td.device == 'gpu' + td._scf.with_solvent.tdscf = True + td._scf.with_solvent.eps = 1.78 + td._scf.with_solvent.build() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot + print(es_gound) ref = np.array([-75.61072291, -75.54419399, -75.51949191, -75.45219025, -75.40975027]) assert np.allclose(es_gound, ref) td = mf.TDA() - assert td.device == 'gpu' es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot + print(es_gound) ref = np.array([-75.60864828, -75.54169327, -75.51738767, -75.44915784, -75.40839714]) assert np.allclose(es_gound, ref) def test_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf td = mf.TDDFT() - assert td.device == 'gpu' + td._scf.with_solvent.tdscf = True + td._scf.with_solvent.eps = 1.78 + td._scf.with_solvent.build() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06898428, -75.99630982, -75.98765186, -75.91045133, -75.84783748]) assert np.allclose(es_gound, ref) td = mf.TDA() - assert td.device == 'gpu' es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06789176, -75.99609709, -75.98589720, -75.90894600, -75.84699115]) @@ -150,14 +148,15 @@ def test_b3lyp_CPCM(self): def test_b3lyp_IEFPCM(self): mf = self.mf_b3lyp_nodf_iefpcm td = mf.TDDFT() - assert td.device == 'gpu' + td._scf.with_solvent.tdscf = True + td._scf.with_solvent.eps = 1.78 + td._scf.with_solvent.build() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06881645, -75.99631929, -75.98713725, -75.91015704, -75.84668800]) assert np.allclose(es_gound, ref) td = mf.TDA() - assert td.device == 'gpu' es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06773319, -75.99610928, -75.98534912, -75.90861455, -75.84576041]) @@ -166,7 +165,9 @@ def test_b3lyp_IEFPCM(self): def test_unrestricted_hf_CPCM(self): mf = self.mfu td = mf.TDHF() - assert td.device == 'gpu' + td._scf.with_solvent.tdscf = True + td._scf.with_solvent.eps = 1.78 + td._scf.with_solvent.build() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-75.64482315, -75.61072291, -75.57156784, -75.56769949, -75.54419399]) @@ -175,7 +176,9 @@ def test_unrestricted_hf_CPCM(self): def test_unrestricted_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf_u td = mf.TDDFT() - assert td.device == 'gpu' + td._scf.with_solvent.tdscf = True + td._scf.with_solvent.eps = 1.78 + td._scf.with_solvent.build() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.09301571, -76.06898428, -76.01822101, -76.01369024, -75.99630982]) diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py new file mode 100644 index 000000000..8a78f41fa --- /dev/null +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py @@ -0,0 +1,20 @@ +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import unittest +import numpy as np +import cupy as cp +from pyscf import lib, gto +from gpu4pyscf.tdscf import rhf, rks +from gpu4pyscf import tdscf \ No newline at end of file From 0edc465b82dcfd42acadebf167d0872212273828 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Tue, 8 Apr 2025 16:32:58 +0800 Subject: [PATCH 11/31] add example and unit test --- examples/27-tddft_with_solvent.py | 63 +++++++++++++++++++++++++++++++ gpu4pyscf/grad/tdrhf.py | 9 +++-- gpu4pyscf/grad/tdrks.py | 4 ++ gpu4pyscf/grad/tduhf.py | 4 ++ gpu4pyscf/grad/tduks.py | 4 ++ gpu4pyscf/solvent/tdscf/pcm.py | 7 +++- 6 files changed, 86 insertions(+), 5 deletions(-) create mode 100644 examples/27-tddft_with_solvent.py diff --git a/examples/27-tddft_with_solvent.py b/examples/27-tddft_with_solvent.py new file mode 100644 index 000000000..1cd599f71 --- /dev/null +++ b/examples/27-tddft_with_solvent.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +################################### +# Example of TDDFT +################################### + +import pyscf +import gpu4pyscf +from gpu4pyscf import tdscf + + +atom =''' +O 0.0000000000 -0.0000000000 0.1174000000 +H -0.7570000000 -0.0000000000 -0.4696000000 +H 0.7570000000 0.0000000000 -0.4696000000 +''' + +mol = pyscf.M(atom=atom, basis='def2-tzvpp') + +xc = 'b3lyp' +mf = gpu4pyscf.dft.RKS(mol, xc=xc) +mf.grids.level = 5 +mf.kernel() # -76.4666495331835 + +# Compute TDDFT and TDA excitation energy +print('------------------- TDDFT -----------------------------') +td = mf.TDDFT().set(nstates=5) +assert td.device == 'gpu' +e_tddft = td.kernel()[0] # [ 7.51062554 9.42244962 9.76602191 11.74385565 13.59746453] +# print('5 TDDFT excitation energy by GPU4PySCF') +# print(e_tddft) + +print('------------------- TDA -----------------------------') +td = mf.TDA().set(nstates=5) +assert td.device == 'gpu' +e_tda = td.kernel()[0] # [ 7.53381449 9.42805412 9.81321061 11.78177395 13.62814798] +# print('5 TDA excitation energy by GPU4PySCF') +# print(e_tda) + +print('The gradient of first TDA excitation energy by GPU4PySCF') +g = td.nuc_grad_method() +g.kernel() +""" +--------- TDA gradients for state 1 ---------- + x y z +0 O -0.0000000000 -0.0000000000 -0.0901095287 +1 H 0.0598456498 -0.0000000000 0.0450549873 +2 H -0.0598456498 0.0000000000 0.0450549873 +---------------------------------------------- +""" \ No newline at end of file diff --git a/gpu4pyscf/grad/tdrhf.py b/gpu4pyscf/grad/tdrhf.py index 31f2c6992..e12f57d38 100644 --- a/gpu4pyscf/grad/tdrhf.py +++ b/gpu4pyscf/grad/tdrhf.py @@ -82,15 +82,13 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): vj = cp.stack((vj0, vj1, vj2)) vk = cp.stack((vk0, vk1, vk2)) veff0doo = vj[0] * 2 - vk[0] # 2 for alpha and beta - if getattr(mf, 'with_solvent', None): - veff0doo += mf.with_solvent._B_dot_x(dmzoo)* 2.0 + veff0doo += td_grad.solvent_response(dmzoo) wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2 if singlet: veff = vj[1] * 2 - vk[1] else: veff = -vk[1] - if getattr(mf, 'with_solvent', None): - veff += mf.with_solvent._B_dot_x((dmxpy + dmxpy.T)*1.0)*2.0 + veff += td_grad.solvent_response(dmxpy + dmxpy.T) veff0mop = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff)) wvo -= contract("ki,ai->ak", veff0mop[:nocc, :nocc], xpy) * 2 # 2 for dm + dm.T wvo += contract("ac,ai->ci", veff0mop[nocc:, nocc:], xpy) * 2 @@ -344,6 +342,9 @@ def _finalize(self): self._write(self.mol, self.de, self.atmlst) logger.note(self, "----------------------------------------------") + def solvent_response(self, dm): + return 0.0 + as_scanner = NotImplemented to_gpu = lib.to_gpu diff --git a/gpu4pyscf/grad/tdrks.py b/gpu4pyscf/grad/tdrks.py index 842b15eb8..7e48302c9 100644 --- a/gpu4pyscf/grad/tdrks.py +++ b/gpu4pyscf/grad/tdrks.py @@ -69,6 +69,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): dmxmy = reduce(cp.dot, (orbv, xmy, orbo.T)) # (X-Y) in ao basis dmzoo = reduce(cp.dot, (orbo, doo, orbo.T)) # T_{ij}*2 in ao basis dmzoo += reduce(cp.dot, (orbv, dvv, orbv.T)) # T_{ij}*2 + T_{ab}*2 in ao basis + td_grad.dmxpy = dmxpy ni = mf._numint ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True) @@ -106,11 +107,13 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): vk2 = cp.asarray(vk2) vk += cp.stack((vk0, vk1, vk2)) * (alpha - hyb) veff0doo = vj[0] * 2 - vk[0] + f1oo[0] + k1ao[0] * 2 + veff0doo += td_grad.solvent_response(dmzoo) wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2 if singlet: veff = vj[1] * 2 - vk[1] + f1vo[0] * 2 else: veff = f1vo[0] - vk[1] + veff += td_grad.solvent_response(dmxpy + dmxpy.T) veff0mop = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff)) wvo -= contract("ki,ai->ak", veff0mop[:nocc, :nocc], xpy) * 2 wvo += contract("ac,ai->ci", veff0mop[nocc:, nocc:], xpy) * 2 @@ -184,6 +187,7 @@ def fvind(x): s1 = mf_grad.get_ovlp(mol) dmz1doo = z1ao + dmzoo + td_grad.dmz1doo = dmz1doo oo0 = reduce(cp.dot, (orbo, orbo.T)) if atmlst is None: diff --git a/gpu4pyscf/grad/tduhf.py b/gpu4pyscf/grad/tduhf.py index 932e48568..7b9f49e90 100644 --- a/gpu4pyscf/grad/tduhf.py +++ b/gpu4pyscf/grad/tduhf.py @@ -85,6 +85,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): dmzoob = reduce(cp.dot, (orbob, doob, orbob.T)) dmzooa += reduce(cp.dot, (orbva, dvva, orbva.T)) dmzoob += reduce(cp.dot, (orbvb, dvvb, orbvb.T)) + td_grad.dmxpy = (dmxpya + dmxpyb)*0.5 vj0, vk0 = mf.get_jk(mol, cp.stack((dmzooa, dmzoob)), hermi=0) vj1, vk1 = mf.get_jk(mol, cp.stack((dmxpya + dmxpya.T, dmxpyb + dmxpyb.T)), hermi=0) @@ -107,9 +108,11 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): vk = vk.reshape(2, 3, nao, nao) veff0doo = vj[0, 0] + vj[1, 0] - vk[:, 0] + veff0doo += td_grad.solvent_response((dmzooa + dmzoob)*0.5) wvoa = reduce(cp.dot, (orbva.T, veff0doo[0], orboa)) * 2 wvob = reduce(cp.dot, (orbvb.T, veff0doo[1], orbob)) * 2 veff = vj[0, 1] + vj[1, 1] - vk[:, 1] + veff += td_grad.solvent_response((td_grad.dmxpy + td_grad.dmxpy.T)) veff0mopa = reduce(cp.dot, (mo_coeff[0].T, veff[0], mo_coeff[0])) veff0mopb = reduce(cp.dot, (mo_coeff[1].T, veff[1], mo_coeff[1])) wvoa -= contract("ki,ai->ak", veff0mopa[:nocca, :nocca], xpya) * 2 @@ -192,6 +195,7 @@ def fvind(x): dmz1dooa = z1ao[0] + dmzooa dmz1doob = z1ao[1] + dmzoob + td_grad.dmz1doo = (dmz1dooa + dmz1doob)*0.5 oo0a = reduce(cp.dot, (orboa, orboa.T)) oo0b = reduce(cp.dot, (orbob, orbob.T)) diff --git a/gpu4pyscf/grad/tduks.py b/gpu4pyscf/grad/tduks.py index a6f0efe39..b4c9a3d76 100644 --- a/gpu4pyscf/grad/tduks.py +++ b/gpu4pyscf/grad/tduks.py @@ -91,6 +91,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): dmzoob = reduce(cp.dot, (orbob, doob, orbob.T)) dmzooa += reduce(cp.dot, (orbva, dvva, orbva.T)) dmzoob += reduce(cp.dot, (orbvb, dvvb, orbvb.T)) + td_grad.dmxpy = (dmxpya + dmxpyb)*0.5 ni = mf._numint ni.libxc.test_deriv_order(mf.xc, 3, raise_error=True) @@ -136,9 +137,11 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): vk = vk.reshape(2, 3, nao, nao) veff0doo = vj[0, 0] + vj[1, 0] - vk[:, 0] + f1oo[:, 0] + k1ao[:, 0] * 2 + veff0doo += td_grad.solvent_response((dmzooa + dmzoob)*0.5) wvoa = reduce(cp.dot, (orbva.T, veff0doo[0], orboa)) * 2 wvob = reduce(cp.dot, (orbvb.T, veff0doo[1], orbob)) * 2 veff = vj[0, 1] + vj[1, 1] - vk[:, 1] + f1vo[:, 0] * 2 + veff += td_grad.solvent_response((td_grad.dmxpy + td_grad.dmxpy.T)) veff0mopa = reduce(cp.dot, (mo_coeff[0].T, veff[0], mo_coeff[0])) veff0mopb = reduce(cp.dot, (mo_coeff[1].T, veff[1], mo_coeff[1])) wvoa -= contract("ki,ai->ak", veff0mopa[:nocca, :nocca], xpya) * 2 @@ -242,6 +245,7 @@ def fvind(x): dmz1dooa = z1ao[0] + dmzooa dmz1doob = z1ao[1] + dmzoob + td_grad.dmz1doo = (dmz1dooa + dmz1doob)*0.5 oo0a = reduce(cp.dot, (orboa, orboa.T)) oo0b = reduce(cp.dot, (orbob, orbob.T)) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index bbe8427f2..a07279460 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -231,6 +231,9 @@ class WithSolventTDSCFGradient: def __init__(self, tda_grad_method): self.__dict__.update(tda_grad_method.__dict__) + + def solvent_response(self, dm): + return self.base._scf.with_solvent._B_dot_x(dm)*2.0 def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): de = super().grad_elec(xy, singlet, atmlst, verbose) @@ -238,9 +241,11 @@ def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): assert self.base._scf.with_solvent.equilibrium_solvation dm = self.base._scf.make_rdm1(ao_repr=True) + if dm.ndim == 3: + dm = dm[0] + dm[1] # TODO: add unrestricted case support dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T) - dmxpy = 1.0 * (self.dmxpy + self.dmxpy.T) + dmxpy = self.dmxpy + self.dmxpy.T pcmobj = self.base._scf.with_solvent de += grad_qv(pcmobj, dm) de += grad_solver(pcmobj, dm) From 3837033d313dcd8a3d3f6a58dcdc4eb0a9b57e98 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Wed, 9 Apr 2025 11:27:19 +0800 Subject: [PATCH 12/31] Finish write the gradient. --- examples/27-tddft_with_solvent.py | 32 ++- gpu4pyscf/solvent/grad/pcm.py | 2 +- gpu4pyscf/solvent/tdscf/pcm.py | 166 +------------- .../solvent/tests/test_pcm_tdscf_grad.py | 209 +++++++++++++++++- 4 files changed, 228 insertions(+), 181 deletions(-) diff --git a/examples/27-tddft_with_solvent.py b/examples/27-tddft_with_solvent.py index 1cd599f71..00232c05e 100644 --- a/examples/27-tddft_with_solvent.py +++ b/examples/27-tddft_with_solvent.py @@ -31,22 +31,30 @@ mol = pyscf.M(atom=atom, basis='def2-tzvpp') xc = 'b3lyp' -mf = gpu4pyscf.dft.RKS(mol, xc=xc) +mf = gpu4pyscf.dft.RKS(mol, xc=xc).PCM() +mf.with_solvent.method = 'IEFPCM' +mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids +mf.with_solvent.eps = 78 mf.grids.level = 5 -mf.kernel() # -76.4666495331835 +mf.kernel() # -76.476456106979 # Compute TDDFT and TDA excitation energy -print('------------------- TDDFT -----------------------------') +print('------------------- vertical exitation TDA -----------------------------') td = mf.TDDFT().set(nstates=5) -assert td.device == 'gpu' -e_tddft = td.kernel()[0] # [ 7.51062554 9.42244962 9.76602191 11.74385565 13.59746453] +td._scf.with_solvent.tdscf = True +td._scf.with_solvent.eps = 1.78 +td._scf.with_solvent.build() +e_tddft = td.kernel()[0] # [ 8.03553827 10.07361783 10.20203523 12.36009792 13.83374455] # print('5 TDDFT excitation energy by GPU4PySCF') # print(e_tddft) -print('------------------- TDA -----------------------------') +print('------------------- adiabatic excitation TDA -----------------------------') td = mf.TDA().set(nstates=5) -assert td.device == 'gpu' -e_tda = td.kernel()[0] # [ 7.53381449 9.42805412 9.81321061 11.78177395 13.62814798] +td._scf.with_solvent.tdscf = True +td._scf.with_solvent.eps = 78.0 +td._scf.with_solvent.equilibrium_solvation = True +td._scf.with_solvent.build() +e_tda = td.kernel()[0] # [ 7.99456759 10.0632959 10.08523494 12.30675282 13.64298125] # print('5 TDA excitation energy by GPU4PySCF') # print(e_tda) @@ -54,10 +62,10 @@ g = td.nuc_grad_method() g.kernel() """ ---------- TDA gradients for state 1 ---------- +--------- PCMTDA gradients for state 1 ---------- x y z -0 O -0.0000000000 -0.0000000000 -0.0901095287 -1 H 0.0598456498 -0.0000000000 0.0450549873 -2 H -0.0598456498 0.0000000000 0.0450549873 +0 O -0.0000000000 0.0000000000 -0.0836461430 +1 H 0.0601539533 -0.0000000000 0.0418232965 +2 H -0.0601539533 -0.0000000000 0.0418232965 ---------------------------------------------- """ \ No newline at end of file diff --git a/gpu4pyscf/solvent/grad/pcm.py b/gpu4pyscf/solvent/grad/pcm.py index dd1297aab..ddb300b54 100644 --- a/gpu4pyscf/solvent/grad/pcm.py +++ b/gpu4pyscf/solvent/grad/pcm.py @@ -297,7 +297,7 @@ def grad_solver(pcmobj, dm, v_grids = None, v_grids_1 = None, q = None): if v_grids_1 is None: v_grids_1 = pcmobj._intermediates['v_grids'] if q is None: - q = pcmobj._intermediates['q'] + q = pcmobj._intermediates['q'] f_epsilon = pcmobj._intermediates['f_epsilon'] if not pcmobj.if_method_in_CPCM_category: A = pcmobj._intermediates['A'] diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index a07279460..85ba75b09 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -63,169 +63,6 @@ def nuc_grad_method(self): return make_tdscf_gradient_object(grad_method) -def grad_solver_td(pcmobj, dm, dm1, with_nuc_v = False, with_nuc_q = False): - ''' - dE = 0.5*v* d(K^-1 R) *v + q*dv - v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR v - dK q) - ''' - mol = pcmobj.mol - log = logger.new_logger(mol, mol.verbose) - t1 = log.init_timer() - if not pcmobj._intermediates: - pcmobj.build() - dm_cache = pcmobj._intermediates.get('dm', None) - if dm_cache is not None and cp.linalg.norm(dm_cache - dm) < 1e-10: - pass - else: - pcmobj._get_vind(dm) - - gridslice = pcmobj.surface['gslice_by_atom'] - f_epsilon = pcmobj._intermediates['f_epsilon'] - if not pcmobj.if_method_in_CPCM_category: - A = pcmobj._intermediates['A'] - D = pcmobj._intermediates['D'] - S = pcmobj._intermediates['S'] - - v_grids = pcmobj._get_vgrids(dm, with_nuc_q) - q = pcmobj._get_qsym(dm, with_nuc_q)[1] - v_grids_1 = pcmobj._get_vgrids(dm1, with_nuc_v) - - vK_1 = pcmobj.left_solve_K(v_grids_1, K_transpose = True) - - def contract_bra(a, B, c): - ''' i,xij,j->jx ''' - tmp = a.dot(B) - return (tmp*c).T - - def contract_ket(a, B, c): - ''' i,xij,j->ix ''' - tmp = B.dot(c) - return (a*tmp).T - - de = cp.zeros([pcmobj.mol.natm,3]) - if pcmobj.method.upper() in ['C-PCM', 'CPCM', 'COSMO']: - # dR = 0, dK = dS - de_dS = 0.5 * vK_1.reshape(-1, 1) * left_multiply_dS(pcmobj.surface, q, stream=None) - de_dS -= 0.5 * q.reshape(-1, 1) * right_multiply_dS(pcmobj.surface, vK_1, stream=None) - de -= cp.asarray([cp.sum(de_dS[p0:p1], axis=0) for p0,p1 in gridslice]) - - dF, _ = get_dF_dA(pcmobj.surface, with_dA = False) - dSii = get_dSii(pcmobj.surface, dF) - de -= 0.5*contract('i,xij->jx', vK_1*q, dSii) # 0.5*cp.einsum('i,xij,i->jx', vK_1, dSii, q) - - elif pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SMD']: - dF, dA = get_dF_dA(pcmobj.surface) - dSii = get_dSii(pcmobj.surface, dF) - dF = None - - dD, dS = get_dD_dS(pcmobj.surface, with_D=True, with_S=True) - - # dR = f_eps/(2*pi) * (dD*A + D*dA), - # dK = dS - f_eps/(2*pi) * (dD*A*S + D*dA*S + D*A*dS) - fac = f_epsilon/(2.0*PI) - - Av = A*v_grids - de_dR = 0.5*fac * contract_ket(vK_1, dD, Av) - de_dR -= 0.5*fac * contract_bra(vK_1, dD, Av) - de_dR = cp.asarray([cp.sum(de_dR[p0:p1], axis=0) for p0,p1 in gridslice]) - - vK_1_D = vK_1.dot(D) - vK_1_Dv = vK_1_D * v_grids - de_dR += 0.5*fac * contract('j,xjn->nx', vK_1_Dv, dA) - - de_dS0 = 0.5*contract_ket(vK_1, dS, q) - de_dS0 -= 0.5*contract_bra(vK_1, dS, q) - de_dS0 = cp.asarray([cp.sum(de_dS0[p0:p1], axis=0) for p0,p1 in gridslice]) - - vK_1_q = vK_1 * q - de_dS0 += 0.5*contract('i,xin->nx', vK_1_q, dSii) - - vK_1_DA = vK_1_D*A - de_dS1 = 0.5*contract_ket(vK_1_DA, dS, q) - de_dS1 -= 0.5*contract_bra(vK_1_DA, dS, q) - de_dS1 = cp.asarray([cp.sum(de_dS1[p0:p1], axis=0) for p0,p1 in gridslice]) - - vK_1_DAq = vK_1_DA*q - de_dS1 += 0.5*contract('j,xjn->nx', vK_1_DAq, dSii) - - Sq = cp.dot(S,q) - ASq = A*Sq - de_dD = 0.5*contract_ket(vK_1, dD, ASq) - de_dD -= 0.5*contract_bra(vK_1, dD, ASq) - de_dD = cp.asarray([cp.sum(de_dD[p0:p1], axis=0) for p0,p1 in gridslice]) - - de_dA = 0.5*contract('j,xjn->nx', vK_1_D*Sq, dA) # 0.5*cp.einsum('j,xjn,j->nx', vK_1_D, dA, Sq) - - de_dK = de_dS0 - fac * (de_dD + de_dA + de_dS1) - de += de_dR - de_dK - - elif pcmobj.method.upper() in [ 'SS(V)PE' ]: - dF, dA = get_dF_dA(pcmobj.surface) - dSii = get_dSii(pcmobj.surface, dF) - dF = None - - dD, dS = get_dD_dS(pcmobj.surface, with_D=True, with_S=True) - - # dR = f_eps/(2*pi) * (dD*A + D*dA), - # dK = dS - f_eps/(2*pi) * (dD*A*S + D*dA*S + D*A*dS) - fac = f_epsilon/(2.0*PI) - - Av = A*v_grids - de_dR = 0.5*fac * contract_ket(vK_1, dD, Av) - de_dR -= 0.5*fac * contract_bra(vK_1, dD, Av) - de_dR = cp.asarray([cp.sum(de_dR[p0:p1], axis=0) for p0,p1 in gridslice]) - - vK_1_D = vK_1.dot(D) - vK_1_Dv = vK_1_D * v_grids - de_dR += 0.5*fac * contract('j,xjn->nx', vK_1_Dv, dA) - - de_dS0 = 0.5*contract_ket(vK_1, dS, q) - de_dS0 -= 0.5*contract_bra(vK_1, dS, q) - de_dS0 = cp.asarray([cp.sum(de_dS0[p0:p1], axis=0) for p0,p1 in gridslice]) - - vK_1_q = vK_1 * q - de_dS0 += 0.5*contract('i,xin->nx', vK_1_q, dSii) - - vK_1_DA = vK_1_D*A - de_dS1 = 0.5*contract_ket(vK_1_DA, dS, q) - de_dS1 -= 0.5*contract_bra(vK_1_DA, dS, q) - de_dS1 = cp.asarray([cp.sum(de_dS1[p0:p1], axis=0) for p0,p1 in gridslice]) - vK_1_DAq = vK_1_DA*q - de_dS1 += 0.5*contract('j,xjn->nx', vK_1_DAq, dSii) - - DT_q = cp.dot(D.T, q) - ADT_q = A * DT_q - de_dS1_T = 0.5*contract_ket(vK_1, dS, ADT_q) - de_dS1_T -= 0.5*contract_bra(vK_1, dS, ADT_q) - de_dS1_T = cp.asarray([cp.sum(de_dS1_T[p0:p1], axis=0) for p0,p1 in gridslice]) - vK_1_ADT_q = vK_1 * ADT_q - de_dS1_T += 0.5*contract('j,xjn->nx', vK_1_ADT_q, dSii) - - Sq = cp.dot(S,q) - ASq = A*Sq - de_dD = 0.5*contract_ket(vK_1, dD, ASq) - de_dD -= 0.5*contract_bra(vK_1, dD, ASq) - de_dD = cp.asarray([cp.sum(de_dD[p0:p1], axis=0) for p0,p1 in gridslice]) - - vK_1_S = cp.dot(vK_1, S) - vK_1_SA = vK_1_S * A - de_dD_T = 0.5*contract_ket(vK_1_SA, -dD.transpose(0,2,1), q) - de_dD_T -= 0.5*contract_bra(vK_1_SA, -dD.transpose(0,2,1), q) - de_dD_T = cp.asarray([cp.sum(de_dD_T[p0:p1], axis=0) for p0,p1 in gridslice]) - - de_dA = 0.5*contract('j,xjn->nx', vK_1_D*Sq, dA) # 0.5*cp.einsum('j,xjn,j->nx', vK_1_D, dA, Sq) - - de_dA_T = 0.5*contract('j,xjn->nx', vK_1_S*DT_q, dA) - - de_dK = de_dS0 - 0.5 * fac * (de_dD + de_dA + de_dS1 + de_dD_T + de_dA_T + de_dS1_T) - de += de_dR - de_dK - - else: - raise RuntimeError(f"Unknown implicit solvent model: {pcmobj.method}") - t1 = log.timer_debug1('grad solver', *t1) - return de.get() - - class WithSolventTDSCFGradient: from gpu4pyscf.lib.utils import to_gpu, device @@ -262,8 +99,7 @@ def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): de += grad_qv(pcmobj, dmxpy, q_sym = qE_sym_dmxpy) * 2.0 v_grids = pcmobj._get_vgrids(dmxpy, with_nuc = False) q = pcmobj._get_qsym(dmxpy, with_nuc = False)[1] - v_grids_1 = pcmobj._get_vgrids(dmxpy, with_nuc = False) - de += grad_solver(pcmobj, dmxpy, v_grids=v_grids, v_grids_1=v_grids_1, q=q) * 2.0 + de += grad_solver(pcmobj, dmxpy, v_grids=v_grids, v_grids_1=v_grids, q=q) * 2.0 return de diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py index 8a78f41fa..73635eec6 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py @@ -15,6 +15,209 @@ import unittest import numpy as np import cupy as cp -from pyscf import lib, gto -from gpu4pyscf.tdscf import rhf, rks -from gpu4pyscf import tdscf \ No newline at end of file +import pyscf +import gpu4pyscf +from pyscf import lib, gto, scf, dft +from gpu4pyscf import tdscf + +atom = """ +O 0.0000000000 0.0000000000 0.0000000000 +H 0.0000000000 -0.7570000000 0.5870000000 +H 0.0000000000 0.7570000000 0.5870000000 +""" + +bas0 = "def2svp" +bas1 = "631g" + + +def cal_analytic_gradient(td, tdgrad): + td.kernel(nstates=5) + tdgrad.base._scf.with_solvent.tdscf = True + tdgrad.kernel() + return tdgrad.de + + +def cal_td(td, tda): + td._scf.with_solvent.tdscf = True + td.kernel() + return td.e + + +def cal_mf(mol, xc, solvent, unrestrict): + if unrestrict: + if xc == 'hf': + mf = scf.UHF(mol).PCM().to_gpu() + else: + mf = dft.UKS(mol, xc=xc).PCM().to_gpu() + mf.grids.atom_grid = (99,590) + mf.grids.prune = None + else: + if xc == 'hf': + mf = scf.RHF(mol).PCM().to_gpu() + else: + mf = dft.RKS(mol, xc=xc).PCM().to_gpu() + mf.grids.atom_grid = (99,590) + mf.grids.prune = None + mf.with_solvent.method = solvent + mf.with_solvent.tdscf = True + mf.with_solvent.equilibrium_solvation = True + mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids + mf.with_solvent.eps = 78 + mf.run() + return mf + + +def get_new_mf(mol, coords, i, j, factor, delta, xc, solvent, unrestrict): + coords_new = coords*1.0 + coords_new[i, j] += delta*factor + mol.set_geom_(coords_new, unit='Ang') + mol.build() + mf = cal_mf(mol, xc, solvent, unrestrict) + return mf + + +def get_td(mf, tda, xc): + if xc == 'hf': + if tda: + td = mf.TDA() + else: + td = mf.TDHF() + else: + if tda: + td = mf.TDA() + else: + td = mf.TDDFT() + + return td + + +def setUpModule(): + global mol, molu + mol = pyscf.M( + atom=atom, basis=bas0, max_memory=32000, output="/dev/null", verbose=1) + molu = pyscf.M( + atom=atom, charge=1, spin=1, basis=bas1, max_memory=32000, output="/dev/null", verbose=1) + + +def tearDownModule(): + global mol, molu + del mol, molu + + +def benchmark_with_finite_diff(mol_input, delta=0.1, xc='b3lyp', tda=False, solvent='CPCM', unrestrict=False, num=True): + mol = mol_input.copy() + mf = cal_mf(mol, xc, solvent, unrestrict) + td = get_td(mf, tda, xc) + tdgrad = td.nuc_grad_method() + + gradient_ana = cal_analytic_gradient(td, tdgrad) + if not num: + return gradient_ana, None + if num: + coords = mol.atom_coords(unit='Ang')*1.0 + natm = coords.shape[0] + grad = np.zeros((natm, 3)) + for i in range(natm): + for j in range(3): + mf_add = get_new_mf(mol, coords, i, j, 1.0, delta, xc, solvent, unrestrict) + td_add = get_td(mf_add, tda, xc) + e1 = cal_td(td_add, tda) + e_add = e1[0] + mf_add.e_tot + + mf_minus = get_new_mf(mol, coords, i, j, -1.0, delta, xc, solvent, unrestrict) + td_minus = get_td(mf_minus, tda, xc) + e1 = cal_td(td_minus, tda) + e_minus = e1[0] + mf_minus.e_tot + grad[i, j] = (e_add - e_minus)/(delta*2.0)*0.52917721092 + return gradient_ana, grad + + + +def _check_grad_numerical(mol, tol=1e-6, xc='hf', disp=None, tda=False, solvent='CPCM', unrestrict=False, num=True): + grad_ana, grad = benchmark_with_finite_diff( + mol, delta=0.005, xc=xc, tda=tda, solvent=solvent, unrestrict=unrestrict, num=num) + if num: + norm_diff = np.linalg.norm(grad_ana - grad) + print(norm_diff) + assert norm_diff < tol + return grad_ana + + +class KnownValues(unittest.TestCase): + def test_grad_tda_singlet_hf_CPCM(self): + """ + $rem + JOBTYPE force + METHOD hf + BASIS def2-svp + CIS_N_ROOTS 5 + CIS_STATE_DERIV 1 + CIS_SINGLETS TRUE + CIS_TRIPLETS FALSE + SYMMETRY FALSE + SYM_IGNORE TRUE + ! RPA 2 + BASIS_LIN_DEP_THRESH 12 + SOLVENT_METHOD PCM + $end + + $PCM + Theory CPCM + HeavyPoints 302 + HPoints 302 + $end + + $solvent + dielectric 78 + $end + -- total gradient after adding PCM contribution -- + --------------------------------------------------- + Atom X Y Z + --------------------------------------------------- + 1 -0.000000 0.000000 0.089607 + 2 -0.000000 0.067883 -0.044803 + 3 0.000000 -0.067883 -0.044803 + --------------------------------------------------- + """ + grad_pyscf = _check_grad_numerical(mol, tol=1e-4, xc='hf', tda=True, solvent='CPCM') + ref = np.array([[-0.000000, 0.000000, 0.089607], + [ 0.000000, 0.067883, -0.044803], + [ 0.000000, -0.067883, -0.044803]]) + norm_diff = np.linalg.norm(grad_pyscf - ref) + assert norm_diff < 1e-5 + + def test_grad_tda_singlet_b3lyp_CPCM(self): + grad_pyscf = _check_grad_numerical(mol, tol=1e-4, xc='b3lyp', tda=True, solvent='CPCM') + ref = np.array([[-0.000000, 0.000000, 0.106714], + [ 0.000000, 0.073132, -0.053357], + [ 0.000000, -0.073132, -0.053357]]) + norm_diff = np.linalg.norm(grad_pyscf - ref) + assert norm_diff < 2e-5 + + def test_grad_tda_singlet_b3lyp_IEPPCM(self): + _check_grad_numerical(mol, tol=1e-4, xc='b3lyp', tda=True, solvent='IEFPCM') + + def test_grad_tda_singlet_b3lyp_COSMO(self): + _check_grad_numerical(mol, tol=1e-4, xc='b3lyp', tda=True, solvent='COSMO') + + def test_grad_tda_unrestrict_hf_CPCM(self): + grad_pyscf = _check_grad_numerical(molu, tol=1e-4, xc='hf', tda=True, unrestrict=True, solvent='CPCM') + print(grad_pyscf) + ref = np.array([[-0.000000, 0.000000, -0.066532], + [ 0.000000, 0.073344, 0.033266], + [ 0.000000, -0.073344, 0.033266]]) + norm_diff = np.linalg.norm(grad_pyscf - ref) + assert norm_diff < 2e-5 + + def test_grad_tda_unrestrict_b3lyp_CPCM(self): + grad_pyscf = _check_grad_numerical(molu, tol=1e-4, xc='b3lyp', tda=True, unrestrict=True, solvent='CPCM') + ref = np.array([[-0.000000, 0.000000, -0.037576], + [ 0.000000, 0.083399, 0.018788], + [ 0.000000, -0.083399, 0.018788]]) + norm_diff = np.linalg.norm(grad_pyscf - ref) + assert norm_diff < 2e-5 + + +if __name__ == "__main__": + print("Full Tests for TDHF and TDDFT Gradient with PCM") + unittest.main() From 26bc40100c5b1783d5581193855f84acdabdefe1 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 10 Apr 2025 14:53:41 +0800 Subject: [PATCH 13/31] Add get_ab method for PCM TDDFT Add more unit tests --- gpu4pyscf/solvent/tdscf/pcm.py | 2 +- gpu4pyscf/solvent/tests/test_pcm_tdscf.py | 92 +++++++- .../solvent/tests/test_pcm_tdscf_grad.py | 202 ++++++++++++++++-- gpu4pyscf/tdscf/rhf.py | 56 ++++- gpu4pyscf/tdscf/uhf.py | 73 ++++++- 5 files changed, 405 insertions(+), 20 deletions(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 85ba75b09..af51db5ee 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -72,7 +72,7 @@ def __init__(self, tda_grad_method): def solvent_response(self, dm): return self.base._scf.with_solvent._B_dot_x(dm)*2.0 - def grad_elec(self, xy, singlet, atmlst=None, verbose=logger.INFO): + def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.INFO): de = super().grad_elec(xy, singlet, atmlst, verbose) assert self.base._scf.with_solvent.equilibrium_solvation diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py index 9246b2cfa..0dc651750 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py @@ -19,6 +19,67 @@ from gpu4pyscf.tdscf import rhf, rks from gpu4pyscf import tdscf + +def diagonalize_tda(a, nroots=5): + nocc, nvir = a.shape[:2] + nov = nocc * nvir + a = a.reshape(nov, nov) + e, xy = np.linalg.eig(np.asarray(a)) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + +def diagonalize(a, b, nroots=5): + nocc, nvir = a.shape[:2] + nov = nocc * nvir + a = a.reshape(nov, nov) + b = b.reshape(nov, nov) + h = np.block([[a , b ], + [-b.conj(),-a.conj()]]) + e, xy = np.linalg.eig(np.asarray(h)) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + +def diagonalize_u(a, b, nroots=5): + a_aa, a_ab, a_bb = a + b_aa, b_ab, b_bb = b + nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape + a_aa = a_aa.reshape((nocc_a*nvir_a,nocc_a*nvir_a)) + a_ab = a_ab.reshape((nocc_a*nvir_a,nocc_b*nvir_b)) + a_bb = a_bb.reshape((nocc_b*nvir_b,nocc_b*nvir_b)) + b_aa = b_aa.reshape((nocc_a*nvir_a,nocc_a*nvir_a)) + b_ab = b_ab.reshape((nocc_a*nvir_a,nocc_b*nvir_b)) + b_bb = b_bb.reshape((nocc_b*nvir_b,nocc_b*nvir_b)) + a = np.block([[ a_aa , a_ab], + [ a_ab.T, a_bb]]) + b = np.block([[ b_aa , b_ab], + [ b_ab.T, b_bb]]) + abba = np.asarray(np.block([[a , b ], + [-b.conj(),-a.conj()]])) + e, xy = np.linalg.eig(abba) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + class KnownValues(unittest.TestCase): @classmethod def setUpClass(cls): @@ -117,17 +178,22 @@ def test_hf_CPCM(self): td._scf.with_solvent.build() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot - print(es_gound) ref = np.array([-75.61072291, -75.54419399, -75.51949191, -75.45219025, -75.40975027]) assert np.allclose(es_gound, ref) + a, b = td.get_ab() + es_get_ab = diagonalize(a, b)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + td = mf.TDA() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot - print(es_gound) ref = np.array([-75.60864828, -75.54169327, -75.51738767, -75.44915784, -75.40839714]) assert np.allclose(es_gound, ref) + es_get_ab = diagonalize_tda(a)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + def test_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf td = mf.TDDFT() @@ -139,12 +205,19 @@ def test_b3lyp_CPCM(self): ref = np.array([-76.06898428, -75.99630982, -75.98765186, -75.91045133, -75.84783748]) assert np.allclose(es_gound, ref) + a, b = td.get_ab() + es_get_ab = diagonalize(a, b)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + td = mf.TDA() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06789176, -75.99609709, -75.98589720, -75.90894600, -75.84699115]) assert np.allclose(es_gound, ref) + es_get_ab = diagonalize_tda(a)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + def test_b3lyp_IEFPCM(self): mf = self.mf_b3lyp_nodf_iefpcm td = mf.TDDFT() @@ -156,12 +229,19 @@ def test_b3lyp_IEFPCM(self): ref = np.array([-76.06881645, -75.99631929, -75.98713725, -75.91015704, -75.84668800]) assert np.allclose(es_gound, ref) + a, b = td.get_ab() + es_get_ab = diagonalize(a, b)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + td = mf.TDA() es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06773319, -75.99610928, -75.98534912, -75.90861455, -75.84576041]) assert np.allclose(es_gound, ref) + es_get_ab = diagonalize_tda(a)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + def test_unrestricted_hf_CPCM(self): mf = self.mfu td = mf.TDHF() @@ -173,6 +253,10 @@ def test_unrestricted_hf_CPCM(self): ref = np.array([-75.64482315, -75.61072291, -75.57156784, -75.56769949, -75.54419399]) assert np.allclose(es_gound, ref) + a, b = td.get_ab() + es_get_ab = diagonalize_u(a, b)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + def test_unrestricted_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf_u td = mf.TDDFT() @@ -184,6 +268,10 @@ def test_unrestricted_b3lyp_CPCM(self): ref = np.array([-76.09301571, -76.06898428, -76.01822101, -76.01369024, -75.99630982]) assert np.allclose(es_gound, ref) + a, b = td.get_ab() + es_get_ab = diagonalize_u(a, b)[0] + assert np.linalg.norm(es_get_ab - es) < 1e-10 + if __name__ == "__main__": print("Full Tests for PCM TDDFT") diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py index 73635eec6..bae9ea134 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py @@ -30,17 +30,168 @@ bas1 = "631g" -def cal_analytic_gradient(td, tdgrad): - td.kernel(nstates=5) - tdgrad.base._scf.with_solvent.tdscf = True - tdgrad.kernel() - return tdgrad.de +def diagonalize(a, b, nroots=5): + nocc, nvir = a.shape[:2] + nov = nocc * nvir + a = a.reshape(nov, nov) + b = b.reshape(nov, nov) + h = np.block([[a , b ], + [-b.conj(),-a.conj()]]) + e, xy = np.linalg.eig(np.asarray(h)) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + +def diagonalize_tda(a, nroots=5): + nocc, nvir = a.shape[:2] + nov = nocc * nvir + a = a.reshape(nov, nov) + e, xy = np.linalg.eig(np.asarray(a)) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + +def diagonalize_u(a, b, nroots=5): + a_aa, a_ab, a_bb = a + b_aa, b_ab, b_bb = b + nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape + a_aa = a_aa.reshape((nocc_a*nvir_a,nocc_a*nvir_a)) + a_ab = a_ab.reshape((nocc_a*nvir_a,nocc_b*nvir_b)) + a_bb = a_bb.reshape((nocc_b*nvir_b,nocc_b*nvir_b)) + b_aa = b_aa.reshape((nocc_a*nvir_a,nocc_a*nvir_a)) + b_ab = b_ab.reshape((nocc_a*nvir_a,nocc_b*nvir_b)) + b_bb = b_bb.reshape((nocc_b*nvir_b,nocc_b*nvir_b)) + a = np.block([[ a_aa , a_ab], + [ a_ab.T, a_bb]]) + b = np.block([[ b_aa , b_ab], + [ b_ab.T, b_bb]]) + abba = np.asarray(np.block([[a , b ], + [-b.conj(),-a.conj()]])) + e, xy = np.linalg.eig(abba) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + +def diagonalize_tda_u(a, nroots=5): + a_aa, a_ab, a_bb = a + nocc_a, nvir_a, nocc_b, nvir_b = a_ab.shape + a_aa = a_aa.reshape((nocc_a*nvir_a,nocc_a*nvir_a)) + a_ab = a_ab.reshape((nocc_a*nvir_a,nocc_b*nvir_b)) + a_bb = a_bb.reshape((nocc_b*nvir_b,nocc_b*nvir_b)) + a = np.block([[ a_aa , a_ab], + [ a_ab.T, a_bb]]) + e, xy = np.linalg.eig(a) + sorted_indices = np.argsort(e) + + e_sorted = e[sorted_indices] + xy_sorted = xy[:, sorted_indices] + + e_sorted_final = e_sorted[e_sorted > 1e-3] + xy_sorted = xy_sorted[:, e_sorted > 1e-3] + return e_sorted_final[:nroots], xy_sorted[:, :nroots] + + +def cal_analytic_gradient(mol, td, tdgrad, nocc, nvir, grad_elec, tda): + a, b = td.get_ab() + assert hasattr(tdgrad.base._scf, 'with_solvent') + + if tda: + atmlst = range(mol.natm) + e_diag, xy_diag = diagonalize_tda(a) + x = xy_diag[:, 0].reshape(nocc, nvir)*np.sqrt(0.5) + # de_td = grad_elec(tdgrad, (x, 0)) + de_td = grad_elec((x, 0)) + gradient_ana = de_td + tdgrad.grad_nuc(atmlst=atmlst) + else: + atmlst = range(mol.natm) + e_diag, xy_diag = diagonalize(a, b) + nsize = xy_diag.shape[0]//2 + norm_1 = np.linalg.norm(xy_diag[:nsize,0]) + norm_2 = np.linalg.norm(xy_diag[nsize:,0]) + x = xy_diag[:nsize,0]*np.sqrt(0.5/(norm_1**2-norm_2**2)) + y = xy_diag[nsize:,0]*np.sqrt(0.5/(norm_1**2-norm_2**2)) + x = x.reshape(nocc, nvir) + y = y.reshape(nocc, nvir) + + # de_td = grad_elec(tdgrad, (x, y)) + de_td = grad_elec((x, y)) + gradient_ana = de_td + tdgrad.grad_nuc(atmlst=atmlst) + + return gradient_ana + + +def cal_analytic_gradient_u(mol, td, tdgrad, nocc_a, nvir_a, nocc_b, nvir_b, grad_elec, tda): + a, b = td.get_ab() + + if tda: + atmlst = range(mol.natm) + e_diag, xy_diag = diagonalize_tda_u(a) + nsize = nocc_a*nvir_a + norm1 = np.linalg.norm(xy_diag[:nsize, 0]) + norm2 = np.linalg.norm(xy_diag[nsize:, 0]) + x_aa = xy_diag[:nsize, 0].reshape(nocc_a, nvir_a)*np.sqrt(1/(norm1**2+norm2**2)) + x_bb = xy_diag[nsize:, 0].reshape(nocc_b, nvir_b)*np.sqrt(1/(norm1**2+norm2**2)) + de_td = grad_elec(((x_aa, x_bb), (0, 0))) + gradient_ana = de_td + tdgrad.grad_nuc(atmlst=atmlst) + else: + atmlst = range(mol.natm) + e_diag, xy_diag = diagonalize_u(a, b) + nsize1 = nocc_a*nvir_a + nsize2 = nocc_b*nvir_b + norm_1 = np.linalg.norm(xy_diag[: nsize1,0]) + norm_2 = np.linalg.norm(xy_diag[nsize1: nsize1+nsize2,0]) + norm_3 = np.linalg.norm(xy_diag[nsize1+nsize2: nsize1+nsize2+nsize1,0]) + norm_4 = np.linalg.norm(xy_diag[nsize1+nsize2+nsize1: ,0]) + norm_factor = np.sqrt(1/(norm_1**2 + norm_2**2 - norm_3**2 - norm_4**2)) + x_aa = xy_diag[: nsize1,0]*norm_factor + x_bb = xy_diag[nsize1: nsize1+nsize2,0]*norm_factor + y_aa = xy_diag[nsize1+nsize2: nsize1+nsize2+nsize1,0]*norm_factor + y_bb = xy_diag[nsize1+nsize2+nsize1: ,0]*norm_factor + x_aa = x_aa.reshape(nocc_a, nvir_a) + x_bb = x_bb.reshape(nocc_b, nvir_b) + y_aa = y_aa.reshape(nocc_a, nvir_a) + y_bb = y_bb.reshape(nocc_b, nvir_b) + x = (x_aa, x_bb) + y = (y_aa, y_bb) + + de_td = grad_elec((x, y)) + gradient_ana = de_td + tdgrad.grad_nuc(atmlst=atmlst) + + return gradient_ana -def cal_td(td, tda): +def cal_td(td, tda, unrestrict): td._scf.with_solvent.tdscf = True - td.kernel() - return td.e + a, b = td.get_ab() + if unrestrict: + if tda: + e1 = diagonalize_tda_u(a)[0] + else: + e1 = diagonalize_u(a, b)[0] + else: + if tda: + e1 = diagonalize_tda(a)[0] + else: + e1 = diagonalize(a, b)[0] + return e1 def cal_mf(mol, xc, solvent, unrestrict): @@ -109,8 +260,26 @@ def benchmark_with_finite_diff(mol_input, delta=0.1, xc='b3lyp', tda=False, solv mf = cal_mf(mol, xc, solvent, unrestrict) td = get_td(mf, tda, xc) tdgrad = td.nuc_grad_method() + grad_elec = tdgrad.grad_elec - gradient_ana = cal_analytic_gradient(td, tdgrad) + if unrestrict: + mo_occ = mf.mo_occ + occidxa = np.where(mo_occ[0]>0)[0] + occidxb = np.where(mo_occ[1]>0)[0] + viridxa = np.where(mo_occ[0]==0)[0] + viridxb = np.where(mo_occ[1]==0)[0] + nocca = len(occidxa) + noccb = len(occidxb) + nvira = len(viridxa) + nvirb = len(viridxb) + gradient_ana = cal_analytic_gradient_u(mol, td, tdgrad, nocca, nvira, noccb, nvirb, grad_elec, tda) + else: + mo_coeff = mf.mo_coeff + mo_occ = mf.mo_occ + nao, nmo = mo_coeff.shape + nocc = int((mo_occ>0).sum()) + nvir = nmo - nocc + gradient_ana = cal_analytic_gradient(mol, td, tdgrad, nocc, nvir, grad_elec, tda) if not num: return gradient_ana, None if num: @@ -121,12 +290,12 @@ def benchmark_with_finite_diff(mol_input, delta=0.1, xc='b3lyp', tda=False, solv for j in range(3): mf_add = get_new_mf(mol, coords, i, j, 1.0, delta, xc, solvent, unrestrict) td_add = get_td(mf_add, tda, xc) - e1 = cal_td(td_add, tda) + e1 = cal_td(td_add, tda, unrestrict) e_add = e1[0] + mf_add.e_tot mf_minus = get_new_mf(mol, coords, i, j, -1.0, delta, xc, solvent, unrestrict) td_minus = get_td(mf_minus, tda, xc) - e1 = cal_td(td_minus, tda) + e1 = cal_td(td_minus, tda, unrestrict) e_minus = e1[0] + mf_minus.e_tot grad[i, j] = (e_add - e_minus)/(delta*2.0)*0.52917721092 return gradient_ana, grad @@ -138,7 +307,6 @@ def _check_grad_numerical(mol, tol=1e-6, xc='hf', disp=None, tda=False, solvent= mol, delta=0.005, xc=xc, tda=tda, solvent=solvent, unrestrict=unrestrict, num=num) if num: norm_diff = np.linalg.norm(grad_ana - grad) - print(norm_diff) assert norm_diff < tol return grad_ana @@ -200,9 +368,11 @@ def test_grad_tda_singlet_b3lyp_IEPPCM(self): def test_grad_tda_singlet_b3lyp_COSMO(self): _check_grad_numerical(mol, tol=1e-4, xc='b3lyp', tda=True, solvent='COSMO') + def test_grad_tda_singlet_b3lyp_ssvpe(self): + _check_grad_numerical(mol, tol=5e-4, xc='b3lyp', tda=True, solvent='ss(v)pe') + def test_grad_tda_unrestrict_hf_CPCM(self): grad_pyscf = _check_grad_numerical(molu, tol=1e-4, xc='hf', tda=True, unrestrict=True, solvent='CPCM') - print(grad_pyscf) ref = np.array([[-0.000000, 0.000000, -0.066532], [ 0.000000, 0.073344, 0.033266], [ 0.000000, -0.073344, 0.033266]]) @@ -217,6 +387,12 @@ def test_grad_tda_unrestrict_b3lyp_CPCM(self): norm_diff = np.linalg.norm(grad_pyscf - ref) assert norm_diff < 2e-5 + def test_grad_tda_unrestrict_b3lyp_IEFPCM(self): + _check_grad_numerical(molu, tol=1e-4, xc='b3lyp', tda=True, unrestrict=True, solvent='IEFPCM') + + def test_grad_tda_unrestrict_b3lyp_ssvpe(self): + _check_grad_numerical(molu, tol=8e-4, xc='b3lyp', tda=True, unrestrict=True, solvent='ss(v)pe') + if __name__ == "__main__": print("Full Tests for TDHF and TDDFT Gradient with PCM") diff --git a/gpu4pyscf/tdscf/rhf.py b/gpu4pyscf/tdscf/rhf.py index 0e9353d13..2d95b45bf 100644 --- a/gpu4pyscf/tdscf/rhf.py +++ b/gpu4pyscf/tdscf/rhf.py @@ -23,6 +23,7 @@ from gpu4pyscf.lib.cupy_helper import contract, tag_array from gpu4pyscf.lib import utils from gpu4pyscf.lib import logger +from gpu4pyscf.gto.int3c1e import int1e_grids from gpu4pyscf.scf import _response_functions # noqa from pyscf import __config__ @@ -43,8 +44,8 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True): Ref: Chem Phys Lett, 256, 454 ''' - if getattr(mf, 'with_solvent', None): - raise NotImplementedError("PCM is not supported for get_ab") + # if getattr(mf, 'with_solvent', None): + # raise NotImplementedError("PCM is not supported for get_ab") if mo_energy is None: mo_energy = mf.mo_energy if mo_coeff is None: @@ -69,6 +70,53 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True): a = cp.diag(e_ia.ravel()).reshape(nocc,nvir,nocc,nvir) b = cp.zeros_like(a) + def add_solvent_(a, b, pcmobj): + charge_exp = pcmobj.surface['charge_exp'] + grid_coords = pcmobj.surface['grid_coords'] + + vmat = -int1e_grids(pcmobj.mol, grid_coords, charge_exponents = charge_exp**2, intopt = pcmobj.intopt) + K_LU = pcmobj._intermediates['K_LU'] + K_LU_pivot = pcmobj._intermediates['K_LU_pivot'] + if not isinstance(K_LU, cp.ndarray): + K_LU = cp.asarray(K_LU) + if not isinstance(K_LU_pivot, cp.ndarray): + K_LU_pivot = cp.asarray(K_LU_pivot) + ngrid_surface = K_LU.shape[0] + L = cp.tril(K_LU, k=-1) + cp.eye(ngrid_surface) + U = cp.triu(K_LU) + + P = cp.eye(ngrid_surface) + for i in range(ngrid_surface): + pivot = int(K_LU_pivot[i].get()) + if K_LU_pivot[i] != i: + P[[i, pivot]] = P[[pivot, i]] + K = P.T @ L @ U + Kinv = cp.linalg.inv(K) + f_epsilon = pcmobj._intermediates['f_epsilon'] + if pcmobj.if_method_in_CPCM_category: + R = -f_epsilon * cp.eye(K.shape[0]) + else: + A = pcmobj._intermediates['A'] + D = pcmobj._intermediates['D'] + DA = D*A + R = -f_epsilon * (cp.eye(K.shape[0]) - 1.0/(2.0*np.pi)*DA) + Q = Kinv @ R + Qs = 0.5*(Q+Q.T) + + q_sym = cp.einsum('gh,hkl->gkl', Qs, vmat) + kEao = contract('gij,gkl->ijkl', vmat, q_sym) + kEmo = contract('pjkl,pi->ijkl', kEao, orbo.conj()) + kEmo = contract('ipkl,pj->ijkl', kEmo, mo) + kEmo = contract('ijpl,pk->ijkl', kEmo, mo.conj()) + kEmo = contract('ijkp,pl->ijkl', kEmo, mo) + kEmo = kEmo.reshape(nocc,nmo,nmo,nmo) + if singlet: + a += cp.einsum('iabj->iajb', kEmo[:nocc,nocc:,nocc:,:nocc])*2 + b += cp.einsum('iajb->iajb', kEmo[:nocc,nocc:,:nocc,nocc:])*2 + else: + raise RuntimeError("There is no solvent response for singlet-triplet excitat") + + def add_hf_(a, b, hyb=1): if getattr(mf, 'with_df', None): from gpu4pyscf.df import int3c2e @@ -99,6 +147,10 @@ def add_hf_(a, b, hyb=1): a -= cp.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb b -= cp.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb + if getattr(mf, 'with_solvent', None): + pcmobj = mf.with_solvent + add_solvent_(a, b, pcmobj) + if isinstance(mf, scf.hf.KohnShamDFT): grids = mf.grids ni = mf._numint diff --git a/gpu4pyscf/tdscf/uhf.py b/gpu4pyscf/tdscf/uhf.py index c9281ea64..af660eb39 100644 --- a/gpu4pyscf/tdscf/uhf.py +++ b/gpu4pyscf/tdscf/uhf.py @@ -23,6 +23,7 @@ from gpu4pyscf.lib import logger from gpu4pyscf.lib.cupy_helper import contract, tag_array from gpu4pyscf.tdscf._uhf_resp_sf import gen_uhf_response_sf +from gpu4pyscf.gto.int3c1e import int1e_grids from gpu4pyscf.tdscf import rhf as tdhf_gpu from gpu4pyscf.dft import KohnShamDFT from pyscf import __config__ @@ -42,8 +43,6 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None): B has three items: (B_aaaa, B_aabb, B_bbbb). B_bbaa = B_aabb.transpose(2,3,0,1). ''' - if getattr(mf, 'with_solvent', None): - raise NotImplementedError("PCM is not supported for get_ab") if mo_energy is None: mo_energy = mf.mo_energy if mo_coeff is None: @@ -84,6 +83,72 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None): a = (a_aa, a_ab, a_bb) b = (b_aa, b_ab, b_bb) + def add_solvent_(a, b, pcmobj): + charge_exp = pcmobj.surface['charge_exp'] + grid_coords = pcmobj.surface['grid_coords'] + + vmat = -int1e_grids(pcmobj.mol, grid_coords, charge_exponents = charge_exp**2, intopt = pcmobj.intopt) + K_LU = pcmobj._intermediates['K_LU'] + K_LU_pivot = pcmobj._intermediates['K_LU_pivot'] + if not isinstance(K_LU, cp.ndarray): + K_LU = cp.asarray(K_LU) + if not isinstance(K_LU_pivot, cp.ndarray): + K_LU_pivot = cp.asarray(K_LU_pivot) + ngrid_surface = K_LU.shape[0] + L = cp.tril(K_LU, k=-1) + cp.eye(ngrid_surface) + U = cp.triu(K_LU) + + P = cp.eye(ngrid_surface) + for i in range(ngrid_surface): + pivot = int(K_LU_pivot[i].get()) + if K_LU_pivot[i] != i: + P[[i, pivot]] = P[[pivot, i]] + K = P.T @ L @ U + Kinv = cp.linalg.inv(K) + f_epsilon = pcmobj._intermediates['f_epsilon'] + if pcmobj.if_method_in_CPCM_category: + R = -f_epsilon * cp.eye(K.shape[0]) + else: + A = pcmobj._intermediates['A'] + D = pcmobj._intermediates['D'] + DA = D*A + R = -f_epsilon * (cp.eye(K.shape[0]) - 1.0/(2.0*np.pi)*DA) + Q = Kinv @ R + Qs = 0.5*(Q+Q.T) + + q_sym = cp.einsum('gh,hkl->gkl', Qs, vmat) + kEao = contract('gij,gkl->ijkl', vmat, q_sym) + + kEmo_aa = contract('pjkl,pi->ijkl', kEao, orbo_a.conj()) + kEmo_aa = contract('ipkl,pj->ijkl', kEmo_aa, mo_a) + kEmo_aa = contract('ijpl,pk->ijkl', kEmo_aa, mo_a.conj()) + kEmo_aa = contract('ijkp,pl->ijkl', kEmo_aa, mo_a) + + kEmo_ab = contract('pjkl,pi->ijkl', kEao, orbo_a.conj()) + kEmo_ab = contract('ipkl,pj->ijkl', kEmo_ab, mo_a) + kEmo_ab = contract('ijpl,pk->ijkl', kEmo_ab, mo_b.conj()) + kEmo_ab = contract('ijkp,pl->ijkl', kEmo_ab, mo_b) + + kEmo_bb = contract('pjkl,pi->ijkl', kEao, orbo_b.conj()) + kEmo_bb = contract('ipkl,pj->ijkl', kEmo_bb, mo_b) + kEmo_bb = contract('ijpl,pk->ijkl', kEmo_bb, mo_b.conj()) + kEmo_bb = contract('ijkp,pl->ijkl', kEmo_bb, mo_b) + + kEmo_aa = kEmo_aa.reshape(nocc_a,nmo_a,nmo_a,nmo_a) + kEmo_ab = kEmo_ab.reshape(nocc_a,nmo_a,nmo_b,nmo_b) + kEmo_bb = kEmo_bb.reshape(nocc_b,nmo_b,nmo_b,nmo_b) + a_aa, a_ab, a_bb = a + b_aa, b_ab, b_bb = b + + a_aa += cp.einsum('iabj->iajb', kEmo_aa[:nocc_a,nocc_a:,nocc_a:,:nocc_a]) + b_aa += cp.einsum('iajb->iajb', kEmo_aa[:nocc_a,nocc_a:,:nocc_a,nocc_a:]) + + a_bb += cp.einsum('iabj->iajb', kEmo_bb[:nocc_b,nocc_b:,nocc_b:,:nocc_b]) + b_bb += cp.einsum('iajb->iajb', kEmo_bb[:nocc_b,nocc_b:,:nocc_b,nocc_b:]) + + a_ab += cp.einsum('iabj->iajb', kEmo_ab[:nocc_a,nocc_a:,nocc_b:,:nocc_b]) + b_ab += cp.einsum('iajb->iajb', kEmo_ab[:nocc_a,nocc_a:,:nocc_b,nocc_b:]) + def add_hf_(a, b, hyb=1): if getattr(mf, 'with_df', None): from gpu4pyscf.df import int3c2e @@ -135,6 +200,10 @@ def add_hf_(a, b, hyb=1): a_ab += cp.einsum('iabj->iajb', eri_ab[:nocc_a,nocc_a:,nocc_b:,:nocc_b]) b_ab += cp.einsum('iajb->iajb', eri_ab[:nocc_a,nocc_a:,:nocc_b,nocc_b:]) + if getattr(mf, 'with_solvent', None): + pcmobj = mf.with_solvent + add_solvent_(a, b, pcmobj) + if isinstance(mf, scf.hf.KohnShamDFT): ni = mf._numint if mf.do_nlc(): From 01726100c7db0c3e5e8064bc562065d4f0c43dd3 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Fri, 11 Apr 2025 08:51:41 +0800 Subject: [PATCH 14/31] in writting --- gpu4pyscf/solvent/pcm.py | 3 ++- gpu4pyscf/solvent/tdscf/pcm.py | 24 ++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index fe68fc8d7..ae65966cf 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -245,7 +245,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'tdscf' + 'equilibrium_solvation', 'e', 'v', 'tdscf', 'state_specific' } from gpu4pyscf.lib.utils import to_gpu, device @@ -265,6 +265,7 @@ def __init__(self, mol): self._intermediates = {} self.eps = 78.3553 self.tdscf = False + self.state_specific = False self.max_cycle = 20 self.conv_tol = 1e-7 diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index af51db5ee..a813977d5 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -25,6 +25,7 @@ from gpu4pyscf.lib import logger from gpu4pyscf import scf from gpu4pyscf.gto import int3c1e +from functools import wraps def make_tdscf_object(tda_method): @@ -43,6 +44,22 @@ def make_tdscf_gradient_object(tda_grad_method): (WithSolventTDSCFGradient, tda_grad_method.__class__), name) +def add_prefix(prefix): + def decorator(func): + def wrapper(*args, ​**kwargs): + original_result = func(*args, **kwargs) + return f"[{prefix}] {original_result}" + return wrapper + return decorator + + +def state_specific(td, x0=None, nstates=None): + td.kernel(x0=x0, nstates=nstates) + for icyc in range(50): + pass + # A.a = decorator(A.a) + + class WithSolventTDSCF: from gpu4pyscf.lib.utils import to_gpu, device @@ -54,6 +71,13 @@ def undo_solvent(self): name_mixin = self.base.with_solvent.__class__.__name__ obj = lib.view(self, lib.drop_class(cls, WithSolventTDSCF, name_mixin)) return obj + + def kernel(self, *args, **kwargs): + pcmobj = self._scf.with_solvent + if pcmobj.state_specific: + pass + else: + return super().kernel(*args, **kwargs) def _finalize(self): super()._finalize() From 47d688f9f28801c57f334a3de8ac34c61fef9470 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 14 Apr 2025 09:54:12 +0800 Subject: [PATCH 15/31] rewrite the LRPCM tdscf --- ...th_solvent.py => 28-tddft_with_solvent.py} | 21 +++---- gpu4pyscf/grad/tdrhf.py | 2 +- gpu4pyscf/grad/tdrks.py | 2 +- gpu4pyscf/grad/tduhf.py | 2 +- gpu4pyscf/grad/tduks.py | 2 +- gpu4pyscf/solvent/_attach_solvent.py | 35 +++--------- gpu4pyscf/solvent/pcm.py | 20 +++---- gpu4pyscf/solvent/tdscf/pcm.py | 55 +++++++++++++++---- gpu4pyscf/solvent/tests/test_pcm_tdscf.py | 31 +++-------- .../solvent/tests/test_pcm_tdscf_grad.py | 11 ++-- gpu4pyscf/tdscf/rhf.py | 28 ++++++---- gpu4pyscf/tdscf/rks.py | 2 +- gpu4pyscf/tdscf/tests/test_tdrks.py | 49 ++++++++++------- gpu4pyscf/tdscf/tests/test_tduks.py | 49 ++++++++++------- gpu4pyscf/tdscf/uhf.py | 24 ++++---- gpu4pyscf/tdscf/uks.py | 2 +- 16 files changed, 175 insertions(+), 160 deletions(-) rename examples/{27-tddft_with_solvent.py => 28-tddft_with_solvent.py} (73%) diff --git a/examples/27-tddft_with_solvent.py b/examples/28-tddft_with_solvent.py similarity index 73% rename from examples/27-tddft_with_solvent.py rename to examples/28-tddft_with_solvent.py index 00232c05e..2fed5eabb 100644 --- a/examples/27-tddft_with_solvent.py +++ b/examples/28-tddft_with_solvent.py @@ -40,21 +40,14 @@ # Compute TDDFT and TDA excitation energy print('------------------- vertical exitation TDA -----------------------------') -td = mf.TDDFT().set(nstates=5) -td._scf.with_solvent.tdscf = True -td._scf.with_solvent.eps = 1.78 -td._scf.with_solvent.build() -e_tddft = td.kernel()[0] # [ 8.03553827 10.07361783 10.20203523 12.36009792 13.83374455] +td = mf.TDDFT(equilibrium_solvation=False).set(nstates=5) +e_tddft = td.kernel()[0] # [ 8.03553836 10.073618 10.20203543 12.36009825 13.83374465] # print('5 TDDFT excitation energy by GPU4PySCF') # print(e_tddft) print('------------------- adiabatic excitation TDA -----------------------------') -td = mf.TDA().set(nstates=5) -td._scf.with_solvent.tdscf = True -td._scf.with_solvent.eps = 78.0 -td._scf.with_solvent.equilibrium_solvation = True -td._scf.with_solvent.build() -e_tda = td.kernel()[0] # [ 7.99456759 10.0632959 10.08523494 12.30675282 13.64298125] +td = mf.TDA(equilibrium_solvation=True).set(nstates=5) +e_tda = td.kernel()[0] # [ 7.99456768 10.06329607 10.08523514 12.30675317 13.64298135] # print('5 TDA excitation energy by GPU4PySCF') # print(e_tda) @@ -64,8 +57,8 @@ """ --------- PCMTDA gradients for state 1 ---------- x y z -0 O -0.0000000000 0.0000000000 -0.0836461430 -1 H 0.0601539533 -0.0000000000 0.0418232965 -2 H -0.0601539533 -0.0000000000 0.0418232965 +0 O -0.0000000000 0.0000000000 -0.0836461563 +1 H 0.0601539524 -0.0000000000 0.0418233032 +2 H -0.0601539524 -0.0000000000 0.0418233032 ---------------------------------------------- """ \ No newline at end of file diff --git a/gpu4pyscf/grad/tdrhf.py b/gpu4pyscf/grad/tdrhf.py index e12f57d38..08a7c8371 100644 --- a/gpu4pyscf/grad/tdrhf.py +++ b/gpu4pyscf/grad/tdrhf.py @@ -98,7 +98,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): wvo += contract("ac,ai->ci", veff0mom[nocc:, nocc:], xmy) * 2 # set singlet=None, generate function for CPHF type response kernel - vresp = mf.gen_response(singlet=None, hermi=1) + vresp = td_grad.base.gen_response(singlet=None, hermi=1) def fvind(x): # For singlet, closed shell ground state dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # 2 for double occupancy diff --git a/gpu4pyscf/grad/tdrks.py b/gpu4pyscf/grad/tdrks.py index 7e48302c9..5933dcbc1 100644 --- a/gpu4pyscf/grad/tdrks.py +++ b/gpu4pyscf/grad/tdrks.py @@ -142,7 +142,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): veff0mom = cp.zeros((nmo, nmo)) # set singlet=None, generate function for CPHF type response kernel - vresp = mf.gen_response(singlet=None, hermi=1) + vresp = td_grad.base.gen_response(singlet=None, hermi=1) def fvind(x): dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) diff --git a/gpu4pyscf/grad/tduhf.py b/gpu4pyscf/grad/tduhf.py index 7b9f49e90..e6849ee30 100644 --- a/gpu4pyscf/grad/tduhf.py +++ b/gpu4pyscf/grad/tduhf.py @@ -127,7 +127,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): wvoa += contract("ac,ai->ci", veff0moma[nocca:, nocca:], xmya) * 2 wvob += contract("ac,ai->ci", veff0momb[noccb:, noccb:], xmyb) * 2 - vresp = mf.gen_response(hermi=1) + vresp = td_grad.base.gen_response(hermi=1) def fvind(x): dm1 = cp.empty((2, nao, nao)) diff --git a/gpu4pyscf/grad/tduks.py b/gpu4pyscf/grad/tduks.py index b4c9a3d76..3f6d7a3ea 100644 --- a/gpu4pyscf/grad/tduks.py +++ b/gpu4pyscf/grad/tduks.py @@ -178,7 +178,7 @@ def grad_elec(td_grad, x_y, singlet=True, atmlst=None, verbose=logger.INFO): veff0moma = cp.zeros((nmoa, nmoa)) veff0momb = cp.zeros((nmob, nmob)) - vresp = mf.gen_response(hermi=1) + vresp = td_grad.base.gen_response(hermi=1) def fvind(x): dm1 = cp.empty((2, nao, nao)) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index ad9794db9..a38322538 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -137,29 +137,21 @@ def nuc_grad_method(self): grad_method = super().nuc_grad_method() return self.with_solvent.nuc_grad_method(grad_method) - def TDA(self): - # if isinstance(self, dft.rks.RKS): - # tda_method = + def TDA(self, equilibrium_solvation, eps_optical=1.78): tda_method = super().TDA() - return self.with_solvent.TDA(tda_method) + return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation) - def TDDFT(self): - # if isinstance(self, dft.rks.RKS): - # tda_method = + def TDDFT(self, equilibrium_solvation, eps_optical=1.78): tda_method = super().TDDFT() - return self.with_solvent.TDDFT(tda_method) + return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation) - def TDHF(self): - # if isinstance(self, dft.rks.RKS): - # tda_method = + def TDHF(self, equilibrium_solvation, eps_optical=1.78): tda_method = super().TDHF() - return self.with_solvent.TDHF(tda_method) + return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation) - def CasidaTDDFT(self): - # if isinstance(self, dft.rks.RKS): - # tda_method = + def CasidaTDDFT(self, equilibrium_solvation, eps_optical=1.78): tda_method = super().CasidaTDDFT() - return self.with_solvent.CasidaTDDFT(tda_method) + return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation) Gradients = nuc_grad_method @@ -175,21 +167,12 @@ def gen_response(self, *args, **kwargs): singlet = singlet or singlet is None def vind_with_solvent(dm1): v = vind(dm1) - if self.with_solvent.tdscf: + if self.with_solvent.equilibrium_solvation: if is_uhf: v_solvent = self.with_solvent._B_dot_x(dm1) v += v_solvent[0] + v_solvent[1] elif singlet: v += self.with_solvent._B_dot_x(dm1) - else: - logger.warn(self, 'Singlet-Triplet has no LR-PCM contribution!') - else: - if self.with_solvent.equilibrium_solvation: - if is_uhf: - v_solvent = self.with_solvent._B_dot_x(dm1) - v += v_solvent[0] + v_solvent[1] - elif singlet: - v += self.with_solvent._B_dot_x(dm1) return v return vind_with_solvent diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index fe68fc8d7..5252da3c8 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -245,7 +245,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'tdscf' + 'equilibrium_solvation', 'e', 'v', 'v_grids_n' } from gpu4pyscf.lib.utils import to_gpu, device @@ -264,7 +264,6 @@ def __init__(self, mol): self.lebedev_order = 29 self._intermediates = {} self.eps = 78.3553 - self.tdscf = False self.max_cycle = 20 self.conv_tol = 1e-7 @@ -275,6 +274,7 @@ def __init__(self, mol): self.e = None self.v = None + self.v_grids_n = None def dump_flags(self, verbose=None): logger.info(self, '******** %s ********', self.__class__) @@ -468,29 +468,29 @@ def nuc_grad_method(self, grad_method): else: raise RuntimeError('Only SCF gradient is supported') - def TDA(self, td): + def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td if self.frozen: raise RuntimeError('Frozen solvent model is not supported') - return pcm_td.make_tdscf_object(td) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) - def TDHF(self, td): + def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td if self.frozen: raise RuntimeError('Frozen solvent model is not supported') - return pcm_td.make_tdscf_object(td) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) - def TDDFT(self, td): + def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td if self.frozen: raise RuntimeError('Frozen solvent model is not supported') - return pcm_td.make_tdscf_object(td) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) - def CasidaTDDFT(self, td): + def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td if self.frozen: raise RuntimeError('Frozen solvent model is not supported') - return pcm_td.make_tdscf_object(td) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def Hessian(self, hess_method): from gpu4pyscf.solvent.hessian import pcm as pcm_hess diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index af51db5ee..c2d722123 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -18,20 +18,25 @@ import cupy as cp from pyscf import lib -from gpu4pyscf.solvent.pcm import PI, switch_h, libsolvent +from gpu4pyscf.solvent.pcm import PCM from gpu4pyscf.solvent.grad.pcm import grad_nuc, grad_qv, grad_solver -from gpu4pyscf.solvent.grad.pcm import left_multiply_dS, right_multiply_dS, get_dF_dA, get_dSii, get_dD_dS -from gpu4pyscf.lib.cupy_helper import contract from gpu4pyscf.lib import logger from gpu4pyscf import scf -from gpu4pyscf.gto import int3c1e -def make_tdscf_object(tda_method): +class TDPCM(PCM): + def __init__(self, mfpcmobj, eps_optical=1.78, equilium_solvation=False): + self.__dict__.update(mfpcmobj.__dict__) + self.equilibrium_solvation = equilium_solvation + if not equilium_solvation: + self.eps = eps_optical + + +def make_tdscf_object(tda_method, eps_optical=1.78, equilibrium_solvation=False): '''For td_method in vacuum, add td of solvent pcmobj''' name = (tda_method._scf.with_solvent.__class__.__name__ + tda_method.__class__.__name__) - return lib.set_class(WithSolventTDSCF(tda_method), + return lib.set_class(WithSolventTDSCF(tda_method, eps_optical, equilibrium_solvation), (WithSolventTDSCF, tda_method.__class__), name) @@ -46,8 +51,33 @@ def make_tdscf_gradient_object(tda_grad_method): class WithSolventTDSCF: from gpu4pyscf.lib.utils import to_gpu, device - def __init__(self, tda_method): + _keys = {'with_solvent'} + + def gen_response(self, *args, **kwargs): + pcmobj = self.with_solvent + mf = self._scf + vind = super().gen_response(*args, **kwargs) + is_uhf = isinstance(mf, scf.uhf.UHF) + # singlet=None is orbital hessian or CPHF type response function + singlet = kwargs.get('singlet', True) + singlet = singlet or singlet is None + def vind_with_solvent(dm1): + v = vind(dm1) + if is_uhf: + v_solvent = pcmobj._B_dot_x(dm1) + v += v_solvent[0] + v_solvent[1] + elif singlet: + v += pcmobj._B_dot_x(dm1) + else: + logger.warn(pcmobj, 'Singlet-Triplet has no LR-PCM contribution!') + return v + return vind_with_solvent + + def __init__(self, tda_method, eps_optical=1.78, equilibrium_solvation=False): self.__dict__.update(tda_method.__dict__) + self.with_solvent = TDPCM(tda_method._scf.with_solvent, eps_optical, equilibrium_solvation) + if not self.with_solvent.equilibrium_solvation: + self.with_solvent.build() def undo_solvent(self): cls = self.__class__ @@ -57,6 +87,11 @@ def undo_solvent(self): def _finalize(self): super()._finalize() + if self.with_solvent.equilibrium_solvation: + logger.info(self.with_solvent, 'equilibrium solvation NOT for vertical excitation') + else: + logger.info(self.with_solvent, 'Not equilibrium solvation NOT for vertical excitation,\n\ + eps_optical = %s', self.with_solvent.eps) def nuc_grad_method(self): grad_method = super().nuc_grad_method() @@ -70,12 +105,12 @@ def __init__(self, tda_grad_method): self.__dict__.update(tda_grad_method.__dict__) def solvent_response(self, dm): - return self.base._scf.with_solvent._B_dot_x(dm)*2.0 + return self.base.with_solvent._B_dot_x(dm)*2.0 def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.INFO): de = super().grad_elec(xy, singlet, atmlst, verbose) - assert self.base._scf.with_solvent.equilibrium_solvation + assert self.base.with_solvent.equilibrium_solvation dm = self.base._scf.make_rdm1(ao_repr=True) if dm.ndim == 3: @@ -83,7 +118,7 @@ def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.INFO): # TODO: add unrestricted case support dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T) dmxpy = self.dmxpy + self.dmxpy.T - pcmobj = self.base._scf.with_solvent + pcmobj = self.base.with_solvent de += grad_qv(pcmobj, dm) de += grad_solver(pcmobj, dm) de += grad_nuc(pcmobj, dm) diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py index 0dc651750..b3c87f796 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf.py @@ -172,10 +172,7 @@ def test_hf_CPCM(self): $end """ mf = self.mf - td = mf.TDHF() - td._scf.with_solvent.tdscf = True - td._scf.with_solvent.eps = 1.78 - td._scf.with_solvent.build() + td = mf.TDHF(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-75.61072291, -75.54419399, -75.51949191, -75.45219025, -75.40975027]) @@ -185,7 +182,7 @@ def test_hf_CPCM(self): es_get_ab = diagonalize(a, b)[0] assert np.linalg.norm(es_get_ab - es) < 1e-10 - td = mf.TDA() + td = mf.TDA(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-75.60864828, -75.54169327, -75.51738767, -75.44915784, -75.40839714]) @@ -196,10 +193,7 @@ def test_hf_CPCM(self): def test_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf - td = mf.TDDFT() - td._scf.with_solvent.tdscf = True - td._scf.with_solvent.eps = 1.78 - td._scf.with_solvent.build() + td = mf.TDDFT(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06898428, -75.99630982, -75.98765186, -75.91045133, -75.84783748]) @@ -209,7 +203,7 @@ def test_b3lyp_CPCM(self): es_get_ab = diagonalize(a, b)[0] assert np.linalg.norm(es_get_ab - es) < 1e-10 - td = mf.TDA() + td = mf.TDA(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06789176, -75.99609709, -75.98589720, -75.90894600, -75.84699115]) @@ -220,10 +214,7 @@ def test_b3lyp_CPCM(self): def test_b3lyp_IEFPCM(self): mf = self.mf_b3lyp_nodf_iefpcm - td = mf.TDDFT() - td._scf.with_solvent.tdscf = True - td._scf.with_solvent.eps = 1.78 - td._scf.with_solvent.build() + td = mf.TDDFT(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06881645, -75.99631929, -75.98713725, -75.91015704, -75.84668800]) @@ -233,7 +224,7 @@ def test_b3lyp_IEFPCM(self): es_get_ab = diagonalize(a, b)[0] assert np.linalg.norm(es_get_ab - es) < 1e-10 - td = mf.TDA() + td = mf.TDA(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.06773319, -75.99610928, -75.98534912, -75.90861455, -75.84576041]) @@ -244,10 +235,7 @@ def test_b3lyp_IEFPCM(self): def test_unrestricted_hf_CPCM(self): mf = self.mfu - td = mf.TDHF() - td._scf.with_solvent.tdscf = True - td._scf.with_solvent.eps = 1.78 - td._scf.with_solvent.build() + td = mf.TDHF(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-75.64482315, -75.61072291, -75.57156784, -75.56769949, -75.54419399]) @@ -259,10 +247,7 @@ def test_unrestricted_hf_CPCM(self): def test_unrestricted_b3lyp_CPCM(self): mf = self.mf_b3lyp_nodf_u - td = mf.TDDFT() - td._scf.with_solvent.tdscf = True - td._scf.with_solvent.eps = 1.78 - td._scf.with_solvent.build() + td = mf.TDDFT(equilibrium_solvation=False) es = td.kernel(nstates=5)[0] es_gound = es + mf.e_tot ref = np.array([-76.09301571, -76.06898428, -76.01822101, -76.01369024, -75.99630982]) diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py index bae9ea134..c4623b76c 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py @@ -179,7 +179,6 @@ def cal_analytic_gradient_u(mol, td, tdgrad, nocc_a, nvir_a, nocc_b, nvir_b, gra def cal_td(td, tda, unrestrict): - td._scf.with_solvent.tdscf = True a, b = td.get_ab() if unrestrict: if tda: @@ -210,8 +209,6 @@ def cal_mf(mol, xc, solvent, unrestrict): mf.grids.atom_grid = (99,590) mf.grids.prune = None mf.with_solvent.method = solvent - mf.with_solvent.tdscf = True - mf.with_solvent.equilibrium_solvation = True mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids mf.with_solvent.eps = 78 mf.run() @@ -230,14 +227,14 @@ def get_new_mf(mol, coords, i, j, factor, delta, xc, solvent, unrestrict): def get_td(mf, tda, xc): if xc == 'hf': if tda: - td = mf.TDA() + td = mf.TDA(equilibrium_solvation=True) else: - td = mf.TDHF() + td = mf.TDHF(equilibrium_solvation=True) else: if tda: - td = mf.TDA() + td = mf.TDA(equilibrium_solvation=True) else: - td = mf.TDDFT() + td = mf.TDDFT(equilibrium_solvation=True) return td diff --git a/gpu4pyscf/tdscf/rhf.py b/gpu4pyscf/tdscf/rhf.py index 2d95b45bf..72a35ebc2 100644 --- a/gpu4pyscf/tdscf/rhf.py +++ b/gpu4pyscf/tdscf/rhf.py @@ -36,7 +36,7 @@ ] -def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True): +def get_ab(td, mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True): r'''A and B matrices for TDDFT response function. A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ai||jb) @@ -44,8 +44,7 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True): Ref: Chem Phys Lett, 256, 454 ''' - # if getattr(mf, 'with_solvent', None): - # raise NotImplementedError("PCM is not supported for get_ab") + if mo_energy is None: mo_energy = mf.mo_energy if mo_coeff is None: @@ -147,8 +146,8 @@ def add_hf_(a, b, hyb=1): a -= cp.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb b -= cp.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb - if getattr(mf, 'with_solvent', None): - pcmobj = mf.with_solvent + if getattr(td, 'with_solvent', None): + pcmobj = td.with_solvent add_solvent_(a, b, pcmobj) if isinstance(mf, scf.hf.KohnShamDFT): @@ -281,7 +280,7 @@ def add_hf_(a, b, hyb=1): return a.get(), b.get() -def gen_tda_operation(mf, fock_ao=None, singlet=True, wfnsym=None): +def gen_tda_operation(td, mf, fock_ao=None, singlet=True, wfnsym=None): '''Generate function to compute A x ''' assert fock_ao is None @@ -299,7 +298,7 @@ def gen_tda_operation(mf, fock_ao=None, singlet=True, wfnsym=None): e_ia = hdiag = mo_energy[viridx] - mo_energy[occidx,None] hdiag = hdiag.ravel() - vresp = mf.gen_response(singlet=singlet, hermi=0) + vresp = td.gen_response(singlet=singlet, hermi=0) nocc, nvir = e_ia.shape def vind(zs): @@ -342,10 +341,15 @@ class TDBase(lib.StreamObject): _finalize = tdhf_cpu.TDBase._finalize gen_vind = NotImplemented + + def gen_response(self, singlet=True, hermi=0): + '''Generate function to compute A x''' + return self._scf.gen_response(singlet=singlet, hermi=hermi) + def get_ab(self, mf=None): if mf is None: mf = self._scf - return get_ab(mf, singlet=self.singlet) + return get_ab(self, mf, singlet=self.singlet) def get_precond(self, hdiag): threshold_t=1.0e-4 @@ -444,7 +448,7 @@ def gen_vind(self, mf=None): '''Generate function to compute Ax''' if mf is None: mf = self._scf - return gen_tda_operation(mf, singlet=self.singlet) + return gen_tda_operation(self, mf, singlet=self.singlet) def init_guess(self, mf=None, nstates=None, wfnsym=None, return_symmetry=False): ''' @@ -522,7 +526,7 @@ def pickeig(w, v, nroots, envs): CIS = TDA -def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None): +def gen_tdhf_operation(td, mf, fock_ao=None, singlet=True, wfnsym=None): '''Generate function to compute [ A B ][X] @@ -540,7 +544,7 @@ def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None): orbo = mo_coeff[:,occidx] e_ia = hdiag = mo_energy[viridx] - mo_energy[occidx,None] - vresp = mf.gen_response(singlet=singlet, hermi=0) + vresp = td.gen_response(singlet=singlet, hermi=0) nocc, nvir = e_ia.shape def vind(zs): @@ -574,7 +578,7 @@ class TDHF(TDBase): def gen_vind(self, mf=None): if mf is None: mf = self._scf - return gen_tdhf_operation(mf, singlet=self.singlet) + return gen_tdhf_operation(self, mf, singlet=self.singlet) def init_guess(self, mf=None, nstates=None, wfnsym=None, return_symmetry=False): assert not return_symmetry diff --git a/gpu4pyscf/tdscf/rks.py b/gpu4pyscf/tdscf/rks.py index fcb761b7f..70a5b8f6f 100644 --- a/gpu4pyscf/tdscf/rks.py +++ b/gpu4pyscf/tdscf/rks.py @@ -69,7 +69,7 @@ def gen_vind(self, mf=None): d_ia = e_ia ** .5 ed_ia = e_ia * d_ia hdiag = e_ia.ravel() ** 2 - vresp = mf.gen_response(singlet=singlet, hermi=1) + vresp = self.gen_response(singlet=singlet, hermi=1) nocc, nvir = e_ia.shape def vind(zs): diff --git a/gpu4pyscf/tdscf/tests/test_tdrks.py b/gpu4pyscf/tdscf/tests/test_tdrks.py index 8eb7dc8d2..8859468f8 100644 --- a/gpu4pyscf/tdscf/tests/test_tdrks.py +++ b/gpu4pyscf/tdscf/tests/test_tdrks.py @@ -384,9 +384,10 @@ def test_casida_tddft_vind(self): def test_ab_hf(self): mf = self.mf - a, b = rhf.get_ab(mf) - ftda = rhf.gen_tda_operation(mf, singlet=True)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rhf.TDHF(mf) + a, b = td.get_ab() + ftda = rhf.gen_tda_operation(td, mf, singlet=True)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) @@ -403,9 +404,10 @@ def test_ab_hf(self): def test_ab_lda(self): mf = self.mf_lda_nodf - a, b = rhf.get_ab(mf) - ftda = rhf.gen_tda_operation(mf, singlet=True)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rks.TDDFT(mf) + a, b = td.get_ab() + ftda = rhf.gen_tda_operation(td, mf, singlet=True)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) @@ -422,9 +424,10 @@ def test_ab_lda(self): def test_ab_lda_df(self): mf = self.mf_lda - a, b = rhf.get_ab(mf) - ftda = rhf.gen_tda_operation(mf, singlet=True)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rks.TDDFT(mf) + a, b = td.get_ab(mf) + ftda = rhf.gen_tda_operation(td, mf, singlet=True)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) @@ -441,9 +444,10 @@ def test_ab_lda_df(self): def test_ab_b3lyp(self): mf = self.mf_b3lyp_nodf - a, b = rks.TDDFT(mf).get_ab() - ftda = rhf.gen_tda_operation(mf, singlet=None)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rks.TDDFT(mf) + a, b = td.get_ab() + ftda = rhf.gen_tda_operation(td, mf, singlet=None)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) @@ -460,9 +464,10 @@ def test_ab_b3lyp(self): def test_ab_b3lyp_df(self): mf = self.mf_b3lyp - a, b = rks.TDDFT(mf).get_ab() - ftda = rhf.gen_tda_operation(mf, singlet=None)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rks.TDDFT(mf) + a, b = td.get_ab() + ftda = rhf.gen_tda_operation(td, mf, singlet=None)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) @@ -479,9 +484,10 @@ def test_ab_b3lyp_df(self): def test_ab_mgga(self): mf = self.mf_m06l_nodf - a, b = rks.TDDFT(mf).get_ab() - ftda = rhf.gen_tda_operation(mf, singlet=None)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rks.TDDFT(mf) + a, b = td.get_ab() + ftda = rhf.gen_tda_operation(td, mf, singlet=None)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) @@ -498,9 +504,10 @@ def test_ab_mgga(self): def test_ab_mgga_df(self): mf = self.mf_m06l - a, b = rks.TDDFT(mf).get_ab() - ftda = rhf.gen_tda_operation(mf, singlet=None)[0] - ftdhf = rhf.gen_tdhf_operation(mf, singlet=True)[0] + td = rks.TDDFT(mf) + a, b = td.get_ab() + ftda = rhf.gen_tda_operation(td, mf, singlet=None)[0] + ftdhf = rhf.gen_tdhf_operation(td, mf, singlet=True)[0] nocc = int(np.count_nonzero(mf.mo_occ == 2)) nvir = int(np.count_nonzero(mf.mo_occ == 0)) np.random.seed(2) diff --git a/gpu4pyscf/tdscf/tests/test_tduks.py b/gpu4pyscf/tdscf/tests/test_tduks.py index 529e26e70..dfe051e9f 100644 --- a/gpu4pyscf/tdscf/tests/test_tduks.py +++ b/gpu4pyscf/tdscf/tests/test_tduks.py @@ -284,9 +284,10 @@ def test_casida_tddft_vind(self): def test_ab_hf(self): mf = self.mf_uhf - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) @@ -330,9 +331,10 @@ def test_ab_hf(self): def test_ab_lda(self): mf = self.mf_lda_nodf - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) @@ -376,9 +378,10 @@ def test_ab_lda(self): def test_ab_lda_df(self): mf = self.mf_lda - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) @@ -422,9 +425,10 @@ def test_ab_lda_df(self): def test_ab_b3lyp(self): mf = self.mf_b3lyp_nodf - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) @@ -468,9 +472,10 @@ def test_ab_b3lyp(self): def test_ab_b3lyp_df(self): mf = self.mf_b3lyp - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) @@ -514,9 +519,10 @@ def test_ab_b3lyp_df(self): def test_ab_mgga(self): mf = self.mf_m06l_nodf - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) @@ -560,9 +566,10 @@ def test_ab_mgga(self): def test_ab_mgga_df(self): mf = self.mf_m06l - a, b = tdscf.uks.TDDFT(mf).get_ab() - ftda = tdscf.uhf.gen_tda_operation(mf)[0] - ftdhf = tdscf.uhf.gen_tdhf_operation(mf)[0] + td = tdscf.uks.TDDFT(mf) + a, b = td.get_ab() + ftda = tdscf.uhf.gen_tda_operation(td, mf)[0] + ftdhf = tdscf.uhf.gen_tdhf_operation(td, mf)[0] nocc_a = int(np.count_nonzero(mf.mo_occ[0] == 1)) nvir_a = int(np.count_nonzero(mf.mo_occ[0] == 0)) nocc_b = int(np.count_nonzero(mf.mo_occ[1] == 1)) diff --git a/gpu4pyscf/tdscf/uhf.py b/gpu4pyscf/tdscf/uhf.py index af660eb39..4069e7d1a 100644 --- a/gpu4pyscf/tdscf/uhf.py +++ b/gpu4pyscf/tdscf/uhf.py @@ -32,7 +32,7 @@ 'TDA', 'CIS', 'TDHF', 'TDUHF', 'TDBase' ] -def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None): +def get_ab(td, mf, mo_energy=None, mo_coeff=None, mo_occ=None): r'''A and B matrices for TDDFT response function. A[i,a,j,b] = \delta_{ab}\delta_{ij}(E_a - E_i) + (ai||jb) @@ -200,8 +200,8 @@ def add_hf_(a, b, hyb=1): a_ab += cp.einsum('iabj->iajb', eri_ab[:nocc_a,nocc_a:,nocc_b:,:nocc_b]) b_ab += cp.einsum('iajb->iajb', eri_ab[:nocc_a,nocc_a:,:nocc_b,nocc_b:]) - if getattr(mf, 'with_solvent', None): - pcmobj = mf.with_solvent + if getattr(td, 'with_solvent', None): + pcmobj = td.with_solvent add_solvent_(a, b, pcmobj) if isinstance(mf, scf.hf.KohnShamDFT): @@ -400,7 +400,7 @@ def add_hf_(a, b, hyb=1): REAL_EIG_THRESHOLD = tdhf_cpu.REAL_EIG_THRESHOLD -def gen_tda_operation(mf, fock_ao=None, wfnsym=None): +def gen_tda_operation(td, mf, fock_ao=None, wfnsym=None): '''A x ''' assert fock_ao is None @@ -433,7 +433,7 @@ def gen_tda_operation(mf, fock_ao=None, wfnsym=None): nocca, nvira = e_ia_a.shape noccb, nvirb = e_ia_b.shape - vresp = mf.gen_response(hermi=0) + vresp = td.gen_response(hermi=0) def vind(zs): nz = len(zs) @@ -461,9 +461,13 @@ def vind(zs): class TDBase(tdhf_gpu.TDBase): + def gen_response(self, mo_coeff=None, mo_occ=None, hermi=0): + '''Generate function to compute A x''' + return self._scf.gen_response(mo_coeff=None, mo_occ=None, hermi=hermi) + def get_ab(self, mf=None): if mf is None: mf = self._scf - return get_ab(mf) + return get_ab(self, mf) def nuc_grad_method(self): if getattr(self._scf, 'with_df', None): @@ -512,7 +516,7 @@ def gen_vind(self, mf=None): '''Generate function to compute Ax''' if mf is None: mf = self._scf - return gen_tda_operation(mf) + return gen_tda_operation(self, mf) def init_guess(self, mf=None, nstates=None, wfnsym=None, return_symmetry=False): if mf is None: mf = self._scf @@ -771,7 +775,7 @@ def all_eigs(w, v, nroots, envs): return self.e, self.xy -def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None): +def gen_tdhf_operation(td, mf, fock_ao=None, singlet=True, wfnsym=None): '''Generate function to compute [ A B ][X] @@ -805,7 +809,7 @@ def gen_tdhf_operation(mf, fock_ao=None, singlet=True, wfnsym=None): nocca, nvira = e_ia_a.shape noccb, nvirb = e_ia_b.shape - vresp = mf.gen_response(hermi=0) + vresp = td.gen_response(hermi=0) def vind(zs): nz = len(zs) @@ -854,7 +858,7 @@ class TDHF(TDBase): def gen_vind(self, mf=None): if mf is None: mf = self._scf - return gen_tdhf_operation(mf, singlet=self.singlet) + return gen_tdhf_operation(self, mf, singlet=self.singlet) get_precond = tdhf_gpu.TDHF.get_precond diff --git a/gpu4pyscf/tdscf/uks.py b/gpu4pyscf/tdscf/uks.py index 0428e29b7..9970630a8 100644 --- a/gpu4pyscf/tdscf/uks.py +++ b/gpu4pyscf/tdscf/uks.py @@ -87,7 +87,7 @@ def gen_vind(self, mf=None): ed_ia = e_ia * d_ia hdiag = e_ia ** 2 hdiag = hdiag - vresp = mf.gen_response(mo_coeff, mo_occ, hermi=1) + vresp = self.gen_response(mo_coeff, mo_occ, hermi=1) nocca, nvira = e_ia_a.shape noccb, nvirb = e_ia_b.shape From 6180862ca9499de5d0359a83d61af9719c3c9f26 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 14 Apr 2025 13:54:17 +0800 Subject: [PATCH 16/31] Review the codes --- examples/28-tddft_with_solvent.py | 12 ++++++---- gpu4pyscf/solvent/_attach_solvent.py | 16 +++++++++---- gpu4pyscf/solvent/grad/pcm.py | 8 +++---- gpu4pyscf/solvent/pcm.py | 23 +++++-------------- gpu4pyscf/solvent/tdscf/pcm.py | 9 ++++---- .../solvent/tests/test_pcm_tdscf_grad.py | 2 -- 6 files changed, 35 insertions(+), 35 deletions(-) diff --git a/examples/28-tddft_with_solvent.py b/examples/28-tddft_with_solvent.py index 2fed5eabb..41c3d62b6 100644 --- a/examples/28-tddft_with_solvent.py +++ b/examples/28-tddft_with_solvent.py @@ -39,19 +39,19 @@ mf.kernel() # -76.476456106979 # Compute TDDFT and TDA excitation energy -print('------------------- vertical exitation TDA -----------------------------') +print('------------------- vertical exitation TDA (non-equilibrium solvent)-----------------------------') td = mf.TDDFT(equilibrium_solvation=False).set(nstates=5) e_tddft = td.kernel()[0] # [ 8.03553836 10.073618 10.20203543 12.36009825 13.83374465] # print('5 TDDFT excitation energy by GPU4PySCF') # print(e_tddft) -print('------------------- adiabatic excitation TDA -----------------------------') +print('------------------- vertical excitation TDA (equilibrium solvent)-----------------------------') td = mf.TDA(equilibrium_solvation=True).set(nstates=5) e_tda = td.kernel()[0] # [ 7.99456768 10.06329607 10.08523514 12.30675317 13.64298135] # print('5 TDA excitation energy by GPU4PySCF') # print(e_tda) -print('The gradient of first TDA excitation energy by GPU4PySCF') +print('The gradient of first TDA excitation energy by GPU4PySCF with equilibrium solvent') g = td.nuc_grad_method() g.kernel() """ @@ -61,4 +61,8 @@ 1 H 0.0601539524 -0.0000000000 0.0418233032 2 H -0.0601539524 -0.0000000000 0.0418233032 ---------------------------------------------- -""" \ No newline at end of file +""" + +# ValueError: equilibrium_solvation must be specified +# td = mf.TDA().set(nstates=5) +# e_tda = td.kernel()[0] \ No newline at end of file diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index a38322538..84637bb2b 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -137,19 +137,27 @@ def nuc_grad_method(self): grad_method = super().nuc_grad_method() return self.with_solvent.nuc_grad_method(grad_method) - def TDA(self, equilibrium_solvation, eps_optical=1.78): + def TDA(self, equilibrium_solvation=None, eps_optical=1.78): + if equilibrium_solvation is None: + raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDA() return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation) - def TDDFT(self, equilibrium_solvation, eps_optical=1.78): + def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78): + if equilibrium_solvation is None: + raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDDFT() return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation) - def TDHF(self, equilibrium_solvation, eps_optical=1.78): + def TDHF(self, equilibrium_solvation=None, eps_optical=1.78): + if equilibrium_solvation is None: + raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDHF() return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation) - def CasidaTDDFT(self, equilibrium_solvation, eps_optical=1.78): + def CasidaTDDFT(self, equilibrium_solvation=None, eps_optical=1.78): + if equilibrium_solvation is None: + raise ValueError('equilibrium_solvation must be specified') tda_method = super().CasidaTDDFT() return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation) diff --git a/gpu4pyscf/solvent/grad/pcm.py b/gpu4pyscf/solvent/grad/pcm.py index 4d339c633..567b8b284 100644 --- a/gpu4pyscf/solvent/grad/pcm.py +++ b/gpu4pyscf/solvent/grad/pcm.py @@ -273,7 +273,7 @@ def grad_qv(pcmobj, dm, q_sym = None): t1 = log.timer_debug1('grad qv', *t1) return de.get() -def grad_solver(pcmobj, dm, v_grids = None, v_grids_1 = None, q = None): +def grad_solver(pcmobj, dm, v_grids = None, v_grids_l = None, q = None): ''' dE = 0.5*v* d(K^-1 R) *v + q*dv v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR v - dK q) @@ -292,8 +292,8 @@ def grad_solver(pcmobj, dm, v_grids = None, v_grids_1 = None, q = None): gridslice = pcmobj.surface['gslice_by_atom'] if v_grids is None: v_grids = pcmobj._intermediates['v_grids'] - if v_grids_1 is None: - v_grids_1 = pcmobj._intermediates['v_grids'] + if v_grids_l is None: + v_grids_l = pcmobj._intermediates['v_grids'] if q is None: q = pcmobj._intermediates['q'] f_epsilon = pcmobj._intermediates['f_epsilon'] @@ -302,7 +302,7 @@ def grad_solver(pcmobj, dm, v_grids = None, v_grids_1 = None, q = None): D = pcmobj._intermediates['D'] S = pcmobj._intermediates['S'] - vK_1 = pcmobj.left_solve_K(v_grids_1, K_transpose = True) + vK_1 = pcmobj.left_solve_K(v_grids_l, K_transpose = True) def contract_bra(a, B, c): ''' i,xij,j->jx ''' diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 5252da3c8..e132a8c57 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -470,26 +470,18 @@ def nuc_grad_method(self, grad_method): def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - if self.frozen: - raise RuntimeError('Frozen solvent model is not supported') return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - if self.frozen: - raise RuntimeError('Frozen solvent model is not supported') return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - if self.frozen: - raise RuntimeError('Frozen solvent model is not supported') return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - if self.frozen: - raise RuntimeError('Frozen solvent model is not supported') return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def Hessian(self, hess_method): @@ -516,7 +508,7 @@ def _B_dot_x(self, dms): out_shape = dms.shape nao = dms.shape[-1] dms = dms.reshape(-1,nao,nao) - + v_grids = -self._get_v(dms) b = self.left_multiply_R(v_grids.T) @@ -533,9 +525,8 @@ def _B_dot_x(self, dms): def if_method_in_CPCM_category(self): return self.method.upper() in ['C-PCM', 'CPCM', "COSMO"] - def left_multiply_R(self, right_vector, f_epsilon = None, R_transpose = False): - if f_epsilon is None: - f_epsilon = self._intermediates['f_epsilon'] + def left_multiply_R(self, right_vector, R_transpose = False): + f_epsilon = self._intermediates['f_epsilon'] if self.if_method_in_CPCM_category: # R = -f_epsilon * cupy.eye(K.shape[0]) return -f_epsilon * right_vector @@ -548,11 +539,9 @@ def left_multiply_R(self, right_vector, f_epsilon = None, R_transpose = False): DA = DA.T return -f_epsilon * (right_vector - 1.0/(2.0*PI) * cupy.dot(DA, right_vector)) - def left_solve_K(self, right_vector, K_LU = None, K_LU_pivot = None, K_transpose = False): + def left_solve_K(self, right_vector, K_transpose = False): ''' K^{-1} @ right_vector ''' - if K_LU is None: - K_LU = self._intermediates['K_LU'] - if K_LU_pivot is None: - K_LU_pivot = self._intermediates['K_LU_pivot'] + K_LU = self._intermediates['K_LU'] + K_LU_pivot = self._intermediates['K_LU_pivot'] return lu_solve((K_LU, K_LU_pivot), right_vector, trans = K_transpose, overwrite_b = False, check_finite = False) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index c2d722123..3b1431dde 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -69,7 +69,7 @@ def vind_with_solvent(dm1): elif singlet: v += pcmobj._B_dot_x(dm1) else: - logger.warn(pcmobj, 'Singlet-Triplet has no LR-PCM contribution!') + logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') return v return vind_with_solvent @@ -88,9 +88,9 @@ def undo_solvent(self): def _finalize(self): super()._finalize() if self.with_solvent.equilibrium_solvation: - logger.info(self.with_solvent, 'equilibrium solvation NOT for vertical excitation') + logger.info(self.with_solvent, 'equilibrium solvation NOT suitable for vertical excitation') else: - logger.info(self.with_solvent, 'Not equilibrium solvation NOT for vertical excitation,\n\ + logger.info(self.with_solvent, 'Non equilibrium solvation NOT suitable for adiabatic excitation,\n\ eps_optical = %s', self.with_solvent.eps) def nuc_grad_method(self): @@ -111,11 +111,12 @@ def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.INFO): de = super().grad_elec(xy, singlet, atmlst, verbose) assert self.base.with_solvent.equilibrium_solvation + if self.base.with_solvent.frozen: + raise RuntimeError('Frozen solvent model is not supported') dm = self.base._scf.make_rdm1(ao_repr=True) if dm.ndim == 3: dm = dm[0] + dm[1] - # TODO: add unrestricted case support dmP = 0.5 * (self.dmz1doo + self.dmz1doo.T) dmxpy = self.dmxpy + self.dmxpy.T pcmobj = self.base.with_solvent diff --git a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py index c4623b76c..986cc6f32 100644 --- a/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py +++ b/gpu4pyscf/solvent/tests/test_pcm_tdscf_grad.py @@ -117,7 +117,6 @@ def cal_analytic_gradient(mol, td, tdgrad, nocc, nvir, grad_elec, tda): atmlst = range(mol.natm) e_diag, xy_diag = diagonalize_tda(a) x = xy_diag[:, 0].reshape(nocc, nvir)*np.sqrt(0.5) - # de_td = grad_elec(tdgrad, (x, 0)) de_td = grad_elec((x, 0)) gradient_ana = de_td + tdgrad.grad_nuc(atmlst=atmlst) else: @@ -131,7 +130,6 @@ def cal_analytic_gradient(mol, td, tdgrad, nocc, nvir, grad_elec, tda): x = x.reshape(nocc, nvir) y = y.reshape(nocc, nvir) - # de_td = grad_elec(tdgrad, (x, y)) de_td = grad_elec((x, y)) gradient_ana = de_td + tdgrad.grad_nuc(atmlst=atmlst) From 905e7aa979899182d94d6f7009621f41f737fe1e Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 14 Apr 2025 17:17:17 +0800 Subject: [PATCH 17/31] in writting --- gpu4pyscf/solvent/_attach_solvent.py | 16 ++++++++-------- gpu4pyscf/solvent/pcm.py | 18 +++++++++--------- gpu4pyscf/solvent/tdscf/pcm.py | 27 ++++++++++++++------------- 3 files changed, 31 insertions(+), 30 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 84637bb2b..659df3c45 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -137,29 +137,29 @@ def nuc_grad_method(self): grad_method = super().nuc_grad_method() return self.with_solvent.nuc_grad_method(grad_method) - def TDA(self, equilibrium_solvation=None, eps_optical=1.78): + def TDA(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDA() - return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation) + return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation, linear_resposne) - def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78): + def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDDFT() - return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation) + return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation, linear_resposne) - def TDHF(self, equilibrium_solvation=None, eps_optical=1.78): + def TDHF(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDHF() - return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation) + return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation, linear_resposne) - def CasidaTDDFT(self, equilibrium_solvation=None, eps_optical=1.78): + def CasidaTDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().CasidaTDDFT() - return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation) + return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation, linear_resposne) Gradients = nuc_grad_method diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 2764e8316..3c54da230 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -469,21 +469,21 @@ def nuc_grad_method(self, grad_method): else: raise RuntimeError('Only SCF gradient is supported') - def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78): + def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) - def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78): + def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) - def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): + def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) - def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): + def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) def Hessian(self, hess_method): from gpu4pyscf.solvent.hessian import pcm as pcm_hess @@ -526,7 +526,7 @@ def _B_dot_x(self, dms): def if_method_in_CPCM_category(self): return self.method.upper() in ['C-PCM', 'CPCM', "COSMO"] - def left_multiply_R(self, right_vector, R_transpose = False): + def left_multiply_R(self, right_vector, f_epsilon = None, R_transpose = False): f_epsilon = self._intermediates['f_epsilon'] if self.if_method_in_CPCM_category: # R = -f_epsilon * cupy.eye(K.shape[0]) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index b37deba52..3a6a31a74 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -32,11 +32,11 @@ def __init__(self, mfpcmobj, eps_optical=1.78, equilium_solvation=False): self.eps = eps_optical -def make_tdscf_object(tda_method, eps_optical=1.78, equilibrium_solvation=False): +def make_tdscf_object(tda_method, eps_optical=1.78, equilibrium_solvation=False, linear_resposne=True): '''For td_method in vacuum, add td of solvent pcmobj''' name = (tda_method._scf.with_solvent.__class__.__name__ + tda_method.__class__.__name__) - return lib.set_class(WithSolventTDSCF(tda_method, eps_optical, equilibrium_solvation), + return lib.set_class(WithSolventTDSCF(tda_method, eps_optical, equilibrium_solvation, linear_resposne), (WithSolventTDSCF, tda_method.__class__), name) @@ -48,19 +48,20 @@ def make_tdscf_gradient_object(tda_grad_method): (WithSolventTDSCFGradient, tda_grad_method.__class__), name) -def add_prefix(prefix): - def decorator(func): - def wrapper(*args, ​**kwargs): - original_result = func(*args, **kwargs) - return f"[{prefix}] {original_result}" - return wrapper - return decorator +# def add_prefix(prefix): +# def decorator(func): +# def wrapper(*args, ​**kwargs): +# original_result = func(*args, **kwargs) +# return f"[{prefix}] {original_result}" +# return wrapper +# return decorator -def state_specific(td, x0=None, nstates=None): - td.kernel(x0=x0, nstates=nstates) - for icyc in range(50): - pass + +# def state_specific(td, x0=None, nstates=None): +# td.kernel(x0=x0, nstates=nstates) +# for icyc in range(50): +# pass # A.a = decorator(A.a) From 4024ae28290907088c977f257772dcaab19f12c6 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 14 Apr 2025 17:18:50 +0800 Subject: [PATCH 18/31] fix a typo --- gpu4pyscf/solvent/tdscf/pcm.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 3b1431dde..2be881418 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -53,6 +53,12 @@ class WithSolventTDSCF: _keys = {'with_solvent'} + def __init__(self, tda_method, eps_optical=1.78, equilibrium_solvation=False): + self.__dict__.update(tda_method.__dict__) + self.with_solvent = TDPCM(tda_method._scf.with_solvent, eps_optical, equilibrium_solvation) + if not self.with_solvent.equilibrium_solvation: + self.with_solvent.build() + def gen_response(self, *args, **kwargs): pcmobj = self.with_solvent mf = self._scf @@ -73,12 +79,6 @@ def vind_with_solvent(dm1): return v return vind_with_solvent - def __init__(self, tda_method, eps_optical=1.78, equilibrium_solvation=False): - self.__dict__.update(tda_method.__dict__) - self.with_solvent = TDPCM(tda_method._scf.with_solvent, eps_optical, equilibrium_solvation) - if not self.with_solvent.equilibrium_solvation: - self.with_solvent.build() - def undo_solvent(self): cls = self.__class__ name_mixin = self.base.with_solvent.__class__.__name__ From 55a70bcd7465f537a33680537d9ef551139f2228 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Tue, 15 Apr 2025 08:28:20 +0800 Subject: [PATCH 19/31] in writting --- gpu4pyscf/solvent/_attach_solvent.py | 16 ++++++++-------- gpu4pyscf/solvent/pcm.py | 16 ++++++++-------- gpu4pyscf/solvent/tdscf/pcm.py | 26 ++++++++++++++------------ 3 files changed, 30 insertions(+), 28 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 659df3c45..354f7c72e 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -137,29 +137,29 @@ def nuc_grad_method(self): grad_method = super().nuc_grad_method() return self.with_solvent.nuc_grad_method(grad_method) - def TDA(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): + def TDA(self, equilibrium_solvation=None, eps_optical=1.78, linear_response=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDA() - return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation, linear_resposne) + return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation, linear_response) - def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): + def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_response=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDDFT() - return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation, linear_resposne) + return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation, linear_response) - def TDHF(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): + def TDHF(self, equilibrium_solvation=None, eps_optical=1.78, linear_response=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().TDHF() - return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation, linear_resposne) + return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation, linear_response) - def CasidaTDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_resposne=True): + def CasidaTDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_response=True): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') tda_method = super().CasidaTDDFT() - return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation, linear_resposne) + return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation, linear_response) Gradients = nuc_grad_method diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 3c54da230..f2f92584d 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -469,21 +469,21 @@ def nuc_grad_method(self, grad_method): else: raise RuntimeError('Only SCF gradient is supported') - def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): + def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_response=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_response) - def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): + def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_response=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_response) - def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): + def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_response=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_response) - def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_resposne=True): + def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78, linear_response=True): from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_resposne) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_response) def Hessian(self, hess_method): from gpu4pyscf.solvent.hessian import pcm as pcm_hess diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 0548a3841..c9d78bd25 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -32,11 +32,11 @@ def __init__(self, mfpcmobj, eps_optical=1.78, equilium_solvation=False): self.eps = eps_optical -def make_tdscf_object(tda_method, eps_optical=1.78, equilibrium_solvation=False, linear_resposne=True): +def make_tdscf_object(tda_method, eps_optical=1.78, equilibrium_solvation=False, linear_response=True): '''For td_method in vacuum, add td of solvent pcmobj''' name = (tda_method._scf.with_solvent.__class__.__name__ + tda_method.__class__.__name__) - return lib.set_class(WithSolventTDSCF(tda_method, eps_optical, equilibrium_solvation, linear_resposne), + return lib.set_class(WithSolventTDSCF(tda_method, eps_optical, equilibrium_solvation, linear_response), (WithSolventTDSCF, tda_method.__class__), name) @@ -68,12 +68,13 @@ def make_tdscf_gradient_object(tda_grad_method): class WithSolventTDSCF: from gpu4pyscf.lib.utils import to_gpu, device - _keys = {'with_solvent'} + _keys = {'with_solvent', 'linear_response'} - def __init__(self, tda_method, eps_optical=1.78, equilibrium_solvation=False): + def __init__(self, tda_method, eps_optical=1.78, equilibrium_solvation=False, linear_response=True): self.__dict__.update(tda_method.__dict__) self.with_solvent = TDPCM(tda_method._scf.with_solvent, eps_optical, equilibrium_solvation) - if not self.with_solvent.equilibrium_solvation: + self.linear_response = linear_response + if not self.with_solvent.equilibrium_solvation and linear_response: self.with_solvent.build() def gen_response(self, *args, **kwargs): @@ -86,13 +87,14 @@ def gen_response(self, *args, **kwargs): singlet = singlet or singlet is None def vind_with_solvent(dm1): v = vind(dm1) - if is_uhf: - v_solvent = pcmobj._B_dot_x(dm1) - v += v_solvent[0] + v_solvent[1] - elif singlet: - v += pcmobj._B_dot_x(dm1) - else: - logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') + if self.linear_response: + if is_uhf: + v_solvent = pcmobj._B_dot_x(dm1) + v += v_solvent[0] + v_solvent[1] + elif singlet: + v += pcmobj._B_dot_x(dm1) + else: + logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') return v return vind_with_solvent From 4bdcfcc11dde043d9103a454f2cb6aa9e253ff5b Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Tue, 15 Apr 2025 08:37:51 +0800 Subject: [PATCH 20/31] fix a typo --- gpu4pyscf/solvent/tdscf/pcm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 2be881418..e5a905ca9 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -130,12 +130,12 @@ def grad_elec(self, xy, singlet=None, atmlst=None, verbose=logger.INFO): de += grad_qv(pcmobj, dm, q_sym = qE_sym_dmP) de += grad_nuc(pcmobj, dm, q_sym = qE_sym_dmP.get()) de += grad_qv(pcmobj, dmP, q_sym = q_sym_dm) - v_grids_1 = pcmobj._get_vgrids(dmP, with_nuc = False) - de += grad_solver(pcmobj, dm, v_grids_1 = v_grids_1) * 2.0 + v_grids_l = pcmobj._get_vgrids(dmP, with_nuc = False) + de += grad_solver(pcmobj, dm, v_grids_l = v_grids_l) * 2.0 de += grad_qv(pcmobj, dmxpy, q_sym = qE_sym_dmxpy) * 2.0 v_grids = pcmobj._get_vgrids(dmxpy, with_nuc = False) q = pcmobj._get_qsym(dmxpy, with_nuc = False)[1] - de += grad_solver(pcmobj, dmxpy, v_grids=v_grids, v_grids_1=v_grids, q=q) * 2.0 + de += grad_solver(pcmobj, dmxpy, v_grids=v_grids, v_grids_l=v_grids, q=q) * 2.0 return de From 66394307cec52437745200a026ba54513db46abd Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Tue, 15 Apr 2025 12:56:41 +0800 Subject: [PATCH 21/31] in writting --- gpu4pyscf/solvent/_attach_solvent.py | 10 ++++++++++ gpu4pyscf/solvent/pcm.py | 28 +++++++++++++++++++++------- gpu4pyscf/solvent/tdscf/pcm.py | 2 +- 3 files changed, 32 insertions(+), 8 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 354f7c72e..2b2077232 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -85,6 +85,15 @@ def get_veff(self, mol=None, dm_or_wfn=None, *args, **kwargs): else: dm = dm_or_wfn with_solvent.e, with_solvent.v = with_solvent.kernel(dm) + else: + if with_solvent.state_specific: + if dm_or_wfn is None: + dm = self.make_rdm1() + elif isinstance(dm_or_wfn, WaveFunction): + dm = dm_or_wfn.make_rdm1() + else: + dm = dm_or_wfn + # with_solvent.e, with_solvent.v = with_solvent.kernel(dm) e_solvent, v_solvent = with_solvent.e, with_solvent.v if veff.shape[-1] != v_solvent.shape[-1]: # lowmem mode, only lower triangular part of Fock matrix is stored @@ -95,6 +104,7 @@ def get_veff(self, mol=None, dm_or_wfn=None, *args, **kwargs): veff = lib.tag_array(veff, e_solvent=e_solvent, v_solvent=v_solvent) else: veff = tag_array(veff, e_solvent=e_solvent, v_solvent=v_solvent) + print(f"Solvent Energy = {e_solvent}") return veff def get_fock(self, h1e=None, s1e=None, vhf=None, dm_or_wfn=None, cycle=-1, diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index f2f92584d..97ec7f46f 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -245,7 +245,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'v_grids_n', 'state_specific' + 'equilibrium_solvation', 'e', 'v', 'v_grids_n', 'state_specific', 'dm_prev' } from gpu4pyscf.lib.utils import to_gpu, device @@ -265,6 +265,7 @@ def __init__(self, mol): self._intermediates = {} self.eps = 78.3553 self.state_specific = False + self.dm_prev = None self.max_cycle = 20 self.conv_tol = 1e-7 @@ -387,13 +388,26 @@ def _get_vind(self, dms): v_grids_e = self._get_v(dms) v_grids = self.v_grids_n - v_grids_e - b = self.left_multiply_R(v_grids.T) - q = self.left_solve_K(b).T - - vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) - qt = self.left_multiply_R(vK_1, R_transpose = True).T - q_sym = (q + qt)/2.0 + if self.dm_prev is not None: + if not isinstance(self.dm_prev, cupy.ndarray): + self.dm_prev = cupy.asarray(self.dm_prev) + self.dm_prev = self.dm_prev.reshape(-1,nao,nao) + if self.dm_prev.shape[0] == 2: + self.dm_prev = (self.dm_prev[0] + self.dm_prev[1]).reshape(-1,nao,nao) + v_grids_e_prev = self._get_v(self.dm_prev) + vgrids_prev = self.v_grids_n - v_grids_e_prev + b = self.left_multiply_R(vgrids_prev.T) + q = self.left_solve_K(b).T + vK_1 = self.left_solve_K(vgrids_prev.T, K_transpose = True) + qt = self.left_multiply_R(vK_1, R_transpose = True).T + q_sym = (q + qt)/2.0 + else: + b = self.left_multiply_R(v_grids.T) + q = self.left_solve_K(b).T + vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) + qt = self.left_multiply_R(vK_1, R_transpose = True).T + q_sym = (q + qt)/2.0 vmat = self._get_vmat(q_sym) epcm = 0.5 * cupy.dot(v_grids[0], q_sym[0]) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index ac7956d9a..c49201497 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -107,7 +107,7 @@ def undo_solvent(self): def kernel(self, *args, **kwargs): pcmobj = self._scf.with_solvent if pcmobj.state_specific: - pass + return super().kernel(*args, **kwargs) else: return super().kernel(*args, **kwargs) From cf35527a21c5ce73e8d4ed1a69ec47c2fa4c7439 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Wed, 16 Apr 2025 13:38:14 +0800 Subject: [PATCH 22/31] in writting --- gpu4pyscf/solvent/_attach_solvent.py | 11 +---------- gpu4pyscf/solvent/pcm.py | 28 ++++++---------------------- gpu4pyscf/solvent/tdscf/pcm.py | 12 ++++-------- 3 files changed, 11 insertions(+), 40 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 2b2077232..a4eb7db43 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -85,15 +85,6 @@ def get_veff(self, mol=None, dm_or_wfn=None, *args, **kwargs): else: dm = dm_or_wfn with_solvent.e, with_solvent.v = with_solvent.kernel(dm) - else: - if with_solvent.state_specific: - if dm_or_wfn is None: - dm = self.make_rdm1() - elif isinstance(dm_or_wfn, WaveFunction): - dm = dm_or_wfn.make_rdm1() - else: - dm = dm_or_wfn - # with_solvent.e, with_solvent.v = with_solvent.kernel(dm) e_solvent, v_solvent = with_solvent.e, with_solvent.v if veff.shape[-1] != v_solvent.shape[-1]: # lowmem mode, only lower triangular part of Fock matrix is stored @@ -104,7 +95,7 @@ def get_veff(self, mol=None, dm_or_wfn=None, *args, **kwargs): veff = lib.tag_array(veff, e_solvent=e_solvent, v_solvent=v_solvent) else: veff = tag_array(veff, e_solvent=e_solvent, v_solvent=v_solvent) - print(f"Solvent Energy = {e_solvent}") + # print(f"Solvent Energy = {e_solvent}") return veff def get_fock(self, h1e=None, s1e=None, vhf=None, dm_or_wfn=None, cycle=-1, diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index 97ec7f46f..c43ac66e6 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -245,7 +245,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'v_grids_n', 'state_specific', 'dm_prev' + 'equilibrium_solvation', 'e', 'v', 'v_grids_n' } from gpu4pyscf.lib.utils import to_gpu, device @@ -264,8 +264,6 @@ def __init__(self, mol): self.lebedev_order = 29 self._intermediates = {} self.eps = 78.3553 - self.state_specific = False - self.dm_prev = None self.max_cycle = 20 self.conv_tol = 1e-7 @@ -388,26 +386,12 @@ def _get_vind(self, dms): v_grids_e = self._get_v(dms) v_grids = self.v_grids_n - v_grids_e - if self.dm_prev is not None: - if not isinstance(self.dm_prev, cupy.ndarray): - self.dm_prev = cupy.asarray(self.dm_prev) - self.dm_prev = self.dm_prev.reshape(-1,nao,nao) - if self.dm_prev.shape[0] == 2: - self.dm_prev = (self.dm_prev[0] + self.dm_prev[1]).reshape(-1,nao,nao) - v_grids_e_prev = self._get_v(self.dm_prev) - vgrids_prev = self.v_grids_n - v_grids_e_prev - b = self.left_multiply_R(vgrids_prev.T) - q = self.left_solve_K(b).T - vK_1 = self.left_solve_K(vgrids_prev.T, K_transpose = True) - qt = self.left_multiply_R(vK_1, R_transpose = True).T - q_sym = (q + qt)/2.0 - else: - b = self.left_multiply_R(v_grids.T) - q = self.left_solve_K(b).T + b = self.left_multiply_R(v_grids.T) + q = self.left_solve_K(b).T - vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) - qt = self.left_multiply_R(vK_1, R_transpose = True).T - q_sym = (q + qt)/2.0 + vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) + qt = self.left_multiply_R(vK_1, R_transpose = True).T + q_sym = (q + qt)/2.0 vmat = self._get_vmat(q_sym) epcm = 0.5 * cupy.dot(v_grids[0], q_sym[0]) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index c49201497..093065265 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -25,10 +25,10 @@ class TDPCM(PCM): - def __init__(self, mfpcmobj, eps_optical=1.78, equilium_solvation=False): + def __init__(self, mfpcmobj, eps_optical=1.78, equilibrium_solvation=False): self.__dict__.update(mfpcmobj.__dict__) - self.equilibrium_solvation = equilium_solvation - if not equilium_solvation: + self.equilibrium_solvation = equilibrium_solvation + if not equilibrium_solvation: self.eps = eps_optical @@ -105,11 +105,7 @@ def undo_solvent(self): return obj def kernel(self, *args, **kwargs): - pcmobj = self._scf.with_solvent - if pcmobj.state_specific: - return super().kernel(*args, **kwargs) - else: - return super().kernel(*args, **kwargs) + super().kernel(*args, **kwargs) def _finalize(self): super()._finalize() From 001fa3f22ec1a1dc9c520436704f5af0e2bdf01b Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Wed, 16 Apr 2025 15:53:42 +0800 Subject: [PATCH 23/31] polarizability test --- gpu4pyscf/properties/polarizability.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/gpu4pyscf/properties/polarizability.py b/gpu4pyscf/properties/polarizability.py index 7949b4f5a..ec4b4712e 100644 --- a/gpu4pyscf/properties/polarizability.py +++ b/gpu4pyscf/properties/polarizability.py @@ -49,7 +49,7 @@ def fx(mo1): return fx -def eval_polarizability(mf): +def eval_polarizability(mf, max_cycle=20, tol=1e-10): """main function to calculate the polarizability Args: @@ -77,7 +77,7 @@ def eval_polarizability(mf): h1 = cupy.array(h1) for idirect in range(3): h1ai = -mvir.T.conj()@h1[idirect]@mocc - mo1 = cphf.solve(fx, mo_energy, mo_occ, h1ai, max_cycle=20, tol=1e-10)[0] + mo1 = cphf.solve(fx, mo_energy, mo_occ, h1ai, max_cycle=max_cycle, tol=tol)[0] for jdirect in range(idirect, 3): p10 = np.trace(mo1.conj().T@mvir.conj().T@h1[jdirect]@mocc)*2 p01 = np.trace(mocc.conj().T@h1[jdirect]@mvir@mo1)*2 @@ -87,4 +87,3 @@ def eval_polarizability(mf): polarizability[2, 1] = polarizability[1, 2] return polarizability - From 915f255e17c9f84e0b5da5ae14703bc39b157cbb Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 17 Apr 2025 14:28:44 +0800 Subject: [PATCH 24/31] change codes, according to comments --- gpu4pyscf/solvent/_attach_solvent.py | 24 ++++++++++++++---------- gpu4pyscf/solvent/pcm.py | 16 ---------------- gpu4pyscf/solvent/tdscf/pcm.py | 14 ++++++++++---- 3 files changed, 24 insertions(+), 30 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index 84637bb2b..3cd0fa814 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -140,26 +140,30 @@ def nuc_grad_method(self): def TDA(self, equilibrium_solvation=None, eps_optical=1.78): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') - tda_method = super().TDA() - return self.with_solvent.TDA(tda_method, eps_optical, equilibrium_solvation) + td = super().TDA() + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') - tda_method = super().TDDFT() - return self.with_solvent.TDDFT(tda_method, eps_optical, equilibrium_solvation) + td = super().TDDFT() + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def TDHF(self, equilibrium_solvation=None, eps_optical=1.78): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') - tda_method = super().TDHF() - return self.with_solvent.TDHF(tda_method, eps_optical, equilibrium_solvation) + td = super().TDHF() + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def CasidaTDDFT(self, equilibrium_solvation=None, eps_optical=1.78): if equilibrium_solvation is None: raise ValueError('equilibrium_solvation must be specified') - tda_method = super().CasidaTDDFT() - return self.with_solvent.CasidaTDDFT(tda_method, eps_optical, equilibrium_solvation) + td = super().CasidaTDDFT() + from gpu4pyscf.solvent.tdscf import pcm as pcm_td + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) Gradients = nuc_grad_method @@ -177,8 +181,8 @@ def vind_with_solvent(dm1): v = vind(dm1) if self.with_solvent.equilibrium_solvation: if is_uhf: - v_solvent = self.with_solvent._B_dot_x(dm1) - v += v_solvent[0] + v_solvent[1] + v_solvent = self.with_solvent._B_dot_x(dm1[0]+dm1[1]) + v += v_solvent elif singlet: v += self.with_solvent._B_dot_x(dm1) return v diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index e132a8c57..39301d505 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -467,22 +467,6 @@ def nuc_grad_method(self, grad_method): return pcm_grad.make_grad_object(grad_method) else: raise RuntimeError('Only SCF gradient is supported') - - def TDA(self, td, equilibrium_solvation=False, eps_optical=1.78): - from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) - - def TDHF(self, td, equilibrium_solvation=False, eps_optical=1.78): - from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) - - def TDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): - from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) - - def CasidaTDDFT(self, td, equilibrium_solvation=False, eps_optical=1.78): - from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) def Hessian(self, hess_method): from gpu4pyscf.solvent.hessian import pcm as pcm_hess diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index e5a905ca9..c66477f0b 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -32,7 +32,7 @@ def __init__(self, mfpcmobj, eps_optical=1.78, equilium_solvation=False): self.eps = eps_optical -def make_tdscf_object(tda_method, eps_optical=1.78, equilibrium_solvation=False): +def make_tdscf_object(tda_method, equilibrium_solvation=False, eps_optical=1.78): '''For td_method in vacuum, add td of solvent pcmobj''' name = (tda_method._scf.with_solvent.__class__.__name__ + tda_method.__class__.__name__) @@ -70,10 +70,16 @@ def gen_response(self, *args, **kwargs): def vind_with_solvent(dm1): v = vind(dm1) if is_uhf: - v_solvent = pcmobj._B_dot_x(dm1) - v += v_solvent[0] + v_solvent[1] + v_solvent = pcmobj._B_dot_x(dm1[0]+dm1[1]) + if self._scf.with_solvent.equilibrium_solvation: + pass + else: + v += v_solvent elif singlet: - v += pcmobj._B_dot_x(dm1) + if self._scf.with_solvent.equilibrium_solvation: + pass + else: + v += pcmobj._B_dot_x(dm1) else: logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') return v From 34af53a0819f106a92ba0fbcc99bdabc7486dce8 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 17 Apr 2025 14:36:40 +0800 Subject: [PATCH 25/31] fix a typo --- gpu4pyscf/solvent/tdscf/pcm.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index f83184b42..816749e09 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -90,14 +90,10 @@ def vind_with_solvent(dm1): if self.linear_response: if is_uhf: v_solvent = pcmobj._B_dot_x(dm1[0]+dm1[1]) - if self._scf.with_solvent.equilibrium_solvation: - pass - else: + if not self._scf.with_solvent.equilibrium_solvation: v += v_solvent elif singlet: - if self._scf.with_solvent.equilibrium_solvation: - pass - else: + if not self._scf.with_solvent.equilibrium_solvation: v += pcmobj._B_dot_x(dm1) else: logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') From a5afe1cb5d89c6190e827f02ec015586cd1dc9fb Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 17 Apr 2025 14:42:52 +0800 Subject: [PATCH 26/31] fix a typo : --- gpu4pyscf/solvent/tdscf/pcm.py | 19 +------------------ 1 file changed, 1 insertion(+), 18 deletions(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 816749e09..2d1f5e214 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -32,7 +32,7 @@ def __init__(self, mfpcmobj, eps_optical=1.78, equilibrium_solvation=False): self.eps = eps_optical -def make_tdscf_object(tda_method, equilibrium_solvation=False, eps_optical=1.78): +def make_tdscf_object(tda_method, equilibrium_solvation=False, eps_optical=1.78, linear_response=True): '''For td_method in vacuum, add td of solvent pcmobj''' name = (tda_method._scf.with_solvent.__class__.__name__ + tda_method.__class__.__name__) @@ -48,23 +48,6 @@ def make_tdscf_gradient_object(tda_grad_method): (WithSolventTDSCFGradient, tda_grad_method.__class__), name) - -# def add_prefix(prefix): -# def decorator(func): -# def wrapper(*args, ​**kwargs): -# original_result = func(*args, **kwargs) -# return f"[{prefix}] {original_result}" -# return wrapper -# return decorator - - -# def state_specific(td, x0=None, nstates=None): -# td.kernel(x0=x0, nstates=nstates) -# for icyc in range(50): -# pass - # A.a = decorator(A.a) - - class WithSolventTDSCF: from gpu4pyscf.lib.utils import to_gpu, device From 7f334bd2e291cc8ac90b0c4e03d7efccd9dce51f Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Thu, 17 Apr 2025 17:43:02 +0800 Subject: [PATCH 27/31] in writting --- gpu4pyscf/solvent/_attach_solvent.py | 2 +- gpu4pyscf/solvent/pcm.py | 30 ++++++++++++++++++++++------ gpu4pyscf/solvent/tdscf/pcm.py | 3 ++- 3 files changed, 27 insertions(+), 8 deletions(-) diff --git a/gpu4pyscf/solvent/_attach_solvent.py b/gpu4pyscf/solvent/_attach_solvent.py index d02c4188d..492913a17 100644 --- a/gpu4pyscf/solvent/_attach_solvent.py +++ b/gpu4pyscf/solvent/_attach_solvent.py @@ -143,7 +143,7 @@ def TDA(self, equilibrium_solvation=None, eps_optical=1.78, linear_response=True raise ValueError('equilibrium_solvation must be specified') td = super().TDA() from gpu4pyscf.solvent.tdscf import pcm as pcm_td - return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical) + return pcm_td.make_tdscf_object(td, equilibrium_solvation, eps_optical, linear_response) def TDDFT(self, equilibrium_solvation=None, eps_optical=1.78, linear_response=True): if equilibrium_solvation is None: diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index d05d28b46..a2b97c5c0 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -245,7 +245,7 @@ class PCM(lib.StreamObject): 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', 'mol', 'radii_table', 'atom_radii', 'lebedev_order', 'lmax', 'eta', 'eps', 'grids', 'max_cycle', 'conv_tol', 'state_id', 'frozen', - 'equilibrium_solvation', 'e', 'v', 'v_grids_n' + 'equilibrium_solvation', 'e', 'v', 'v_grids_n', 'dmprev' } from gpu4pyscf.lib.utils import to_gpu, device @@ -275,6 +275,7 @@ def __init__(self, mol): self.e = None self.v = None self.v_grids_n = None + self.dmprev = None def dump_flags(self, verbose=None): logger.info(self, '******** %s ********', self.__class__) @@ -383,15 +384,32 @@ def _get_vind(self, dms): dms = (dms[0] + dms[1]).reshape(-1,nao,nao) if not isinstance(dms, cupy.ndarray): dms = cupy.asarray(dms) + v_grids_e = self._get_v(dms) v_grids = self.v_grids_n - v_grids_e - b = self.left_multiply_R(v_grids.T) - q = self.left_solve_K(b).T + if self.dmprev is not None: + dmprev = self.dmprev + dmprev = dmprev.reshape(-1,nao,nao) + if dmprev.shape[0] == 2: + dmprev = (dmprev[0] + dmprev[1]).reshape(-1,nao,nao) + if not isinstance(dmprev, cupy.ndarray): + dmprev = cupy.asarray(dmprev) + v_grids_e_prev = self._get_v(dmprev) + v_grids_prev = self.v_grids_n - v_grids_e_prev + b = self.left_multiply_R(v_grids_prev.T) + q = self.left_solve_K(b).T + + vK_1 = self.left_solve_K(v_grids_prev.T, K_transpose = True) + qt = self.left_multiply_R(vK_1, R_transpose = True).T + q_sym = (q + qt)/2.0 + else: + b = self.left_multiply_R(v_grids.T) + q = self.left_solve_K(b).T - vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) - qt = self.left_multiply_R(vK_1, R_transpose = True).T - q_sym = (q + qt)/2.0 + vK_1 = self.left_solve_K(v_grids.T, K_transpose = True) + qt = self.left_multiply_R(vK_1, R_transpose = True).T + q_sym = (q + qt)/2.0 vmat = self._get_vmat(q_sym) epcm = 0.5 * cupy.dot(v_grids[0], q_sym[0]) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index 2d1f5e214..b4de8325a 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -68,6 +68,7 @@ def gen_response(self, *args, **kwargs): # singlet=None is orbital hessian or CPHF type response function singlet = kwargs.get('singlet', True) singlet = singlet or singlet is None + assert not self._scf.with_solvent.equilibrium_solvation def vind_with_solvent(dm1): v = vind(dm1) if self.linear_response: @@ -79,7 +80,7 @@ def vind_with_solvent(dm1): if not self._scf.with_solvent.equilibrium_solvation: v += pcmobj._B_dot_x(dm1) else: - logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') + logger.warn(pcmobj, 'Singlet-Triplet excitation has no LR-PCM contribution!') return v return vind_with_solvent From 480f6b540f9a4c5314f9b2bb1f12d5b6c41de3f5 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Mon, 21 Apr 2025 17:50:01 +0800 Subject: [PATCH 28/31] add an example for SS --- examples/29-tddft_sspcm.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) create mode 100644 examples/29-tddft_sspcm.py diff --git a/examples/29-tddft_sspcm.py b/examples/29-tddft_sspcm.py new file mode 100644 index 000000000..57d126711 --- /dev/null +++ b/examples/29-tddft_sspcm.py @@ -0,0 +1,22 @@ +#!/usr/bin/env python +# Copyright 2021-2025 The PySCF Developers. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +################################### +# Example of TDDFT with State-Specific Solvent +################################### + +""" + +""" \ No newline at end of file From b9b4d66a0f6e1d8044b031c362f7b14365363861 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Tue, 22 Apr 2025 13:40:24 +0800 Subject: [PATCH 29/31] fix a typo --- gpu4pyscf/solvent/tdscf/pcm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/gpu4pyscf/solvent/tdscf/pcm.py b/gpu4pyscf/solvent/tdscf/pcm.py index b4de8325a..0376e44a7 100644 --- a/gpu4pyscf/solvent/tdscf/pcm.py +++ b/gpu4pyscf/solvent/tdscf/pcm.py @@ -91,7 +91,7 @@ def undo_solvent(self): return obj def kernel(self, *args, **kwargs): - super().kernel(*args, **kwargs) + return super().kernel(*args, **kwargs) def _finalize(self): super()._finalize() From 130a9ad6eb169f7717f0d8c36de1380bf854befc Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Wed, 23 Apr 2025 08:25:53 +0800 Subject: [PATCH 30/31] in writting --- examples/29-tddft_sspcm.py | 158 ++++++++++++++++++++++++++++++++++++- gpu4pyscf/solvent/pcm.py | 53 +++++++++++++ 2 files changed, 210 insertions(+), 1 deletion(-) diff --git a/examples/29-tddft_sspcm.py b/examples/29-tddft_sspcm.py index 57d126711..0f5e1f253 100644 --- a/examples/29-tddft_sspcm.py +++ b/examples/29-tddft_sspcm.py @@ -18,5 +18,161 @@ ################################### """ +[Experimental Feature Notice] +----------------------------------------------------------------------- +This implementation is EXPERIMENTAL and subject to change without warning. +Users must fully understand the theoretical assumptions and limitations +before application in production systems. Misuse may lead to incorrect results. +""" +import numpy as np +import pyscf +import cupy as cp +import gpu4pyscf +from pyscf import gto, scf, dft, tddft +from gpu4pyscf.solvent._attach_solvent import SCFWithSolvent +from gpu4pyscf.lib.cupy_helper import tag_array, pack_tril +from gpu4pyscf.scf.hf_lowmem import WaveFunction +from gpu4pyscf.lib.cupy_helper import contract +from pyscf import lib +import copy +from functools import reduce +from gpu4pyscf.scf import cphf + +def get_total_density(td_grad, mf, x_y, singlet=True, relaxed=True): + mol = mf.mol + mo_coeff = cp.asarray(mf.mo_coeff) + mo_energy = cp.asarray(mf.mo_energy) + mo_occ = cp.asarray(mf.mo_occ) + nmo = mo_coeff.shape[1] + nocc = int((mo_occ > 0).sum()) + nvir = nmo - nocc + orbv = mo_coeff[:, nocc:] + orbo = mo_coeff[:, :nocc] + x, y = x_y + x = cp.asarray(x) + y = cp.asarray(y) + xpy = (x + y).reshape(nocc, nvir).T + xmy = (x - y).reshape(nocc, nvir).T + dvv = contract("ai,bi->ab", xpy, xpy) + contract("ai,bi->ab", xmy, xmy) # 2 T_{ab} + doo = -contract("ai,aj->ij", xpy, xpy) - contract("ai,aj->ij", xmy, xmy) # 2 T_{ij} + dmxpy = reduce(cp.dot, (orbv, xpy, orbo.T)) # (X+Y) in ao basis + dmxmy = reduce(cp.dot, (orbv, xmy, orbo.T)) # (X-Y) in ao basis + dmzoo = reduce(cp.dot, (orbo, doo, orbo.T)) # T_{ij}*2 in ao basis + dmzoo += reduce(cp.dot, (orbv, dvv, orbv.T)) # T_{ij}*2 + T_{ab}*2 in ao basis + vj0, vk0 = mf.get_jk(mol, dmzoo, hermi=0) + vj1, vk1 = mf.get_jk(mol, dmxpy + dmxpy.T, hermi=0) + vj2, vk2 = mf.get_jk(mol, dmxmy - dmxmy.T, hermi=0) + if not isinstance(vj0, cp.ndarray): + vj0 = cp.asarray(vj0) + if not isinstance(vk0, cp.ndarray): + vk0 = cp.asarray(vk0) + if not isinstance(vj1, cp.ndarray): + vj1 = cp.asarray(vj1) + if not isinstance(vk1, cp.ndarray): + vk1 = cp.asarray(vk1) + if not isinstance(vj2, cp.ndarray): + vj2 = cp.asarray(vj2) + if not isinstance(vk2, cp.ndarray): + vk2 = cp.asarray(vk2) + vj = cp.stack((vj0, vj1, vj2)) + vk = cp.stack((vk0, vk1, vk2)) + veff0doo = vj[0] * 2 - vk[0] # 2 for alpha and beta + veff0doo += td_grad.solvent_response(dmzoo) + wvo = reduce(cp.dot, (orbv.T, veff0doo, orbo)) * 2 + if singlet: + veff = vj[1] * 2 - vk[1] + else: + veff = -vk[1] + veff += td_grad.solvent_response(dmxpy + dmxpy.T) + veff0mop = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff)) + wvo -= contract("ki,ai->ak", veff0mop[:nocc, :nocc], xpy) * 2 # 2 for dm + dm.T + wvo += contract("ac,ai->ci", veff0mop[nocc:, nocc:], xpy) * 2 + veff = -vk[2] + veff0mom = reduce(cp.dot, (mo_coeff.T, veff, mo_coeff)) + wvo -= contract("ki,ai->ak", veff0mom[:nocc, :nocc], xmy) * 2 + wvo += contract("ac,ai->ci", veff0mom[nocc:, nocc:], xmy) * 2 + + # set singlet=None, generate function for CPHF type response kernel + vresp = td_grad.base.gen_response(singlet=None, hermi=1) + + def fvind(x): # For singlet, closed shell ground state + dm = reduce(cp.dot, (orbv, x.reshape(nvir, nocc) * 2, orbo.T)) # 2 for double occupancy + v1ao = vresp(dm + dm.T) # for the upused 2 + return reduce(cp.dot, (orbv.T, v1ao, orbo)).ravel() + + z1 = cphf.solve( + fvind, + mo_energy, + mo_occ, + wvo, + max_cycle=td_grad.cphf_max_cycle, + tol=td_grad.cphf_conv_tol)[0] + z1 = z1.reshape(nvir, nocc) + + z1ao = reduce(cp.dot, (orbv, z1, orbo.T)) + + if relaxed: + dmz1doo = z1ao + dmzoo + else: + dmz1doo = dmzoo + return (dmz1doo + dmz1doo.T) * 0.5 + mf.make_rdm1() + + +def get_phi(pcmobj, sigma): + mol = pcmobj.mol + int2c2e = mol._add_suffix('int2c2e') + grid_coords = pcmobj.surface['grid_coords'] + charge_exp = pcmobj.surface['charge_exp'] + fakemol_charge = gto.fakemol_for_charges(grid_coords.get(), expnt=charge_exp.get()**2) + v_ng = gto.mole.intor_cross(int2c2e, fakemol_charge, fakemol_charge) + v_ng = cp.asarray(v_ng) + phi_s = sigma@v_ng + return phi_s + + +def get_deltaG(mfeq, mfneq, tdneq, tdgneq, eps_optical=1.78): + pcmobj = mfeq.with_solvent + eps = pcmobj.eps + + pcmobj_optical = copy.copy(pcmobj) + pcmobj_optical.dmprev = None + pcmobj_optical.eps = eps_optical + pcmobj_optical.build() + dm1_eq = mfeq.make_rdm1() + + dm2 = get_total_density(tdgneq, mfneq, tdneq.xy[0]) + chi_e = (eps - 1)/(4 * np.pi) + chi_fast = (eps_optical - 1)/(4 * np.pi) + chi_slow = (eps - eps_optical)/(4 * np.pi) + q_1 = pcmobj._get_qsym(dm1_eq, with_nuc=True)[0] + q_1_s = chi_slow/chi_e * q_1 + q_2 = pcmobj._get_qsym(dm2, with_nuc=True)[0] + q_2_s = chi_slow/chi_e * q_2 + v1 = pcmobj._get_vgrids(dm1_eq, with_nuc=True) + v2 = pcmobj._get_vgrids(dm2, with_nuc=True) + phi_1_s = get_phi(pcmobj, q_1_s) + vgrids_1 = v1 + phi_1_s + vgrids_2 = v2 + phi_1_s + b = pcmobj_optical.left_multiply_R(vgrids_1.T) + q = pcmobj_optical.left_solve_K(b).T + vK_1 = pcmobj_optical.left_solve_K(vgrids_1.T, K_transpose = True) + qt = pcmobj_optical.left_multiply_R(vK_1, R_transpose = True).T + q_1_f = (q + qt)/2.0 + b = pcmobj_optical.left_multiply_R(vgrids_2.T) + q = pcmobj_optical.left_solve_K(b).T + vK_1 = pcmobj_optical.left_solve_K(vgrids_2.T, K_transpose = True) + qt = pcmobj_optical.left_multiply_R(vK_1, R_transpose = True).T + q_2_f = (q + qt)/2.0 + v2_rho = pcmobj._get_vgrids(dm2) + v1_rho = pcmobj._get_vgrids(dm1_eq) + phi2_f = get_phi(pcmobj, q_2_f) + phi1_f = get_phi(pcmobj, q_1_f) + + delta_G = 0.5*q_2_f.T@v2_rho + q_1_s.T@v2_rho - 0.5*q_1_s.T@v1_rho + 0.5*q_1_s@phi2_f - 0.5*q_1_s@phi1_f + nuc_cor = 0.5*q_1_s.T@pcmobj.v_grids_n + 0.5*q_2_f.T@pcmobj.v_grids_n + e_ss = tdneq.e[0] + mfneq.e_tot + delta_G + mfeq.with_solvent.e + nuc_cor + e_ss, delta_G + mfeq.with_solvent.e + nuc_cor + -""" \ No newline at end of file +def external_iteration(): + pass \ No newline at end of file diff --git a/gpu4pyscf/solvent/pcm.py b/gpu4pyscf/solvent/pcm.py index a2b97c5c0..08f3ffcb1 100644 --- a/gpu4pyscf/solvent/pcm.py +++ b/gpu4pyscf/solvent/pcm.py @@ -240,6 +240,21 @@ def get_D_S(surface, with_S=True, with_D=False, stream=None): raise RuntimeError('Failed in generating PCM D and S matrices.') return D, S + +# def get_phi(pcmobj, sigma): +# """ +# """ +# mol = pcmobj.mol +# int2c2e = mol._add_suffix('int2c2e') +# grid_coords = pcmobj.surface['grid_coords'] +# charge_exp = pcmobj.surface['charge_exp'] +# fakemol_charge = gto.fakemol_for_charges(grid_coords.get(), expnt=charge_exp.get()**2) +# v_ng = gto.mole.intor_cross(int2c2e, fakemol_charge, fakemol_charge) +# v_ng = cupy.asarray(v_ng) +# phi_s = sigma@v_ng +# return phi_s + + class PCM(lib.StreamObject): _keys = { 'method', 'vdw_scale', 'surface', 'r_probe', 'intopt', @@ -276,6 +291,8 @@ def __init__(self, mol): self.v = None self.v_grids_n = None self.dmprev = None + # self.dm0 = None + # self.eps_optical = 1.78 def dump_flags(self, verbose=None): logger.info(self, '******** %s ********', self.__class__) @@ -388,6 +405,12 @@ def _get_vind(self, dms): v_grids_e = self._get_v(dms) v_grids = self.v_grids_n - v_grids_e + # eps = self.eps + # eps_optical = self.eps_optical + # chi_e = (eps - 1)/(4 * numpy.pi) + # chi_fast = (eps_optical - 1)/(4 * numpy.pi) + # chi_slow = (eps - eps_optical)/(4 * numpy.pi) + if self.dmprev is not None: dmprev = self.dmprev dmprev = dmprev.reshape(-1,nao,nao) @@ -395,6 +418,12 @@ def _get_vind(self, dms): dmprev = (dmprev[0] + dmprev[1]).reshape(-1,nao,nao) if not isinstance(dmprev, cupy.ndarray): dmprev = cupy.asarray(dmprev) + # dm0 = self.dm0 + # dm0 = dm0.reshape(-1,nao,nao) + # if dm0.shape[0] == 2: + # dm0 = (dm0[0] + dm0[1]).reshape(-1,nao,nao) + # if not isinstance(dm0, cupy.ndarray): + # dm0 = cupy.asarray(dm0) v_grids_e_prev = self._get_v(dmprev) v_grids_prev = self.v_grids_n - v_grids_e_prev b = self.left_multiply_R(v_grids_prev.T) @@ -403,6 +432,30 @@ def _get_vind(self, dms): vK_1 = self.left_solve_K(v_grids_prev.T, K_transpose = True) qt = self.left_multiply_R(vK_1, R_transpose = True).T q_sym = (q + qt)/2.0 + + # v_grids_e_0 = self._get_v(dm0) + # v_grids_0 = self.v_grids_n - v_grids_e_0 + # b_0 = self.left_multiply_R(v_grids_0.T) + # q_0 = self.left_solve_K(b_0).T + + # vK_1_0 = self.left_solve_K(v_grids_0.T, K_transpose = True) + # qt_0 = self.left_multiply_R(vK_1_0, R_transpose = True).T + # q_sym_0 = (q_0 + qt_0)/2.0 + + # q_sym_0_s = chi_slow/chi_e * q_sym_0 + # phi_1_s = get_phi(self, q_sym_0_s) + # vgrids_2 = v_grids_prev + phi_1_s + # b_2 = self.left_multiply_R(vgrids_2.T) + # q_2 = self.left_solve_K(b_2).T + # vK_1_2 = self.left_solve_K(vgrids_2.T, K_transpose = True) + # qt_2 = self.left_multiply_R(vK_1_2, R_transpose = True).T + # q_sym_2_f = (q_2 + qt_2)/2.0 + # q_sym = q_sym_0_s + q_sym_2_f + # q = q_0 * chi_slow/chi_e + q_2 + # q_sym = q_sym_0_s + # q = q_0 * chi_slow/chi_e + + else: b = self.left_multiply_R(v_grids.T) q = self.left_solve_K(b).T From eab2941d0c22826dd66722728eb6d1cfbc5c0768 Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Wed, 23 Apr 2025 17:08:13 +0800 Subject: [PATCH 31/31] add an example of SSPCM-TDDFT --- examples/29-tddft_sspcm.py | 118 ++++++++++++++++++++++++++++++------- 1 file changed, 98 insertions(+), 20 deletions(-) diff --git a/examples/29-tddft_sspcm.py b/examples/29-tddft_sspcm.py index 0f5e1f253..600de362d 100644 --- a/examples/29-tddft_sspcm.py +++ b/examples/29-tddft_sspcm.py @@ -17,28 +17,32 @@ # Example of TDDFT with State-Specific Solvent ################################### -""" -[Experimental Feature Notice] ------------------------------------------------------------------------ -This implementation is EXPERIMENTAL and subject to change without warning. -Users must fully understand the theoretical assumptions and limitations -before application in production systems. Misuse may lead to incorrect results. -""" import numpy as np import pyscf import cupy as cp -import gpu4pyscf from pyscf import gto, scf, dft, tddft -from gpu4pyscf.solvent._attach_solvent import SCFWithSolvent -from gpu4pyscf.lib.cupy_helper import tag_array, pack_tril -from gpu4pyscf.scf.hf_lowmem import WaveFunction from gpu4pyscf.lib.cupy_helper import contract -from pyscf import lib import copy from functools import reduce from gpu4pyscf.scf import cphf +""" +[Experimental Feature Notice] +----------------------------------------------------------------------- +This implementation is EXPERIMENTAL and subject to change without warning. +Users must fully understand the theoretical assumptions and limitations +before application in production systems. Misuse may lead to incorrect results. + +Users may refer to the following papers for more details: +1. The State-specific approach and how to do a SS-PCM TDDFT calculation: + Exploring Chemistry with Electronic Structure Methods +2. 10.1063/1.2222364 +3. 10.1063/1.2757168 +""" + def get_total_density(td_grad, mf, x_y, singlet=True, relaxed=True): + """ + """ mol = mf.mol mo_coeff = cp.asarray(mf.mo_coeff) mo_energy = cp.asarray(mf.mo_energy) @@ -130,7 +134,7 @@ def get_phi(pcmobj, sigma): return phi_s -def get_deltaG(mfeq, mfneq, tdneq, tdgneq, eps_optical=1.78): +def get_deltaG(mfeq, mfneq, tdneq, dm2, eps_optical=1.78): pcmobj = mfeq.with_solvent eps = pcmobj.eps @@ -140,14 +144,10 @@ def get_deltaG(mfeq, mfneq, tdneq, tdgneq, eps_optical=1.78): pcmobj_optical.build() dm1_eq = mfeq.make_rdm1() - dm2 = get_total_density(tdgneq, mfneq, tdneq.xy[0]) chi_e = (eps - 1)/(4 * np.pi) - chi_fast = (eps_optical - 1)/(4 * np.pi) chi_slow = (eps - eps_optical)/(4 * np.pi) q_1 = pcmobj._get_qsym(dm1_eq, with_nuc=True)[0] q_1_s = chi_slow/chi_e * q_1 - q_2 = pcmobj._get_qsym(dm2, with_nuc=True)[0] - q_2_s = chi_slow/chi_e * q_2 v1 = pcmobj._get_vgrids(dm1_eq, with_nuc=True) v2 = pcmobj._get_vgrids(dm2, with_nuc=True) phi_1_s = get_phi(pcmobj, q_1_s) @@ -171,8 +171,86 @@ def get_deltaG(mfeq, mfneq, tdneq, tdgneq, eps_optical=1.78): delta_G = 0.5*q_2_f.T@v2_rho + q_1_s.T@v2_rho - 0.5*q_1_s.T@v1_rho + 0.5*q_1_s@phi2_f - 0.5*q_1_s@phi1_f nuc_cor = 0.5*q_1_s.T@pcmobj.v_grids_n + 0.5*q_2_f.T@pcmobj.v_grids_n e_ss = tdneq.e[0] + mfneq.e_tot + delta_G + mfeq.with_solvent.e + nuc_cor - e_ss, delta_G + mfeq.with_solvent.e + nuc_cor + return e_ss + + +def get_mf(mol, xc, solvent_model=None): + if xc.lower() == 'hf': + mf = scf.RHF(mol).PCM().to_gpu() + if solvent_model is not None: + mf.with_solvent = solvent_model + else: + mf.with_solvent.method = 'cpcm' + mf.with_solvent.equilibrium_solvation = False + mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids + mf.with_solvent.eps = 78 + radii_array = pyscf.data.radii.UFF*1.1 + mf.with_solvent.radii_table = radii_array + else: + raise NotImplementedError + + return mf + +def get_td(mf, xc, tda=False, equilibrium_solvation=True, linear_response=False): + if tda: + raise NotImplementedError + else: + if xc.lower() == 'hf': + td = mf.TDHF(equilibrium_solvation=equilibrium_solvation, + linear_response=linear_response).set(nstates=5) + else: + raise NotImplementedError + return td + + +def external_iteration(mol, xc='hf', tda=False, max_cycle=20, conv_tol=1e-6, + relaxed=True, equilibrium_solvation=True, linear_response=False): + mf1 = get_mf(mol, xc) + mf1.kernel() + td1 = get_td(mf1, xc, tda, equilibrium_solvation, linear_response) + td1.kernel() + g1 = td1.nuc_grad_method() + e0 = mf1.e_tot + td1.e[0] + + for icycle in range(max_cycle): + if icycle == 0: + dmprev = get_total_density(g1, mf1, td1.xy[0], relaxed=relaxed) + pcmobj = copy.copy(mf1.with_solvent) + else: + dmprev = get_total_density(gcycle, mfcycle, tdcycle.xy[0], relaxed=relaxed) + pcmobj = copy.copy(mfcycle.with_solvent) + + pcmobj.dmprev = dmprev + pcmobj.build() + pcmobj.kernel(dmprev) + + mfcycle = get_mf(mol, xc, pcmobj) + mfcycle.kernel(mf1.make_rdm1()) + tdcycle = get_td(mfcycle, xc, tda, equilibrium_solvation, linear_response) + tdcycle.kernel() + gcycle = tdcycle.nuc_grad_method() + e_ex = tdcycle.e[0]+mfcycle.e_tot + if abs(e_ex-e0)