add integration by range for Cp2kCube class

.github/workflows/pub-pypi.yml

@@ -35,7 +35,7 @@ jobs:
         --wheel
         --outdir dist/
         .
-        
+
     - name: Publish package distributions to PyPI
       if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
       uses: pypa/gh-action-pypi-publish@release/v1

cp2kdata/cell.py

@@ -1,10 +1,15 @@
-from ase.geometry.cell import cellpar_to_cell
-from ase.geometry.cell import cell_to_cellpar
-import numpy.typing as npt
+from copy import deepcopy
+
+
 import numpy as np
+import numpy.typing as npt
 from numpy.linalg import LinAlgError
-from copy import deepcopy
+from ase.geometry.cell import cellpar_to_cell
+from ase.geometry.cell import cell_to_cellpar
+
+from cp2kdata.log import get_logger
 
+logger = get_logger(__name__)
 
 class Cp2kCell:
     def __init__(
@@ -37,7 +42,7 @@ def __init__(
                     [0, 0, cell_param]
                 ]
             )
-            print("input cell_param is a float, the cell is assumed to be cubic")
+            logger.info("Input cell_param is a float, the cell is assumed to be cubic")
         elif cell_param.shape == (3,):
             self.cell_matrix = np.array(
                 [
@@ -46,24 +51,24 @@ def __init__(
                     [0, 0, cell_param[2]]
                 ]
             )
-            print("the length of input cell_param is 3, "
+            logger.info("The length of input cell_param is 3, "
                   "the cell is assumed to be orthorhombic")
         elif cell_param.shape == (6,):
             self.cell_matrix = cellpar_to_cell(cell_param)
-            print("the length of input cell_param is 6, "
+            logger.info("The length of input cell_param is 6, "
                   "the Cp2kCell assumes it is [a, b, c, alpha, beta, gamma], "
                   "which will be converted to cell matrix")
         elif cell_param.shape == (3, 3):
             self.cell_matrix = cell_param
-            print("input cell_param is a matrix with shape of (3,3), "
+            logger.info("Input cell_param is a matrix with shape of (3,3), "
                   "the cell is read as is")
         else:
             raise ValueError("The input cell_param is not supported")
 
         if (grid_point is None) and (grid_spacing_matrix is None):
             self.grid_point = None
             self.grid_spacing_matrix = None
-            print("No grid point information")
+            logger.info("No grid point information")
         elif (grid_point is None) and (grid_spacing_matrix is not None):
             self.grid_spacing_matrix = grid_spacing_matrix
             self.grid_point = np.round(
@@ -91,7 +96,7 @@ def get_dv(self):
         try:
             return np.linalg.det(self.grid_spacing_matrix)
         except LinAlgError as ae:
-            print("No grid point information is available")
+            logger.exception("No grid point information is available")
 
     def get_cell_param(self):
         return self.cell_param

cp2kdata/cube/cube.py

@@ -1,4 +1,3 @@
-import os
 from copy import deepcopy
 
 import numpy as np
@@ -10,8 +9,10 @@
 import asciichartpy as acp
 
 from cp2kdata.log import get_logger
-from cp2kdata.utils import file_content, interpolate_spline
-from cp2kdata.utils import au2A, au2eV
+from cp2kdata.utils import file_content
+from cp2kdata.utils import interpolate_spline
+from cp2kdata.utils import find_closet_idx_by_value
+from cp2kdata.units import au2A, au2eV
 from cp2kdata.cell import Cp2kCell
 
 logger = get_logger(__name__)
@@ -231,9 +232,72 @@ def write_cube(self, fname, comments='#'):
                     if grid_point[2] % 6 != 0:
                         fw.write('\n')
 
-    def get_integration(self):
+    def get_integration(self,
+                        start_x: float=None,
+                        end_x: float=None,
+                        start_y: float=None,
+                        end_y: float=None,
+                        start_z: float=None,
+                        end_z: float=None
+                        )-> float:
+
+        logger.info("Start to calculate the integration of the cube file")
+
+
+        _cell_angles = self.cell.get_cell_angles()
+        _grid_point = self.cell.grid_point
+        _gs_matrix = self.cell.grid_spacing_matrix
+
+        _x_array = np.arange(0, _grid_point[0])*_gs_matrix[0][0]
+        _y_array = np.arange(0, _grid_point[1])*_gs_matrix[1][1]
+        _z_array = np.arange(0, _grid_point[2])*_gs_matrix[2][2]
+
+        if (start_x is not None) or (end_x is not None) or (start_y is not None) or (end_y is not None) or (start_z is not None) or (end_z is not None):
+            if np.all(_cell_angles == 90.0):
+                if start_x is None:
+                    _idx_start_x = None
+                else:
+                    _idx_start_x = find_closet_idx_by_value(_x_array, start_x)
+                if end_x is None:
+                    _idx_end_x = None
+                else:
+                    _idx_end_x = find_closet_idx_by_value(_x_array, end_x)
+                if start_y is None:
+                    _idx_start_y = None
+                else:
+                    _idx_start_y = find_closet_idx_by_value(_y_array, start_y)
+                if end_y is None:
+                    _idx_end_y = None
+                else:
+                    _idx_end_y = find_closet_idx_by_value(_y_array, end_y)
+                if start_z is None:
+                    _idx_start_z = None
+                else:
+                    _idx_start_z = find_closet_idx_by_value(_z_array, start_z)
+                if end_z is None:
+                    _idx_end_z = None
+                else:
+                    _idx_end_z = find_closet_idx_by_value(_z_array, end_z)
+
+                _slice_x = slice(_idx_start_x, _idx_end_x)
+                _slice_y = slice(_idx_start_y, _idx_end_y)
+                _slice_z = slice(_idx_start_z, _idx_end_z)
+            else:
+                raise ValueError("To use integration by range, all cell angles should be 90 degree")
+        else:
+            _slice_x = slice(None)
+            _slice_y = slice(None)
+            _slice_z = slice(None)
+
+
+        logger.info(f"The integration range for x is from {_x_array[_slice_x][0]:.3f} Bohr to {_x_array[_slice_x][-1]:.3f} Bohr")
+        logger.info(f"The integration range for y is from {_y_array[_slice_y][0]:.3f} Bohr to {_y_array[_slice_y][-1]:.3f} Bohr")
+        logger.info(f"The integration range for z is from {_z_array[_slice_z][0]:.3f} Bohr to {_z_array[_slice_z][-1]:.3f} Bohr")
+
+        _cube_vals_to_integrate = self.cube_vals[_slice_x, _slice_y, _slice_z]
         dv = self.cell.get_dv()
-        result = np.sum(self.cube_vals)*dv
+        result = np.sum(_cube_vals_to_integrate)*dv
+
         return result
 
     def get_cell(self):

cp2kdata/log.py

@@ -5,7 +5,10 @@
 level = logging._nameToLevel.get(level_name, logging.INFO)
 
 # format to include timestamp and module
-logging.basicConfig(format='CP2KDATA| %(levelname)-8s %(name)-40s: %(message)s', level=level)
+if level_name == 'DEBUG':
+    logging.basicConfig(format='CP2KDATA| %(asctime)s - %(levelname)-8s %(name)-40s: %(message)s', level=level)
+else:
+    logging.basicConfig(format='CP2KDATA| %(message)s', level=level)
 # suppress transitions logging
 # logging.getLogger('transitions.core').setLevel(logging.WARNING)
 

cp2kdata/pdos/pdos.py

@@ -8,7 +8,7 @@
 from matplotlib.gridspec import GridSpec
 from scipy.ndimage import gaussian_filter1d
 
-from cp2kdata.utils import au2eV
+from cp2kdata.units import au2eV
 
 atomic_color_map = {
     'Ac': (0.39344, 0.62101, 0.45034),

cp2kdata/utils.py

@@ -1,21 +1,18 @@
 """
-this script put misc function here.
+this module sotres misc functions
 """
-#from cp2kdata import Cp2kOutput
-import numpy as np
 import os
 import shutil
+from typing import Any
 
+import numpy as np
+import numpy.typing as npt
 from ase.geometry.analysis import Analysis
 from ase.io import read, write
 
-from .log import get_logger
+from cp2kdata.log import get_logger
 
 logger = get_logger(__name__)
-# frequently used unit convertion
-au2eV = 27.211386245988
-au2A = 0.529177210903
-
 
 def create_path(path):
     path += '/'
@@ -30,13 +27,24 @@ def create_path(path):
             counter += 1
     os.makedirs(path)
 
-
 def interpolate_spline(old_x, old_y, new_x):
     from scipy import interpolate
     f = interpolate.splrep(old_x, old_y, s=0, per=True)
     new_y = interpolate.splev(new_x, f)
     return new_x, new_y
 
+def find_closet_idx_by_value(arr: npt.NDArray, value: Any) -> int:
+    """
+    Find the index of the closest value in the given array to the specified value.
+
+    Parameters:
+    arr (numpy.ndarray): The input array.
+    value (Any): The value to find the closest index for.
+
+    Returns:
+    int: The index of the closest value in the array.
+    """
+    return np.abs(arr - value).argmin()
 
 def set_pbc(pos, cell):
     """set pbc for a list of Atoms object"""

docs/cube/README.md

@@ -1,6 +1,6 @@
 # Manipulate CP2K Cube Files
 
-The `CP2KData` Python package provides tools for working with cube files generated by the CP2K quantum chemistry software. One of the standout features of this package is its ability to handle CP2K cube files and perform various analyses.
+The `CP2KData` Python package provides tools for working with cube files generated by the CP2K software. One of the features of this package is its ability to handle CP2K cube files and perform various analyses.
 
 ## Getting Started
 Import the necessary modules and load a cube file:
@@ -35,7 +35,7 @@ ase.atoms.Atoms
 ```
 
 ## Integration over space
-User can obtain the integration over space using the `get_integration()` method, for example, if you get integration using density cube, you will get the total number electrons in this cube
+User can obtain the integration over the full unit cell using the `get_integration()` method, for example, if you get integration using density cube, you will get the total number electrons in this cube
 ```python
 mycube.get_integration()
 ```
@@ -45,6 +45,21 @@ mycube.get_integration()
 ```
 The result is not exactly an integer. User should round it to integer manually.
 
+Sometimes, we need to integrate over a certain area instead of the full unit cell.
+One can use `start_x`, `end_x`, `start_y`, `end_y`, `start_z`, and `end_z` variables to define the desired range of values for the x, y, and z axes.
+```python
+mycube = Cp2kCube("example.cube")
+print(mycube.get_integration(start_z = 41.76))
+print(mycube.get_integration(end_z = 41.76))
+mycube.get_integration()
+```
+```shell
+# output
+0.8702837077294211
+0.12970844915317598
+```
+Note that the integration by area is not applied to non-orthogonal cells.
+
 ## Planar Averaging
 You can calculate planar average data from the cube file, both with and without interpolation:
 ```python

tests/test_cell/test_cell.py

@@ -70,13 +70,12 @@ def test_get_volume(self, sample_data):
         expected_volume = np.linalg.det(self._create_expected_cell_matrix(cell_param))
         assert cell.get_volume() == expected_volume
 
-    def test_get_dv(self, sample_data, capsys):
+    def test_get_dv(self, sample_data, caplog):
         cell_param, grid_point, grid_spacing_matrix = sample_data
         cell = Cp2kCell(cell_param, grid_point, grid_spacing_matrix)
         if (grid_point is None) and (grid_spacing_matrix is None):
             cell.get_dv() 
-            captured = capsys.readouterr()
-            assert captured.out.splitlines()[-1] == "No grid point information is available"
+            assert caplog.messages[-1] == "No grid point information is available"
         else:
             print(cell.cell_matrix, grid_point, grid_spacing_matrix)
             expected_grid_spacing_matrix = self._create_expected_grid_spacing_matrix(cell.cell_matrix, grid_point, grid_spacing_matrix)

tests/test_cube/gen_test_answer.py

@@ -13,6 +13,9 @@
     np.save(os.path.join(dir_name, "grid_point.npy"), cube.cell.grid_point)
     np.save(os.path.join(dir_name, "gs_matrix.npy"), cube.cell.grid_spacing_matrix)
     np.save(os.path.join(dir_name, "num_atoms.npy"), cube.num_atoms)
-    np.save(os.path.join(dir_name, "integral.npy"), cube.get_integration())
+    np.save(os.path.join(dir_name, "integration_all.npy"), cube.get_integration())
+    np.save(os.path.join(dir_name, "integration_x_5.0_7.0.npy"), cube.get_integration(start_x=5.0, end_x=7.0))
+    np.save(os.path.join(dir_name, "integration_y_5.0_7.0.npy"), cube.get_integration(start_y=5.0, end_y=7.0))
+    np.save(os.path.join(dir_name, "integration_z_5.0_7.0.npy"), cube.get_integration(start_z=5.0, end_z=7.0))
     np.save(os.path.join(dir_name, "pav.npy"), cube.get_pav())
     np.save(os.path.join(dir_name, "mav.npy"), cube.get_mav(l1=1, l2=1, ncov=2))

tests/test_cube/test_cube.py

@@ -42,6 +42,35 @@ def test_cube_vals(self, cube_and_answer):
         cube_vals = cube.cube_vals
         cube_vals_answer = np.load(os.path.join(answer_dir, "cube_vals.npy"))
         assert np.all(cube_vals == cube_vals_answer)
+
+    def test_integration(self, cube_and_answer):
+        cube = cube_and_answer[0]
+        answer_dir = cube_and_answer[1]
+        integration = cube.get_integration()
+        integration_answer = np.load(os.path.join(answer_dir, "integration_all.npy"))
+        assert np.all(integration == integration_answer)
+
+    def test_integration_by_range_x(self, cube_and_answer):
+        cube = cube_and_answer[0]
+        answer_dir = cube_and_answer[1]
+        integration = cube.get_integration(start_x=5.0, end_x=7.0)
+        integration_answer = np.load(os.path.join(answer_dir, "integration_x_5.0_7.0.npy"))
+        assert np.all(integration == integration_answer)
+
+    def test_integration_by_range_y(self, cube_and_answer):
+        cube = cube_and_answer[0]
+        answer_dir = cube_and_answer[1]
+        integration = cube.get_integration(start_y=5.0, end_y=7.0)
+        integration_answer = np.load(os.path.join(answer_dir, "integration_y_5.0_7.0.npy"))
+        assert np.all(integration == integration_answer)
+
+    def test_integration_by_range_z(self, cube_and_answer):
+        cube = cube_and_answer[0]
+        answer_dir = cube_and_answer[1]
+        integration = cube.get_integration(start_z=5.0, end_z=7.0)
+        integration_answer = np.load(os.path.join(answer_dir, "integration_z_5.0_7.0.npy"))
+        assert np.all(integration == integration_answer)
+
     def test_pav(self, cube_and_answer):
         cube = cube_and_answer[0]
         answer_dir = cube_and_answer[1]

tests/test_pdos/pdos_files/case_1_pdos/answer/answer.json

@@ -4,7 +4,7 @@
     "listidx": null,
     "kind": "1",
     "timestep": "0",
-    "fermi": 0.6222699806732536,
-    "homo_ener": 0.6222699806732536,
-    "lumo_ener": 4.414503190686633
+    "fermi": 0.6222699260317295,
+    "homo_ener": 0.6222699260317295,
+    "lumo_ener": 4.414502803049129
 }