update some type

Author: jxxcr

src/mdkits/util/arg_type.py

@@ -1,5 +1,8 @@
 import click, os
-from . import cp2k_input_parsing
+from ase.io import read
+import numpy as np
+from ase.collections import g2
+from mdkits.util import os_operation, cp2k_input_parsing, out_err
 
 
 class CellType(click.ParamType):
@@ -8,16 +11,17 @@ class CellType(click.ParamType):
     def convert(self, value, param, ctx):
         if isinstance(value, str):
             if ',' not in value:
-                cell = cp2k_input_parsing.parse_cell(value)
+                cell = cp2k_input_parsing.parse_cell()
                 return cell
             else:
                 cell = [float(x) for x in value.split(',')]
 
                 if len(cell) == 3:
-                    click.echo(f"system cell: x = {cell[0]}, y = {cell[1]}, z = {cell[2]}, a = {90}\u00B0, b = {90}\u00B0, c = {90}\u00B0")
-                    return cell + [90, 90, 90]
+                    cell += [90, 90, 90]
+                    out_err.cell_output(cell)
+                    return cell
                 elif len(cell) == 6:
-                    click.echo(f"system cell: x = {cell[0]}, y = {cell[1]}, z = {cell[2]}, a = {cell[3]}\u00B0, b = {cell[4]}\u00B0, c = {cell[5]}\u00B0")
+                    out_err.cell_output(cell)
                     return cell
                 else:
                     self.fail(f"{value} is not a valid cell parameter", param, ctx)
@@ -37,16 +41,45 @@ def convert(self, value, param, ctx):
                 self.fail(f"{value} is not a valid frame range", param, ctx)
 
 
-from ase.collections import g2
+class StructureType(click.ParamType):
+    name = "structure file type"
+    def convert(self, value, param, ctx):
+        no_cell=np.array([0., 0., 0., 90., 90., 90.])
+        if isinstance(value, str):
+            if os.path.exists(value):
+                try:
+                    atoms = read(value)
+                except:
+                    self.fail(f"{value} is not a valid structure file", param, ctx)
+
+                if np.array_equal(atoms.cell.cellpar(), no_cell):
+                    cell = cp2k_input_parsing.parse_cell()
+                    atoms.set_cell(cell)
+
+                return atoms
+            else:
+                self.fail(f"{value} is not exists", param, ctx)
+
+
+
 class MoleculeType(click.Choice):
     name = "mocular type"
     def __init__(self):
         super().__init__(self)
         g2.names.append(click.Path(exists=True))
         self.choices = tuple(g2.names)
 
+class AdsSiteType(click.Choice):
+    name = "adsorption site"
+    def __init__(self):
+        super().__init__(self)
+        site = ['ontop', 'hollow','fcc', 'hcp', 'bridge', 'shortbridge', 'longbridge']
+        self.choices = tuple(site)
+
 
 
 Cell = CellType()
 FrameRange = FrameRangeType()
 Molecule = MoleculeType()
+AdsSite = AdsSiteType()
+Structure = StructureType()

src/mdkits/util/cp2k_input_parsing.py

@@ -5,34 +5,35 @@
 
 
 import sys
+from mdkits.util import os_operation, out_err
 
 
-def parse_cell(cp2k_input_file):
+def parse_cell():
     """
     function: parse cell information from cp2k input file
     parameter:
         cp2k_input_file: filename of cp2k input
     return:
         cell: list with 6 number
     """
-    try:
-        with open(cp2k_input_file, 'r') as f:
-            cell = []
-            for line in f:
-                if "ABC" in line:
-                    xyz = line.split()[-3:]
-                    cell.extend(xyz)
-                if "ALPHA_BETA_GAMMA" in line:
-                    abc = line.split()[-3:]
-                    cell.extend(abc)
-            if len(cell) == 3:
-                cell.extend([90.0, 90.0, 90.0])
+    for file in os_operation.default_input:
+        try:
+            with open(file, 'r') as f:
+                cell = []
+                for line in f:
+                    if "ABC" in line:
+                        xyz = line.split()[-3:]
+                        cell.extend(xyz)
+                    if "ALPHA_BETA_GAMMA" in line:
+                        abc = line.split()[-3:]
+                        cell.extend(abc)
+                if len(cell) == 3:
+                    cell.extend([90.0, 90.0, 90.0])
 
-            print(f"system cell: x = {cell[0]}, y = {cell[1]}, z = {cell[2]}, a = {cell[3]}\u00B0, b = {cell[4]}\u00B0, c = {cell[5]}\u00B0")
-            return cell
-    except FileNotFoundError:
-        print(f'cp2k input file name "{cp2k_input_file}" is not found, assign a cp2k input file or assign a "cell"')
-        sys.exit(1)
+                out_err.cell_output(cell)
+                return cell
+        except FileNotFoundError:
+            sys.exit(f"cant parse cell information from {','.join(os_operation.default_input)}, assign a cell")
 
 
 #def get_cell(cp2k_input_file, cell=None):

src/mdkits/util/os_operation.py

@@ -33,3 +33,8 @@ def sort_word_and_number(unsort_list):
     sorted_list = sorted(unsort_list, key=fns)
 
     return sorted_list
+
+
+def default_input():
+    default_input_name = os.environ.get("DEFAULT_INPUT", "input.inp,setup.inp,cell.inc").split(',')
+    return default_input_name

src/mdkits/util/out_err.py

@@ -0,0 +1,8 @@
+"""
+output and error for cli
+"""
+
+
+def cell_output(cell: list):
+    print(f"system cell: x = {cell[0]}, y = {cell[1]}, z = {cell[2]}, a = {cell[3]}\u00B0, b = {cell[4]}\u00B0, c = {cell[5]}\u00B0")
+