Rename main directory

alphafold/np/relax/amber_minimize.py → openfold/np/relax/amber_minimize.py

@@ -20,11 +20,11 @@ import time
 from typing import Collection, Optional, Sequence
 
 from absl import logging
-from alphafold.np import (
+from openfold.np import (
     protein,
     residue_constants,
 )
-from alphafold.utils.loss import (
+from openfold.utils.loss import (
     find_structural_violations_np, 
     compute_violation_metrics_np,
 )
@@ -32,7 +32,7 @@ from alphafold.utils.loss import (
 find_structural_violations = find_structural_violations_np
 compute_violation_metrics = compute_violation_metrics_np
 
-from alphafold.np.relax import cleanup, utils
+from openfold.np.relax import cleanup, utils
 import ml_collections
 import numpy as np
 from simtk import openmm

alphafold/np/relax/relax.py → openfold/np/relax/relax.py

@@ -15,8 +15,8 @@
 
 """Amber relaxation."""
 from typing import Any, Dict, Sequence, Tuple
-from alphafold.np import protein
-from alphafold.np.relax import amber_minimize, utils
+from openfold.np import protein
+from openfold.np.relax import amber_minimize, utils
 import numpy as np
 
 

alphafold/np/relax/utils.py → openfold/np/relax/utils.py

@@ -15,7 +15,7 @@
 
 """Utils for minimization."""
 import io
-from alphafold.np import residue_constants
+from openfold.np import residue_constants
 from Bio import PDB
 import numpy as np
 from simtk.openmm import app as openmm_app

alphafold/np/residue_constants.py → openfold/np/residue_constants.py

@@ -405,7 +405,7 @@ def load_stereo_chemical_props() -> Tuple[Mapping[str, List[Bond]],
   """
   # TODO: this file should be downloaded in a setup script
   stereo_chemical_props_path = (
-      'alphafold/resources/stereo_chemical_props.txt')
+      'openfold/resources/stereo_chemical_props.txt')
   with open(stereo_chemical_props_path, 'rt') as f:
     stereo_chemical_props = f.read()
   lines_iter = iter(stereo_chemical_props.splitlines())

alphafold/utils/feats.py → openfold/utils/feats.py

@@ -18,9 +18,9 @@ import torch
 import torch.nn as nn
 from typing import Dict
 
-import alphafold.np.residue_constants as residue_constants
-from alphafold.utils.affine_utils import T
-from alphafold.utils.tensor_utils import (
+import openfold.np.residue_constants as residue_constants
+from openfold.utils.affine_utils import T
+from openfold.utils.tensor_utils import (
     batched_gather, 
     one_hot,
 )

alphafold/utils/loss.py → openfold/utils/loss.py

@@ -20,10 +20,10 @@ import torch
 import torch.nn as nn
 from typing import Dict, Optional
 
-from alphafold.np import residue_constants
-from alphafold.model.primitives import Linear
-from alphafold.utils.affine_utils import T
-from alphafold.utils.tensor_utils import (
+from openfold.np import residue_constants
+from openfold.model.primitives import Linear
+from openfold.utils.affine_utils import T
+from openfold.utils.tensor_utils import (
     tree_map, 
     tensor_tree_map, 
     masked_mean,
@@ -992,7 +992,7 @@ def find_structural_violations_np(
     batch = tree_map(to_tensor, batch, np.ndarray)
     atom14_pred_positions = to_tensor(atom14_pred_positions)
 
-    out = find_structural_violations(batch, atom14_pred_positions, config)
+    out = find_structural_violations(batch, atom14_pred_positions, **config)
 
     to_np = lambda x: np.array(x)
     np_out = tensor_tree_map(to_np, out)
@@ -1246,7 +1246,7 @@ def experimentally_resolved_loss(
 def masked_msa_loss(logits, true_msa, bert_mask, eps=1e-8):
     errors = softmax_cross_entropy(
         logits,
-        torch.nn.functional.one_hot(true_msa, num_classes=23,
+        torch.nn.functional.one_hot(true_msa, num_classes=23)
     )
     loss = (
         torch.sum(errors * bert_mask, dim=(-1, -2)) /
@@ -1280,14 +1280,14 @@ class AlphaFoldLoss(nn.Module):
         loss_fns = {
             "distogram": 
                 lambda: distogram_loss(
-                    logits=out["distogram_logits"],
-                    {**batch,
+                    out["distogram_logits"],
+                    **{**batch,
                      **self.config.distogram},
                 ),
             "experimentally_resolved": 
                 lambda: experimentally_resolved_loss(
-                    logits=out["experimentally_resolved"],
-                    {**batch,
+                    out["experimentally_resolved"],
+                    **{**batch,
                      **self.config.experimentally_resolved},
                 ),
             "fape":
@@ -1298,22 +1298,22 @@ class AlphaFoldLoss(nn.Module):
                 ),
             "lddt": 
                 lambda: lddt_loss(
-                    logits=out["lddt_logits"],
+                    out["lddt_logits"],
                     all_atom_pred_pos=out["final_atom_positions"]
-                    {**batch,
+                    **{**batch,
                      **self.config.lddt},
                 ),
             "masked_msa": 
                 lambda: masked_msa_loss(
-                    logits=out["masked_msa_logits"],
-                    {**batch,
+                    out["masked_msa_logits"],
+                    **{**batch,
                      **self.config.masked_msa},
                 ),
             "supervised_chi":
                 lambda: supervised_chi_loss(
                     out["sm"]["angles"],
                     out["sm"]["unnormalized_angles"],
-                    {**batch,
+                    **{**batch,
                      **self.config.supervised_chi},
                 ),
             "violation":

run_pretrained_alphafold.py

@@ -25,14 +25,14 @@ import torch.nn as nn
 import numpy as np
 
 from config import model_config
-from alphafold.model.model import AlphaFold
-import alphafold.np.protein as protein
-import alphafold.np.relax.relax as relax
-from alphafold.np import residue_constants
-from alphafold.utils.import_weights import (
+from openfold.model.model import AlphaFold
+import openfold.np.protein as protein
+import openfold.np.relax.relax as relax
+from openfold.np import residue_constants
+from openfold.utils.import_weights import (
     import_jax_weights_,
 )
-from alphafold.utils.tensor_utils import (
+from openfold.utils.tensor_utils import (
     tree_map,
     tensor_tree_map,
 )
@@ -40,7 +40,7 @@ from alphafold.utils.tensor_utils import (
 
 MODEL_NAME = "model_1"
 MODEL_DEVICE = "cuda:1"
-PARAM_PATH = "alphafold/resources/params/params_model_1.npz"
+PARAM_PATH = "openfold/resources/params/params_model_1.npz"
 FEAT_PATH = "tests/test_data/sample_feats.pickle"
 
 

scripts/install_third_party_dependencies.sh

@@ -23,11 +23,11 @@ pushd lib/conda/lib/python3.9/site-packages/ \
     && popd
 
 # Download folding resources
-wget -q -P alphafold/resources \
+wget -q -P openfold/resources \
     https://git.scicore.unibas.ch/schwede/openstructure/-/raw/7102c63615b64735c4941278d92b554ec94415f8/modules/mol/alg/src/stereo_chemical_props.txt
 
-# Download pretrained Alphafold weights
-scripts/download_alphafold_params.sh alphafold/resources
+# Download pretrained openfold weights
+scripts/download_alphafold_params.sh openfold/resources
 
 # Decompress test data
 gunzip tests/test_data/sample_feats.pickle.gz