diff --git a/basisopt/basis/atomic.py b/basisopt/basis/atomic.py index 8559b75..7bff3f2 100644 --- a/basisopt/basis/atomic.py +++ b/basisopt/basis/atomic.py @@ -13,10 +13,11 @@ from basisopt.opt.eventemper import EvenTemperedStrategy from basisopt.opt.optimizers import optimize from basisopt.opt.strategies import Strategy +from basisopt.opt.welltemper import WellTemperedStrategy from basisopt.util import bo_logger from . import zetatools as zt -from .basis import Basis, even_temper_expansion +from .basis import Basis, even_temper_expansion, well_temper_expansion def needs_element(func: Callable) -> Callable: @@ -40,6 +41,7 @@ class AtomicBasis(Basis): Attributes: et_params (data.ETParams): even tempered expansion parameters + wt_params (data.WTParams): well tempered expansion parameters charge (int): net charge on atom multiplicity (int): spin multiplicity of atom config (dict): configuration of basis, (k, v) pairs @@ -65,6 +67,7 @@ def __init__(self, name: str = 'H', charge: int = 0, mult: int = 1): self.element = name self._done_setup = False self.et_params = None + self.wt_params = None if self._element is not None: self.charge = charge @@ -83,6 +86,7 @@ def as_dict(self) -> dict[str, Any]: d["@module"] = type(self).__module__ d["@class"] = type(self).__name__ d["et_params"] = self.et_params + d["wt_params"] = self.wt_params if hasattr(self, 'strategy'): if isinstance(self.strategy, Strategy): @@ -105,6 +109,7 @@ def from_dict(cls, d: dict[str, Any]) -> object: instance._tests = basis._tests instance._molecule = basis._molecule instance.et_params = d.get("et_params", None) + instance.wt_params = d.get("wt_params", None) instance.strategy = d.get("strategy", None) if instance.strategy: instance._done_setup = d.get("done_setup", False) @@ -301,6 +306,45 @@ def set_even_tempered( else: self._molecule.basis[self._symbol] = even_temper_expansion(self.et_params) + @needs_element + def set_well_tempered( + self, + method: str = 'hf', + accuracy: float = 1e-5, + max_n: int = 18, + max_l: int = -1, + exact_ref: bool = True, + params: dict[str, Any] = {}, + ): + """Looks up or computes a well tempered basis expansion for the atom + + Arguments: + method (str): method to use; possibilities can be found through Wrapper object + accuracy (float): the tolerance to optimize to, compared to reference value + max_n (int): max number of primitives per shell + max_l (int): angular momentum to go up to; if -1, will use max l in minimal config + exact_ref (bool): uses exact numerical HF energy if True, + calculates cc-pV5Z reference value if False + params (dict): dictionary of parameters to pass to the backend - + see the relevant Wrapper object for options + + Sets: + self.wt_params + """ + self.wt_params = data.get_well_temper_params(atom=self._symbol.title(), accuracy=accuracy) + if len(self.wt_params) == 0: + # optimize new params + if exact_ref: + reference = ('exact', data._ATOMIC_HF_ENERGIES[self._element.atomic_number]) + else: + reference = ('cc-pV5Z', None) + strategy = WellTemperedStrategy(max_n=max_n, max_l=max_l) + self.setup(method=method, strategy=strategy, reference=reference, params=params) + self.optimize(algorithm='Nelder-Mead', params=params) + self.wt_params = strategy.shells + else: + self._molecule.basis[self._symbol] = well_temper_expansion(self.wt_params) + @needs_element def optimize(self, algorithm: str = 'Nelder-Mead', params: dict[str, Any] = {}) -> OptResult: """Runs the basis optimization diff --git a/basisopt/basis/basis.py b/basisopt/basis/basis.py index 7a09add..af1227d 100644 --- a/basisopt/basis/basis.py +++ b/basisopt/basis/basis.py @@ -7,7 +7,7 @@ from basisopt import data from basisopt.containers import InternalBasis, Result, Shell -from basisopt.data import ETParams +from basisopt.data import ETParams, WTParams from basisopt.testing import Test from basisopt.util import bo_logger, dict_decode @@ -67,6 +67,27 @@ def even_temper_expansion(params: ETParams) -> list[Shell]: return el_basis +def well_temper_expansion(params: WTParams) -> list[Shell]: + """Forms a basis for an element from well tempered expansion parameters + + Arguments: + params (list): list of tuples corresponding to shells + e.g. [(c_s, x_s, g_s, d_s, n_s), (c_p, x_p, g_p, d_p, n_p), ...] where each shell + is expanded as c_l * (x_l**k)*(1 + g_l*((k+1)/n_l)**d_l) for k=0,...,n_l + + Returns: + list of Shell objects for the well tempered expansion + """ + el_basis = [] + for ix, (c, x, g, d, n) in enumerate(params): + new_shell = Shell() + new_shell.l = data.INV_AM_DICT[ix] + new_shell.exps = np.array([c * (x**p) * (1.0 + g * ((p + 1) / n) ** d) for p in range(n)]) + uncontract_shell(new_shell) + el_basis.append(new_shell) + return el_basis + + def fix_ratio(exps: np.ndarray, ratio: float = 1.4) -> np.ndarray: """Returns a sorted numpy array of exponents where x_{i+1}/x_i >= ratio diff --git a/basisopt/basis/guesses.py b/basisopt/basis/guesses.py index b66e95c..9a2fb3a 100644 --- a/basisopt/basis/guesses.py +++ b/basisopt/basis/guesses.py @@ -5,7 +5,12 @@ from basisopt.bse_wrapper import fetch_basis from basisopt.containers import Shell -from .basis import even_temper_expansion, fix_ratio, uncontract_shell +from .basis import ( + even_temper_expansion, + fix_ratio, + uncontract_shell, + well_temper_expansion, +) # All guess functions need this signature # func(atomic, params={}), where atomic is an AtomicBasis object @@ -57,3 +62,14 @@ def even_tempered_guess(atomic, params={}): if atomic.et_params is None: atomic.set_even_tempered(**params) return even_temper_expansion(atomic.et_params) + + +def well_tempered_guess(atomic, params={}): + """Takes guess from a well-tempered expansion + + Params: + see signature for AtomicBasis.set_well_tempered + """ + if atomic.wt_params is None: + atomic.set_well_tempered(**params) + return well_temper_expansion(atomic.wt_params) diff --git a/basisopt/data.py b/basisopt/data.py index 25d7b17..cc36894 100644 --- a/basisopt/data.py +++ b/basisopt/data.py @@ -29,6 +29,11 @@ def atomic_number(element: str) -> int: ETParams = list[tuple[float, float, int]] +"""Dictionary with pre-optimised well-tempered expansions for atoms""" +_WELL_TEMPERED_DATA = {} + +WTParams = list[tuple[float, float, float, float, int]] + def get_even_temper_params(atom: str = 'H', accuracy: float = 1e-5) -> ETParams: """Searches for the relevant even tempered expansion @@ -43,6 +48,19 @@ def get_even_temper_params(atom: str = 'H', accuracy: float = 1e-5) -> ETParams: return [] +def get_well_temper_params(atom: str = 'H', accuracy: float = 1e-5) -> WTParams: + """Searches for the relevant well tempered expansion + from _WELL_TEMPERED_DATA + """ + if atom in _WELL_TEMPERED_DATA: + log_acc = -np.log10(accuracy) + index = max(4, log_acc) - 4 + index = int(min(index, 3)) + return _WELL_TEMPERED_DATA[atom][index] + else: + return [] + + """Essentially exact numerical Hartree-Fock energies for all atoms in Hartree. Ref: Saito 2009, doi.org/10.1016/j.adt.2009.06.001""" _ATOMIC_HF_ENERGIES = { diff --git a/basisopt/opt/welltemper.py b/basisopt/opt/welltemper.py new file mode 100644 index 0000000..61a8808 --- /dev/null +++ b/basisopt/opt/welltemper.py @@ -0,0 +1,178 @@ +from typing import Any, Final + +import numpy as np +from mendeleev import element as md_element + +from basisopt.basis.basis import well_temper_expansion +from basisopt.basis.guesses import null_guess +from basisopt.containers import InternalBasis + +from .preconditioners import unit +from .strategies import Strategy + +_INITIAL_GUESS: Final = (0.1, 2.2, 20.0, 10.0, 8) +_MIN_RATIO: Final = 1.01 + + +class WellTemperedStrategy(Strategy): + """Implements a strategy for a well tempered basis set, where each angular + momentum shell is described by five parameters: (c, x, g, d, n) + Each exponent in that shell is then given by + y_k = c*(x**k)*(1 + g_l*((k+1)/n_l)**d_l) for k=0,...,n + + Algorithm: + Evaluate: energy (can change to any RMSE-compatible property) + Loss: root-mean-square error + Guess: null, uses _INITIAL_GUESS above + Pre-conditioner: None + + Initialisation: + - Find minimum no. of shells needed + - max_l >= min_l + - generate initial parameters for each shell + + First run: + - optimize parameters for each shell once, sequentially + + Next shell in list not marked finished: + - re-optimise + - below threshold or n=max_n: mark finished + - above threshold: increment n + Repeat until all shells are marked finished. + + Uses iteration, limited by two parameters: + max_n: max number of exponents in shell + target: threshold for objective function + + Additional attributes: + shells (list): list of (c, x, g, d, n) parameter tuples + shell_done (list): list of flags for whether shell is finished (0) or not (1) + target (float): threshold for optimization delta + max_n (int): maximum number of primitives in shell expansion + max_l (int): maximum angular momentum shell to do; + if -1, does minimal configuration + """ + + def __init__( + self, eval_type: str = 'energy', target: float = 1e-5, max_n: int = 18, max_l: int = -1 + ): + super().__init__(eval_type=eval_type, pre=unit) + self.name = 'WellTemper' + self.shells = [] + self.shell_done = [] + self.target = target + self.guess = null_guess + self.guess_params = {} + self.max_n = max_n + self.max_l = max_l + + def as_dict(self) -> dict[str, Any]: + """Returns MSONable dictionary of object""" + d = super().as_dict() + d["@module"] = type(self).__module__ + d["@class"] = type(self).__name__ + d["shells"] = self.shells + d["shell_done"] = self.shell_done + d["target"] = self.target + d["max_n"] = self.max_n + d["max_l"] = self.max_l + return d + + @classmethod + def from_dict(cls, d: dict[str, Any]) -> object: + """Creates WellTemperedStrategy from MSONable dictionary""" + strategy = Strategy.from_dict(d) + instance = cls( + eval_type=d.get("eval_type", 'energy'), + target=d.get("target", 1e-5), + max_n=d.get("max_n", 18), + max_l=d.get("max_l", -1), + ) + instance.name = strategy.name + instance.params = strategy.params + instance.first_run = strategy.first_run + instance._step = strategy._step + instance.last_objective = strategy.last_objective + instance.delta_objective = strategy.delta_objective + instance.shells = d.get("shells", []) + instance.shell_done = d.get("shell_done", []) + return instance + + def set_basis_shells(self, basis: InternalBasis, element: str): + """Expands well tempered parameters into a basis set + + Arguments: + basis (InternalBasis): the basis set to expand + element (str): the atom type + """ + basis[element] = well_temper_expansion(self.shells) + + def initialise(self, basis: InternalBasis, element: str): + """Initialises the strategy by determing the initial + parameters for each angular momentum shell for the + given element. + + Arguments: + basis (InternalBasis): the basis set being optimized + element (str): the atom type of interest + """ + if self.max_l < 0: + el = md_element(element.title()) + l_list = [l for (n, l) in el.ec.conf.keys()] + min_l = len(set(l_list)) + + self.max_l = max(min_l, self.max_l) + self.shells = [_INITIAL_GUESS] * self.max_l + self.shell_done = [1] * self.max_l + self.set_basis_shells(basis, element) + self.last_objective = 0.0 + self.delta_objective = 0.0 + self.first_run = True + + def get_active(self, basis: InternalBasis, element: str) -> np.ndarray: + """Returns the well temper params for the current shell""" + (c, x, g, d, _) = self.shells[self._step] + return np.array([c, x, g, d]) + + def set_active(self, values: np.ndarray, basis: InternalBasis, element: str): + """Given the well temper params for a shell, expands the basis + Checks that the smallest exponent is >= 1e-5 + that the ratio is >= 1.01, to prevent impossible exponents + that the gamma spacing parameter is >= 1.01 to add well tempering + and that the delta parameter is also >= 1.01 to aid in deviation + from a rigid geometric series. + """ + (c, x, g, d, n) = self.shells[self._step] + c = max(values[0], 1e-5) + x = max(values[1], _MIN_RATIO) + g = max(values[2], _MIN_RATIO) + d = max(values[3], _MIN_RATIO) + self.shells[self._step] = (c, x, g, d, n) + self.set_basis_shells(basis, element) + + def next(self, basis: InternalBasis, element: str, objective: float) -> bool: + self.delta_objective = np.abs(self.last_objective - objective) + self.last_objective = objective + + carry_on = True + if self.first_run: + self._step = self._step + 1 + if self._step == self.max_l: + self.first_run = False + self._step = 0 + (c, x, g, d, n) = self.shells[self._step] + self.shells[self._step] = (c, x, g, d, min(n + 1, self.max_n)) + else: + if self.delta_objective < self.target: + self.shell_done[self._step] = 0 + + self._step = (self._step + 1) % self.max_l + (c, x, g, d, n) = self.shells[self._step] + if n == self.max_n: + self.shell_done[self._step] = 0 + elif self.shell_done[self._step] != 0: + self.shells[self._step] = (c, x, g, d, n + 1) + + carry_on = np.sum(self.shell_done) != 0 + + return carry_on diff --git a/doc/src/_tutorials/compare.rst b/doc/src/_tutorials/compare.rst new file mode 100644 index 0000000..2e36fe4 --- /dev/null +++ b/doc/src/_tutorials/compare.rst @@ -0,0 +1,114 @@ +:orphan: +.. _`sec:cpmareetwt`: + +==================================================== +Comparing even-tempered and well-tempered basis sets +==================================================== + +This tutorial compares an even-tempered basis set for the neon atom with an equivalent well-tempered basis set, both in terms of +the SCF energy and the exponents that arrise from the expansions. An example script can also be found in +``examples/et_wt_compare.py``. + +An even-tempered basis set +-------------------------- + +We use the same procedure for optimizing an even-tempered basis set as in the :ref:`sec:eventemper` example, +please see that tutorial for further details of how to perform this in BasisOpt. An even-tempered expansion of :math:`n` exponents +is defined as: + +.. math:: + + \alpha_n = \alpha_0 k^{n-1} + +where :math:`\alpha_0` and :math:`k` are the two variables to be optimized. These can be thought of as the initial (most diffuse) +exponent and the spacing between all subsequent exponents, respectively. By following the even-tempered example +and running the following optimization with the Psi4 backend + +.. code-block:: python + + ne_et.set_even_tempered(method='scf', accuracy=1e-4) + +We should get results that look very similar to:: + + l = 0: (0.24303219077706045, 2.295167745669027, 18) + l = 1: (0.12199362316666561, 2.180056874861101, 13) + +where the first value in parentheses corresponds to the optimized :math:`\alpha_0`, the second value is :math:`k`, +and the final integer is the number of functions that were required to reach the desired accuracy. + +A well-tempered basis set +------------------------- + +A well-tempered basis set builds upon the even-tempered basis by introducing two additional parameters, such that +the expansion is defined as: + +.. math:: + + \alpha_n = \alpha_0 k^{n-1} [1 + \gamma(n / N)^{\delta}] + +where :math:`N` is the total number of primitive exponents in the expansion, and :math:`\gamma` and :math:`\delta` are the +additional variables to be optimized. + +A well-tempered Atomic Basis can be optimized in a similar fashion as the even-tempered example. A full example input is +provided in ``examples/wt_strategy.py``, but perhaps the most important part is: + +.. code-block:: python + + ne_wt.set_well_tempered(method='scf', accuracy=1e-4) + + +The results should look similar to:: + + l = 0: (0.21397184263141658, 2.129617983476232, 16.3786689342292, 11.79947009544172, 18) + l = 1: (0.11831666016132977, 2.130099585140868, 2.7477960324041435, 10.480981519305843, 12) + +A comparison of the above with the optimized even-tempered parameters shows that there are +some minor changes in the initial exponent and first spacing parameter. Perhaps more significant is that only 12 p-type +functions were required to reach the desired accuracy, rather than 13 in the even-tempered case. The effect of the third +and fourth parameters (those introduced in the well-tempered expansion) are perhaps better visualised. + +Visualising the exponents +------------------------- + +In this tutorial example, we create a Internal Basis object as a data structure for storing and comparing the even- and +well-tempered basis sets: + +.. code-block:: python + + viz_basis = InternalBasis() + viz_basis['even'] = ne_et.get_basis()['ne'] + viz_basis['well'] = ne_wt.get_basis()['ne'] + +We then plot the exponents using a similar process as for the :ref:`sec:visualize` example. + +.. code-block:: python + + fig, ax = plot_exponents(viz_basis, atoms=['even', 'well'], split_by_shell=True) + ax[0].set_title('Even-tempered') + ax[1].set_title('Well-tempered') + +Once the plot is saved, the resulting image should look like the following. + +.. image:: comparison_plot.png + :scale: 75 % + :alt: An event plot comparing even- and well-tempered basis sets. The well-tempered version has increased spacing between the lines for larger exponent values. + +The effect of using a well-tempered expansion is particularly clear for the s shell - it can be seen that as the exponent value +increases (becomes tighter), the spacing between the exponents increases. For the p shell the exponents from the well-tempered expansion +still span a similar range as from even-tempered, even though there is one fewer function. + +Effect on energy +---------------- + +The resulting neon SCF energies are:: + + Even-tempered energy: -128.54704832 + Well-tempered energy: -128.54708446 + Numerical SCF energy: -128.54709811 + +We can see that switching to a well-tempered scheme lowers the energy by roughly 36 :math:`\mu E_h` and is significantly closer to the +numerical SCF result, despite having one fewer p function. + +.. toctree:: + :hidden: + \ No newline at end of file diff --git a/doc/src/_tutorials/comparison_plot.png b/doc/src/_tutorials/comparison_plot.png new file mode 100644 index 0000000..9fd27dc Binary files /dev/null and b/doc/src/_tutorials/comparison_plot.png differ diff --git a/doc/src/tutorials.rst b/doc/src/tutorials.rst index df2af97..42d8033 100644 --- a/doc/src/tutorials.rst +++ b/doc/src/tutorials.rst @@ -12,6 +12,7 @@ Tutorials - :doc:`_tutorials/multimol` - :doc:`_tutorials/reduce` - :doc:`_tutorials/visualize` +- :doc:`_tutorials/compare` .. toctree:: :hidden: diff --git a/examples/et_wt_compare.py b/examples/et_wt_compare.py new file mode 100644 index 0000000..9a1305b --- /dev/null +++ b/examples/et_wt_compare.py @@ -0,0 +1,32 @@ +import basisopt as bo +from basisopt.basis.atomic import AtomicBasis +from basisopt.basis.basis import InternalBasis +from basisopt.util import bo_logger + +bo.set_backend('psi4') +bo.set_tmp_dir('tmp/') + +bo_logger.info("Optimizing even-tempered basis set") +ne_et = AtomicBasis('Ne') +ne_et.set_even_tempered(method='scf', accuracy=1e-4) + +bo_logger.info("Optimizing well-tempered basis set") +ne_wt = AtomicBasis('Ne') +ne_wt.set_well_tempered(method='scf', accuracy=1e-4) + +bo_logger.info("Visualizing the difference between the sets") +viz_basis = InternalBasis() +viz_basis['even'] = ne_et.get_basis()['ne'] +viz_basis['well'] = ne_wt.get_basis()['ne'] + +import matplotlib.pyplot as plt + +from basisopt.viz.basis import plot_exponents + +fig, ax = plot_exponents(viz_basis, atoms=['even', 'well'], split_by_shell=True) +ax[0].set_title('Even-tempered') +ax[1].set_title('Well-tempered') + +filename = "comparison_plot.png" +bo_logger.info("Writing exponents plot to %s", filename) +plt.savefig(filename) diff --git a/examples/wt_strategy.py b/examples/wt_strategy.py new file mode 100644 index 0000000..62203a0 --- /dev/null +++ b/examples/wt_strategy.py @@ -0,0 +1,14 @@ +import basisopt as bo + +bo.set_backend('psi4') +bo.set_tmp_dir('tmp/') + +from basisopt.basis.atomic import AtomicBasis + +ne = AtomicBasis('Ne') +ne.set_well_tempered(method='scf', accuracy=1e-4) +print(ne.wt_params) +print(ne._molecule.get_result('energy'), ne._molecule.get_reference('energy')) +for i, p in enumerate(ne.get_basis()['ne']): + print(f"l = {i}:", p.exps) +ne.save('tmp/neon-wt.obj') diff --git a/tests/test_basis.py b/tests/test_basis.py index 773c3d8..7054f3d 100644 --- a/tests/test_basis.py +++ b/tests/test_basis.py @@ -40,6 +40,20 @@ def test_even_temper_expansion(): assert almost_equal(p_shell.exps[11], 474.989023, thresh=1e-6) +def test_well_temper_expansion(): + wt_params = [(0.1, 2.3, 12.2, 8.6, 13), (0.2, 2.1, 32.8, 9.9, 10)] + wt_basis = basis.well_temper_expansion(wt_params) + assert len(wt_basis) == 2 + + s_shell = wt_basis[0] + p_shell = wt_basis[1] + assert len(s_shell.exps) == 13 + assert almost_equal(s_shell.exps[0], 0.1, thresh=1e-6) + assert almost_equal(s_shell.exps[10], 1615.706126, thresh=1e-6) + assert len(p_shell.exps) == 10 + assert almost_equal(p_shell.exps[6], 33.623095, thresh=1e-6) + + def test_fix_ratio(): exps = np.array([0.1, 0.2, 0.58, 1.3, 3.0, 8.2]) new_exps = basis.fix_ratio(exps)