Upload current version

2025-11-20 12:36:30 +01:00
parent dfdd200867
commit 4b071e68b5
11 changed files with 367965 additions and 459 deletions
@@ -1,5 +1,275 @@
 import pandas as pd
-import tensorflow as tf
+import torch
+from torch.utils.data import Dataset, DataLoader
 import numpy as np
 import pathlib
+from typing import Sequence
+import random

+class BaseDataset(Dataset):
+    def __init__(
+        self,
+        filepaths,
+        n_input,
+        n_output,
+        n_windows_per_file,
+        step=1,
+        feature_columns=("lat", "lon", "alt", "ias"),
+        context_columns=("r","t"),
+        time_columns=("timestamp",),
+        target_columns=("lat", "lon", "alt"),
+        device="cpu",
+    ):
+        super().__init__()
+        self.filepaths = [pathlib.Path(fp) for fp in filepaths]
+        self.n_input = n_input
+        self.n_output = n_output
+        self.window_size = n_input + n_output
+        self.device = device
+        self.n_windows_per_file = n_windows_per_file
+        self.step = step
+
+        # Which columns to use
+        sample = pd.read_csv(self.filepaths[0], nrows=5)
+        self.feature_cols = list(feature_columns)
+        self.context_cols = list(context_columns)
+        self.time_cols = list(time_columns)
+        self.target_cols = list(target_columns)
+
+        # Total samples = windows_per_file * number_of_files
+        self.total_windows = n_windows_per_file * len(self.filepaths)
+
+        for col_list in [self.feature_cols, self.context_cols, self.time_cols, self.target_cols]:
+            assert all(
+                col in sample.columns for col in col_list
+            ), "Some specified columns not found in data."
+
+    def __len__(self):
+        return self.total_windows
+    
+    def _get_start_end_indices(self, idx, n_rows):
+        # Max valid start index so window stays inside the file
+        max_start = n_rows - self.window_size * self.step
+        if max_start < 0:
+            raise ValueError("File too small for one window.")
+
+        # Evenly spaced index:
+        # raw_start = (n_rows / n_windows_per_file) * local_idx
+        # mapped into integer space
+        raw_start = (n_rows / self.n_windows_per_file) * idx
+        start_idx = int(raw_start // self.step) * self.step
+
+        # Clamp to safe zone
+        if start_idx > max_start:
+            start_idx = max_start
+
+        end_idx = start_idx + self.window_size * self.step
+        return start_idx, end_idx
+
+    def _modify_df(self, df):
+        """Hook for subclasses to modify dataframe before slicing windows."""
+        return df
+
+
+class EvenlySpacedDataset(BaseDataset):
+    def __init__(
+        self,
+        filepaths,
+        n_input,
+        n_output,
+        n_windows_per_file,
+        step=1,
+        feature_columns=("lat", "lon", "alt", "ias"),
+        context_columns=("r","t"),
+        time_columns=("timestamp",),
+        target_columns=("lat", "lon", "alt"),
+        device="cpu",
+    ):
+        super().__init__(
+            filepaths=filepaths,
+            n_input=n_input,
+            n_output=n_output,
+            feature_columns=feature_columns,
+            context_columns=context_columns,
+            time_columns=time_columns,
+            target_columns=target_columns,
+            device=device,
+            step=step,
+            n_windows_per_file=n_windows_per_file,
+        )
+
+        # Preload all data into memory
+        self.data = []
+        for fp in self.filepaths:
+            df = pd.read_csv(fp)
+            n_rows = len(df)
+            # Extract windows
+            for w in range(n_windows_per_file):
+                try:
+                    start_idx, end_idx = self._get_start_end_indices(w, n_rows)
+                    window = df.iloc[start_idx:end_idx:self.step]
+                    self.data.append(window)
+                except ValueError:
+                    self.total_windows -= 1
+                    continue
+
+    def __len__(self):
+        return self.total_windows
+
+    def __getitem__(self, idx):
+        window = self.data[idx]
+
+        X_feat = window.iloc[: self.n_input][self.feature_cols].to_numpy(np.float32)
+        if len(self.context_cols) > 0:
+            X_context = window.iloc[: self.n_input][self.context_cols].to_numpy(np.float32)
+        else:
+            X_context = None
+        X_time = window.iloc[: self.window_size][self.time_cols].to_numpy(np.float32)
+        Y_out = window.iloc[self.n_input :][self.target_cols].to_numpy(np.float32)
+        return (
+            torch.tensor(X_feat, device=self.device),
+            torch.tensor(X_time, device=self.device),
+            torch.tensor(Y_out, device=self.device),
+            torch.tensor(X_context, device=self.device) if X_context is not None else None,
+        )
+
+class EvenlySpacedStreamingDataset(BaseDataset):
+    def __init__(
+        self,
+        filepaths,
+        n_input,
+        n_output,
+        n_windows_per_file,
+        step=1,
+        feature_columns=("lat", "lon", "alt", "ias"),
+        context_columns=("r","t"),
+        time_columns=("timestamp",),
+        target_columns=("lat", "lon", "alt"),
+        device="cpu",
+    ):
+        super().__init__(
+            filepaths=filepaths,
+            n_input=n_input,
+            n_output=n_output,
+            feature_columns=feature_columns,
+            context_columns=context_columns,
+            time_columns=time_columns,
+            target_columns=target_columns,
+            device=device,
+            step=step,
+            n_windows_per_file=n_windows_per_file,
+        )
+        
+        # Precompute row counts for each file
+        self.row_counts = [self._count_rows(fp) for fp in self.filepaths]
+
+
+    def _count_rows(self, fp):
+        with open(fp, "r") as f:
+            return sum(1 for _ in f) - 1  # minus header
+
+    def __len__(self):
+        return self.total_windows
+
+    def __getitem__(self, idx):
+        # Figure out which file we belong to
+        file_idx = idx // self.n_windows_per_file
+        local_idx = idx % self.n_windows_per_file
+        fp = self.filepaths[file_idx]
+
+        n_rows = self.row_counts[file_idx]
+
+        start_idx, end_idx = self._get_start_end_indices(local_idx, n_rows)
+
+        # Read only the required rows
+        skip = list(set(range(1, start_idx + 1)))
+        df = pd.read_csv(fp, skiprows=skip, nrows=self.window_size * self.step).iloc[::self.step]
+
+        # Slice into input / decoder-input / targets
+        X_feat = df.iloc[: self.n_input][self.feature_cols].to_numpy(np.float32)
+        if len(self.context_cols) > 0:
+            X_context = df.iloc[: self.n_input][self.context_cols].to_numpy(np.float32)
+        else:
+            X_context = None
+        X_time = df.iloc[: self.window_size][self.time_cols].to_numpy(np.float32)
+        Y_out = df.iloc[self.n_input :][self.target_cols].to_numpy(np.float32)
+        return (
+            torch.tensor(X_feat, device=self.device),
+            torch.tensor(X_time, device=self.device),
+            torch.tensor(Y_out, device=self.device),
+            torch.tensor(X_context, device=self.device) if X_context is not None else None,
+        )
+
+def get_datasets(
+    files: Sequence[str],
+    n_input: int,
+    n_output: int,
+    n_windows_per_file: int,
+    step: int = 1,
+    feature_columns=("lat", "lon", "alt", "ias"),
+    context_columns=("r","t"),
+    time_columns=("timestamp",),
+    target_columns=("lat", "lon", "alt"),
+    device="cpu",
+    test_split: float = 0.2,
+    val_split: float = 0.1,
+    seed: int = 42,
+    streaming: bool = False,
+) -> tuple[Dataset, Dataset, Dataset]:
+    """Utility to create train/val/test datasets from file list."""
+    random.seed(seed)
+    files = list(files)
+    random.shuffle(files)
+
+    n_total = len(files)
+    n_test = int(n_total * test_split)
+    n_val = int(n_total * val_split)
+
+    test_files = files[:n_test]
+    val_files = files[n_test : n_test + n_val]
+    train_files = files[n_test + n_val :]
+
+    if not streaming:
+        ds = EvenlySpacedDataset
+    else:
+        ds = EvenlySpacedStreamingDataset
+
+    train_dataset = ds(
+        train_files,
+        n_input=n_input,
+        n_output=n_output,
+        n_windows_per_file=n_windows_per_file,
+        step=step,
+        feature_columns=feature_columns,
+        context_columns=context_columns,
+        time_columns=time_columns,
+        target_columns=target_columns,
+        device=device,
+    )
+    val_dataset = ds(
+        val_files,
+        n_input=n_input,
+        n_output=n_output,
+        n_windows_per_file=n_windows_per_file,
+        step=step,
+        feature_columns=feature_columns,
+        context_columns=context_columns,
+        time_columns=time_columns,
+        target_columns=target_columns,
+        device=device,
+    )
+    test_dataset = ds(
+        test_files,
+        n_input=n_input,
+        n_output=n_output,
+        n_windows_per_file=n_windows_per_file,
+        step=step,
+        feature_columns=feature_columns,
+        context_columns=context_columns,
+        time_columns=time_columns,
+        target_columns=target_columns,
+        device=device,
+    )
+
+    return train_dataset, val_dataset, test_dataset
+