1 Star 0 Fork 0

饭喵/torchsponge_pyscf

加入 Gitee
与超过 1200万 开发者一起发现、参与优秀开源项目,私有仓库也完全免费 :)
免费加入
文件
克隆/下载
metad.py 2.55 KB
一键复制 编辑 原始数据 按行查看 历史
饭喵 提交于 2024-11-04 13:06 . AIMD MetaD
import torch
from torch import Tensor
import numpy as np
import h5py
from torchsponge import Sponge
from torchsponge.potential import EnergyCell
from torchsponge.potential.pyscfcell import PySCFForceCell
from torchsponge import Protein, WithForceCell, UpdaterMD, ForceField, WithEnergyCell, RunOneStepCell, ForceFieldBase
from torchsponge.callback import WriteH5MD
from torchsponge.sampling import Metadynamics
from torchsponge.function import VelocityGenerator
from torchsponge.colvar import Torsion
from torchsponge.function import PI
mol = Protein(pdb='alad.pdb')
potential = ForceField(mol, 'AMBER.FF14SB')
phi = Torsion([4, 6, 8, 14])
psi = Torsion([6, 8, 14, 16])
min_opt = torch.optim.Adam(mol.parameters(), 1e-4)
mini = Sponge(mol, potential, min_opt, metrics={'phi': phi, 'psi': psi})
mini.run(100, 10)
class BlankEnergy(EnergyCell):
def __init__(self, name: str = 'energy', length_unit: str = 'nm', energy_unit: str = 'kj/mol', use_pbc: bool = None, **kwargs):
super().__init__(name=name, length_unit=length_unit, energy_unit=energy_unit, use_pbc=use_pbc, **kwargs)
def forward(self, coordinate: Tensor, neighbour_index: Tensor = None, neighbour_mask: Tensor = None,
neighbour_vector: Tensor = None, neighbour_distance: Tensor = None, pbc_box: Tensor = None):
b = coordinate.shape[0]
energy = torch.zeros(b, 1, dtype=torch.float32, device=coordinate.device)
return energy
metad = Metadynamics(
colvar=[phi, psi],
update_pace=10,
height=2.5,
sigma=0.05,
grid_min=-PI,
grid_max=PI,
grid_bin=360,
temperature=300,
bias_factor=100,
)
blank = BlankEnergy()
ff = ForceFieldBase(blank)
bs = WithEnergyCell(mol,ff,bias=metad)
pyscf_frc = PySCFForceCell(system=mol, pyscf_option='ks', mol_parameters={'basis': '6-31g','verbose':0}, ks_parameters={'xc':'b3lyp'}, use_pbc=False)
wf = WithForceCell(mol, pyscf_frc)
vgen = VelocityGenerator(300)
velocity = vgen(mol.shape, mol.atom_mass)
opt = UpdaterMD(
mol,
temperature = 300,
integrator='leap_frog',
thermostat='langevin',
time_step=2e-3,
velocity=velocity
)
onestep = RunOneStepCell(energy=bs, force=wf, optimizer=opt)
md = Sponge(network=onestep, write_h5md=WriteH5MD(mol, 'test.h5md', 100, write_metrics=True, write_bias=True), metrics={'phi': phi, 'psi': psi})
md.run(100,10,10) # total steps, print_interval, write_info_interval
with h5py.File('test.h5md') as f:
md_hills = np.array(f['parameters']['meta_potential']['hills0'])
md_grids = np.array(f['parameters']['meta_potential']['grids0'])
print(md_hills)
print(md_grids)
Loading...
马建仓 AI 助手
尝试更多
代码解读
代码找茬
代码优化
Python
1
https://gitee.com/strongfc/torchsponge_pyscf.git
git@gitee.com:strongfc/torchsponge_pyscf.git
strongfc
torchsponge_pyscf
torchsponge_pyscf
master

搜索帮助