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])