"""Lightning data module and batch utilities for TE regression training."""
from __future__ import annotations
# Import Python Utilities
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import Any
# Import YAML Utilities
import yaml
# Import PyTorch Lightning Utilities
from lightning.pytorch import LightningDataModule
# Import PyTorch Utilities
import torch
from torch.utils.data import DataLoader
# Import Dataset Utilities
from scripts.datasets.transmission_error_dataset import TransmissionErrorCurveDataset
from scripts.datasets.transmission_error_dataset import build_directional_file_manifest
from scripts.datasets.transmission_error_dataset import load_dataset_processing_config
from scripts.datasets.transmission_error_dataset import normalize_dataset_name
from scripts.datasets.transmission_error_dataset import normalize_input_mode
from scripts.datasets.transmission_error_dataset import resolve_dataset_schema
from scripts.datasets.transmission_error_dataset import resolve_input_mode_selection
from scripts.datasets.transmission_error_dataset import resolve_dataset_selection
from scripts.datasets.transmission_error_dataset import resolve_project_relative_path
from scripts.datasets.transmission_error_dataset import split_directional_file_manifest
COMMON_SPLIT_NAME_LIST = ("train", "validation", "test")
COMMON_SPLIT_DIRECTION_MAP = {
"Fw": "forward",
"Bw": "backward",
}
def split_directional_file_manifest_from_common_manifest(
directional_file_manifest: list[tuple[Path, str]],
split_manifest_path: str | Path,
excluded_condition_id_list: list[str] | None = None,
) -> tuple[list[tuple[Path, str]], list[tuple[Path, str]], list[tuple[Path, str]]]:
"""Apply a repository-owned paired-condition split manifest.
Args:
directional_file_manifest: Direction-filtered dataset file manifest.
split_manifest_path: YAML manifest containing paired Fw/Bw split
assignments.
excluded_condition_id_list: Optional condition identifiers quarantined
after the common split is applied.
Returns:
Train, validation, and test directional file manifests.
"""
# Load Common Split Manifest
resolved_split_manifest_path = resolve_project_relative_path(split_manifest_path)
assert resolved_split_manifest_path.is_file(), (
f"Common Split Manifest does not exist | {resolved_split_manifest_path}"
)
split_manifest_payload = yaml.safe_load(
resolved_split_manifest_path.read_text(encoding="utf-8")
)
assert isinstance(split_manifest_payload, dict), (
"Common Split Manifest must contain a YAML mapping"
)
manifest_entry_list = split_manifest_payload.get("entry_list")
assert isinstance(manifest_entry_list, list) and manifest_entry_list, (
"Common Split Manifest entry_list is empty"
)
# Index The Direction-Filtered Dataset Paths
directional_entry_by_path: dict[Path, tuple[Path, str]] = {}
for csv_file_path, direction_label in directional_file_manifest:
resolved_csv_file_path = Path(csv_file_path).resolve()
assert resolved_csv_file_path not in directional_entry_by_path, (
f"Directional dataset path appears more than once | {resolved_csv_file_path}"
)
directional_entry_by_path[resolved_csv_file_path] = (
resolved_csv_file_path,
direction_label,
)
# Apply Paired Split Assignments And Quarantine
excluded_condition_id_set = set(excluded_condition_id_list or [])
observed_condition_id_set: set[str] = set()
matched_directional_path_set: set[Path] = set()
split_directional_manifest_map: dict[str, list[tuple[Path, str]]] = {
split_name: [] for split_name in COMMON_SPLIT_NAME_LIST
}
for manifest_entry in manifest_entry_list:
condition_id = str(manifest_entry["condition_id"])
split_name = str(manifest_entry["split"])
assert split_name in COMMON_SPLIT_NAME_LIST, (
f"Unsupported Common Split name | {split_name}"
)
if condition_id in excluded_condition_id_set:
assert split_name == "train", (
"Quarantined conditions must belong to the training split | "
f"{condition_id} | {split_name}"
)
observed_condition_id_set.add(condition_id)
continue
direction_file_map = manifest_entry["direction_files"]
for common_direction_name, dataset_direction_label in COMMON_SPLIT_DIRECTION_MAP.items():
source_record = direction_file_map[common_direction_name]
resolved_source_path = resolve_project_relative_path(source_record["path"])
if resolved_source_path not in directional_entry_by_path:
continue
dataset_entry = directional_entry_by_path[resolved_source_path]
assert dataset_entry[1] == dataset_direction_label, (
"Common Split direction differs from dataset direction | "
f"{resolved_source_path} | {dataset_entry[1]} vs {dataset_direction_label}"
)
split_directional_manifest_map[split_name].append(dataset_entry)
matched_directional_path_set.add(resolved_source_path)
# Require Exact Quarantine And Dataset Coverage
assert observed_condition_id_set == excluded_condition_id_set, (
"Requested quarantined conditions were not found exactly in the Common Split | "
f"observed={sorted(observed_condition_id_set)} | "
f"expected={sorted(excluded_condition_id_set)}"
)
unmatched_directional_path_set = (
set(directional_entry_by_path) - matched_directional_path_set
)
quarantined_directional_path_count = len(excluded_condition_id_set) * (
int(any(label == "forward" for _, label in directional_file_manifest))
+ int(any(label == "backward" for _, label in directional_file_manifest))
)
assert len(unmatched_directional_path_set) == quarantined_directional_path_count, (
"Common Split does not cover the direction-filtered dataset exactly after "
f"quarantine | unmatched={len(unmatched_directional_path_set)} | "
f"expected_quarantined={quarantined_directional_path_count}"
)
return (
split_directional_manifest_map["train"],
split_directional_manifest_map["validation"],
split_directional_manifest_map["test"],
)
[docs]
@dataclass
class NormalizationStatistics:
"""Feature and target normalization tensors for TE regression."""
input_feature_mean: torch.Tensor
input_feature_std: torch.Tensor
target_mean: torch.Tensor
target_std: torch.Tensor
[docs]
@dataclass
class DatasetSplitSummary:
"""Curve counts for the train, validation, and test splits."""
train_curve_count: int
validation_curve_count: int
test_curve_count: int
dataset_name: str
dataset_schema: str
input_mode: str
input_feature_name_list: list[str]
target_feature_name_list: list[str]
input_feature_dim: int
target_feature_dim: int
[docs]
def move_batch_tensor_collection_to_device(batch_value: Any, device: torch.device, use_non_blocking_transfer: bool = False) -> Any:
"""Recursively move a nested batch structure to the requested device.
Args:
batch_value: Tensor, list, tuple, dictionary, or scalar batch value.
device: Target device for tensor transfer.
use_non_blocking_transfer: Whether CUDA transfers should request
non-blocking behavior when possible.
Returns:
Any: Batch structure mirrored on the target device.
"""
# Move Tensor To Requested Device
if isinstance(batch_value, torch.Tensor):
non_blocking_transfer_is_enabled = use_non_blocking_transfer and device.type == "cuda"
return batch_value.to(device, non_blocking=non_blocking_transfer_is_enabled)
# Recursively Move Dictionary Values
if isinstance(batch_value, dict):
return {key: move_batch_tensor_collection_to_device(value, device, use_non_blocking_transfer) for key, value in batch_value.items()}
# Recursively Move Tuple Values
if isinstance(batch_value, tuple):
return tuple(move_batch_tensor_collection_to_device(value, device, use_non_blocking_transfer) for value in batch_value)
# Recursively Move List Values
if isinstance(batch_value, list):
return [move_batch_tensor_collection_to_device(value, device, use_non_blocking_transfer) for value in batch_value]
return batch_value
[docs]
def collate_transmission_error_points(batch_dictionary_list: list[dict[str, Any]], point_stride: int = 1, maximum_points_per_curve: int | None = None, shuffle_points: bool = True) -> dict[str, Any]:
"""Collate curve samples into one point-level training batch.
Args:
batch_dictionary_list: Curve-level samples produced by the dataset.
point_stride: Step used to subsample points from each curve.
maximum_points_per_curve: Optional cap on sampled points per curve.
shuffle_points: Whether to shuffle the concatenated point batch.
Returns:
dict[str, Any]: Batch dictionary containing concatenated tensors and
lightweight curve-level metadata.
"""
# Validate Batch Input
assert len(batch_dictionary_list) > 0, "Batch Dictionary List is empty"
# Initialize Point Lists
input_tensor_list: list[torch.Tensor] = []
target_tensor_list: list[torch.Tensor] = []
angular_position_tensor_list: list[torch.Tensor] = []
point_count_per_curve: list[int] = []
# Extract Point Samples From Each Curve
for curve_sample_dictionary in batch_dictionary_list:
# Extract Point Sample Dictionary From Curve Sample Dictionary
point_sample_dictionary = extract_point_tensor_from_curve_sample(
curve_sample_dictionary,
point_stride,
maximum_points_per_curve,
)
# Append Point Tensors To Batch Lists
input_tensor_list.append(point_sample_dictionary["input_tensor"])
target_tensor_list.append(point_sample_dictionary["target_tensor"])
angular_position_tensor_list.append(point_sample_dictionary["angular_position_deg"])
point_count_per_curve.append(point_sample_dictionary["input_tensor"].shape[0])
# Concatenate Point Tensors
input_tensor = torch.cat(input_tensor_list, dim=0)
target_tensor = torch.cat(target_tensor_list, dim=0)
angular_position_deg = torch.cat(angular_position_tensor_list, dim=0)
# Shuffle Points Inside Batch
if shuffle_points and input_tensor.shape[0] > 1:
permutation_indices = torch.randperm(input_tensor.shape[0])
input_tensor = input_tensor.index_select(0, permutation_indices)
target_tensor = target_tensor.index_select(0, permutation_indices)
angular_position_deg = angular_position_deg.index_select(0, permutation_indices)
return {
"input_tensor": input_tensor,
"target_tensor": target_tensor,
"angular_position_deg": angular_position_deg,
"point_count_per_curve": torch.tensor(point_count_per_curve, dtype=torch.long),
"curve_count": len(batch_dictionary_list),
"direction_label": [curve_sample_dictionary["direction_label"] for curve_sample_dictionary in batch_dictionary_list],
"source_file_path": [curve_sample_dictionary["source_file_path"] for curve_sample_dictionary in batch_dictionary_list],
}
def extract_sequence_tensor_from_curve_sample(
curve_sample_dictionary: dict[str, Any],
point_stride: int = 1,
sequence_length: int = 17,
sequence_stride: int = 1,
target_position: str = "center",
maximum_sequences_per_curve: int | None = None,
) -> dict[str, torch.Tensor]:
"""Convert one curve sample into overlapping temporal sequence tensors.
Args:
curve_sample_dictionary: Dataset sample containing full-curve tensors.
point_stride: Step used before sequence-window construction.
sequence_length: Number of sampled points included in each window.
sequence_stride: Step between consecutive temporal windows.
target_position: Window position used as the regression target point.
maximum_sequences_per_curve: Optional cap on windows per curve.
Returns:
dict[str, torch.Tensor]: Sequence-level input tensors plus target and
angular-position tensors aligned with the configured target point.
"""
# Validate Sequence Sampling Parameters
assert sequence_length > 0, f"Sequence Length must be positive | {sequence_length}"
assert sequence_stride > 0, f"Sequence Stride must be positive | {sequence_stride}"
if maximum_sequences_per_curve is not None: assert maximum_sequences_per_curve > 0, (f"Maximum Sequences Per Curve must be positive | {maximum_sequences_per_curve}")
normalized_target_position = target_position.strip().lower()
assert normalized_target_position in ["center", "last"], f"Unsupported Sequence Target Position | {target_position}"
if normalized_target_position == "center": assert sequence_length % 2 == 1, f"Center target requires an odd Sequence Length | {sequence_length}"
# Start From Point-Sampled Curve Tensors
point_sample_dictionary = extract_point_tensor_from_curve_sample(
curve_sample_dictionary=curve_sample_dictionary,
point_stride=point_stride,
maximum_points_per_curve=None,
)
input_tensor = point_sample_dictionary["input_tensor"]
target_tensor = point_sample_dictionary["target_tensor"]
angular_position_deg = point_sample_dictionary["angular_position_deg"]
# Validate Enough Sampled Points Exist For At Least One Window
sampled_point_count = int(input_tensor.shape[0])
assert sampled_point_count >= sequence_length, (
f"Sampled curve is shorter than Sequence Length | {sampled_point_count} < {sequence_length}"
)
# Build Window Start Indices
sequence_start_index_tensor = torch.arange(
0,
sampled_point_count - sequence_length + 1,
sequence_stride,
dtype=torch.long,
)
assert len(sequence_start_index_tensor) > 0, "Sequence Index Tensor is empty after applying Sequence Stride"
# Reduce Number Of Windows Per Curve
if maximum_sequences_per_curve is not None and len(sequence_start_index_tensor) > maximum_sequences_per_curve:
reduced_index_positions = torch.linspace(
0,
len(sequence_start_index_tensor) - 1,
steps=maximum_sequences_per_curve,
dtype=torch.float32,
).round().long()
sequence_start_index_tensor = sequence_start_index_tensor.index_select(0, reduced_index_positions)
# Stack Temporal Windows
sequence_input_tensor_list: list[torch.Tensor] = []
target_index_list: list[int] = []
target_offset = sequence_length // 2 if normalized_target_position == "center" else sequence_length - 1
for sequence_start_index in sequence_start_index_tensor.tolist():
sequence_end_index = sequence_start_index + sequence_length
sequence_input_tensor_list.append(input_tensor[sequence_start_index:sequence_end_index])
target_index_list.append(sequence_start_index + target_offset)
sequence_input_tensor = torch.stack(sequence_input_tensor_list, dim=0)
target_index_tensor = torch.tensor(target_index_list, dtype=torch.long)
sequence_target_tensor = target_tensor.index_select(0, target_index_tensor)
sequence_angular_position_deg = angular_position_deg.index_select(0, target_index_tensor)
return {
"input_tensor": sequence_input_tensor,
"target_tensor": sequence_target_tensor,
"angular_position_deg": sequence_angular_position_deg,
}
def collate_transmission_error_sequences(
batch_dictionary_list: list[dict[str, Any]],
point_stride: int = 1,
sequence_length: int = 17,
sequence_stride: int = 1,
target_position: str = "center",
maximum_sequences_per_curve: int | None = None,
shuffle_sequences: bool = True,
) -> dict[str, Any]:
"""Collate curve samples into one sequence-window training batch."""
# Validate Batch Input
assert len(batch_dictionary_list) > 0, "Batch Dictionary List is empty"
# Initialize Sequence Lists
input_tensor_list: list[torch.Tensor] = []
target_tensor_list: list[torch.Tensor] = []
angular_position_tensor_list: list[torch.Tensor] = []
sequence_count_per_curve: list[int] = []
# Extract Sequence Samples From Each Curve
for curve_sample_dictionary in batch_dictionary_list:
sequence_sample_dictionary = extract_sequence_tensor_from_curve_sample(
curve_sample_dictionary=curve_sample_dictionary,
point_stride=point_stride,
sequence_length=sequence_length,
sequence_stride=sequence_stride,
target_position=target_position,
maximum_sequences_per_curve=maximum_sequences_per_curve,
)
input_tensor_list.append(sequence_sample_dictionary["input_tensor"])
target_tensor_list.append(sequence_sample_dictionary["target_tensor"])
angular_position_tensor_list.append(sequence_sample_dictionary["angular_position_deg"])
sequence_count_per_curve.append(sequence_sample_dictionary["input_tensor"].shape[0])
# Concatenate Sequence Tensors
input_tensor = torch.cat(input_tensor_list, dim=0)
target_tensor = torch.cat(target_tensor_list, dim=0)
angular_position_deg = torch.cat(angular_position_tensor_list, dim=0)
# Shuffle Windows Inside Batch
if shuffle_sequences and input_tensor.shape[0] > 1:
permutation_indices = torch.randperm(input_tensor.shape[0])
input_tensor = input_tensor.index_select(0, permutation_indices)
target_tensor = target_tensor.index_select(0, permutation_indices)
angular_position_deg = angular_position_deg.index_select(0, permutation_indices)
return {
"input_tensor": input_tensor,
"target_tensor": target_tensor,
"angular_position_deg": angular_position_deg,
"sequence_count_per_curve": torch.tensor(sequence_count_per_curve, dtype=torch.long),
"curve_count": len(batch_dictionary_list),
"direction_label": [curve_sample_dictionary["direction_label"] for curve_sample_dictionary in batch_dictionary_list],
"source_file_path": [curve_sample_dictionary["source_file_path"] for curve_sample_dictionary in batch_dictionary_list],
}
[docs]
class TransmissionErrorDataModule(LightningDataModule):
"""LightningDataModule for TE curve splits, sampling, and normalization."""
[docs]
def __init__(
self,
dataset_config_path: str | Path,
dataset_name: str | None = None,
input_mode: str | None = None,
curve_batch_size: int = 2,
point_stride: int = 20,
maximum_points_per_curve: int | None = None,
collate_mode: str = "point",
sequence_length: int = 17,
sequence_stride: int = 1,
sequence_target_position: str = "center",
maximum_sequences_per_curve: int | None = None,
maximum_normalization_curve_count: int | None = None,
shuffle_training_batch_elements: bool = True,
num_workers: int = 0,
pin_memory: bool = False,
use_non_blocking_transfer: bool = False,
use_forward_direction: bool = True,
use_backward_direction: bool = True,
split_manifest_path: str | Path | None = None,
excluded_condition_id_list: list[str] | None = None,
expected_curve_count_by_split: dict[str, int] | None = None,
) -> None:
"""Initialize the reusable TE training datamodule.
Args:
dataset_config_path: Dataset YAML configuration path.
dataset_name: Optional dataset selector overriding the YAML value.
input_mode: Optional selector for setpoints or actual values.
curve_batch_size: Number of curves loaded per dataloader batch.
point_stride: Subsampling stride applied inside each curve.
maximum_points_per_curve: Optional cap on sampled points per curve.
collate_mode: Batch collation mode, either ``point`` or
``sequence``.
sequence_length: Number of point-sampled timesteps per sequence.
sequence_stride: Step between consecutive temporal windows.
sequence_target_position: Window point aligned with the target.
maximum_sequences_per_curve: Optional cap on windows per curve.
maximum_normalization_curve_count: Optional cap on training curves
scanned for normalization statistics.
shuffle_training_batch_elements: Whether training batches should
shuffle points or windows after curve-level collation.
num_workers: PyTorch dataloader worker count.
pin_memory: Whether dataloaders should pin host memory.
use_non_blocking_transfer: Whether device transfer should request
non-blocking CUDA copies when possible.
use_forward_direction: Whether forward-direction curves are
included in the dataset manifest.
use_backward_direction: Whether backward-direction curves are
included in the dataset manifest.
split_manifest_path: Optional paired-condition split manifest that
replaces the generated random split.
excluded_condition_id_list: Optional manifest condition IDs to
quarantine after assigning splits.
expected_curve_count_by_split: Optional exact post-quarantine curve
counts for train, validation, and test.
"""
super().__init__()
# Validate Dataloader Parameters
assert curve_batch_size > 0, f"Curve Batch Size must be positive | {curve_batch_size}"
assert point_stride > 0, f"Point Stride must be positive | {point_stride}"
assert collate_mode in ["point", "sequence"], f"Unsupported Collate Mode | {collate_mode}"
assert sequence_length > 0, f"Sequence Length must be positive | {sequence_length}"
assert sequence_stride > 0, f"Sequence Stride must be positive | {sequence_stride}"
if maximum_normalization_curve_count is not None: assert maximum_normalization_curve_count > 0, (f"Maximum Normalization Curve Count must be positive | {maximum_normalization_curve_count}")
assert num_workers >= 0, f"Num Workers must be non-negative | {num_workers}"
assert use_forward_direction or use_backward_direction, "At least one direction must be enabled"
# Save Dataset Parameters
self.dataset_config_path = resolve_project_relative_path(dataset_config_path)
self.requested_dataset_name = normalize_dataset_name(dataset_name) if dataset_name is not None else None
self.requested_input_mode = input_mode
self.curve_batch_size = curve_batch_size
self.point_stride = point_stride
self.maximum_points_per_curve = maximum_points_per_curve
self.collate_mode = collate_mode
self.sequence_length = sequence_length
self.sequence_stride = sequence_stride
self.sequence_target_position = sequence_target_position
self.maximum_sequences_per_curve = maximum_sequences_per_curve
self.maximum_normalization_curve_count = maximum_normalization_curve_count
self.shuffle_training_batch_elements = shuffle_training_batch_elements
self.num_workers = num_workers
self.pin_memory = pin_memory
self.use_non_blocking_transfer = use_non_blocking_transfer
self.use_forward_direction = use_forward_direction
self.use_backward_direction = use_backward_direction
self.split_manifest_path = (
resolve_project_relative_path(split_manifest_path)
if split_manifest_path is not None
else None
)
self.excluded_condition_id_list = list(excluded_condition_id_list or [])
self.expected_curve_count_by_split = (
{
split_name: int(curve_count)
for split_name, curve_count in expected_curve_count_by_split.items()
}
if expected_curve_count_by_split is not None
else None
)
if self.expected_curve_count_by_split is not None:
assert set(self.expected_curve_count_by_split) == set(COMMON_SPLIT_NAME_LIST), (
"Expected Curve Count map must define train, validation, and test"
)
# Initialize Runtime Attributes
self.train_dataset: TransmissionErrorCurveDataset | None = None
self.validation_dataset: TransmissionErrorCurveDataset | None = None
self.test_dataset: TransmissionErrorCurveDataset | None = None
self.input_feature_dim: int | None = None
self.target_feature_dim: int | None = None
self.normalization_statistics: NormalizationStatistics | None = None
self.dataset_name: str | None = None
self.dataset_schema: str | None = None
self.input_mode: str | None = None
self.input_feature_name_list: list[str] = []
self.target_feature_name_list: list[str] = []
[docs]
def setup(self, stage: str | None = None) -> None:
"""Create dataset splits and compute normalization statistics.
Args:
stage: Lightning stage hint. The current implementation uses the
same prepared splits for fit, validation, and test access.
"""
# Skip Repeated Setup If Datasets Are Already Available
if self.train_dataset is not None and self.validation_dataset is not None and self.test_dataset is not None: return
# Load Dataset Processing Configuration
dataset_processing_config = load_dataset_processing_config(self.dataset_config_path)
# Resolve Dataset Selection And Root
selected_dataset_name, dataset_root = resolve_dataset_selection(
dataset_processing_config,
self.requested_dataset_name,
)
selected_input_mode = resolve_input_mode_selection(
dataset_processing_config,
selected_dataset_name,
self.requested_input_mode,
)
dataset_schema = resolve_dataset_schema(selected_dataset_name, selected_input_mode)
self.dataset_name = dataset_schema.dataset_name
self.dataset_schema = dataset_schema.schema_name
self.input_mode = normalize_input_mode(selected_dataset_name, selected_input_mode)
self.input_feature_name_list = list(dataset_schema.input_feature_name_list)
self.target_feature_name_list = list(dataset_schema.target_feature_name_list)
# Build Directional File Manifest
directional_file_manifest = build_directional_file_manifest(
dataset_root,
self.use_forward_direction,
self.use_backward_direction,
selected_dataset_name,
)
# Split Directional File Manifest
if self.split_manifest_path is not None:
(
train_directional_file_manifest,
validation_directional_file_manifest,
test_directional_file_manifest,
) = split_directional_file_manifest_from_common_manifest(
directional_file_manifest=directional_file_manifest,
split_manifest_path=self.split_manifest_path,
excluded_condition_id_list=self.excluded_condition_id_list,
)
else:
(
train_directional_file_manifest,
validation_directional_file_manifest,
test_directional_file_manifest,
) = split_directional_file_manifest(
directional_file_manifest,
float(dataset_processing_config["split"]["validation_split"]),
float(dataset_processing_config["split"].get("test_split", 0.0)),
int(dataset_processing_config["split"]["random_seed"]),
)
# Validate Exact Post-Quarantine Split Counts
if self.expected_curve_count_by_split is not None:
observed_curve_count_by_split = {
"train": len(train_directional_file_manifest),
"validation": len(validation_directional_file_manifest),
"test": len(test_directional_file_manifest),
}
assert observed_curve_count_by_split == self.expected_curve_count_by_split, (
"Dataset split differs from the configured exact curve-count contract | "
f"observed={observed_curve_count_by_split} | "
f"expected={self.expected_curve_count_by_split}"
)
# Build Train Dataset Object
self.train_dataset = TransmissionErrorCurveDataset(
dataset_root=dataset_root,
dataset_name=selected_dataset_name,
input_mode=selected_input_mode,
directional_file_manifest=train_directional_file_manifest,
)
# Build Validation Dataset Object
self.validation_dataset = TransmissionErrorCurveDataset(
dataset_root=dataset_root,
dataset_name=selected_dataset_name,
input_mode=selected_input_mode,
directional_file_manifest=validation_directional_file_manifest,
)
# Build Test Dataset Object
self.test_dataset = TransmissionErrorCurveDataset(
dataset_root=dataset_root,
dataset_name=selected_dataset_name,
input_mode=selected_input_mode,
directional_file_manifest=test_directional_file_manifest,
)
# Read Feature Dimensions
first_train_sample = self.train_dataset[0]
self.input_feature_dim = int(first_train_sample["input_tensor"].shape[-1])
self.target_feature_dim = int(first_train_sample["target_tensor"].shape[-1])
# Compute Normalization Statistics
self.normalization_statistics = self.compute_normalization_statistics(self.train_dataset)
[docs]
def compute_normalization_statistics(self, curve_dataset: TransmissionErrorCurveDataset) -> NormalizationStatistics:
"""Compute point-level normalization statistics from the train split.
Args:
curve_dataset: Dataset used to accumulate sampled point statistics.
Returns:
NormalizationStatistics: Mean and standard deviation tensors for
model inputs and targets.
"""
# Initialize Accumulators
input_feature_sum = torch.zeros(self.input_feature_dim, dtype=torch.float64)
input_feature_squared_sum = torch.zeros(self.input_feature_dim, dtype=torch.float64)
target_sum = torch.zeros(self.target_feature_dim, dtype=torch.float64)
target_squared_sum = torch.zeros(self.target_feature_dim, dtype=torch.float64)
total_point_count = 0
normalization_curve_count = len(curve_dataset)
if self.maximum_normalization_curve_count is not None:
normalization_curve_count = min(normalization_curve_count, self.maximum_normalization_curve_count)
# Scan Training Curves
for curve_index in range(normalization_curve_count):
curve_sample_dictionary = curve_dataset[curve_index]
point_sample_dictionary = extract_point_tensor_from_curve_sample(
curve_sample_dictionary,
self.point_stride,
self.maximum_points_per_curve,
)
input_tensor = point_sample_dictionary["input_tensor"].double()
target_tensor = point_sample_dictionary["target_tensor"].double()
input_feature_sum += input_tensor.sum(dim=0)
input_feature_squared_sum += torch.square(input_tensor).sum(dim=0)
target_sum += target_tensor.sum(dim=0)
target_squared_sum += torch.square(target_tensor).sum(dim=0)
total_point_count += int(input_tensor.shape[0])
# Validate Point Count
assert total_point_count > 0, "Normalization Statistics cannot be computed from an empty training set"
# Compute Means
input_feature_mean = input_feature_sum / total_point_count
target_mean = target_sum / total_point_count
# Compute Variances
input_feature_variance = (input_feature_squared_sum / total_point_count) - torch.square(input_feature_mean)
target_variance = (target_squared_sum / total_point_count) - torch.square(target_mean)
# Clamp Variances
input_feature_variance = torch.clamp(input_feature_variance, min=1.0e-12)
target_variance = torch.clamp(target_variance, min=1.0e-12)
# Compute Standard Deviations
input_feature_std = torch.sqrt(input_feature_variance)
target_std = torch.sqrt(target_variance)
return NormalizationStatistics(
input_feature_mean=input_feature_mean.float(),
input_feature_std=input_feature_std.float(),
target_mean=target_mean.float(),
target_std=target_std.float(),
)
[docs]
def get_target_feature_dim(self) -> int:
"""Return the resolved target feature dimension after setup."""
assert self.target_feature_dim is not None, "Target Feature Dim is not available before setup"
return self.target_feature_dim
[docs]
def get_normalization_statistics(self) -> NormalizationStatistics:
"""Return the cached normalization statistics after setup."""
assert self.normalization_statistics is not None, "Normalization Statistics are not available before setup"
return self.normalization_statistics
[docs]
def get_dataset_split_summary(self) -> DatasetSplitSummary:
"""Return the current split sizes in number of curves."""
assert self.train_dataset is not None, "Train Dataset is not initialized"
assert self.validation_dataset is not None, "Validation Dataset is not initialized"
assert self.test_dataset is not None, "Test Dataset is not initialized"
return DatasetSplitSummary(
train_curve_count=len(self.train_dataset),
validation_curve_count=len(self.validation_dataset),
test_curve_count=len(self.test_dataset),
dataset_name=str(self.dataset_name),
dataset_schema=str(self.dataset_schema),
input_mode=str(self.input_mode),
input_feature_name_list=list(self.input_feature_name_list),
target_feature_name_list=list(self.target_feature_name_list),
input_feature_dim=self.get_input_feature_dim(),
target_feature_dim=self.get_target_feature_dim(),
)
[docs]
def train_dataloader(self) -> DataLoader:
"""Build the training dataloader with configured collation."""
assert self.train_dataset is not None, "Train Dataset is not initialized"
return DataLoader(
self.train_dataset,
batch_size=self.curve_batch_size,
shuffle=True,
num_workers=self.num_workers,
persistent_workers=(self.num_workers > 0),
pin_memory=self.pin_memory,
collate_fn=self._build_collate_function(shuffle_batch_elements=self.shuffle_training_batch_elements),
)
[docs]
def val_dataloader(self) -> DataLoader:
"""Build the validation dataloader with deterministic ordering."""
assert self.validation_dataset is not None, "Validation Dataset is not initialized"
return DataLoader(
self.validation_dataset,
batch_size=self.curve_batch_size,
shuffle=False,
num_workers=self.num_workers,
persistent_workers=(self.num_workers > 0),
pin_memory=self.pin_memory,
collate_fn=self._build_collate_function(shuffle_batch_elements=False),
)
[docs]
def test_dataloader(self) -> DataLoader:
"""Build the test dataloader with deterministic ordering."""
assert self.test_dataset is not None, "Test Dataset is not initialized"
return DataLoader(
self.test_dataset,
batch_size=self.curve_batch_size,
shuffle=False,
num_workers=self.num_workers,
persistent_workers=(self.num_workers > 0),
pin_memory=self.pin_memory,
collate_fn=self._build_collate_function(shuffle_batch_elements=False),
)
def _build_collate_function(self, shuffle_batch_elements: bool):
"""Return the configured point or sequence collate function."""
# Build Point-Level Collation
if self.collate_mode == "point":
return partial(
collate_transmission_error_points,
point_stride=self.point_stride,
maximum_points_per_curve=self.maximum_points_per_curve,
shuffle_points=shuffle_batch_elements,
)
# Build Sequence-Level Collation
return partial(
collate_transmission_error_sequences,
point_stride=self.point_stride,
sequence_length=self.sequence_length,
sequence_stride=self.sequence_stride,
target_position=self.sequence_target_position,
maximum_sequences_per_curve=self.maximum_sequences_per_curve,
shuffle_sequences=shuffle_batch_elements,
)
[docs]
def transfer_batch_to_device(self, batch: Any, device: torch.device, dataloader_idx: int) -> Any:
"""Move a collated batch to the target accelerator device.
Args:
batch: Batch emitted by one of the TE dataloaders.
device: Target Lightning device.
dataloader_idx: Dataloader index supplied by Lightning.
Returns:
Any: Batch moved to the requested device.
"""
return move_batch_tensor_collection_to_device(
batch,
device,
self.use_non_blocking_transfer,
)