|
|
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
|
|
|
|
