|
|
| import torch |
| import numpy as np |
|
|
| from .utils import kabsch, aldp_diff, tic_diff |
|
|
|
|
| class Metric: |
| def __init__(self, args, mds): |
| self.device = args.device |
| self.molecule = args.molecule |
| self.save_dir = args.save_dir |
| self.timestep = args.timestep |
| self.friction = args.friction |
| self.num_samples = args.num_samples |
| self.m = mds.m |
| self.std = mds.std |
| self.log_prob = mds.log_prob |
| self.heavy_atoms = mds.heavy_atoms |
| self.energy_function = mds.energy_function |
| self.target_position = mds.target_position |
|
|
| def __call__(self): |
| positions, forces, potentials = [], [], [] |
| for i in range(self.num_samples): |
| position = np.load(f"{self.save_dir}/positions/{i}.npy").astype(np.float32) |
| force, potential = self.energy_function(position) |
| positions.append(torch.from_numpy(position).to(self.device)) |
| forces.append(torch.from_numpy(force).to(self.device)) |
| potentials.append(torch.from_numpy(potential).to(self.device)) |
| final_position = torch.stack([position[-1] for position in positions]) |
| rmsd, rmsd_std = self.rmsd( |
| final_position[:, self.heavy_atoms], |
| self.target_position[:, self.heavy_atoms], |
| ) |
| thp, hit = self.thp(final_position, self.target_position) |
| ets, ets_std = self.ets(hit, potentials) |
| metrics = { |
| "rmsd": 10 * rmsd, |
| "thp": 100 * thp, |
| "ets": ets, |
| "rmsd_std": 10 * rmsd_std, |
| "ets_std": ets_std, |
| } |
| return metrics |
|
|
| def rmsd(self, position, target_position): |
| R, t = kabsch(position, target_position) |
| position = torch.matmul(position, R.transpose(-2, -1)) + t |
| rmsd = (position - target_position).square().sum(-1).mean(-1).sqrt() |
| rmsd, std_rmsd = rmsd.mean().item(), rmsd.std().item() |
| return rmsd, std_rmsd |
|
|
| def thp(self, position, target_position): |
| if self.molecule == "aldp": |
| psi_diff, phi_diff = aldp_diff(position, target_position) |
| hit = psi_diff.square() + phi_diff.square() < 0.75 ** 2 |
| else: |
| tic1_diff, tic2_diff = tic_diff(self.molecule, position, target_position) |
| hit = tic1_diff.square() + tic2_diff.square() < 0.75 ** 2 |
| hit = hit.squeeze() |
| thp = hit.sum().float() / len(hit) |
| return thp.item(), hit |
|
|
| def ets(self, hit, potentials): |
| etss = [] |
| for i, hit_idx in enumerate(hit): |
| if hit_idx: |
| ets = potentials[i].max(0)[0] |
| etss.append(ets) |
| if len(etss) > 0: |
| etss = torch.tensor(etss) |
| ets, std_ets = etss.mean().item(), etss.std().item() |
| return ets, std_ets |
| else: |
| return None, None |
|
|
