New Cp2kCube and Cp2kCell

Devel

cp2kdata/cell.py

@@ -0,0 +1,95 @@
+from ase.geometry.cell import cellpar_to_cell
+from ase.geometry.cell import cell_to_cellpar
+import numpy.typing as npt
+import numpy as np
+from copy import deepcopy
+
+class Cp2kCell:
+    def __init__(
+        self, 
+        cell_param: npt.NDArray[np.float64], 
+        grid_point: npt.NDArray[np.int_] = None, 
+        grid_spacing_matrix: npt.NDArray[np.float64] = None
+        ):
+        """
+        The documentation of Cell class used in cp2kdata
+        Parameters
+        ----------
+        cell_param: np. array
+            The length of the cell in bohr
+        divi: np. array
+            The number of grid points in each direction
+        h: float
+            The grid spacing in bohr
+
+        a b c alpha beta gamma
+        grid_point   : array
+        grid_spacing_matrix : matrix 3x3
+        """
+
+
+        if len(cell_param) == 1 and isinstance(cell_param, float):
+            self.cell_matrix = np.array(
+                [
+                    [cell_param, 0, 0], 
+                    [0, cell_param, 0], 
+                    [0, 0, cell_param]
+                ]
+            )
+            print("input cell_param is a float, the cell is assumed to be cubic")
+        elif cell_param.shape == 3:
+            self.cell_matrix = np.array(
+                [
+                    [cell_param[0], 0, 0], 
+                    [0, cell_param[1], 0], 
+                    [0, 0, cell_param[2]]
+                ]
+            )
+            print("input cell_param is a list or array, the cell is assumed to be orthorhombic")
+        elif cell_param.shape == 6:
+            self.cell_matrix = cellpar_to_cell(cell_param)
+            print("input cell_param is in [a, b, c, alpha, beta, gamma] form, it is converted to cell matrix")
+        elif cell_param.shape == (3, 3):
+            self.cell_matrix = cell_param
+            print("input cell_param is a matrix, 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):
+            print("No grid point generated")
+        elif (grid_point is None) and (grid_spacing_matrix is not None):
+            self.grid_spacing_matrix = grid_spacing_matrix
+            self.grid_point = np.round(self.cell_matrix/self.grid_spacing_matrix)
+        elif (grid_point is not None) and (grid_spacing_matrix is None):
+            self.grid_point = np.array(grid_point)
+            self.grid_spacing_matrix = self.cell_matrix/self.grid_point[:, np.newaxis]
+        elif (grid_point is not None) and (grid_spacing_matrix is not None):
+            self.grid_point = np.array(grid_point)
+            self.grid_spacing_matrix = np.array(grid_spacing_matrix)
+
+        self.grid_point = self.grid_point.astype(int)
+        self.volume = np.linalg.det(self.cell_matrix)
+        self.dv = np.linalg.det(self.grid_spacing_matrix)
+
+        self.cell_param = cell_to_cellpar(self.cell_matrix)
+
+    def copy(self):
+        return deepcopy(self)
+
+    def get_volume(self):
+        return self.volume
+
+    def get_dv(self):
+        return self.dv
+
+    def get_cell_param(self):
+        return self.cell_param
+
+    def get_cell_angles(self):
+        return self.cell_param[3:]
+
+    def get_cell_lengths(self):
+        return self.cell_param[:3]
+
+

cp2kdata/cli/cmd.py

@@ -1,5 +1,6 @@
 from email.policy import default
 import click
+from cp2kdata import Cp2kCube
 from .funcs import *
 
 
@@ -19,6 +20,11 @@ def plot():
 
 cli.add_command(plot)
 
+@click.group("cube")
+def cube():
+    click.echo('Manipulate Cube Files')
+cli.add_command(cube)
+
 #-- for gen test --# 
 #-- Cutoff --#
 @click.command()
@@ -221,4 +227,60 @@ def hubbardU(target_dir, exp_yaml):
 def ti(fig_name):
     plot_ti(fig_name=fig_name)
 
-plot.add_command(ti)
+plot.add_command(ti)
+
+
+# -- for cube -- #
+@click.command()
+@click.option(
+    '--cube_file', 
+    type=str, 
+    default=".", 
+    help='cube file'
+    )
+@click.option(
+    '--axis', 
+    type=str, 
+    default="z", 
+    help='axis'
+    )
+@click.option(
+    '--mav', 
+    type=bool, 
+    default=False, 
+    help='switch on macro average or not'
+    )
+@click.option(
+    '--l1', 
+    type=float, 
+    default=1, 
+    help='l1'
+    )
+@click.option(
+    '--l2', 
+    type=float, 
+    default=1, 
+    help='l2'
+    )
+@click.option(
+    '--ncov', 
+    type=int, 
+    default=1, 
+    help='ncov'
+    )
+@click.option(
+    '--unit', 
+    type=str, 
+    default="eV", 
+    help='unit'
+    )
+@click.option(
+    '--width', 
+    type=int, 
+    default=135, 
+    help='width'
+    )
+def view(cube_file, axis, mav, l1, l2, ncov, unit, width):
+    cube = Cp2kCube(cube_file)
+    cube.view_cube_acsii(axis=axis, mav=mav, l1=l1, l2=l2, ncov=ncov, unit=unit, width=width)
+cube.add_command(view)

cp2kdata/cli/funcs.py

@@ -2,4 +2,5 @@
 from ..test_input import write_basis_test_inp
 from ..test_input import write_hubbard_U_test_inp
 from ..plots.test_plot import plot_U_test, get_exp_collect_from_yaml, plot_cutoff_test, plot_basis_test
-from ..plots.fep_plot import plot_ti
+from ..plots.fep_plot import plot_ti
+