fixed linting errors that came with new range of python versions

Author: LukeTheHecker

invert/adapters/context_lstm.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 

invert/auto_inverse/auto_inverse.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Optional, Union
 
 import mne
 import numpy as np
@@ -88,15 +87,15 @@ def generate_summary(self) -> str:
 
         return summary
 
-    def get_recommended_solver(self) -> Optional[str]:
+    def get_recommended_solver(self) -> str | None:
         """Get the recommended solver name."""
         return self.recommended_solver
 
 
 class AutoInverse:
     def __init__(
         self,
-        prior: Union[str, PriorEnum] = "patch",
+        prior: str | PriorEnum = "patch",
         snr="auto",
         alpha="auto",
         n_samples=100,
@@ -148,16 +147,16 @@ def __init__(
         self.n_samples = n_samples
         self.n_timepoints = n_timepoints
         self.verbose = verbose
-        self.report: Optional[Report] = None
+        self.report: Report | None = None
         self.sim_params = None
         self.solvers_to_test = None
         self.simulation_config = None
 
     def fit(
         self,
-        data: Union[EvokedArray, Evoked],
+        data: EvokedArray | Evoked,
         forward: Forward,
-    ) -> Optional[Report]:
+    ) -> Report | None:
         """
         Fit the auto inverse model to determine the best solver.
 

invert/benchmark/datasets.py

@@ -1,13 +1,12 @@
-from typing import Union
 
 from pydantic import BaseModel
 
 
 class DatasetConfig(BaseModel):
     name: str
     description: str
-    n_sources: Union[int, tuple[int, int]]
-    n_orders: Union[int, tuple[int, int]]
+    n_sources: int | tuple[int, int]
+    n_orders: int | tuple[int, int]
     snr_range: tuple[float, float]
     n_timepoints: int
     n_samples: int = 50

invert/benchmark/runner.py

@@ -87,13 +87,18 @@
     "SAM": ("invert.solvers.beamformers.sam", "SolverSAM"),
     # "EBB": ("invert.solvers.beamformers.ebb", "SolverEBB"),
     "ESMV": ("invert.solvers.beamformers.esmv", "SolverESMV"),
+    "ESMV-MVPURE": ("invert.solvers.beamformers.esmv_mvpure", "SolverESMVMVPURE"),
     "ESMV2": ("invert.solvers.beamformers.esmv2", "SolverESMV2"),
     "ESMV3": ("invert.solvers.beamformers.esmv3", "SolverESMV3"),
     "DeblurFlexESMV": (
         "invert.solvers.beamformers.deblur_flex_esmv",
         "SolverDeblurFlexESMV",
     ),
     "FlexESMV": ("invert.solvers.beamformers.flex_esmv", "SolverFlexESMV"),
+    "FlexESMV-MVPURE": (
+        "invert.solvers.beamformers.flex_esmv_mvpure",
+        "SolverFlexESMVMVPURE",
+    ),
     "SafeFlexESMV": ("invert.solvers.beamformers.safe_flex_esmv", "SolverSafeFlexESMV"),
     "SharpFlexESMV": (
         "invert.solvers.beamformers.sharp_flex_esmv",
@@ -226,12 +231,16 @@ def _expects_simulation_config(solver_cls: type[BaseSolver]) -> bool:
 
 
 def _default_nn_batch_size(forward: mne.Forward) -> int:
-    """Default ANN training batch size: 2x number of dipoles in the source model."""
+    """Default ANN training batch size: ~12x number of dipoles (min 4096).
+
+    Larger batches provide more diverse training data per epoch, which
+    significantly improves ANN solver quality (see ann_testbench round-2/3).
+    """
     try:
         n_dipoles = int(forward["sol"]["data"].shape[1])
     except Exception:
         return 1
-    return max(1, 2 * n_dipoles)
+    return max(4096, 12 * n_dipoles)
 
 
 def _generate_solver_categories() -> dict[str, list[str]]:
@@ -579,7 +588,7 @@ def run(self) -> list[BenchmarkResult]:
                             update={"batch_size": _default_nn_batch_size(self.forward)}
                         )
                         logger.info(
-                            "ANN training batch_size=%d (default=2*n_dipoles) for %s",
+                            "ANN training batch_size=%d for %s",
                             int(train_sim_config.batch_size),
                             solver_name,
                         )

invert/benchmark/visualize.py

@@ -1,6 +1,5 @@
 import json
 from pathlib import Path
-from typing import Optional, Union
 
 import matplotlib.pyplot as plt
 import numpy as np
@@ -36,9 +35,9 @@ def _solver_to_category(solver_name: str) -> str:
 
 
 def visualize_results(
-    results_path_or_data: Union[str, Path, list[BenchmarkResult]],
-    metrics: Optional[list[str]] = None,
-    save_path: Optional[Union[str, Path]] = None,
+    results_path_or_data: str | Path | list[BenchmarkResult],
+    metrics: list[str] | None = None,
+    save_path: str | Path | None = None,
 ) -> list[plt.Figure]:
     if isinstance(results_path_or_data, (str, Path)):
         path = Path(results_path_or_data)
@@ -120,7 +119,7 @@ def visualize_results(
 
         # Build legend grouped by category, placed below the plot
         handles, labels = ax.get_legend_handles_labels()
-        label_to_handle = dict(zip(labels, handles))
+        label_to_handle = dict(zip(labels, handles, strict=False))
 
         legend_handles = []
         legend_labels = []

invert/config.py

@@ -43,17 +43,20 @@
     "Source-MAP-MSP",
     "MVAB",
     "LCMV",
+    "LCMV-MVPURE",
     "DICS",
     "SMV",
     "WNMV",
     "HOCMV",
     "ESMV",
+    "ESMV-MVPURE",
     "MCMV",
     "HOCMCMV",
     "ReciPSIICOS",
     "SAM",
     "Adapt-Flex-ESMV",
     "Flex-ESMV",
+    "Flex-ESMV-MVPURE",
     "Flex-ESMV2",
     "Deblur-Flex-ESMV",
     "Safe-Flex-ESMV",

invert/denoising/ASR.py

@@ -1,5 +1,4 @@
 import logging
-from typing import Union
 
 import mne
 import numpy as np
@@ -116,7 +115,7 @@ def run(self, data: np.ndarray, calibration=None) -> np.ndarray:
 
     # -------------------------- mne wrapper --------------------------
 
-    def run_mne(self, mne_obj: Union[mne.io.BaseRaw, mne.Epochs, mne.Evoked]):
+    def run_mne(self, mne_obj: mne.io.BaseRaw | mne.Epochs | mne.Evoked):
         """
         Run ASR on an MNE object and return a cleaned copy.
         """

invert/denoising/lgsp.py

@@ -1,13 +1,13 @@
 from dataclasses import dataclass
-from typing import Literal, Optional, Union
+from typing import Literal
 
 import mne
 import numpy as np
 from scipy.linalg import eigh
 from sklearn.covariance import OAS
 
 ArrayLike = np.ndarray
-MNEObj = Union[mne.io.BaseRaw, mne.Epochs, mne.Evoked, mne.EvokedArray]
+MNEObj = mne.io.BaseRaw | mne.Epochs | mne.Evoked | mne.EvokedArray
 
 
 # ============================
@@ -51,12 +51,12 @@ def _inv_sqrt_spd(Q: np.ndarray, d: np.ndarray) -> np.ndarray:
 
 @dataclass
 class DenoiserBase:
-    sfreq: Optional[float] = None
+    sfreq: float | None = None
     window_size: float = 0.5  # seconds
     step_frac: float = 0.5  # overlap = 1 - step_frac
     use_windows: bool = True  # set False to process full signal
 
-    def _ensure_sfreq(self, sfreq: Optional[float]):
+    def _ensure_sfreq(self, sfreq: float | None):
         if self.sfreq is None:
             if sfreq is None:
                 raise ValueError(
@@ -65,7 +65,7 @@ def _ensure_sfreq(self, sfreq: Optional[float]):
             self.sfreq = sfreq
 
     # ---- public API on NumPy ----
-    def run(self, X: ArrayLike, sfreq: Optional[float] = None) -> ArrayLike:
+    def run(self, X: ArrayLike, sfreq: float | None = None) -> ArrayLike:
         """
         X shape must be (n_channels, n_times).
         """
@@ -137,12 +137,12 @@ def _process_window(self, Xw: ArrayLike) -> ArrayLike:
 
 @dataclass
 class LGSP(DenoiserBase):
-    L: Optional[np.ndarray] = None  # (m, n_sources)
+    L: np.ndarray | None = None  # (m, n_sources)
     sigma_prior: Literal["identity"] = "identity"
     lambda_ref: float = 1e-3  # ridge on C_ref
     alpha_model: float = 0.2  # blend with scaled I: Cref'=(1-a)Cref+a*trace(Cref)/m*I
-    rank: Optional[int] = None  # keep r smallest generalized eigenvalues
-    tau: Optional[float] = None  # or threshold on generalized eigenvalues
+    rank: int | None = None  # keep r smallest generalized eigenvalues
+    tau: float | None = None  # or threshold on generalized eigenvalues
     center: bool = True
     shrink_data: bool = True
 
@@ -164,7 +164,7 @@ def _build_Cref(self, n_ch: int) -> np.ndarray:
         return Cref
 
     def _gevd_brain_projector(
-        self, Cx: np.ndarray, Cref: np.ndarray, r: Optional[int], tau: Optional[float]
+        self, Cx: np.ndarray, Cref: np.ndarray, r: int | None, tau: float | None
     ) -> np.ndarray:
         """
         Return the sensor-space projector P onto the 'brain-like' subspace
@@ -214,7 +214,7 @@ def _process_window(self, Xw: ArrayLike) -> ArrayLike:
 
 @dataclass
 class SRB(DenoiserBase):
-    L: Optional[np.ndarray] = None  # (m, n_sources) required for meaningful use
+    L: np.ndarray | None = None  # (m, n_sources) required for meaningful use
     mu: float = 1e-1  # Tikhonov on sources (||S||^2)
     beta: float = 0.5  # blend factor: X_clean = beta*L*W*X + (1-beta)*X
     center: bool = True

invert/ensemble/ensemble.py

@@ -46,7 +46,7 @@ def apply_inverse_operator(self, mne_obj, *args, **kwargs):
 
         logger.info("Neg Log Likelihoods:")
         for solver_name, neg_log_likelihood in zip(
-            self.solver_names, self.neg_log_likelihoods
+            self.solver_names, self.neg_log_likelihoods, strict=False
         ):
             logger.info(f"{solver_name}: {neg_log_likelihood}")
         logger.info(f"Final likelihood: {final_neg_log_likelihood}")

invert/evaluate/evaluate.py

@@ -259,7 +259,7 @@ def nmse(y_true, y_pred):
     if np.any(np.isnan(y_true)) or np.any(np.isnan(y_pred)):
         return np.nan
     error = np.zeros(y_true.shape[1])
-    for i, (y_true_slice, y_pred_slice) in enumerate(zip(y_true.T, y_pred.T)):
+    for i, (y_true_slice, y_pred_slice) in enumerate(zip(y_true.T, y_pred.T, strict=False)):
         error[i] = np.mean(
             (
                 (y_true_slice / abs(y_true_slice).max())
@@ -274,7 +274,7 @@ def corr(y_true, y_pred):
     if np.any(np.isnan(y_true)) or np.any(np.isnan(y_pred)):
         return np.nan
     error = np.zeros(y_true.shape[1])
-    for i, (y_true_slice, y_pred_slice) in enumerate(zip(y_true.T, y_pred.T)):
+    for i, (y_true_slice, y_pred_slice) in enumerate(zip(y_true.T, y_pred.T, strict=False)):
         error[i] = pearsonr(y_true_slice, y_pred_slice)[0]
     return error