hamtools.py

#!/usr/bin/python
# coding=utf-8
# Copyright Dmitry Korotin dmitry@korotin.name
# Based on original source by Nikolay Skorikov nskorikov@gmail.com

__author__ = "Dmitry Korotin"
__author_email__ = "dmitry@korotin.name"

import re
import numpy as np
import os.path

def distance(a, b):
    d = float(0.0)
    for ai, bi in zip(a, b):
        d += (ai - bi) ** 2
    return np.sqrt(d)

def read_cell(f):
    alat = float(next(f).strip())
    plat = [next(f).split(), next(f).split(), next(f).split()]
    for p in plat:
        for i in range(3):
            p[i] = float(p[i])
    qlat = np.linalg.inv(plat)
    return alat, plat, qlat

def read_list_of_vectors(f, n=0):
    tlist = []
    # If we get nonzero length of list of vectors, read It
    if n != 0:
        for i in range(n):
            try:
                tlist.append(next(f).strip().split())
            except StopIteration:
                break
    # let's try to read basis of unknown length
    else:
        while True:
            try:
                tmp = next(f).strip().split()
                if tmp:
                    tlist.append(tmp)
                else:
                    break
            except StopIteration:
                break
    return tlist

def read_atoms(f):
    nat = int(next(f).strip())
    atom = read_list_of_vectors(f, nat)
    return nat, atom

def read_basis(f):
    hdim, nblocks = next(f).strip(" ,\'").split()
    nblocks = int(nblocks)
    basis = read_list_of_vectors(f, nblocks)
    return basis

def read_single(f, t):
    if t == 'integer':
        return int(next(f).strip())
    if t == 'float':
        return float(next(f).strip())
    
def print_square_matrix(mtrx, key, key2='border'):
    sss = str('')
    if key == 'both' or key == 'real':
        if key == 'both':
            print('           Real part\n')
        for ss in mtrx:
            for s in ss:
                sss += '{:8.4f} '.format(s.real)
            print(sss)
            sss = ''
        if key2 == 'border':
            print('---------------------------------\n')
    if key == 'both' or key == 'imag':
        if key == 'both':
            print('           Imag part\n')
        for ss in mtrx:
            for s in ss:
                sss += '{:8.4f} '.format(s.imag)
            print(sss)
            sss = ''
        if key2 == 'border':
            print('---------------------------------\n')
            
class HermitianMatrix():
    """ 
    Hermitian matrix to be read from *.am files. Parent class for Hamiltonian and OccupationMatrix classes
    """
    def __init__(self, normalize_k_points: bool = False):
        self.timestamp = 0
        self.nspin = 1
        self.nkp = 0
        self.dim = 0
        self.kpoints = []
        self.matrix = []
        self.normalize_k_points = normalize_k_points

    def read_kpoints(self, f, nkp):
        points = []
        for i in range(nkp):
            kp = tuple( [float(t) for t in next(f).strip().split()] )
            weight = kp[0]/(2.0/self.nspin)
            coords = tuple(kp[1:4])
            points.append((weight, coords))
        if self.normalize_k_points:
            total_weight = sum(p[0] for p in points)
            if total_weight > 0:
                points = [(w / total_weight, coords) for w, coords in points]
        return points
    
    def get_k_sum(self):
        ksum = np.zeros((self.nspin, self.dim, self.dim), dtype=np.complex64)
        for isp in range(self.nspin):
            for ikp in range(self.nkp):
                ksum[isp, :, :] += self.matrix[isp, ikp, :, :] * self.kpoints[ikp][0]
                
        return ksum
    
    def read_hermitian_matrix(self, f):
        m = np.zeros((self.nspin, self.nkp, self.dim, self.dim), dtype=np.complex64)
        for isp in range(self.nspin):
            for ikp in range(self.nkp):
                for i in range(self.dim):
                    for j in range(i, self.dim):
                        tmp = [float(t) for t in next(f).strip().split()]
                        m[isp][ikp][i][j] = (tmp[0] + 1j * tmp[1])
                        m[isp][ikp][j][i] = (tmp[0] - 1j * tmp[1])
        return m

class Hamiltonian(HermitianMatrix):
    """
    Contains content of hamilt.am
    """

    def __init__(self, path='hamilt.am', normalize_k_points: bool = False):
        super().__init__(normalize_k_points=normalize_k_points)
        self.read_hamilt(path)

    def read_hamilt(self, path):
        with open(path, 'r') as f:
            for line in f:
                if re.match(r'&hash', line):
                    self.timestamp = read_single(f, 'integer')
                if re.match(r'&nspin', line):
                    self.nspin = read_single(f, 'integer')
                if re.match(r'&nkp', line):
                    self.nkp = read_single(f, 'integer')
                if re.match(r'&dim', line):
                    self.dim = read_single(f, 'integer')
                if re.match(r'&kpoints', line):
                    self.kpoints = self.read_kpoints(f, self.nkp)
                if re.match(r'&hamiltonian', line):
                    self.matrix = self.read_hermitian_matrix(f)
                    
    def print_ham(self, s):
        hsum = self.get_k_sum()
        print('\n      Sum of Hamiltonian over BZ (real part):')
        for b in s.basis:
            print('\n{:s}{:3s}-{:s}:'.format(19*' ', b[0], b[2]))
            print('{:s}Majority spin:'.format(15*' '))
            istart = b[4]-1
            iend = istart + b[3]
            print_square_matrix(hsum[0, istart:iend, istart:iend], 'real', 'suppress_border')
            if self.nspin ==2 :
                print(' ')
                print('{:s}Minority spin:'.format(15*' '))
                istart = b[4] - 1
                iend = istart + b[3]
                print_square_matrix(hsum[1, istart:iend, istart:iend], 'real', 'do_not_underline_block')
    
class OccupationMatrix(HermitianMatrix):
    """
    Contains content of occm.am
    """

    def __init__(self, path='occm.am', normalize_k_points: bool = False):
        super().__init__(normalize_k_points=normalize_k_points)
        self.read_occm(path)

    def read_occm(self, path):
        with open(path, 'r') as f:
            for line in f:
                if re.match(r'&hash', line):
                    self.timestamp = read_single(f, 'integer')
                if re.match(r'&nspin', line):
                    self.nspin = read_single(f, 'integer')
                if re.match(r'&nkp', line):
                    self.nkp = read_single(f, 'integer')
                if re.match(r'&dim', line):
                    self.dim = read_single(f, 'integer')
                if re.match(r'&kpoints', line):
                    self.kpoints = self.read_kpoints(f, self.nkp)
                if re.match(r'&occupation_matrix', line):
                    self.matrix = self.read_hermitian_matrix(f)
    
class Crystall():
    """Class Crystal contains description of considered system taken from system.am"""
    def __init__(self, path='./'):
        self.efermi = 0.0
        self.datoms = []
        self.dbasis = []
        self.read_system(path)

    def prettyprint(self):
        print( ' \nSystem under consideratin:')
        print( '')
        print( '  ALAT={:8.4f}'.format(self.cell[0]))
        print( '  PLAT=  {:10.7f} {:10.7f} {:10.7f}'.format(self.cell[1][0][0], self.cell[1][0][1], self.cell[1][0][2]))
        print( '         {:10.7f} {:10.7f} {:10.7f}'.format(self.cell[1][1][0], self.cell[1][1][1], self.cell[1][1][2]))
        print( '         {:10.7f} {:10.7f} {:10.7f}'.format(self.cell[1][2][0], self.cell[1][2][1], self.cell[1][2][2]))
        print( '         =================================')
        print( '  QLAT=  {:10.7f} {:10.7f} {:10.7f}'.format(self.cell[2][0][0], self.cell[2][0][1], self.cell[2][0][2]))
        print( '         {:10.7f} {:10.7f} {:10.7f}'.format(self.cell[2][1][0], self.cell[2][1][1], self.cell[2][1][2]))
        print( '         {:10.7f} {:10.7f} {:10.7f}'.format(self.cell[2][2][0], self.cell[2][2][1], self.cell[2][2][2]))
        print( '')
        print( '  Unit cell contains:')
        for a in self.atoms:
            print('  {:5s}  {:10.7f} {:10.7f} {:10.7f}'.format(a[0], a[1][0], a[1][1], a[1][2]))
        print('  We will consider only:')
        for a, b in zip(self.datoms, self.dbasis):
            print('  {:4s}-{:1s} {:10.7f} {:10.7f} {:10.7f}'.format(a[0], b[2], a[1][0], a[1][1], a[1][2]))

    def read_system(self,path):
        """
        Reads and parses system.am and creates appropriate attributes of class Crystall
        """
        with open(path+'system.am', 'r') as f:
            for line in f:
                if re.match(r'&cell', line):
                    self.cell = read_cell(f)
                if re.match(r'&atom', line):
                    self.nat, self.atoms = read_atoms(f)
                if re.match(r'&nelec', line):
                    self.nelec = read_single(f, 'float')
                if re.match(r'&efermi', line):
                    self.efermi = read_single(f, 'float')
                if re.match(r'&basis', line):
                    self.basis = read_basis(f)
                if re.match(r'&hash', line):
                    self.timestamp = read_single(f, 'integer')
                if re.match(r'&crystcoor', line):
                    self.crystcoord = True
        # Convert list of atoms
        for a in self.atoms:
            for i in range(1, 4):
                a[i] = float(a[i])
            a[1] = np.array(a[1:])
            del (a[2:])
        # Convert Basis
        for a in self.basis:
            a[1] = int(a[1])
            for i in range(3, len(a)):
                a[i] = int(a[i])

    def extract_d_atoms(self):
        import copy
        """
        Method extracts from self.atoms atoms with presence of d-orbital in WF basis
        self.datoms - analog of self.atoms but with d atoms only
        self.dbasis - analog of self.basis but with d states only
        """
        na = 0
        for ia, a in enumerate(self.atoms):
            for b in self.basis:
                if b[2] == 'd' and b[1]-1 == ia:
                    self.datoms.append(a)
                    c =copy.deepcopy(b)
                    self.dbasis.append(c)
                    self.dbasis[na][1] = na
                    na += 1


    def crys2cart(self,ham):
        pass

def check_rotated_orbitals(orbitals, l):
    """
    Check and correct orbital ordering for rotated coordinate systems.
    
    Based on the TB (Tight Binding) convention from the file format specification:
    - s: [1]
    - p: [2(y), 3(z), 4(x)] -> standard order [3(z), 4(x), 2(y)]
    - d: [5(xy), 6(yz), 7(3z²-r²), 8(xz), 9(x²-y²)] -> standard order [7, 8, 6, 9, 5]
    - f: [10-16] in specification order
    
    Parameters:
    -----------
    orbitals : list of int
        List of orbital indices as read from file
    l : str
        Orbital angular momentum quantum number ('s', 'p', 'd', 'f')
        
    Returns:
    --------
    list of int
        Corrected orbital indices in standard order
    """
    # Check if orbitals contain duplicates (indicating rotation/reordering needed)
    if len(set(orbitals)) != len(orbitals):
        # Standard orbital orderings for rotated systems
        if l == 'p':
            # Standard p orbital order: z, x, y (indices 3, 4, 2)
            return [3, 4, 2]
        elif l == 'd':
            # Standard d orbital order: 3z²-r², xz, yz, x²-y², xy (indices 7, 8, 6, 9, 5)
            return [7, 8, 6, 9, 5]
        elif l == 'f':
            # Standard f orbital order: y(3x²-y²), xyz, y(5z²-1), z(5z²-3), x(5z²-1), z(x²-y²), x(3y²-x²)
            # Corresponding to indices 10, 11, 12, 13, 14, 15, 16
            return [10, 11, 12, 13, 14, 15, 16]
        elif l == 's':
            # s orbital - no rotation needed
            return [1]
    
    # If no duplicates found, return original ordering
    return orbitals

def split_basis_into_blocks(path='./', merge_atom_blocks=True):
    """
    Parse system.am file to extract atomic orbital block information.
    
    Parameters:
    -----------
    path : str, default './'
        Directory path containing the system.am file
    merge_atom_blocks : bool, default True
        If True, merge blocks belonging to the same atom
        If False, keep each block separate
    
    Returns:
    --------
    tuple
        atoms_label : list of str
            Atom labels/symbols for each block
        atoms_num : list of int  
            Atom numbers for each block
        block_dims : numpy.ndarray
            Dimensions of each block
        block_start : numpy.ndarray
            Starting indices of each block
        block_orbitals : list of list
            Orbital names for each block
    """
    # Orbital names according to TB convention (indices 1-16)
    orbitals = [
        "s",                           # 1
        "p_y", "p_z", "p_x",          # 2, 3, 4
        "d_{xy}", "d_{yz}", "d_{3z^2-r^2}", "d_{xz}", "d_{x^2-y^2}",  # 5-9
        "f_{y(3x^2-y^2)}", "f_{xyz}", "f_{y(5z^2-1)}", "f_{z(5z^2-3)}", # 10-13
        "f_{x(5z^2-1)}", "f_{z(x^2-y^2)}", "f_{x(3y^2-x^2)}"          # 14-16
    ]
    
    # Initialize lists to store block information
    block_dims = []
    block_start = []
    atoms_label = []
    atoms_num = []
    block_orbitals = []
    
    # Construct full path to system.am file
    system_file = os.path.join(path, 'system.am')
    
    try:
        with open(system_file, mode='r') as f:
            # Skip lines until we find '&basis'
            line = f.readline()
            while line and line.strip() != '&basis':
                line = f.readline()
                
            if not line:  # End of file reached
                raise ValueError("'&basis' section not found in system.am file")
            
            # Read dimension and number of blocks
            header_line = f.readline().split()
            if len(header_line) < 2:
                raise ValueError("Invalid format for basis header line")
                
            total_dim = int(header_line[0])
            n_blocks = int(header_line[1])
            
            # Process all blocks
            current_atom = None
            
            for iblock in range(n_blocks):
                line = f.readline().split()
                if len(line) < 6:
                    raise ValueError(f"Invalid format for block {iblock + 1} line: {line}")
                
                atom_symbol = line[0]      # atom_sym
                atom_number = int(line[1]) # atom_num
                l_quantum = line[2]        # l_sym
                block_dim = int(line[3])   # block_dim
                start_index = int(line[4]) # block_start
                
                # Process orbitals for current block
                orbital_indices = [int(k) for k in line[5:]]
                if len(orbital_indices) != block_dim:
                    raise ValueError(f"Block {iblock + 1}: orbital count ({len(orbital_indices)}) "
                                   f"doesn't match block dimension ({block_dim})")
                
                # Check and correct orbital ordering
                corrected_orbs = check_rotated_orbitals(orbital_indices, l_quantum)
                current_block_orbs = [orbitals[k-1] for k in corrected_orbs]  # Convert to 0-based indexing
                
                # Decide whether to merge with previous block or create new one
                should_merge = (merge_atom_blocks and 
                               current_atom is not None and 
                               atom_number == current_atom)
                
                if should_merge:
                    # Merge with previous block (same atom)
                    block_dims[-1] += block_dim
                    block_orbitals[-1] += current_block_orbs
                else:
                    # Create new block
                    current_atom = atom_number
                    atoms_label.append(atom_symbol)
                    atoms_num.append(atom_number)
                    block_dims.append(block_dim)
                    block_start.append(start_index)
                    block_orbitals.append(current_block_orbs)
    
    except FileNotFoundError:
        raise FileNotFoundError(f"system.am file not found at path: {system_file}")
    except (ValueError, IndexError) as e:
        raise ValueError(f"Error parsing system.am file: {str(e)}")
    
    # Convert to numpy arrays for consistency
    block_dims = np.array(block_dims, dtype=int)
    block_start = np.array(block_start, dtype=int)
    
    # Validate total dimensions
    if sum(block_dims) != total_dim:
        print(f"Warning: Sum of block dimensions ({sum(block_dims)}) "
              f"doesn't match total dimension ({total_dim})")
    
    return atoms_label, atoms_num, block_dims, block_start, block_orbitals

def print_basis_blocks_info(atoms_label, atoms_num, block_dims, block_start, block_orbitals, merge_atom_blocks=True):
    """
    Helper function to print block information in a readable format.
    """
    merge_status = "merged" if merge_atom_blocks else "separate"
    print(f"Block Information (atom blocks {merge_status}):")
    print("=" * 80)
    print(f"{'Block':<5} {'Atom':<8} {'Atom#':<6} {'Dim':<4} {'Start':<6} {'Orbitals'}")
    print("-" * 80)
    
    for i in range(len(atoms_label)):
        # Format orbital display
        if len(block_orbitals[i]) <= 8:
            orb_str = ', '.join(block_orbitals[i])
        else:
            orb_str = ', '.join(block_orbitals[i][:6]) + f", ... ({len(block_orbitals[i])} total)"
        
        print(f"{i:<5} {atoms_label[i]:<8} {atoms_num[i]:<6} {block_dims[i]:<4} {block_start[i]:<6} {orb_str}")


# Begin main
if __name__ == "__main__":

    # np.show_config()

    s = Crystall()
    h = Hamiltonian()

    s.extract_d_atoms()
    s.prettyprint()

    # h.make_second_spin()
    h.print_ham(s)

    if (os.path.exists('./occm.am')):
        occm = OccupationMatrix()
        occm_ksum = occm.get_k_sum()
        print("\nOccupation matrix:")
        print_square_matrix(occm_ksum[0],'real','')