Add files using upload-large-folder tool

config/callbacks/model_checkpoint.yaml

@@ -0,0 +1,16 @@
+# https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html
+
+model_checkpoint:
+  # _target_: utils.modelcheckpoint_timestamp.DynamicTimestampModelCheckpoint
+  _target_: lightning.pytorch.callbacks.ModelCheckpoint
+  dirpath: artifacts/checkpoints # We need to modify this base on distributed or not
+  filename: "model-best-step{step:08d}-loss{loss/mse_epoch:.6f}" # checkpoint filename
+  monitor: "trainer/global_step" # name of the logged metric which determines when model is improving
+  verbose: True # verbosity mode
+  save_last: True # additionally always save an exact copy of the last checkpoint to a file last.ckpt
+  save_top_k: 50 # save k best models (determined by above metric)
+  mode: "max" # "max" means higher metric value is better, can be also "min"
+  auto_insert_metric_name: False # when True, the checkpoints filenames will contain the metric name
+  save_weights_only: False # if True, then only the model’s weights will be saved
+  save_on_train_epoch_end: False # whether to run checkpointing at the end of the training epoch or the end of validation
+  every_n_train_steps: ${trainer.val_check_interval} 

config/data/cameo22.yaml

@@ -0,0 +1,3 @@
+_target_: onescience.datapipes.simplefold.test_datamodule.SimpleFoldInferenceDataModule
+target_dir: data/cameo22/
+num_workers: 1

config/data/casp14.yaml

@@ -0,0 +1,3 @@
+_target_: onescience.datapipes.simplefold.test_datamodule.SimpleFoldInferenceDataModule
+target_dir: data/casp14/
+num_workers: 1

config/data/pdb_sp.yaml

@@ -0,0 +1,56 @@
+_target_: onescience.datapipes.simplefold.train_datamodule.SimpleFoldTrainingDataModule
+
+datasets:
+  # - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+  #   data_name: rcsb_protein
+  #   tokenized_dir: data/rcsb_protein_tokenized
+  #   target_dir: data/rcsb_processed_targets
+  #   manifest_path: data/rcsb_protein_tokenized/manifest.json
+  #   cropper:
+  #     _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+  #     min_neighborhood: 0
+  #     max_neighborhood: 40
+  #   filters:
+  #     - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.resolution.ResolutionFilter
+  #       resolution: 5.0
+  #     - _target_: onescience.datapipes.boltz_data_pipeline.filter.dynamic.date.DateFilter
+  #       date: "2020-05-01"
+  #       ref: released
+
+  - _target_: onescience.datapipes.simplefold.train_datamodule.DatasetConfig
+    data_name: swissprot
+    tokenized_dir: ./datasets/tokenized
+    target_dir: ./datasets
+    manifest_path: ./datasets/manifest.json
+    record_list: 
+    cropper:
+      _target_: onescience.datapipes.boltz_data_pipeline.crop.boltz.BoltzCropper
+      min_neighborhood: 0
+      max_neighborhood: 40
+
+filters:
+    - _target_:  onescience.datapipes.boltz_data_pipeline.filter.dynamic.size.SizeFilter
+      min_chains: 1
+      max_chains: 300
+
+tokenizer:
+  _target_: onescience.datapipes.boltz_data_pipeline.tokenize.boltz_protein.BoltzTokenizer
+
+featurizer:
+  _target_: onescience.datapipes.boltz_data_pipeline.feature.featurizer.BoltzFeaturizer
+
+symmetries: data/symmetry.pkl
+max_tokens: 256
+max_atoms: 2304
+pad_to_max_tokens: False
+pad_to_max_atoms: False
+batch_size: 8
+num_workers: 1
+pin_memory: True
+return_train_symmetries: False
+min_dist: 2.0
+max_dist: 22.0
+num_bins: 64
+atoms_per_window_queries: 32
+rotation_augment_ref_pos: False
+rotation_augment_coords: True

config/hydra/default.yaml

@@ -0,0 +1,4 @@
+# https://hydra.cc/docs/configure_hydra/intro/
+
+# Use Hydra's built-in logging so the packaged runtime does not require
+# the optional hydra-colorlog plugin.

config/logger/tensorboard.yaml

@@ -0,0 +1,9 @@
+# https://www.tensorflow.org/tensorboard/
+
+tensorboard:
+  _target_: lightning.pytorch.loggers.tensorboard.TensorBoardLogger
+  save_dir: "${paths.output_dir}/tensorboard/"
+  name: null
+  log_graph: False
+  default_hp_metric: True
+  prefix: ""

config/trainer/default.yaml

@@ -0,0 +1,24 @@
+_target_: lightning.pytorch.trainer.Trainer
+default_root_dir: ${paths.output_dir}
+
+# set True to to ensure deterministic results
+# makes training slower but gives more reproducibility than just setting seeds
+deterministic: False
+
+accelerator: gpu
+devices: auto
+strategy: 
+  _target_: lightning.pytorch.strategies.DDPStrategy
+  find_unused_parameters: True
+  gradient_as_bucket_view: True
+  timeout:
+    _target_: datetime.timedelta
+    seconds: 1800 # 30 minutes timeout for each training step
+num_nodes: 1
+min_epochs: 1 # prevents early stopping
+max_steps: 100000
+precision: bf16-mixed
+val_check_interval: 5000
+check_val_every_n_epoch: null
+num_sanity_val_steps: 0
+limit_val_batches: 0.0

configuration.json

@@ -0,0 +1,9 @@
+{
+  "framework": "Pytorch",
+  "task": "protein-structure-prediction",
+  "model": {
+    "type": "SimpleFold"
+  },
+  "allow_remote_download": false,
+  "entry_file": "scripts/run_inference.py"
+}

models/esm/inverse_folding/gvp_encoder.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from argparse import Namespace
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .features import GVPGraphEmbedding
+from .gvp_modules import GVPConvLayer, LayerNorm 
+from .gvp_utils import unflatten_graph
+
+
+
+class GVPEncoder(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.embed_graph = GVPGraphEmbedding(args)
+
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        
+        conv_activations = (F.relu, torch.sigmoid)
+        self.encoder_layers = nn.ModuleList(
+                GVPConvLayer(
+                    node_hidden_dim,
+                    edge_hidden_dim,
+                    drop_rate=args.dropout,
+                    vector_gate=True,
+                    attention_heads=0,
+                    n_message=3,
+                    conv_activations=conv_activations,
+                    n_edge_gvps=0,
+                    eps=1e-4,
+                    layernorm=True,
+                ) 
+            for i in range(args.num_encoder_layers)
+        )
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        node_embeddings, edge_embeddings, edge_index = self.embed_graph(
+                coords, coord_mask, padding_mask, confidence)
+        
+        for i, layer in enumerate(self.encoder_layers):
+            node_embeddings, edge_embeddings = layer(node_embeddings,
+                    edge_index, edge_embeddings)
+
+        node_embeddings = unflatten_graph(node_embeddings, coords.shape[0])
+        return node_embeddings

models/esm/inverse_folding/util.py

@@ -0,0 +1,323 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import math
+
+import biotite.structure
+from biotite.structure.io import pdbx, pdb
+from biotite.structure.residues import get_residues
+from biotite.structure import filter_backbone
+from biotite.structure import get_chains
+from biotite.sequence import ProteinSequence
+import numpy as np
+from scipy.spatial import transform
+from scipy.stats import special_ortho_group
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data as data
+from typing import Sequence, Tuple, List
+
+from onescience.datapipes.esm import BatchConverter
+
+
+def load_structure(fpath, chain=None):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id or list of chain ids to load
+    Returns:
+        biotite.structure.AtomArray
+    """
+    if fpath.endswith('cif'):
+        with open(fpath) as fin:
+            pdbxf = pdbx.PDBxFile.read(fin)
+        structure = pdbx.get_structure(pdbxf, model=1)
+    elif fpath.endswith('pdb'):
+        with open(fpath) as fin:
+            pdbf = pdb.PDBFile.read(fin)
+        structure = pdb.get_structure(pdbf, model=1)
+    bbmask = filter_backbone(structure)
+    structure = structure[bbmask]
+    all_chains = get_chains(structure)
+    if len(all_chains) == 0:
+        raise ValueError('No chains found in the input file.')
+    if chain is None:
+        chain_ids = all_chains
+    elif isinstance(chain, list):
+        chain_ids = chain
+    else:
+        chain_ids = [chain] 
+    for chain in chain_ids:
+        if chain not in all_chains:
+            raise ValueError(f'Chain {chain} not found in input file')
+    chain_filter = [a.chain_id in chain_ids for a in structure]
+    structure = structure[chain_filter]
+    return structure
+
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: An instance of biotite AtomArray
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+    residue_identities = get_residues(structure)[1]
+    seq = ''.join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+    return coords, seq
+
+
+def load_coords(fpath, chain):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    structure = load_structure(fpath, chain)
+    return extract_coords_from_structure(structure)
+
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+    """
+    Example for atoms argument: ["N", "CA", "C"]
+    """
+    def filterfn(s, axis=None):
+        filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+        sum = filters.sum(0)
+        if not np.all(sum <= np.ones(filters.shape[1])):
+            raise RuntimeError("structure has multiple atoms with same name")
+        index = filters.argmax(0)
+        coords = s[index].coord
+        coords[sum == 0] = float("nan")
+        return coords
+
+    return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+
+def get_sequence_loss(model, alphabet, coords, seq):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, seq)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+
+    prev_output_tokens = tokens[:, :-1].to(device)
+    target = tokens[:, 1:]
+    target_padding_mask = (target == alphabet.padding_idx)
+    logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens)
+    loss = F.cross_entropy(logits, target, reduction='none')
+    loss = loss[0].cpu().detach().numpy()
+    target_padding_mask = target_padding_mask[0].cpu().numpy()
+    return loss, target_padding_mask
+
+
+def score_sequence(model, alphabet, coords, seq):
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, coords, seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(~target_padding_mask)
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output(model, alphabet, coords):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, None)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+    encoder_out = model.encoder.forward(coords, padding_mask, confidence,
+            return_all_hiddens=False)
+    # remove beginning and end (bos and eos tokens)
+    return encoder_out['encoder_out'][0][1:-1, 0]
+
+
+def rotate(v, R):
+    """
+    Rotates a vector by a rotation matrix.
+    
+    Args:
+        v: 3D vector, tensor of shape (length x batch_size x channels x 3)
+        R: rotation matrix, tensor of shape (length x batch_size x 3 x 3)
+
+    Returns:
+        Rotated version of v by rotation matrix R.
+    """
+    R = R.unsqueeze(-3)
+    v = v.unsqueeze(-1)
+    return torch.sum(v * R, dim=-2)
+
+
+def get_rotation_frames(coords):
+    """
+    Returns a local rotation frame defined by N, CA, C positions.
+
+    Args:
+        coords: coordinates, tensor of shape (batch_size x length x 3 x 3)
+        where the third dimension is in order of N, CA, C
+
+    Returns:
+        Local relative rotation frames in shape (batch_size x length x 3 x 3)
+    """
+    v1 = coords[:, :, 2] - coords[:, :, 1]
+    v2 = coords[:, :, 0] - coords[:, :, 1]
+    e1 = normalize(v1, dim=-1)
+    u2 = v2 - e1 * torch.sum(e1 * v2, dim=-1, keepdim=True)
+    e2 = normalize(u2, dim=-1)
+    e3 = torch.cross(e1, e2, dim=-1)
+    R = torch.stack([e1, e2, e3], dim=-2)
+    return R
+
+
+def nan_to_num(ts, val=0.0):
+    """
+    Replaces nans in tensor with a fixed value.    
+    """
+    val = torch.tensor(val, dtype=ts.dtype, device=ts.device)
+    return torch.where(~torch.isfinite(ts), val, ts)
+
+
+def rbf(values, v_min, v_max, n_bins=16):
+    """
+    Returns RBF encodings in a new dimension at the end.
+    """
+    rbf_centers = torch.linspace(v_min, v_max, n_bins, device=values.device)
+    rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
+    rbf_std = (v_max - v_min) / n_bins
+    v_expand = torch.unsqueeze(values, -1)
+    z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
+    return torch.exp(-z ** 2)
+
+
+def norm(tensor, dim, eps=1e-8, keepdim=False):
+    """
+    Returns L2 norm along a dimension.
+    """
+    return torch.sqrt(
+            torch.sum(torch.square(tensor), dim=dim, keepdim=keepdim) + eps)
+
+
+def normalize(tensor, dim=-1):
+    """
+    Normalizes a tensor along a dimension after removing nans.
+    """
+    return nan_to_num(
+        torch.div(tensor, norm(tensor, dim=dim, keepdim=True))
+    )
+
+
+class CoordBatchConverter(BatchConverter):
+    def __call__(self, raw_batch: Sequence[Tuple[Sequence, str]], device=None):
+        """
+        Args:
+            raw_batch: List of tuples (coords, confidence, seq)
+            In each tuple,
+                coords: list of floats, shape L x 3 x 3
+                confidence: list of floats, shape L; or scalar float; or None
+                seq: string of length L
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        self.alphabet.cls_idx = self.alphabet.get_idx("<cath>") 
+        batch = []
+        for coords, confidence, seq in raw_batch:
+            if confidence is None:
+                confidence = 1.
+            if isinstance(confidence, float) or isinstance(confidence, int):
+                confidence = [float(confidence)] * len(coords)
+            if seq is None:
+                seq = 'X' * len(coords)
+            batch.append(((coords, confidence), seq))
+
+        coords_and_confidence, strs, tokens = super().__call__(batch)
+
+        # pad beginning and end of each protein due to legacy reasons
+        coords = [
+            F.pad(torch.tensor(cd), (0, 0, 0, 0, 1, 1), value=np.inf)
+            for cd, _ in coords_and_confidence
+        ]
+        confidence = [
+            F.pad(torch.tensor(cf), (1, 1), value=-1.)
+            for _, cf in coords_and_confidence
+        ]
+        coords = self.collate_dense_tensors(coords, pad_v=np.nan)
+        confidence = self.collate_dense_tensors(confidence, pad_v=-1.)
+        if device is not None:
+            coords = coords.to(device)
+            confidence = confidence.to(device)
+            tokens = tokens.to(device)
+        padding_mask = torch.isnan(coords[:,:,0,0])
+        coord_mask = torch.isfinite(coords.sum(-2).sum(-1))
+        confidence = confidence * coord_mask + (-1.) * padding_mask
+        return coords, confidence, strs, tokens, padding_mask
+
+    def from_lists(self, coords_list, confidence_list=None, seq_list=None, device=None):
+        """
+        Args:
+            coords_list: list of length batch_size, each item is a list of
+            floats in shape L x 3 x 3 to describe a backbone
+            confidence_list: one of
+                - None, default to highest confidence
+                - list of length batch_size, each item is a scalar
+                - list of length batch_size, each item is a list of floats of
+                    length L to describe the confidence scores for the backbone
+                    with values between 0. and 1.
+            seq_list: either None or a list of strings
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        batch_size = len(coords_list)
+        if confidence_list is None:
+            confidence_list = [None] * batch_size
+        if seq_list is None:
+            seq_list = [None] * batch_size
+        raw_batch = zip(coords_list, confidence_list, seq_list)
+        return self.__call__(raw_batch, device)
+
+    @staticmethod
+    def collate_dense_tensors(samples, pad_v):
+        """
+        Takes a list of tensors with the following dimensions:
+            [(d_11,       ...,           d_1K),
+             (d_21,       ...,           d_2K),
+             ...,
+             (d_N1,       ...,           d_NK)]
+        and stack + pads them into a single tensor of:
+        (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+        """
+        if len(samples) == 0:
+            return torch.Tensor()
+        if len(set(x.dim() for x in samples)) != 1:
+            raise RuntimeError(
+                f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+            )
+        (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+        max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+        result = torch.empty(
+            len(samples), *max_shape, dtype=samples[0].dtype, device=device
+        )
+        result.fill_(pad_v)
+        for i in range(len(samples)):
+            result_i = result[i]
+            t = samples[i]
+            result_i[tuple(slice(0, k) for k in t.shape)] = t
+        return result

weight/boltz1_conf.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fea245d912c570ec117b2277c2719f312a6fc109c07b6f6ef741690ee775c2f5
+size 3595352714

weight/esm_models/esm2_t36_3B_UR50D-contact-regression.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4da500eab246481dc9c8c95bc7b1d02f2803d761c380b0e95186d4a07d0fc84e
+size 6759

weight/esm_models/esm2_t36_3B_UR50D.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7de8b4082ba15891959ab368b77ce3886697af1efb16d3c9e9e7b0c5d3f07500
+size 5678116398

weight/plddt.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb32fa9cdc9e80406b793a8c09a929077534d9991a1d08f4c159d2e4ed81315f
+size 462812900

weight/plddt_module_1.6B.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb32fa9cdc9e80406b793a8c09a929077534d9991a1d08f4c159d2e4ed81315f
+size 462812900

weight/simplefold_100M.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4cd0b8a0b317a6ab8634444fffd78ce84cfd49c20fe927b83c76c36fda5f54bd
+size 386772550

weight/simplefold_700M.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:77a770e2dd1695397ab5943d5b31c99b9039b8a35ed86db12f2b2db1be9db7f8
+size 2764152946