boltz-vsynthes 1.0.0__py3-none-any.whl

boltz/model/modules/utils.py

@@ -0,0 +1,303 @@
+# started from code from https://github.com/lucidrains/alphafold3-pytorch, MIT License, Copyright (c) 2024 Phil Wang
+
+from functools import partial
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch.nn import (
+    Linear,
+    Module,
+)
+from torch.types import Device
+
+LinearNoBias = partial(Linear, bias=False)
+
+
+def exists(v):
+    return v is not None
+
+
+def default(v, d):
+    return v if exists(v) else d
+
+
+def log(t, eps=1e-20):
+    return torch.log(t.clamp(min=eps))
+
+
+class SwiGLU(Module):
+    def forward(
+        self,
+        x,  #: Float['... d']
+    ):  # -> Float[' ... (d//2)']:
+        x, gates = x.chunk(2, dim=-1)
+        return F.silu(gates) * x
+
+
+def center(atom_coords, atom_mask):
+    atom_mean = torch.sum(
+        atom_coords * atom_mask[:, :, None], dim=1, keepdim=True
+    ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True)
+    atom_coords = atom_coords - atom_mean
+    return atom_coords
+
+
+def compute_random_augmentation(
+    multiplicity, s_trans=1.0, device=None, dtype=torch.float32
+):
+    R = random_rotations(multiplicity, dtype=dtype, device=device)
+    random_trans = (
+        torch.randn((multiplicity, 1, 3), dtype=dtype, device=device) * s_trans
+    )
+    return R, random_trans
+
+
+def randomly_rotate(coords, return_second_coords=False, second_coords=None):
+    R = random_rotations(len(coords), coords.dtype, coords.device)
+
+    if return_second_coords:
+        return torch.einsum("bmd,bds->bms", coords, R), torch.einsum(
+            "bmd,bds->bms", second_coords, R
+        ) if second_coords is not None else None
+
+    return torch.einsum("bmd,bds->bms", coords, R)
+
+
+def center_random_augmentation(
+    atom_coords,
+    atom_mask,
+    s_trans=1.0,
+    augmentation=True,
+    centering=True,
+    return_second_coords=False,
+    second_coords=None,
+):
+    """Algorithm 19"""
+    if centering:
+        atom_mean = torch.sum(
+            atom_coords * atom_mask[:, :, None], dim=1, keepdim=True
+        ) / torch.sum(atom_mask[:, :, None], dim=1, keepdim=True)
+        atom_coords = atom_coords - atom_mean
+
+        if second_coords is not None:
+            # apply same transformation also to this input
+            second_coords = second_coords - atom_mean
+
+    if augmentation:
+        atom_coords, second_coords = randomly_rotate(
+            atom_coords, return_second_coords=True, second_coords=second_coords
+        )
+        random_trans = torch.randn_like(atom_coords[:, 0:1, :]) * s_trans
+        atom_coords = atom_coords + random_trans
+
+        if second_coords is not None:
+            second_coords = second_coords + random_trans
+
+    if return_second_coords:
+        return atom_coords, second_coords
+
+    return atom_coords
+
+
+class ExponentialMovingAverage:
+    """from https://github.com/yang-song/score_sde_pytorch/blob/main/models/ema.py, Apache-2.0 license
+    Maintains (exponential) moving average of a set of parameters."""
+
+    def __init__(self, parameters, decay, use_num_updates=True):
+        """
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; usually the result of
+            `model.parameters()`.
+          decay: The exponential decay.
+          use_num_updates: Whether to use number of updates when computing
+            averages.
+        """
+        if decay < 0.0 or decay > 1.0:
+            raise ValueError("Decay must be between 0 and 1")
+        self.decay = decay
+        self.num_updates = 0 if use_num_updates else None
+        self.shadow_params = [p.clone().detach() for p in parameters if p.requires_grad]
+        self.collected_params = []
+
+    def update(self, parameters):
+        """
+        Update currently maintained parameters.
+        Call this every time the parameters are updated, such as the result of
+        the `optimizer.step()` call.
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; usually the same set of
+            parameters used to initialize this object.
+        """
+        decay = self.decay
+        if self.num_updates is not None:
+            self.num_updates += 1
+            decay = min(decay, (1 + self.num_updates) / (10 + self.num_updates))
+        one_minus_decay = 1.0 - decay
+        with torch.no_grad():
+            parameters = [p for p in parameters if p.requires_grad]
+            for s_param, param in zip(self.shadow_params, parameters):
+                s_param.sub_(one_minus_decay * (s_param - param))
+
+    def compatible(self, parameters):
+        if len(self.shadow_params) != len(parameters):
+            print(
+                f"Model has {len(self.shadow_params)} parameter tensors, the incoming ema {len(parameters)}"
+            )
+            return False
+
+        for s_param, param in zip(self.shadow_params, parameters):
+            if param.data.shape != s_param.data.shape:
+                print(
+                    f"Model has parameter tensor of shape {s_param.data.shape} , the incoming ema {param.data.shape}"
+                )
+                return False
+        return True
+
+    def copy_to(self, parameters):
+        """
+        Copy current parameters into given collection of parameters.
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+            updated with the stored moving averages.
+        """
+        parameters = [p for p in parameters if p.requires_grad]
+        for s_param, param in zip(self.shadow_params, parameters):
+            if param.requires_grad:
+                param.data.copy_(s_param.data)
+
+    def store(self, parameters):
+        """
+        Save the current parameters for restoring later.
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+            temporarily stored.
+        """
+        self.collected_params = [param.clone() for param in parameters]
+
+    def restore(self, parameters):
+        """
+        Restore the parameters stored with the `store` method.
+        Useful to validate the model with EMA parameters without affecting the
+        original optimization process. Store the parameters before the
+        `copy_to` method. After validation (or model saving), use this to
+        restore the former parameters.
+        Args:
+          parameters: Iterable of `torch.nn.Parameter`; the parameters to be
+            updated with the stored parameters.
+        """
+        for c_param, param in zip(self.collected_params, parameters):
+            param.data.copy_(c_param.data)
+
+    def state_dict(self):
+        return dict(
+            decay=self.decay,
+            num_updates=self.num_updates,
+            shadow_params=self.shadow_params,
+        )
+
+    def load_state_dict(self, state_dict, device):
+        self.decay = state_dict["decay"]
+        self.num_updates = state_dict["num_updates"]
+        self.shadow_params = [
+            tensor.to(device) for tensor in state_dict["shadow_params"]
+        ]
+
+    def to(self, device):
+        self.shadow_params = [tensor.to(device) for tensor in self.shadow_params]
+
+
+# the following is copied from Torch3D, BSD License, Copyright (c) Meta Platforms, Inc. and affiliates.
+
+
+def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    # pyre-fixme[58]: `/` is not supported for operand types `float` and `Tensor`.
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def random_quaternions(
+    n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
+) -> torch.Tensor:
+    """
+    Generate random quaternions representing rotations,
+    i.e. versors with nonnegative real part.
+
+    Args:
+        n: Number of quaternions in a batch to return.
+        dtype: Type to return.
+        device: Desired device of returned tensor. Default:
+            uses the current device for the default tensor type.
+
+    Returns:
+        Quaternions as tensor of shape (N, 4).
+    """
+    if isinstance(device, str):
+        device = torch.device(device)
+    o = torch.randn((n, 4), dtype=dtype, device=device)
+    s = (o * o).sum(1)
+    o = o / _copysign(torch.sqrt(s), o[:, 0])[:, None]
+    return o
+
+
+def random_rotations(
+    n: int, dtype: Optional[torch.dtype] = None, device: Optional[Device] = None
+) -> torch.Tensor:
+    """
+    Generate random rotations as 3x3 rotation matrices.
+
+    Args:
+        n: Number of rotation matrices in a batch to return.
+        dtype: Type to return.
+        device: Device of returned tensor. Default: if None,
+            uses the current device for the default tensor type.
+
+    Returns:
+        Rotation matrices as tensor of shape (n, 3, 3).
+    """
+    quaternions = random_quaternions(n, dtype=dtype, device=device)
+    return quaternion_to_matrix(quaternions)

boltz/model/optim/ema.py

@@ -0,0 +1,389 @@
+# --------------------------------------------------------------------------------------
+# Modified from Bio-Diffusion (https://github.com/BioinfoMachineLearning/bio-diffusion):
+# Modified from : https://github.com/BioinfoMachineLearning/bio-diffusion/blob/main/src/utils/__init__.py
+# --------------------------------------------------------------------------------------
+
+from typing import Any, Optional
+
+import torch
+from pytorch_lightning import Callback, LightningModule, Trainer
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.types import STEP_OUTPUT
+
+
+class EMA(Callback):
+    """Implements Exponential Moving Averaging (EMA).
+
+    When training a model, this callback maintains moving averages
+    of the trained parameters. When evaluating, we use the moving
+    averages copy of the trained parameters. When saving, we save
+    an additional set of parameters with the prefix `ema`.
+
+    Adapted from:
+    https://github.com/NVIDIA/NeMo/blob/main/nemo/collections/common/callbacks/ema.py
+    https://github.com/BioinfoMachineLearning/bio-diffusion/blob/main/src/utils/__init__.py
+
+    """
+
+    def __init__(
+        self,
+        decay: float = 0.999,
+        apply_ema_every_n_steps: int = 1,
+        start_step: int = 0,
+        eval_with_ema: bool = True,
+        warm_start: bool = True,
+    ) -> None:
+        """Initialize the EMA callback.
+
+        Parameters
+        ----------
+        decay: float
+            The exponential decay, has to be between 0-1.
+        apply_ema_every_n_steps: int, optional (default=1)
+            Apply EMA every n global steps.
+        start_step: int, optional (default=0)
+            Start applying EMA from ``start_step`` global step onwards.
+        eval_with_ema: bool, optional (default=True)
+            Validate the EMA weights instead of the original weights.
+            Note this means that when saving the model, the
+            validation metrics are calculated with the EMA weights.
+
+        """
+        if not (0 <= decay <= 1):
+            msg = "EMA decay value must be between 0 and 1"
+            raise MisconfigurationException(msg)
+
+        self._ema_weights: Optional[dict[str, torch.Tensor]] = None
+        self._cur_step: Optional[int] = None
+        self._weights_buffer: Optional[dict[str, torch.Tensor]] = None
+        self.apply_ema_every_n_steps = apply_ema_every_n_steps
+        self.start_step = start_step
+        self.eval_with_ema = eval_with_ema
+        self.decay = decay
+        self.warm_start = warm_start
+
+    @property
+    def ema_initialized(self) -> bool:
+        """Check if EMA weights have been initialized.
+
+        Returns
+        -------
+        bool
+            Whether the EMA weights have been initialized.
+
+        """
+        return self._ema_weights is not None
+
+    def state_dict(self) -> dict[str, Any]:
+        """Return the current state of the callback.
+
+        Returns
+        -------
+        dict[str, Any]
+            The current state of the callback.
+
+        """
+        return {
+            "cur_step": self._cur_step,
+            "ema_weights": self._ema_weights,
+        }
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """Load the state of the callback.
+
+        Parameters
+        ----------
+        state_dict: dict[str, Any]
+            The state of the callback to load.
+
+        """
+        self._cur_step = state_dict["cur_step"]
+        self._ema_weights = state_dict["ema_weights"]
+
+    def should_apply_ema(self, step: int) -> bool:
+        """Check if EMA should be applied at the current step.
+
+        Parameters
+        ----------
+        step: int
+            The current global step.
+
+        Returns
+        -------
+        bool
+            True if EMA should be applied, False otherwise.
+
+        """
+        return (
+            step != self._cur_step
+            and step >= self.start_step
+            and step % self.apply_ema_every_n_steps == 0
+        )
+
+    def apply_ema(self, pl_module: LightningModule) -> None:
+        """Apply EMA to the model weights.
+
+        Parameters
+        ----------
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        decay = self.decay
+        if self.warm_start:
+            decay = min(decay, (1 + self._cur_step) / (10 + self._cur_step))
+
+        for k, orig_weight in pl_module.state_dict().items():
+            ema_weight = self._ema_weights[k]
+            if (
+                ema_weight.data.dtype != torch.long  # noqa: PLR1714
+                and orig_weight.data.dtype != torch.long  # skip non-trainable weights
+            ):
+                diff = ema_weight.data - orig_weight.data
+                diff.mul_(1.0 - decay)
+                ema_weight.sub_(diff)
+
+    def on_load_checkpoint(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,  # noqa: ARG002
+        checkpoint: dict[str, Any],
+    ) -> None:
+        """Load the EMA weights from the checkpoint.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+        checkpoint: dict[str, Any]
+            The checkpoint to load.
+
+        """
+        if "ema" in checkpoint:
+            print("LOADING CHECKPOINT RUNNING")
+            self.load_state_dict(checkpoint["ema"])
+
+    def on_save_checkpoint(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,  # noqa: ARG002
+        checkpoint: dict[str, Any],
+    ) -> None:
+        """Save the EMA weights to the checkpoint.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+        checkpoint: dict[str, Any]
+            The checkpoint to save.
+
+        """
+        if self.ema_initialized:
+            checkpoint["ema"] = self.state_dict()
+
+    def on_train_start(self, trainer: Trainer, pl_module: LightningModule) -> None:  # noqa: ARG002
+        """Initialize EMA weights and move to device.
+
+        Parameters
+        ----------
+        trainer: pl.Trainer
+            The Trainer instance.
+        pl_module: pl.LightningModule
+            The LightningModule instance.
+
+        """
+        # Create EMA weights if not already initialized
+        if not self.ema_initialized:
+            self._ema_weights = {
+                k: p.detach().clone() for k, p in pl_module.state_dict().items()
+            }
+
+        # Move EMA weights to the correct device
+        self._ema_weights = {
+            k: p.to(pl_module.device) for k, p in self._ema_weights.items()
+        }
+
+    def on_train_batch_end(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        outputs: STEP_OUTPUT,  # noqa: ARG002
+        batch: Any,  # noqa: ARG002
+        batch_idx: int,  # noqa: ARG002
+    ) -> None:
+        """Apply EMA to the model weights at the end of each training batch.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+        outputs: STEP_OUTPUT
+            The outputs of the model.
+        batch: Any
+            The current batch.
+        batch_idx: int
+            The index of the current batch.
+
+        """
+        if self.should_apply_ema(trainer.global_step):
+            self._cur_step = trainer.global_step
+            self.apply_ema(pl_module)
+
+    def replace_model_weights(self, pl_module: LightningModule) -> None:
+        """Replace model weights with EMA weights.
+
+        Parameters
+        ----------
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._weights_buffer = {
+            k: p.detach().clone().to("cpu") for k, p in pl_module.state_dict().items()
+        }
+        pl_module.load_state_dict(self._ema_weights, strict=False)
+
+    def restore_original_weights(self, pl_module: LightningModule) -> None:
+        """Restore model weights to original weights.
+
+        Parameters
+        ----------
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        pl_module.load_state_dict(self._weights_buffer, strict=False)
+        del self._weights_buffer
+
+    def _on_eval_start(self, pl_module: LightningModule) -> None:
+        """Use EMA weights for evaluation.
+
+        Parameters
+        ----------
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        if self.ema_initialized and self.eval_with_ema:
+            self.replace_model_weights(pl_module)
+
+    def _on_eval_end(self, pl_module: LightningModule) -> None:
+        """Restore original weights after evaluation.
+
+        Parameters
+        ----------
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        if self.ema_initialized and self.eval_with_ema:
+            self.restore_original_weights(pl_module)
+
+    def on_validation_start(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        """Use EMA weights for validation.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._on_eval_start(pl_module)
+
+    def on_validation_end(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        """Restore original weights after validation.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._on_eval_end(pl_module)
+
+    def on_test_start(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        """Use EMA weights for testing.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._on_eval_start(pl_module)
+
+    def on_test_end(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        """Restore original weights after testing.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._on_eval_end(pl_module)
+
+    def on_predict_start(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        """Use EMA weights for prediction.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._on_eval_start(pl_module)
+
+    def on_predict_end(
+        self,
+        trainer: Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        """Restore original weights after prediction.
+
+        Parameters
+        ----------
+        trainer: Trainer
+            The Trainer instance.
+        pl_module: LightningModule
+            The LightningModule instance.
+
+        """
+        self._on_eval_end(pl_module)