boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/model/optim/scheduler.py

@@ -0,0 +1,99 @@
+import torch
+
+
+class AlphaFoldLRScheduler(torch.optim.lr_scheduler._LRScheduler):
+    """Implements the learning rate schedule defined AF3.
+
+    A linear warmup is followed by a plateau at the maximum
+    learning rate and then exponential decay. Note that the
+    initial learning rate of the optimizer in question is
+    ignored; use this class' base_lr parameter to specify
+    the starting point of the warmup.
+
+    """
+
+    def __init__(
+        self,
+        optimizer: torch.optim.Optimizer,
+        last_epoch: int = -1,
+        base_lr: float = 0.0,
+        max_lr: float = 1.8e-3,
+        warmup_no_steps: int = 1000,
+        start_decay_after_n_steps: int = 50000,
+        decay_every_n_steps: int = 50000,
+        decay_factor: float = 0.95,
+    ) -> None:
+        """Initialize the learning rate scheduler.
+
+        Parameters
+        ----------
+        optimizer : torch.optim.Optimizer
+            The optimizer.
+        last_epoch : int, optional
+            The last epoch, by default -1
+        base_lr : float, optional
+            The base learning rate, by default 0.0
+        max_lr : float, optional
+            The maximum learning rate, by default 1.8e-3
+        warmup_no_steps : int, optional
+            The number of warmup steps, by default 1000
+        start_decay_after_n_steps : int, optional
+            The number of steps after which to start decay, by default 50000
+        decay_every_n_steps : int, optional
+            The number of steps after which to decay, by default 50000
+        decay_factor : float, optional
+            The decay factor, by default 0.95
+
+        """
+        step_counts = {
+            "warmup_no_steps": warmup_no_steps,
+            "start_decay_after_n_steps": start_decay_after_n_steps,
+        }
+
+        for k, v in step_counts.items():
+            if v < 0:
+                msg = f"{k} must be nonnegative"
+                raise ValueError(msg)
+
+        if warmup_no_steps > start_decay_after_n_steps:
+            msg = "warmup_no_steps must not exceed start_decay_after_n_steps"
+            raise ValueError(msg)
+
+        self.optimizer = optimizer
+        self.last_epoch = last_epoch
+        self.base_lr = base_lr
+        self.max_lr = max_lr
+        self.warmup_no_steps = warmup_no_steps
+        self.start_decay_after_n_steps = start_decay_after_n_steps
+        self.decay_every_n_steps = decay_every_n_steps
+        self.decay_factor = decay_factor
+
+        super().__init__(optimizer, last_epoch=last_epoch)
+
+    def state_dict(self) -> dict:
+        state_dict = {k: v for k, v in self.__dict__.items() if k not in ["optimizer"]}
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        self.__dict__.update(state_dict)
+
+    def get_lr(self):
+        if not self._get_lr_called_within_step:
+            msg = (
+                "To get the last learning rate computed by the scheduler, use "
+                "get_last_lr()"
+            )
+            raise RuntimeError(msg)
+
+        step_no = self.last_epoch
+
+        if step_no <= self.warmup_no_steps:
+            lr = self.base_lr + (step_no / self.warmup_no_steps) * self.max_lr
+        elif step_no > self.start_decay_after_n_steps:
+            steps_since_decay = step_no - self.start_decay_after_n_steps
+            exp = (steps_since_decay // self.decay_every_n_steps) + 1
+            lr = self.max_lr * (self.decay_factor**exp)
+        else:  # plateau
+            lr = self.max_lr
+
+        return [lr for group in self.optimizer.param_groups]

boltz/model/potentials/potentials.py

@@ -0,0 +1,497 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Optional, Union
+
+import torch
+
+from boltz.data import const
+from boltz.model.potentials.schedules import (
+    ExponentialInterpolation,
+    ParameterSchedule,
+    PiecewiseStepFunction,
+)
+
+
+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 = 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",
+            )
+        value = self.compute_variable(coords, index, compute_gradient=False)
+        energy = self.compute_function(value, *args)
+        return energy.sum(dim=-1)
+
+    def compute_gradient(self, coords, feats, parameters):
+        index, args, com_args = self.compute_args(feats, parameters)
+        if com_args is not None:
+            com_index, atom_pad_mask = com_args
+        else:
+            com_index, atom_pad_mask = None, None
+
+        if index.shape[1] == 0:
+            return torch.zeros_like(coords)
+
+        if com_index is not None:
+            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])
+
+        value, grad_value = self.compute_variable(coords, index, compute_gradient=True)
+        energy, dEnergy = self.compute_function(value, *args, compute_derivative=True)
+
+        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)),
+            dEnergy.tile(grad_value.shape[-3]).unsqueeze(-1)
+            * grad_value.flatten(start_dim=-3, end_dim=-2),
+            "sum",
+        )
+
+        if com_index is not None:
+            grad_atom = grad_atom[..., com_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, 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
+
+
+class FlatBottomPotential(Potential):
+    def compute_function(
+        self, value, k, lower_bounds, upper_bounds, 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"))
+
+        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 DistancePotential(Potential):
+    def compute_variable(self, coords, index, 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, 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, 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["angle_buffer"], parameters["angle_buffer"]
+        )
+        upper_bounds[bond_mask * angle_mask] *= 1.0 + min(
+            parameters["angle_buffer"], parameters["angle_buffer"]
+        )
+        lower_bounds[~bond_mask * ~angle_mask] *= 1.0 - parameters["clash_buffer"]
+        upper_bounds[~bond_mask * ~angle_mask] = float("inf")
+
+        k = torch.ones_like(lower_bounds)
+
+        return pair_index, (k, lower_bounds, upper_bounds), 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
+
+
+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
+
+
+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),
+        )
+
+
+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
+
+
+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
+
+
+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
+
+
+def get_potentials():
+    potentials = [
+        SymmetricChainCOMPotential(
+            parameters={
+                "guidance_interval": 4,
+                "guidance_weight": 0.5,
+                "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]
+                ),
+                "resampling_weight": PiecewiseStepFunction(
+                    thresholds=[0.6], values=[0.01, 0.0]
+                ),
+                "buffer": 0.225,
+            }
+        ),
+        ConnectionsPotential(
+            parameters={
+                "guidance_interval": 1,
+                "guidance_weight": 0.15,
+                "resampling_weight": 1.0,
+                "buffer": 2.0,
+            }
+        ),
+        PoseBustersPotential(
+            parameters={
+                "guidance_interval": 1,
+                "guidance_weight": 0.05,
+                "resampling_weight": 0.1,
+                "bond_buffer": 0.20,
+                "angle_buffer": 0.20,
+                "clash_buffer": 0.15,
+            }
+        ),
+        ChiralAtomPotential(
+            parameters={
+                "guidance_interval": 1,
+                "guidance_weight": 0.10,
+                "resampling_weight": 1.0,
+                "buffer": 0.52360,
+            }
+        ),
+        StereoBondPotential(
+            parameters={
+                "guidance_interval": 1,
+                "guidance_weight": 0.05,
+                "resampling_weight": 1.0,
+                "buffer": 0.52360,
+            }
+        ),
+        PlanarBondPotential(
+            parameters={
+                "guidance_interval": 1,
+                "guidance_weight": 0.05,
+                "resampling_weight": 1.0,
+                "buffer": 0.26180,
+            }
+        ),
+    ]
+    return potentials
+
+
+@dataclass
+class GuidanceConfig:
+    """Guidance configuration."""
+
+    potentials: Optional[list[Potential]] = None
+    guidance_update: Optional[bool] = None
+    num_guidance_gd_steps: Optional[int] = None
+    guidance_gd_step_size: Optional[int] = None
+    fk_steering: Optional[bool] = None
+    fk_resampling_interval: Optional[int] = 1
+    fk_lambda: Optional[float] = 1.0
+    fk_method: Optional[str] = None
+    fk_batch_size: Optional[int] = 2

boltz/model/potentials/schedules.py

@@ -0,0 +1,32 @@
+import math
+from abc import ABC
+
+class ParameterSchedule(ABC):
+    def compute(self, t):
+        raise NotImplementedError
+
+class ExponentialInterpolation(ParameterSchedule):
+    def __init__(self, start, end, alpha):
+        self.start = start
+        self.end = end
+        self.alpha = alpha
+
+    def compute(self, t):
+        if self.alpha != 0:
+            return self.start + (self.end - self.start) * (math.exp(self.alpha * t) - 1) / (math.exp(self.alpha) - 1)
+        else:
+            return self.start + (self.end - self.start) * t
+
+class PiecewiseStepFunction(ParameterSchedule):
+    def __init__(self, thresholds, values):
+        self.thresholds = thresholds
+        self.values = values
+
+    def compute(self, t):
+        assert len(self.thresholds) > 0
+        assert len(self.values) == len(self.thresholds) + 1
+
+        idx = 0
+        while idx < len(self.thresholds) and t > self.thresholds[idx]:
+            idx += 1
+        return self.values[idx]