Вычисление параметров молекулы водорода

0. Импорты

from hfscf import *
import numpy as np
import multiprocessing as mp
from tqdm import trange
from collections import namedtuple
from jupyterthemes import jtplot
import sys

1. Задаем функцию расчёта орбиталей на основе подхода Рутана Холла

SfcOut = namedtuple('SfcOut', 'S H X F C P')


def SCF(r=1.4632, Z=(1, 1), b1=GTO["H"], b2=GTO["H"], b=3, cores=mp.cpu_count()):
    """Solve self-consistent-field equations for two atoms."""

    R = [0, r]

    # Calculate overlap matrix (beacause basis functions are not in general orthogonal).
    S = get_S(R, b1=b1, b2=b2, b=b)

    # Calculate core-hamiltonian matrix (kinetic energy and nuclear attraction).
    H = get_H(R, Z, b1=b1, b2=b2, b=b)

    # Transformation matrix (orthogonalization of basis).
    X = diagon(S)
    # Transformation matrix adjoint.
    X_adj = X.getH()

    # Guess charge-density matrix.
    P_old = np.zeros((2, 2), dtype=np.float64)

    # Iteratively try to optimize P.
    with mp.Pool(abs(cores)) as pool, trange(50, desc=f'cores={cores}', ncols=0) as pbar:
        for _ in pbar:
            # Calculate two-electron part of the Fock matrix.
            if cores < 0:
                pool = None
            G = get_G(r, P_old, b1=b1, b2=b2, b=b, pool=pool)

            # Calculate Fock matrix.
            # F = H + G
            F = H + G

            # Calculate transformed Fock matrix.
            F_trans = X_adj * F * X

            # Diagonalize F' to obtain C'.
            C_trans = diagon2(F_trans)

            # Calculate matrix of expansion coefficients.
            # C = X * C'
            C = X * C_trans

            # Recalculate charge-density matrix.
            P_new = get_P(C)

            # Check convergence.
            if np.linalg.norm(P_new - P_old) < 1E-4:
                return SfcOut(S=S, H=H, X=X, F=F, C=C, P=P_new)
            else:
                P_old = P_new

    print('Energy minimization has stopped, but the forces have not converged to the '
          'requested precision Fmax < 10 (which may not be possible for your system).')
    return SfcOut(S=S, H=H, X=X, F=F, C=C, P=P_new)

2. Считаем и смотрим на результат

1) Считаем (и заодно сравним скорости с небольшим распараллеливанием и без него)

r = 1.398
Z = (1, 1)
b1 = GTO["H"]
b2 = GTO["H"]
b = 6

out = SCF(r, Z, b1, b2, b, cores=-1)
out = SCF(r, Z, b1, b2, b, cores=6)

cores=-1:  12% 6/50 [01:16<09:21, 12.75s/it]
cores=6:  12% 6/50 [00:19<02:23,  3.25s/it]

Ускорились почти в 4 раза -- весьма неплохо, учитывая что:

• мы поколдовали лишь над одной функцией (get_G)
• CPython концептуально не распараллеливаемый (GIL и все такое)

2) Энергии (электронная и полная; я подозреваю, в Хартри)

E_elec = ener_elec(out.P, out.H, out.F)
E_tot = ener_tot(r, Z, E_elec)
print(E_elec)
print(E_tot)

-1.8486174862021283
-1.1333099039417562

3) Строим 1D и 2D плоты орбиталей

jtplot.style(figsize=(14, 9))

orbital(out.C, r, b1=b1, b2=b2, b=b)
orbital2D(out.C, out.X, out.F, r, b1=b1, b2=b2, b=b)

Тут хорошо видно связывающую и разрыхляющую орбитали (1 и 2 соответственно).