#!/usr/bin/env python # ./cube2mayavi.py -i 050Li3PS4-050LiI_PS4_I.cube -atom I -iso 1.26483e-09 import numpy as np from mayavi import mlab ang_borr = 0.5291772109217 def read_cube(filename): """ .cubeファイルを読み込み、原子座標と3次元ボクセルデータを返す関数 """ with open(filename, 'r') as f: # ヘッダーの読み飛ばし（最初の2行はコメント） f.readline() f.readline() # 3行目: 原子の数 と 原点座標 line = f.readline().split() natoms = abs(int(line[0])) # 原子の数 origin = np.array([float(x) for x in line[1:4]]) # 原点 # 4-6行目: グリッドの分割数(N)とベクトル(X, Y, Z軸) line = f.readline().split() nx, x_vec = int(line[0]), np.array([float(x) for x in line[1:4]]) line = f.readline().split() ny, y_vec = int(line[0]), np.array([float(x) for x in line[1:4]]) line = f.readline().split() nz, z_vec = int(line[0]), np.array([float(x) for x in line[1:4]]) # 原子座標の読み込み atoms = [] for _ in range(natoms): line = f.readline().split() atoms.append([float(x) for x in line[2:5]]) atoms = np.array(atoms) # ボクセルデータの読み込み data = [] for line in f: data.extend([float(x) for x in line.split()]) vol_data = np.array(data).reshape(nx, ny, nz) return vol_data, atoms, origin, (x_vec, y_vec, z_vec) def draw_radial_scale(max_dist=10.0, step=1.0, axis='x'): """ 原点から動径方向に目盛り（定規）を描画する """ if axis == 'y': vec = np.array([0, 1, 0]) tick_dir = np.array([0.2, 0, 0]) elif axis == 'z': vec = np.array([0, 0, 1]) tick_dir = np.array([0.2, 0, 0]) else: # default x vec = np.array([1, 0, 0]) tick_dir = np.array([0, 0, 0.2]) # メインの直線を引く end_point = vec * max_dist mlab.plot3d([0, end_point[0]], [0, end_point[1]], [0, end_point[2]], tube_radius=0.02, color=(0, 0, 0)) # 目盛りと数字を配置 for r in np.arange(step, max_dist + step, step): pos = vec * r t_start = pos - tick_dir t_end = pos + tick_dir mlab.plot3d([t_start[0], t_end[0]], [t_start[1], t_end[1]], [t_start[2], t_end[2]], tube_radius=0.02, color=(0, 0, 0)) text_pos = pos + tick_dir * 1.5 mlab.text3d(text_pos[0], text_pos[1], text_pos[2], f"{int(r)}", scale=0.25, color=(0, 0, 0)) def change_view(azimuth=30, elevation=60): mlab.view(azimuth=azimuth, elevation=elevation) def Li_visualize(filename, iso_val=None, radius_cutoff=5.0): print(f"Reading {filename}...") vol_data, atoms, origin, vectors = read_cube(filename) atoms, origin_ang = atoms * ang_borr, origin * ang_borr dx = np.linalg.norm(vectors[0]) * ang_borr dy = np.linalg.norm(vectors[1]) * ang_borr dz = np.linalg.norm(vectors[2]) * ang_borr nx, ny, nz = vol_data.shape i, j, k = np.mgrid[0:nx, 0:ny, 0:nz] X = origin_ang[0] + i * dx Y = origin_ang[1] + j * dy Z = origin_ang[2] + k * dz R = np.sqrt(X**2 + Y**2 + Z**2) vol_inner = np.where(R <= 5, vol_data, 0.0) vol_outer = np.where((R < 9) & (R > 5), vol_data, 0.0) mlab.figure(bgcolor=(1, 1, 1), size=(800, 600)) if args.iso_val is None: iso_val = np.std(vol_data) * 2.0 else: iso_val = args.iso_val print(f"Plotting isosurface at value: +/- {iso_val:.4f}") print("max_value: ", np.max(vol_data)) print("min_value: ", np.min(vol_data)) # --- 確率密度 内側 --- src_in = mlab.pipeline.scalar_field(vol_inner) src_in.spacing = [dx, dy, dz] src_in.origin = origin_ang obj_in = mlab.pipeline.iso_surface(src_in, contours=[iso_val], opacity=0.4, color=(0, 0.2, 0.8)) # --- 確率密度 外側 --- src_out = mlab.pipeline.scalar_field(vol_outer) src_out.spacing = [dx, dy, dz] src_out.origin = origin_ang obj_out = mlab.pipeline.iso_surface(src_out, contours=[iso_val], opacity=0.2, color=(0.7, 0.1, 0.1)) # --- 原子の描画 --- mlab.points3d(atoms[0, 0], atoms[0, 1], atoms[0, 2], scale_factor=0.8, color=(0.576, 0.439, 0.8), resolution=20) mlab.points3d(atoms[1:, 0], atoms[1:, 1], atoms[1:, 2], scale_factor=0.8, color=(1, 1, 0), resolution=20) # --- 結合の描画 --- p_coord = atoms[0] s_coords = atoms[1:] for s_coord in s_coords: mlab.plot3d([p_coord[0], s_coord[0]], [p_coord[1], s_coord[1]], [p_coord[2], s_coord[2]], tube_radius=0.1, color=(0.6, 0.6, 0.6), opacity=1.0) mlab.points3d(0, 0, 0, scale_factor=6.6, mode='sphere', color=(0, 0, 0), opacity=0.05, resolution=50) mlab.points3d(0, 0, 0, scale_factor=14, mode='sphere', color=(0, 0, 0), opacity=0.05, resolution=50) mlab.axes() return obj_in, obj_out def I_visualize(filename, iso_val=None, cutoff=8): """ 1つのIのcubeファイルを読み込んで可視化する """ print(f"Reading {filename}...") vol_data, atoms, origin, vectors = read_cube(filename) atoms, origin_ang = atoms * ang_borr, origin * ang_borr dx = np.linalg.norm(vectors[0]) * ang_borr dy = np.linalg.norm(vectors[1]) * ang_borr dz = np.linalg.norm(vectors[2]) * ang_borr mlab.figure(bgcolor=(1, 1, 1), size=(800, 600)) if iso_val is None: if 'args' in globals() and hasattr(args, 'iso_val') and args.iso_val is not None: iso_val = args.iso_val else: iso_val = np.std(vol_data) * 2.0 print(f"Plotting isosurface at value: +/- {iso_val:.4f}") # --- カットオフによるデータの球状マスク処理 --- if cutoff is not None: nx, ny, nz = vol_data.shape x = origin_ang[0] + np.arange(nx) * dx y = origin_ang[1] + np.arange(ny) * dy z = origin_ang[2] + np.arange(nz) * dz X, Y, Z = np.meshgrid(x, y, z, indexing='ij') R = np.sqrt(X**2 + Y**2 + Z**2) vol_data[R > cutoff] = 0.0 print("max_value: ", np.max(vol_data)) print("min_value: ", np.min(vol_data)) # --- 確率密度 --- src_in = mlab.pipeline.scalar_field(vol_data) src_in.spacing = [dx, dy, dz] src_in.origin = origin_ang # 赤色で表示 obj = mlab.pipeline.iso_surface(src_in, contours=[iso_val], opacity=0.3, color=(0.850, 0.058, 0.058)) # --- 原子の描画 --- mlab.points3d(atoms[0, 0], atoms[0, 1], atoms[0, 2], scale_factor=0.8, color=(0.576, 0.439, 0.8), resolution=20) mlab.points3d(atoms[1:, 0], atoms[1:, 1], atoms[1:, 2], scale_factor=0.8, color=(1, 1, 0), resolution=20) # --- 結合の描画 --- p_coord = atoms[0] s_coords = atoms[1:] for s_coord in s_coords: mlab.plot3d([p_coord[0], s_coord[0]], [p_coord[1], s_coord[1]], [p_coord[2], s_coord[2]], tube_radius=0.1, color=(0.6, 0.6, 0.6), opacity=1.0) mlab.points3d(0, 0, 0, scale_factor=11, mode='sphere', color=(0, 0, 0), opacity=0.05, resolution=50) return obj # --- 実行部分 --- if __name__ == "__main__": import argparse description = """This is a test program""" par = argparse.ArgumentParser(description=description) par.add_argument('-i', '--infiles', default="", required=True, nargs="+", help='input file') par.add_argument('-atom', '--atom', default="", required=True, choices=['Li', 'I'], help='atom') par.add_argument('-iso', '--iso_val', required=False, type=float, help='しきい値') args = par.parse_args() if args.atom == "Li": obj_in, obj_out = Li_visualize(args.infiles[0]) mlab.show() elif args.atom == "I": if len(args.infiles) > 1: print("Warning: Multiple files were passed, but only the first one will be read.") obj = I_visualize(args.infiles[0]) mlab.show() else: print("atom error")