TYP: Change | to 3.9 syntax, move stuff to _typing

pysindy/_typing.py

@@ -4,4 +4,11 @@
 # In python 3.12, use type statement
 # https://docs.python.org/3/reference/simple_stmts.html#the-type-statement
 NpFlt = np.floating[npt.NBitBase]
-Float2D = np.ndarray[tuple[int, int], np.dtype[NpFlt]]
+FloatDType = np.dtype[np.floating[npt.NBitBase]]
+Int1D = np.ndarray[tuple[int], np.dtype[np.int_]]
+Float1D = np.ndarray[tuple[int], FloatDType]
+Float2D = np.ndarray[tuple[int, int], FloatDType]
+Float3D = np.ndarray[tuple[int, int, int], FloatDType]
+Float4D = np.ndarray[tuple[int, int, int, int], FloatDType]
+Float5D = np.ndarray[tuple[int, int, int, int, int], FloatDType]
+FloatND = npt.NDArray[NpFlt]

pysindy/optimizers/base.py

@@ -9,7 +9,6 @@
 from typing import Tuple
 
 import numpy as np
-from numpy.typing import NBitBase
 from scipy import sparse
 from sklearn.base import BaseEstimator
 from sklearn.linear_model import LinearRegression
@@ -19,10 +18,10 @@
 from sklearn.utils.validation import check_X_y
 
 from .._typing import Float2D
+from .._typing import FloatDType
 from ..utils import AxesArray
 from ..utils import drop_nan_samples
 
-AnyFloat = np.dtype[np.floating[NBitBase]]
 NFeat = NewType("NFeat", int)
 NTarget = NewType("NTarget", int)
 
@@ -41,8 +40,8 @@ def _rescale_data(X, y, sample_weight):
 
 
 class _BaseOptimizer(BaseEstimator, abc.ABC):
-    coef_: np.ndarray[tuple[NTarget, NFeat], AnyFloat]
-    intercept_: np.ndarray[tuple[NTarget], AnyFloat]
+    coef_: np.ndarray[tuple[NTarget, NFeat], FloatDType]
+    intercept_: np.ndarray[tuple[NTarget], FloatDType]
 
     @property
     def complexity(self):
@@ -102,7 +101,7 @@ class BaseOptimizer(LinearRegression, _BaseOptimizer):
 
     max_iter: int
     normalize_columns: bool
-    initial_guess: Optional[np.ndarray[tuple[NTarget, NFeat], AnyFloat]]
+    initial_guess: Optional[np.ndarray[tuple[NTarget, NFeat], FloatDType]]
     copy_X: bool
     unbias: bool
 

pysindy/optimizers/trapping_sr3.py

@@ -5,33 +5,29 @@
 from itertools import repeat
 from math import comb
 from typing import cast
-from typing import NewType
 from typing import Optional
 from typing import TypeVar
 from typing import Union
 
 import cvxpy as cp
 import numpy as np
-from numpy.typing import NBitBase
 from numpy.typing import NDArray
 from sklearn.exceptions import ConvergenceWarning
 
+from .._typing import Float1D
+from .._typing import Float2D
+from .._typing import Float3D
+from .._typing import Float4D
+from .._typing import Float5D
+from .._typing import Int1D
 from ..feature_library.polynomial_library import n_poly_features
 from ..feature_library.polynomial_library import PolynomialLibrary
 from ..utils import reorder_constraints
-from .base import AnyFloat
+from .base import FloatDType
 from .base import NFeat
 from .base import NTarget
 from .constrained_sr3 import ConstrainedSR3
 
-Int1D = np.ndarray[tuple[int], np.dtype[np.int_]]
-Float1D = np.ndarray[tuple[int], AnyFloat]
-Float2D = np.ndarray[tuple[int, int], AnyFloat]
-Float3D = np.ndarray[tuple[int, int, int], AnyFloat]
-Float4D = np.ndarray[tuple[int, int, int, int], AnyFloat]
-Float5D = np.ndarray[tuple[int, int, int, int, int], AnyFloat]
-FloatND = NDArray[np.floating[NBitBase]]
-
 
 class EnstrophyMat:
     """Pre-compute some useful factors of an enstrophy matrix
@@ -601,7 +597,7 @@ def _solve_m_relax_and_split(
         self,
         trap_ctr: Float1D,
         prev_A: Float2D,
-        coef_sparse: np.ndarray[tuple[NFeat, NTarget], AnyFloat],
+        coef_sparse: np.ndarray[tuple[NFeat, NTarget], FloatDType],
     ) -> tuple[Float1D, Float2D]:
         r"""Updates the trap center
 
@@ -693,7 +689,7 @@ def _reduce(self, x, y):
             self.constraint_lhs = reorder_constraints(
                 self.constraint_lhs, n_features, output_order="feature"
             )
-        coef_sparse: np.ndarray[tuple[NFeat, NTarget], AnyFloat] = self.coef_.T
+        coef_sparse: np.ndarray[tuple[NFeat, NTarget], FloatDType] = self.coef_.T
 
         # Print initial values for each term in the optimization
         if self.verbose: