record shell normalization and allow basisset renormalization

qcelemental/models/basis.py

@@ -1,11 +1,15 @@
+import math
 from enum import Enum
+from functools import lru_cache
 from typing import Dict, List, Optional
 
 from pydantic import Field, constr, validator
 
 from ..exceptions import ValidationError
 from .basemodels import ProtoModel
 
+M_SQRTPI_CUBED = math.pi * math.sqrt(math.pi)
+
 
 class HarmonicType(str, Enum):
     """
@@ -16,6 +20,13 @@ class HarmonicType(str, Enum):
     cartesian = "cartesian"
 
 
+class NormalizationScheme(str, Enum):
+
+    cca = "cca"
+    erd = "erd"
+    bse = "bse"
+
+
 class ElectronShell(ProtoModel):
     """
     Information for a single electronic shell
@@ -28,6 +39,9 @@ class ElectronShell(ProtoModel):
         ...,
         description="General contraction coefficients for this shell, individual list components will be the individual segment contraction coefficients.",
     )
+    normalization_scheme: Optional[NormalizationScheme] = Field(
+        None, description="Normalization scheme for this shell."
+    )
 
     @validator("coefficients")
     def _check_coefficient_length(cls, v, values):
@@ -46,6 +60,45 @@ def _check_general_contraction_or_fused(cls, v, values):
 
         return v
 
+    @validator("normalization_scheme", always=True)
+    def _check_normsch(cls, v, values):
+
+        # Bad construction, pass on errors
+        try:
+            normsch = cls._calculate_normsch(values["angular_momentum"], values["exponents"], values["coefficients"])
+        except KeyError:
+            return v
+
+        if v is None:
+            v = normsch
+        else:
+            if v != normsch:
+                raise ValidationError(f"Calculated normalization_scheme ({normsch}) does not match supplied ({v}).")
+
+        return v
+
+    @classmethod
+    def _calculate_normsch(cls, am: int, exps: List[float], coefs: List[List[float]]) -> NormalizationScheme:
+        def single_am(idx, l):
+            m = l + 1.5
+
+            # try CCA
+            candidate_already_normalized_coefs = coefs[idx]
+            norm = cls._cca_contraction_normalization(l, exps, candidate_already_normalized_coefs)
+            if abs(norm - 1) < 1.0e-10:
+                return NormalizationScheme.cca
+
+            # try ERD
+            candidate_already_normalized_coefs = [coefs[idx][i] / pow(exps[i], 0.5 * m) for i in range(len(exps))]
+            norm = cls._erd_normalization(l, exps, candidate_already_normalized_coefs)
+            if abs(norm - 1) < 1.0e-10:
+                return NormalizationScheme.erd
+
+            # BSE not confirmable
+            return NormalizationScheme.bse
+
+        return _collapse_equal_list(single_am(idx, l) for idx, l in enumerate(am))
+
     def nfunctions(self) -> int:
         """
         Computes the number of basis functions on this shell.
@@ -73,6 +126,126 @@ def is_contracted(self) -> bool:
 
         return (len(self.coefficients) != 1) and (len(self.angular_momentum) == 1)
 
+    def normalize_shell(self, dtype: NormalizationScheme) -> "ElectronShell":
+        """Construct new ElectronShell with coefficients normalized by ``dtype``."""
+
+        naive_coefs = self._denormalize_to_bse()
+
+        bse_shell = self.dict()
+        bse_shell["coefficients"] = naive_coefs
+        bse_shell["normalization_scheme"] = "bse"
+        bse_shell = ElectronShell(**bse_shell)
+
+        if dtype == "bse":
+            return bse_shell
+        elif dtype in ["cca", "psi4"]:
+            norm_coef = bse_shell._cca_normalize_shell()
+        elif dtype == "erd":
+            norm_coef = bse_shell._erd_normalize_shell()
+
+        new_shell = self.dict()
+        new_shell["coefficients"] = norm_coef
+        new_shell["normalization_scheme"] = dtype
+        return ElectronShell(**new_shell)
+
+    def _denormalize_to_bse(self) -> List[List[float]]:
+        """Compute replacement coefficients for any-normalization shell ``self`` that are within a scale factor of BSE unnormalization."""
+
+        def single_am(idx, l):
+
+            if self.normalization_scheme == "cca":
+                prim_norm = self._cca_primitive_normalization(l, self.exponents)
+                return [self.coefficients[idx][i] / prim_norm[i] for i in range(len(self.exponents))]
+
+            elif self.normalization_scheme == "erd":
+                m = l + 1.5
+                return [self.coefficients[idx][i] / pow(self.exponents[i], 0.5 * m) for i in range(len(self.exponents))]
+
+            elif self.normalization_scheme == "bse":
+                return self.coefficients[idx]
+
+        return [single_am(idx, l) for idx, l in enumerate(self.angular_momentum)]
+
+    @staticmethod
+    def _cca_primitive_normalization(l: int, exps: List[float]) -> List[float]:
+        """Compute CCA normalization factor for primitive shell using angular momentum ``l`` and exponents ``exps``."""
+        m = l + 1.5
+        prim_norm = [
+            math.sqrt((pow(2.0, l) * pow(2.0 * exps[p], m)) / (M_SQRTPI_CUBED * _df(2 * l))) for p in range(len(exps))
+        ]
+
+        return prim_norm
+
+    @staticmethod
+    def _cca_contraction_normalization(l: int, exps: List[float], coefs: List[List[float]]) -> float:
+        """Compute CCA normalization factor for coefficients ``coefs`` using angular momentum ``l`` and exponents ``exps``."""
+
+        m = l + 1.5
+        summ = 0.0
+        for i in range(len(exps)):
+            for j in range(len(exps)):
+                z = pow(exps[i] + exps[j], m)
+                summ += (coefs[i] * coefs[j]) / z
+
+        tmp = (M_SQRTPI_CUBED * _df(2 * l)) / pow(2.0, l)
+        norm = math.sqrt(1.0 / (tmp * summ))
+        # except (ZeroDivisionError, ValueError): [idx][i] = 1.0
+
+        return norm
+
+    def _cca_normalize_shell(self) -> List[List[float]]:
+        """Compute replacement coefficients for unnormalized (BSE-normalized) shell ``self`` that fulfill CCA normalization."""
+
+        if self.normalization_scheme != "bse":
+            raise TypeError('Unnormalized shells expected. Use ``normalize_shell(dtype="cca")`` for flexibility.')
+
+        def single_am(idx, l):
+            prim_norm = ElectronShell._cca_primitive_normalization(l, self.exponents)
+            norm = ElectronShell._cca_contraction_normalization(
+                l, self.exponents, [self.coefficients[idx][i] * prim_norm[i] for i in range(len(self.exponents))]
+            )
+            return [self.coefficients[idx][i] * norm * prim_norm[i] for i in range(len(self.exponents))]
+
+        return [single_am(idx, l) for idx, l in enumerate(self.angular_momentum)]
+
+    def _erd_normalization(l: int, exps: List[float], coefs: List[List[float]]) -> float:
+        """Compute ERD normalization factor for coefficients ``coefs`` using angular momentum ``l`` and exponents ``exps``."""
+
+        m = l + 1.5
+        summ = 0.0
+        for i in range(len(exps)):
+            for j in range(i + 1):
+                temp = coefs[i] * coefs[j]
+                temp3 = 2.0 * math.sqrt(exps[i] * exps[j]) / (exps[i] + exps[j])
+                temp *= pow(temp3, m)
+
+                summ += temp
+                if i != j:
+                    summ += temp
+
+        prefac = 1.0
+        if l > 1:
+            prefac = pow(2.0, 2 * l) / _df(2 * l)
+        norm = math.sqrt(prefac / summ)
+
+        return norm
+
+    def _erd_normalize_shell(self) -> List[List[float]]:
+        """Compute replacement coefficients for unnormalized (BSE-normalized) shell ``self`` that fulfill ERD normalization."""
+
+        if self.normalization_scheme != "bse":
+            raise TypeError('Unnormalized shells expected. Use ``normalize_shell(dtype="erd")`` for flexibility.')
+
+        def single_am(idx, l):
+            m = l + 1.5
+
+            norm = ElectronShell._erd_normalization(l, self.exponents, self.coefficients[idx])
+            return [
+                self.coefficients[idx][i] * norm * pow(self.exponents[i], 0.5 * m) for i in range(len(self.exponents))
+            ]
+
+        return [single_am(idx, l) for idx, l in enumerate(self.angular_momentum)]
+
 
 class ECPType(str, Enum):
     """
@@ -189,3 +362,52 @@ def _calculate_nbf(cls, atom_map, center_data) -> int:
             ret += center_count[center]
 
         return ret
+
+    def normalize_shell(self, dtype: NormalizationScheme) -> "BasisSet":
+        """Construct new BasisSet with coefficients of all shells in center_data normalized by ``dtype``."""
+
+        new_bs = self.dict()
+
+        for lbl, center in self.center_data.items():
+            for ish, sh in enumerate(center.electron_shells):
+                new_bs["center_data"][lbl]["electron_shells"][ish] = sh.normalize_shell(dtype)
+
+        return BasisSet(**new_bs)
+
+    def normalization_scheme(self) -> NormalizationScheme:
+        """Identify probable normalization scheme governing shell ``coefficients`` in center_data.
+
+        Returns
+        -------
+        NormalizationScheme
+            Satisfied by all ElectronShells.
+
+        Raises
+        ------
+        TypeError
+            If the BasisSet's ElectronShells are detected to have inconsistent normalization schemes.
+
+        """
+        shell_norm = []
+        for lbl, center in self.center_data.items():
+            for ish, sh in enumerate(center.electron_shells):
+                shell_norm.append(sh.normalization_scheme)
+
+        return _collapse_equal_list(shell_norm)
+
+
+@lru_cache(maxsize=500)
+def _df(i):
+    if i in [0, 1, 2]:
+        return 1.0
+    else:
+        return (i - 1) * _df(i - 2)
+
+
+def _collapse_equal_list(lst):
+    lst = list(lst)
+    first = lst[0]
+    if lst.count(first) == len(lst):
+        return first
+    else:
+        raise TypeError(f"Inconsistent members in list: {lst}")