Fix import ordering.

PiperOrigin-RevId: 387085679
Change-Id: I73287fcd0a29e899543b64c596e306195a2f435e

alphafold/model/prng_test.py

@@ -15,9 +15,8 @@
 """Tests for prng."""
 
 from absl.testing import absltest
-import jax
-
 from alphafold.model import prng
+import jax
 
 
 class PrngTest(absltest.TestCase):

alphafold/model/quat_affine_test.py

@@ -16,10 +16,10 @@
 
 from absl import logging
 from absl.testing import absltest
+from alphafold.model import quat_affine
 import jax
 import jax.numpy as jnp
 import numpy as np
-from alphafold.model import quat_affine
 
 VERBOSE = False
 np.set_printoptions(precision=3, suppress=True)

alphafold/model/r3.py

@@ -33,12 +33,10 @@ unintended use of these cores on both GPUs and TPUs.
 
 import collections
 from typing import List
-
+from alphafold.model import quat_affine
 import jax.numpy as jnp
 import tree
 
-from alphafold.model import quat_affine
-
 # Array of 3-component vectors, stored as individual array for
 # each component.
 Vecs = collections.namedtuple('Vecs', ['x', 'y', 'z'])

alphafold/model/tf/data_transforms.py

@@ -13,13 +13,13 @@
 # limitations under the License.
 
 """Data for AlphaFold."""
-import numpy as np
-import tensorflow.compat.v1 as tf
 
 from alphafold.common import residue_constants
 from alphafold.model.tf import shape_helpers
 from alphafold.model.tf import shape_placeholders
 from alphafold.model.tf import utils
+import numpy as np
+import tensorflow.compat.v1 as tf
 
 # Pylint gets confused by the curry1 decorator because it changes the number
 #   of arguments to the function.
@@ -623,4 +623,3 @@ def make_atom14_masks(protein):
   protein['atom37_atom_exists'] = residx_atom37_mask
 
   return protein
-

alphafold/model/tf/input_pipeline.py

@@ -13,11 +13,11 @@
 # limitations under the License.
 
 """Feature pre-processing input pipeline for AlphaFold."""
-import tensorflow.compat.v1 as tf
-import tree
 
 from alphafold.model.tf import data_transforms
 from alphafold.model.tf import shape_placeholders
+import tensorflow.compat.v1 as tf
+import tree
 
 # Pylint gets confused by the curry1 decorator because it changes the number
 #   of arguments to the function.

alphafold/model/tf/protein_features.py

@@ -15,10 +15,8 @@
 """Contains descriptions of various protein features."""
 import enum
 from typing import Dict, Optional, Sequence, Tuple, Union
-
-import tensorflow.compat.v1 as tf
-
 from alphafold.common import residue_constants
+import tensorflow.compat.v1 as tf
 
 # Type aliases.
 FeaturesMetadata = Dict[str, Tuple[tf.dtypes.DType, Sequence[Union[str, int]]]]
@@ -129,4 +127,3 @@ def shape(feature_name: str,
       raise ValueError("Could not parse %s (shape: %s) with values: %s" % (
           feature_name, raw_sizes, replacements))
   return sizes
-

alphafold/model/tf/protein_features_test.py

@@ -17,9 +17,8 @@ import uuid
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import tensorflow.compat.v1 as tf
-
 from alphafold.model.tf import protein_features
+import tensorflow.compat.v1 as tf
 
 
 def _random_bytes():

alphafold/model/tf/proteins_dataset.py

@@ -14,12 +14,10 @@
 
 """Datasets consisting of proteins."""
 from typing import Dict, Mapping, Optional, Sequence
-
+from alphafold.model.tf import protein_features
 import numpy as np
 import tensorflow.compat.v1 as tf
 
-from alphafold.model.tf import protein_features
-
 TensorDict = Dict[str, tf.Tensor]
 
 

alphafold/model/tf/shape_helpers_test.py

@@ -14,11 +14,10 @@
 
 """Tests for shape_helpers."""
 
+from alphafold.model.tf import shape_helpers
 import numpy as np
 import tensorflow.compat.v1 as tf
 
-from alphafold.model.tf import shape_helpers
-
 
 class ShapeTest(tf.test.TestCase):
 

alphafold/relax/amber_minimize.py

@@ -19,6 +19,11 @@ import time
 from typing import Collection, Optional, Sequence
 
 from absl import logging
+from alphafold.common import protein
+from alphafold.common import residue_constants
+from alphafold.model import folding
+from alphafold.relax import cleanup
+from alphafold.relax import utils
 import ml_collections
 import numpy as np
 from simtk import openmm
@@ -26,12 +31,6 @@ from simtk import unit
 from simtk.openmm import app as openmm_app
 from simtk.openmm.app.internal.pdbstructure import PdbStructure
 
-from alphafold.common import protein
-from alphafold.common import residue_constants
-from alphafold.model import folding
-from alphafold.relax import cleanup
-from alphafold.relax import utils
-
 
 ENERGY = unit.kilocalories_per_mole
 LENGTH = unit.angstroms

alphafold/relax/amber_minimize_test.py

@@ -16,10 +16,9 @@
 import os
 
 from absl.testing import absltest
-import numpy as np
-
 from alphafold.common import protein
 from alphafold.relax import amber_minimize
+import numpy as np
 # Internal import (7716).
 
 

alphafold/relax/cleanup_test.py

@@ -16,9 +16,8 @@
 import io
 
 from absl.testing import absltest
-from simtk.openmm.app.internal import pdbstructure
-
 from alphafold.relax import cleanup
+from simtk.openmm.app.internal import pdbstructure
 
 
 def _pdb_to_structure(pdb_str):

alphafold/relax/relax.py

@@ -14,12 +14,10 @@
 
 """Amber relaxation."""
 from typing import Any, Dict, Sequence, Tuple
-
-import numpy as np
-
 from alphafold.common import protein
 from alphafold.relax import amber_minimize
 from alphafold.relax import utils
+import numpy as np
 
 
 class AmberRelaxation(object):
@@ -80,4 +78,3 @@ class AmberRelaxation(object):
     violations = out['structural_violations'][
         'total_per_residue_violations_mask']
     return min_pdb, debug_data, violations
-

alphafold/relax/relax_test.py

@@ -16,9 +16,9 @@
 import os
 
 from absl.testing import absltest
-import numpy as np
 from alphafold.common import protein
 from alphafold.relax import relax
+import numpy as np
 # Internal import (7716).
 
 

alphafold/relax/utils.py

@@ -14,14 +14,12 @@
 
 """Utils for minimization."""
 import io
-
+from alphafold.common import residue_constants
 from Bio import PDB
 import numpy as np
 from simtk.openmm import app as openmm_app
 from simtk.openmm.app.internal.pdbstructure import PdbStructure
 
-from alphafold.common import residue_constants
-
 
 def overwrite_pdb_coordinates(pdb_str: str, pos) -> str:
   pdb_file = io.StringIO(pdb_str)
@@ -80,4 +78,3 @@ def assert_equal_nonterminal_atom_types(
   no_oxt_mask[..., oxt] = False
   np.testing.assert_almost_equal(ref_atom_mask[no_oxt_mask],
                                  atom_mask[no_oxt_mask])
-

alphafold/relax/utils_test.py

@@ -17,10 +17,9 @@
 import os
 
 from absl.testing import absltest
-import numpy as np
-
 from alphafold.common import protein
 from alphafold.relax import utils
+import numpy as np
 # Internal import (7716).
 
 

run_alphafold.py

@@ -25,8 +25,6 @@ from typing import Dict
 from absl import app
 from absl import flags
 from absl import logging
-import numpy as np
-
 from alphafold.common import protein
 from alphafold.data import pipeline
 from alphafold.data import templates
@@ -34,6 +32,7 @@ from alphafold.model import data
 from alphafold.model import config
 from alphafold.model import model
 from alphafold.relax import relax
+import numpy as np
 # Internal import (7716).
 
 flags.DEFINE_list('fasta_paths', None, 'Paths to FASTA files, each containing '

run_alphafold_test.py

@@ -18,10 +18,9 @@ import os
 
 from absl.testing import absltest
 from absl.testing import parameterized
+import run_alphafold
 import mock
 import numpy as np
-
-import run_alphafold
 # Internal import (7716).