PySCFとGPUの連携で中・大規模分子を解析【Pythonで始める量子化学計算】

import tensorflow as tftf.test.gpu_device_name() 

!pip install gpu4pyscf-cuda12x
!pip install pyscf
!pip install numpy
!pip install geometric
!pip install rdkit
!pip install pubchempy
!pip install py3Dmol

!nvcc --version

Cuda compilation tools, release 12.0, V12.x.x

!nvidia-smi

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 530.30.02   Driver Version: 530.30.02   CUDA Version: 12.1       |
+-----------------------------------------------------------------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
+-----------------------------------------------------------------------------+

import time
from pyscf import gto, scf
from gpu4pyscf import scf as gpu_scf
from rdkit import Chem
from rdkit.Chem import AllChem
import pubchempy as pcp

# 化合物名を使用してPubChemから取得
compound = pcp.get_compounds("carotenoid", "name")[0]
smiles = compound.isomeric_smiles
mol = Chem.MolFromSmiles(smiles)

# 水素原子を追加
mol = Chem.AddHs(mol)

# 3D構造の生成と最適化
AllChem.EmbedMolecule(mol)
AllChem.UFFOptimizeMolecule(mol)

# XYZ形式で分子構造を取得
xyz_coords = Chem.MolToXYZBlock(mol)

# 原子の数とヘッダー情報を追加
num_atoms = mol.GetNumAtoms()
xyz_full = f"{num_atoms}\\nCarotene molecule\\n" + xyz_coords

# 必要な行のみを抽出（5行目以降）
xyz_lines = xyz_full.splitlines()
xyz_trimmed = "\\n".join(xyz_lines[4:])

print("使用するXYZデータ:")
print(xyz_trimmed)

# PySCFの分子設定
mol = gto.Mole()
mol.atom = xyz_trimmed
mol.basis = "sto-3g"
mol.build()

# CPU計算
start_cpu = time.time()
mf_cpu = scf.RHF(mol)
mf_cpu.kernel()
end_cpu = time.time()
cpu_time = end_cpu - start_cpu

# GPU計算
start_gpu = time.time()
mf_gpu = gpu_scf.RHF(mol)
mf_gpu.kernel()
end_gpu = time.time()
gpu_time = end_gpu - start_gpu

print(f"CPU計算時間: {cpu_time:.2f}秒")
print(f"GPU計算時間: {gpu_time:.2f}秒")

import os

folder_path = "/content/drive/MyDrive/pyscf_GPU"
base_filename = "Epitalon"
path = os.path.join(folder_path, base_filename)

if not os.path.exists(folder_path):
    os.makedirs(folder_path)
    print(f"Created folder: {folder_path}")

checkpoint_file = path + ".chk"
print(f"Checkpoint file will be saved as: {checkpoint_file}")

from rdkit import Chem
from rdkit.Chem import AllChem
import pubchempy as pcp

compound = pcp.get_compounds("Epitalon", "name")[0]
smiles = compound.isomeric_smiles
mol = Chem.MolFromSmiles(smiles)
mol = Chem.AddHs(mol)
AllChem.EmbedMolecule(mol)
AllChem.UFFOptimizeMolecule(mol)

# XYZ形式に変換
xyz_coords = Chem.MolToXYZBlock(mol)
num_atoms = mol.GetNumAtoms()
xyz_full = f"{num_atoms}\\nEpitalon molecule\\n" + xyz_coords
xyz_trimmed = "\\n".join(xyz_full.splitlines()[4:])
print("使用するXYZデータ:\\n", xyz_trimmed)

from pyscf import gto
from gpu4pyscf import dft
from pyscf.geomopt import geometric_solver

# 分子の設定
mol = gto.Mole()
mol.atom = xyz_trimmed
mol.basis = "sto-3g"
mol.build()

# GPUでのDFT計算設定
mf = dft.RKS(mol, xc='B3LYP').density_fit().to_gpu()
mf.chkfile = path + '.chk'
mf.verbose = 1

energies = []  # エネルギー記録用
geometries = []  # 座標記録用

def cb(envs):
    mf = envs['g_scanner'].base  # 現在の計算オブジェクトを取得
    energies.append(mf.e_tot)  # エネルギーをリストに追加
    geometries.append(mol.atom_coords(unit="ANG"))  # ジオメトリーをリストに追加

conv_params = {
    'convergence_energy': 1e-5,
    'convergence_grms': 3e-4,
    'convergence_gmax': 4.5e-4,
    'convergence_drms': 1.2e-3,
    'convergence_dmax': 1.8e-3
}

# 構造最適化実行
geometric_solver.optimize(mf, callback=cb, **conv_params)

import matplotlib.pyplot as plt

plt.figure(figsize=(10, 5))
plt.plot(energies, marker='o')
plt.title('Energy Convergence During Geometry Optimization')
plt.xlabel('Optimization Step')
plt.ylabel('Total Energy (Hartree)')
plt.grid(True)
plt.show()

import py3Dmol
from ipywidgets import interact, IntSlider, Dropdown

# ジオメトリーの保存
for i, coords in enumerate(geometries):
    filename = f"{path}_geometry_step_{i+1}.xyz"
    with open(filename, 'w') as f:
        f.write(f"{mol.natm}\\n")
        f.write(f"Step {i+1} Energy: {energies[i]}\\n")
        for j in range(mol.natm):
            symbol = mol.atom_pure_symbol(j)
            x, y, z = coords[j]
            f.write(f"{symbol} {x:.6f} {y:.6f} {z:.6f}\\n")

print("All geometries have been saved.")

# 保存されたXYZファイルを読み込む
def read_xyz(file_path):
    with open(file_path, 'r') as f:
        return f.read()

# 3D可視化
def show_structure(xyz_content, style):
    view = py3Dmol.view(width=800, height=400)
    view.addModel(xyz_content, 'xyz')
    if style == "ball and stick":
        view.setStyle({'stick': {'radius': 0.1}, 'sphere': {'radius': 0.3}})
    elif style == "stick":
        view.setStyle({'stick': {}})
    elif style == "sphere":
        view.setStyle({'sphere': {}})
    view.zoomTo()
    view.show()

# ウィジェットを使ったインタラクション
def visualize_structure(idx, style):
    file_path = f"{path}_geometry_step_{idx}.xyz"
    xyz_content = read_xyz(file_path)
    show_structure(xyz_content, style)

max_steps = len(geometries)
interact(visualize_structure,
         idx=IntSlider(min=1, max=max_steps, step=1, value=1),
         style=Dropdown(options=["ball and stick", "stick", "sphere"], value="ball and stick"))

# 最適化後のジオメトリーを取得
final_geometry = mol.atom_coords(unit="ANG")

# SCF計算の再実行
mf = dft.RKS(mol)
mf.chkfile = 'checkpoint_file.chk'
mf.xc = 'B3LYP'
mf.kernel()

# 計算結果の検証
if mf.mo_occ is not None and len(mf.mo_occ) > 0:
    print("SCF calculation completed successfully.")
    print(f"Number of occupied orbitals: {sum(mf.mo_occ > 0)}")
else:
    raise ValueError(
        "SCF calculation failed or mo_occ is not properly assigned."
        )

# HOMOおよびLUMOのインデックスを取得
homo_index = (mf.mo_occ > 0).nonzero()[0][-1]
lumo_index = (mf.mo_occ == 0).nonzero()[0][0]

from pyscf import tools
import py3Dmol
from IPython.display import display, HTML
import ipywidgets as widgets
import numpy as np

# 指定した軌道のCUBEファイルを生成する関数
def generate_cube_file(orbital_index, filename):
    # Convert mf.mo_coeff to a NumPy array before passing it to tools.cubegen.orbital
    coeff = mf.mo_coeff[:, orbital_index].get() # Use .get() to move the data to CPU
    tools.cubegen.orbital(mol, filename, coeff) 

# 可視化する軌道のリストを作成
orbitals = {
    "LUMO+3": lumo_index + 3,
    "LUMO+2": lumo_index + 2,
    "LUMO+1": lumo_index + 1,
    "LUMO": lumo_index,
    "HOMO": homo_index,
    "HOMO-1": homo_index - 1,
    "HOMO-2": homo_index - 2,
    "HOMO-3": homo_index - 3,
}

# エネルギーをeVに変換
hartree_to_ev = 27.2114
orbital_energies = {
    name: mf.mo_energy[idx] * hartree_to_ev for name, idx in orbitals.items()
}

# ウィジェットの作成
dropdown = widgets.Dropdown(
    options=[
        (f"{name} ({orbital_energies[name]:.6f} eV)", name)
        for name in orbitals.keys()
    ],
    description="Orbital:",
    disabled=False,
)

output = widgets.Output()

# 可視化関数
def visualize_orbital(orbital_name):
    orbital_index = orbitals[orbital_name]
    cube_filename = f"{orbital_name}.cube"

    generate_cube_file(orbital_index, cube_filename)

    view = py3Dmol.view(width=800, height=600)
    with open(cube_filename) as f:
        cube = f.read()
    view.addVolumetricData(
        cube, "cube", {"isoval": 0.02, "color": "red", "opacity": 0.75}
    )
    view.addVolumetricData(
        cube, "cube", {"isoval": -0.02, "color": "blue", "opacity": 0.75}
    )

    # 最適化後の分子形状を追加
    xyz = "\\n".join(
    [
        (
            f"{mol.atom_symbol(i)} "
            f"{final_geometry[i][0]:.6f} "
            f"{final_geometry[i][1]:.6f} "
            f"{final_geometry[i][2]:.6f}"
        )
        for i in range(mol.natm)
    ]
    )
    xyz = f"{mol.natm}\\n\\n" + xyz  # xyzフォーマットのヘッダー追加
    view.addModel(xyz, "xyz")
    view.setStyle({"stick": {}})
    view.zoomTo()

    # エネルギーのラベルを追加
    orbital_energy = orbital_energies[orbital_name]
    view.addLabel(
        f"{orbital_name}: {orbital_energy:.6f} eV",
        {
            "fontSize": 12,
            "fontColor": "black",
            "backgroundColor": "white",
            "backgroundOpacity": 0.8,
            "showBackground": True,
            "position": {"x": 0, "y": 0, "z": 0},
        },
    )

    with output:
        output.clear_output(wait=True)
        display(HTML(view._make_html()))

# ドロップダウンの変更イベントに可視化関数をバインド
def on_dropdown_change(change):
    with output:
        output.clear_output(wait=True)
        visualize_orbital(change["new"])

dropdown.observe(on_dropdown_change, names="value")

# インターフェースの表示
display(dropdown)
display(output)

# 初期表示
visualize_orbital(list(orbitals.keys())[0])