357 lines
13 KiB
Python
357 lines
13 KiB
Python
import pandas as pd
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from torch.utils.data import Dataset, DataLoader
|
|
import numpy as np
|
|
import pathlib
|
|
from typing import Sequence
|
|
import random
|
|
from aiRNN import preprocessors
|
|
|
|
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=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
|
|
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)
|
|
sample = self._modify_df(sample)
|
|
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."""
|
|
lat_norm, lon_norm, alt_norm = preprocessors.norm_coords(
|
|
df["lat"].values, df["lon"].values, df["alt"].values
|
|
)
|
|
df.loc[:, "lat"] = lat_norm
|
|
df.loc[:, "lon"] = lon_norm
|
|
df.loc[:, "alt"] = alt_norm
|
|
df.loc[:, "ias"] = preprocessors.norm_ias(df["ias"].values)
|
|
df.loc[:, "dt"] = df["timestamp"].diff()
|
|
df.loc[:, "dt"] = preprocessors.fillna_with_mean(df["dt"], allow_nan_mean=False)
|
|
df.loc[:, "timestamp"] = preprocessors.norm_time(df["timestamp"].values)
|
|
for col in ["lat", "lon", "alt", "ias"]:
|
|
if df[col].isnull().all() and col != "ias":
|
|
raise ValueError(f"All values in column {col} are NaN.")
|
|
df.loc[:, col] = preprocessors.fillna_with_mean(df[col], allow_nan_mean=(col == "ias"))
|
|
|
|
mapped_df = preprocessors.map_categories(df, preprocessors.category_mappings)
|
|
max_values = list(preprocessors.category_max_limits.values())
|
|
X_a_1 = F.one_hot(torch.tensor(mapped_df["R_1_IDX"].values), max_values[0]).float()
|
|
X_a_2 = F.one_hot(torch.tensor(mapped_df["R_2_IDX"].values), max_values[1]).float()
|
|
X_t = F.one_hot(torch.tensor(mapped_df["T_IDX"].values), max_values[2]).float()
|
|
df.loc[:, [f"type_encoding_{i}" for i in range(4)]] = preprocessors.encode_features(
|
|
X_a_1, X_a_2, X_t
|
|
)
|
|
last_row = df.iloc[-1][["lat", "lon", "alt", "ias", "timestamp"]]
|
|
|
|
df.loc[:, "last_lat"] = last_row["lat"]
|
|
df.loc[:, "last_lon"] = last_row["lon"]
|
|
df.loc[:, "last_alt"] = last_row["alt"]
|
|
df.loc[:, "last_ias"] = last_row["ias"]
|
|
df.loc[:, "last_timestamp"] = last_row["timestamp"]
|
|
|
|
return df
|
|
|
|
def save_entire_dataset(self, filepath):
|
|
"""Utility to save the entire dataset to a single pth file."""
|
|
X_feat_list = []
|
|
X_time_list = []
|
|
Y_out_list = []
|
|
X_context_list = []
|
|
for i in range(len(self)):
|
|
X_f, X_t, Y, X_c = self[i]
|
|
X_feat_list.append(X_f.cpu())
|
|
X_time_list.append(X_t.cpu())
|
|
Y_out_list.append(Y.cpu())
|
|
if X_c is not None:
|
|
X_context_list.append(X_c.cpu())
|
|
data_dict = {
|
|
"X_feat": torch.stack(X_feat_list),
|
|
"X_time": torch.stack(X_time_list),
|
|
"Y_out": torch.stack(Y_out_list),
|
|
}
|
|
if len(X_context_list) > 0:
|
|
data_dict["X_context"] = torch.stack(X_context_list)
|
|
torch.save(data_dict, filepath)
|
|
|
|
class SaveDataset(Dataset):
|
|
def __init__(self, data_dict, device="cpu"):
|
|
super().__init__()
|
|
self.X_feat = data_dict["X_feat"].to(device)
|
|
self.X_time = data_dict["X_time"].to(device)
|
|
self.Y_out = data_dict["Y_out"].to(device)
|
|
self.device = device
|
|
if "X_context" in data_dict:
|
|
self.X_context = data_dict["X_context"].to(device)
|
|
else:
|
|
self.X_context = None
|
|
|
|
def __len__(self):
|
|
return self.X_feat.shape[0]
|
|
|
|
def __getitem__(self, idx):
|
|
X_f = self.X_feat[idx]
|
|
X_t = self.X_time[idx]
|
|
Y = self.Y_out[idx]
|
|
if self.X_context is not None:
|
|
X_c = self.X_context[idx]
|
|
else:
|
|
X_c = None
|
|
return X_f, X_t, Y, X_c
|
|
|
|
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=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
|
|
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, dtype={"icao": str, "r": str, "t": str, "timestamp": float, "lat": float, "lon": float, "alt": float, "ias": float})
|
|
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]
|
|
window = self._modify_df(window)
|
|
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[0][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=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
|
|
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, dtype={"icao": str, "r": str, "t": str, "timestamp": float, "lat": float, "lon": float, "alt": float, "ias": float}).iloc[::self.step]
|
|
df = self._modify_df(df)
|
|
|
|
# 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=(*[f"type_encoding_{i}" for i in range(4)], "last_lat", "last_lon", "last_alt", "last_ias"),
|
|
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
|
|
|