[1034] lint

Author: simone99n

src/weathergen/datasets/data_reader_eobs.py

@@ -11,14 +11,12 @@
 from pathlib import Path
 from typing import override
 
-import dask.array as da
 import numpy as np
 import xarray as xr
 from numpy.typing import NDArray
 
 from weathergen.datasets.data_reader_base import (
     DataReaderTimestep,
-    DTRange,
     ReaderData,
     TimeWindowHandler,
     TIndex,
@@ -28,14 +26,15 @@
 
 _logger = logging.getLogger(__name__)
 
+
 # TODO make this datareader works with multiple datasets in ZARR format
 class DataReaderEObs(DataReaderTimestep):
     """
     Data reader for gridded Zarr datasets with regular lat/lon structure.
-    
+
     This reader handles datasets stored as Zarr with dimensions (time, latitude, longitude)
     and converts the gridded data to point-wise format required by the framework.
-    
+
     The reader implements lazy initialization to work efficiently with multiple dataloader workers.
     """
 
@@ -76,7 +75,7 @@ def __init__(
         self.geoinfo_channels = []
         self.geoinfo_idx = []
         self.properties = {}
-        
+
         # Grid properties
         self.latitudes: NDArray | None = None
         self.longitudes: NDArray | None = None
@@ -104,12 +103,7 @@ def _lazy_init(self) -> None:
 
         try:
             # Open the Zarr dataset with xarray
-            self.ds = xr.open_zarr(
-                self._filename,
-                consolidated=True,
-                chunks=None,
-                zarr_format=2
-                )
+            self.ds = xr.open_zarr(self._filename, consolidated=True, chunks=None, zarr_format=2)
         except Exception as e:
             name = self._stream_info["name"]
             _logger.error(f"Failed to open {name} at {self._filename}: {e}")
@@ -123,14 +117,9 @@ def _lazy_init(self) -> None:
         data_end_time = np.datetime64(time_coord[-1])
 
         # Check if dataset overlaps with requested time window
-        if (
-            self._tw_handler.t_start >= data_end_time
-            or self._tw_handler.t_end <= data_start_time
-        ):
+        if self._tw_handler.t_start >= data_end_time or self._tw_handler.t_end <= data_start_time:
             name = self._stream_info["name"]
-            _logger.warning(
-                f"{name} is not supported over data loader window. Stream is skipped."
-            )
+            _logger.warning(f"{name} is not supported over data loader window. Stream is skipped.")
             self.init_empty()
             self._initialized = True
             return
@@ -156,9 +145,7 @@ def _lazy_init(self) -> None:
         )
 
         # Calculate valid time range indices
-        time_mask = (time_coord >= self._tw_handler.t_start) & (
-            time_coord < self._tw_handler.t_end
-        )
+        time_mask = (time_coord >= self._tw_handler.t_start) & (time_coord < self._tw_handler.t_end)
         self.len = int(np.sum(time_mask))
 
         if self.len <= 0:
@@ -183,9 +170,7 @@ def _lazy_init(self) -> None:
                 f"Longitude values outside valid range [-180, 180] in stream "
                 f"'{self._stream_info['name']}'. Converting from [0, 360] format."
             )
-            self.longitudes = ((self.longitudes + 180.0) % 360.0 - 180.0).astype(
-                np.float32
-            )
+            self.longitudes = ((self.longitudes + 180.0) % 360.0 - 180.0).astype(np.float32)
 
         self.n_lat = len(self.latitudes)
         self.n_lon = len(self.longitudes)
@@ -405,9 +390,9 @@ def _get(self, idx: TIndex, channels_idx: list[int]) -> ReaderData:
 
         # Create coordinate grid
         lon_grid, lat_grid = np.meshgrid(self.longitudes, self.latitudes)
-        coords_single = np.stack(
-            [lat_grid.flatten(), lon_grid.flatten()], axis=1
-        ).astype(np.float32)
+        coords_single = np.stack([lat_grid.flatten(), lon_grid.flatten()], axis=1).astype(
+            np.float32
+        )
 
         # Repeat coordinates for each timestep
         coords = np.tile(coords_single, (len(t_idxs), 1))
@@ -427,4 +412,4 @@ def _get(self, idx: TIndex, channels_idx: list[int]) -> ReaderData:
 
         check_reader_data(rd, dtr)
 
-        return rd
+        return rd

src/weathergen/datasets/data_reader_fesom.py

@@ -43,7 +43,7 @@ def __init__(
         tw_handler: TimeWindowHandler,
         filename: Path,
         stream_info: dict,
-        ) -> None:
+    ) -> None:
         # Store configuration but DO NOT open files here
         self.filenames = sorted(glob.glob(str(filename) + "/*"))
         self._tw_handler = tw_handler

src/weathergen/datasets/icon_dataset.py

@@ -63,7 +63,7 @@ def __init__(
         step_hrs: int,
         filename: Path,
         stream_info: dict,
-        ):
+    ):
         self.len_hrs = len_hrs
 
         format_str = "%Y%m%d%H%M"

src/weathergen/datasets/multi_stream_data_sampler.py

@@ -21,9 +21,9 @@
     TIndex,
     str_to_datetime64,
 )
+from weathergen.datasets.data_reader_eobs import DataReaderEObs
 from weathergen.datasets.data_reader_fesom import DataReaderFesom
 from weathergen.datasets.data_reader_obs import DataReaderObs
-from weathergen.datasets.data_reader_eobs import DataReaderEObs
 from weathergen.datasets.icon_dataset import IconDataset
 from weathergen.datasets.masking import Masker
 from weathergen.datasets.stream_data import StreamData, spoof