| from abc import ABC, abstractmethod |
| from typing import Optional, Dict, Any, Set, List, Union |
|
|
| import torch |
| import numpy as np |
| from . import vb_const as const |
| from .vb_potentials_schedules import ( |
| ParameterSchedule, |
| ExponentialInterpolation, |
| PiecewiseStepFunction, |
| ) |
| from .vb_loss_diffusionv2 import weighted_rigid_align |
|
|
|
|
| class Potential(ABC): |
| def __init__( |
| self, |
| parameters: Optional[ |
| Dict[str, Union[ParameterSchedule, float, int, bool]] |
| ] = None, |
| ): |
| self.parameters = parameters |
|
|
| def compute(self, coords, feats, parameters): |
| index, args, com_args, ref_args, operator_args = self.compute_args( |
| feats, parameters |
| ) |
|
|
| if index.shape[1] == 0: |
| return torch.zeros(coords.shape[:-2], device=coords.device) |
|
|
| if com_args is not None: |
| com_index, atom_pad_mask = com_args |
| unpad_com_index = com_index[atom_pad_mask] |
| unpad_coords = coords[..., atom_pad_mask, :] |
| coords = torch.zeros( |
| (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3), |
| device=coords.device, |
| ).scatter_reduce( |
| -2, |
| unpad_com_index.unsqueeze(-1).expand_as(unpad_coords), |
| unpad_coords, |
| "mean", |
| ) |
| else: |
| com_index, atom_pad_mask = None, None |
|
|
| if ref_args is not None: |
| ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args |
| coords = coords[..., ref_atom_index, :] |
| else: |
| ref_coords, ref_mask, ref_atom_index, ref_token_index = ( |
| None, |
| None, |
| None, |
| None, |
| ) |
|
|
| if operator_args is not None: |
| negation_mask, union_index = operator_args |
| else: |
| negation_mask, union_index = None, None |
|
|
| value = self.compute_variable( |
| coords, |
| index, |
| ref_coords=ref_coords, |
| ref_mask=ref_mask, |
| compute_gradient=False, |
| ) |
| energy = self.compute_function( |
| value, *args, negation_mask=negation_mask, compute_derivative=False |
| ) |
|
|
| if union_index is not None: |
| neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy) |
| Z = torch.zeros( |
| (*energy.shape[:-1], union_index.max() + 1), device=union_index.device |
| ).scatter_reduce( |
| -1, |
| union_index.expand_as(neg_exp_energy), |
| neg_exp_energy, |
| "sum", |
| ) |
| softmax_energy = neg_exp_energy / Z[..., union_index] |
| softmax_energy[Z[..., union_index] == 0] = 0 |
| return (energy * softmax_energy).sum(dim=-1) |
|
|
| return energy.sum(dim=tuple(range(1, energy.dim()))) |
|
|
| def compute_gradient(self, coords, feats, parameters): |
| index, args, com_args, ref_args, operator_args = self.compute_args( |
| feats, parameters |
| ) |
| if index.shape[1] == 0: |
| return torch.zeros_like(coords) |
|
|
| if com_args is not None: |
| com_index, atom_pad_mask = com_args |
| unpad_coords = coords[..., atom_pad_mask, :] |
| unpad_com_index = com_index[atom_pad_mask] |
| coords = torch.zeros( |
| (*unpad_coords.shape[:-2], unpad_com_index.max() + 1, 3), |
| device=coords.device, |
| ).scatter_reduce( |
| -2, |
| unpad_com_index.unsqueeze(-1).expand_as(unpad_coords), |
| unpad_coords, |
| "mean", |
| ) |
| com_counts = torch.bincount(com_index[atom_pad_mask]) |
| else: |
| com_index, atom_pad_mask = None, None |
|
|
| if ref_args is not None: |
| ref_coords, ref_mask, ref_atom_index, ref_token_index = ref_args |
| coords = coords[..., ref_atom_index, :] |
| else: |
| ref_coords, ref_mask, ref_atom_index, ref_token_index = ( |
| None, |
| None, |
| None, |
| None, |
| ) |
|
|
| if operator_args is not None: |
| negation_mask, union_index = operator_args |
| else: |
| negation_mask, union_index = None, None |
|
|
| value, grad_value = self.compute_variable( |
| coords, |
| index, |
| ref_coords=ref_coords, |
| ref_mask=ref_mask, |
| compute_gradient=True, |
| ) |
| energy, dEnergy = self.compute_function( |
| value, |
| *args, negation_mask=negation_mask, compute_derivative=True |
| ) |
| if union_index is not None: |
| neg_exp_energy = torch.exp(-1 * parameters["union_lambda"] * energy) |
| Z = torch.zeros( |
| (*energy.shape[:-1], union_index.max() + 1), device=union_index.device |
| ).scatter_reduce( |
| -1, |
| union_index.expand_as(energy), |
| neg_exp_energy, |
| "sum", |
| ) |
| softmax_energy = neg_exp_energy / Z[..., union_index] |
| softmax_energy[Z[..., union_index] == 0] = 0 |
| f = torch.zeros( |
| (*energy.shape[:-1], union_index.max() + 1), device=union_index.device |
| ).scatter_reduce( |
| -1, |
| union_index.expand_as(energy), |
| energy * softmax_energy, |
| "sum", |
| ) |
| dSoftmax = ( |
| dEnergy |
| * softmax_energy |
| * (1 + parameters["union_lambda"] * (energy - f[..., union_index])) |
| ) |
| prod = dSoftmax.tile(grad_value.shape[-3]).unsqueeze( |
| -1 |
| ) * grad_value.flatten(start_dim=-3, end_dim=-2) |
| if prod.dim() > 3: |
| prod = prod.sum(dim=list(range(1, prod.dim() - 2))) |
| grad_atom = torch.zeros_like(coords).scatter_reduce( |
| -2, |
| index.flatten(start_dim=0, end_dim=1) |
| .unsqueeze(-1) |
| .expand((*coords.shape[:-2], -1, 3)), |
| prod, |
| "sum", |
| ) |
| else: |
| prod = dEnergy.tile(grad_value.shape[-3]).unsqueeze( |
| -1 |
| ) * grad_value.flatten(start_dim=-3, end_dim=-2) |
| if prod.dim() > 3: |
| prod = prod.sum(dim=list(range(1, prod.dim() - 2))) |
| grad_atom = torch.zeros_like(coords).scatter_reduce( |
| -2, |
| index.flatten(start_dim=0, end_dim=1) |
| .unsqueeze(-1) |
| .expand((*coords.shape[:-2], -1, 3)), |
| prod, |
| "sum", |
| ) |
|
|
| if com_index is not None: |
| grad_atom = grad_atom[..., com_index, :] |
| elif ref_token_index is not None: |
| grad_atom = grad_atom[..., ref_token_index, :] |
|
|
| return grad_atom |
|
|
| def compute_parameters(self, t): |
| if self.parameters is None: |
| return None |
| parameters = { |
| name: parameter |
| if not isinstance(parameter, ParameterSchedule) |
| else parameter.compute(t) |
| for name, parameter in self.parameters.items() |
| } |
| return parameters |
|
|
| @abstractmethod |
| def compute_function( |
| self, value, *args, negation_mask=None, compute_derivative=False |
| ): |
| raise NotImplementedError |
|
|
| @abstractmethod |
| def compute_variable(self, coords, index, compute_gradient=False): |
| raise NotImplementedError |
|
|
| @abstractmethod |
| def compute_args(self, t, feats, **parameters): |
| raise NotImplementedError |
|
|
| def get_reference_coords(self, feats, parameters): |
| return None, None |
|
|
|
|
| class FlatBottomPotential(Potential): |
| def compute_function( |
| self, |
| value, |
| k, |
| lower_bounds, |
| upper_bounds, |
| negation_mask=None, |
| compute_derivative=False, |
| ): |
| if lower_bounds is None: |
| lower_bounds = torch.full_like(value, float("-inf")) |
| if upper_bounds is None: |
| upper_bounds = torch.full_like(value, float("inf")) |
| lower_bounds = lower_bounds.expand_as(value).clone() |
| upper_bounds = upper_bounds.expand_as(value).clone() |
|
|
| if negation_mask is not None: |
| unbounded_below_mask = torch.isneginf(lower_bounds) |
| unbounded_above_mask = torch.isposinf(upper_bounds) |
| unbounded_mask = unbounded_below_mask + unbounded_above_mask |
| assert torch.all(unbounded_mask + negation_mask) |
| lower_bounds[~unbounded_above_mask * ~negation_mask] = upper_bounds[ |
| ~unbounded_above_mask * ~negation_mask |
| ] |
| upper_bounds[~unbounded_above_mask * ~negation_mask] = float("inf") |
| upper_bounds[~unbounded_below_mask * ~negation_mask] = lower_bounds[ |
| ~unbounded_below_mask * ~negation_mask |
| ] |
| lower_bounds[~unbounded_below_mask * ~negation_mask] = float("-inf") |
|
|
| neg_overflow_mask = value < lower_bounds |
| pos_overflow_mask = value > upper_bounds |
|
|
| energy = torch.zeros_like(value) |
| energy[neg_overflow_mask] = (k * (lower_bounds - value))[neg_overflow_mask] |
| energy[pos_overflow_mask] = (k * (value - upper_bounds))[pos_overflow_mask] |
| if not compute_derivative: |
| return energy |
|
|
| dEnergy = torch.zeros_like(value) |
| dEnergy[neg_overflow_mask] = ( |
| -1 * k.expand_as(neg_overflow_mask)[neg_overflow_mask] |
| ) |
| dEnergy[pos_overflow_mask] = ( |
| 1 * k.expand_as(pos_overflow_mask)[pos_overflow_mask] |
| ) |
|
|
| return energy, dEnergy |
|
|
|
|
| class ReferencePotential(Potential): |
| def compute_variable( |
| self, coords, index, ref_coords, ref_mask, compute_gradient=False |
| ): |
| aligned_ref_coords = weighted_rigid_align( |
| ref_coords.float(), |
| coords[:, index].float(), |
| ref_mask, |
| ref_mask, |
| ) |
|
|
| r = coords[:, index] - aligned_ref_coords |
| r_norm = torch.linalg.norm(r, dim=-1) |
|
|
| if not compute_gradient: |
| return r_norm |
|
|
| r_hat = r / r_norm.unsqueeze(-1) |
| grad = (r_hat * ref_mask.unsqueeze(-1)).unsqueeze(1) |
| return r_norm, grad |
|
|
|
|
| class DistancePotential(Potential): |
| def compute_variable( |
| self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False |
| ): |
| r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1]) |
| r_ij_norm = torch.linalg.norm(r_ij, dim=-1) |
| r_hat_ij = r_ij / r_ij_norm.unsqueeze(-1) |
|
|
| if not compute_gradient: |
| return r_ij_norm |
|
|
| grad_i = r_hat_ij |
| grad_j = -1 * r_hat_ij |
| grad = torch.stack((grad_i, grad_j), dim=1) |
| return r_ij_norm, grad |
|
|
|
|
| class DihedralPotential(Potential): |
| def compute_variable( |
| self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False |
| ): |
| r_ij = coords.index_select(-2, index[0]) - coords.index_select(-2, index[1]) |
| r_kj = coords.index_select(-2, index[2]) - coords.index_select(-2, index[1]) |
| r_kl = coords.index_select(-2, index[2]) - coords.index_select(-2, index[3]) |
|
|
| n_ijk = torch.cross(r_ij, r_kj, dim=-1) |
| n_jkl = torch.cross(r_kj, r_kl, dim=-1) |
|
|
| r_kj_norm = torch.linalg.norm(r_kj, dim=-1) |
| n_ijk_norm = torch.linalg.norm(n_ijk, dim=-1) |
| n_jkl_norm = torch.linalg.norm(n_jkl, dim=-1) |
|
|
| sign_phi = torch.sign( |
| r_kj.unsqueeze(-2) @ torch.cross(n_ijk, n_jkl, dim=-1).unsqueeze(-1) |
| ).squeeze(-1, -2) |
| phi = sign_phi * torch.arccos( |
| torch.clamp( |
| (n_ijk.unsqueeze(-2) @ n_jkl.unsqueeze(-1)).squeeze(-1, -2) |
| / (n_ijk_norm * n_jkl_norm), |
| -1 + 1e-8, |
| 1 - 1e-8, |
| ) |
| ) |
|
|
| if not compute_gradient: |
| return phi |
|
|
| a = ( |
| (r_ij.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2) |
| ).unsqueeze(-1) |
| b = ( |
| (r_kl.unsqueeze(-2) @ r_kj.unsqueeze(-1)).squeeze(-1, -2) / (r_kj_norm**2) |
| ).unsqueeze(-1) |
|
|
| grad_i = n_ijk * (r_kj_norm / n_ijk_norm**2).unsqueeze(-1) |
| grad_l = -1 * n_jkl * (r_kj_norm / n_jkl_norm**2).unsqueeze(-1) |
| grad_j = (a - 1) * grad_i - b * grad_l |
| grad_k = (b - 1) * grad_l - a * grad_i |
| grad = torch.stack((grad_i, grad_j, grad_k, grad_l), dim=1) |
| return phi, grad |
|
|
|
|
| class AbsDihedralPotential(DihedralPotential): |
| def compute_variable( |
| self, coords, index, ref_coords=None, ref_mask=None, compute_gradient=False |
| ): |
| if not compute_gradient: |
| phi = super().compute_variable( |
| coords, index, compute_gradient=compute_gradient |
| ) |
| phi = torch.abs(phi) |
| return phi |
|
|
| phi, grad = super().compute_variable( |
| coords, index, compute_gradient=compute_gradient |
| ) |
| grad[(phi < 0)[..., None, :, None].expand_as(grad)] *= -1 |
| phi = torch.abs(phi) |
|
|
| return phi, grad |
|
|
|
|
| class PoseBustersPotential(FlatBottomPotential, DistancePotential): |
| def compute_args(self, feats, parameters): |
| pair_index = feats["rdkit_bounds_index"][0] |
| lower_bounds = feats["rdkit_lower_bounds"][0].clone() |
| upper_bounds = feats["rdkit_upper_bounds"][0].clone() |
| bond_mask = feats["rdkit_bounds_bond_mask"][0] |
| angle_mask = feats["rdkit_bounds_angle_mask"][0] |
|
|
| lower_bounds[bond_mask * ~angle_mask] *= 1.0 - parameters["bond_buffer"] |
| upper_bounds[bond_mask * ~angle_mask] *= 1.0 + parameters["bond_buffer"] |
| lower_bounds[~bond_mask * angle_mask] *= 1.0 - parameters["angle_buffer"] |
| upper_bounds[~bond_mask * angle_mask] *= 1.0 + parameters["angle_buffer"] |
| lower_bounds[bond_mask * angle_mask] *= 1.0 - min( |
| parameters["bond_buffer"], parameters["angle_buffer"] |
| ) |
| upper_bounds[bond_mask * angle_mask] *= 1.0 + min( |
| parameters["bond_buffer"], parameters["angle_buffer"] |
| ) |
| lower_bounds[~bond_mask * ~angle_mask] *= 1.0 - parameters["clash_buffer"] |
| upper_bounds[~bond_mask * ~angle_mask] = float("inf") |
|
|
| vdw_radii = torch.zeros( |
| const.num_elements, dtype=torch.float32, device=pair_index.device |
| ) |
| vdw_radii[1:119] = torch.tensor( |
| const.vdw_radii, dtype=torch.float32, device=pair_index.device |
| ) |
| atom_vdw_radii = ( |
| feats["ref_element"].float() @ vdw_radii.unsqueeze(-1) |
| ).squeeze(-1)[0] |
| bond_cutoffs = 0.35 + atom_vdw_radii[pair_index].mean(dim=0) |
| lower_bounds[~bond_mask] = torch.max(lower_bounds[~bond_mask], bond_cutoffs[~bond_mask]) |
| upper_bounds[bond_mask] = torch.min(upper_bounds[bond_mask], bond_cutoffs[bond_mask]) |
|
|
| k = torch.ones_like(lower_bounds) |
|
|
| return pair_index, (k, lower_bounds, upper_bounds), None, None, None |
|
|
|
|
| class ConnectionsPotential(FlatBottomPotential, DistancePotential): |
| def compute_args(self, feats, parameters): |
| pair_index = feats["connected_atom_index"][0] |
| lower_bounds = None |
| upper_bounds = torch.full( |
| (pair_index.shape[1],), parameters["buffer"], device=pair_index.device |
| ) |
| k = torch.ones_like(upper_bounds) |
|
|
| return pair_index, (k, lower_bounds, upper_bounds), None, None, None |
|
|
|
|
| class VDWOverlapPotential(FlatBottomPotential, DistancePotential): |
| def compute_args(self, feats, parameters): |
| atom_chain_id = ( |
| torch.bmm( |
| feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float() |
| ) |
| .squeeze(-1) |
| .long() |
| )[0] |
| atom_pad_mask = feats["atom_pad_mask"][0].bool() |
| chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask]) |
| single_ion_mask = (chain_sizes > 1)[atom_chain_id] |
|
|
| vdw_radii = torch.zeros( |
| const.num_elements, dtype=torch.float32, device=atom_chain_id.device |
| ) |
| vdw_radii[1:119] = torch.tensor( |
| const.vdw_radii, dtype=torch.float32, device=atom_chain_id.device |
| ) |
| atom_vdw_radii = ( |
| feats["ref_element"].float() @ vdw_radii.unsqueeze(-1) |
| ).squeeze(-1)[0] |
|
|
| pair_index = torch.triu_indices( |
| atom_chain_id.shape[0], |
| atom_chain_id.shape[0], |
| 1, |
| device=atom_chain_id.device, |
| ) |
|
|
| pair_pad_mask = atom_pad_mask[pair_index].all(dim=0) |
| pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]] |
|
|
| num_chains = atom_chain_id.max() + 1 |
| connected_chain_index = feats["connected_chain_index"][0] |
| connected_chain_matrix = torch.eye( |
| num_chains, device=atom_chain_id.device, dtype=torch.bool |
| ) |
| connected_chain_matrix[connected_chain_index[0], connected_chain_index[1]] = ( |
| True |
| ) |
| connected_chain_matrix[connected_chain_index[1], connected_chain_index[0]] = ( |
| True |
| ) |
| connected_chain_mask = connected_chain_matrix[ |
| atom_chain_id[pair_index[0]], atom_chain_id[pair_index[1]] |
| ] |
|
|
| pair_index = pair_index[ |
| :, pair_pad_mask * pair_ion_mask * ~connected_chain_mask |
| ] |
|
|
| lower_bounds = atom_vdw_radii[pair_index].sum(dim=0) * ( |
| 1.0 - parameters["buffer"] |
| ) |
| upper_bounds = None |
| k = torch.ones_like(lower_bounds) |
|
|
| return pair_index, (k, lower_bounds, upper_bounds), None, None, None |
|
|
|
|
| class SymmetricChainCOMPotential(FlatBottomPotential, DistancePotential): |
| def compute_args(self, feats, parameters): |
| atom_chain_id = ( |
| torch.bmm( |
| feats["atom_to_token"].float(), feats["asym_id"].unsqueeze(-1).float() |
| ) |
| .squeeze(-1) |
| .long() |
| )[0] |
| atom_pad_mask = feats["atom_pad_mask"][0].bool() |
| chain_sizes = torch.bincount(atom_chain_id[atom_pad_mask]) |
| single_ion_mask = chain_sizes > 1 |
|
|
| pair_index = feats["symmetric_chain_index"][0] |
| pair_ion_mask = single_ion_mask[pair_index[0]] * single_ion_mask[pair_index[1]] |
| pair_index = pair_index[:, pair_ion_mask] |
| lower_bounds = torch.full( |
| (pair_index.shape[1],), |
| parameters["buffer"], |
| dtype=torch.float32, |
| device=pair_index.device, |
| ) |
| upper_bounds = None |
| k = torch.ones_like(lower_bounds) |
|
|
| return ( |
| pair_index, |
| (k, lower_bounds, upper_bounds), |
| (atom_chain_id, atom_pad_mask), |
| None, |
| None, |
| ) |
|
|
|
|
| class StereoBondPotential(FlatBottomPotential, AbsDihedralPotential): |
| def compute_args(self, feats, parameters): |
| stereo_bond_index = feats["stereo_bond_index"][0] |
| stereo_bond_orientations = feats["stereo_bond_orientations"][0].bool() |
|
|
| lower_bounds = torch.zeros( |
| stereo_bond_orientations.shape, device=stereo_bond_orientations.device |
| ) |
| upper_bounds = torch.zeros( |
| stereo_bond_orientations.shape, device=stereo_bond_orientations.device |
| ) |
| lower_bounds[stereo_bond_orientations] = torch.pi - parameters["buffer"] |
| upper_bounds[stereo_bond_orientations] = float("inf") |
| lower_bounds[~stereo_bond_orientations] = float("-inf") |
| upper_bounds[~stereo_bond_orientations] = parameters["buffer"] |
|
|
| k = torch.ones_like(lower_bounds) |
|
|
| return stereo_bond_index, (k, lower_bounds, upper_bounds), None, None, None |
|
|
|
|
| class ChiralAtomPotential(FlatBottomPotential, DihedralPotential): |
| def compute_args(self, feats, parameters): |
| chiral_atom_index = feats["chiral_atom_index"][0] |
| chiral_atom_orientations = feats["chiral_atom_orientations"][0].bool() |
|
|
| lower_bounds = torch.zeros( |
| chiral_atom_orientations.shape, device=chiral_atom_orientations.device |
| ) |
| upper_bounds = torch.zeros( |
| chiral_atom_orientations.shape, device=chiral_atom_orientations.device |
| ) |
| lower_bounds[chiral_atom_orientations] = parameters["buffer"] |
| upper_bounds[chiral_atom_orientations] = float("inf") |
| upper_bounds[~chiral_atom_orientations] = -1 * parameters["buffer"] |
| lower_bounds[~chiral_atom_orientations] = float("-inf") |
|
|
| k = torch.ones_like(lower_bounds) |
| return chiral_atom_index, (k, lower_bounds, upper_bounds), None, None, None |
|
|
|
|
| class PlanarBondPotential(FlatBottomPotential, AbsDihedralPotential): |
| def compute_args(self, feats, parameters): |
| double_bond_index = feats["planar_bond_index"][0].T |
| double_bond_improper_index = torch.tensor( |
| [ |
| [1, 2, 3, 0], |
| [4, 5, 0, 3], |
| ], |
| device=double_bond_index.device, |
| ).T |
| improper_index = ( |
| double_bond_index[:, double_bond_improper_index] |
| .swapaxes(0, 1) |
| .flatten(start_dim=1) |
| ) |
| lower_bounds = None |
| upper_bounds = torch.full( |
| (improper_index.shape[1],), |
| parameters["buffer"], |
| device=improper_index.device, |
| ) |
| k = torch.ones_like(upper_bounds) |
|
|
| return improper_index, (k, lower_bounds, upper_bounds), None, None, None |
|
|
|
|
| class TemplateReferencePotential(FlatBottomPotential, ReferencePotential): |
| def compute_args(self, feats, parameters): |
| if "template_mask_cb" not in feats or "template_force" not in feats: |
| return torch.empty([1, 0]), None, None, None, None |
|
|
| template_mask = feats["template_mask_cb"][feats["template_force"]] |
| if template_mask.shape[0] == 0: |
| return torch.empty([1, 0]), None, None, None, None |
|
|
| ref_coords = feats["template_cb"][feats["template_force"]].clone() |
| ref_mask = feats["template_mask_cb"][feats["template_force"]].clone() |
| ref_atom_index = ( |
| torch.bmm( |
| feats["token_to_rep_atom"].float(), |
| torch.arange( |
| feats["atom_pad_mask"].shape[1], |
| device=feats["atom_pad_mask"].device, |
| dtype=torch.float32, |
| )[None, :, None], |
| ) |
| .squeeze(-1) |
| .long() |
| )[0] |
| ref_token_index = ( |
| torch.bmm( |
| feats["atom_to_token"].float(), |
| feats["token_index"].unsqueeze(-1).float(), |
| ) |
| .squeeze(-1) |
| .long() |
| )[0] |
|
|
| index = torch.arange( |
| template_mask.shape[-1], dtype=torch.long, device=template_mask.device |
| )[None] |
| upper_bounds = torch.full( |
| template_mask.shape, float("inf"), device=index.device, dtype=torch.float32 |
| ) |
| ref_idxs = torch.argwhere(template_mask).T |
| upper_bounds[ref_idxs.unbind()] = feats["template_force_threshold"][ |
| feats["template_force"] |
| ][ref_idxs[0]] |
|
|
| lower_bounds = None |
| k = torch.ones_like(upper_bounds) |
| return ( |
| index, |
| (k, lower_bounds, upper_bounds), |
| None, |
| (ref_coords, ref_mask, ref_atom_index, ref_token_index), |
| None, |
| ) |
|
|
|
|
| class ContactPotentital(FlatBottomPotential, DistancePotential): |
| def compute_args(self, feats, parameters): |
| index = feats["contact_pair_index"][0] |
| union_index = feats["contact_union_index"][0] |
| negation_mask = feats["contact_negation_mask"][0] |
| lower_bounds = None |
| upper_bounds = feats["contact_thresholds"][0].clone() |
| k = torch.ones_like(upper_bounds) |
| return ( |
| index, |
| (k, lower_bounds, upper_bounds), |
| None, |
| None, |
| (negation_mask, union_index), |
| ) |
|
|
|
|
| def get_potentials(steering_args, boltz2=False): |
| potentials = [] |
| if steering_args["fk_steering"] or steering_args["physical_guidance_update"]: |
| potentials.extend( |
| [ |
| SymmetricChainCOMPotential( |
| parameters={ |
| "guidance_interval": 4, |
| "guidance_weight": 0.5 |
| if steering_args["physical_guidance_update"] |
| else 0.0, |
| "resampling_weight": 0.5, |
| "buffer": ExponentialInterpolation( |
| start=1.0, end=5.0, alpha=-2.0 |
| ), |
| } |
| ), |
| VDWOverlapPotential( |
| parameters={ |
| "guidance_interval": 5, |
| "guidance_weight": ( |
| PiecewiseStepFunction(thresholds=[0.4], values=[0.125, 0.0]) |
| if steering_args["physical_guidance_update"] |
| else 0.0 |
| ), |
| "resampling_weight": PiecewiseStepFunction( |
| thresholds=[0.6], values=[0.01, 0.0] |
| ), |
| "buffer": 0.225, |
| } |
| ), |
| ConnectionsPotential( |
| parameters={ |
| "guidance_interval": 1, |
| "guidance_weight": 0.15 |
| if steering_args["physical_guidance_update"] |
| else 0.0, |
| "resampling_weight": 1.0, |
| "buffer": 2.0, |
| } |
| ), |
| PoseBustersPotential( |
| parameters={ |
| "guidance_interval": 1, |
| "guidance_weight": 0.01 |
| if steering_args["physical_guidance_update"] |
| else 0.0, |
| "resampling_weight": 0.1, |
| "bond_buffer": 0.125, |
| "angle_buffer": 0.125, |
| "clash_buffer": 0.10, |
| } |
| ), |
| ChiralAtomPotential( |
| parameters={ |
| "guidance_interval": 1, |
| "guidance_weight": 0.1 |
| if steering_args["physical_guidance_update"] |
| else 0.0, |
| "resampling_weight": 1.0, |
| "buffer": 0.52360, |
| } |
| ), |
| StereoBondPotential( |
| parameters={ |
| "guidance_interval": 1, |
| "guidance_weight": 0.05 |
| if steering_args["physical_guidance_update"] |
| else 0.0, |
| "resampling_weight": 1.0, |
| "buffer": 0.52360, |
| } |
| ), |
| PlanarBondPotential( |
| parameters={ |
| "guidance_interval": 1, |
| "guidance_weight": 0.05 |
| if steering_args["physical_guidance_update"] |
| else 0.0, |
| "resampling_weight": 1.0, |
| "buffer": 0.26180, |
| } |
| ), |
| ] |
| ) |
| if boltz2 and ( |
| steering_args["fk_steering"] or steering_args["contact_guidance_update"] |
| ): |
| potentials.extend( |
| [ |
| ContactPotentital( |
| parameters={ |
| "guidance_interval": 4, |
| "guidance_weight": ( |
| PiecewiseStepFunction( |
| thresholds=[0.25, 0.75], values=[0.0, 0.5, 1.0] |
| ) |
| if steering_args["contact_guidance_update"] |
| else 0.0 |
| ), |
| "resampling_weight": 1.0, |
| "union_lambda": ExponentialInterpolation( |
| start=8.0, end=0.0, alpha=-2.0 |
| ), |
| } |
| ), |
| TemplateReferencePotential( |
| parameters={ |
| "guidance_interval": 2, |
| "guidance_weight": 0.1 |
| if steering_args["contact_guidance_update"] |
| else 0.0, |
| "resampling_weight": 1.0, |
| } |
| ), |
| ] |
| ) |
| return potentials |
|
|
