mlmm-toolkit 0.2.2.dev0__py3-none-any.whl

pysisyphus/wavefunction/Basis.py

@@ -0,0 +1,239 @@
+import functools
+import itertools as it
+import json
+from pathlib import Path
+from typing import Tuple
+
+import numpy as np
+
+from pysisyphus.config import BASIS_LIB_DIR
+from pysisyphus.elem_data import ATOMIC_NUMBERS
+from pysisyphus.Geometry import Geometry
+from pysisyphus.helpers_pure import file_or_str
+from pysisyphus.wavefunction import Shell, Shells
+from pysisyphus.wavefunction.helpers import L_MAP
+
+
+def basis_from_json(name):
+    basis_path = Path(name).with_suffix(".json")
+    if not basis_path.is_absolute():
+        basis_path = BASIS_LIB_DIR / basis_path
+
+    with open(basis_path) as handle:
+        data = json.load(handle)
+    elements = data["elements"]
+    return elements
+
+
+def basis_from_orca_str(basis_str):
+    """
+    %basis
+     newgto 1
+     s 2
+     1 0.7 0.5
+     1 0.8 0.6
+     p 2
+     1 0.7 0.5
+     1 0.8 0.6
+     d 2
+     1 0.7 0.5
+     1 0.8 0.6
+     d 3
+     1 0.7 0.5
+     1 0.8 0.6
+     1 0.6 0.4
+     end
+    end
+    """
+    basis = dict()
+    for line in basis_str.strip().split("\n"):
+        line = line.strip().lower()
+        # Skip comments
+        if any([line.startswith(key) for key in ("#", "%basis", "end")]) or line == "":
+            continue
+        tokens = line.split()
+        if tokens[0] == "newgto":
+            atomic_num = int(tokens[1])
+            basis.setdefault(atomic_num, dict())
+        elif tokens[0].isalpha():
+            ang_mom, nprims = tokens
+            ang_mom = L_MAP[ang_mom]
+            nprims = int(nprims)
+            prev_counter = 0
+            shell = {
+                "angular_momentum": ang_mom,
+                "coefficients": list(),
+                "exponents": list(),
+            }
+            basis[atomic_num].setdefault("electron_shells", list()).append(shell)
+        else:
+            counter, exponent, coefficient = [float(t) for t in tokens]
+            counter = int(counter)
+            if coefficient <= 1e-8:
+                continue
+            assert counter == (prev_counter + 1)
+            shell["coefficients"].append(coefficient)
+            shell["exponents"].append(exponent)
+    return basis
+
+
+def basis_from_pyscf_str(basis_str):
+    """
+    # Comment1
+    He    S
+         13.6267000              0.1752300
+          1.9993500              0.8934830
+          0.3829930              0.0000000
+    He    S
+         13.6267000              0.0000000
+          1.9993500              0.0000000
+          0.3829930              1.0000000
+    """
+    basis = dict()
+    for line in basis_str.strip().split("\n"):
+        line = line.strip().lower()
+        # Skip comments
+        if line.startswith("#") or (line == "") or ("basis" in line) or ("end" in line):
+            continue
+        tokens = line.split()
+        if tokens[0].isalpha():
+            element, ang_mom = [t.lower() for t in tokens]
+            ang_mom = L_MAP[ang_mom]
+            if ang_mom == "sp":
+                raise Exception("SP shells are not yet handled")
+            atomic_num = ATOMIC_NUMBERS[element]
+            shell = {
+                "angular_momentum": ang_mom,
+                "coefficients": list(),
+                "exponents": list(),
+            }
+            basis.setdefault(atomic_num, dict()).setdefault(
+                "electron_shells", list()
+            ).append(shell)
+        else:
+            exponent, coefficient = [float(t) for t in tokens]
+            if coefficient <= 1e-8:
+                continue
+            shell["coefficients"].append(coefficient)
+            shell["exponents"].append(exponent)
+    return basis
+
+
+@file_or_str(".gbs", ".GBS")
+def basis_from_gbs(text):
+    """
+    ****
+    Be 0
+    S 6 1.00
+    1.101996057600000e-02 -3.733117104118833e-01
+    3.570467226624002e-02 -9.552914830540544e-01
+    1.156831381426176e-01  9.464046666860781e-02
+    3.748133675820811e-01 -1.194926861090494e+00
+    1.214395310965943e+00  2.762798723086632e-01
+    3.934640807529656e+00 -1.847390284420840e+00
+    """
+    basis = dict()
+    for line in text.strip().split("\n"):
+        line = line.strip().lower()
+        # Skip empty lines
+        if (line == "") or ("*" in line):
+            continue
+        tokens = line.split()
+        # New center identifier line
+        if (tokens[0].isalpha()) and (len(tokens) == 2):
+            atom, zero = tokens
+            atomic_num = str(ATOMIC_NUMBERS[atom.lower()])
+            assert int(zero) == 0, "Only one basis per atom type is supported!"
+            basis.setdefault(atomic_num, dict())
+        # New shell
+        elif tokens[0].isalpha() and (len(tokens) == 3):
+            ang_mom, nprims, scale_fact = tokens
+            ang_mom = L_MAP[ang_mom.lower()]
+            nprims = int(nprims)
+            shell = {
+                "angular_momentum": [
+                    ang_mom,
+                ],
+                "coefficients": [[]],
+                "exponents": list(),
+            }
+            basis[atomic_num].setdefault("electron_shells", list()).append(shell)
+        else:
+            exponent, coefficient = [float(t) for t in tokens]
+            shell["coefficients"][0].append(coefficient)
+            shell["exponents"].append(exponent)
+    return basis
+
+
+@functools.singledispatch
+def shells_with_basis(
+    atoms: Tuple,
+    coords,
+    basis=None,
+    name=None,
+    shells_cls=None,
+    min_coeff=1e-8,
+    **kwargs
+):
+    assert (basis is not None) or (name is not None)
+    if shells_cls is None:
+        shells_cls = Shells
+    if name is not None:
+        basis = basis_from_json(name)
+
+    coords3d = np.reshape(coords, (len(atoms), 3))
+    shells = list()
+    for i, (atom, c3d) in enumerate(zip(atoms, coords3d)):
+        Zs = str(ATOMIC_NUMBERS[atom.lower()])
+        basis_shells = basis[Zs]["electron_shells"]
+        for bshell in basis_shells:
+            L = bshell["angular_momentum"]
+            # Two special cases can appear:
+            # SP shells where L has two instead of only one entry
+            # General contractions where L has length one, but multiple sets of coefficients
+            # are given.
+            #
+            # With SP shells, there will be as many L values (2), as there are sets
+            # of contraction coefficients (2). With general contractions, there will be
+            # only one L value, but multiple sets of contraction coefficients.
+            #
+            # Segemented contractions will always have one L value and one set of contraction
+            # coefficients. Long story short: We can treat everything on an equal footing if
+            # we just repeat the L value as many times as required, when only one L value is
+            # present.
+            if len(L) == 1:
+                L = it.cycle(L)
+            exponents = np.array(bshell["exponents"], dtype=float)
+            # Unpack possible general contractions or SP shells
+            for L, coeffs in zip(L, bshell["coefficients"]):
+                # For now we filter out 0.0 contraction coefficients, as present in
+                # generally contracted basis sets.
+                coeffs = np.array(coeffs, dtype=float)
+                valid = np.abs(coeffs) >= min_coeff
+                shell = Shell(
+                    L=L,
+                    center=c3d,
+                    coeffs=coeffs[valid],
+                    exps=exponents[valid],
+                    atomic_num=Zs,
+                    center_ind=i,
+                )
+                shells.append(shell)
+    shells = shells_cls(shells, **kwargs)
+    return shells
+
+
+@shells_with_basis.register
+def _(geom: Geometry, **kwargs):
+    return shells_with_basis(geom.atoms, geom.coords, **kwargs)
+
+
+class Basis:
+    """
+    Read basis sets from files.
+    Bring them in a suitable order. 1s2s2p3s3p4s3d etc.
+
+    NOT YET USED
+    """
+
+    pass

pysisyphus/wavefunction/DIIS.py

@@ -0,0 +1,76 @@
+from typing import Dict, Optional
+
+import numpy as np
+
+
+ArrDict = Dict[str, np.ndarray]
+
+
+class DIIS:
+    def __init__(self, arr_specs: ArrDict):
+        """DIIS container.
+
+        Parameters
+        ----------
+        arr_specs
+            Dict with string-keys and array-values. The key 'err_vecs' must always be
+            present. Additional keys can be provided, depending on what must be stored
+            next to the error vectors.
+            The arrays should be at least 2d, with the first dimension having size 'nvecs'.
+        """
+        self.nvecs = arr_specs["err_vecs"].shape[0]
+        for k, v in arr_specs.items():
+            assert v.shape[0] == self.nvecs
+            setattr(self, k, v)
+
+        self.ref_keys = set(arr_specs.keys())
+        self.cur_nvecs = 0
+        self.cur_index = 0
+
+    def store(self, arrs: ArrDict) -> int:
+        """Store arrays.
+
+        Returns the index, where the arrays were stored in their respective arrays.
+        """
+        assert set(arrs.keys()) == self.ref_keys
+        index = self.cur_index % self.nvecs
+
+        for key in self.ref_keys:
+            getattr(self, key)[index] = arrs[key].copy()
+        self.cur_nvecs = min(self.cur_nvecs + 1, self.nvecs)
+        self.cur_index += 1
+        # print(f"\tstored at {index=}, {self.cur_nvecs=}, {self.cur_index=}")
+        return index
+
+    def reset(self):
+        """Reset DIIS container."""
+        self.cur_index = 0
+        self.cur_nvecs = 0
+        # Also zero the arrays?!
+
+    def get_coeffs(self) -> Optional[np.ndarray]:
+        """Try to calculate DIIS coefficients from the error matrix."""
+        if not self.can_get_coeffs():
+            return None
+
+        err_vecs = self.err_vecs[: self.cur_nvecs]
+        err_mat = np.einsum("ij,kj->ik", err_vecs, err_vecs)
+        try:
+            coeffs = np.linalg.solve(err_mat, np.ones(self.cur_nvecs))
+            coeffs /= coeffs.sum()
+        except np.linalg.LinAlgError:
+            coeffs = None
+            self.reset()
+        return coeffs
+
+    def get(self, key: str) -> np.ndarray:
+        return getattr(self, key)[: self.cur_nvecs]
+
+    def can_get_coeffs(self) -> bool:
+        return self.cur_nvecs >= 2
+
+    def __str__(self):
+        return f"DIIS(nvecs={self.nvecs}, {self.cur_nvecs} stored)"
+
+    def __repr__(self):
+        return self.__str__()

pysisyphus/wavefunction/__init__.py

@@ -0,0 +1,25 @@
+import logging
+
+from pysisyphus import logger as pysis_logger
+
+logger = pysis_logger.getChild("wavefunction")
+logger.setLevel(logging.DEBUG)
+
+from pysisyphus.wavefunction.shells import (
+    get_l,
+    AOMixShells,
+    MoldenShells,
+    ORCAShells,
+    ORCAMoldenShells,
+    Shell,
+    Shells,
+)
+
+from pysisyphus.wavefunction.excited_states import norm_ci_coeffs
+from pysisyphus.wavefunction.wavefunction import Wavefunction
+from pysisyphus.wavefunction.localization import (
+    cholesky,
+    edmiston_ruedenberg,
+    foster_boys,
+    pipek_mezey,
+)

pysisyphus/wavefunction/build_ext.py

@@ -0,0 +1,42 @@
+from cffi import FFI
+
+
+def load_h(h_fn):
+    with open(h_fn) as handle:
+        text = handle.read().strip()
+    return text
+
+
+def run():
+    ffibuilder = FFI()
+    dir_ = "devel_ints/"
+    names = ("ovlp3d", "dipole3d", "diag_quadrupole3d", "quadrupole3d", "kinetic3d")
+
+    all_sources = list()
+    all_includes = list()
+    all_cdefs = list()
+    for name in names:
+        c_name = dir_ + name + ".c"
+        all_sources.append(c_name)
+        h_name = dir_ + name + ".h"
+        all_includes.append(h_name)
+        cdefs = load_h(h_name)
+        all_cdefs.append(cdefs)
+
+    cdefs = "\n\n".join(all_cdefs)
+    includes = "\n".join([f'#include "{h_name}"' for h_name in all_includes])
+
+    ffibuilder.cdef(cdefs)
+    ffibuilder.set_source(
+        "_ints1el",
+        includes,
+        sources=all_sources,
+        libraries=[
+            "m",
+        ],
+    )
+    ffibuilder.compile(verbose=True)
+
+
+if __name__ == "__main__":
+    run()

pysisyphus/wavefunction/cart2sph.py

@@ -0,0 +1,190 @@
+#!/usr/bin/env python
+
+# [1] https://doi.org/10.1002/qua.560540202
+#     Transformation between Cartesian and pure spherical harmonic Gaussians
+#     Schlegel, Frisch, 1995
+
+import argparse
+import functools
+from cmath import sqrt
+from math import floor
+import sys
+from typing import Dict, List, Tuple
+
+import numpy as np
+from numpy.typing import NDArray
+from scipy.special import factorial as fact
+from scipy.special import binom
+
+from pysisyphus.wavefunction.helpers import canonical_order
+from pysisyphus.wavefunction.normalization import get_lmn_factors
+
+
+ZERO_THRESH = 1e-14
+
+
+@functools.cache
+def cart2sph_coeff(lx: int, ly: int, lz: int, m: int) -> complex:
+    """Coefficient to convert Cartesian to spherical harmonics.
+
+    Eq. (15) in [1].
+
+    Paramters
+    ---------
+    lx
+        Cartesian quantum number.
+    ly
+        Cartesian quantum number.
+    lz
+        Cartesian quantum number.
+    m
+        Magnetic quantum number.
+
+    Returns
+    -------
+    coeff
+        Contribution of the given Cartesian basis function to the spherical
+        harmonic Y^(lx + ly + lz)_m.
+    """
+
+    l = lx + ly + lz
+    assert -l <= m <= l
+
+    j = (lx + ly - abs(m)) / 2
+    if (j % 1) != 0:  # Return when j is not integer
+        return 0
+    j = int(j)
+
+    lfact = fact(l)  # l!
+    lmm = l - abs(m)  # (l - |m|)
+    sign = 1 if m >= 0 else -1
+
+    coeff = 0.0j
+    for i in range(floor(lmm / 2) + 1):
+        i_fact = (
+            binom(l, i)
+            * binom(i, j)
+            * (-1) ** i
+            * fact(2 * l - 2 * i)
+            / fact(lmm - 2 * i)
+        )
+        k_sum = 0.0
+        for k in range(j + 1):
+            prefact = (-1) ** (sign * (abs(m) - lx + 2 * k) / 2)
+            k_sum += prefact * binom(j, k) * binom(abs(m), lx - 2 * k)
+        coeff += i_fact * k_sum
+    coeff *= sqrt(
+        fact(2 * lx)
+        * fact(2 * ly)
+        * fact(2 * lz)
+        * lfact
+        * fact(lmm)
+        / (fact(2 * l) * fact(lx) * fact(ly) * fact(lz) * fact(l + abs(m)))
+    )
+    coeff *= 1 / (2**l * lfact)
+
+    return coeff
+
+
+@functools.cache
+def cart2sph_coeffs_for(
+    l: int,
+    real: bool = True,
+    zero_small: bool = True,
+    zero_thresh: float = ZERO_THRESH,
+) -> NDArray:
+    all_coeffs = list()
+    for lx, ly, lz in canonical_order(l):
+        coeffs = list()
+        for m in range(-l, l + 1):
+            coeff = cart2sph_coeff(lx, ly, lz, m)
+            # Form linear combination for m != 0 to get real spherical orbitals.
+            if real and m != 0:
+                coeff_minus = cart2sph_coeff(lx, ly, lz, -m)
+                sign = 1 if m < 0 else -1
+                coeff = (coeff + sign * coeff_minus) / sqrt(sign * 2)
+            coeffs.append(coeff)
+        all_coeffs.append(coeffs)
+    C = np.array(all_coeffs, dtype=complex)
+    if real:
+        assert real == ~np.iscomplex(C).all()  # C should be real
+        C = np.real(C)
+    if zero_small:
+        mask = np.abs(C) <= zero_thresh
+        C[mask] = 0.0
+    # Final shape will be (spherical, Cartesian)
+    # Return read-only view; otherwise it would be too easy to mess up the
+    # cached result by inplace operations on C.
+    C = C.T.view()
+    C.flags.writeable = False
+    return C
+
+
+@functools.cache
+def expand_sph_quantum_numbers(
+    Lm, zero_thresh=ZERO_THRESH, with_lmn_factors=False
+) -> Tuple[NDArray[float], List[Tuple[int, int, int]]]:
+    """Factors and Cart. angular momentum vectors for given sph. quantum numbers."""
+    L, m = Lm
+    m += L  # Map m from [-L, L] onto [0, 2*L]
+    C = cart2sph_coeffs_for(L)  # shape (spherical, Cartesian)
+    # Include lmn-factor that appears when contracted functions are normalized
+    # according to 'normalization.norm_cgto_lmn(L).
+    if with_lmn_factors:
+        lmn_factors = get_lmn_factors(L)
+        C = C * lmn_factors[None, :]
+    indices, *_ = np.where(np.abs(C[m]) > zero_thresh)
+    factors = C[m, indices].copy()
+
+    cart_lmn = [lmn for i, lmn in enumerate(canonical_order(L)) if i in indices]
+    return factors, cart_lmn
+
+
+def cart2sph_coeffs(
+    l_max: int,
+    **kwargs,
+) -> Dict[int, NDArray]:
+    coeffs = {l: cart2sph_coeffs_for(l, **kwargs) for l in range(l_max + 1)}
+    return coeffs
+
+
+def cart2sph_nlms(l_max: int) -> Dict[int, Tuple[Tuple[int, int, int]]]:
+    nlms = dict()
+    for l in range(l_max + 1):
+        n = l
+        nlms[l] = tuple([(n, l, m) for m in range(-l, l + 1)])
+    return nlms
+
+
+def print_coeffs(l: int, C: NDArray):
+    print("".join([f"{m: >13d}" for m in range(-l, l + 1)]))
+    for (lx, ly, lz), line in zip(canonical_order(l), C):
+        fmt = "{:>12.4f}" * len(line)
+        verb = "".join(["x"] * lx + ["y"] * ly + ["z"] * lz)
+        print(f"{verb}: {fmt.format(*line)}")
+
+
+def parse_args(args):
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("l_max", type=int, default=3)
+    parser.add_argument("--not-normalized", action="store_true")
+
+    return parser.parse_args(args)
+
+
+def run():
+    args = parse_args(sys.argv[1:])
+
+    l_max = args.l_max
+    normalized = not args.not_normalized
+
+    Cs = cart2sph_coeffs(l_max, normalized=normalized)
+    for l, C in Cs.items():
+        print(f"{l=}")
+        print_coeffs(l, C.T)
+        print()
+
+
+if __name__ == "__main__":
+    run()