Upload current version
This commit is contained in:
Binary file not shown.
@@ -0,0 +1 @@
|
||||
{"r_1": {"02": 0, "0264": 1, "0265": 2, "0268": 3, "0270": 4, "0273": 5, "0275": 6, "0276": 7, "0278": 8, "0281": 9, "0283": 10, "0284": 11, "100001": 12, "100003": 13, "100004": 14, "100008": 15, "102002": 16, "102003": 17, "102004": 18, "102005": 19, "2": 20, "40199": 21, "40200": 22, "43": 23, "4K": 24, "4X": 25, "5601": 26, "5629": 27, "5699": 28, "83": 29, "84": 30, "84006": 31, "84007": 32, "9582": 33, "9585": 34, "9H": 35, "A6": 36, "A7": 37, "AP": 38, "B": 39, "C": 40, "CC": 41, "CS": 42, "D": 43, "EC": 44, "EI": 45, "ES": 46, "ET": 47, "EW": 48, "F": 49, "G": 50, "GND": 51, "GO": 52, "HA": 53, "HB": 54, "I": 55, "KUNTO13": 56, "KUNTO3": 57, "KUNTO7": 58, "KUNTO72": 59, "LJ": 60, "LN": 61, "LX": 62, "LY": 63, "LZ": 64, "M": 65, "MR7": 66, "N441FF": 67, "N737WK": 68, "OE": 69, "OH": 70, "OK": 71, "OM": 72, "OO": 73, "OY": 74, "Other": 75, "PH": 76, "PI": 77, "RA": 78, "SAHKO4": 79, "SAHKO8": 80, "SE": 81, "SLP": 82, "SP": 83, "SU": 84, "SX": 85, "T7": 86, "TC": 87, "TF": 88, "TWR": 89, "UK": 90, "UP": 91, "UR": 92, "VIESTI52": 93, "VN": 94, "VP": 95, "VT": 96, "YL": 97, "YR": 98, "YU": 99}, "r_2": {"ATE": 0, "ATG": 1, "ATH": 2, "ATI": 3, "DYU": 4, "ENL": 5, "ENM": 6, "ENN": 7, "ENO": 8, "ENP": 9, "ENR": 10, "ENS": 11, "ENT": 12, "ENU": 13, "ENV": 14, "Empty": 15, "LKL": 16, "LKO": 17, "LVH": 18, "LVI": 19, "LVL": 20, "LXC": 21, "LZE": 22, "LZF": 23, "LZG": 24, "LZH": 25, "LZI": 26, "LZK": 27, "LZM": 28, "LZN": 29, "LZO": 30, "LZP": 31, "LZR": 32, "LZT": 33, "LZU": 34, "Other": 35, "RDY": 36, "RGL": 37, "ROK": 38, "ROL": 39, "ROY": 40, "RUA": 41, "RUC": 42, "RUD": 43, "RUF": 44, "RUO": 45, "RZX": 46, "WEC": 47, "WIW": 48, "WIX": 49}, "t": {"A139": 0, "A169": 1, "A20N": 2, "A21N": 3, "A319": 4, "A320": 5, "A321": 6, "A332": 7, "A333": 8, "A339": 9, "A359": 10, "A35K": 11, "A388": 12, "AS32": 13, "AS50": 14, "AT43": 15, "AT45": 16, "AT72": 17, "AT75": 18, "AT76": 19, "B190": 20, "B350": 21, "B38M": 22, "B429": 23, "B734": 24, "B737": 25, "B738": 26, "B739": 27, "B744": 28, "B752": 29, "B763": 30, "B77L": 31, "B77W": 32, "B788": 33, "B789": 34, "BCS3": 35, "BE20": 36, "C130": 37, "C152": 38, "C172": 39, "C208": 40, "C25A": 41, "C25B": 42, "C25C": 43, "C30J": 44, "C56X": 45, "C68A": 46, "CL35": 47, "CL60": 48, "CRJ2": 49, "CRJ9": 50, "D228": 51, "DA40": 52, "DA42": 53, "DH8A": 54, "DH8B": 55, "DH8C": 56, "DH8D": 57, "DHC6": 58, "E190": 59, "E195": 60, "E290": 61, "E295": 62, "E35L": 63, "E550": 64, "E55P": 65, "E75L": 66, "E75S": 67, "EC35": 68, "EC45": 69, "EH10": 70, "F2TH": 71, "F50": 72, "FA7X": 73, "FA8X": 74, "G115": 75, "GLEX": 76, "GLF4": 77, "GLF5": 78, "GLF6": 79, "H25B": 80, "LJ35": 81, "LJ45": 82, "NG5": 83, "Other": 84, "P28A": 85, "PA31": 86, "PC12": 87, "PC24": 88, "R44": 89, "RV7": 90, "S92": 91, "SB20": 92, "SERV": 93, "SF34": 94, "SR22": 95, "SU95": 96, "SW4": 97, "TWR": 98, "WT9": 99}}
|
||||
@@ -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
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import math
|
||||
|
||||
|
||||
class HaversineMSEAltitudeLoss(nn.Module):
|
||||
def __init__(self, alt_const=1.0, earth_radius_km=6371.0):
|
||||
super().__init__()
|
||||
self.alt_const = alt_const
|
||||
self.R = earth_radius_km
|
||||
|
||||
def forward(self, pred, target):
|
||||
"""
|
||||
pred: (..., 3) -> [lat, lon, alt]
|
||||
target: (..., 3) -> [lat, lon, alt]
|
||||
|
||||
lat/lon expected in degrees.
|
||||
alt in meters (or same unit in pred/target).
|
||||
"""
|
||||
lat1 = torch.deg2rad(pred[..., 0])
|
||||
lon1 = torch.deg2rad(pred[..., 1])
|
||||
lat2 = torch.deg2rad(target[..., 0])
|
||||
lon2 = torch.deg2rad(target[..., 1])
|
||||
|
||||
dlat = lat2 - lat1
|
||||
dlon = lon2 - lon1
|
||||
|
||||
# Haversine formula
|
||||
a = (
|
||||
torch.sin(dlat / 2) ** 2
|
||||
+ torch.cos(lat1) * torch.cos(lat2) * torch.sin(dlon / 2) ** 2
|
||||
)
|
||||
c = 2 * torch.atan2(torch.sqrt(a), torch.sqrt(1 - a))
|
||||
dist_km = self.R * c # great-circle distance
|
||||
|
||||
# MSE on distance
|
||||
haversine_mse = torch.mean(dist_km ** 2)
|
||||
|
||||
# Altitude penalty
|
||||
delta_alt = pred[..., 2] - target[..., 2]
|
||||
alt_penalty = self.alt_const * torch.mean(delta_alt ** 2)
|
||||
|
||||
return haversine_mse + alt_penalty
|
||||
|
||||
class WeightedMSELoss(nn.Module):
|
||||
def __init__(self, weights):
|
||||
super().__init__()
|
||||
self.weights = torch.tensor(weights).float()
|
||||
|
||||
def forward(self, pred, target):
|
||||
"""
|
||||
pred: (..., D)
|
||||
target: (..., D)
|
||||
"""
|
||||
diff = pred - target
|
||||
weighted_diff = diff * self.weights.to(diff.device)
|
||||
return torch.mean(weighted_diff ** 2)
|
||||
@@ -0,0 +1,146 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class BaseRNN(nn.Module):
|
||||
def __init__(self, time_in, feat_in, context_in, hidden_size, out_size, rnn_type="RNN"):
|
||||
super().__init__()
|
||||
self.time_in = time_in
|
||||
self.feat_in = feat_in
|
||||
self.context_in = context_in
|
||||
self.hidden_size = hidden_size
|
||||
self.out_size = out_size
|
||||
|
||||
# Projecters to a common RNN input size
|
||||
self.time_proj = nn.Linear(time_in, hidden_size)
|
||||
self.feat_proj = nn.Linear(feat_in, hidden_size)
|
||||
if context_in is not None:
|
||||
self.context_proj = nn.Linear(context_in, hidden_size)
|
||||
|
||||
# RNN cell
|
||||
if rnn_type == "RNN":
|
||||
self.rnn = nn.RNNCell(input_size=hidden_size, hidden_size=hidden_size)
|
||||
elif rnn_type == "LSTM":
|
||||
self.rnn = nn.LSTMCell(input_size=hidden_size, hidden_size=hidden_size)
|
||||
else:
|
||||
raise ValueError(f"Unsupported rnn_type: {rnn_type}")
|
||||
|
||||
# Readout: predict from hidden state (optionally conditioned on time and context too)
|
||||
readout_in_size = hidden_size + time_in
|
||||
if context_in is not None:
|
||||
readout_in_size += context_in
|
||||
self.readout = nn.Linear(readout_in_size, out_size)
|
||||
|
||||
def forward(self, time_seq, feat_init=None, context=None, init_steps=0, pred_steps=0, offset=0, hidden=None):
|
||||
b, total_len, _ = time_seq.shape
|
||||
assert init_steps <= total_len
|
||||
|
||||
if feat_init is not None:
|
||||
assert feat_init.shape[0] == b and feat_init.shape[1] == init_steps
|
||||
|
||||
if hidden is None:
|
||||
hidden = torch.zeros(b, self.hidden_size, device=time_seq.device)
|
||||
|
||||
if self.context_in is not None:
|
||||
assert context is not None and context.shape[0] == b
|
||||
context_proj = self.context_proj(context) # (b, hidden_size)
|
||||
else:
|
||||
context_proj = torch.zeros(b, self.hidden_size, device=time_seq.device)
|
||||
|
||||
# --- Initialization phase: feed time + feature in parallel for init_steps ---
|
||||
for t in range(init_steps + offset):
|
||||
t_in = time_seq[:, t, :] # (b, time_in)
|
||||
t_proj = self.time_proj(t_in) # (b, hidden_size)
|
||||
|
||||
rnn_in = t_proj + context_proj # start with time + context
|
||||
if feat_init is not None and t < init_steps:
|
||||
f_in = feat_init[:, t, :] # (b, feat_in)
|
||||
f_proj = self.feat_proj(f_in) # (b, hidden_size)
|
||||
rnn_in = rnn_in + f_proj # combine with feature input
|
||||
|
||||
if isinstance(self.rnn, nn.LSTMCell):
|
||||
if isinstance(hidden, tuple):
|
||||
h_t, c_t = hidden
|
||||
else:
|
||||
h_t = hidden
|
||||
c_t = torch.zeros(b, self.hidden_size, device=time_seq.device)
|
||||
h_t, c_t = self.rnn(rnn_in, (h_t, c_t))
|
||||
hidden = (h_t, c_t)
|
||||
else:
|
||||
hidden = self.rnn(rnn_in, hidden)
|
||||
|
||||
# --- Prediction phase: feed only time (+ context) inputs for pred_steps ---
|
||||
preds = []
|
||||
for p in range(pred_steps):
|
||||
t_idx = init_steps + p
|
||||
assert t_idx < total_len, "time_seq too short for requested pred_steps"
|
||||
t_in = time_seq[:, t_idx, :] # (b, time_in)
|
||||
t_proj = self.time_proj(t_in) # project time input
|
||||
|
||||
rnn_in = t_proj + context_proj # RNN input: time + context
|
||||
|
||||
if isinstance(self.rnn, nn.LSTMCell):
|
||||
h_t, c_t = hidden
|
||||
h_t, c_t = self.rnn(rnn_in, (h_t, c_t)) # update hidden
|
||||
hidden = (h_t, c_t)
|
||||
else:
|
||||
hidden = self.rnn(rnn_in, hidden) # update hidden
|
||||
|
||||
# Readout uses hidden + raw time (+ raw context input)
|
||||
read_inputs = [hidden if not isinstance(hidden, tuple) else hidden[0], t_in]
|
||||
if self.context_in is not None:
|
||||
read_inputs.append(context)
|
||||
read = torch.cat(read_inputs, dim=-1)
|
||||
out = self.readout(read) # (b, out_size)
|
||||
preds.append(out.unsqueeze(1))
|
||||
preds = torch.cat(preds, dim=1) # (b, pred_steps, out_size)
|
||||
return preds, hidden
|
||||
|
||||
|
||||
|
||||
|
||||
class TwoInputRNN(BaseRNN):
|
||||
"""
|
||||
Many-to-many RNN with an initialization phase that consumes both:
|
||||
- time inputs (always present)
|
||||
- feature inputs (only during init_steps)
|
||||
After init_steps, only time inputs are provided and the model predicts a sequence.
|
||||
Predictions can be compared to targets with a specified `offset`.
|
||||
"""
|
||||
def __init__(self, time_in, feat_in, hidden_size, out_size):
|
||||
super().__init__(time_in, feat_in, context_in=None, hidden_size=hidden_size, out_size=out_size)
|
||||
|
||||
class ThreeInputRNN(BaseRNN):
|
||||
"""
|
||||
Many-to-many RNN with an initialization phase that consumes three inputs:
|
||||
- time inputs (always present)
|
||||
- feature inputs (only during init_steps)
|
||||
- context inputs (always present + constant over time)
|
||||
After init_steps, only time and context inputs are provided and the model predicts a sequence.
|
||||
Predictions can be compared to targets with a specified `offset`.
|
||||
"""
|
||||
def __init__(self, time_in, feat_in, context_in, hidden_size, out_size):
|
||||
super().__init__(time_in, feat_in, context_in=context_in, hidden_size=hidden_size, out_size=out_size)
|
||||
|
||||
class TwoInputLSTM(BaseRNN):
|
||||
"""
|
||||
Many-to-many LSTM with an initialization phase that consumes both:
|
||||
- time inputs (always present)
|
||||
- feature inputs (only during init_steps)
|
||||
After init_steps, only time inputs are provided and the model predicts a sequence.
|
||||
Predictions can be compared to targets with a specified `offset`.
|
||||
"""
|
||||
def __init__(self, time_in, feat_in, hidden_size, out_size):
|
||||
super().__init__(time_in, feat_in, context_in=None, hidden_size=hidden_size, out_size=out_size, rnn_type="LSTM")
|
||||
|
||||
class ThreeInputLSTM(BaseRNN):
|
||||
"""
|
||||
Many-to-many LSTM with an initialization phase that consumes three inputs:
|
||||
- time inputs (always present)
|
||||
- feature inputs (only during init_steps)
|
||||
- context inputs (always present + constant over time)
|
||||
After init_steps, only time and context inputs are provided and the model predicts a sequence.
|
||||
Predictions can be compared to targets with a specified `offset`.
|
||||
"""
|
||||
def __init__(self, time_in, feat_in, context_in, hidden_size, out_size):
|
||||
super().__init__(time_in, feat_in, context_in=context_in, hidden_size=hidden_size, out_size=out_size, rnn_type="LSTM")
|
||||
@@ -0,0 +1,150 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
import pathlib
|
||||
import json
|
||||
|
||||
|
||||
MIN_LAT, MAX_LAT = 57.0, 72.0
|
||||
MIN_LON, MAX_LON = 3.0, 32.0
|
||||
MIN_ALT, MAX_ALT = -500.0, 50000.0 # in feet
|
||||
MIN_IAS, MAX_IAS = 0.0, 800.0 # in knots
|
||||
|
||||
def min_max_normalize(value, min_val, max_val):
|
||||
return (value - min_val) / (max_val - min_val)
|
||||
|
||||
def min_max_denormalize(norm_value, min_val, max_val):
|
||||
return norm_value * (max_val - min_val) + min_val
|
||||
|
||||
def norm_coords(lat, lon, alt):
|
||||
lat_norm = min_max_normalize(lat, MIN_LAT, MAX_LAT)
|
||||
lon_norm = min_max_normalize(lon, MIN_LON, MAX_LON)
|
||||
alt_norm = min_max_normalize(alt, MIN_ALT, MAX_ALT)
|
||||
return lat_norm, lon_norm, alt_norm
|
||||
|
||||
def denorm_coords(lat_norm, lon_norm, alt_norm):
|
||||
lat = min_max_denormalize(lat_norm, MIN_LAT, MAX_LAT)
|
||||
lon = min_max_denormalize(lon_norm, MIN_LON, MAX_LON)
|
||||
alt = min_max_denormalize(alt_norm, MIN_ALT, MAX_ALT)
|
||||
return lat, lon, alt
|
||||
|
||||
def norm_ias(ias):
|
||||
return min_max_normalize(ias, MIN_IAS, MAX_IAS)
|
||||
|
||||
def denorm_ias(ias_norm):
|
||||
return min_max_denormalize(ias_norm, MIN_IAS, MAX_IAS)
|
||||
|
||||
class AutoEncoder(nn.Module):
|
||||
def __init__(self, input_a_1_size, input_a_2_size, input_b_size, latent_size):
|
||||
super(AutoEncoder, self).__init__()
|
||||
self.input_size = input_a_1_size + input_a_2_size + input_b_size
|
||||
self.input_a_1_size = input_a_1_size
|
||||
self.input_a_2_size = input_a_2_size
|
||||
self.input_b_size = input_b_size
|
||||
|
||||
self.encoder = nn.Sequential(
|
||||
nn.Linear(self.input_size, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, latent_size),
|
||||
)
|
||||
self.decoder_A_1 = nn.Sequential(
|
||||
nn.Linear(latent_size, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, self.input_a_1_size),
|
||||
nn.Softmax(dim=1),
|
||||
)
|
||||
self.decoder_A_2 = nn.Sequential(
|
||||
nn.Linear(latent_size, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, self.input_a_2_size),
|
||||
nn.Softmax(dim=1),
|
||||
)
|
||||
self.decoder_B = nn.Sequential(
|
||||
nn.Linear(latent_size, 128),
|
||||
nn.ReLU(),
|
||||
nn.Linear(128, self.input_b_size),
|
||||
nn.Softmax(dim=1),
|
||||
)
|
||||
def forward(self, X_a_1, X_a_2, X_b):
|
||||
X = torch.cat((X_a_1, X_a_2, X_b), dim=1)
|
||||
latent = self.encoder(X)
|
||||
reconstructed_A_1 = self.decoder_A_1(latent)
|
||||
reconstructed_A_2 = self.decoder_A_2(latent)
|
||||
reconstructed_B = self.decoder_B(latent)
|
||||
return reconstructed_A_1, reconstructed_A_2, reconstructed_B
|
||||
|
||||
category_max_limits = {
|
||||
"r_1": 100,
|
||||
"r_2": 50,
|
||||
"t": 100,
|
||||
}
|
||||
autoencoder = AutoEncoder(input_a_1_size=category_max_limits["r_1"], input_a_2_size=category_max_limits["r_2"], input_b_size=category_max_limits["t"], latent_size=4)
|
||||
auto_encoder_path = pathlib.Path(__file__).parent / "autoencoder.pth"
|
||||
if auto_encoder_path.exists():
|
||||
autoencoder.load_state_dict(torch.load(auto_encoder_path))
|
||||
autoencoder.eval()
|
||||
|
||||
def train_autoencoder(X_a_1, X_a_2, X_b, num_epochs=100, learning_rate=1e-3):
|
||||
criterion = nn.CrossEntropyLoss()
|
||||
optimizer = torch.optim.SGD(autoencoder.parameters(), lr=learning_rate)
|
||||
|
||||
for epoch in range(num_epochs):
|
||||
autoencoder.train()
|
||||
optimizer.zero_grad()
|
||||
rec_A_1, rec_A_2, rec_B = autoencoder(X_a_1, X_a_2, X_b)
|
||||
loss_A_1 = criterion(rec_A_1, torch.argmax(X_a_1, dim=1))
|
||||
loss_A_2 = criterion(rec_A_2, torch.argmax(X_a_2, dim=1))
|
||||
loss_B = criterion(rec_B, torch.argmax(X_b, dim=1))
|
||||
loss = loss_A_1 + loss_A_2 + loss_B
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
if (epoch + 1) % 10 == 0:
|
||||
print(f"Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}")
|
||||
|
||||
torch.save(autoencoder.state_dict(), auto_encoder_path)
|
||||
|
||||
def encode_features(X_a, X_b):
|
||||
autoencoder.eval()
|
||||
with torch.no_grad():
|
||||
X = torch.cat((X_a, X_b), dim=1)
|
||||
latent = autoencoder.encoder(X)
|
||||
return latent
|
||||
|
||||
mappings_path = pathlib.Path(__file__).parent / "category_mappings.json"
|
||||
|
||||
|
||||
def create_category_mappings(df):
|
||||
df = df.loc[:, ["r", "t"]].dropna().reset_index(drop=True)
|
||||
df.loc[:, "r_1"] = df["r"].str.split("-").str[0]
|
||||
df.loc[:, "r_2"] = df["r"].str.split("-").str[1].fillna("Empty")
|
||||
|
||||
for col, max_unique in category_max_limits.items():
|
||||
top_values = df[col].value_counts().head(max_unique-1).index.tolist()
|
||||
df.loc[~df[col].isin(top_values), col] = "Other"
|
||||
|
||||
for col in category_max_limits.keys():
|
||||
df.loc[:, col.upper()] = df[col].astype("category")
|
||||
|
||||
category_mappings = {}
|
||||
for col in category_max_limits.keys():
|
||||
category_mappings[col] = {cat: idx for idx, cat in enumerate(df[col.upper()].cat.categories)}
|
||||
with open(mappings_path, "w") as f:
|
||||
json.dump(category_mappings, f)
|
||||
return category_mappings
|
||||
|
||||
def map_categories(df, category_mappings):
|
||||
df = df.loc[:, ["r", "t"]].dropna().reset_index(drop=True)
|
||||
df.loc[:, "r_1"] = df["r"].str.split("-").str[0]
|
||||
df.loc[:, "r_2"] = df["r"].str.split("-").str[1].fillna("Empty")
|
||||
|
||||
for col in category_mappings.keys():
|
||||
mapping = category_mappings[col]
|
||||
df.loc[~df[col].isin(mapping.keys()), col] = "Other"
|
||||
df.loc[:, f"{col.upper()}_IDX"] = df[col].map(mapping).astype("int64")
|
||||
return df
|
||||
|
||||
|
||||
if mappings_path.exists():
|
||||
with open(mappings_path, "r") as f:
|
||||
category_mappings = json.load(f)
|
||||
else:
|
||||
print("Category mappings file not found. Please create mappings using 'create_category_mappings' function.")
|
||||
Reference in New Issue
Block a user