Added multi satellite based fetching for training

main.py

@@ -1,54 +1,147 @@
 import os
 import logging
 import torch
+
 from src.config.settings import load_settings
 from src.utils import setup_logging
-from src.data.s3_manager import S3Manager
+from src.data.data_manager import DataManager
 from src.model.ifnet import IFNet
 from src.training.trainer import Trainer
 
+
 def main():
     logger = setup_logging()
     settings = load_settings()
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    
-    logger.info("Starting Autonomous Satellite Interpolation Pipeline...")
-    
-    s3_manager = S3Manager(settings)
+
+    device = torch.device(
+        "cuda" if torch.cuda.is_available() else "cpu"
+    )
+
+    logger.info(
+        "Starting Universal Satellite Interpolation Pipeline..."
+    )
+
+    # Universal Data Manager
+    data_manager = DataManager(settings)
+
+    # Model init
     model = IFNet()
-    
-    # Load previous checkpoint if it exists
+
+    # Resume checkpoint if exists
     checkpoint_path = settings.training.load_model_path
-    
+
     if os.path.exists(checkpoint_path):
-        model.load_state_dict(torch.load(checkpoint_path, map_location=device))
-        logger.info(f"✅ Loaded EXISTING model from Kaggle Input. Resuming training!")
+        model.load_state_dict(
+            torch.load(
+                checkpoint_path,
+                map_location=device
+            )
+        )
+
+        logger.info(
+            f"Loaded existing checkpoint: {checkpoint_path}"
+        )
     else:
-        logger.warning(f"⚠️ Checkpoint not found at {checkpoint_path}. Starting from scratch.")
-    trainer = Trainer(settings, model, device)
-    
-    # 🤖 100% Config-Driven Chunk Generation
+        logger.warning(
+            f"No checkpoint found at {checkpoint_path}. "
+            f"Starting from scratch."
+        )
+
+    trainer = Trainer(
+        settings=settings,
+        model=model,
+        device=device
+    )
+
+    # Config-driven chunk generation
+    sat_type = settings.data.satellite_type.lower()
     year = settings.data.year
-    start_day = settings.data.start_day
-    end_day = settings.data.end_day
-    prefix_type = settings.data.prefix_type  # 🚨 Config se prefix uthaya
-    
-    # 🚨 Hardcoded RadC hata diya!
-    chunks = [f"{prefix_type}/{year}/{day:03d}/" for day in range(start_day, end_day + 1)]
-    
+    start_day = settings.data.start_unit
+    end_day = settings.data.end_unit
+    prefix_type = settings.data.prefix_type
+
+    chunks = []
+
+    if sat_type == "goes":
+        # Example:
+        # ABI-L1b-RadC/2026/150/
+        chunks = [
+            f"{prefix_type}/{year}/{day:03d}/"
+            for day in range(
+                start_day,
+                end_day + 1
+            )
+        ]
+
+    elif sat_type == "himawari":
+        # Example:
+        # AHI-L1b-FLDK/2026/06/18/
+        month = settings.data.month
+
+        chunks = [
+            f"{prefix_type}/{year}/{month:02d}/{day:02d}/"
+            for day in range(
+                start_day,
+                end_day + 1
+            )
+        ]
+
+    else:
+        raise ValueError(
+            f"Unsupported satellite type: {sat_type}"
+        )
+
+    # Main chunk loop
     for chunk_idx, chunk in enumerate(chunks):
-        logger.info(f"=== Fetching and Processing Chunk {chunk_idx + 1}/{len(chunks)}: {chunk} ===")
-        
-        s3_manager.download_chunk(chunk)
-        
-        for epoch in range(1, settings.training.epochs + 1):
-            logger.info(f"--- Chunk {chunk_idx + 1} | Starting Epoch {epoch}/{settings.training.epochs} ---")
-            trainer.train_chunk(settings.data.download_dir, epoch)
-            trainer.save_checkpoint("latest_model.pth")
-            
-        s3_manager.purge_chunk()
-    trainer.shutdown() 
-    logger.info("🔥 Config-Driven Multi-Month Training Complete!")
-            
+        logger.info(
+            f"=== Processing Chunk "
+            f"{chunk_idx + 1}/{len(chunks)}: {chunk} ==="
+        )
+
+        # Fetch -> Standardize -> Crop -> Save .pt triplets
+        data_manager.process_chunk(chunk)
+
+        # Train on generated triplets
+        for epoch in range(
+            1,
+            settings.training.epochs + 1
+        ):
+            logger.info(
+                f"--- Chunk {chunk_idx + 1} | "
+                f"Epoch {epoch}/{settings.training.epochs} ---"
+            )
+
+            trainer.train_chunk(
+                settings.data.download_dir,
+                epoch
+            )
+
+            trainer.save_checkpoint(
+                "latest_model.pth"
+            )
+
+        # Purge .pt files after training chunk
+        logger.info(
+            "Purging processed .pt triplets..."
+        )
+
+        for f in os.listdir(
+            settings.data.download_dir
+        ):
+            if f.endswith(".pt"):
+                os.remove(
+                    os.path.join(
+                        settings.data.download_dir,
+                        f
+                    )
+                )
+
+    trainer.shutdown()
+
+    logger.info(
+        "Universal multi-satellite training complete."
+    )
+
+
 if __name__ == "__main__":
     main()

pyproject.toml

@@ -10,6 +10,7 @@ dependencies = [
     "numpy>=2.4.6",
     "pytest>=9.1.0",
     "pyyaml>=6.0.3",
+    "satpy>=0.60.0",
     "tensorboard>=2.20.0",
     "torch>=2.12.0",
     "torchvision>=0.27.0",

src/config/config.yaml

@@ -2,14 +2,33 @@ training:
   epochs: 100
   batch_size: 4
   learning_rate: 0.0001
+  weight_decay: 0.001
+  num_workers: 4
   checkpoints_dir: "/kaggle/working/checkpoints"
   load_model_path: ""
+
 data:
-  s3_bucket: "noaa-goes19" # Agar GOES-19 chahiye toh yahan update kar lena
-  download_dir: "/kaggle/working/goes_data"
-  prefix_type: "ABI-L1b-RadF"  # 🚨 Naya prefix variable
+  satellite_type: "himawari"
+  s3_bucket: "noaa-himawari9"
+  download_dir: "/kaggle/working/data"
+
+  # Satellite path prefix
+  prefix_type: "AHI-L1b-FLDK"
+
+  # Universal time config
   year: 2024
-  start_day: 264   
-  end_day: 294     
+  month: 9
+  start_unit: 21
+  end_unit: 30
+
+  # Temporal triplet spacing
   frame_step: 3
-  crop_size: 512
+
+  # Crop logic
+  crop_size: 512
+  crop_stride_divisor: 4
+  static_motion_threshold: 0.005
+
+  # Brightness temperature normalization
+  min_bt: 180.0
+  max_bt: 330.0

src/config/settings.py

@@ -1,35 +1,60 @@
 from dataclasses import dataclass
+from typing import Optional
 import yaml
 
+
 @dataclass
 class TrainingConfig:
     epochs: int
     batch_size: int
     learning_rate: float
+    weight_decay: float
+    num_workers: int
     checkpoints_dir: str
     load_model_path: str = ""
 
+
 @dataclass
 class DataConfig:
+    satellite_type: str
     s3_bucket: str
     download_dir: str
-    prefix_type: str   # 🚨 Yahan prefix_type add kar diya
+    prefix_type: str
+
     year: int
-    start_day: int
-    end_day: int
-    frame_step: int
-    crop_size: int
+    month: Optional[int] = None
+
+    start_unit: int = 1
+    end_unit: int = 1
+
+    frame_step: int = 1
+
+    crop_size: int = 256
+    crop_stride_divisor: int = 4
+    static_motion_threshold: float = 0.005
+
+    min_bt: float = 180.0
+    max_bt: float = 330.0
+
 
 @dataclass
 class Settings:
     training: TrainingConfig
     data: DataConfig
 
-def load_settings(config_path: str = "src/config/config.yaml") -> Settings:
-    with open(config_path, 'r') as file:
+
+def load_settings(
+    config_path: str = "src/config/config.yaml"
+) -> Settings:
+
+    with open(config_path, "r") as file:
         raw_config = yaml.safe_load(file)
-        
+
     return Settings(
-        training=TrainingConfig(**raw_config.get("training", {})),
-        data=DataConfig(**raw_config.get("data", {}))
+        training=TrainingConfig(
+            **raw_config["training"]
+        ),
+        data=DataConfig(
+            **raw_config["data"]
+        )
     )

src/data/data_manager.py

@@ -0,0 +1,174 @@
+import os
+import glob
+import shutil
+import logging
+import torch
+import torch.nn.functional as F
+
+from src.config.settings import Settings
+from src.data.fetchers.goes_fetcher import GOESFetcher
+from src.data.fetchers.himawari_fetcher import HimawariFetcher
+from src.data.standardizer import UniversalStandardizer
+
+logger = logging.getLogger(__name__)
+
+
+class DataManager:
+    """
+    Universal multi-satellite data pipeline manager.
+    """
+
+    def __init__(self, settings: Settings):
+        self.settings = settings
+        self.pt_dir = settings.data.download_dir
+        self.raw_dir = os.path.join(self.pt_dir, "raw_data")
+
+        os.makedirs(self.pt_dir, exist_ok=True)
+        os.makedirs(self.raw_dir, exist_ok=True)
+
+        sat_type = getattr(settings.data, "satellite_type", "goes").lower()
+
+        if sat_type == "goes":
+            self.fetcher = GOESFetcher(
+                bucket_name=settings.data.s3_bucket
+            )
+        elif sat_type == "himawari":
+            self.fetcher = HimawariFetcher(
+                bucket_name=settings.data.s3_bucket
+            )
+        else:
+            raise ValueError(f"Unsupported satellite type: {sat_type}")
+
+    def process_chunk(self, chunk_prefix: str) -> None:
+        logger.info(
+            f"Processing chunk {chunk_prefix} "
+            f"using {self.fetcher.__class__.__name__}"
+        )
+
+        raw_files = self.fetcher.fetch_chunk(chunk_prefix, self.raw_dir)
+
+        if len(raw_files) < 3:
+            logger.warning("Not enough frames for triplets.")
+            return
+
+        frame_step = self.settings.data.frame_step
+
+        for i in range(len(raw_files) - 2 * frame_step):
+            try:
+                t0_path = raw_files[i]
+                t1_path = raw_files[i + frame_step]
+                t2_path = raw_files[i + 2 * frame_step]
+
+                img0_raw = self.fetcher.apply_planck_function(t0_path)
+                gt_raw = self.fetcher.apply_planck_function(t1_path)
+                img1_raw = self.fetcher.apply_planck_function(t2_path)
+                
+                
+
+                img0 = UniversalStandardizer.normalize_bt(
+                    img0_raw,
+                    self.settings.data.min_bt,
+                    self.settings.data.max_bt
+                )
+                gt = UniversalStandardizer.normalize_bt(
+                    gt_raw,
+                    self.settings.data.min_bt,
+                    self.settings.data.max_bt
+                )
+                img1 = UniversalStandardizer.normalize_bt(
+                    img1_raw,
+                    self.settings.data.min_bt,
+                    self.settings.data.max_bt
+                )
+
+                img0_crop, img1_crop, gt_crop = (
+                    self._motion_guided_argmax_crop(
+                        img0, img1, gt
+                    )
+                )
+
+                safe_prefix = chunk_prefix.replace("/", "_")
+                pt_filename = os.path.join(
+                    self.pt_dir,
+                    f"triplet_{safe_prefix}_{i:03d}.pt"
+                )
+
+                triplet_tensor = torch.stack(
+                    [img0_crop, gt_crop, img1_crop],
+                    dim=0
+                )
+
+                torch.save(triplet_tensor, pt_filename)
+
+            except Exception as e:
+                logger.error(
+                    f"Triplet processing failed ({i}): {e}"
+                )
+                continue
+
+        self.purge_raw_files()
+
+    def _motion_guided_argmax_crop(
+        self,
+        img0: torch.Tensor,
+        img1: torch.Tensor,
+        gt: torch.Tensor
+    ):
+        crop_size = self.settings.data.crop_size
+        stride = crop_size // self.settings.data.crop_stride_divisor
+
+        _, h, w = img0.shape
+
+        if h < crop_size or w < crop_size:
+            raise ValueError(
+                f"Image smaller than crop size: {h}x{w}"
+            )
+
+        motion_map = torch.abs(img1 - img0)
+
+        # Mask out outer-space / invalid regions
+        space_mask = (img0 > 0.0).float()
+        motion_map = motion_map * space_mask
+
+        pooled_motion = F.avg_pool2d(
+            motion_map.unsqueeze(0),
+            kernel_size=crop_size,
+            stride=stride
+        )
+
+        _, _, h_out, w_out = pooled_motion.shape
+
+        flat_idx = torch.argmax(pooled_motion).item()
+
+        y_out = flat_idx // w_out
+        x_out = flat_idx % w_out
+
+        y = y_out * stride
+        x = x_out * stride
+
+        y = max(0, min(y, h - crop_size))
+        x = max(0, min(x, w - crop_size))
+
+        img0_crop = img0[:, y:y+crop_size, x:x+crop_size]
+        img1_crop = img1[:, y:y+crop_size, x:x+crop_size]
+        gt_crop = gt[:, y:y+crop_size, x:x+crop_size]
+
+        crop_motion = torch.abs(
+            img1_crop - img0_crop
+        ).mean().item()
+
+        if crop_motion < self.settings.data.static_motion_threshold:
+            raise ValueError(
+                f"Static crop rejected: {crop_motion:.5f}"
+            )
+
+        return img0_crop, img1_crop, gt_crop
+
+    def purge_raw_files(self):
+        logger.info("Purging raw files...")
+
+        for f in glob.glob(os.path.join(self.raw_dir, "*")):
+            if os.path.isfile(f):
+                os.remove(f)
+            elif os.path.isdir(f):
+                shutil.rmtree(f)

src/data/dataset.py

@@ -6,7 +6,7 @@ from torch.utils.data import Dataset
 from src.data.transforms import augment_triplet
 
 
-class GOESTripletDataset(Dataset):
+class SatelliteTripletDataset(Dataset):
     """PyTorch Dataset for loading pre-processed Satellite TIR triplets.
 
     Expects data to be stored as `.pt` files containing tensors of shape [3, 1, H, W],

src/data/fetchers/__init__.py

@@ -0,0 +1 @@
+from .base_fetcher import SatelliteFetcher

src/data/fetchers/base_fetcher.py

@@ -0,0 +1,48 @@
+import logging
+from abc import ABC, abstractmethod
+import torch
+
+logger = logging.getLogger(__name__)
+
+class SatelliteFetcher(ABC):
+    """
+    Abstract base class for all satellite data fetchers (GOES, Himawari, INSAT, etc.).
+    Enforces a strict interface for fetching and standardizing physical meteorological data.
+    """
+
+    def __init__(self, bucket_name: str):
+        """
+        Initializes the fetcher with the specific AWS or local bucket/source.
+        
+        Args:
+            bucket_name (str): The name of the data repository/bucket.
+        """
+        self.bucket_name = bucket_name
+
+    @abstractmethod
+    def fetch_chunk(self, chunk_prefix: str, output_dir: str) -> list[str]:
+        """
+        Downloads a chunk of raw satellite files from the source to a local directory.
+
+        Args:
+            chunk_prefix (str): The path/prefix for the specific time chunk.
+            output_dir (str): Local temporary directory to save the raw files.
+
+        Returns:
+            list[str]: A list of local file paths that were downloaded.
+        """
+        pass
+
+    @abstractmethod
+    def apply_planck_function(self, raw_data_path: str) -> torch.Tensor:
+        """
+        Reads the raw satellite file, applies the satellite-specific inverse 
+        Planck function, and returns a physical Brightness Temperature tensor in Kelvin.
+
+        Args:
+            raw_data_path (str): The local path to the downloaded raw file (.nc, .h5, etc.).
+
+        Returns:
+            torch.Tensor: The calculated Brightness Temperature in Kelvin.
+        """
+        pass

src/data/fetchers/goes_fetcher.py

@@ -0,0 +1,78 @@
+import os
+import logging
+import boto3
+from botocore import UNSIGNED
+from botocore.config import Config
+import xarray as xr
+import numpy as np
+import torch
+
+from src.data.fetchers.base_fetcher import SatelliteFetcher
+
+logger = logging.getLogger(__name__)
+
+
+class GOESFetcher(SatelliteFetcher):
+    def __init__(self, bucket_name: str = "noaa-goes16"):
+        super().__init__(bucket_name)
+        self.s3_client = boto3.client(
+            "s3",
+            config=Config(signature_version=UNSIGNED)
+        )
+
+    def fetch_chunk(self, chunk_prefix: str, output_dir: str) -> list[str]:
+        os.makedirs(output_dir, exist_ok=True)
+
+        paginator = self.s3_client.get_paginator("list_objects_v2")
+
+        c13_files = []
+
+        for page in paginator.paginate(
+            Bucket=self.bucket_name,
+            Prefix=chunk_prefix
+        ):
+            for obj in page.get("Contents", []):
+                key = obj["Key"]
+                if "M6C13" in key and key.endswith(".nc"):
+                    c13_files.append(key)
+
+        c13_files = sorted(c13_files)
+
+        downloaded_paths = []
+
+        for file_key in c13_files:
+            filename = os.path.basename(file_key)
+            local_path = os.path.join(output_dir, filename)
+
+            if not os.path.exists(local_path):
+                self.s3_client.download_file(
+                    self.bucket_name,
+                    file_key,
+                    local_path
+                )
+
+            downloaded_paths.append(local_path)
+
+        return downloaded_paths
+
+    def apply_planck_function(self, raw_data_path: str) -> torch.Tensor:
+        with xr.open_dataset(raw_data_path, mask_and_scale=True) as ds:
+            rad = ds["Rad"]
+
+            fk1 = ds["planck_fk1"].values
+            fk2 = ds["planck_fk2"].values
+            bc1 = ds["planck_bc1"].values
+            bc2 = ds["planck_bc2"].values
+
+            rad_safe = rad.where(rad > 0)
+
+            bt = (fk2 / np.log((fk1 / rad_safe) + 1) - bc1) / bc2
+
+            clean = np.nan_to_num(
+                bt.values,
+                nan=0.0,
+                posinf=330.0,
+                neginf=180.0
+            )
+
+            return torch.from_numpy(clean).float().unsqueeze(0)

src/data/fetchers/himawari_fetcher.py

@@ -0,0 +1,165 @@
+import os
+import glob
+import logging
+import boto3
+from botocore import UNSIGNED
+from botocore.config import Config
+import numpy as np
+import torch
+from satpy import Scene
+
+from src.data.fetchers.base_fetcher import SatelliteFetcher
+
+logger = logging.getLogger(__name__)
+
+
+class HimawariFetcher(SatelliteFetcher):
+    """
+    Himawari-8/9 AHI Full Disk Fetcher.
+    Band 14 (11.2 μm IR)
+    """
+
+    def __init__(self, bucket_name: str = "noaa-himawari9"):
+        super().__init__(bucket_name)
+
+        self.s3_client = boto3.client(
+            "s3",
+            config=Config(signature_version=UNSIGNED)
+        )
+
+    def fetch_chunk(
+        self,
+        chunk_prefix: str,
+        output_dir: str
+    ) -> list[str]:
+
+        os.makedirs(output_dir, exist_ok=True)
+
+        paginator = self.s3_client.get_paginator(
+            "list_objects_v2"
+        )
+        all_b14_files = []
+
+        for page in paginator.paginate(
+            Bucket=self.bucket_name,
+            Prefix=chunk_prefix
+        ):
+            for obj in page.get("Contents", []):
+                key = obj["Key"]
+
+                if (
+                    "B14" in key and
+                    key.endswith(".DAT.bz2")
+                ):
+                    all_b14_files.append(key)
+
+        if not all_b14_files:
+            raise ValueError(
+                f"No B14 files found for {chunk_prefix}"
+            )
+
+        # Group files by timestamp
+        timestamp_groups = {}
+
+        for key in all_b14_files:
+            filename = os.path.basename(key)
+
+            # Example:
+            # HS_H09_20260618_1200_B14_FLDK_R20_S0110.DAT.bz2
+            parts = filename.split("_")
+
+            timestamp = f"{parts[2]}_{parts[3]}"
+
+            if timestamp not in timestamp_groups:
+                timestamp_groups[timestamp] = []
+
+            timestamp_groups[timestamp].append(key)
+
+        frame_dirs = []
+
+        for timestamp in sorted(timestamp_groups.keys()):
+            segment_files = sorted(
+                timestamp_groups[timestamp]
+            )
+
+            if len(segment_files) != 10:
+                logger.warning(
+                    f"Skipping incomplete frame {timestamp} "
+                    f"({len(segment_files)} segments)"
+                )
+                continue
+
+            timestamp_dir = os.path.join(
+                output_dir,
+                timestamp
+            )
+
+            os.makedirs(timestamp_dir, exist_ok=True)
+
+            for file_key in segment_files:
+                filename = os.path.basename(file_key)
+                local_path = os.path.join(
+                    timestamp_dir,
+                    filename
+                )
+
+                if not os.path.exists(local_path):
+                    self.s3_client.download_file(
+                        self.bucket_name,
+                        file_key,
+                        local_path
+                    )
+
+            frame_dirs.append(timestamp_dir)
+
+        return frame_dirs
+
+
+    def apply_planck_function(
+        self,
+        raw_data_path: str
+    ) -> torch.Tensor:
+
+        logger.debug(
+            f"[Himawari] Loading frame from {raw_data_path}"
+        )
+
+        file_paths = sorted(
+            glob.glob(
+                os.path.join(
+                    raw_data_path,
+                    "*_B14_FLDK_R20_S*.DAT.bz2"
+                )
+            )
+        )
+
+        if len(file_paths) != 10:
+            raise ValueError(
+                f"Incomplete Himawari frame: "
+                f"{len(file_paths)} segments found"
+            )
+
+        scn = Scene(
+            reader="ahi_hsd",
+            filenames=file_paths
+        )
+
+        scn.load(
+            ["B14"],
+            calibration="brightness_temperature"
+        )
+
+        bt_data = scn["B14"].values
+
+        clean_numpy_array = np.nan_to_num(
+            bt_data,
+            nan=0.0,
+            posinf=330.0,
+            neginf=180.0
+        )
+
+        tensor = torch.from_numpy(
+            clean_numpy_array
+        ).float().unsqueeze(0)
+
+        return tensor

src/data/s3_manager.py

@@ -1,150 +0,0 @@
-import glob
-import logging
-import os
-import concurrent.futures
-
-import boto3
-import numpy as np
-import torch
-import torch.nn.functional as F
-import xarray as xr
-from botocore import UNSIGNED
-from botocore.config import Config
-
-from src.config.settings import Settings
-
-logger = logging.getLogger(__name__)
-
-class S3Manager:
-    """Manages downloading GOES data, Argmax Motion-Guided Cropping, and purging."""
-
-    def __init__(self, settings: Settings):
-        self.settings = settings
-        self.s3_client = boto3.client("s3", config=Config(signature_version=UNSIGNED))
-        self.bucket_name = settings.data.s3_bucket
-        
-        self.pt_dir = settings.data.download_dir
-        self.raw_dir = os.path.join(self.pt_dir, "raw_nc")
-        
-        os.makedirs(self.pt_dir, exist_ok=True)
-        os.makedirs(self.raw_dir, exist_ok=True)
-
-    def download_chunk(self, prefix: str) -> None:
-        logger.info(f"Fetching chunk from S3: {prefix}")
-        
-        response = self.s3_client.list_objects_v2(Bucket=self.bucket_name, Prefix=prefix)
-        c13_files = [obj['Key'] for obj in response.get('Contents', []) if 'M6C13' in obj['Key'] and obj['Key'].endswith('.nc')]
-        c13_files = sorted(c13_files)
-        
-        if len(c13_files) < 3:
-            logger.warning(f"Not enough files in prefix {prefix} to form a triplet.")
-            return
-
-        step = self.settings.data.frame_step
-        crop_size = self.settings.data.crop_size
-
-        def fetch_file(key):
-            filename = key.split('/')[-1]
-            local_path = os.path.join(self.raw_dir, filename)
-            if not os.path.exists(local_path):
-                logger.info(f"Downloading {filename}...")
-                thread_s3 = boto3.client("s3", config=Config(signature_version=UNSIGNED))
-                thread_s3.download_file(self.bucket_name, key, local_path)
-            return local_path
-
-        # Concurrent downloading (Fixes network bottleneck)
-        with concurrent.futures.ThreadPoolExecutor(max_workers=8) as executor:
-            local_nc_paths = list(executor.map(fetch_file, c13_files))
-
-        # 3. Create Triplets using Argmax Motion Map (Solves Problem 3 & 4)
-        for i in range(len(local_nc_paths) - 2 * step):
-            try:
-                # SOLVING PROBLEM 3: Load full array into RAM exactly ONCE per file.
-                t0_full = self._load_full_tensor(local_nc_paths[i])
-                t2_full = self._load_full_tensor(local_nc_paths[i + 2 * step])
-                
-                # SOLVING PROBLEM 4: Create full motion map abs(t2 - t0)
-                motion_map = torch.abs(t2_full - t0_full) # Shape: [1, H, W]
-                
-                # Ignore empty space (Outer space padding = 0.0)
-                space_mask = (t0_full > 0.0).float()
-                motion_map = motion_map * space_mask
-                
-                # Add batch dimension for pooling: [1, 1, H, W]
-                motion_map_4d = motion_map.unsqueeze(0) 
-                
-                # Use PyTorch AvgPool2d to find the 256x256 window with maximum average motion
-                stride = 64
-                pooled_motion = F.avg_pool2d(motion_map_4d, kernel_size=crop_size, stride=stride)
-                
-                # Find argmax (highest motion density) - BUG FIXED WITH .item()
-                b, c, h_out, w_out = pooled_motion.shape
-                flat_idx = torch.argmax(pooled_motion).item() # <--- Added .item() here
-                y_out = flat_idx // w_out
-                x_out = flat_idx % w_out
-                
-                # Map pooled coordinates back to actual image coordinates
-                best_start_y = y_out * stride
-                best_start_x = x_out * stride
-                
-                # SOLVING PROBLEM 1: Log the exact motion score to study distribution instead of magic numbers
-                best_motion = pooled_motion.flatten()[flat_idx].item()
-                logger.info(f"Motion-Guided Argmax Crop mapped at Y:{best_start_y}, X:{best_start_x} | Motion Score: {best_motion:.5f}")
-                
-                # Slice the exact best crop directly from RAM
-                t0_crop = t0_full[:, best_start_y:best_start_y+crop_size, best_start_x:best_start_x+crop_size]
-                t2_crop = t2_full[:, best_start_y:best_start_y+crop_size, best_start_x:best_start_x+crop_size]
-                
-                # SAFETY FILTER: Agar argmax ne bhi low-motion kachra uthaya hai, toh reject karo
-                crop_motion = torch.abs(t2_crop - t0_crop).mean().item()
-                if crop_motion < 0.005:
-                    logger.warning(f"⏩ Argmax crop too static (Score: {crop_motion:.5f}). Skipping.")
-                    continue
-                    
-              
-                # Load T1 (Present) and slice using the exact same coordinates
-                t1_full = self._load_full_tensor(local_nc_paths[i + step])
-                t1_crop = t1_full[:, best_start_y:best_start_y+crop_size, best_start_x:best_start_x+crop_size]                
-                # Save Triplet
-                triplet_tensor = torch.stack([t0_crop, t1_crop, t2_crop], dim=0)
-                safe_prefix = prefix.replace("/", "_")
-                pt_filename = os.path.join(self.pt_dir, f'triplet_{safe_prefix}_{i:03d}.pt')
-                
-                torch.save(triplet_tensor, pt_filename)
-                
-            except Exception as e:
-                logger.error(f"Error processing argmax triplet {i}: {e}")
-                continue
-
-        self.purge_raw_files()
-        logger.info(f"Chunk converted to .pt triplets. Raw .nc files purged.")
-
-    def _load_full_tensor(self, nc_path: str) -> torch.Tensor:
-        """Loads the ENTIRE NetCDF image into memory ONCE to prevent 30x disk reads."""
-        with xr.open_dataset(nc_path) as ds:
-            rad = ds["Rad"]
-            fk1 = ds["planck_fk1"].values
-            fk2 = ds["planck_fk2"].values
-            bc1 = ds["planck_bc1"].values
-            bc2 = ds["planck_bc2"].values
-
-            rad_safe = rad.where(rad > 0)
-            bt = (fk2 / np.log((fk1 / rad_safe) + 1) - bc1) / bc2
-
-            min_bt, max_bt = 180.0, 330.0
-            bt_norm = (bt - min_bt) / (max_bt - min_bt)
-            bt_norm = bt_norm.clip(0, 1)
-
-            clean_numpy_array = np.nan_to_num(bt_norm.values, nan=0.0, posinf=1.0, neginf=0.0)
-            tensor = torch.from_numpy(clean_numpy_array).float().unsqueeze(0)
-            
-        return tensor
-
-    def purge_raw_files(self) -> None:
-        for f in glob.glob(os.path.join(self.raw_dir, "*.nc")):
-            os.remove(f)
-            
-    def purge_chunk(self) -> None:
-        for f in glob.glob(os.path.join(self.pt_dir, "*.pt")):
-            os.remove(f)
-        logger.info("Purged trained .pt chunk from disk.")

src/data/standardizer.py

@@ -0,0 +1,37 @@
+import logging
+import torch
+
+logger = logging.getLogger(__name__)
+
+class UniversalStandardizer:
+    """
+    Standardizes physical radiometric satellite data into unified AI-ready tensors.
+    Ensures that irrespective of the satellite source, the model receives uniformly 
+    scaled inputs.
+    """
+
+    @staticmethod
+    def normalize_bt(bt_tensor: torch.Tensor, min_bt: float = 180.0, max_bt: float = 330.0) -> torch.Tensor:
+        """
+        Clips Brightness Temperature (K) and normalizes it to a [0, 1] range.
+
+        Args:
+            bt_tensor (torch.Tensor): The physical brightness temperature tensor in Kelvin.
+            min_bt (float, optional): Lower bound for clipping (typically overshooting tops). Defaults to 180.0.
+            max_bt (float, optional): Upper bound for clipping (typically hot desert). Defaults to 330.0.
+
+        Returns:
+            torch.Tensor: Normalized tensor bounded strictly between [0, 1].
+        """
+        logger.debug(f"Normalizing tensor of shape {bt_tensor.shape} with bounds [{min_bt}K, {max_bt}K]")
+        
+        # Apply min-max normalization
+        bt_norm = (bt_tensor - min_bt) / (max_bt - min_bt)
+        
+        # Clip values strictly between 0 and 1
+        bt_norm = torch.clamp(bt_norm, 0.0, 1.0)
+        
+        # Clean any remaining NaNs or Infs that might have survived the math
+        bt_norm = torch.nan_to_num(bt_norm, nan=0.0, posinf=1.0, neginf=0.0)
+        
+        return bt_norm

src/data/transforms.py

@@ -19,6 +19,10 @@ def augment_triplet(img0: torch.Tensor, img1: torch.Tensor, gt: torch.Tensor, cr
         Tuple of augmented tensors (img0, img1, gt).
     """
     _, h, w = img0.shape
+    if h < crop_size or w < crop_size:
+        raise ValueError(
+            f"Input smaller than crop size: {h}x{w}"
+        )
     
     # Random Spatial Cropping
     top = random.randint(0, h - crop_size)

src/training/trainer.py

@@ -9,7 +9,7 @@ from torch.utils.tensorboard import SummaryWriter
 from src.config.settings import Settings
 from src.model.ifnet import IFNet
 from src.model.loss import CompositeLoss
-from src.data.dataset import GOESTripletDataset
+from src.data.dataset import SatelliteTripletDataset
 
 logger = logging.getLogger(__name__)
 
@@ -58,9 +58,11 @@ class Trainer:
         self.settings = settings
         self.device = device
         self.model = model.to(device)
-        self.optimizer = AdamW(self.model.parameters(), lr=settings.training.learning_rate, weight_decay=1e-3)
+        self.optimizer = AdamW(self.model.parameters(), lr=settings.training.learning_rate, weight_decay=settings.training.weight_decay)
         self.criterion = CompositeLoss()
-        
+        self.scaler = torch.amp.GradScaler(
+            enabled=self.device.type == "cuda"
+        )
         self.writer = SummaryWriter(log_dir=os.path.join(settings.training.checkpoints_dir, 'logs'))
         os.makedirs(settings.training.checkpoints_dir, exist_ok=True)
         self.global_step = 0
@@ -70,12 +72,19 @@ class Trainer:
         
     def train_chunk(self, data_dir: str, epoch: int) -> None:
         """Trains the model for one epoch on the current data chunk."""
-        dataset = GOESTripletDataset(data_dir=data_dir, augment=True)
+        dataset = SatelliteTripletDataset(data_dir=data_dir, augment=True)
         if len(dataset) == 0:
             logger.warning(f"No data found in {data_dir}. Skipping training chunk.")
             return
             
-        dataloader = DataLoader(dataset, batch_size=self.settings.training.batch_size, shuffle=True, num_workers=2)
+        dataloader = DataLoader(
+            dataset,
+            batch_size=self.settings.training.batch_size,
+            shuffle=True,
+            num_workers=self.settings.training.num_workers,
+            pin_memory=True,
+            persistent_workers=True
+        )
         
         self.model.train()
         for batch_idx, (img0, img1, gt) in enumerate(dataloader):
@@ -102,8 +111,9 @@ class Trainer:
             
             loss = loss_student + loss_teacher + (loss_distill * 0.01)
             
-            loss.backward()
-            self.optimizer.step()
+            self.scaler.scale(loss).backward()
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
             if batch_idx % 10 == 0:
                 logger.info(f"Epoch {epoch} | Batch {batch_idx}/{len(dataloader)} | Loss: {loss.item():.4f}")
                 self.writer.add_scalar('Loss/train', loss.item(), self.global_step)

tests/test_config.py

@@ -1,36 +1,39 @@
-import logging
-import os
-
 from src.config.settings import load_settings
-from src.utils import setup_logging
+
 
 def test_load_settings(tmp_path):
     yaml_content = """
-    training:
-      epochs: 100
-      batch_size: 4
-      learning_rate: 0.0001
-      checkpoints_dir: "/kaggle/working/checkpoints"
-      load_model_path: "/kaggle/input/models/1week_model.pth"
-      
-    data:
-      s3_bucket: "noaa-goes16"
-      download_dir: "goes_data"
-      year: 2024
-      start_day: 264
-      end_day: 294
-      frame_step: 3
-      crop_size: 512
-    """
+training:
+  epochs: 100
+  batch_size: 4
+  learning_rate: 0.0001
+  weight_decay: 0.001
+  num_workers: 4
+  checkpoints_dir: "/tmp/checkpoints"
+  load_model_path: ""
+
+data:
+  satellite_type: "goes"
+  s3_bucket: "noaa-goes16"
+  download_dir: "/tmp/data"
+  prefix_type: "ABI-L1b-RadF"
+  year: 2024
+  start_unit: 264
+  end_unit: 294
+  frame_step: 3
+  crop_size: 512
+  crop_stride_divisor: 4
+  static_motion_threshold: 0.005
+  min_bt: 180.0
+  max_bt: 330.0
+"""
+
     config_file = tmp_path / "config.yaml"
     config_file.write_text(yaml_content)
 
     settings = load_settings(str(config_file))
-    assert settings.training.epochs == 100
-    assert settings.data.year == 2024
-    assert settings.training.load_model_path == "/kaggle/input/models/1week_model.pth"
 
-def test_setup_logging():
-    logger = setup_logging(level=logging.DEBUG)
-    assert logger.level == logging.DEBUG
-    assert len(logger.handlers) > 0
+    assert settings.training.weight_decay == 0.001
+    assert settings.training.num_workers == 4
+    assert settings.data.satellite_type == "goes"
+    assert settings.data.crop_stride_divisor == 4

tests/test_data_manager.py

@@ -0,0 +1,62 @@
+import torch
+import pytest
+
+from src.data.data_manager import DataManager
+from src.config.settings import Settings, TrainingConfig, DataConfig
+
+
+@pytest.fixture
+def settings():
+    return Settings(
+        training=TrainingConfig(
+            epochs=1,
+            batch_size=1,
+            learning_rate=1e-4,
+            weight_decay=0.001,
+            num_workers=2,
+            checkpoints_dir="chkpt"
+        ),
+        data=DataConfig(
+            satellite_type="goes",
+            s3_bucket="test",
+            download_dir="tmp",
+            prefix_type="ABI",
+            year=2024,
+            start_unit=1,
+            end_unit=2,
+            frame_step=1,
+            crop_size=64,
+            crop_stride_divisor=4,
+            static_motion_threshold=0.005
+        )
+    )
+
+
+def test_motion_crop_detects_motion(settings):
+    manager = DataManager(settings)
+
+    img0 = torch.ones((1, 256, 256)) * 0.5
+    img1 = img0.clone()
+    gt = img0.clone()
+
+    img1[:, 100:164, 100:164] = 1.0
+
+    crop0, crop1, crop_gt = manager._motion_guided_argmax_crop(
+        img0, img1, gt
+    )
+
+    motion = torch.abs(crop1 - crop0).mean().item()
+
+    assert crop0.shape == (1, 64, 64)
+    assert motion > settings.data.static_motion_threshold
+
+
+def test_static_crop_rejection(settings):
+    manager = DataManager(settings)
+
+    img0 = torch.zeros((1, 256, 256))
+    img1 = torch.zeros((1, 256, 256))
+    gt = torch.zeros((1, 256, 256))
+
+    with pytest.raises(ValueError):
+        manager._motion_guided_argmax_crop(img0, img1, gt)

tests/test_dataset.py

@@ -1,24 +1,30 @@
 import torch
-import pytest
-from src.data.transforms import augment_triplet
-from src.data.dataset import GOESTripletDataset
-
-def test_dataset(tmp_path):
-    # Setup mock data
-    tensor_data = torch.zeros((3, 1, 512, 512))
+from src.data.dataset import SatelliteTripletDataset
+def test_dataset_augment(tmp_path):
+    tensor_data = torch.ones((3, 1, 512, 512))
     torch.save(tensor_data, tmp_path / "triplet_001.pt")
-    
-    dataset = GOESTripletDataset(str(tmp_path), augment=False)
+
+    dataset = SatelliteTripletDataset(str(tmp_path), augment=True)
+
+    img0, img1, gt = dataset[0]
+
+    assert img0.shape == (1, 256, 256)
+    assert img1.shape == (1, 256, 256)
+    assert gt.shape == (1, 256, 256)
+
+def test_dataset_loading(tmp_path):
+    triplet = torch.zeros((3, 1, 256, 256))
+    torch.save(triplet, tmp_path / "triplet_001.pt")
+
+    dataset = SatelliteTripletDataset(
+        str(tmp_path),
+        augment=False
+    )
+
     assert len(dataset) == 1
-    
+
     img0, img1, gt = dataset[0]
-    assert img0.shape == (1, 512, 512)
-    assert gt.shape == (1, 512, 512)
-
-def test_augment_triplet():
-    img0 = torch.zeros((1, 512, 512))
-    img1 = torch.zeros((1, 512, 512))
-    gt = torch.zeros((1, 512, 512))
-    
-    a_img0, a_img1, a_gt = augment_triplet(img0, img1, gt, crop_size=256)
-    assert a_img0.shape == (1, 256, 256)
+
+    assert img0.shape == (1, 256, 256)
+    assert img1.shape == (1, 256, 256)
+    assert gt.shape == (1, 256, 256)

tests/test_fetchers.py

@@ -0,0 +1,86 @@
+import numpy as np
+import xarray as xr
+import torch
+from src.data.fetchers.goes_fetcher import GOESFetcher
+import pytest
+from unittest.mock import MagicMock
+
+from src.data.fetchers.himawari_fetcher import HimawariFetcher
+
+
+
+def test_goes_planck(tmp_path):
+    ds = xr.Dataset(
+        {
+            "Rad": (
+                ("y", "x"),
+                np.ones((32, 32)) * 10
+            )
+        }
+    )
+
+    ds["planck_fk1"] = 1.0
+    ds["planck_fk2"] = 100.0
+    ds["planck_bc1"] = 0.1
+    ds["planck_bc2"] = 1.0
+
+    nc_path = tmp_path / "test.nc"
+    ds.to_netcdf(nc_path)
+
+    fetcher = GOESFetcher()
+
+    tensor = fetcher.apply_planck_function(
+        str(nc_path)
+    )
+
+    assert isinstance(tensor, torch.Tensor)
+    assert tensor.shape == (1, 32, 32)
+
+
+
+
+def test_himawari_missing_segments(tmp_path):
+    fetcher = HimawariFetcher("dummy")
+
+    fetcher.s3_client.get_paginator = MagicMock()
+    fetcher.s3_client.get_paginator.return_value.paginate.return_value = [
+        {
+            "Contents": [
+                {
+                    "Key": f"HS_H09_20260618_1200_B14_FLDK_R20_S{i:04d}.DAT.bz2"
+                }
+                for i in range(8)
+            ]
+        }
+    ]
+
+    result = fetcher.fetch_chunk("prefix", str(tmp_path))
+
+    assert result == []
+        
+
+
+
+def test_goes_corrupt_file(tmp_path):
+    corrupt_file = tmp_path / "bad.nc"
+    corrupt_file.write_text("corrupted")
+
+    fetcher = GOESFetcher("dummy")
+
+    with pytest.raises(Exception):
+        fetcher.apply_planck_function(str(corrupt_file))
+        
+
+
+
+def test_himawari_corrupt_frame(tmp_path):
+    frame_dir = tmp_path / "frame"
+    frame_dir.mkdir()
+
+    for i in range(5):
+        (frame_dir / f"seg_{i}.DAT.bz2").write_text("bad")
+
+    fetcher = HimawariFetcher("dummy")
+
+    with pytest.raises(ValueError):
+        fetcher.apply_planck_function(str(frame_dir))

tests/test_pipeline.py

@@ -0,0 +1,101 @@
+import torch
+from unittest.mock import MagicMock
+
+from src.data.data_manager import DataManager
+from src.config.settings import Settings, TrainingConfig, DataConfig
+import pytest
+
+
+
+@pytest.fixture
+def settings():
+    return Settings(
+        training=TrainingConfig(
+            epochs=1,
+            batch_size=1,
+            learning_rate=1e-4,
+            weight_decay=0.001,
+            num_workers=2,
+            checkpoints_dir="chkpt"
+        ),
+        data=DataConfig(
+            satellite_type="goes",
+            s3_bucket="dummy",
+            download_dir="tmp",
+            prefix_type="ABI",
+            year=2024,
+            start_unit=1,
+            end_unit=2,
+            frame_step=1,
+            crop_size=64,
+            crop_stride_divisor=4,
+            static_motion_threshold=0.005
+        )
+    )
+
+def test_pipeline_integration(tmp_path):
+    settings = Settings(
+        training=TrainingConfig(
+            epochs=1,
+            batch_size=1,
+            learning_rate=1e-4,
+            weight_decay=0.001,
+            num_workers=2,
+            checkpoints_dir="chkpt"
+        ),
+        data=DataConfig(
+            satellite_type="goes",
+            s3_bucket="dummy",
+            download_dir=str(tmp_path),
+            prefix_type="ABI",
+            year=2024,
+            start_unit=1,
+            end_unit=2,
+            frame_step=1,
+            crop_size=64,
+            crop_stride_divisor=4,
+            static_motion_threshold=0.005
+        )
+    )
+
+    manager = DataManager(settings)
+
+    manager.fetcher.fetch_chunk = MagicMock(
+        return_value=["a.nc", "b.nc", "c.nc"]
+    )
+
+    img0 = torch.ones((1, 256, 256)) * 250.0
+    gt = img0.clone()
+    img1 = img0.clone()
+    img1[:, 100:164, 100:164] = 300.0
+
+    fake_tensors = [img0, gt, img1]
+
+    manager.fetcher.apply_planck_function = MagicMock(
+        side_effect=fake_tensors
+    )
+
+    manager.process_chunk("dummy_prefix")
+
+    saved_files = list(tmp_path.glob("*.pt"))
+
+    assert len(saved_files) == 1
+    
+
+
+
+
+def test_factory_goes(settings):
+    settings.data.satellite_type = "goes"
+
+    manager = DataManager(settings)
+
+    assert manager.fetcher.__class__.__name__ == "GOESFetcher"
+
+
+def test_factory_himawari(settings):
+    settings.data.satellite_type = "himawari"
+
+    manager = DataManager(settings)
+
+    assert manager.fetcher.__class__.__name__ == "HimawariFetcher"

tests/test_s3_manager.py

@@ -1,46 +0,0 @@
-from unittest.mock import MagicMock, patch
-import numpy as np
-import pytest
-import xarray as xr
-import torch
-
-from src.config.settings import DataConfig, Settings, TrainingConfig
-from src.data.s3_manager import S3Manager
-
-@pytest.fixture
-def mock_settings():
-    # Updated to match the exact schema
-    return Settings(
-        training=TrainingConfig(
-            epochs=1, batch_size=1, learning_rate=0.001, 
-            checkpoints_dir="chkpt", load_model_path=""
-        ),
-        data=DataConfig(
-            s3_bucket="test-bucket", download_dir="test_dir", 
-            year=2024, start_day=200, end_day=201, frame_step=3, crop_size=256
-        ),
-    )
-
-def test_load_full_tensor(mock_settings, tmp_path):
-    manager = S3Manager(mock_settings)
-    
-    # Create mock NetCDF dataset simulating GOES-16/19 data
-    ds = xr.Dataset(
-        {"Rad": (("y", "x"), np.ones((512, 512)) * 10)},
-        coords={"y": np.arange(512), "x": np.arange(512)},
-    )
-    ds["planck_fk1"] = 1.0
-    ds["planck_fk2"] = 100.0
-    ds["planck_bc1"] = 0.1
-    ds["planck_bc2"] = 1.0
-
-    nc_path = tmp_path / "test.nc"
-    ds.to_netcdf(nc_path)
-
-    # 🚨 FIX: Call the ACTUAL method that exists in your s3_manager.py
-    tensor = manager._load_full_tensor(str(nc_path))
-
-    # Assertions to ensure it returns a valid PyTorch tensor with the correct shape
-    assert isinstance(tensor, torch.Tensor)
-    assert tensor.shape == (1, 512, 512)
-    assert not torch.isnan(tensor).any()  # Ensure no NaN values exist in the output

tests/test_standardizer.py

@@ -0,0 +1,21 @@
+import torch
+from src.data.standardizer import UniversalStandardizer
+
+
+def test_normalize_bt():
+    bt = torch.tensor([180.0, 255.0, 330.0])
+
+    out = UniversalStandardizer.normalize_bt(bt)
+
+    assert out[0] == 0.0
+    assert out[2] == 1.0
+    assert 0.0 <= out[1] <= 1.0
+
+
+def test_normalize_bt_clipping():
+    bt = torch.tensor([100.0, 400.0])
+
+    out = UniversalStandardizer.normalize_bt(bt)
+
+    assert out[0] == 0.0
+    assert out[1] == 1.0

tests/test_trainer.py

@@ -1,23 +1,77 @@
 import torch
-import pytest
 from src.training.trainer import Trainer
 from src.model.ifnet import IFNet
 from src.config.settings import Settings, TrainingConfig, DataConfig
+import os
+
 
 def test_trainer_init():
-    # 🚨 FIX: Updated DataConfig and TrainingConfig parameters
     settings = Settings(
         training=TrainingConfig(
-            epochs=1, batch_size=2, learning_rate=1e-4, 
-            checkpoints_dir="chkpt", load_model_path=""
+            epochs=1,
+            batch_size=2,
+            learning_rate=1e-4,
+            weight_decay=0.001,
+            num_workers=2,
+            checkpoints_dir="chkpt"
         ),
         data=DataConfig(
-            s3_bucket="b", download_dir="d", year=2024, 
-            start_day=200, end_day=201, frame_step=3, crop_size=256
+            satellite_type="goes",
+            s3_bucket="bucket",
+            download_dir="data",
+            prefix_type="ABI",
+            year=2024,
+            start_unit=1,
+            end_unit=2,
+            frame_step=1,
+            crop_size=256
         )
     )
-    model = IFNet()
-    device = torch.device("cpu")
-    trainer = Trainer(settings, model, device)
+
+    trainer = Trainer(
+        settings,
+        IFNet(),
+        torch.device("cpu")
+    )
+
+    assert trainer.global_step == 0
     
-    assert trainer.global_step == 0
+
+
+def test_checkpoint_save(tmp_path):
+    settings = Settings(
+        training=TrainingConfig(
+            epochs=1,
+            batch_size=2,
+            learning_rate=1e-4,
+            weight_decay=0.001,
+            num_workers=2,
+            checkpoints_dir=str(tmp_path)
+        ),
+        data=DataConfig(
+            satellite_type="goes",
+            s3_bucket="b",
+            download_dir="d",
+            prefix_type="ABI",
+            year=2024,
+            start_unit=1,
+            end_unit=2,
+            frame_step=1,
+            crop_size=256,
+            crop_stride_divisor=4,
+            static_motion_threshold=0.005
+        )
+    )
+
+    trainer = Trainer(
+        settings,
+        model=IFNet(),
+        device=torch.device("cpu")
+    )
+
+    trainer.save_checkpoint("test_model.pth")
+    trainer.shutdown()
+
+    assert os.path.exists(
+        os.path.join(tmp_path, "test_model.pth")
+    )