diff --git a/dqc/api/properties.py b/dqc/api/properties.py index 7583215..b9a4831 100644 --- a/dqc/api/properties.py +++ b/dqc/api/properties.py @@ -13,7 +13,8 @@ convert_raman_ints __all__ = ["hessian_pos", "vibration", "edipole", "equadrupole", "is_orb_min", - "lowest_eival_orb_hessian", "ir_spectrum", "raman_spectrum"] + "lowest_eival_orb_hessian", "ir_spectrum", "raman_spectrum", + "optimal_geometry"] # This file contains functions to calculate the perturbation properties of systems. @@ -317,6 +318,28 @@ def is_orb_min(qc: BaseQCCalc, threshold: float = -1e-3) -> bool: eival = lowest_eival_orb_hessian(qc) return bool(torch.all(eival > threshold)) +def optimal_geometry(qc: BaseQCCalc, length_unit: Optional[str] = None) -> torch.Tensor: + """ + Compute the optimal atomic positions of the system. + + Arguments + --------- + qc: BaseQCCalc + Quantum Chemistry calculation that has run. + + length_unit: str or None + The returned unit. If ``None``, returns in atomic unit. + + Returns + ------- + torch.Tensor + Tensor with shape ``(natoms, ndim)`` represents the position + of atoms at the optimal geometry. + """ + atompos = _optimal_geometry(qc) + atompos = convert_length(atompos, to_unit=length_unit) + return atompos + @memoize_method def _hessian_pos(qc: BaseQCCalc) -> torch.Tensor: # calculate the hessian in atomic unit @@ -460,6 +483,28 @@ def _equadrupole(qc: BaseQCCalc) -> torch.Tensor: return quadrupole + ion_quadrupole +@memoize_method +def _optimal_geometry(qc: BaseQCCalc) -> torch.Tensor: + # calculate the optimal geometry + system = qc.get_system() + atompos = system.atompos + + # check if the atompos requires grad + _check_differentiability(atompos, "atom positions", "optimal geometry") + + # get the energy for a given geometry + def _get_energy(atompos: torch.Tensor) -> torch.Tensor: + new_system = system.make_copy(moldesc=(system.atomzs, atompos)) + new_qc = qc.__class__(new_system).run() + ene = new_qc.energy() # calculate the energy + return ene + + # get the minimal enery position + minpos = xitorch.optimize.minimize(_get_energy, atompos, method="gd", maxiter=200, + step=1e-2) + + return minpos + ########### helper functions ########### def _jac(a: torch.Tensor, b: torch.Tensor, create_graph: Optional[bool] = None, diff --git a/dqc/system/base_system.py b/dqc/system/base_system.py index 19b1a36..27bed25 100644 --- a/dqc/system/base_system.py +++ b/dqc/system/base_system.py @@ -80,6 +80,14 @@ def requires_grid(self) -> bool: def getparamnames(self, methodname: str, prefix: str = "") -> List[str]: pass + @abstractmethod + def make_copy(self, **kwargs) -> BaseSystem: + """ + Returns a copy of the system identical to the orginal except for new + parameters set in the kwargs. + """ + pass + ####################### system properties ####################### @abstractproperty def atompos(self) -> torch.Tensor: @@ -129,4 +137,4 @@ def efield(self) -> Optional[Tuple[torch.Tensor, ...]]: Returns the external electric field of the system, or None if there is no electric field. """ - pass + pass \ No newline at end of file diff --git a/dqc/system/mol.py b/dqc/system/mol.py index 439fd51..dd107e2 100644 --- a/dqc/system/mol.py +++ b/dqc/system/mol.py @@ -1,3 +1,4 @@ +from __future__ import annotations from typing import List, Union, Optional, Tuple, Dict import warnings import torch @@ -294,6 +295,36 @@ def getparamnames(self, methodname: str, prefix: str = "") -> List[str]: else: raise KeyError("Unknown methodname: %s" % methodname) + def make_copy(self, **kwargs) -> Mol: + """ + Returns a copy of the system identical to the orginal except for new + parameters set in the kwargs. + + Arguments + --------- + **kwargs + Must be the same kwargs as Mol. + """ + # create dictionary of all parameters + parameters = { + 'moldesc': (self.atomzs, self.atompos), + 'basis': self._basis_inp, + 'orthogonalize_basis': self._orthogonalize_basis, + 'ao_parameterizer': self._aoparamzer, + 'grid': self._grid_inp, + 'spin': self._spin, + 'charge': self._charge, + 'orb_weights': None, + 'efield': self._efield, + 'vext': self._vext, + 'dtype': self._dtype, + 'device': self._device + } + # update dictionary with provided kwargs + parameters.update(kwargs) + # create new system + return Mol(**parameters) + ################### properties ################### @property def atompos(self) -> torch.Tensor: diff --git a/dqc/system/sol.py b/dqc/system/sol.py index 722cd8b..a12bc05 100644 --- a/dqc/system/sol.py +++ b/dqc/system/sol.py @@ -1,3 +1,4 @@ +from __future__ import annotations from typing import Optional, Tuple, Union, List, Dict import torch import numpy as np @@ -68,6 +69,7 @@ def __init__(self, self._dtype = dtype self._device = device self._grid_inp = grid + self._basis_inp = basis self._grid: Optional[BaseGrid] = None charge = 0 # we can't have charged solids for now @@ -99,7 +101,8 @@ def __init__(self, self._orb_weights = _orb_weights self._orb_weights_u = _orb_weights_u self._orb_weights_d = _orb_weights_d - self._lattice = Lattice(alattice) + self._alattice_inp = alattice + self._lattice = Lattice(self._alattice_inp) self._lattsum_opt = PBCIntOption.get_default(lattsum_opt) def densityfit(self, method: Optional[str] = None, @@ -240,6 +243,32 @@ def requires_grid(self) -> bool: def getparamnames(self, methodname: str, prefix: str = "") -> List[str]: pass + def make_copy(self, **kwargs) -> Sol: + """ + Returns a copy of the system identical to the orginal except for new + parameters set in the kwargs. + + Arguments + --------- + **kwargs + Must be the same kwargs as Sol. + """ + # create dictionary of all parameters + parameters = { + 'soldesc': (self.atomzs, self.atompos), + 'alattice': self._alattice_inp, + 'basis': self._basis_inp, + 'grid': self._grid_inp, + 'spin': self._spin, + 'lattsum_opt': self._lattsum_opt, + 'dtype': self._dtype, + 'device': self._device + } + # update dictionary with provided kwargs + parameters.update(kwargs) + # create new system + return Sol(**parameters) + ################### properties ################### @property def atompos(self) -> torch.Tensor: diff --git a/dqc/test/test_properties.py b/dqc/test/test_properties.py index 18abf23..1694e5c 100644 --- a/dqc/test/test_properties.py +++ b/dqc/test/test_properties.py @@ -5,7 +5,7 @@ import psutil from dqc.api.properties import hessian_pos, vibration, edipole, equadrupole, \ ir_spectrum, raman_spectrum, is_orb_min, \ - lowest_eival_orb_hessian + lowest_eival_orb_hessian, optimal_geometry from dqc.system.mol import Mol from dqc.qccalc.hf import HF from dqc.xc.base_xc import BaseXC @@ -441,3 +441,31 @@ def get_jac_ene(atomposs, efield, grad_efield): # raman spectra intensities torch.autograd.gradgradcheck(get_jac_ene, (atomposs.detach(), efield, grad_efield.detach()), atol=3e-4) + +def test_optimal_geometry(h2o_qc): + # test if the optimal geometry of h2o similar to pyscf + # from CCCBDB (calculated geometry for H2O) + pyscf_h2o_opt = h2o_qc.get_system().atompos + + # create a new h2o with poor initial geometry + h2o_init = torch.tensor([ + [0.0, 0.0, 0.214], + [0.0, 1.475, -0.863], + [0.0, -1.475, -0.863], + ], dtype=dtype).requires_grad_() + + # use bond length to assess optimal geometry as they are rotation invariant + def bond_length(h2o): + # Get the bond lengths of an h20 molecule + return torch.stack([(h2o[0] - h2o[1]).norm(), (h2o[0] - h2o[2]).norm()]) + + # check starting geometry is not optimal + assert not torch.allclose(bond_length(h2o_init), bond_length(pyscf_h2o_opt), rtol=2e-4) + + # optimize geometry + system = h2o_qc.get_system() + new_system = system.make_copy(moldesc=(system.atomzs, system.atompos)) + new_qc = h2o_qc.__class__(new_system).run() + h2o_opt = optimal_geometry(new_qc) + + assert torch.allclose(bond_length(h2o_opt), bond_length(pyscf_h2o_opt), rtol=2e-4) \ No newline at end of file diff --git a/dqc/test/test_system.py b/dqc/test/test_system.py index aeec15e..59ecd06 100644 --- a/dqc/test/test_system.py +++ b/dqc/test/test_system.py @@ -127,6 +127,18 @@ def test_mol_cache(): if os.path.exists(cache_fname): os.remove(cache_fname) +def test_mol_copy(): + # test if copy is computed correctly + moldesc = "H 0 0 0; H 1 0 0" + mol = Mol(moldesc, basis="3-21G") + + mol_copy = mol.make_copy() + assert torch.allclose(mol_copy.atompos, mol.atompos) + + new_pos = torch.tensor([[0.0, 0.0, 0.01], [1.01, 0.0, 0.0]], dtype=dtype) + mol_copy_2 = mol.make_copy(moldesc=(mol.atomzs, new_pos)) + assert torch.allclose(mol_copy_2.atompos, new_pos) + def test_sol_cache(): # test if cache is stored correctly @@ -168,6 +180,19 @@ def test_sol_cache(): j2c1 = h1.df.j2c assert torch.allclose(j2c, j2c1) +def test_sol_copy(): + # test if copy is computed correctly + soldesc = "H 0 0 0" + a = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=dtype) * 3 + sol = Sol(soldesc, alattice=a, basis="3-21G") + + sol_copy = sol.make_copy() + assert torch.allclose(sol_copy.atompos, sol.atompos) + + new_pos = torch.tensor([[0.0, 0.0, 0.01]], dtype=dtype) + sol_copy_2 = sol.make_copy(soldesc=(sol.atomzs, new_pos)) + assert torch.allclose(sol_copy_2.atompos, new_pos) + ##################### pbc ##################### def test_mol_pbc_nuclei_energy(): # test the calculation of ion-ion interaction energy (+ gradients w.r.t. pos)