Skip to content

Commit

Permalink
Update class structures of localization modules
Browse files Browse the repository at this point in the history
  • Loading branch information
sunqm committed Sep 10, 2023
1 parent 66b246f commit fa74064
Show file tree
Hide file tree
Showing 5 changed files with 101 additions and 81 deletions.
6 changes: 6 additions & 0 deletions pyscf/lib/misc.py
Original file line number Diff line number Diff line change
Expand Up @@ -660,6 +660,12 @@ def add_keys(self, **kwargs):
self._keys = keys
return self

@classmethod
def dir_attributes(cls):
'''Returns all attributes defined in class'''
return {key for key, val in vars(cls).items()
if not (key.startswith('__') or callable(val))}

_warn_once_registry = {}
def check_sanity(obj, keysref, stdout=sys.stdout):
'''Check misinput of class attributes, check whether a class method is
Expand Down
157 changes: 86 additions & 71 deletions pyscf/lo/boys.py
Original file line number Diff line number Diff line change
Expand Up @@ -127,9 +127,93 @@ def atomic_init_guess(mol, mo_coeff):
u, w, vh = numpy.linalg.svd(mo[idx])
return lib.dot(u, vh).conj().T

class Boys(ciah.CIAHOptimizer):
class OrbitalLocalizer(lib.StreamObject, ciah.CIAHOptimizerMixin):

conv_tol = getattr(__config__, 'lo_boys_Boys_conv_tol', 1e-6)
conv_tol_grad = getattr(__config__, 'lo_boys_Boys_conv_tol_grad', None)
max_cycle = getattr(__config__, 'lo_boys_Boys_max_cycle', 100)
max_iters = getattr(__config__, 'lo_boys_Boys_max_iters', 20)
max_stepsize = getattr(__config__, 'lo_boys_Boys_max_stepsize', .05)
ah_trust_region = getattr(__config__, 'lo_boys_Boys_ah_trust_region', 3)
ah_start_tol = getattr(__config__, 'lo_boys_Boys_ah_start_tol', 1e9)
ah_max_cycle = getattr(__config__, 'lo_boys_Boys_ah_max_cycle', 40)
init_guess = getattr(__config__, 'lo_boys_Boys_init_guess', 'atomic')

_keys = {
'conv_tol', 'conv_tol_grad', 'max_cycle', 'max_iters',
'max_stepsize', 'ah_trust_region', 'ah_start_tol',
'ah_max_cycle', 'init_guess', 'mol', 'mo_coeff',
}

def __init__(self, mol, mo_coeff=None):
self.mol = mol
self.stdout = mol.stdout
self.verbose = mol.verbose
self.mo_coeff = mo_coeff

self._keys = set.union(OrbitalLocalizer.dir_attributes(),
ciah.CIAHOptimizerMixin.dir_attributes(),
self.__dict__)

def dump_flags(self, verbose=None):
log = logger.new_logger(self, verbose)
log.info('\n')
log.info('******** %s ********', self.__class__)
log.info('conv_tol = %s' , self.conv_tol )
log.info('conv_tol_grad = %s' , self.conv_tol_grad )
log.info('max_cycle = %s' , self.max_cycle )
log.info('max_stepsize = %s' , self.max_stepsize )
log.info('max_iters = %s' , self.max_iters )
log.info('kf_interval = %s' , self.kf_interval )
log.info('kf_trust_region = %s', self.kf_trust_region)
log.info('ah_start_tol = %s' , self.ah_start_tol )
log.info('ah_start_cycle = %s' , self.ah_start_cycle )
log.info('ah_level_shift = %s' , self.ah_level_shift )
log.info('ah_conv_tol = %s' , self.ah_conv_tol )
log.info('ah_lindep = %s' , self.ah_lindep )
log.info('ah_max_cycle = %s' , self.ah_max_cycle )
log.info('ah_trust_region = %s', self.ah_trust_region)
log.info('init_guess = %s' , self.init_guess )

def get_init_guess(self, key='atomic'):
'''Generate initial guess for localization.
Kwargs:
key : str or bool
If key is 'atomic', initial guess is based on the projected
atomic orbitals. False
'''
nmo = self.mo_coeff.shape[1]
if isinstance(key, str) and key.lower() == 'atomic':
u0 = atomic_init_guess(self.mol, self.mo_coeff)
elif isinstance(key, str) and key.lower().startswith('cho'):
mo_init = cholesky_mos(self.mo_coeff)
S = self.mol.intor_symmetric('int1e_ovlp')
u0 = numpy.linalg.multi_dot([self.mo_coeff.T, S, mo_init])
else:
u0 = numpy.eye(nmo)
if (isinstance(key, str) and key.lower().startswith('rand')
or numpy.linalg.norm(self.get_grad(u0)) < 1e-5):
# Add noise to kick initial guess out of saddle point
dr = numpy.cos(numpy.arange((nmo-1)*nmo//2)) * 1e-3
u0 = self.extract_rotation(dr)
return u0

def gen_g_hop(self, u):
raise NotImplementedError

def get_grad(self, u=None):
raise NotImplementedError

def cost_function(self, u=None):
raise NotImplementedError

kernel = kernel


class Boys(OrbitalLocalizer):
r'''
The Foster-Boys localization optimizer that maximizes the orbital dipole
Base class oflocalization optimizer that maximizes the orbital dipole
\sum_i | <i| r |i> |^2
Expand Down Expand Up @@ -182,49 +266,6 @@ class Boys(ciah.CIAHOptimizer):
'''

conv_tol = getattr(__config__, 'lo_boys_Boys_conv_tol', 1e-6)
conv_tol_grad = getattr(__config__, 'lo_boys_Boys_conv_tol_grad', None)
max_cycle = getattr(__config__, 'lo_boys_Boys_max_cycle', 100)
max_iters = getattr(__config__, 'lo_boys_Boys_max_iters', 20)
max_stepsize = getattr(__config__, 'lo_boys_Boys_max_stepsize', .05)
ah_trust_region = getattr(__config__, 'lo_boys_Boys_ah_trust_region', 3)
ah_start_tol = getattr(__config__, 'lo_boys_Boys_ah_start_tol', 1e9)
ah_max_cycle = getattr(__config__, 'lo_boys_Boys_ah_max_cycle', 40)
init_guess = getattr(__config__, 'lo_boys_Boys_init_guess', 'atomic')

_keys = set((
'conv_tol', 'conv_tol_grad', 'max_cycle', 'max_iters',
'max_stepsize', 'ah_trust_region', 'ah_start_tol',
'ah_max_cycle', 'init_guess', 'mol', 'mo_coeff',
))

def __init__(self, mol, mo_coeff=None):
ciah.CIAHOptimizer.__init__(self)
self.mol = mol
self.stdout = mol.stdout
self.verbose = mol.verbose
self.mo_coeff = mo_coeff

def dump_flags(self, verbose=None):
log = logger.new_logger(self, verbose)
log.info('\n')
log.info('******** %s ********', self.__class__)
log.info('conv_tol = %s' , self.conv_tol )
log.info('conv_tol_grad = %s' , self.conv_tol_grad )
log.info('max_cycle = %s' , self.max_cycle )
log.info('max_stepsize = %s' , self.max_stepsize )
log.info('max_iters = %s' , self.max_iters )
log.info('kf_interval = %s' , self.kf_interval )
log.info('kf_trust_region = %s', self.kf_trust_region)
log.info('ah_start_tol = %s' , self.ah_start_tol )
log.info('ah_start_cycle = %s' , self.ah_start_cycle )
log.info('ah_level_shift = %s' , self.ah_level_shift )
log.info('ah_conv_tol = %s' , self.ah_conv_tol )
log.info('ah_lindep = %s' , self.ah_lindep )
log.info('ah_max_cycle = %s' , self.ah_max_cycle )
log.info('ah_trust_region = %s', self.ah_trust_region)
log.info('init_guess = %s' , self.init_guess )

def gen_g_hop(self, u):
mo_coeff = lib.dot(self.mo_coeff, u)
dip = dipole_integral(self.mol, mo_coeff)
Expand Down Expand Up @@ -303,32 +344,6 @@ def cost_function(self, u=None):
val = r2 - numpy.einsum('xii,xii->', dip, dip)
return val

def get_init_guess(self, key='atomic'):
'''Generate initial guess for localization.
Kwargs:
key : str or bool
If key is 'atomic', initial guess is based on the projected
atomic orbitals. False
'''
nmo = self.mo_coeff.shape[1]
if isinstance(key, str) and key.lower() == 'atomic':
u0 = atomic_init_guess(self.mol, self.mo_coeff)
elif isinstance(key, str) and key.lower().startswith('cho'):
mo_init = cholesky_mos(self.mo_coeff)
S = self.mol.intor_symmetric('int1e_ovlp')
u0 = numpy.linalg.multi_dot([self.mo_coeff.T, S, mo_init])
else:
u0 = numpy.eye(nmo)
if (isinstance(key, str) and key.lower().startswith('rand')
or numpy.linalg.norm(self.get_grad(u0)) < 1e-5):
# Add noise to kick initial guess out of saddle point
dr = numpy.cos(numpy.arange((nmo-1)*nmo//2)) * 1e-3
u0 = self.extract_rotation(dr)
return u0

kernel = kernel

FB = BF = Boys


Expand Down
3 changes: 1 addition & 2 deletions pyscf/lo/edmiston.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,12 +22,11 @@

import numpy
from functools import reduce

from pyscf.scf import hf
from pyscf.lo import boys


class EdmistonRuedenberg(boys.Boys):
class EdmistonRuedenberg(boys.OrbitalLocalizer):

def get_jk(self, u):
mo_coeff = numpy.dot(self.mo_coeff, u)
Expand Down
6 changes: 3 additions & 3 deletions pyscf/lo/pipek.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,7 @@ def atomic_pops(mol, mo_coeff, method='meta_lowdin', mf=None):
return proj


class PipekMezey(boys.Boys):
class PipekMezey(boys.OrbitalLocalizer):
'''The Pipek-Mezey localization optimizer that maximizes the orbital
population
Expand Down Expand Up @@ -189,13 +189,13 @@ class PipekMezey(boys.Boys):
_keys = set(['pop_method', 'conv_tol', 'exponent'])

def __init__(self, mol, mo_coeff=None, mf=None, pop_method=None):
boys.Boys.__init__(self, mol, mo_coeff)
boys.OrbitalLocalizer.__init__(self, mol, mo_coeff)
self._scf = mf
if pop_method is not None:
self.pop_method = pop_method

def dump_flags(self, verbose=None):
boys.Boys.dump_flags(self, verbose)
boys.OrbitalLocalizer.dump_flags(self, verbose)
logger.info(self, 'pop_method = %s',self.pop_method)

def gen_g_hop(self, u):
Expand Down
10 changes: 5 additions & 5 deletions pyscf/soscf/ciah.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@
def expmat(a):
return scipy.linalg.expm(a)

class CIAHOptimizer(lib.StreamObject):
class CIAHOptimizerMixin:

conv_tol_grad = getattr(__config__, 'soscf_ciah_CIAHOptimizer_conv_tol_grad', 1e-4)
max_stepsize = getattr(__config__, 'soscf_ciah_CIAHOptimizer_max_stepsize', .05)
Expand All @@ -43,11 +43,11 @@ class CIAHOptimizer(lib.StreamObject):
ah_max_cycle = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_max_cycle', 30)
ah_trust_region = getattr(__config__, 'soscf_ciah_CIAHOptimizer_ah_trust_region', 3.)

_keys = set((
_keys = {
'conv_tol_grad', 'max_stepsize', 'max_iters', 'kf_interval',
'kf_trust_region', 'ah_start_tol', 'ah_start_cycle', 'ah_level_shift',
'ah_conv_tol', 'ah_lindep', 'ah_max_cycle', 'ah_trust_region',
))
}

def gen_g_hop(self, u):
raise NotImplementedError
Expand All @@ -69,10 +69,10 @@ def extract_rotation(self, dr, u0=1):
return numpy.dot(u0, expmat(dr))

def get_grad(self, u):
pass
raise NotImplementedError

def cost_function(self, u):
pass
raise NotImplementedError


def rotate_orb_cc(iah, u0, conv_tol_grad=None, verbose=logger.NOTE):
Expand Down

0 comments on commit fa74064

Please sign in to comment.