Psi4Writeinput.py

#!/usr/bin/env python
"""
Psi4WriteInput.py - Generate Psi4 input file from XYZ or SDF file

Usage:
    Psi4WriteInput.py -i <infile> -o <outfile> [options]

Examples:
    # 默认方法 r2scan-3c/def2-mtzvpp
    Psi4WriteInput.py -i molecule.xyz -o input.dat

    # 指定方法和基组
    Psi4WriteInput.py -i molecule.xyz -o input.dat --method wb97x-d3bj --basis def2-tzvp

    # 溶剂效应 (water)
    Psi4WriteInput.py -i molecule.xyz -o input.dat --solvent water
"""

import os
import re
import sys
import argparse
from typing import Optional, Tuple, List, Dict, Any


def parse_xyz_file(xyz_file: str) -> Tuple[List[str], List[Tuple[float, float, float]], Optional[str]]:
    """解析 XYZ 文件"""
    atoms = []
    coords = []
    comment = None

    with open(xyz_file, 'r') as f:
        lines = f.readlines()

    if len(lines) < 3:
        raise ValueError(f"XYZ 文件格式错误: {xyz_file}，至少需要3行")

    try:
        natoms = int(lines[0].strip())
    except ValueError:
        raise ValueError(f"XYZ 文件第1行应为原子数: {lines[0].strip()}")

    if len(lines) > 1:
        comment = lines[1].strip()

    for i in range(2, min(2 + natoms, len(lines))):
        parts = lines[i].strip().split()
        if len(parts) < 4:
            raise ValueError(f"XYZ 文件第{i+1}行格式错误")

        element = parts[0]
        x, y, z = float(parts[1]), float(parts[2]), float(parts[3])
        atoms.append(element)
        coords.append((x, y, z))

    if len(atoms) != natoms:
        raise ValueError(f"原子数不匹配: 声明 {natoms}，实际 {len(atoms)}")

    return atoms, coords, comment


def parse_sdf_file(sdf_file: str) -> Tuple[List[str], List[Tuple[float, float, float]], int, int]:
    """解析 SDF 文件"""
    from rdkit import Chem

    suppl = Chem.SDMolSupplier(sdf_file)
    mol = suppl[0]
    if mol is None:
        raise ValueError(f"无法读取 SDF 文件: {sdf_file}")

    conf = mol.GetConformer()
    atoms = [atom.GetSymbol() for atom in mol.GetAtoms()]
    coords = [(conf.GetAtomPosition(i).x,
               conf.GetAtomPosition(i).y,
               conf.GetAtomPosition(i).z)
              for i in range(mol.GetNumAtoms())]

    charge = 0
    multiplicity = 1

    if mol.HasProp('FormalCharge'):
        charge = int(mol.GetProp('FormalCharge'))
    if mol.HasProp('SpinMultiplicity'):
        multiplicity = int(mol.GetProp('SpinMultiplicity'))

    return atoms, coords, charge, multiplicity


def parse_comment_for_charge_multiplicity(comment: Optional[str]) -> Tuple[int, int]:
    """从注释行解析电荷和自旋多重度"""
    charge = 0
    multiplicity = 1

    if comment:
        parts = comment.strip().split()
        if len(parts) >= 2:
            try:
                charge = int(parts[0])
                multiplicity = int(parts[1])
                return charge, multiplicity
            except ValueError:
                pass

    return charge, multiplicity


def generate_psi4_input_python_api(
    atoms: List[str],
    coords: List[Tuple[float, float, float]],
    charge: int,
    multiplicity: int,
    method: str,
    basis: str,
    solvent: Optional[str] = None,
    solvent_model: str = "pcm",
    memory_gb: int = 1,
    num_threads: int = 1,
    reference: Optional[str] = None,
    ddx_options: Optional[Dict[str, str]] = None,
    comments: Optional[List[str]] = None
) -> str:
    """
    生成 Psi4 输入文件内容（使用 Python API 风格，支持复合方法名称）
    """
    lines = []

    # 添加注释
    lines.append("#" + "=" * 68)
    lines.append("# Psi4 Input File Generated by Psi4WriteInput.py")
    lines.append(f"# Method: {method}")
    lines.append(f"# Basis: {basis}")
    if solvent:
        lines.append(f"# Solvent: {solvent} ({solvent_model})")
    lines.append(f"# Charge: {charge}, Multiplicity: {multiplicity}")
    lines.append("#" + "=" * 68)
    lines.append("")

    if comments:
        for comment in comments:
            lines.append(f"# {comment}")
        lines.append("")

    # 内存和线程设置
    lines.append(f"psi4.set_memory('{memory_gb} GB')")
    lines.append(f"psi4.set_num_threads({num_threads})")
    lines.append("")

    # 分子定义
    lines.append(f"mol = psi4.geometry(\"\"\"")
    lines.append(f"  {charge} {multiplicity}")

    for atom, (x, y, z) in zip(atoms, coords):
        lines.append(f"  {atom} {x:.10f} {y:.10f} {z:.10f}")

    lines.append(f"\"\"\")")
    lines.append("")

    # 构建复合方法名称
    method_lower = method.lower()
    basis_upper = basis.upper()

    # 复合方法名称映射
    if method_lower == 'r2scan-3c' or method_lower == 'r2scan3c':
        composite_method = "R2SCAN-3C/DEF2-mTZVPP"
        if basis_upper != "DEF2-MTZVPP":
            composite_method = f"{composite_method}/{basis_upper}"
    elif method_lower == 'wb97x-d3bj' or method_lower == 'wb97x_d3bj':
        composite_method = f"WB97X-D3BJ/{basis_upper}"
    elif method_lower == 'b3lyp-d3bj' or method_lower == 'b3lyp_d3bj':
        composite_method = f"B3LYP-D3BJ/{basis_upper}"
    elif method_lower == 'pbe0-d3bj' or method_lower == 'pbe0_d3bj':
        composite_method = f"PBE0-D3BJ/{basis_upper}"
    elif method_lower == 'hf':
        composite_method = f"HF/{basis_upper}"
    else:
        composite_method = f"{method.upper()}/{basis_upper}"

    # 构建 options 字典
    options = {
        "scf_type": "df",
    }

    # 参考波函数
    if reference is None:
        if multiplicity == 1:
            reference = "rks" if method_lower != 'hf' else "rhf"
        else:
            reference = "uks" if method_lower != 'hf' else "uhf"
    options["reference"] = reference

    # 溶剂模型设置（使用 DDX）
    if solvent:
        options["ddx"] = True
        options["ddx_model"] = solvent_model
        options["ddx_solvent"] = solvent
        # 默认参数
        options["ddx_radii_set"] = "uff"
        options["ddx_radii_scaling"] = 1.1

        # 用户自定义参数
        if ddx_options:
            for key, value in ddx_options.items():
                # 确保键名格式正确
                if not key.startswith("ddx_"):
                    options[f"ddx_{key}"] = value
                else:
                    options[key] = value

    # 输出 options 设置
    lines.append("# Set options")
    lines.append("psi4.set_options({")
    for key, value in options.items():
        if isinstance(value, bool):
            lines.append(f"    '{key}': {str(value).lower()},")
        elif isinstance(value, str):
            lines.append(f"    '{key}': '{value}',")
        else:
            lines.append(f"    '{key}': {value},")
    lines.append("})")
    lines.append("")

    # 执行能量计算 - 使用复合方法名称
    lines.append("# Compute single point energy")
    lines.append(f"energy = psi4.energy('{composite_method}')")
    lines.append("")
    lines.append("# Print result")
    lines.append("print(f'Total Energy: {energy:.8f} Hartree')")
    lines.append("print(f'Total Energy: {energy * 627.509474:.3f} kcal/mol')")

    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(
        description="Generate Psi4 input file from XYZ or SDF file (Python API format)",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # 生成输入文件（默认 R2SCAN-3C/DEF2-MTZVPP）
  %(prog)s -i molecule.xyz -o input.dat

  # 指定方法和基组（自动组合为复合方法名）
  %(prog)s -i molecule.xyz -o input.dat --method wb97x-d3bj --basis def2-tzvp
  # 生成: energy = psi4.energy('WB97X-D3BJ/DEF2-TZVP')

  # 溶剂效应 (water)
  %(prog)s -i molecule.xyz -o input.dat --solvent water

  # 使用 COSMO 溶剂模型
  %(prog)s -i molecule.xyz -o input.dat --solvent water --solvent-model cosmo

  # 自定义 DDX 参数
  %(prog)s -i molecule.xyz -o input.dat --solvent water --ddx-options "radii_set,Bondi,radii_scaling,1.2"
        """
    )

    # 必需参数
    parser.add_argument("-i", "--infile", required=True, help="输入文件 (XYZ, SDF, SD)")
    parser.add_argument("-o", "--outfile", required=True, help="输出 Psi4 输入文件 (.dat)")

    # 分子参数
    parser.add_argument("-c", "--charge", type=int, help="分子电荷（默认从文件读取或 0）")
    parser.add_argument("-m", "--multiplicity", type=int, help="自旋多重度（默认从文件读取或 1）")

    # 计算参数
    parser.add_argument("--method", default="r2scan-3c",
                        help="DFT 方法（默认: r2scan-3c）")
    parser.add_argument("--basis", default="def2-mtzvpp",
                        help="基组（默认: def2-mtzvpp）")
    parser.add_argument("--reference", choices=['rhf', 'uhf', 'rohf', 'rks', 'uks'],
                        help="参考波函数（自动检测）")

    # 溶剂参数
    parser.add_argument("--solvent", help="溶剂名称（如 water, ethanol）")
    parser.add_argument("--solvent-model", default="cosmo", choices=['pcm', 'cosmo'],
                        help="溶剂模型（默认: pcm）")

    # 计算资源
    parser.add_argument("--memory-gb", type=int, default=24, help="内存 (GB, 默认: 24)")
    parser.add_argument("--num-threads", type=int, default=24, help="线程数 (默认: 24")

    # DDX 选项
    parser.add_argument("--ddx-options",
                        help="DDX 选项，逗号分隔，如 'radii_set,UFF,radii_scaling,1.1'")

    # 注释
    parser.add_argument("--comment", action="append", help="添加自定义注释行（可多次使用）")

    # 其他
    parser.add_argument("--overwrite", action="store_true", help="覆盖已存在的输出文件")
    parser.add_argument("-v", "--verbose", action="store_true", help="详细输出")

    args = parser.parse_args()

    # 检查输出文件是否已存在
    if not args.overwrite and os.path.exists(args.outfile):
        print(f"错误: 输出文件已存在: {args.outfile}（使用 --overwrite 覆盖）")
        sys.exit(1)

    # 判断输入文件类型并解析
    infile_ext = os.path.splitext(args.infile)[1].lower()

    if infile_ext == '.xyz':
        if args.verbose:
            print(f"读取 XYZ 文件: {args.infile}")

        try:
            atoms, coords, comment = parse_xyz_file(args.infile)
        except Exception as e:
            print(f"错误: 解析 XYZ 文件失败 - {e}")
            sys.exit(1)

        if args.verbose:
            print(f"  原子数: {len(atoms)}")

        if args.charge is None or args.multiplicity is None:
            xyz_charge, xyz_multiplicity = parse_comment_for_charge_multiplicity(comment)
            charge = args.charge if args.charge is not None else xyz_charge
            multiplicity = args.multiplicity if args.multiplicity is not None else xyz_multiplicity
        else:
            charge = args.charge
            multiplicity = args.multiplicity

    elif infile_ext in ['.sdf', '.sd']:
        if args.verbose:
            print(f"读取 SDF 文件: {args.infile}")

        try:
            from rdkit import Chem
            atoms, coords, sdf_charge, sdf_multiplicity = parse_sdf_file(args.infile)
        except ImportError:
            print("错误: 读取 SDF 文件需要 RDKit，请安装: conda install rdkit")
            sys.exit(1)
        except Exception as e:
            print(f"错误: 解析 SDF 文件失败 - {e}")
            sys.exit(1)

        charge = args.charge if args.charge is not None else sdf_charge
        multiplicity = args.multiplicity if args.multiplicity is not None else sdf_multiplicity

    else:
        print(f"错误: 不支持的文件格式: {infile_ext}（支持: .xyz, .sdf, .sd）")
        sys.exit(1)

    # 解析 DDX 选项
    ddx_options = {}
    if args.ddx_options:
        parts = args.ddx_options.split(',')
        for i in range(0, len(parts), 2):
            if i + 1 < len(parts):
                ddx_options[parts[i]] = parts[i+1]

    # 生成输入文件
    if args.verbose:
        print(f"\n生成 Psi4 输入文件:")
        print(f"  方法: {args.method}")
        print(f"  基组: {args.basis}")
        print(f"  复合方法: {args.method.upper()}/{args.basis.upper()}")
        print(f"  电荷: {charge}, 自旋多重度: {multiplicity}")
        if args.solvent:
            print(f"  溶剂: {args.solvent} ({args.solvent_model})")

    psi4_input = generate_psi4_input_python_api(
        atoms=atoms,
        coords=coords,
        charge=charge,
        multiplicity=multiplicity,
        method=args.method,
        basis=args.basis,
        solvent=args.solvent,
        solvent_model=args.solvent_model,
        memory_gb=args.memory_gb,
        num_threads=args.num_threads,
        reference=args.reference,
        ddx_options=ddx_options,
        comments=args.comment
    )

    # 写入文件
    with open(args.outfile, 'w') as f:
        f.write(psi4_input)

    print(f"\n✅ Psi4 输入文件已生成: {args.outfile}")

    if args.verbose:
        print("\n" + "-" * 50)
        print("文件内容预览:")
        print("-" * 50)
        preview_lines = psi4_input.split('\n')[:40]
        for line in preview_lines:
            print(line)
        if len(psi4_input.split('\n')) > 40:
            print(f"...（共 {len(psi4_input.split('\n'))} 行）")
        print("-" * 50)


if __name__ == "__main__":
    main()