mdkits 0.1a1__py3-none-any.whl

mdkits/cli/matplot.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+
+import numpy as np
+import matplotlib.pyplot as plt
+import argparse
+import os
+from matplotlib import use as muse
+muse('Agg')
+
+
+def parse_slice(s):
+    if s == None:
+        return None
+    return [int(x) for x in s.replace(':', ' ').split()]
+
+
+def parse_argument():
+    parser = argparse.ArgumentParser(description='extract pos file from output file')
+
+    parser.add_argument('input_file_name', type=str, help='input file name')
+    parser.add_argument('-o', type=str, help='output file name, default is "out.xyz"', default='out.png')
+    parser.add_argument('-u', type=parse_slice, help='which clume to use, look like gnuplot u 1:2', default=[0, 1])
+    parser.add_argument('--label', type=str, help='line name')
+    parser.add_argument('--dpi', type=int, help='image dpi, default is 300', default=300)
+    parser.add_argument('--err',  help='plot err line', action='store_true')
+
+    return parser.parse_args()
+
+
+def plot_data(x, y, dpi, label, name):
+    fig, ax = plt.subplots(dpi=dpi)
+    ax.plot(x, y, label=label)
+    ax.legend()
+    fig.savefig(name)
+
+
+def plot_err(x, y, dpi, label, name):
+    fig, ax = plt.subplots(dpi=dpi)
+    ax.scatter(x, y, label=label)
+    line = np.linspace(np.amax(x), np.amin(x), 1000)
+    ax.plot(line, line)
+    ax.legend()
+    fig.savefig(name)
+
+
+def main():
+    args = parse_argument()
+    if args.label == None:
+        args.label = args.input_file_name
+    data = np.loadtxt(args.input_file_name)
+    if args.err:
+        plot_err(data[:, 0], data[:, 1], dpi=args.dpi, label=args.label, name=args.o)
+    else:
+        plot_data(data[:, 0], data[:, 1], dpi=args.dpi, label=args.label, name=args.o)
+
+    print(os.path.abspath(args.o))
+
+
+if __name__ == '__main__':
+    main()

mdkits/cli/packmol_input.py

@@ -0,0 +1,76 @@
+import argparse
+
+
+def parse_cell(s):
+    return [float(x) for x in s.replace(',', ' ').split()]
+
+
+def parse_argument():
+    parser = argparse.ArgumentParser(description='generate packmol input file with give parameter')
+    parser.add_argument('--size', type=int, help='water size default is 30', default=30)
+    parser.add_argument('--cell', type=parse_cell, help='input box size(a,b,c)')
+    parser.add_argument('--addwat', type=int, help='add some additional water, default is 0', default=0)
+    parser.add_argument('--ioncon', type=float, help='concentration of sol box, default is 0.0', default=0.0)
+    parser.add_argument('--tolerance', type=float, help='tolerance of packmol, default is 2.5', default=2.5)
+    parser.add_argument('--watpath', type=str, help='water xyz file path', default='C:\\home\\.can\\temp\\packmol\\default\\water.xyz')
+    parser.add_argument('--ionpath', type=str, help='ion xyz file path')
+    parser.add_argument('-o', type=str, help='output file name, default is "input.pm"', default='input.pm')
+    parser.add_argument('--output', type=str, help='output file name of packmol, default is "solbox.xyz"', default='solbox.xyz')
+
+    return parser.parse_args()
+
+
+def get_water_number():
+    water_number = water_volume / water_size
+
+    return int(round(water_number, 0))
+
+
+def get_ion_number(concentration):
+    ion_number = ( (concentration * avogadro) / 1e+27 ) * water_volume
+
+    return int(round(ion_number, 0))
+
+
+def main():
+    global water_volume, water_size, avogadro
+    args = parse_argument()
+    water_volume = args.cell[0] * args.cell[1] * args.cell[2]
+    water_size = args.size
+    avogadro = 6.02214179e+23
+    water_number = get_water_number() + args.addwat
+    ion_number = get_ion_number(args.ioncon)
+
+    if ion_number == 0:
+        packmol_input_str = f"""
+        tolerance {args.tolerance}
+        filetype xyz
+        output {args.output}
+        pbc {args.cell[3]} {args.cell[4]} {args.cell[5]}
+        structure {args.watpath}
+          number {water_number}
+          inside box 2. 2. 2. {args.cell[0]-2} {args.cell[1]-2} {args.cell[2]-2}
+        end structure
+        """
+    else:
+        packmol_input_str = f"""
+        tolerance {args.tolerance}
+        filetype xyz
+        output {args.output}
+        pbc {args.cell[3]} {args.cell[4]} {args.cell[5]}
+        structure {args.watpath}
+          number {water_number}
+          inside box 2. 2. 2. {args.cell[0]-2} {args.cell[1]-2} {args.cell[2]-2}
+        end structure
+        structure {args.ionpath}
+          number {ion_number}
+          inside box 2. 2. 2. {args.cell[0]-2} {args.cell[1]-2} {args.cell[2]-2}
+        end structure
+        """
+
+    with open(args.o, 'w') as f:
+        f.write(packmol_input_str)
+
+
+if __name__ == '__main__':
+    main()

mdkits/cli/pdos.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3.9
+
+# caculate pdos
+
+from cp2kdata import Cp2kPdos
+import argparse
+import numpy as np
+from util import os_operation
+
+
+def to_file(filename, ener, dos):
+	to_file = np.column_stack((ener, dos))
+	np.savetxt(filename, to_file, delimiter=" ")
+
+
+# set argument
+parser = argparse.ArgumentParser(description='calculate pdos')
+parser.add_argument('filename', type=str, nargs='?', default=os_operation.default_file_name('*-k*.pdos'))
+parser.add_argument('--type', type=str, default='total')
+args = parser.parse_args()
+
+
+# pdos
+#dos_types = ['0']
+#for dos_type in dos_types:
+#	for atom_type in range(1, args.type_number+1):
+#		dos_obj = Cp2kPdos(f'{args.project_name}-k{atom_type}-1_{args.run_step}.pdos')
+#		dos, ener = dos_obj.get_raw_dos()
+#		to_file(f"{dos_obj.read_dos_element()}_{dos_type}.pdos", ener, dos)
+
+# total dos
+print(args.filename)
+for file in args.filename:
+    dos_obj = Cp2kPdos(file)
+    dos, ener = dos_obj.get_raw_dos(dos_type=args.type)
+    to_file(f'{dos_obj.read_dos_element()}_{args.type}.pdos', ener, dos)

mdkits/cli/plot.py

@@ -0,0 +1,289 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+
+import numpy as np
+from matplotlib import pyplot as plt
+from matplotlib import rcParams
+import matplotlib.ticker as ticker
+import yaml, click
+from mdtool.util import fig_operation
+
+DEFAULT_CONFIG = {
+        'data': {},
+        'x': {},
+        'y': {},
+        'figsize': None,
+        'x_range': None,
+        'y_range': None,
+        'legend_fontsize': None,
+        'ticks_fontsize': None,
+        'line': {},
+        'fold': None
+    }
+
+
+def load_config(file_path):
+    try:
+        with open(file_path, 'r') as file:
+            data = yaml.safe_load(file)
+    except FileNotFoundError:
+        print(f"no file {file_path}")
+        exit(1)
+    except yaml.YAMLError as exc:
+        print(f"yaml parse error: {exc}")
+        exit(1)
+    config = data.copy()
+    #for key, default in DEFAULT_CONFIG.items():
+    #    config[key] = data.get(key, default)
+    for key, value in config.items():
+        for dkey, dvalue in DEFAULT_CONFIG.items():
+            value[dkey] = data[key].get(dkey, dvalue)
+
+    return config
+
+
+def check_figsize(figsize):
+    match figsize:
+        case list() if len(figsize) == 2:
+            rcParams['figure.figsize'] = figsize
+        case None:
+            pass
+        case _:
+            print("wrong yaml structure")
+            exit(1)
+
+
+def check_int_ticks(ax, int_ticks):
+    match int_ticks:
+        case 'x':
+            ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))
+        case 'y':
+            ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True))
+        case 'a':
+            ax.xaxis.set_major_locator(ticker.MaxNLocator(integer=True))
+            ax.yaxis.set_major_locator(ticker.MaxNLocator(integer=True))
+        case _:
+            pass
+
+
+def check_axis_range(ax, x_range=None, y_range=None):
+    if x_range:
+        match x_range:
+            case list() if len(x_range) == 2:
+                ax.set_xlim(x_range[0], x_range[1])
+            case None:
+                pass
+            case _:
+                print("wrong yaml structure")
+                exit(1)
+    if y_range:
+        match y_range:
+            case list() if len(y_range) == 2:
+                ax.set_ylim(y_range[0], y_range[1])
+            case None:
+                pass
+            case _:
+                print("wrong yaml structure")
+                exit(1)
+
+
+def check_lenend_fontsize(ax, legend_fontsize, rmse=None):
+    if rmse is not None:
+        match legend_fontsize:
+            case int():
+                l = ax.legend(fontsize=legend_fontsize, markerscale=0, handlelength=0)
+            case None:
+                l = ax.legend(markerscale=0, handlelength=0)
+            case _:
+                pass
+
+        #p = l.get_window_extent()
+        #ax.annotate(f"RMSE:{rmse:.4f}", (p.p0[0], p.p1[1]), (p.p0[0], p.p1[1]-300), xycoords='figure pixels', fontsize=legend_fontsize)
+    else:
+        match legend_fontsize:
+            case int():
+                ax.legend(fontsize=legend_fontsize)
+            case None:
+                ax.legend()
+            case _:
+                pass
+
+
+def check_output_mode(fig, plt, m, figure, fold=None):
+    match m:
+        case 'show':
+            plt.show()
+        case 'save':
+            fig_operation.savefig(fig, f"{figure}", fold)
+            plt.close()
+        case 'all':
+            fig_operation.savefig(fig, f"{figure}", fold)
+            plt.show()
+            plt.close()
+        case _:
+            pass
+
+
+def check_ticks_fontsize(ax, ticks_fontsize):
+    match ticks_fontsize:
+        case int():
+            ax.tick_params(axis='both', labelsize=ticks_fontsize)
+        case None:
+            pass
+        case _:
+            pass
+
+
+def mode_1(data_dict, figure, **kwargs):
+    # init config
+    check_figsize(data_dict['figsize'])
+
+    fig, ax = plt.subplots()
+
+    # ticks config
+    check_int_ticks(ax, kwargs['int_ticks'])
+
+    # plot part
+    for label, data in data_dict['data'].items():
+        label = r"$\mathrm{ "+ label +" }$"
+        data = np.loadtxt(data)
+        x_slice, xlabel = next(iter(data_dict['x'].items()))
+        for y_slice, ylabel in data_dict['y'].items():
+            ax.plot(data[:, x_slice], data[:, y_slice], label=label)
+
+        # range config
+        check_axis_range(ax, x_range=data_dict['x_range'], y_range=data_dict['y_range'])
+
+        # axis label part
+        ax.set_xlabel(r"$\mathrm{ "+ xlabel +" }$")
+        ax.set_ylabel(r"$\mathrm{ "+ ylabel +" }$")
+
+        # h or v line
+        match data_dict['line']:
+            case _:
+                pass
+
+        # legend config
+        check_lenend_fontsize(ax, data_dict['legend_fontsize'])
+
+    # output part
+    check_output_mode(fig, plt, kwargs['m'], figure, data_dict['fold'])
+
+
+def mode_2(data_dict, figure, **kwargs):
+    # init config
+    check_figsize(data_dict['figsize'])
+
+    fig, ax = plt.subplots()
+
+    # ticks config
+    check_int_ticks(ax, kwargs['int_ticks'])
+
+    # plot part
+    ylabel, data = next(iter(data_dict['data'].items()))
+    data = np.loadtxt(data)
+    x_slice, xlabel = next(iter(data_dict['x'].items()))
+    for y_slice, label in data_dict['y'].items():
+        label = r"$\mathrm{ "+ label +" }$"
+        ax.plot(data[:, x_slice], data[:, y_slice], label=label)
+
+        # range config
+        check_axis_range(ax, x_range=data_dict['x_range'], y_range=data_dict['y_range'])
+
+        # axis label part
+        ax.set_xlabel(r"$\mathrm{ "+ xlabel +" }$")
+        ax.set_ylabel(r"$\mathrm{ "+ ylabel +" }$")
+
+    # legend config
+    check_lenend_fontsize(ax, data_dict['legend_fontsize'])
+
+    # output part
+    check_output_mode(fig, plt, kwargs['m'], figure, data_dict['fold'])
+
+
+def mode_error(data_dict, figure, **kwargs):
+    # init config
+    check_figsize(data_dict['figsize'])
+
+    fig, ax = plt.subplots(figsize=(4.8, 4.8))
+
+    # ticks config
+    check_int_ticks(ax, kwargs['int_ticks'])
+
+    # plot part
+    label, data = next(iter(data_dict['data'].items()))
+    data = np.loadtxt(data)
+    x_slice, xlabel = next(iter(data_dict['x'].items()))
+    for y_slice, ylabel in data_dict['y'].items():
+        label = r"$\mathrm{ "+ label +" }$"
+        ax.xaxis.set_major_formatter(ticker.FormatStrFormatter('%.2f'))
+        ax.yaxis.set_major_formatter(ticker.FormatStrFormatter('%.2f'))
+        ax.scatter(data[:, x_slice], data[:, y_slice], label=label)
+        data_diff = data[:, x_slice] - data[:, y_slice]
+        data_rmse = np.sqrt(np.average(data_diff*data_diff))
+        ax.axline((data[0, x_slice], data[0, x_slice]), slope=1, ls='--', c='black')
+        ax.set_xticks([])
+        ax.set_yticks([])
+        check_lenend_fontsize(ax, data_dict['ticks_fontsize'])
+
+        # range config
+        check_axis_range(ax, x_range=data_dict['x_range'], y_range=data_dict['y_range'])
+
+        # axis label part
+        ax.set_xlabel(r"$\mathrm{ "+ xlabel +" }$")
+        ax.set_ylabel(r"$\mathrm{ "+ ylabel +" }$")
+
+        lims = [np.min([ax.get_xlim(), ax.get_ylim()]), np.max([ax.get_xlim(), ax.get_ylim()])]
+        ax.text(lims[0]+(lims[1]-lims[0])*.6, lims[0]+(lims[1]-lims[0])*.1, f"RMSE: {data_rmse:.4f}", fontsize=data_dict['legend_fontsize'])
+        ax.set_xlim(lims)
+        ax.set_ylim(lims)
+        #ax.set_aspect('equal')
+    #plt.axis('equal')
+    # legend config
+    check_lenend_fontsize(ax, data_dict['legend_fontsize'], rmse=data_rmse)
+
+    # output part
+    check_output_mode(fig, plt, kwargs['m'], figure, data_dict['fold'])
+
+
+@click.command(name='plot')
+@click.argument('yaml_file', type=click.Path(exists=True), nargs=-1)
+@click.option('--int_ticks', help='set x ticks or y ticks with int, choices: x, y, a(all)', type=click.Choice(['x', 'y', 'a']))
+@click.option('--error', help='set error mode', is_flag=True)
+@click.option('-m', help='output mode: show, save, all, default is save', type=click.Choice(['show', 'save', 'all']), default='save')
+def main(yaml_file, int_ticks, error, m):
+    kwargs = locals()
+    figure_dict = {}
+    for yaml_file in yaml_file:
+        yaml_dict = load_config(yaml_file)
+        figure_dict.update(yaml_dict)
+
+        for figure, data_dict in figure_dict.items():
+            check_data = len(data_dict['data'])
+
+            # mode error
+            if error:
+                mode = 'error'
+                mode_error(data_dict, figure, **kwargs)
+
+            # mode 1
+            elif check_data > 1:
+                if len(data_dict['y']) > 1:
+                    print("wrong yaml structure")
+                    exit(1)
+                mode = 1
+                mode_1(data_dict, figure, **kwargs)
+
+            # mode 2
+            elif check_data == 1:
+                mode = 2
+                mode_2(data_dict, figure, **kwargs)
+
+            else:
+                print("wrong yaml structure")
+                exit(1)
+            print(f"{figure} is done with mode {mode}")
+
+
+if __name__ == '__main__':
+    main()

mdkits/cli/supercell.py

@@ -0,0 +1,72 @@
+#!/usr/bin/env python3
+
+from ase.io import read
+import argparse, os
+import numpy as np
+from ase.build import make_supercell
+from util import (
+    structure_parsing,
+    encapsulated_ase,
+    cp2k_input_parsing
+    )
+
+
+def supercell(atom, x, y, z):
+    P = [ [x, 0, 0], [0, y, 0], [0, 0, z] ]
+    super_atom = make_supercell(atom, P)
+    return super_atom
+
+
+def parse_cell(s):
+    return [float(x) for x in s.replace(',', ' ').split()]
+
+
+def super_cell(s):
+    super_cell = [int(x) for x in s.replace(',', ' ').split()]
+    leng = len(super_cell)
+    if leng == 2:
+        super_cell.append(1)
+    elif leng == 3:
+        pass
+    else:
+        print('wrong super cell size')
+        exit(1)
+
+    return super_cell
+
+
+def parse_argument():
+    parser = argparse.ArgumentParser(description='make a supercell')
+
+    parser.add_argument('input_file_name', type=str, help='input file name')
+    parser.add_argument('super', type=super_cell, help='super cell size, a,b,c')
+    parser.add_argument('-o', type=str, help='output file name, default is "super.cif"', default='super.cif')
+    parser.add_argument('--cp2k_input_file', type=str, help='input file name of cp2k, default is "input.inp"', default='input.inp')
+    parser.add_argument('--coord', help='coord format', action='store_true')
+    parser.add_argument('--cell', type=parse_cell, help='set cell, a list of lattice, --cell x,y,z or x,y,z,a,b,c')
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_argument()
+    if args.input_file_name == None:
+        print('give a xyz file')
+        sys.exit()
+
+    atom = encapsulated_ase.atoms_read_with_cell(args.input_file_name, cell=args.cell, coord_mode=args.coord)
+    atom.set_pbc(True)
+    atom.wrap()
+    super_atom = supercell(atom, args.super[0], args.super[1], args.super[2])
+
+    suffix = args.o.split('.')[-1]
+    if suffix == 'data':
+        super_atom.write(f'{args.o}', format='lammps-data', atom_style='atomic')
+    else:
+        super_atom.write(f'{args.o}')
+
+    print(os.path.abspath(args.o))
+
+
+if __name__ == '__main__':
+    main()

mdkits/cli/wrap.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+
+import os
+from MDAnalysis import Universe
+import MDAnalysis, click
+from mdtool.util import (
+    arg_type,
+    os_operation,
+    cp2k_input_parsing
+    )
+
+
+@click.command(name='wrap')
+@click.argument('filename', type=click.Path(exists=True), default=os_operation.default_file_name('*-pos-1.xyz', last=True))
+@click.option('-o', type=str, help='output file name', default='wraped.xyz', show_default=True)
+@click.option('--cell', type=arg_type.Cell, help='set cell from cp2k input file or a list of lattice: --cell x,y,z or x,y,z,a,b,c', default='input.inp', show_default=True)
+def main(filename, o, cell):
+    """
+    wrap the coordinates in a cell from a trajectory file
+    """
+
+    u = Universe(filename)
+    u.dimensions = cell
+    ag = u.select_atoms("all")
+
+    with MDAnalysis.Writer(o, ag.n_atoms) as W:
+        for ts in u.trajectory:
+            ag.wrap()
+            W.write(ag)
+
+    click.echo("\nwrap is done, output file {o} is:")
+    click.echo(os.path.abspath(o))
+
+
+if __name__ == '__main__':
+    main()

mdkits/config/__init__.py

@@ -0,0 +1,35 @@
+"""
+Configuration center.
+Use https://www.dynaconf.com/
+"""""
+import os
+import sys
+from pathlib import Path
+
+from dynaconf import Dynaconf
+
+
+_base_dir = Path(__file__).parent.parent
+
+_settings_files = [
+    # All config file will merge.
+    Path(__file__).parent / 'settings.yml',  # Load default config.
+]
+
+# User configuration. It will be created automatically by the pip installer .
+_external_files = [
+    Path(sys.prefix, 'etc', 'mdtool', 'settings.yml')
+]
+
+
+settings = Dynaconf(
+    # Set env `MDTOOL_FOO='bar'`，use `settings.FOO` .
+    envvar_prefix='MDTOOL',
+    settings_files=_settings_files,  # load user configuration.
+    # environments=True,  # Enable multi-level configuration，eg: default, development, production
+    load_dotenv=True,  # Enable load .env
+    # env_switcher='MDTOOL_ENV',
+    lowercase_read=False,  # If true, can't use `settings.foo`, but can only use `settings.FOO`
+    includes=_external_files,  # Customs settings.
+    base_dir=_base_dir,  # `settings.BASE_DIR`
+)

mdkits/config/settings.yml

@@ -0,0 +1,4 @@
+verbose: false
+debug: false
+loglevel: warning
+logpath: /tmp/mdtool

mdkits/mdtool.py

@@ -0,0 +1,28 @@
+import click
+from mdtool.cli import (
+    convert,
+    wrap,
+    extract,
+    data,
+    plot,
+)
+
+
+@click.group(name='cli')
+@click.pass_context
+@click.version_option()
+def cli(ctx):
+    print(ctx)
+    """tools for md or dft"""
+    pass
+
+
+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)
+
+
+if __name__ == '__main__':
+    cli()

mdkits/util/arg_type.py

@@ -0,0 +1,41 @@
+import click, os
+from . import cp2k_input_parsing
+
+
+class CellType(click.ParamType):
+    name = "pbc cell type"
+
+    def convert(self, value, param, ctx):
+        if isinstance(value, str):
+            if ',' not in value:
+                cell = cp2k_input_parsing.parse_cell(value)
+                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]
+                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")
+                    return cell
+                else:
+                    self.fail(f"{value} is not a valid cell parameter", param, ctx)
+
+
+class FrameRangeType(click.ParamType):
+    name = "frame range"
+    def convert(self, value, param, ctx):
+        if isinstance(value, str):
+            parts = value.split(':')
+
+            range_list = [int(x) if x else None for x in parts]
+
+            if len(range_list) > 0 and len(range_list) <= 3:
+                return range_list
+            else:
+                self.fail(f"{value} is not a valid frame range", param, ctx)
+
+
+Cell = CellType()
+FrameRange = FrameRangeType()