Spaces:
Sleeping
Sleeping
| # Copyright Generate Biomedicines, Inc. | |
| # | |
| # 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. | |
| """XCS represents protein structure as a tuple of PyTorch tensors. | |
| The tensors in an XCS representation are: | |
| `X` (FloatTensor), the Cartesian coordinates representing the protein | |
| structure with shape `(num_batch, num_residues, num_atoms, 3)`. The | |
| `num_atoms` dimension can be one of two sizes: `num_atoms=4` for | |
| backbone-only structures or `num_atoms=14` for all-atom structures | |
| (excluding hydrogens). The first four atoms will always be | |
| `N, CA, C, O`, and the meaning of the optional 10 additional atom | |
| positions will vary based on the residue identity at | |
| a given position. Atom orders for each amino acid are defined in | |
| `constants.AA_GEOMETRY[TRIPLET_CODE]["atoms"]`. | |
| `C` (LongTensor), the chain map encoding per-residue chain assignments with | |
| shape `(num_batch, num_residues)`.The chain map codes positions as `0` | |
| when masked, poitive integers for chain indices, and negative integers | |
| to represent missing residues (of the corresponding positive integers). | |
| `S` (LongTensor), the sequence of the protein as alphabet indices with | |
| shape `(num_batch, num_residues)`. The standard alphabet is | |
| `ACDEFGHIKLMNPQRSTVWY`, also defined in `constants.AA20`. | |
| """ | |
| from functools import partial, wraps | |
| from inspect import getfullargspec | |
| import torch | |
| from torch.nn import functional as F | |
| try: | |
| pass | |
| except ImportError: | |
| print("MST not installed!") | |
| def validate_XCS(all_atom=None, sequence=True): | |
| """Decorator factory that adds XCS validation to any function. | |
| Args: | |
| all_atom (bool, optional): If True, requires that input structure | |
| tensors have 14 residues per atom. If False, reduces to 4 residues | |
| per atom. If None, applies no transformation on input structures. | |
| sequence (bool, optional): If True, makes sure that if S and O are both | |
| provided, that they match, i.e. that O is a one-hot version of S. | |
| If only one of S or O is provided, the other is generated, and both | |
| are passed. | |
| """ | |
| def decorator(func): | |
| def new_func(*args, **kwargs): | |
| args = list(args) | |
| arg_list = getfullargspec(func)[0] | |
| tensors = {} | |
| for var in ["X", "C", "S", "O"]: | |
| try: | |
| if var in kwargs: | |
| tensors[var] = kwargs[var] | |
| else: | |
| tensors[var] = args[arg_list.index(var)] | |
| except IndexError: # empty args_list | |
| tensors[var] = None | |
| except ValueError: # variable not an argument of function | |
| if not sequence and var in ["S", "O"]: | |
| pass | |
| else: | |
| raise Exception( | |
| f"Variable {var} is required by validation but not defined!" | |
| ) | |
| if tensors["X"] is not None and tensors["C"] is not None: | |
| if tensors["X"].shape[:2] != tensors["C"].shape[:2]: | |
| raise ValueError( | |
| f"X shape {tensors['X'].shape} does not match C shape" | |
| f" {tensors['C'].shape}" | |
| ) | |
| if all_atom is not None and tensors["X"] is not None: | |
| if all_atom and tensors["X"].shape[2] != 14: | |
| raise ValueError("Side chain atoms missing!") | |
| elif not all_atom: | |
| if "X" in kwargs: | |
| kwargs["X"] = tensors["X"][:, :, :4] | |
| else: | |
| args[arg_list.index("X")] = tensors["X"][:, :, :4] | |
| if sequence and (tensors["S"] is not None or tensors["O"] is not None): | |
| if tensors["O"] is None: | |
| if "O" in kwargs: | |
| kwargs["O"] = F.one_hot(tensors["S"], 20).float() | |
| else: | |
| args[arg_list.index("O")] = F.one_hot(tensors["S"], 20).float() | |
| elif tensors["S"] is None: | |
| if "S" in kwargs: | |
| kwargs["S"] = tensors["O"].argmax(dim=2) | |
| else: | |
| args[arg_list.index("S")] = tensors["O"].argmax(dim=2) | |
| else: | |
| if not torch.allclose(tensors["O"].argmax(dim=2), tensors["S"]): | |
| raise ValueError("S and O are both provided but don't match!") | |
| return func(*args, **kwargs) | |
| return new_func | |
| return decorator | |
| validate_XC = partial(validate_XCS, sequence=False) | |