Update docstrings to numpydoc format in read_input.py, read_input_k.p…

…y, and wan90.py

src/hwave/qlmsio/read_input_k.py

@@ -1,3 +1,15 @@
+"""Input file reader for k-space calculations.
+
+This module provides functionality to read and parse input files for k-space 
+calculations in QLMS. It handles geometry, transfer integrals, and various 
+interaction terms.
+
+Classes
+-------
+QLMSkInput
+    Reads and parses input files for k-space calculations.
+"""
+
 import logging
 import numpy as np
 import sys
@@ -10,6 +22,46 @@
 
 
 class QLMSkInput():
+    """Input file reader for k-space calculations.
+    
+    Parameters
+    ----------
+    info_inputfile : dict
+        Dictionary containing input file information with keys:
+        - path_to_input : str
+            Path to input file directory
+        - interaction : dict
+            Dictionary specifying interaction files
+        - initial : str, optional
+            Filename for initial Green's function in k-space (NPZ format)
+        - initial_uhf : str, optional
+            Filename for initial Green's function in real space
+        - geometry_uhf : str, optional
+            Filename for geometry in real space
+        - onebodyg_uhf : str, optional
+            Filename for one-body Green's function indices
+    solver_type : str, optional
+        Type of solver to use (default: "UHFk")
+        
+    Attributes
+    ----------
+    valid_namelist : list
+        List of valid input file section names
+    ham_param : CaseInsensitiveDict
+        Dictionary storing Hamiltonian parameters
+    green : CaseInsensitiveDict
+        Dictionary storing Green's function data
+        
+    Methods
+    -------
+    get_param(key)
+        Get parameters by key
+    _read_data(file_name, value_type)
+        Read data from file into dictionary
+    _read_green(file_name) 
+        Read Green's function indices from file
+    """
+
     valid_namelist = [s.lower() for s in ["path_to_input", "Geometry", "Transfer", "CoulombIntra", "CoulombInter", "Hund", "Ising", "PairLift", "Exchange", "PairHop", "Extern"]]
 
     def __init__(self, info_inputfile, solver_type="UHFk"):
@@ -91,6 +143,25 @@ def __init__(self, info_inputfile, solver_type="UHFk"):
             self.green["onebodyg_uhf"] = self._read_green(file_name)
 
     def _read_data(self, file_name, value_type="real"):
+        """Read data from file into dictionary.
+        
+        Parameters
+        ----------
+        file_name : str
+            Name of file to read
+        value_type : str, optional
+            Type of values to read ("real" or "complex")
+            
+        Returns
+        -------
+        dict
+            Dictionary containing data read from file
+            
+        Raises
+        ------
+        FileNotFoundError
+            If input file is not found
+        """
         info = {}
         try:
             data = np.loadtxt(file_name, skiprows = 5)
@@ -110,6 +181,23 @@ def _read_data(self, file_name, value_type="real"):
         return info
 
     def _read_green(self, file_name):
+        """Read Green's function indices from file.
+        
+        Parameters
+        ----------
+        file_name : str
+            Name of file to read
+            
+        Returns
+        -------
+        ndarray
+            Array of Green's function indices
+            
+        Raises
+        ------
+        FileNotFoundError
+            If input file is not found
+        """
         try:
             _data = np.loadtxt(file_name, dtype=np.int32, skiprows = 5)
         except FileNotFoundError:
@@ -118,6 +206,21 @@ def _read_green(self, file_name):
         return _data
 
     def get_param(self, key):
+        """Get parameters by key.
+        
+        Parameters
+        ----------
+        key : str
+            Key to retrieve parameters for:
+            - "mod"/"parameter": Returns None
+            - "ham"/"hamiltonian": Returns Hamiltonian parameters
+            - "output"/"green": Returns Green's function data
+            
+        Returns
+        -------
+        CaseInsensitiveDict or None
+            Requested parameters or None if key is invalid
+        """
         if key == "mod" or key == "parameter":
             return None
         elif key == "ham" or key == "hamiltonian":

src/hwave/qlmsio/wan90.py

@@ -1,3 +1,22 @@
+"""Functions for reading/writing Wannier90 format files.
+
+This module provides functions to read and write geometry and Hamiltonian data
+in Wannier90 format.
+
+Functions
+---------
+read_geom(name_in)
+    Read geometry data from file.
+read_geometry(name_in) 
+    Read extended geometry data from file.
+write_geom(name_out, info_geometry)
+    Write geometry data to file.
+read_w90(name_in)
+    Read Wannier90 Hamiltonian data from file.
+write_w90(name_in, info_int, info_geometry, interact_shape)
+    Write Wannier90 Hamiltonian data to file.
+"""
+
 from __future__ import print_function
 
 import itertools
@@ -7,6 +26,29 @@
 logger = logging.getLogger("qlms").getChild("wan90")
 
 def read_geom(name_in):
+    """Read geometry data from file.
+    
+    Parameters
+    ----------
+    name_in : str
+        Input filename
+        
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - norb : int
+            Number of orbitals
+        - rvec : ndarray
+            Real space lattice vectors (3x3)
+        - center : ndarray
+            Orbital center positions (norb x 3)
+            
+    Raises
+    ------
+    SystemExit
+        If file not found
+    """
     try:
         with open(name_in, 'r') as f:
             # skip header
@@ -31,6 +73,33 @@ def read_geom(name_in):
     return data
 
 def read_geometry(name_in):
+    """Read extended geometry data from file.
+    
+    Parameters
+    ----------
+    name_in : str
+        Input filename
+        
+    Returns
+    -------
+    dict
+        Dictionary containing:
+        - unit_vec : ndarray
+            Unit cell vectors
+        - degree : ndarray
+            Degrees of freedom
+        - cell_vec : ndarray
+            Cell vectors
+        - site2vec : dict
+            Mapping of sites to vectors
+        - n_orb : int
+            Number of orbitals
+            
+    Raises
+    ------
+    SystemExit
+        If file not found
+    """
     info_geometry = {}
     unit_vec_line_end = 3
     cell_vec_line_start = 4
@@ -62,6 +131,15 @@ def read_geometry(name_in):
     return info_geometry
 
 def write_geom(name_out, info_geometry):
+    """Write geometry data to file.
+    
+    Parameters
+    ----------
+    name_out : str
+        Output filename
+    info_geometry : dict
+        Geometry information dictionary
+    """
     with open(name_out, "w") as fw:
         #Unit_cell
         for vec in info_geometry["unit_vec"]:
@@ -75,6 +153,27 @@ def write_geom(name_out, info_geometry):
             fw.write("{} {} {}\n".format(pos, pos, pos))
 
 def read_w90(name_in):
+    """Read Wannier90 Hamiltonian data from file.
+    
+    Parameters
+    ----------
+    name_in : str
+        Input filename
+        
+    Returns
+    -------
+    dict
+        Dictionary mapping (irvec, orbvec) tuples to complex values where:
+        - irvec : tuple
+            Real space lattice vector indices
+        - orbvec : tuple 
+            Orbital indices
+            
+    Raises
+    ------
+    SystemExit
+        If file not found
+    """
     try:
         with open(name_in, 'r') as f:
             # skip header
@@ -104,23 +203,50 @@ def read_w90(name_in):
     return data
 
 def write_w90(name_in, info_int, info_geometry, interact_shape):
+    """Write Wannier90 Hamiltonian data to file.
+    
+    Parameters
+    ----------
+    name_in : str
+        Output filename
+    info_int : dict
+        Interaction information dictionary
+    info_geometry : dict
+        Geometry information dictionary
+    interact_shape : tuple
+        Shape of interaction array
+    """
+    # Build output content as a list of strings
+    output = []
+
+    # Header information
+    norb = info_int["n_orb"]
+    Nlattice = np.prod(interact_shape)
+    n_total = (Nlattice * norb) ** 2
+
+    output.append("# wannier90 format")
+    output.append(str(norb))
+    output.append(str(n_total))
+
+    # Write list of ones, 15 per line
+    ones = np.ones(n_total, dtype=int)
+    for i in range(0, n_total, 15):
+        output.append(" ".join(map(str, ones[i:i+15])))
+
+    # Interaction data
+    fmt = "{} {} {} {} {} {} {}"
+    for idx, (site_org, int_value) in info_int["interaction"].items():
+        site = info_geometry["site2vec"][idx]
+        output.append(fmt.format(
+            site_org,
+            site[0], site[1], 
+            site[2]+1, site[3]+1,
+            int_value.real, int_value.imag
+        ))
+
+    # Write all output at once
     with open(name_in, "w") as fw:
-        fw.write("# wannier90 format\n")
-        norb = info_int["n_orb"]
-        fw.write("{}\n".format(norb))
-        Nlattice = interact_shape[0]*interact_shape[1]*interact_shape[2]
-        fw.write("{}\n".format((Nlattice*norb)**2))
-        for idx in range((Nlattice*norb)**2):
-            fw.write("1 ")
-            if (idx+1)%15 == 0:
-                fw.write("\n")
-        if (idx+1)%15 != 0:
-            fw.write("\n")
-        for idx, value in info_int["interaction"].items():
-            site = info_geometry["site2vec"][idx]
-            site_org = value[0]
-            int_value = value[1]
-            fw.write("{} {} {} {} {} {} {}\n".format(site_org, site[0], site[1], site[2]+1, site[3]+1, int_value.real, int_value.imag))
+        fw.write("\n".join(output) + "\n")
 
 if __name__ == "__main__":
     path_to_sample = "../../sample/dir-model"