supercell positions

from pyscf import lib

import numpy as np


from pyscf.pbc import gto, dft
import os

import numpy as np

import scipy as sp

import ase.io as io

import time

bohr_to_angstrom = 0.529177


def cell_plus_imgs(cell, nimgs):
    '''Create a supercell via nimgs[i] in each +/- direction, as in get_lattice_Ls().
    Note this function differs from :fun:`super_cell` that super_cell only
    stacks the images in + direction.

    Args:
        cell : instance of :class:`Cell`
        nimgs : (3,) array

    Returns:
        supcell : instance of :class:`Cell`
    '''
    supcell = cell.copy()
    a = cell.lattice_vectors()
    Ts = lib.cartesian_prod((np.arange(-nimgs[0],nimgs[0]+1),
                             np.arange(-nimgs[1],nimgs[1]+1),
                             np.arange(-nimgs[2],nimgs[2]+1)))
    Ls = np.dot(Ts, a)
    symbs = [atom[0] for atom in cell.atom] * len(Ls)
    coords = Ls.reshape(-1,1,3) + cell.atom_coords()
    supcell.atom = list(zip(symbs, coords.reshape(-1,3)))
    supcell.unit = 'B'
    supcell.a = np.einsum('i,ij->ij', nimgs, a)
    supcell.build(False, False, verbose=0)
    supcell.verbose = cell.verbose
    return supcell


def main():
    poscar_file = os.path.abspath('POSCAR')
    ase_atoms = io.read(poscar_file, format='vasp')

    print(ase_atoms)
    cell = gto.Cell()
    cell.from_ase(ase_atoms)
    cell.basis = 'gth-szv-molopt-sr'
    cell.pseudo = {'Ni': gto.pseudo.parse(
"""
Ni GTH-PADE-q10 GTH-LDA-q10
    2    0    8
     0.56000000    0
    3
     0.42539870    3     3.61965071    -1.19635099     0.74622158
                                        3.08896496    -1.92673583
                                                       3.05859831
     0.58408076    2     1.74222007    -0.16325873
                                        0.38634067
     0.27811348    1   -11.60842823
"""),
                  'O': 'gth-pade'
    }
    # cell.ke_cutoff = 15
    cell.unit = 'angstrom'
    cell.build()

    z = ase_atoms.get_scaled_positions()
    a = np.asarray(ase_atoms.get_cell().array)
    print("Unit cell geometry")
    for row in a:
        print("{: 5.10f} {: 5.10f} {: 5.10f}".format(row[0], row[1], row[2]))

    test_positions = []
    for ii in range(len(z)):
        test_positions.append(a[0, :] * z[ii, 0] + a[1, :] * z[ii, 1] + a[2, :] * z[ii, 2])
    test_positions = np.array(test_positions)
    assert np.allclose(test_positions, ase_atoms.get_positions())
    print("Atomic Positions in Angstroms")
    print(ase_atoms.get_positions())

    print("PYSCF Positions")
    for row in cell.atom:
        print("{} {: 5.10f} {: 5.10f} {: 5.10f}".format(row[0], row[1][0], row[1][1], row[1][2]))

    scell = cell_plus_imgs(cell, [1, 1, 1])
    for row in scell.atom:
        print("{}\t{: 5.10f}\t{: 5.10f}\t{: 5.10f}".format(row[0], row[1][0] * bohr_to_angstrom, row[1][1] * bohr_to_angstrom, row[1][2] * bohr_to_angstrom))



if  __name__ == '__main__':
    main()