Fix import ordering.

PiperOrigin-RevId: 387085679
Change-Id: I73287fcd0a29e899543b64c596e306195a2f435e

alphafold/common/protein.py

@@ -13,15 +13,13 @@
 # limitations under the License.
 
 """Protein data type."""
+import dataclasses
 import io
 from typing import Any, Mapping, Optional
-
+from alphafold.common import residue_constants
 from Bio.PDB import PDBParser
-import dataclasses
 import numpy as np
 
-from alphafold.common import residue_constants
-
 FeatureDict = Mapping[str, np.ndarray]
 ModelOutput = Mapping[str, Any]  # Is a nested dict.
 

alphafold/common/protein_test.py

@@ -18,10 +18,9 @@ import os
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import numpy as np
-
 from alphafold.common import protein
 from alphafold.common import residue_constants
+import numpy as np
 # Internal import (7716).
 
 TEST_DATA_DIR = 'alphafold/common/testdata/'

alphafold/common/residue_constants.py

@@ -16,7 +16,7 @@
 
 import collections
 import functools
-from typing import Mapping, List, Tuple
+from typing import List, Mapping, Tuple
 
 import numpy as np
 import tree

alphafold/common/residue_constants_test.py

@@ -16,9 +16,8 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import numpy as np
-
 from alphafold.common import residue_constants
+import numpy as np
 
 
 class ResidueConstantsTest(parameterized.TestCase):

alphafold/data/mmcif_parsing.py

@@ -14,13 +14,13 @@
 
 """Parses the mmCIF file format."""
 import collections
+import dataclasses
 import io
 from typing import Any, Mapping, Optional, Sequence, Tuple
 
 from absl import logging
 from Bio import PDB
 from Bio.Data import SCOPData
-import dataclasses
 
 # Type aliases:
 ChainId = str

alphafold/data/parsers.py

@@ -14,11 +14,10 @@
 
 """Functions for parsing various file formats."""
 import collections
+import dataclasses
 import re
 import string
-from typing import Iterable, List, Optional, Sequence, Tuple, Dict
-
-import dataclasses
+from typing import Dict, Iterable, List, Optional, Sequence, Tuple
 
 DeletionMatrix = Sequence[Sequence[int]]
 

alphafold/data/pipeline.py

@@ -16,17 +16,15 @@
 
 import os
 from typing import Mapping, Optional, Sequence
-
-import numpy as np
-
-# Internal import (7716).
-
 from alphafold.common import residue_constants
 from alphafold.data import parsers
 from alphafold.data import templates
 from alphafold.data.tools import hhblits
 from alphafold.data.tools import hhsearch
 from alphafold.data.tools import jackhmmer
+import numpy as np
+
+# Internal import (7716).
 
 FeatureDict = Mapping[str, np.ndarray]
 

alphafold/data/templates.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Functions for getting templates and calculating template features."""
+import dataclasses
 import datetime
 import glob
 import os
@@ -20,15 +21,13 @@ import re
 from typing import Any, Dict, Mapping, Optional, Sequence, Tuple
 
 from absl import logging
-import dataclasses
-import numpy as np
-
-# Internal import (7716).
-
 from alphafold.common import residue_constants
 from alphafold.data import mmcif_parsing
 from alphafold.data import parsers
 from alphafold.data.tools import kalign
+import numpy as np
+
+# Internal import (7716).
 
 
 class Error(Exception):

alphafold/data/tools/hmmbuild.py

@@ -19,10 +19,8 @@ import re
 import subprocess
 
 from absl import logging
-
-# Internal import (7716).
-
 from alphafold.data.tools import utils
+# Internal import (7716).
 
 
 class Hmmbuild(object):

alphafold/data/tools/hmmsearch.py

@@ -19,10 +19,8 @@ import subprocess
 from typing import Optional, Sequence
 
 from absl import logging
-
-# Internal import (7716).
-
 from alphafold.data.tools import utils
+# Internal import (7716).
 
 
 class Hmmsearch(object):

alphafold/model/all_atom.py

@@ -34,15 +34,13 @@ the network to facilitate easier conversion to existing protein datastructures.
 """
 
 from typing import Dict, Optional
-
-import jax
-import jax.numpy as jnp
-import numpy as np
-
 from alphafold.common import residue_constants
 
 from alphafold.model import r3
 from alphafold.model import utils
+import jax
+import jax.numpy as jnp
+import numpy as np
 
 
 def squared_difference(x, y):

alphafold/model/all_atom_test.py

@@ -16,9 +16,9 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import numpy as np
 from alphafold.model import all_atom
 from alphafold.model import r3
+import numpy as np
 
 L1_CLAMP_DISTANCE = 10
 

alphafold/model/config.py

@@ -14,9 +14,8 @@
 """Model config."""
 
 import copy
-import ml_collections
-
 from alphafold.model.tf import shape_placeholders
+import ml_collections
 
 
 NUM_RES = shape_placeholders.NUM_RES

alphafold/model/data.py

@@ -17,11 +17,9 @@
 import io
 import os
 from typing import List
-
+from alphafold.model import utils
 import haiku as hk
 import numpy as np
-
-from alphafold.model import utils
 # Internal import (7716).
 
 

alphafold/model/features.py

@@ -15,14 +15,12 @@
 """Code to generate processed features."""
 import copy
 from typing import List, Mapping, Tuple
-
+from alphafold.model.tf import input_pipeline
+from alphafold.model.tf import proteins_dataset
 import ml_collections
 import numpy as np
 import tensorflow.compat.v1 as tf
 
-from alphafold.model.tf import input_pipeline
-from alphafold.model.tf import proteins_dataset
-
 FeatureDict = Mapping[str, np.ndarray]
 
 

alphafold/model/folding.py

@@ -16,13 +16,6 @@
 
 import functools
 from typing import Dict
-
-import haiku as hk
-import jax
-import jax.numpy as jnp
-import ml_collections
-import numpy as np
-
 from alphafold.common import residue_constants
 from alphafold.model import all_atom
 from alphafold.model import common_modules
@@ -30,6 +23,11 @@ from alphafold.model import prng
 from alphafold.model import quat_affine
 from alphafold.model import r3
 from alphafold.model import utils
+import haiku as hk
+import jax
+import jax.numpy as jnp
+import ml_collections
+import numpy as np
 
 
 def squared_difference(x, y):

alphafold/model/layer_stack_test.py

@@ -17,14 +17,13 @@
 import functools
 from absl.testing import absltest
 from absl.testing import parameterized
+from alphafold.model import layer_stack
 import haiku as hk
 import jax
 import jax.numpy as jnp
 import numpy as np
 import scipy
 
-from alphafold.model import layer_stack
-
 
 # Suffixes applied by Haiku for repeated module names.
 suffixes = [''] + [f'_{i}' for i in range(1, 100)]

alphafold/model/lddt_test.py

@@ -16,8 +16,8 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import numpy as np
 from alphafold.model import lddt
+import numpy as np
 
 
 class LddtTest(parameterized.TestCase, absltest.TestCase):

alphafold/model/model.py

@@ -16,6 +16,9 @@
 from typing import Any, Mapping, Optional, Union
 
 from absl import logging
+from alphafold.common import confidence
+from alphafold.model import features
+from alphafold.model import modules
 import haiku as hk
 import jax
 import ml_collections
@@ -23,10 +26,6 @@ import numpy as np
 import tensorflow.compat.v1 as tf
 import tree
 
-from alphafold.common import confidence
-from alphafold.model import features
-from alphafold.model import modules
-
 
 def get_confidence_metrics(
     prediction_result: Mapping[str, Any]) -> Mapping[str, Any]:
@@ -140,4 +139,3 @@ class RunModel:
     logging.info('Output shape was %s',
                  tree.map_structure(lambda x: x.shape, result))
     return result
-

alphafold/model/modules.py

@@ -17,11 +17,6 @@
 The structure generation code is in 'folding.py'.
 """
 import functools
-
-import haiku as hk
-import jax
-import jax.numpy as jnp
-
 from alphafold.common import residue_constants
 from alphafold.model import all_atom
 from alphafold.model import common_modules
@@ -32,6 +27,9 @@ from alphafold.model import mapping
 from alphafold.model import prng
 from alphafold.model import quat_affine
 from alphafold.model import utils
+import haiku as hk
+import jax
+import jax.numpy as jnp
 
 
 def softmax_cross_entropy(logits, labels):