dockformerpp/data/data_pipeline.py

# Copyright 2021 AlQuraishi Laboratory
# Copyright 2021 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
import time
from typing import List

import numpy as np
import torch
import ml_collections as mlc
from rdkit import Chem

from dockformerpp.data import data_transforms
from dockformerpp.data.data_transforms import get_restype_atom37_mask, get_restypes
from dockformerpp.data.protein_features import make_protein_features
from dockformerpp.data.utils import FeatureTensorDict, FeatureDict
from dockformerpp.utils import protein


def _np_filter_and_to_tensor_dict(np_example: FeatureDict, features_to_keep: List[str]) -> FeatureTensorDict:
    """Creates dict of tensors from a dict of NumPy arrays.

    Args:
        np_example: A dict of NumPy feature arrays.
        features: A list of strings of feature names to be returned in the dataset.

    Returns:
        A dictionary of features mapping feature names to features. Only the given
        features are returned, all other ones are filtered out.
    """
    # torch generates warnings if feature is already a torch Tensor
    to_tensor = lambda t: torch.tensor(t) if type(t) != torch.Tensor else t.clone().detach()
    tensor_dict = {
        k: to_tensor(v) for k, v in np_example.items() if k in features_to_keep
    }

    return tensor_dict


def _add_protein_probablistic_features(features: FeatureDict, cfg: mlc.ConfigDict, mode: str) -> FeatureDict:
    if mode == "train":
        p = torch.rand(1).item()
        use_clamped_fape_value = float(p < cfg.supervised.clamp_prob)
        features["use_clamped_fape"] = np.float32(use_clamped_fape_value)
    else:
        features["use_clamped_fape"] = np.float32(0.0)
    return features


@data_transforms.curry1
def compose(x, fs):
    for f in fs:
        x = f(x)
    return x


def _apply_protein_transforms(tensors: FeatureTensorDict) -> FeatureTensorDict:
    transforms = [
        data_transforms.cast_to_64bit_ints,
        data_transforms.squeeze_features,
        data_transforms.make_atom14_masks,
        data_transforms.make_atom14_positions,
        data_transforms.atom37_to_frames,
        data_transforms.atom37_to_torsion_angles(""),
        data_transforms.make_pseudo_beta(),
        data_transforms.get_backbone_frames,
        data_transforms.get_chi_angles,
    ]

    tensors = compose(transforms)(tensors)

    return tensors


def _apply_protein_probablistic_transforms(tensors: FeatureTensorDict, cfg: mlc.ConfigDict, mode: str) \
        -> FeatureTensorDict:
    transforms = [data_transforms.make_target_feat()]

    crop_feats = dict(cfg.common.feat)

    if cfg[mode].fixed_size:
        transforms.append(data_transforms.select_feat(list(crop_feats)))
        # TODO bshor: restore transforms for training on cropped proteins, need to handle pocket somehow
        # if so, look for random_crop_to_size and make_fixed_size in data_transforms.py

    compose(transforms)(tensors)

    return tensors


class DataPipeline:
    """Assembles input features."""
    def __init__(self, config: mlc.ConfigDict, mode: str):
        self.config = config
        self.mode = mode

        self.feature_names = config.common.unsupervised_features
        if config[mode].supervised:
            self.feature_names += config.supervised.supervised_features

    def process_pdb(self, pdb_path: str) -> FeatureTensorDict:
        """
            Assembles features for a protein in a PDB file.
        """
        with open(pdb_path, 'r') as f:
            pdb_str = f.read()

        protein_object = protein.from_pdb_string(pdb_str)
        description = os.path.splitext(os.path.basename(pdb_path))[0].upper()
        pdb_feats = make_protein_features(protein_object, description)
        pdb_feats = _add_protein_probablistic_features(pdb_feats, self.config, self.mode)

        tensor_feats = _np_filter_and_to_tensor_dict(pdb_feats, self.feature_names)

        tensor_feats = _apply_protein_transforms(tensor_feats)
        tensor_feats = _apply_protein_probablistic_transforms(tensor_feats, self.config, self.mode)

        return tensor_feats

def _prepare_recycles(feat: torch.Tensor, num_recycles: int) -> torch.Tensor:
    return feat.unsqueeze(-1).repeat(*([1] * len(feat.shape)), num_recycles)


def _fit_to_crop(target_tensor: torch.Tensor, crop_size: int, start_ind: int) -> torch.Tensor:
    if len(target_tensor.shape) == 1:
        ret = torch.zeros((crop_size, ), dtype=target_tensor.dtype)
        ret[start_ind:start_ind + target_tensor.shape[0]] = target_tensor
        return ret
    elif len(target_tensor.shape) == 2:
        ret = torch.zeros((crop_size, target_tensor.shape[-1]), dtype=target_tensor.dtype)
        ret[start_ind:start_ind + target_tensor.shape[0], :] = target_tensor
        return ret
    else:
        ret = torch.zeros((crop_size, *target_tensor.shape[1:]), dtype=target_tensor.dtype)
        ret[start_ind:start_ind + target_tensor.shape[0], ...] = target_tensor
        return ret


def parse_input_json(input_path: str, mode: str, config: mlc.ConfigDict, data_pipeline: DataPipeline,
                     data_dir: str, idx: int) -> FeatureTensorDict:
    start_load_time = time.time()
    input_data = json.load(open(input_path, "r"))
    if mode == "train" or mode == "eval":
        print("loading", input_data["pdb_id"], end=" ")

    num_recycles = config.common.max_recycling_iters + 1

    input_protein_r_feats = data_pipeline.process_pdb(pdb_path=os.path.join(data_dir, input_data["input_r_structure"]))
    input_protein_l_feats = data_pipeline.process_pdb(pdb_path=os.path.join(data_dir, input_data["input_l_structure"]))

    n_res_r = input_protein_r_feats["protein_target_feat"].shape[0]
    n_res_l = input_protein_l_feats["protein_target_feat"].shape[0]
    n_res_total = n_res_r + n_res_l
    n_affinity = 1

    # add 1 for affinity token
    crop_size = n_res_total + n_affinity
    if (mode == "train" or mode == "eval") and config.train.fixed_size:
        crop_size = config.train.crop_size

    assert crop_size >= n_res_total + n_affinity, f"crop_size: {crop_size}, n_res_r: {n_res_r}, n_res_l: {n_res_l}"

    token_mask = torch.zeros((crop_size,), dtype=torch.float32)
    token_mask[:n_res_total + n_affinity] = 1

    protein_r_mask = torch.zeros((crop_size,), dtype=torch.float32)
    protein_r_mask[:n_res_r] = 1

    protein_l_mask = torch.zeros((crop_size,), dtype=torch.float32)
    protein_l_mask[n_res_r:n_res_total] = 1

    affinity_mask = torch.zeros((crop_size,), dtype=torch.float32)
    affinity_mask[n_res_total] = 1

    structural_mask = torch.zeros((crop_size,), dtype=torch.float32)
    structural_mask[:n_res_total] = 1

    inter_pair_mask = torch.zeros((crop_size, crop_size), dtype=torch.float32)
    inter_pair_mask[:n_res_r, n_res_r:n_res_total] = 1
    inter_pair_mask[n_res_r:n_res_total, :n_res_r] = 1

    tf_dim = input_protein_r_feats["protein_target_feat"].shape[-1]

    target_feat = torch.zeros((crop_size, tf_dim + 3), dtype=torch.float32)
    target_feat[:n_res_r, :tf_dim] = input_protein_r_feats["protein_target_feat"]
    target_feat[n_res_r:n_res_total, :tf_dim] = input_protein_l_feats["protein_target_feat"]

    target_feat[:n_res_r, tf_dim] = 1  # Set "is_protein_r" flag for protein rows
    target_feat[n_res_r:n_res_total, tf_dim + 1] = 1  # Set "is_protein_l" flag for ligand rows
    target_feat[n_res_total, tf_dim + 2] = 1  # Set "is_affinity" flag for affinity row

    input_positions = torch.zeros((crop_size, 3), dtype=torch.float32)
    input_positions[:n_res_r] = input_protein_r_feats["pseudo_beta"]
    input_positions[n_res_r:n_res_total] = input_protein_l_feats["pseudo_beta"]

    distogram_mask = torch.zeros(crop_size)
    if mode == "train":
        ones_indices = torch.randperm(n_res_total)[:int(n_res_total * config.train.distogram_mask_prob)]
        # print(ones_indices)
        distogram_mask[ones_indices] = 1
        input_positions = input_positions * (1 - distogram_mask).unsqueeze(-1)
    elif mode == "predict":
        # ignore all positions where pseudo_beta is 0, 0, 0
        distogram_mask = (input_positions == 0).all(dim=-1).float()
        # print("Ignoring residues", torch.nonzero(distogram_mask).flatten())

    # Implement ligand as amino acid type 20
    aatype = torch.cat([input_protein_r_feats["aatype"], input_protein_l_feats["aatype"]], dim=0)
    residue_index = torch.cat([input_protein_r_feats["residue_index"], input_protein_l_feats["residue_index"]], dim=0)
    residx_atom37_to_atom14 = torch.cat([input_protein_r_feats["residx_atom37_to_atom14"],
                                         input_protein_l_feats["residx_atom37_to_atom14"]],
                                        dim=0)
    atom37_atom_exists = torch.cat([input_protein_r_feats["atom37_atom_exists"],
                                    input_protein_l_feats["atom37_atom_exists"]], dim=0)

    feats = {
        "token_mask": token_mask,
        "protein_r_mask": protein_r_mask,
        "protein_l_mask": protein_l_mask,
        "affinity_mask": affinity_mask,
        "structural_mask": structural_mask,
        "inter_pair_mask": inter_pair_mask,

        "target_feat": target_feat,
        "input_positions": input_positions,
        "distogram_mask": distogram_mask,
        "residue_index": _fit_to_crop(residue_index, crop_size, 0),
        "aatype": _fit_to_crop(aatype, crop_size, 0),
        "residx_atom37_to_atom14": _fit_to_crop(residx_atom37_to_atom14, crop_size, 0),
        "atom37_atom_exists": _fit_to_crop(atom37_atom_exists, crop_size, 0),
    }

    if mode == "predict":
        feats.update({
            "in_chain_residue_index_r": input_protein_r_feats["in_chain_residue_index"],
            "chain_index_r": input_protein_r_feats["chain_index"],
            "in_chain_residue_index_l": input_protein_l_feats["in_chain_residue_index"],
            "chain_index_l": input_protein_l_feats["chain_index"],
        })

    if mode == 'train' or mode == 'eval':
        gt_protein_r_feats = data_pipeline.process_pdb(pdb_path=os.path.join(data_dir, input_data["gt_r_structure"]))
        gt_protein_l_feats = data_pipeline.process_pdb(pdb_path=os.path.join(data_dir, input_data["gt_l_structure"]))

        affinity_loss_factor = torch.tensor([1.0], dtype=torch.float32)
        if input_data.get("affinity") is None:
            eps = 1e-6
            affinity_loss_factor = torch.tensor([eps], dtype=torch.float32)
            affinity = torch.tensor([0.0], dtype=torch.float32)
        else:
            affinity = torch.tensor([input_data["affinity"]], dtype=torch.float32)

        resolution = torch.tensor(input_data["resolution"], dtype=torch.float32)

        # prepare inter_contacts
        expanded_prot_r_pos = gt_protein_r_feats["pseudo_beta"].unsqueeze(1)  # Shape: (n_res_r, 1, 3)
        expanded_prot_l_pos = gt_protein_l_feats["pseudo_beta"].unsqueeze(0)  # Shape: (1, n_res_l, 3)
        distances = torch.sqrt(torch.sum((expanded_prot_r_pos - expanded_prot_l_pos) ** 2, dim=-1))
        inter_contact = (distances < 8.0).float()
        binding_site_mask_r = inter_contact.any(dim=1).float()
        binding_site_mask_l = inter_contact.any(dim=0).float()
        print("attaching binding masks", binding_site_mask_r.shape, binding_site_mask_l.shape)
        binding_site_mask = torch.cat([binding_site_mask_r, binding_site_mask_l], dim=0)

        inter_contact_reshaped_to_crop = torch.zeros((crop_size, crop_size), dtype=torch.float32)
        inter_contact_reshaped_to_crop[:n_res_r, n_res_r:n_res_total] = inter_contact
        inter_contact_reshaped_to_crop[n_res_r:n_res_total, :n_res_r] = inter_contact.T

        # Use CA positions only
        atom37_gt_positions = torch.cat([gt_protein_r_feats["all_atom_positions"],
                                         gt_protein_l_feats["all_atom_positions"]], dim=0)
        atom37_atom_exists_in_res = torch.cat([gt_protein_r_feats["atom37_atom_exists"],
                                               gt_protein_l_feats["atom37_atom_exists"]], dim=0)
        atom37_atom_exists_in_gt = torch.cat([gt_protein_r_feats["all_atom_mask"],
                                              gt_protein_l_feats["all_atom_mask"]], dim=0)

        atom14_gt_positions = torch.cat([gt_protein_r_feats["atom14_gt_positions"],
                                         gt_protein_l_feats["atom14_gt_positions"]], dim=0)
        atom14_atom_exists_in_res = torch.cat([gt_protein_r_feats["atom14_atom_exists"],
                                               gt_protein_l_feats["atom14_atom_exists"]], dim=0)
        atom14_atom_exists_in_gt = torch.cat([gt_protein_r_feats["atom14_gt_exists"],
                                              gt_protein_l_feats["atom14_gt_exists"]], dim=0)

        gt_pseudo_beta_joined = torch.cat([gt_protein_r_feats["pseudo_beta"], gt_protein_l_feats["pseudo_beta"]], dim=0)
        gt_pseudo_beta_joined_mask = torch.cat([gt_protein_r_feats["pseudo_beta_mask"],
                                         gt_protein_l_feats["pseudo_beta_mask"]], dim=0)

        # IGNORES: residx_atom14_to_atom37, rigidgroups_group_exists,
        # rigidgroups_group_is_ambiguous, pseudo_beta_mask, backbone_rigid_mask, protein_target_feat
        gt_protein_feats = {
            "atom37_gt_positions": atom37_gt_positions,  # torch.Size([n_struct, 37, 3])
            "atom37_atom_exists_in_res": atom37_atom_exists_in_res,  # torch.Size([n_struct, 37])
            "atom37_atom_exists_in_gt": atom37_atom_exists_in_gt,  # torch.Size([n_struct, 37])

            "atom14_gt_positions": atom14_gt_positions,  # torch.Size([n_struct, 14, 3])
            "atom14_atom_exists_in_res": atom14_atom_exists_in_res,  # torch.Size([n_struct, 14])
            "atom14_atom_exists_in_gt": atom14_atom_exists_in_gt,  # torch.Size([n_struct, 14])

            "gt_pseudo_beta_joined": gt_pseudo_beta_joined,  # torch.Size([n_struct, 3])
            "gt_pseudo_beta_joined_mask": gt_pseudo_beta_joined_mask,  # torch.Size([n_struct])

            # These we don't need to add the ligand to, because padding is sufficient (everything should be 0)
            "atom14_alt_gt_positions": torch.cat([gt_protein_r_feats["atom14_alt_gt_positions"],
                                                  gt_protein_l_feats["atom14_alt_gt_positions"]], dim=0),  # torch.Size([n_res, 14, 3])
            "atom14_alt_gt_exists": torch.cat([gt_protein_r_feats["atom14_alt_gt_exists"],
                                               gt_protein_l_feats["atom14_alt_gt_exists"]], dim=0),  # torch.Size([n_res, 14])
            "atom14_atom_is_ambiguous": torch.cat([gt_protein_r_feats["atom14_atom_is_ambiguous"],
                                                   gt_protein_l_feats["atom14_atom_is_ambiguous"]], dim=0),  # torch.Size([n_res, 14])
            "rigidgroups_gt_frames": torch.cat([gt_protein_r_feats["rigidgroups_gt_frames"],
                                                gt_protein_l_feats["rigidgroups_gt_frames"]], dim=0),  # torch.Size([n_res, 8, 4, 4])
            "rigidgroups_gt_exists": torch.cat([gt_protein_r_feats["rigidgroups_gt_exists"],
                                                gt_protein_l_feats["rigidgroups_gt_exists"]], dim=0),  # torch.Size([n_res, 8])
            "rigidgroups_alt_gt_frames": torch.cat([gt_protein_r_feats["rigidgroups_alt_gt_frames"],
                                                    gt_protein_l_feats["rigidgroups_alt_gt_frames"]], dim=0),  # torch.Size([n_res, 8, 4, 4])
            "backbone_rigid_tensor": torch.cat([gt_protein_r_feats["backbone_rigid_tensor"],
                                                gt_protein_l_feats["backbone_rigid_tensor"]], dim=0),  # torch.Size([n_res, 4, 4])
            "backbone_rigid_mask": torch.cat([gt_protein_r_feats["backbone_rigid_mask"],
                                              gt_protein_l_feats["backbone_rigid_mask"]], dim=0),  # torch.Size([n_res])
            "chi_angles_sin_cos": torch.cat([gt_protein_r_feats["chi_angles_sin_cos"],
                                             gt_protein_l_feats["chi_angles_sin_cos"]], dim=0),
            "chi_mask": torch.cat([gt_protein_r_feats["chi_mask"], gt_protein_l_feats["chi_mask"]], dim=0),
        }

        for k, v in gt_protein_feats.items():
            gt_protein_feats[k] = _fit_to_crop(v, crop_size, 0)

        feats = {
            **feats,
            **gt_protein_feats,
            "resolution": resolution,
            "affinity": affinity,
            "affinity_loss_factor": affinity_loss_factor,
            "seq_length": torch.tensor(n_res_total),
            "binding_site_mask": _fit_to_crop(binding_site_mask, crop_size, 0),
            "gt_inter_contacts": inter_contact_reshaped_to_crop,
        }

    for k, v in feats.items():
        # print(k, v.shape)
        feats[k] = _prepare_recycles(v, num_recycles)

    feats["batch_idx"] = torch.tensor(
        [idx for _ in range(crop_size)], dtype=torch.int64, device=feats["aatype"].device
    )

    print("load time", round(time.time() - start_load_time, 4))

    return feats