rc-foundry 0.1.1__py3-none-any.whl

mpnn/utils/probability.py

@@ -0,0 +1,37 @@
+import numpy as np
+
+
+def sample_bernoulli_rv(p):
+    """
+    Given a probability p, representing the success probability of a Bernoulli
+    distribution, sample X ~ Bernoulli(p).
+
+    Arguments:
+        p (float): a float between 0 and 1, representing the success probability
+            of a Bernoulli distribution.
+
+    Returns:
+        x (int): the result of sampling the random variable X ~ Bernoulli(p).
+            P(X = 1) = p
+            P(X = 0) = 1 - p.
+    """
+    # Check that 0 <= p <= 1.
+    if p < 0 or p > 1:
+        raise ValueError("The success probability p must be between 0 and 1 inclusive.")
+
+    # Handle the edge cases, otherwise utilize the numpy uniform distribution.
+    if p == 0:
+        x = 0
+    elif p == 1:
+        x = 1
+    else:
+        # Sample the Y ~ Uniform(0, 1) distribution.
+        uniform_sample = np.random.uniform(0.0, 1.0)
+
+        # P(Y < p) = p.
+        if uniform_sample < p:
+            x = 1
+        else:
+            x = 0
+
+    return x

mpnn/utils/weights.py

@@ -0,0 +1,309 @@
+import torch
+from mpnn.transforms.feature_aggregation.token_encodings import (
+    legacy_token_order,
+    token_order,
+)
+
+
+def load_legacy_weights(model: torch.nn.Module, weights_path: str) -> None:
+    """
+    Load a legacy MPNN checkpoint from 'weights_path' into 'model' (the
+    refactored MPNN implementation).
+
+    This performs several transformations:
+      - Copies certain non-learned registries (e.g., periodic table info) from
+        the new model into the checkpoint state dict (to match the new code).
+      - Renames legacy parameter/buffer names into the new module naming scheme.
+      - Fixes a 120->119 atom-type embedding weight size mismatch by dropping
+        the unused legacy atom type.
+      - Reorders pairwise backbone distance embedding weights to match the new
+        atom-pair ordering.
+      - Reorders token (AA) embeddings/projections weights from the legacy order
+        (alphabetical 1-letter) to the new order (alphabetical 3-letter).
+    """
+    # Load legacy checkpoint state dict.
+    checkpoint = torch.load(weights_path, map_location="cpu", weights_only=False)
+    checkpoint_state_dict = checkpoint["model_state_dict"]
+
+    # Values to copy directly from the *current* model into the state dict.
+    # These are effectively "configuration" tensors/registries, not learned
+    # weights we want to preserve from the legacy model.
+    values_to_copy = [
+        "graph_featurization_module.side_chain_atom_types",
+        "graph_featurization_module.periodic_table_groups",
+        "graph_featurization_module.periodic_table_periods",
+    ]
+    # Copy over some hard-coded registers/values.
+    for value_name in values_to_copy:
+        # Walk the attribute chain.
+        attr_list = value_name.split(".")
+        sub_module = model
+        while len(attr_list) > 1 and sub_module is not None:
+            attr = attr_list.pop(0)
+            if hasattr(sub_module, attr):
+                sub_module = getattr(sub_module, attr)
+            else:
+                sub_module = None
+
+        # If the current sub-module exists, and it has the final attribute,
+        # copy it into the checkpoint state dict.
+        if sub_module is not None:
+            if hasattr(sub_module, attr_list[0]):
+                checkpoint_state_dict[value_name] = getattr(sub_module, attr_list[0])
+
+    # Mapping of legacy weight names to new weight names.
+    # Left side = name in the old checkpoint.
+    # Right side = name expected by the new model implementation.
+    legacy_weight_to_new_weight = {
+        "features.embeddings.linear.weight": "graph_featurization_module.positional_embedding.embed_positional_features.weight",
+        "features.embeddings.linear.bias": "graph_featurization_module.positional_embedding.embed_positional_features.bias",
+        "features.edge_embedding.weight": "graph_featurization_module.edge_embedding.weight",
+        "features.norm_edges.weight": "graph_featurization_module.edge_norm.weight",
+        "features.norm_edges.bias": "graph_featurization_module.edge_norm.bias",
+        "context_encoder_layers.0.norm1.weight": "protein_ligand_context_encoder_layers.0.norm1.weight",
+        "context_encoder_layers.0.norm1.bias": "protein_ligand_context_encoder_layers.0.norm1.bias",
+        "context_encoder_layers.0.norm2.weight": "protein_ligand_context_encoder_layers.0.norm2.weight",
+        "context_encoder_layers.0.norm2.bias": "protein_ligand_context_encoder_layers.0.norm2.bias",
+        "context_encoder_layers.0.W1.weight": "protein_ligand_context_encoder_layers.0.W1.weight",
+        "context_encoder_layers.0.W1.bias": "protein_ligand_context_encoder_layers.0.W1.bias",
+        "context_encoder_layers.0.W2.weight": "protein_ligand_context_encoder_layers.0.W2.weight",
+        "context_encoder_layers.0.W2.bias": "protein_ligand_context_encoder_layers.0.W2.bias",
+        "context_encoder_layers.0.W3.weight": "protein_ligand_context_encoder_layers.0.W3.weight",
+        "context_encoder_layers.0.W3.bias": "protein_ligand_context_encoder_layers.0.W3.bias",
+        "context_encoder_layers.0.dense.W_in.weight": "protein_ligand_context_encoder_layers.0.dense.W_in.weight",
+        "context_encoder_layers.0.dense.W_in.bias": "protein_ligand_context_encoder_layers.0.dense.W_in.bias",
+        "context_encoder_layers.0.dense.W_out.weight": "protein_ligand_context_encoder_layers.0.dense.W_out.weight",
+        "context_encoder_layers.0.dense.W_out.bias": "protein_ligand_context_encoder_layers.0.dense.W_out.bias",
+        "context_encoder_layers.1.norm1.weight": "protein_ligand_context_encoder_layers.1.norm1.weight",
+        "context_encoder_layers.1.norm1.bias": "protein_ligand_context_encoder_layers.1.norm1.bias",
+        "context_encoder_layers.1.norm2.weight": "protein_ligand_context_encoder_layers.1.norm2.weight",
+        "context_encoder_layers.1.norm2.bias": "protein_ligand_context_encoder_layers.1.norm2.bias",
+        "context_encoder_layers.1.W1.weight": "protein_ligand_context_encoder_layers.1.W1.weight",
+        "context_encoder_layers.1.W1.bias": "protein_ligand_context_encoder_layers.1.W1.bias",
+        "context_encoder_layers.1.W2.weight": "protein_ligand_context_encoder_layers.1.W2.weight",
+        "context_encoder_layers.1.W2.bias": "protein_ligand_context_encoder_layers.1.W2.bias",
+        "context_encoder_layers.1.W3.weight": "protein_ligand_context_encoder_layers.1.W3.weight",
+        "context_encoder_layers.1.W3.bias": "protein_ligand_context_encoder_layers.1.W3.bias",
+        "context_encoder_layers.1.dense.W_in.weight": "protein_ligand_context_encoder_layers.1.dense.W_in.weight",
+        "context_encoder_layers.1.dense.W_in.bias": "protein_ligand_context_encoder_layers.1.dense.W_in.bias",
+        "context_encoder_layers.1.dense.W_out.weight": "protein_ligand_context_encoder_layers.1.dense.W_out.weight",
+        "context_encoder_layers.1.dense.W_out.bias": "protein_ligand_context_encoder_layers.1.dense.W_out.bias",
+        "y_context_encoder_layers.0.norm1.weight": "ligand_context_encoder_layers.0.norm1.weight",
+        "y_context_encoder_layers.0.norm1.bias": "ligand_context_encoder_layers.0.norm1.bias",
+        "y_context_encoder_layers.0.norm2.weight": "ligand_context_encoder_layers.0.norm2.weight",
+        "y_context_encoder_layers.0.norm2.bias": "ligand_context_encoder_layers.0.norm2.bias",
+        "y_context_encoder_layers.0.W1.weight": "ligand_context_encoder_layers.0.W1.weight",
+        "y_context_encoder_layers.0.W1.bias": "ligand_context_encoder_layers.0.W1.bias",
+        "y_context_encoder_layers.0.W2.weight": "ligand_context_encoder_layers.0.W2.weight",
+        "y_context_encoder_layers.0.W2.bias": "ligand_context_encoder_layers.0.W2.bias",
+        "y_context_encoder_layers.0.W3.weight": "ligand_context_encoder_layers.0.W3.weight",
+        "y_context_encoder_layers.0.W3.bias": "ligand_context_encoder_layers.0.W3.bias",
+        "y_context_encoder_layers.0.dense.W_in.weight": "ligand_context_encoder_layers.0.dense.W_in.weight",
+        "y_context_encoder_layers.0.dense.W_in.bias": "ligand_context_encoder_layers.0.dense.W_in.bias",
+        "y_context_encoder_layers.0.dense.W_out.weight": "ligand_context_encoder_layers.0.dense.W_out.weight",
+        "y_context_encoder_layers.0.dense.W_out.bias": "ligand_context_encoder_layers.0.dense.W_out.bias",
+        "y_context_encoder_layers.1.norm1.weight": "ligand_context_encoder_layers.1.norm1.weight",
+        "y_context_encoder_layers.1.norm1.bias": "ligand_context_encoder_layers.1.norm1.bias",
+        "y_context_encoder_layers.1.norm2.weight": "ligand_context_encoder_layers.1.norm2.weight",
+        "y_context_encoder_layers.1.norm2.bias": "ligand_context_encoder_layers.1.norm2.bias",
+        "y_context_encoder_layers.1.W1.weight": "ligand_context_encoder_layers.1.W1.weight",
+        "y_context_encoder_layers.1.W1.bias": "ligand_context_encoder_layers.1.W1.bias",
+        "y_context_encoder_layers.1.W2.weight": "ligand_context_encoder_layers.1.W2.weight",
+        "y_context_encoder_layers.1.W2.bias": "ligand_context_encoder_layers.1.W2.bias",
+        "y_context_encoder_layers.1.W3.weight": "ligand_context_encoder_layers.1.W3.weight",
+        "y_context_encoder_layers.1.W3.bias": "ligand_context_encoder_layers.1.W3.bias",
+        "y_context_encoder_layers.1.dense.W_in.weight": "ligand_context_encoder_layers.1.dense.W_in.weight",
+        "y_context_encoder_layers.1.dense.W_in.bias": "ligand_context_encoder_layers.1.dense.W_in.bias",
+        "y_context_encoder_layers.1.dense.W_out.weight": "ligand_context_encoder_layers.1.dense.W_out.weight",
+        "y_context_encoder_layers.1.dense.W_out.bias": "ligand_context_encoder_layers.1.dense.W_out.bias",
+        "features.node_project_down.weight": "graph_featurization_module.node_embedding.weight",
+        "features.node_project_down.bias": "graph_featurization_module.node_embedding.bias",
+        "features.norm_nodes.weight": "graph_featurization_module.node_norm.weight",
+        "features.norm_nodes.bias": "graph_featurization_module.node_norm.bias",
+        "features.type_linear.weight": "graph_featurization_module.embed_atom_type_features.weight",
+        "features.type_linear.bias": "graph_featurization_module.embed_atom_type_features.bias",
+        "features.y_nodes.weight": "graph_featurization_module.ligand_subgraph_node_embedding.weight",
+        "features.y_edges.weight": "graph_featurization_module.ligand_subgraph_edge_embedding.weight",
+        "features.norm_y_edges.weight": "graph_featurization_module.ligand_subgraph_edge_norm.weight",
+        "features.norm_y_edges.bias": "graph_featurization_module.ligand_subgraph_edge_norm.bias",
+        "features.norm_y_nodes.weight": "graph_featurization_module.ligand_subgraph_node_norm.weight",
+        "features.norm_y_nodes.bias": "graph_featurization_module.ligand_subgraph_node_norm.bias",
+        "W_v.weight": "W_protein_to_ligand_edges_embed.weight",
+        "W_v.bias": "W_protein_to_ligand_edges_embed.bias",
+        "W_c.weight": "W_protein_encoding_embed.weight",
+        "W_c.bias": "W_protein_encoding_embed.bias",
+        "W_nodes_y.weight": "W_ligand_nodes_embed.weight",
+        "W_nodes_y.bias": "W_ligand_nodes_embed.bias",
+        "W_edges_y.weight": "W_ligand_edges_embed.weight",
+        "W_edges_y.bias": "W_ligand_edges_embed.bias",
+        "V_C.weight": "W_final_context_embed.weight",
+        "V_C_norm.weight": "final_context_norm.weight",
+        "V_C_norm.bias": "final_context_norm.bias",
+    }
+    # Rename the weights in the checkpoint state dict.
+    for legacy_weight_name, new_weight_name in legacy_weight_to_new_weight.items():
+        if legacy_weight_name in checkpoint_state_dict:
+            checkpoint_state_dict[new_weight_name] = checkpoint_state_dict.pop(
+                legacy_weight_name
+            )
+
+    # Remove unused atom type embedding weight.
+    # - Previous LigandMPNN used 120 atom types, but the last one was unused.
+    # - The new model uses 119 atom types.
+    atom_type_embedding_keys = [
+        "graph_featurization_module.embed_atom_type_features.weight",
+        "graph_featurization_module.ligand_subgraph_node_embedding.weight",
+    ]
+    # For each of these keys, drop the unused atom type embedding.
+    for key in atom_type_embedding_keys:
+        if key in checkpoint_state_dict:
+            legacy_weight = checkpoint_state_dict[key]
+            num_atomic_numbers = model.graph_featurization_module.num_atomic_numbers
+            checkpoint_state_dict[key] = torch.cat(
+                [
+                    legacy_weight[:, :num_atomic_numbers],
+                    legacy_weight[:, num_atomic_numbers + 1 :],
+                ],
+                dim=1,
+            )
+
+    # Permute weights for embedding of pairwise backbone atom distances.
+    # - The legacy model used the order specified in 'legacy_order' dict.
+    # - The new model uses the order specified in 'new_order' list (the
+    # outer product of the atom types in the order N, Ca, C, O, Cb).
+    legacy_order = {
+        "Ca-Ca": 0,
+        "N-N": 1,
+        "C-C": 2,
+        "O-O": 3,
+        "Cb-Cb": 4,
+        "Ca-N": 5,
+        "Ca-C": 6,
+        "Ca-O": 7,
+        "Ca-Cb": 8,
+        "N-C": 9,
+        "N-O": 10,
+        "N-Cb": 11,
+        "Cb-C": 12,
+        "Cb-O": 13,
+        "O-C": 14,
+        "N-Ca": 15,
+        "C-Ca": 16,
+        "O-Ca": 17,
+        "Cb-Ca": 18,
+        "C-N": 19,
+        "O-N": 20,
+        "Cb-N": 21,
+        "C-Cb": 22,
+        "O-Cb": 23,
+        "C-O": 24,
+    }
+    new_order = [
+        "N-N",
+        "N-Ca",
+        "N-C",
+        "N-O",
+        "N-Cb",
+        "Ca-N",
+        "Ca-Ca",
+        "Ca-C",
+        "Ca-O",
+        "Ca-Cb",
+        "C-N",
+        "C-Ca",
+        "C-C",
+        "C-O",
+        "C-Cb",
+        "O-N",
+        "O-Ca",
+        "O-C",
+        "O-O",
+        "O-Cb",
+        "Cb-N",
+        "Cb-Ca",
+        "Cb-C",
+        "Cb-O",
+        "Cb-Cb",
+    ]
+    pairwise_backbone_atom_embeddings_keys = [
+        "graph_featurization_module.edge_embedding.weight",
+    ]
+    for key in pairwise_backbone_atom_embeddings_keys:
+        if key in checkpoint_state_dict:
+            # Grab the legacy weight and shape.
+            legacy_weight = checkpoint_state_dict[key]
+            out_dim, _ = legacy_weight.shape
+
+            # Grab the necessary dimensions from the model.
+            num_positional_embeddings = (
+                model.graph_featurization_module.num_positional_embeddings
+            )
+            num_atoms = (
+                model.graph_featurization_module.num_backbone_atoms
+                + model.graph_featurization_module.num_virtual_atoms
+            )
+            num_rbf = model.graph_featurization_module.num_rbf
+
+            # Split positional and RBF embedding weights.
+            legacy_weight_positional_embeddings = legacy_weight[
+                :, :num_positional_embeddings
+            ]
+            legacy_weight_rbf_embeddings_flat = legacy_weight[
+                :, num_positional_embeddings:
+            ]
+
+            # Reshape the weights to separate atom pairs and the rbf dimension.
+            legacy_weight_rbf_embeddings_atom_pairs = (
+                legacy_weight_rbf_embeddings_flat.view(
+                    out_dim, num_atoms * num_atoms, num_rbf
+                )
+            )
+
+            # Reorder the atom pairs to match the new order.
+            new_weight_rbf_embeddings_atom_pairs = (
+                legacy_weight_rbf_embeddings_atom_pairs[
+                    :, [legacy_order[atom_pair_name] for atom_pair_name in new_order], :
+                ]
+            )
+
+            # Flatten the reordered weights back to 2D.
+            new_weight_rbf_embeddings_flat = (
+                new_weight_rbf_embeddings_atom_pairs.reshape(
+                    out_dim, num_atoms * num_atoms * num_rbf
+                )
+            )
+
+            # Concatenate positional + reordered RBF
+            checkpoint_state_dict[key] = torch.cat(
+                [legacy_weight_positional_embeddings, new_weight_rbf_embeddings_flat],
+                dim=1,
+            )
+
+    # Permute the token order of amino acids coming out of the model to match
+    # the new vocabulary order.
+    # - The legacy model used an order specified by alphabetic order of one-
+    # letter amino acid codes.
+    # - The new model uses an order specified by alphabetic order of three-
+    # letter amino acid codes.
+    token_embedding_keys = ["W_s.weight"]
+    token_projection_keys = ["W_out.weight", "W_out.bias"]
+    # For each of these keys, reorder the embeddings/projections.
+    for key in token_embedding_keys + token_projection_keys:
+        if key in checkpoint_state_dict:
+            # Grab the old weight.
+            legacy_weight = checkpoint_state_dict[key]
+
+            # Reorder the weight/bias according to the new token order.
+            if "weight" in key:
+                checkpoint_state_dict[key] = legacy_weight[
+                    [legacy_token_order.index(aa) for aa in token_order], :
+                ]
+            elif "bias" in key:
+                checkpoint_state_dict[key] = legacy_weight[
+                    [legacy_token_order.index(aa) for aa in token_order]
+                ]
+            else:
+                raise ValueError(f"Unrecognized key for token projection: {key}")
+
+    # Load the modified state dict into the model.
+    model.load_state_dict(checkpoint_state_dict, strict=True)

rc_foundry-0.1.1.dist-info/METADATA

@@ -0,0 +1,239 @@
+Metadata-Version: 2.4
+Name: rc-foundry
+Version: 0.1.1
+Summary: Shared utilities and training infrastructure for biomolecular structure prediction models.
+Author-email: Institute for Protein Design <contact@ipd.uw.edu>
+License: BSD 3-Clause License
+        
+        Copyright (c) 2025, Institute for Protein Design, University of Washington
+        
+        Redistribution and use in source and binary forms, with or without
+        modification, are permitted provided that the following conditions are met:
+        
+        * Redistributions of source code must retain the above copyright notice, this
+          list of conditions and the following disclaimer.
+        
+        * Redistributions in binary form must reproduce the above copyright notice,
+          this list of conditions and the following disclaimer in the documentation
+          and/or other materials provided with the distribution.
+        
+        * Neither the name of the copyright holder nor the names of its
+          contributors may be used to endorse or promote products derived from
+          this software without specific prior written permission.
+        
+        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+        AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+        IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+        DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+        FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+        DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+        SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+        CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+        OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+        OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+License-File: LICENSE.md
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: BSD License
+Classifier: Natural Language :: English
+Classifier: Operating System :: MacOS
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: Implementation :: CPython
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.12
+Requires-Dist: assertpy
+Requires-Dist: atomworks[ml]>=2.1.1
+Requires-Dist: beartype<1,>=0.18.0
+Requires-Dist: dm-tree<1,>=0.1.6
+Requires-Dist: einops<1,>=0.8.0
+Requires-Dist: einx<1,>=0.1.0
+Requires-Dist: environs<12,>=11.0.0
+Requires-Dist: hydra-core<1.4,>=1.3.0
+Requires-Dist: ipykernel>=6.31.0
+Requires-Dist: jaxtyping<1,>=0.2.17
+Requires-Dist: lightning>=2.5.0
+Requires-Dist: loralib>=0.1.1
+Requires-Dist: opt-einsum<4,>=3.4.0
+Requires-Dist: pandas
+Requires-Dist: rich>=13.9.4
+Requires-Dist: rootutils<1.1,>=1.0.7
+Requires-Dist: toolz
+Requires-Dist: torch<3,>=2.2.0
+Requires-Dist: typer<1,>=0.20.0
+Requires-Dist: wandb<1,>=0.15.10
+Requires-Dist: zstandard
+Provides-Extra: all
+Requires-Dist: cuequivariance-ops-cu12>=0.6.1; (sys_platform == 'linux') and extra == 'all'
+Requires-Dist: cuequivariance-ops-torch-cu12>=0.6.1; (sys_platform == 'linux') and extra == 'all'
+Requires-Dist: cuequivariance-torch>=0.6.1; (sys_platform == 'linux') and extra == 'all'
+Requires-Dist: pydantic>=2.8; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: assertpy; extra == 'dev'
+Requires-Dist: atomworks[dev,ml,openbabel]>=2.1.1; extra == 'dev'
+Requires-Dist: debugpy<2,>=1.8.5; extra == 'dev'
+Requires-Dist: ipdb; extra == 'dev'
+Requires-Dist: ipykernel<7,>=6.29.4; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest-benchmark<6,>=5.0.0; extra == 'dev'
+Requires-Dist: pytest-cov<5,>=4.1.0; extra == 'dev'
+Requires-Dist: pytest-dotenv<1,>=0.5.2; extra == 'dev'
+Requires-Dist: pytest-testmon<3,>=2.1.1; extra == 'dev'
+Requires-Dist: pytest-xdist<4,>=3.6.1; extra == 'dev'
+Requires-Dist: pytest<9,>=8.2.0; extra == 'dev'
+Requires-Dist: ruff==0.8.3; extra == 'dev'
+Provides-Extra: rf3
+Requires-Dist: cuequivariance-ops-cu12>=0.6.1; (sys_platform == 'linux') and extra == 'rf3'
+Requires-Dist: cuequivariance-ops-torch-cu12>=0.6.1; (sys_platform == 'linux') and extra == 'rf3'
+Requires-Dist: cuequivariance-torch>=0.6.1; (sys_platform == 'linux') and extra == 'rf3'
+Provides-Extra: rfd3
+Requires-Dist: pydantic>=2.8; extra == 'rfd3'
+Description-Content-Type: text/markdown
+
+# Protein design with Foundry
+
+Foundry provides tooling and infrastructure for using and training all classes of models for protein design, including design (RFD3), inverse folding (ProteinMPNN) and protein folding (RF3).
+
+All models within Foundry rely on [AtomWorks](https://github.com/RosettaCommons/atomworks) - a unified framework for manipulating and processing biomolecular structures - for both training and inference. 
+
+## Getting Started
+### Quickstart guide
+**Installation**
+```bash
+pip install rc-foundry[all]
+```
+
+**Downloading weights** All models can be downloaded to a target folder with:
+
+```
+foundry install all --checkpoint_dir <path/to/ckpt/dir>
+```
+This will download all the models supported (including multiple checkpoints of RF3) but as a beginner you can start with:
+```
+foundry install rfd3 ligandmpnn rf3 --checkpoint_dir  <path/to/ckpt/dir>
+```
+
+>*See `examples/all.ipynb` for how to run each model in a notebook.*
+
+### RFdiffusion3 (RFD3)
+
+[RFdiffusion3](https://www.biorxiv.org/content/10.1101/2025.09.18.676967v2) is an all-atom generative model capable of designing protein structures under complex constraints. 
+
+> *See [models/rfd3/README.md](models/rfd3/README.md) for complete documentation.*
+
+<div align="center">
+  <img src="docs/_static/rfd3_trajectory.png" alt="RFdiffusion3 generation trajectory." width="400">
+</div>
+
+### ProteinMPNN
+[ProteinMPNN](https://www.science.org/doi/10.1126/science.add2187) and [LigandMPNN](https://www.nature.com/articles/s41592-025-02626-1) are lightweight inverse-folding models which can be use to design diverse sequences for backbones under constrained conditions.
+
+> *See [models/mpnn/README.md](models/mpnn/README.md) for complete documentation.*
+
+
+### RosettaFold3 (RF3)
+
+[RF3](https://doi.org/10.1101/2025.08.14.670328) is a structure prediction neural network that narrows the gap between closed-source AF-3 and open-source alternatives.
+
+<div align="center">
+  <img src="docs/_static/prot_dna.png" alt="Protein-DNA complex prediction" width="400">
+</div>
+
+> *See [models/rf3/README.md](models/rf3/README.md) for complete documentation.*
+
+---
+
+## Development
+
+### Code Organization
+
+**Strict dependency flow:** `foundry` → `atomworks`
+
+- **atomworks**: Structure I/O, preprocessing, featurization
+- **foundry**: Model architectures, training, inference endpoints
+- **models/\<model\>:** Released models.
+
+#### For Core Developers (Multiple Packages)
+
+Install both `foundry` and models in editable mode for development:
+
+```bash
+# Install foundry and RF3 in editable mode
+uv pip install -e . -e ./models/rf3 -e ./models/rfd3 -e ./models/mpnn
+
+# Or install only foundry (no models)
+uv pip install -e .
+```
+
+This approach allows you to:
+- Modify `foundry` shared utilities and see changes immediately
+- Work on specific models without installing all models
+- Add new models as independent packages in `models/`
+
+### Adding New Models
+
+To add a new model:
+
+1. Create `models/<model_name>/` directory with its own `pyproject.toml`
+2. Add `foundry` as a dependency
+3. Implement model-specific code in `models/<model_name>/src/`
+4. Users can install with: `uv pip install -e ./models/<model_name>`
+
+### Pre-commit Formatting
+
+We ship a `.pre-commit-config.yaml` that runs `make format` (via `ruff format`) before each commit. Enable it once per clone:
+
+```bash
+pip install pre-commit  # if not already installed
+pre-commit install
+```
+
+After installation the hook automatically formats the repo whenever you `git commit`. Use `pre-commit run --all-files` to apply it manually.
+
+## Citation
+
+If you use this repository code or data in your work, please cite the relavant work as below:
+
+```bibtex
+@article{corley2025accelerating,
+  title={Accelerating biomolecular modeling with atomworks and rf3},
+  author={Corley, Nathaniel and Mathis, Simon and Krishna, Rohith and Bauer, Magnus S and Thompson, Tuscan R and Ahern, Woody and Kazman, Maxwell W and Brent, Rafael I and Didi, Kieran and Kubaney, Andrew and others},
+  journal={bioRxiv},
+  year={2025}
+}
+
+@article {butcher2025_rfdiffusion3,
+    author = {Butcher, Jasper and Krishna, Rohith and Mitra, Raktim and Brent, Rafael Isaac and Li, Yanjing and Corley, Nathaniel and Kim, Paul T and Funk, Jonathan and Mathis, Simon Valentin and Salike, Saman and Muraishi, Aiko and Eisenach, Helen and Thompson, Tuscan Rock and Chen, Jie and Politanska, Yuliya and Sehgal, Enisha and Coventry, Brian and Zhang, Odin and Qiang, Bo and Didi, Kieran and Kazman, Maxwell and DiMaio, Frank and Baker, David},
+    title = {De novo Design of All-atom Biomolecular Interactions with RFdiffusion3},
+    elocation-id = {2025.09.18.676967},
+    year = {2025},
+    doi = {10.1101/2025.09.18.676967},
+    publisher = {Cold Spring Harbor Laboratory},
+    URL = {https://www.biorxiv.org/content/early/2025/11/19/2025.09.18.676967},
+    eprint = {https://www.biorxiv.org/content/early/2025/11/19/2025.09.18.676967.full.pdf},
+    journal = {bioRxiv}
+}
+
+@article{dauparas2022robust,
+  title={Robust deep learning--based protein sequence design using ProteinMPNN},
+  author={Dauparas, Justas and Anishchenko, Ivan and Bennett, Nathaniel and Bai, Hua and Ragotte, Robert J and Milles, Lukas F and Wicky, Basile IM and Courbet, Alexis and de Haas, Rob J and Bethel, Neville and others},
+  journal={Science},
+  volume={378},
+  number={6615},
+  pages={49--56},
+  year={2022},
+  publisher={American Association for the Advancement of Science}
+}
+
+@article{dauparas2025atomic,
+  title={Atomic context-conditioned protein sequence design using LigandMPNN},
+  author={Dauparas, Justas and Lee, Gyu Rie and Pecoraro, Robert and An, Linna and Anishchenko, Ivan and Glasscock, Cameron and Baker, David},
+  journal={Nature Methods},
+  pages={1--7},
+  year={2025},
+  publisher={Nature Publishing Group US New York}
+}
+```
+## Acknowledgments
+We thank Rachel Clune and Hope Woods from the RosettaCommons for their collaboration on the codebase, documentation, tutorials and examples.