From f51643f0ef7c27997d56c0f1e9ee6969816f6ed0 Mon Sep 17 00:00:00 2001
From: robinzyb <38876805+robinzyb@users.noreply.github.com>
Date: Mon, 15 Apr 2024 22:25:03 +0200
Subject: [PATCH] update get_dos() with custom dos type

---
 cp2kdata/pdos/pdos.py | 87 +++++++++++++++++++++++++++++++++----------
 docs/pdos/README.md   | 12 +++++-
 2 files changed, 79 insertions(+), 20 deletions(-)

diff --git a/cp2kdata/pdos/pdos.py b/cp2kdata/pdos/pdos.py
index acc2c06..8f4d28c 100644
--- a/cp2kdata/pdos/pdos.py
+++ b/cp2kdata/pdos/pdos.py
@@ -1,5 +1,6 @@
 import glob
 import os
+from typing import Tuple
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -190,14 +191,48 @@ def __init__(self, file_name, parse_file_name=True):
             self.project_name, self.spin, self.listidx, self.kind, self.timestep = pdos_name_parser(
                 self.file)
 
+    def get_dos(self,
+                sigma: float=1,
+                dos_type: str="total",
+                usecols: Tuple[int]=None):
+        """
+        Get the density of states (DOS) from the CP2K DOS file.
+
+        Parameters:
+        - sigma (float): The standard deviation for Gaussian smoothing of the DOS.
+        - dos_type (str): The type of DOS to retrieve. Can be "total", "s", "p", "d", or "f".
+        - usecols (Tuple[int]): The columns to use for custom DOS. Only applicable if dos_type is "custom".
+
+        Returns:
+        - smth_dos (np.ndarray): The smoothed DOS.
+        - ener (np.ndarray): The energy values corresponding to the DOS.
+        """
+        # smooth the dos data
+        dos, ener = self.get_raw_dos(dos_type=dos_type, usecols=usecols)
+        smth_dos = gaussian_filter1d(dos, sigma)
+        self.smth_dos = smth_dos
+
+        return smth_dos, ener
+
     def read_dos_element(self):
+        """
+        Reads the element from the first line of the file.
+
+        Returns:
+            str: The element extracted from the first line of the file.
+        """
         with open(self.file) as f:
             first_line = f.readline()
             element = first_line.split()[6]
-#        element = ''.join([i for i in self.kind if not i.isdigit()])
         return element
 
     def read_dos_fermi(self):
+        """
+        Reads the Fermi energy from the file.
+
+        Returns:
+            float: The Fermi energy in eV.
+        """
         # this is fermi energy not fermi level!
         # fermi energy is same as HOMO energy
         with open(self.file) as f:
@@ -208,9 +243,15 @@ def read_dos_fermi(self):
         return fermi
 
     def read_dos_energies(self):
-        energies = np.loadtxt(self.file, usecols=1)
-        energies = energies * au2eV
-        return energies
+            """
+            Reads the DOS energies from the file and converts them to electron volts (eV).
+
+            Returns:
+                numpy.ndarray: An array of DOS energies in eV.
+            """
+            energies = np.loadtxt(self.file, usecols=1)
+            energies = energies * au2eV
+            return energies
 
     @property
     def occupation(self):
@@ -219,23 +260,37 @@ def occupation(self):
         return occupation
 
     def get_homo_ener(self):
-        homo_idx = np.where(self.occupation == 0)[0][0]-1
-        homo_ener = self.energies[homo_idx]
+            """
+            Get the energy of the highest occupied molecular orbital (HOMO).
+
+            Returns:
+                float: The energy of the HOMO.
+            """
+            homo_idx = np.where(self.occupation == 0)[0][0]-1
+            homo_ener = self.energies[homo_idx]
 
-        return homo_ener
+            return homo_ener
 
     def get_lumo_ener(self):
-        lumo_ener = self.energies[self.occupation == 0][0]
-        return lumo_ener
+            """
+            Get the energy of the lowest unoccupied molecular orbital (LUMO).
 
-    def get_raw_dos(self, dos_type="total", steplen=0.1):
+            Returns:
+                lumo_ener (float): The energy of the LUMO.
+            """
+            lumo_ener = self.energies[self.occupation == 0][0]
+            return lumo_ener
+
+    def get_raw_dos(self, dos_type="total", steplen=0.1, usecols=None):
 
         file = self.file
         energies = self.energies
         fermi = self.fermi
-        steplen = 0.1
-
-        if dos_type == "total":
+        #steplen = 0.1
+        if (dos_type == "custom") and (usecols is not None):
+            weights = np.loadtxt(file, usecols=usecols).sum(axis=1)
+            print("use customed columns")
+        elif dos_type == "total":
             tmp_len = len(np.loadtxt(file, usecols=2))
             weights = np.ones(tmp_len)
         elif dos_type == "s":
@@ -259,13 +314,7 @@ def get_raw_dos(self, dos_type="total", steplen=0.1):
         self.ener = ener
         return dos, ener
 
-    def get_dos(self, sigma=1, dos_type="total"):
-        # smooth the dos data
-        dos, ener = self.get_raw_dos(dos_type=dos_type)
-        smth_dos = gaussian_filter1d(dos, sigma)
-        self.smth_dos = smth_dos
 
-        return smth_dos, ener
 
 
 PDOS_NAME_RE = re.compile(
diff --git a/docs/pdos/README.md b/docs/pdos/README.md
index 005c44c..e956c49 100644
--- a/docs/pdos/README.md
+++ b/docs/pdos/README.md
@@ -6,9 +6,19 @@
 from cp2kdata import Cp2kPdos
 dosfile = "Universality-ALPHA_k2-1_50.pdos"
 mypdos = Cp2kPdos(dosfile)
-dos, ener = mypdos.get_dos()
+dos, ener = mypdos.get_dos(sigma=1, dos_type="total")
 ```
 
+If DOS is parsed from LDOS files, `dos_type=custom` is preferred, which can be used together with `usecols`. `usecols` accepts tuple with `int` elements. For example, `usecols=(3, 4, 5)`
+
+```python
+dosfile = "water-list1-1_0.pdos
+mypdos = Cp2kPdos(dosfile)
+dos, ener = mypdos.get_dos(dos_type='custom', usecols=(3, 4, 5))
+```
+
+
+
 ## Quickplot of  PDOS Files in Single Point Energy Calculation
 
 ```python