admin
a16d2239a1
- 源码: engines/c/ → engines/src/c/ - 源码: engines/cpp/ → engines/src/cpp/ - 源码: engines/fortran/ → engines/src/fortran/ - 编译产物: engines/release/ (静态编译) - dynamics_c.exe (516KB, C 引擎) - dynamics_cpp.exe (3.3MB, C++ 引擎) - dynamics_f90.exe (988KB, Fortran 引擎) - compute.py: engine_map 指向 release/ - compute.py: param.json 和校准缓存移至 release/ - engine_dll.py: DLL 搜索路径指向 release/ - release/ 目录纳入 git 版本管理，用户克隆后直接可用
425 lines
16 KiB
Python
425 lines
16 KiB
Python
"""
|
||
engines/engine_dll.py
|
||
---------------------
|
||
Python ctypes 包装器：加载 C/C++/Fortran 动态链接库并调用 run_dynamics()。
|
||
|
||
用法（由 compute.py 内部调用，不直接运行）：
|
||
|
||
from engines.engine_dll import load_dll, run_dynamics_dll
|
||
|
||
dll = load_dll("c") # 自动查找 engines/c/build/dynamics_c.dll/.so/.dylib
|
||
arrays = run_dynamics_dll(dll, config, atom_data, bond_data, driver_data)
|
||
# arrays: dict with keys x, y, z, vx, vy, vz shape=(n_frames, n_atoms)
|
||
|
||
DLL 编译（C 版本）：
|
||
Windows: gcc -O3 -shared -o engines/c/build/dynamics_c.dll engines/c/dynamics_lib.c -lm
|
||
Linux: gcc -O3 -shared -fPIC -o engines/c/build/dynamics_c.so engines/c/dynamics_lib.c -lm
|
||
macOS: gcc -O3 -dynamiclib -o engines/c/build/dynamics_c.dylib engines/c/dynamics_lib.c -lm
|
||
"""
|
||
|
||
import ctypes
|
||
import os
|
||
import platform
|
||
import numpy as np
|
||
|
||
# ── DLL 文件名后缀 ─────────────────────────────────────────────
|
||
_SUFFIX = {
|
||
"windows": ".dll",
|
||
"linux": ".so",
|
||
"darwin": ".dylib",
|
||
}
|
||
|
||
# ── method 字符串 → 整数 ID ────────────────────────────────────
|
||
_METHOD_ID = {
|
||
"explicit_euler": 0,
|
||
"euler": 0,
|
||
"implicit_euler": 1,
|
||
"midpoint": 2,
|
||
"leapfrog": 3,
|
||
}
|
||
|
||
_HERE = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
|
||
_DLL_NAME = {
|
||
"c": "dynamics_c",
|
||
"cpp": "dynamics_cpp",
|
||
"c++": "dynamics_cpp",
|
||
"fortran": "dynamics_f90",
|
||
"f90": "dynamics_f90",
|
||
# "python" 引擎通过直接 import 调用，不使用 DLL
|
||
}
|
||
|
||
# 引擎名规范化：将别名统一为目录名
|
||
_ENGINE_DIR = {
|
||
"c": "c",
|
||
"cpp": "cpp",
|
||
"c++": "cpp",
|
||
"fortran": "fortran",
|
||
"f90": "fortran",
|
||
"python": "python",
|
||
}
|
||
|
||
|
||
def _dll_candidates(engine: str) -> list[str]:
|
||
"""返回 DLL 候选路径列表（按优先级）。"""
|
||
sys = platform.system().lower()
|
||
ext = _SUFFIX.get(sys, ".so")
|
||
eng_dir = _ENGINE_DIR.get(engine, engine)
|
||
name = _DLL_NAME.get(engine, f"dynamics_{engine}")
|
||
base = os.path.join(_HERE, "release", name)
|
||
return [
|
||
base + ext,
|
||
base + ".dll",
|
||
base + ".so",
|
||
base + ".dylib",
|
||
]
|
||
|
||
|
||
def load_dll(engine: str = "c"):
|
||
"""加载指定引擎。
|
||
|
||
- C/C++/Fortran: 返回 ctypes.CDLL 对象
|
||
- Python: 返回模块对象（直接 import，无需编译）
|
||
|
||
Args:
|
||
engine: "c", "cpp", "fortran", 或 "python"
|
||
Raises:
|
||
FileNotFoundError: DLL/模块文件不存在
|
||
"""
|
||
if _ENGINE_DIR.get(engine, engine) == "python":
|
||
import importlib.util, sys as _sys
|
||
mod_path = os.path.join(_HERE, "python", "dynamics_lib.py")
|
||
if not os.path.exists(mod_path):
|
||
raise FileNotFoundError(f"Python 引擎未找到: {mod_path}")
|
||
spec = importlib.util.spec_from_file_location(
|
||
"engines.python.dynamics_lib", mod_path)
|
||
mod = importlib.util.module_from_spec(spec)
|
||
spec.loader.exec_module(mod)
|
||
return mod # 返回模块，不是 CDLL
|
||
|
||
for p in _dll_candidates(engine):
|
||
if os.path.exists(p):
|
||
lib = ctypes.CDLL(p)
|
||
_setup_prototype(lib)
|
||
return lib
|
||
raise FileNotFoundError(
|
||
f"DLL 未找到（引擎 {engine}），候选路径：\n" +
|
||
"\n".join(f" {p}" for p in _dll_candidates(engine)) +
|
||
f"\n请先编译：cd engines/{engine} && make dll"
|
||
)
|
||
|
||
|
||
def _setup_prototype(lib: ctypes.CDLL) -> None:
|
||
"""配置 run_dynamics 的参数类型和返回类型。"""
|
||
c_dbl_p = ctypes.POINTER(ctypes.c_double)
|
||
c_int_p = ctypes.POINTER(ctypes.c_int)
|
||
cb_type = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_int)
|
||
|
||
lib.run_dynamics.restype = ctypes.c_int
|
||
lib.run_dynamics.argtypes = [
|
||
ctypes.c_int, # n_atoms
|
||
c_dbl_p, # pos_init [n_atoms*3]
|
||
c_dbl_p, # vel_init [n_atoms*3]
|
||
c_dbl_p, # masses [n_atoms]
|
||
c_int_p, # fixed [n_atoms*3]
|
||
ctypes.c_int, # n_bonds
|
||
c_int_p, # bond_pairs [n_bonds*2]
|
||
c_dbl_p, # bond_k [n_bonds]
|
||
c_dbl_p, # bond_r0 [n_bonds]
|
||
ctypes.c_double, # box_a
|
||
ctypes.c_double, # dt
|
||
ctypes.c_int, # NT
|
||
ctypes.c_int, # NSTEP
|
||
ctypes.c_int, # warmup_steps
|
||
ctypes.c_int, # method_id
|
||
ctypes.c_double, # Gx
|
||
ctypes.c_double, # Gy
|
||
ctypes.c_double, # Gz
|
||
ctypes.c_double, # Bx
|
||
ctypes.c_double, # By
|
||
ctypes.c_double, # Bz
|
||
ctypes.c_int, # gravity_field
|
||
ctypes.c_int, # elastic_force
|
||
ctypes.c_int, # damping_force
|
||
ctypes.c_double, # gravity_strength
|
||
ctypes.c_int, # n_drivers
|
||
c_int_p, # drv_idx [n_drivers]
|
||
c_dbl_p, # drv_amp [n_drivers*3]
|
||
c_dbl_p, # drv_freq [n_drivers*3]
|
||
c_dbl_p, # drv_phi [n_drivers*3]
|
||
c_dbl_p, # drv_eq [n_drivers*3]
|
||
c_dbl_p, # drv_ncycles [n_drivers]
|
||
c_int_p, # drv_has_period [n_drivers]
|
||
ctypes.c_int, # n_frames
|
||
c_dbl_p, # out_x
|
||
c_dbl_p, # out_y
|
||
c_dbl_p, # out_z
|
||
c_dbl_p, # out_vx
|
||
c_dbl_p, # out_vy
|
||
c_dbl_p, # out_vz
|
||
cb_type, # progress_cb (可为 NULL)
|
||
]
|
||
|
||
|
||
def _c_dbl(arr: np.ndarray):
|
||
"""返回 float64 C 连续数组的 ctypes 指针。"""
|
||
a = np.ascontiguousarray(arr, dtype=np.float64)
|
||
return a.ctypes.data_as(ctypes.POINTER(ctypes.c_double)), a
|
||
|
||
|
||
def _c_int(arr: np.ndarray):
|
||
"""返回 int32 C 连续数组的 ctypes 指针。"""
|
||
a = np.ascontiguousarray(arr, dtype=np.int32)
|
||
return a.ctypes.data_as(ctypes.POINTER(ctypes.c_int)), a
|
||
|
||
|
||
def _is_python_module(lib) -> bool:
|
||
"""判断 lib 是否为 Python 引擎模块（而非 ctypes.CDLL）。"""
|
||
return not isinstance(lib, ctypes.CDLL)
|
||
|
||
|
||
def _run_dynamics_python(lib, config, atom_positions, atom_velocities, atom_masses,
|
||
atom_fixed, bond_pairs, bond_stiffness, bond_rest_lengths,
|
||
driver_data, atom_ids, progress_cb=None) -> dict:
|
||
"""调用 Python 引擎的 run_dynamics()，参数/返回值格式与 ctypes 版相同。"""
|
||
n = len(atom_masses)
|
||
NT = int(config["NT"])
|
||
NSTEP = int(config.get("NSTEP", 1))
|
||
warmup = int(config.get("warmup_steps", 0))
|
||
dt = float(config["DT"])
|
||
box_a = float(config.get("box_a", 300.0))
|
||
method_str = str(config.get("method", "leapfrog")).lower().replace(" ", "_")
|
||
method_id = _METHOD_ID.get(method_str, 3)
|
||
|
||
G = config.get("G", [0.0, 0.0, 0.0])
|
||
B = config.get("B", [0.0, 0.0, 0.0])
|
||
if hasattr(G, "tolist"): G = G.tolist()
|
||
if hasattr(B, "tolist"): B = B.tolist()
|
||
|
||
gravity_field = int(config.get("gravity_field", 0))
|
||
elastic_force = int(config.get("elastic_force", 1))
|
||
damping_force = int(config.get("damping_force", 0))
|
||
gravity_strength = float(config.get("gravity_strength", 1.0))
|
||
|
||
record_steps = NT - warmup
|
||
n_frames = max(1, record_steps // NSTEP)
|
||
|
||
nd = len(driver_data) if driver_data else 0
|
||
if nd > 0:
|
||
drv_idx = np.array([d["local_idx"] for d in driver_data], dtype=np.int64)
|
||
drv_amp = np.array([d["amp"] for d in driver_data], dtype=np.float64)
|
||
drv_freq = np.array([d["freq"] for d in driver_data], dtype=np.float64)
|
||
drv_phi = np.array([d["phi"] for d in driver_data], dtype=np.float64)
|
||
drv_eq = np.array([d["eq_pos"] for d in driver_data], dtype=np.float64)
|
||
drv_nc = np.array([d["n_cycles"] for d in driver_data], dtype=np.float64)
|
||
drv_hp = np.array([d["has_period"]for d in driver_data], dtype=np.int32)
|
||
else:
|
||
drv_idx = drv_amp = drv_freq = drv_phi = drv_eq = drv_nc = drv_hp = \
|
||
np.zeros(0, dtype=np.int64)
|
||
|
||
out_x, out_y, out_z, out_vx, out_vy, out_vz = lib.run_dynamics(
|
||
n_atoms=n,
|
||
pos_init=atom_positions,
|
||
vel_init=atom_velocities,
|
||
masses=atom_masses,
|
||
fixed=atom_fixed,
|
||
n_bonds=len(bond_pairs),
|
||
bond_pairs=bond_pairs,
|
||
bond_k=bond_stiffness,
|
||
bond_r0=bond_rest_lengths,
|
||
box_a=box_a,
|
||
dt=dt,
|
||
NT=NT,
|
||
NSTEP=NSTEP,
|
||
warmup_steps=warmup,
|
||
method_id=method_id,
|
||
Gx=float(G[0]), Gy=float(G[1]), Gz=float(G[2]),
|
||
Bx=float(B[0]), By=float(B[1]), Bz=float(B[2]),
|
||
gravity_field=gravity_field,
|
||
elastic_force=elastic_force,
|
||
damping_force=damping_force,
|
||
gravity_strength=gravity_strength,
|
||
n_drivers=nd,
|
||
drv_idx=drv_idx,
|
||
drv_amp=drv_amp,
|
||
drv_freq=drv_freq,
|
||
drv_phi=drv_phi,
|
||
drv_eq=drv_eq,
|
||
drv_ncycles=drv_nc,
|
||
drv_has_period=drv_hp,
|
||
n_frames=n_frames,
|
||
progress_cb=progress_cb,
|
||
)
|
||
|
||
shape = (n_frames, n)
|
||
t_arr = np.arange(n_frames) * NSTEP * dt + warmup * dt
|
||
return {
|
||
"x": out_x.reshape(shape), "y": out_y.reshape(shape),
|
||
"z": out_z.reshape(shape), "vx": out_vx.reshape(shape),
|
||
"vy": out_vy.reshape(shape), "vz": out_vz.reshape(shape),
|
||
"t": t_arr,
|
||
}
|
||
|
||
|
||
def run_dynamics_dll(
|
||
lib,
|
||
config: dict,
|
||
atom_positions: np.ndarray, # (n_atoms, 3)
|
||
atom_velocities: np.ndarray, # (n_atoms, 3)
|
||
atom_masses: np.ndarray, # (n_atoms,)
|
||
atom_fixed: np.ndarray, # (n_atoms, 3) int, 1=固定
|
||
bond_pairs: np.ndarray, # (n_bonds, 2) int 0-based 局部索引
|
||
bond_stiffness: np.ndarray, # (n_bonds,)
|
||
bond_rest_lengths: np.ndarray,# (n_bonds,)
|
||
driver_data: list, # 驱动原子列表（见下文）
|
||
atom_ids: np.ndarray, # (n_atoms,) 全局 atom id（用于驱动原子查找）
|
||
progress_cb=None,
|
||
) -> dict:
|
||
"""调用 DLL 的 run_dynamics()，返回抽帧后的轨迹数组。
|
||
|
||
driver_data 格式（每个元素对应一个驱动原子）：
|
||
{
|
||
"atom_id": int, # 全局 atom id
|
||
"local_idx": int, # 在 atom_ids 数组中的位置（0-based）
|
||
"amp": [ax, ay, az],
|
||
"freq": [fx, fy, fz],
|
||
"phi": [px, py, pz],
|
||
"eq_pos": [ex, ey, ez],
|
||
"n_cycles": float, # 0=不限
|
||
"has_period": int, # 0/1
|
||
}
|
||
|
||
返回：
|
||
{
|
||
"x": np.ndarray (n_frames, n_atoms),
|
||
"y": ...,
|
||
"z": ...,
|
||
"vx": ..., "vy": ..., "vz": ...,
|
||
"t": np.ndarray (n_frames,), # 时间轴
|
||
}
|
||
"""
|
||
# Python 引擎：直接调用模块函数，不走 ctypes
|
||
if _is_python_module(lib):
|
||
return _run_dynamics_python(
|
||
lib, config, atom_positions, atom_velocities, atom_masses,
|
||
atom_fixed, bond_pairs, bond_stiffness, bond_rest_lengths,
|
||
driver_data, atom_ids, progress_cb)
|
||
|
||
n = len(atom_masses)
|
||
NT = int(config["NT"])
|
||
NSTEP = int(config.get("NSTEP", 1))
|
||
warmup = int(config.get("warmup_steps", 0))
|
||
dt = float(config["DT"])
|
||
box_a = float(config.get("box_a", 300.0))
|
||
method_str = str(config.get("method", "leapfrog")).lower().replace(" ", "_")
|
||
method_id = _METHOD_ID.get(method_str, 3)
|
||
|
||
G = config.get("G", [0.0, 0.0, 0.0])
|
||
B = config.get("B", [0.0, 0.0, 0.0])
|
||
if hasattr(G, "tolist"): G = G.tolist()
|
||
if hasattr(B, "tolist"): B = B.tolist()
|
||
|
||
gravity_field = int(config.get("gravity_field", 0))
|
||
elastic_force = int(config.get("elastic_force", 1))
|
||
damping_force = int(config.get("damping_force", 0))
|
||
gravity_strength = float(config.get("gravity_strength", 1.0))
|
||
|
||
# ── 计算帧数 ──────────────────────────────────────────────
|
||
record_steps = NT - warmup
|
||
n_frames = max(1, record_steps // NSTEP)
|
||
|
||
# ── 驱动原子数据 ──────────────────────────────────────────
|
||
nd = len(driver_data) if driver_data else 0
|
||
if nd > 0:
|
||
drv_idx_arr = np.array([d["local_idx"] for d in driver_data], dtype=np.int32)
|
||
drv_amp_arr = np.array([d["amp"] for d in driver_data], dtype=np.float64).ravel()
|
||
drv_freq_arr = np.array([d["freq"] for d in driver_data], dtype=np.float64).ravel()
|
||
drv_phi_arr = np.array([d["phi"] for d in driver_data], dtype=np.float64).ravel()
|
||
drv_eq_arr = np.array([d["eq_pos"] for d in driver_data], dtype=np.float64).ravel()
|
||
drv_nc_arr = np.array([d["n_cycles"] for d in driver_data], dtype=np.float64)
|
||
drv_hp_arr = np.array([d["has_period"] for d in driver_data], dtype=np.int32)
|
||
else:
|
||
drv_idx_arr = np.zeros(1, dtype=np.int32)
|
||
drv_amp_arr = np.zeros(3, dtype=np.float64)
|
||
drv_freq_arr = np.zeros(3, dtype=np.float64)
|
||
drv_phi_arr = np.zeros(3, dtype=np.float64)
|
||
drv_eq_arr = np.zeros(3, dtype=np.float64)
|
||
drv_nc_arr = np.zeros(1, dtype=np.float64)
|
||
drv_hp_arr = np.zeros(1, dtype=np.int32)
|
||
|
||
# ── 输出缓冲区 ────────────────────────────────────────────
|
||
out_x = np.zeros(n_frames * n, dtype=np.float64)
|
||
out_y = np.zeros(n_frames * n, dtype=np.float64)
|
||
out_z = np.zeros(n_frames * n, dtype=np.float64)
|
||
out_vx = np.zeros(n_frames * n, dtype=np.float64)
|
||
out_vy = np.zeros(n_frames * n, dtype=np.float64)
|
||
out_vz = np.zeros(n_frames * n, dtype=np.float64)
|
||
|
||
# ── ctypes 指针（保留 arr 引用防止 GC） ──────────────────
|
||
p_pos, _pos = _c_dbl(atom_positions.ravel())
|
||
p_vel, _vel = _c_dbl(atom_velocities.ravel())
|
||
p_mass, _mass = _c_dbl(atom_masses)
|
||
p_fixed, _fixed = _c_int(atom_fixed.ravel())
|
||
p_bp, _bp = _c_int(bond_pairs.ravel() if len(bond_pairs) else np.zeros(2, dtype=np.int32))
|
||
p_bk, _bk = _c_dbl(bond_stiffness if len(bond_stiffness) else np.zeros(1))
|
||
p_br0, _br0 = _c_dbl(bond_rest_lengths if len(bond_rest_lengths) else np.zeros(1))
|
||
p_didx, _didx = _c_int(drv_idx_arr)
|
||
p_damp, _damp = _c_dbl(drv_amp_arr)
|
||
p_dfrq, _dfrq = _c_dbl(drv_freq_arr)
|
||
p_dphi, _dphi = _c_dbl(drv_phi_arr)
|
||
p_deq, _deq = _c_dbl(drv_eq_arr)
|
||
p_dnc, _dnc = _c_dbl(drv_nc_arr)
|
||
p_dhp, _dhp = _c_int(drv_hp_arr)
|
||
|
||
p_ox = out_x.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
|
||
p_oy = out_y.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
|
||
p_oz = out_z.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
|
||
p_ovx = out_vx.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
|
||
p_ovy = out_vy.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
|
||
p_ovz = out_vz.ctypes.data_as(ctypes.POINTER(ctypes.c_double))
|
||
|
||
# 进度回调
|
||
cb_type = ctypes.CFUNCTYPE(None, ctypes.c_int, ctypes.c_int)
|
||
if progress_cb is not None:
|
||
cb = cb_type(progress_cb)
|
||
else:
|
||
cb = ctypes.cast(None, cb_type)
|
||
|
||
ret = lib.run_dynamics(
|
||
n,
|
||
p_pos, p_vel, p_mass, p_fixed,
|
||
len(bond_pairs), p_bp, p_bk, p_br0,
|
||
box_a, dt, NT, NSTEP, warmup, method_id,
|
||
float(G[0]), float(G[1]), float(G[2]),
|
||
float(B[0]), float(B[1]), float(B[2]),
|
||
gravity_field, elastic_force, damping_force, gravity_strength,
|
||
nd, p_didx, p_damp, p_dfrq, p_dphi, p_deq, p_dnc, p_dhp,
|
||
n_frames,
|
||
p_ox, p_oy, p_oz, p_ovx, p_ovy, p_ovz,
|
||
cb,
|
||
)
|
||
|
||
if ret != 0:
|
||
raise RuntimeError(f"run_dynamics() returned error code {ret}")
|
||
|
||
shape = (n_frames, n)
|
||
t_arr = np.arange(n_frames) * NSTEP * dt + warmup * dt
|
||
|
||
return {
|
||
"x": out_x.reshape(shape),
|
||
"y": out_y.reshape(shape),
|
||
"z": out_z.reshape(shape),
|
||
"vx": out_vx.reshape(shape),
|
||
"vy": out_vy.reshape(shape),
|
||
"vz": out_vz.reshape(shape),
|
||
"t": t_arr,
|
||
}
|
||
|
||
|
||
def is_dll_available(engine: str = "c") -> bool:
|
||
"""检查指定引擎是否可用（DLL 已编译或 Python 模块存在）。"""
|
||
if _ENGINE_DIR.get(engine, engine) == "python":
|
||
return os.path.exists(os.path.join(_HERE, "python", "dynamics_lib.py"))
|
||
return any(os.path.exists(p) for p in _dll_candidates(engine))
|