isort fix

src/nemos/__init__.py

@@ -14,5 +14,5 @@
     styles,
     tree_utils,
     type_casting,
-    utils,
+    utils
 )

src/nemos/_documentation_utils/__init__.py

@@ -19,5 +19,5 @@
     plot_rates_and_smoothed_counts,
     plot_weighted_sum_basis,
     run_animation,
-    tuning_curve_plot,
+    tuning_curve_plot
 )

src/nemos/base_regressor.py

@@ -9,11 +9,10 @@
 from copy import deepcopy
 from typing import Any, Dict, NamedTuple, Optional, Tuple, Union
 
-import jaxopt
-from numpy.typing import ArrayLike, NDArray
-
 import jax
 import jax.numpy as jnp
+import jaxopt
+from numpy.typing import ArrayLike, NDArray
 
 from . import solvers, utils, validation
 from ._regularizer_builder import AVAILABLE_REGULARIZERS, create_regularizer

src/nemos/convolve.py

@@ -8,10 +8,9 @@
 from functools import partial
 from typing import Any, Literal, Optional
 
-from numpy.typing import ArrayLike, NDArray
-
 import jax
 import jax.numpy as jnp
+from numpy.typing import ArrayLike, NDArray
 
 from . import type_casting, utils
 

src/nemos/glm.py

@@ -7,13 +7,12 @@
 from functools import wraps
 from typing import Any, Callable, Literal, NamedTuple, Optional, Tuple, Union
 
+import jax
+import jax.numpy as jnp
 import jaxopt
 from numpy.typing import ArrayLike
 from scipy.optimize import root
 
-import jax
-import jax.numpy as jnp
-
 from . import observation_models as obs
 from . import tree_utils, validation
 from .base_regressor import BaseRegressor

src/nemos/observation_models.py

@@ -3,10 +3,9 @@
 import abc
 from typing import Callable, Literal, Union
 
-from numpy.typing import NDArray
-
 import jax
 import jax.numpy as jnp
+from numpy.typing import NDArray
 
 from . import utils
 from .base_class import Base

src/nemos/regularizer.py

@@ -9,11 +9,10 @@
 import abc
 from typing import Callable, Tuple, Union
 
-import jaxopt
-from numpy.typing import NDArray
-
 import jax
 import jax.numpy as jnp
+import jaxopt
+from numpy.typing import NDArray
 
 from . import tree_utils
 from .base_class import Base

src/nemos/simulation.py

@@ -2,13 +2,12 @@
 
 from typing import Callable, Tuple, Union
 
+import jax
+import jax.numpy as jnp
 import numpy as np
 import scipy.stats as sts
 from numpy.typing import NDArray
 
-import jax
-import jax.numpy as jnp
-
 from . import convolve, validation
 from .pytrees import FeaturePytree
 

src/nemos/solvers.py

@@ -1,14 +1,13 @@
 from functools import partial
 from typing import Callable, NamedTuple, Optional, Union
 
-from jaxopt import OptStep
-from jaxopt._src import loop
-from jaxopt.prox import prox_none
-
 import jax
 import jax.flatten_util
 import jax.numpy as jnp
 from jax import grad, jit, lax, random
+from jaxopt import OptStep
+from jaxopt._src import loop
+from jaxopt.prox import prox_none
 
 from .tree_utils import tree_add_scalar_mul, tree_l2_norm, tree_slice, tree_sub
 from .typing import KeyArrayLike, Pytree

src/nemos/type_casting.py

@@ -12,13 +12,12 @@
 from functools import wraps
 from typing import Any, Callable, List, Literal, Optional, Type, Union
 
+import jax
+import jax.numpy as jnp
 import numpy as np
 import pynapple as nap
 from numpy.typing import NDArray
 
-import jax
-import jax.numpy as jnp
-
 from . import tree_utils
 
 _NAP_TIME_PRECISION = 10 ** (-nap.nap_config.time_index_precision)

src/nemos/typing.py

@@ -2,9 +2,8 @@
 
 from typing import Any, Callable, NamedTuple, Tuple, Union
 
-import jaxopt
-
 import jax.numpy as jnp
+import jaxopt
 from jax._src.typing import ArrayLike
 
 from .pytrees import FeaturePytree

src/nemos/utils.py

@@ -3,11 +3,10 @@
 import warnings
 from typing import Any, Callable, List, Literal, Union
 
-import numpy as np
-from numpy.typing import NDArray
-
 import jax
 import jax.numpy as jnp
+import numpy as np
+from numpy.typing import NDArray
 
 from .tree_utils import pytree_map_and_reduce
 from .type_casting import is_numpy_array_like, support_pynapple

src/nemos/validation.py

@@ -3,10 +3,9 @@
 import warnings
 from typing import Any, Optional, Union
 
-from numpy.typing import DTypeLike, NDArray
-
 import jax
 import jax.numpy as jnp
+from numpy.typing import DTypeLike, NDArray
 
 from .pytrees import FeaturePytree
 from .tree_utils import get_valid_multitree, pytree_map_and_reduce