#!/usr/bin/env python # ./dump2cube.py -i 050Li3PS4-050LiI.lammpstrj -m 160 160 160 -cut 8 import numpy as np import argparse import re import os import time description = """This is a test program""" par = argparse.ArgumentParser(description=description) par.add_argument('-i', '--trjfile', default="", required=True, help='input trjfile') par.add_argument('-m', '--mesh', default=(30, 30, 30), required=False, nargs=3, type=int, help='mesh grid') par.add_argument('-cut', '--cutoff', default=7, required=False, type=float, help='mesh grid') args = par.parse_args() ang_borr = 0.5291772109217 dirname = os.path.dirname(__file__) class LammpsTrj(): def __init__(self): """ trjファイルを読み込んで、原子数、lattice、ステップ数を格納 """ with open(args.trjfile) as o: data = o.read().split() self.atoms = int(data[7]) # 原子数 data = np.array(data).reshape(-1, self.atoms*7+32) self.lattices = data[:, 17:23].astype(float) # (step, 6) self.data = data[:, 32:] # 座標データ self.steps = data.shape[0] self.mesh = args.mesh self.cutoff = args.cutoff def setLattice(self, lat): """ latticeの形によって操作を分岐 self.M, self.M_を作成 """ if lat.shape[0] == 6: M = np.array([[lat[1]-lat[0], 0, 0], [0, lat[3]-lat[2], 0], [0, 0, lat[5]-lat[4]]]).astype(float) if lat.shape[0] == 9: xlo_bound, xhi_bound, xy = lat[0, 0], lat[0, 1], lat[0, 2] ylo_bound, yhi_bound, xz = lat[1, 0], lat[1, 1], lat[1, 2] zlo_bound, zhi_bound, yz = lat[2, 0], lat[2, 1], lat[2, 2] xlo = xlo_bound - np.min([0.0, xy, xz, xy+xz]) xhi = xhi_bound - np.max([0.0, xy, xz, xy+xz]) ylo = ylo_bound - np.min([0.0, yz]) yhi = yhi_bound - np.max([0.0, yz]) zlo = zlo_bound zhi = zhi_bound lx = xhi - xlo ly = yhi - ylo lz = zhi - zlo a = lx b = np.sqrt(ly**2 + xy**2) c = np.sqrt(lz**2 + xz**2 + yz**2) alpha = np.arccos((xy*xz + ly*yz)/b/c) beta = np.arccos(xz/c) gamma = np.arccos(xy/b) v1 = [a, 0, 0] v2 = [b*np.cos(gamma), b*np.sin(gamma), 0] v3 = [c*np.cos(beta), c*(np.cos(alpha)-np.cos(beta)*np.cos(gamma))/np.sin(gamma), c*np.sqrt(1+2*np.cos(alpha)*np.cos(beta)*np.cos(gamma) - np.cos(alpha)**2-np.cos(beta)**2 - np.cos(gamma)**2 / np.sin(gamma))] M = np.array([v1, v2, v3]) return M def getOnestep(self, step): """ 引数のstepにおける座標を、原子ごとにself.elems(辞書)に格納する self.elemsの座標は分率座標 """ step_data = self.data[step, :].reshape(self.atoms, -1) # (atoms, 6) self.M = self.setLattice(self.lattices[step]) self.M_ = np.linalg.inv(self.M) self.elems = {} for e in np.unique(step_data[:, 2]): d = step_data[step_data[:, 2] == e, 3:].astype(float) dxyz = d[:, 0:3] dxyz = dxyz @ self.M_ dxyz = dxyz - np.floor(dxyz) d[:, 0:3] = dxyz self.elems[e] = d # 元素ごとに結晶座標とmol番号を格納 return self.elems def makeCube(self, trjfile): """ すべてのP原子について処理を行う。 P原子ごとに最も近い4つのI原子を選出し、回転・アライメントを行う。 """ # Li用の座標リスト self.li_coords_list = [] # I用のリスト (全てまとめて格納) self.i_coords_list = [] # ヒストグラム結果格納用 self.li_hist = None self.i_hist = None count_p = 0 # 処理したP原子の総数カウント用 for step in range(self.steps): print(f"processing {step} step") self.elems = self.getOnestep(step) elems_ = {k: v.copy() for k, v in self.elems.items()} # copy # 全てのP原子に対してループ処理を行う for p_data in elems_["P"]: # 全てのI原子の中から、このP原子に最も近い4つのI原子を特定する i_diffs = elems_["I"][:, 0:3] - p_data[0:3] i_diffs = i_diffs - np.around(i_diffs) # 周期境界条件の考慮(分率) i_diffs_abs = i_diffs @ self.M # 絶対座標へ変換 distances = np.linalg.norm(i_diffs_abs, axis=1) sorted_indices = np.argsort(distances) closest_Is = elems_["I"][sorted_indices[:4]] # 近い順に4つ抽出 nearest_I_coords = closest_Is[0] # 回転定義に使用する「最も近いI」 # S原子の処理 (同じmol番号のSを取得) s_data = elems_["S"][elems_["S"][:, 3] == p_data[3]] # 1つのPS4に注目して、Pを原点に配置するための距離計算 si_diff = s_data[:, 0:3] - nearest_I_coords[0:3] si_diff = si_diff - np.around(si_diff) si_diff = np.linalg.norm(si_diff, axis=1) # 最も遠いSをz軸上に配置するため、そのインデックスを取得 n_S_idx = np.argmax(si_diff) # その他のインデックスを取得 s_indices = [0, 1, 2, 3] s_indices.remove(n_S_idx) n2_S_idx = s_indices[0] s_xyz = s_data[:, 0:3] - p_data[0:3] s_xyz = s_xyz - np.round(s_xyz) s_xyz = s_xyz @ self.M # 絶対座標へ # 1つのSをyz平面上に配置(z軸周り回転) theta = np.arctan2(s_xyz[n_S_idx][0], s_xyz[n_S_idx][1]) self.Mat_z = np.array([[np.cos(-theta), np.sin(-theta), 0], [-np.sin(-theta), np.cos(-theta), 0], [0, 0, 1]]) s_xyz = (self.Mat_z @ s_xyz.T).T # 1つのSをz軸上に配置(x軸周り回転) theta2 = np.arctan2(s_xyz[n_S_idx][1], s_xyz[n_S_idx][2]) self.Mat_x = np.array([[1, 0, 0], [0, np.cos(-theta2), np.sin(-theta2)], [0, -np.sin(-theta2), np.cos(-theta2)]]) s_xyz = (self.Mat_x @ s_xyz.T).T # もう1つのSをyz平面上に配置(z軸周り回転) theta3 = np.arctan2(s_xyz[n2_S_idx][0], s_xyz[n2_S_idx][1]) self.Mat_z2 = np.array([[np.cos(-theta3), np.sin(-theta3), 0], [-np.sin(-theta3), np.cos(-theta3), 0], [0, 0, 1]]) s_xyz = (self.Mat_z2 @ s_xyz.T).T self.s_xyz = s_xyz # PS4座標保存 (Output用) p_xyz = np.array([0, 0, 0]) self.ps4_coord = np.vstack((p_xyz, s_xyz)) self.ps4_coord = self.ps4_coord / ang_borr # 座標回転関数の定義 def rotate_coords(coords_fractional): if len(coords_fractional) == 0: return np.array([]) rot_xyz = coords_fractional[:, 0:3] - p_data[0:3] rot_xyz = rot_xyz - np.round(rot_xyz) rot_xyz = rot_xyz @ self.M rot_xyz = (self.Mat_z @ rot_xyz.T).T rot_xyz = (self.Mat_x @ rot_xyz.T).T rot_xyz = (self.Mat_z2 @ rot_xyz.T).T return rot_xyz # ヒストグラム範囲 bounds = [[-(self.cutoff)/ang_borr, (self.cutoff)/ang_borr]] * 3 # --- Liの処理: 座標をリストに追加 --- segment_xyz_Li = [] for xyz in self.elems["Li"]: diff = xyz[0:3] - p_data[0:3] diff = diff - np.around(diff) diff = diff @ self.M diff = np.linalg.norm(diff) for xyzdata in xyz[diff < self.cutoff, :]: segment_xyz_Li.append(xyzdata) segment_xyz_Li = np.array(segment_xyz_Li, dtype=float) if len(segment_xyz_Li) > 0: rotationed_Li = rotate_coords(segment_xyz_Li) self.li_coords_list.extend(rotationed_Li.tolist()) # --- Iの処理: 抽出した4つのIを回転させてまとめてリストに追加 --- rotationed_Is = rotate_coords(closest_Is) self.i_coords_list.extend(rotationed_Is.tolist()) count_p += 1 # ループ終了後にまとめてヒストグラム計算 print("Calculating Histograms...") volume = (self.cutoff*2)**3 bounds = [[-(self.cutoff)/ang_borr, (self.cutoff)/ang_borr]] * 3 # Li if len(self.li_coords_list) > 0: print("num_Li: ", len(self.li_coords_list)) li_arr = np.array(self.li_coords_list) / ang_borr li_hist, _ = np.histogramdd(li_arr, bins=self.mesh, range=bounds) self.li_hist = li_hist.ravel() / len(self.li_coords_list) / volume print("sum_Li: ", np.sum(self.li_hist) * volume) # I (All 4 atoms combined) if len(self.i_coords_list) > 0: print("num_I: ", len(self.i_coords_list)) i_arr = np.array(self.i_coords_list) / ang_borr i_hist, _ = np.histogramdd(i_arr, bins=self.mesh, range=bounds) self.i_hist = i_hist.ravel() / len(self.i_coords_list) / volume print("sum_I: ", np.sum(self.i_hist) * volume) def outputCube(self): output_start_time = time.time() base = re.match(r"(\d{3}.*?LiI).*?", args.trjfile.split("/")[-1]).group(1) self.atomsDic = {'I': '53', 'Li': '3', 'P': '15', 'S': '16'} if self.li_hist is None and self.i_hist is None: print("No data found.") return def get_header(): body = f"created from {__file__}, {args}\n" body += "Contains the selected quantity on a FFT grid\n" origin = [-(self.cutoff)/ang_borr] * 3 body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format( 5, *origin) body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format( self.mesh[0], (self.cutoff*2)/self.mesh[0]/ang_borr, 0, 0) body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format( self.mesh[1], 0, (self.cutoff*2)/self.mesh[1]/ang_borr, 0) body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}\n".format( self.mesh[2], 0, 0, (self.cutoff*2)/self.mesh[2]/ang_borr) # Pの座標 body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}{:>12.7f}\n".format( int(self.atomsDic["P"]), float(self.atomsDic["P"]), *self.ps4_coord[0]) # Sの座標 for s_coord in self.ps4_coord[1:5]: body += "{:>5d}{:>12.7f}{:>12.7f}{:>12.7f}{:>12.7f}\n".format( int(self.atomsDic["S"]), float(self.atomsDic["S"]), *s_coord) return body # Liの出力 if self.li_hist is not None: body = get_header() for idx, r in enumerate(self.li_hist): if idx % 6 == 5: body += "{:>13.5E}\n".format(r) else: body += "{:>13.5E}".format(r) outfile = f"{dirname}/{base}_PS4_Li.cube" with open(outfile, "w") as o: o.write(body) print(f"{outfile} was created.") # Iの出力 (まとめて1ファイルに出力) if self.i_hist is not None: body = get_header() for idx, r in enumerate(self.i_hist): if idx % 6 == 5: body += "{:>13.5E}\n".format(r) else: body += "{:>13.5E}".format(r) outfile = f"{dirname}/{base}_PS4_I.cube" with open(outfile, "w") as o: o.write(body) print(f"{outfile} was created.") output_end_time = time.time() print(f"output_time : {output_end_time - output_start_time} s") trj = LammpsTrj() trj.makeCube(args.trjfile) trj.outputCube()