[Feature] Adds a well tempered Strategy and an example

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

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

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)

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 = {

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