| import sys |
| import torch |
| from tqdm import tqdm |
| from torch.distributions import Normal |
|
|
| from .dynamics import dynamics |
| from .utils.utils import kabsch |
|
|
|
|
| class MDs: |
| def __init__(self, args): |
| self.device = args.device |
| self.molecule = args.molecule |
| self.end_state = args.end_state |
| self.num_samples = args.num_samples |
| self.start_state = args.start_state |
| self.get_md_info(args) |
| self.mds = self._init_mds(args) |
| |
| self.log_prob = Normal(0, self.std).log_prob |
| self.target_position = self.target_position - self.target_position[:, self.heavy_atoms].mean(-2, keepdim=True) |
| R, t = kabsch( |
| self.start_position[:, self.heavy_atoms], |
| self.target_position[:, self.heavy_atoms], |
| ) |
| self.start_position = torch.matmul(self.start_position, R.transpose(-2, -1)) + t |
|
|
| def get_md_info(self, args): |
| md = getattr(dynamics, self.molecule.title())(args, self.end_state) |
| self.num_particles = md.num_particles |
| self.heavy_atoms = torch.from_numpy(md.heavy_atoms).to(self.device) |
| self.energy_function = md.energy_function |
| self.target_position = torch.tensor( |
| md.position, dtype=torch.float, device=self.device |
| ).unsqueeze(0) |
| self.std = torch.tensor( |
| md.std, |
| dtype=torch.float, |
| device=args.device, |
| ) |
| self.m = torch.tensor( |
| md.m, |
| dtype=torch.float, |
| device=args.device, |
| ).unsqueeze(-1) |
|
|
| def _init_mds(self, args): |
| mds = [] |
| for _ in tqdm(range(self.num_samples)): |
| md = getattr(dynamics, self.molecule.title())(args, self.start_state) |
| mds.append(md) |
| self.start_position = torch.tensor( |
| md.position, dtype=torch.float, device=self.device |
| ).unsqueeze(0) |
| return mds |
|
|
| def step(self, force): |
| force = force.detach().cpu().numpy() |
| for i in range(self.num_samples): |
| self.mds[i].step(force[i]) |
|
|
| def report(self): |
| positions, forces = [], [] |
| for i in range(self.num_samples): |
| position, force = self.mds[i].report() |
| positions.append(position) |
| forces.append(force) |
| positions = torch.tensor(positions, dtype=torch.float, device=self.device) |
| forces = torch.tensor(forces, dtype=torch.float, device=self.device) |
| return positions, forces |
|
|
| def reset(self): |
| for i in range(self.num_samples): |
| self.mds[i].reset() |
|
|
| def set_temperature(self, temperature): |
| for i in range(self.num_samples): |
| self.mds[i].set_temperature(temperature) |
|
|
