Minor linting fixes

torax/post_processing.py

@@ -17,10 +17,8 @@
 # In torax/post_processing.py
 
 import dataclasses
-
 import jax
 from jax import numpy as jnp
-
 from torax import array_typing
 from torax import constants
 from torax import geometry

torax/sources/tests/bootstrap_current_source.py

@@ -14,10 +14,9 @@
 
 """Tests for bootstrap_current_source."""
 
+from absl.testing import absltest
 import jax.numpy as jnp
 import numpy as np
-from absl.testing import absltest
-
 from torax import geometry
 from torax.sources import bootstrap_current_source
 from torax.sources import runtime_params as runtime_params_lib

torax/sources/tests/electron_cyclotron_source.py

@@ -159,15 +159,7 @@ def test_invalid_source_types_raise_errors(self):
   def test_extraction_of_relevant_profile_from_output(self):
     """Tests that the relevant profile is extracted from the output."""
     geo = geometry.build_circular_geometry()
-    # pylint: disable=missing-kwoa
-    source_builder = self._source_class_builder()  # pytype: disable=missing-parameter
-    # pylint: enable=missing-kwoa
-    source_models_builder = source_models_lib.SourceModelsBuilder(
-        {'foo': source_builder},
-    )
-    source_models = source_models_builder()
-    source = source_models.sources['foo']
-    self.assertIsInstance(source, source_lib.Source)
+    source = self._source_class()
     cell = source_lib.ProfileType.CELL.get_profile_shape(geo)
     fake_profile = jnp.stack((jnp.ones(cell), 2 * jnp.ones(cell)))
     # Check TEMP_EL and PSI are modified

torax/sources/tests/ion_cyclotron_source.py

@@ -18,13 +18,12 @@
 import os
 from unittest import mock
 
-import chex
-import jax
-import numpy as np
 from absl.testing import absltest
 from absl.testing import parameterized
+import chex
+import jax
 from jax import numpy as jnp
-
+import numpy as np
 from torax import core_profile_setters
 from torax import geometry
 from torax.config import runtime_params as general_runtime_params

torax/sources/tests/qei_source.py

@@ -15,9 +15,7 @@
 """Tests for qei_source."""
 
 import dataclasses
-
 from absl.testing import absltest
-
 from torax import core_profile_setters
 from torax import geometry
 from torax.config import runtime_params as general_runtime_params

torax/state.py

@@ -18,13 +18,11 @@
 
 import dataclasses
 import enum
-from typing import Any
-from typing import Optional
+from typing import Any, Optional
 
 import chex
 import jax
 from jax import numpy as jnp
-
 from torax import array_typing
 from torax import geometry
 from torax.config import config_args