mdkits 0.1.13__py3-none-any.whl → 1.2.3__py3-none-any.whl

mdkits/md_cli/vac.py

@@ -0,0 +1,65 @@
+import numpy as np
+import click
+import MDAnalysis
+from MDAnalysis import Universe
+from MDAnalysis.analysis.base import AnalysisBase
+from mdkits.util import os_operation, arg_type
+
+
+class Velocity_AutoCorrelation(AnalysisBase):
+    def __init__(self, filename, select, dt=0.001):
+        u = Universe(filename)
+        u.trajectory.ts.dt = dt
+        self.u = u
+        self.atomgroup = u.select_atoms(select)
+
+        super(Velocity_AutoCorrelation, self).__init__(self.atomgroup.universe.trajectory, verbose=True)
+
+    def _prepare(self):
+        self.cvv = []
+        self.v0 = self.atomgroup.positions
+        self.normalize = 1/np.sum(self.v0*self.v0)
+        self.cvv.append(np.sum(self.v0*self.v0)*self.normalize)
+
+
+    def _append(self, cvv):
+        self.cvv.append(cvv*self.normalize)
+
+    def _single_frame(self):
+        cvv = np.sum(self.atomgroup.positions*self.v0)
+        self._append(cvv)
+
+    def _conclude(self):
+        self.cvv = np.array(self.cvv)
+
+        sf = self.cvv.shape[0]
+        fftraj = np.fft.rfft(self.cvv)
+        fdos = np.abs(fftraj)
+
+        faxis = np.fft.rfftfreq(sf, d=1/sf)
+
+        combine = np.column_stack((np.arange(len(self.cvv)), self.cvv))
+
+        np.savetxt('vac.dat', combine, fmt='%.5f', header="frame\tvac")
+        np.savetxt('freq.dat', np.column_stack((faxis, fdos)), fmt='%.5f', header="freq\tabundance")
+
+
+@click.command(name="vac")
+@click.argument("filename", type=click.Path(exists=True), default=os_operation.default_file_name('*-vel-1.xyz', last=True))
+@click.option("--select", type=str, default="all", help="atom selection", show_default=True)
+@click.option('-r', type=arg_type.FrameRange, help='range of frame to analysis')
+def main(filename, select, r):
+    """analysis velocity autocorrelation function and frequency"""
+    a = Velocity_AutoCorrelation(filename, select)
+
+    if r is not None:
+        if len(r) == 2:
+            a.run(start=r[0], stop=r[1])
+        elif len(r) == 3:
+            a.run(start=r[0], stop=r[1], step=r[2])
+    else:
+        a.run()
+
+
+if __name__ == '__main__':
+    main()

mdkits/{cli → md_cli}/wrap.py

@@ -6,7 +6,8 @@ import MDAnalysis, click
 from mdkits.util import (
     arg_type,
     os_operation,
-    cp2k_input_parsing
+    cp2k_input_parsing,
+    out_err
     )
 
 
@@ -28,8 +29,8 @@ def main(filename, o, cell):
             ag.wrap()
             W.write(ag)
 
-    click.echo("\nwrap is done, output file {o} is:")
-    click.echo(os.path.abspath(o))
+    click.echo(f"\nwrap is done, output file {o} is:")
+    out_err.path_output(o)
 
 
 if __name__ == '__main__':

mdkits/mdkits.py

@@ -1,14 +1,12 @@
 import click
 from mdkits.build_cli import build_cli
+from mdkits.dft_cli import dft_cli
+from mdkits.md_cli import md_cli
 from mdkits.cli import (
     convert,
-    wrap,
     extract,
     data,
     plot,
-    density,
-    cube,
-    pdos,
 )
 
 
@@ -20,15 +18,13 @@ def cli(ctx):
     pass
 
 
+cli.add_command(md_cli.cli)
+cli.add_command(build_cli.cli_build)
+cli.add_command(dft_cli.main)
 cli.add_command(convert.main)
-cli.add_command(wrap.main)
 cli.add_command(extract.main)
 cli.add_command(data.main)
 cli.add_command(plot.main)
-cli.add_command(density.main)
-cli.add_command(cube.main)
-cli.add_command(pdos.main)
-cli.add_command(build_cli.cli_build)
 
 
 if __name__ == '__main__':

mdkits/util/arg_type.py

@@ -6,7 +6,7 @@ from mdkits.util import os_operation, cp2k_input_parsing, out_err
 
 
 class CellType(click.ParamType):
-    name = "pbc cell type"
+    name = "cell type"
 
     def convert(self, value, param, ctx):
         if isinstance(value, str):
@@ -14,10 +14,8 @@ class CellType(click.ParamType):
 
             if len(cell) == 3:
                 cell += [90, 90, 90]
-                out_err.cell_output(cell)
                 return cell
             elif len(cell) == 6:
-                out_err.cell_output(cell)
                 return cell
             else:
                 self.fail(f"{value} is not a valid cell parameter", param, ctx)
@@ -38,7 +36,7 @@ class FrameRangeType(click.ParamType):
 
 
 class StructureType(click.ParamType):
-    name = "structure file type"
+    name = "structure type"
     def convert(self, value, param, ctx):
         no_cell=np.array([0., 0., 0., 90., 90., 90.])
         if isinstance(value, str):
@@ -58,19 +56,30 @@ class StructureType(click.ParamType):
                 self.fail(f"{value} is not exists", param, ctx)
 
 
-
 class MoleculeType(click.Choice):
     name = "mocular type"
     def __init__(self):
         g2.names.append(click.Path(exists=True))
         super().__init__(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)
+        super().__init__(choices=tuple(site))
+
+
+class FileListType(click.ParamType):
+    name = "file list"
+    def convert(self, value, param, ctx):
+        if isinstance(value, str):
+            import glob
+            file_list = glob.glob(value)
+            if file_list:
+                return file_list
+            else:
+                self.fail(f"No files match {value}", param, ctx)
 
 
 
@@ -78,4 +87,5 @@ Cell = CellType()
 FrameRange = FrameRangeType()
 Molecule = MoleculeType()
 AdsSite = AdsSiteType()
-Structure = StructureType()
+Structure = StructureType()
+FileList = FileListType()

mdkits/util/cp2k_input_parsing.py

@@ -30,7 +30,11 @@ def parse_cell():
                 if len(cell) == 3:
                     cell.extend([90.0, 90.0, 90.0])
 
-                out_err.cell_output(cell)
+                if len(cell) == 0:
+                    print("parse failed")
+                    return [0., 0., 0., 90., 90., 90.]
+                else:
+                    print(f"parsed 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:
             return [0., 0., 0., 90., 90., 90.]

mdkits/util/encapsulated_ase.py

@@ -66,24 +66,45 @@ def rdf(chunk, cell, bin_size, name, parallel=True):
     ana.clear_cache()
 
 
-def ave_cube(filepath):
+def ave_cube(filepath, axis):
     """
     function: average hartree file in z_coordinate
     parameter:
         filepath: hartree cube file path
+        axis: average axis, can be 'x','y','z' or 0,1,2
     return:
         z_cube_data: a list of cube data alone z axes
         z_coordinates: a list of coordinates of z axes
     """
+    if isinstance(axis, str):
+        mapping = {'x': 0, 'y': 1, 'z': 2}
+        key = axis.lower()
+        if key not in mapping:
+            raise ValueError(f"axis string must be one of 'x','y','z', got {axis}")
+        axis_idx = mapping[key]
+    else:
+        axis_idx = int(axis)
+    
     # read data from filepath
     data, atoms = read_cube_data(filepath)
-    # define need parameter
-    npoints = data.shape[2]
-    step_size = atoms.cell[2, 2] / ( npoints - 1 )
-    # average hartree file, and calculate z_coordinates
+
+    ndim = data.ndim
+    if axis_idx < 0:
+        axis_idx = ndim + axis_idx
+    if not (0 <= axis_idx < ndim):
+        raise ValueError(f"axis must be between 0 and {ndim-1}, got {axis_idx}")
+
+    npoints = data.shape[axis_idx]
+    step_size = atoms.cell.cellpar()[axis_idx] / ( npoints - 1 )
+
     z_coordinates = [i * step_size for i in range(npoints)]
-    z_cube_data = data[:, :, :].sum(axis=(0, 1)) / ( data.shape[0] * data.shape[1] )
-    return np.column_stack((z_coordinates, z_cube_data))
+
+    other_axes = tuple(i for i in range(ndim) if i != axis_idx)
+    if other_axes:
+        averaged = data.mean(axis=other_axes)
+    else:
+        averaged = data.copy()
+    return np.column_stack((z_coordinates, averaged))
 
 
 def atoms_read_with_cell(file_name, cell=None, coord_mode=False, default_cell=np.array([0., 0., 0., 90., 90., 90.])):

mdkits/util/encapsulated_mda.py

@@ -3,7 +3,7 @@ import numpy as np
 from . import numpy_geo
 
 
-def update_water(self, o_group, h_group, distance_judg=1.2, angle_judg:tuple[float, float]=None, return_index=False):
+def update_water(self, o_group, h_group, distance_judg=1.2, angle_judg:tuple[float, float]=(None, None), return_index=False):
     """
     input: o and h atom
     output: o and two h in this frame
@@ -53,6 +53,9 @@ def update_water(self, o_group, h_group, distance_judg=1.2, angle_judg:tuple[flo
     if return_index:
         return o_index, oh1_index, oh2_index
     else:
+        if len(o_index) == 0:
+            raise ValueError("No water found in this atom group")
+
         o = o_group[o_index]
         oh1 = h_group[oh1_index]
         oh2 = h_group[oh2_index]

mdkits/util/numpy_geo.py

@@ -45,7 +45,10 @@ def vector_between_two_vector(vector1, vector2):
 
 
 def vector_vector_angle(vector, surface_vector):
-    cos = np.dot(vector, surface_vector) / (np.linalg.norm(vector) * np.linalg.norm(surface_vector))
+    if len(vector.shape) == 1:
+        cos = np.dot(vector, surface_vector) / (np.linalg.norm(vector) * np.linalg.norm(surface_vector))
+    else:
+        cos = np.dot(vector, surface_vector) / (np.linalg.norm(vector, axis=1) * np.linalg.norm(surface_vector))
     vector_vector_angle = np.arccos(np.clip(cos, -1.0, 1.0))
     vector_vector_angle = np.degrees(vector_vector_angle)
     return vector_vector_angle
@@ -96,7 +99,7 @@ def unwrap(atom1, atom2, coefficients, max=0, total=False):
         return min_dist, closest_point
 
 
-def find_surface(surface_group:np.ndarray, layer_tolerance=1, surface_tolerance=5):
+def find_surface(surface_group:np.ndarray, layer_tolerance=0.05, surface_tolerance=5):
     sort_group = np.sort(surface_group)
     layer_mean = []
     current_layer = [sort_group[0]]
@@ -108,12 +111,14 @@ def find_surface(surface_group:np.ndarray, layer_tolerance=1, surface_tolerance=
             current_layer = [sort_group[i]]
         layer_mean.append(np.mean(current_layer))
     
-    if len(current_layer) == 1:
-        return layer_mean[0]
+    if len(layer_mean) == 1:
+        return [layer_mean[0], 0]
 
     diff = np.diff(layer_mean)
     if np.any(diff > surface_tolerance):
-        index = np.argmax(diff > 5)
+        index = np.argmax(diff > surface_tolerance)
         return (layer_mean[index], layer_mean[index + 1])
     else:
+        if layer_mean[-1] > layer_mean[0]:
+            return [layer_mean[-1], 0]
         return (layer_mean[-1], layer_mean[0])

mdkits/util/os_operation.py

@@ -13,7 +13,7 @@ def remove_temp_dir(temp_dir):
     shutil.rmtree(temp_dir)
 
 
-def default_file_name(match, last=False):
+def default_file_name(match, last=False, space_split=False):
     import glob
     file_list = glob.glob(match)
     if file_list:
@@ -21,6 +21,8 @@ def default_file_name(match, last=False):
             default_file_name = sort_word_and_number(file_list)[-1]
         else:
             default_file_name = list(file_list)
+            if space_split:
+                default_file_name = ' '.join(default_file_name)
     else:
         default_file_name = None
 

mdkits/util/out_err.py

@@ -6,17 +6,29 @@ import numpy as np
 import sys, os
 
 
-def cell_output(cell):
-    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")
+def cell_output(atoms):
+    cell = atoms.cell.cellpar()
+    if not hasattr(atoms, "name"):
+        atoms.name = "present"
+    print(f"{atoms.name} cell: x = {cell[0]}, y = {cell[1]}, z = {cell[2]}, a = {cell[3]}\u00B0, b = {cell[4]}\u00B0, c = {cell[5]}\u00B0")
 
 
 def path_output(file: str):
-    print(os.path.abspath(file))
+    env_var_name = 'ssh_name'
+    file_path = os.path.abspath(file)
+    if os.environ.get(env_var_name):
+        ssh_name = os.environ.get(env_var_name)
+        file_path = f"{ssh_name}:{file_path}"
+    print(file_path)
 
 def check_cell(atoms, cell=None):
-    if not np.array_equal(atoms.cell.cellpar(), np.array([0., 0., 0., 90., 90., 90.])):
-        cell_output(atoms.cell.cellpar())
+    if cell is not None:
+        atoms.set_cell(cell)
+        cell_output(atoms)
+    elif not np.array_equal(atoms.cell.cellpar(), np.array([0., 0., 0., 90., 90., 90.])):
+        cell_output(atoms)
     elif np.array_equal(atoms.cell.cellpar(), np.array([0., 0., 0., 90., 90., 90.])) and cell is not None:
         atoms.set_cell(cell)
+        cell_output(atoms)
     else:
-        raise ValueError("can't parse cell please use --cell set cell")
+        raise ValueError("can't parse cell please use --cell set cell")