lps-lii-ps4-density-analysis/cube2mayavi.py

#!/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")