Merge pull request #11 from sofroniewn/optimal_geometry

Add optimal_geometry method

Author: mfkasim1

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,

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

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:

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:

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)

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)