---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[20], line 1
----> 1 next(iter(train_loader))[1].shape

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.__next__(self)
    627 if self._sampler_iter is None:
    628     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    629     self._reset()  # type: ignore[call-arg]
--> 630 data = self._next_data()
    631 self._num_yielded += 1
    632 if self._dataset_kind == _DatasetKind.Iterable and \
    633         self._IterableDataset_len_called is not None and \
    634         self._num_yielded > self._IterableDataset_len_called:

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:673, in _SingleProcessDataLoaderIter._next_data(self)
    671 def _next_data(self):
    672     index = self._next_index()  # may raise StopIteration
--> 673     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    674     if self._pin_memory:
    675         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:52, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     50         data = self.dataset.__getitems__(possibly_batched_index)
     51     else:
---> 52         data = [self.dataset[idx] for idx in possibly_batched_index]
     53 else:
     54     data = self.dataset[possibly_batched_index]

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:52, in <listcomp>(.0)
     50         data = self.dataset.__getitems__(possibly_batched_index)
     51     else:
---> 52         data = [self.dataset[idx] for idx in possibly_batched_index]
     53 else:
     54     data = self.dataset[possibly_batched_index]

Cell In[17], line 29, in TimeSeriesDataset.__getitem__(self, idx)
     26 if self.augment:
     27     # Apply time warping to input_series
     28     warp_factor = np.random.uniform(self.warp_range[0], self.warp_range[1])
---> 29     input_series = self.time_warp(input_series, warp_factor)
     30     # Optional: Apply the same warp to target_series if needed
     31     # For sequence-to-sequence tasks, you may choose to warp the target as well
     32     # target_series = self.time_warp(target_series, warp_factor)
     33 
     34 # Convert to torch tensors
     35 input_tensor = torch.tensor(input_series, dtype=torch.float32)

Cell In[17], line 62, in TimeSeriesDataset.time_warp(self, series, warp_factor)
     60 warped_series = np.zeros((warped_length, n_features))
     61 for i in range(n_features):
---> 62     f = interp1d(original_indices, series[:, i], kind='linear', fill_value='extrapolate')
     63     warped_series[:, i] = f(warped_positions)
     64 return warped_series

NameError: name 'interp1d' is not defined

in_sequence_length = 2
out_sequence_length = 5
batch_size_ = 128
train_loader, val_loader, test_loader = create_data_loaders(train_data, val_data, test_data, input_seq_length = in_sequence_length, output_seq_length = out_sequence_length, batch_size = batch_size_, train_only = train_only_)


def create_data_loaders(train_data, val_data, test_data, input_seq_length, output_seq_length, batch_size, train_only=False):
    """
    Creates DataLoaders for training, validation, and test datasets with optional augmentation for training data.

    Parameters:
    - train_data: pd.DataFrame or np.array for training.
    - val_data: pd.DataFrame or np.array for validation.
    - test_data: pd.DataFrame or np.array for testing.
    - input_seq_length: int, the length of input sequences.
    - output_seq_length: int, the length of output sequences.
    - batch_size: int, batch size for DataLoaders.
    - train_only: bool, whether to return only the training DataLoader.

    Returns:
    - train_loader: DataLoader for training data.
    - val_loader: DataLoader for validation data (if train_only=False).
    - test_loader: DataLoader for test data (if train_only=False).
    """
    # Training dataset with augmentation enabled
    train_dataset = TimeSeriesDataset(train_data, input_seq_length, output_seq_length, augment=True)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    
    if not train_only:
        # Validation and test datasets without augmentation
        val_dataset = TimeSeriesDataset(val_data, input_seq_length, output_seq_length, augment=False)
        test_dataset = TimeSeriesDataset(test_data, input_seq_length, output_seq_length, augment=False)
        val_loader = DataLoader(val_dataset, batch_size=batch_size)
        test_loader = DataLoader(test_dataset, batch_size=batch_size)
        return train_loader, val_loader, test_loader

    return train_loader, None, Non



class TimeSeriesDataset(Dataset):
    def __init__(self, data, input_seq_length, output_seq_length, augment=False, warp_range=(0.8, 1.2)):
        """
        TimeSeriesDataset with optional on-the-fly time warping augmentation.

        Parameters:
        - data: pd.DataFrame or np.array containing the time series data.
        - input_seq_length: int, the length of input sequences.
        - output_seq_length: int, the length of output sequences.
        - augment: bool, whether to apply time warping augmentation.
        - warp_range: tuple(float, float), range of warp factors for augmentation.
        """
        self.data = data.values if isinstance(data, pd.DataFrame) else data
        self.input_seq_length = input_seq_length
        self.output_seq_length = output_seq_length
        self.augment = augment
        self.warp_range = warp_range

    def __len__(self):
        return len(self.data) - self.input_seq_length - self.output_seq_length + 1

    def __getitem__(self, idx):
        input_series = self.data[idx : idx + self.input_seq_length].copy()
        target_series = self.data[idx + self.input_seq_length : idx + self.input_seq_length + self.output_seq_length].copy()
        
        if self.augment:
            # Apply time warping to input_series
            warp_factor = np.random.uniform(self.warp_range[0], self.warp_range[1])
            input_series = self.time_warp(input_series, warp_factor)
            # Optional: Apply the same warp to target_series if needed
            # For sequence-to-sequence tasks, you may choose to warp the target as well
            # target_series = self.time_warp(target_series, warp_factor)
        
        # Convert to torch tensors
        input_tensor = torch.tensor(input_series, dtype=torch.float32)
        target_tensor = torch.tensor(target_series, dtype=torch.float32)
        
        return input_tensor, target_tensor

    def time_warp(self, series, warp_factor):
        """
        Applies time warping to a time series by stretching or compressing the time dimension.

        Parameters:
        - series: np.array of shape (seq_length, n_features)
        - warp_factor: float > 0, warp factor where >1 speeds up the series, <1 slows it down

        Returns:
        - warped_series: np.array of shape (seq_length, n_features)
        """
        seq_length, n_features = series.shape
        # Original time indices
        original_indices = np.linspace(0, 1, seq_length)
        # Warped time indices
        warped_length = seq_length  # Keep the length consistent
        warped_indices = np.linspace(0, 1, warped_length)
        # Apply warp factor to the positions
        warped_positions = original_indices ** (1 / warp_factor)
        # Interpolate each feature
        warped_series = np.zeros((warped_length, n_features))
        for i in range(n_features):
            f = interp1d(original_indices, series[:, i], kind='linear', fill_value='extrapolate')
            warped_series[:, i] = f(warped_positions)
        return warped_series