calculate_dielectric.py

from mbd import mbd
from pymbd import get_free_atom_data
import numpy as np
import os
import sys
from contextlib import contextmanager


def get_epsilon_M_for_q(species, xyz, lattice, hirsh, q_dir, scale):
    assert np.abs(np.linalg.norm(q_dir)-1) < 1e-10
    eps = 1e-10
    alpha_0_free = get_free_atom_data(species)[0]
    alpha_0 = alpha_0_free*hirsh**scale
    alpha = mbd.do_scs_k_point(
        'R', 'dip,gg',
        xyz,
        alpha_0,
        unit_cell=lattice,
        k_point=eps*q_dir
    )
    alpha_uc = np.array([[alpha[i::3, j::3].sum() for i in range(3)] for j in range(3)])
    volume_uc = np.abs(np.linalg.det(lattice))
    return np.real(1/(1-4*np.pi*q_dir@alpha_uc@q_dir/volume_uc))


def get_epsilon_M(species, xyz, lattice, hirsh, scale):
    q_dirs = [
        np.array(x)/np.sqrt(2) for x in
        [(1, 1, 0), (1, 0, 1), (0, 1, 1), (1, -1, 0), (-1, 0, 1), (0, -1, 1)]
    ]
    r = [get_epsilon_M_for_q(species, xyz, lattice, hirsh, q_dir, scale) for q_dir in q_dirs]
    epsilon_M = np.array([
        [(r[0]+r[1]-r[2]+r[3]+r[4]-r[5])/4, (r[0]-r[3])/2, (r[1]-r[4])/2],
        [0, (r[0]-r[1]+r[2]+r[3]-r[4]+r[5])/4, (r[2]-r[5])/2],
        [0, 0, (-r[0]+r[1]+r[2]-r[3]+r[4]+r[5])/4]
    ])
    epsilon_M += epsilon_M.T
    return epsilon_M


@contextmanager
def stdout_redirected(to=os.devnull):
    fd = sys.stdout.fileno()

    def _redirect_stdout(to):
        sys.stdout.close()
        os.dup2(to.fileno(), fd)
        sys.stdout = os.fdopen(fd, 'w')

    with os.fdopen(os.dup(fd), 'w') as stdout_old:
        with open(to, 'w') as file:
            _redirect_stdout(to=file)
        try:
            yield
        finally:
            _redirect_stdout(to=stdout_old)


if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--scale', default=4/3, type=float,
        help='Hirshfeld scaling coeff for polarizability [default: 4/3]'
    )
    args = parser.parse_args()

    bohr = mbd.bohr
    n_atoms = int(next(sys.stdin))
    next(sys.stdin)
    species, xyz, hirsh = zip(*(
        (sp, [float(x) for x in r], float(h))
        for sp, *r, h
        in (next(sys.stdin).split() for _ in range(n_atoms))
    ))
    lattice = np.array([
        [float(x) for x in r]
        for _, *r
        in (next(sys.stdin).split() for _ in range(3))
    ])
    xyz = np.array(xyz)/bohr
    lattice = np.array(lattice)/bohr
    hirsh = np.array(hirsh)

    with stdout_redirected():
        mbd.init_grid(15)
        eM = get_epsilon_M(species, xyz, lattice, hirsh, scale=args.scale)

    print(eM)