From 73efa056a7845ae01be26e063b8077fa84576b06 Mon Sep 17 00:00:00 2001
From: Theo Brown <7982453+theo-brown@users.noreply.github.com>
Date: Wed, 4 Dec 2024 11:16:20 +0000
Subject: [PATCH] Minor linting fixes

---
 torax/post_processing.py                         |  2 --
 torax/sources/tests/bootstrap_current_source.py  |  3 +--
 torax/sources/tests/electron_cyclotron_source.py | 10 +---------
 torax/sources/tests/ion_cyclotron_source.py      |  7 +++----
 torax/sources/tests/qei_source.py                |  2 --
 torax/state.py                                   |  4 +---
 6 files changed, 6 insertions(+), 22 deletions(-)

diff --git a/torax/post_processing.py b/torax/post_processing.py
index ee75cfac..30964c5e 100644
--- a/torax/post_processing.py
+++ b/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
diff --git a/torax/sources/tests/bootstrap_current_source.py b/torax/sources/tests/bootstrap_current_source.py
index 37f8947f..079abdd4 100644
--- a/torax/sources/tests/bootstrap_current_source.py
+++ b/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
diff --git a/torax/sources/tests/electron_cyclotron_source.py b/torax/sources/tests/electron_cyclotron_source.py
index 6bbc88ad..c73497e9 100644
--- a/torax/sources/tests/electron_cyclotron_source.py
+++ b/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
diff --git a/torax/sources/tests/ion_cyclotron_source.py b/torax/sources/tests/ion_cyclotron_source.py
index e9249b5c..c7dd2750 100644
--- a/torax/sources/tests/ion_cyclotron_source.py
+++ b/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
diff --git a/torax/sources/tests/qei_source.py b/torax/sources/tests/qei_source.py
index d7947bff..79a46cdd 100644
--- a/torax/sources/tests/qei_source.py
+++ b/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
diff --git a/torax/state.py b/torax/state.py
index c342ed9e..d4ca8d3f 100644
--- a/torax/state.py
+++ b/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