Skip to content

random_windows_pytorch_dataset

RandomWindowPytorchDataset

Bases: PytorchDataset

A PyTorch Dataset class that generates random windows on-the-fly for contrastive learning pretraining.

This class extends PytorchDataset to support random window generation without relying on predefined windows.

Parameters:

Name Type Description Default
cfg DictConfig

Configuration options for the dataset.

 required
split str

The data split to use (e.g., ‘train’, ‘validation’, ‘test’).

 required
min_window_size int

Minimum size of generated windows.

 required
max_window_size int

Maximum size of generated windows.

 required
n_windows int

Number of windows to generate for each sample.

 required
Source code in meds_torch/data/components/random_windows_pytorch_dataset.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
class RandomWindowPytorchDataset(PytorchDataset):
    """A PyTorch Dataset class that generates random windows on-the-fly for contrastive learning pretraining.

    This class extends PytorchDataset to support random window generation without relying on predefined
    windows.

    Args:
        cfg (DictConfig): Configuration options for the dataset.
        split (str): The data split to use (e.g., 'train', 'validation', 'test').
        min_window_size (int): Minimum size of generated windows.
        max_window_size (int): Maximum size of generated windows.
        n_windows (int): Number of windows to generate for each sample.
    """

    def __init__(self, cfg: DictConfig, split: str):
        super().__init__(cfg, split)
        self.min_window_size = cfg.min_window_size
        self.max_window_size = cfg.max_window_size
        self.n_windows = cfg.n_windows
        self.window_cols = [f"window_{i}" for i in range(cfg.n_windows)]
        self.cfg = cfg
        self.filter_index()

    def filter_index(self):
        filtered_index = []
        for i in range(len(self.index)):
            seq_length = self.index[i][2] - self.index[i][1]
            window_size = self.get_random_window_size(seq_length)
            if window_size < self.min_window_size:
                logger.warning(
                    f"Sequence length {seq_length} is too short to accommodate "
                    f"{self.n_windows} windows of minimum size {self.min_window_size}. "
                    f"Skipping subject {self.index[i][0]}."
                )
            else:
                filtered_index.append(self.index[i])
        if len(filtered_index) < len(self.index):
            logger.warning(f"Filtered data to {len(filtered_index) / len(self.index)}% of original data.")
        self.index = filtered_index

    def get_random_window_size(self, seq_length: int) -> int:
        """Calculates the size of the generated random windows."""
        total_window_size = min(seq_length, self.max_window_size * self.n_windows)
        window_size = total_window_size // self.n_windows
        return window_size

    def generate_random_windows(self, seq_length: int) -> list[tuple[int, int]]:
        """Generate random windows within a sequence.

        This method supports two modes of operation: (a) Random windows: Windows can overlap and be in any
        order. (b) Consecutive random windows: Windows are side by side and in order.

        The window sizes are predetermined to be the same and add up to less than the length of the dataset.

        Args:
            seq_length (int): Length of the sequence to generate windows from.

        Returns:
            List[Tuple[int, int]]: List of (start, end) indices for each window.

        Raises:
            ValueError: If the sequence length is too short to accommodate all windows.
        """
        # Use user-defined window names if provided, otherwise use default names
        window_names = (
            self.cfg.window_names
            if hasattr(self.cfg, "window_names")
            else [f"window_{i}" for i in range(self.n_windows)]
        )

        # Ensure we have the correct number of window names
        if len(window_names) != self.n_windows:
            raise ValueError(
                f"Number of window names ({len(window_names)}) does not match n_windows ({self.n_windows})"
            )

        window_size = self.get_random_window_size(seq_length)
        total_window_size = window_size * self.n_windows
        windows = []

        if self.cfg.get("consecutive_windows", False):
            # Mode (b): Consecutive random windows
            start = np.random.randint(0, seq_length - total_window_size + 1)
            for i in range(self.n_windows):
                end = start + window_size
                windows.append((start, end))
                start = end
        else:
            # Mode (a): Random windows
            for _ in range(self.n_windows):
                start = np.random.randint(0, seq_length - window_size + 1)
                end = start + window_size
                windows.append((start, end))

        # Sort windows by start index to maintain order
        windows.sort(key=lambda x: x[0])

        # Create a dictionary with user-defined or default window names
        named_windows = {name: window for name, window in zip(window_names, windows)}

        return named_windows

    def partition_sequence(self, sequence: dict, windows: list[tuple[int, int]]) -> dict:
        """Partition a sequence into multiple windows.

        Args:
            sequence (dict): The full sequence data.
            windows (List[Tuple[int, int]]): List of (start, end) indices for each window.

        Returns:
            dict: A dictionary with partitioned data for each window.
        """
        partitioned = {}
        for window_name, (start, end) in windows.items():
            partitioned[window_name] = {
                k: v[start:end] if isinstance(v, (list, np.ndarray)) else v for k, v in sequence.items()
            }
        return partitioned

    @SeedableMixin.WithSeed
    def _seeded_getitem(self, idx: int) -> dict:
        """Get a randomly windowed item from the dataset.

        Args:
            idx (int): Index of the item to retrieve.

        Returns:
            dict: A dictionary containing randomly generated windows of the sequence.
        """
        subject_id, _, _ = self.index[idx]
        shard = self.subj_map[subject_id]
        subject_idx = self.subj_indices[subject_id]

        subject_dynamic_data = JointNestedRaggedTensorDict(
            tensors_fp=Path(self.config.data_dir) / "data" / f"{shard}.nrt"
        )[subject_idx]

        full_sequence = self.load_subject(subject_dynamic_data, subject_id, 0, len(subject_dynamic_data))

        seq_len = len(full_sequence["dynamic"])
        windows = self.generate_random_windows(seq_len)
        partitioned_sequence = self.partition_sequence(full_sequence, windows)

        if self.config.do_include_subject_id:
            partitioned_sequence["subject_id"] = subject_id

        return partitioned_sequence

    def collate(self, batch: list[dict]) -> dict:
        """Collate a batch of randomly windowed sequences.

        Args:
            batch (List[dict]): A list of dictionaries, each containing windowed sequences.

        Returns:
            dict: A dictionary with collated data for each window.
        """
        out = {}
        for col in self.window_cols:
            if col in batch[0]:
                out[col] = super().collate([x[col] for x in batch])

        # Collate any additional data that's not part of the windows
        for key in batch[0].keys():
            if key not in self.window_cols and key not in out:
                out[key] = [x[key] for x in batch]

        return out

collate(batch)

Collate a batch of randomly windowed sequences.

Parameters:

Name Type Description Default
batch List[dict]

A list of dictionaries, each containing windowed sequences.

 required

Returns:

Name Type Description
dict dict

A dictionary with collated data for each window.
Source code in meds_torch/data/components/random_windows_pytorch_dataset.py 
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
def collate(self, batch: list[dict]) -> dict:
    """Collate a batch of randomly windowed sequences.

    Args:
        batch (List[dict]): A list of dictionaries, each containing windowed sequences.

    Returns:
        dict: A dictionary with collated data for each window.
    """
    out = {}
    for col in self.window_cols:
        if col in batch[0]:
            out[col] = super().collate([x[col] for x in batch])

    # Collate any additional data that's not part of the windows
    for key in batch[0].keys():
        if key not in self.window_cols and key not in out:
            out[key] = [x[key] for x in batch]

    return out

generate_random_windows(seq_length)

Generate random windows within a sequence.

This method supports two modes of operation: (a) Random windows: Windows can overlap and be in any order. (b) Consecutive random windows: Windows are side by side and in order.

The window sizes are predetermined to be the same and add up to less than the length of the dataset.

Parameters:

Name Type Description Default
seq_length int

Length of the sequence to generate windows from.

 required

Returns:

Type Description
list[tuple[int, int]]

List[Tuple[int, int]]: List of (start, end) indices for each window.

Raises:

Type Description
ValueError

If the sequence length is too short to accommodate all windows.
Source code in meds_torch/data/components/random_windows_pytorch_dataset.py 
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
def generate_random_windows(self, seq_length: int) -> list[tuple[int, int]]:
    """Generate random windows within a sequence.

    This method supports two modes of operation: (a) Random windows: Windows can overlap and be in any
    order. (b) Consecutive random windows: Windows are side by side and in order.

    The window sizes are predetermined to be the same and add up to less than the length of the dataset.

    Args:
        seq_length (int): Length of the sequence to generate windows from.

    Returns:
        List[Tuple[int, int]]: List of (start, end) indices for each window.

    Raises:
        ValueError: If the sequence length is too short to accommodate all windows.
    """
    # Use user-defined window names if provided, otherwise use default names
    window_names = (
        self.cfg.window_names
        if hasattr(self.cfg, "window_names")
        else [f"window_{i}" for i in range(self.n_windows)]
    )

    # Ensure we have the correct number of window names
    if len(window_names) != self.n_windows:
        raise ValueError(
            f"Number of window names ({len(window_names)}) does not match n_windows ({self.n_windows})"
        )

    window_size = self.get_random_window_size(seq_length)
    total_window_size = window_size * self.n_windows
    windows = []

    if self.cfg.get("consecutive_windows", False):
        # Mode (b): Consecutive random windows
        start = np.random.randint(0, seq_length - total_window_size + 1)
        for i in range(self.n_windows):
            end = start + window_size
            windows.append((start, end))
            start = end
    else:
        # Mode (a): Random windows
        for _ in range(self.n_windows):
            start = np.random.randint(0, seq_length - window_size + 1)
            end = start + window_size
            windows.append((start, end))

    # Sort windows by start index to maintain order
    windows.sort(key=lambda x: x[0])

    # Create a dictionary with user-defined or default window names
    named_windows = {name: window for name, window in zip(window_names, windows)}

    return named_windows

get_random_window_size(seq_length)

Calculates the size of the generated random windows.

Source code in meds_torch/data/components/random_windows_pytorch_dataset.py 
52
53
54
55
56
def get_random_window_size(self, seq_length: int) -> int:
    """Calculates the size of the generated random windows."""
    total_window_size = min(seq_length, self.max_window_size * self.n_windows)
    window_size = total_window_size // self.n_windows
    return window_size

partition_sequence(sequence, windows)

Partition a sequence into multiple windows.

Parameters:

Name Type Description Default
sequence dict

The full sequence data.

 required
windows List[Tuple[int, int]]

List of (start, end) indices for each window.

 required

Returns:

Name Type Description
dict dict

A dictionary with partitioned data for each window.
Source code in meds_torch/data/components/random_windows_pytorch_dataset.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
def partition_sequence(self, sequence: dict, windows: list[tuple[int, int]]) -> dict:
    """Partition a sequence into multiple windows.

    Args:
        sequence (dict): The full sequence data.
        windows (List[Tuple[int, int]]): List of (start, end) indices for each window.

    Returns:
        dict: A dictionary with partitioned data for each window.
    """
    partitioned = {}
    for window_name, (start, end) in windows.items():
        partitioned[window_name] = {
            k: v[start:end] if isinstance(v, (list, np.ndarray)) else v for k, v in sequence.items()
        }
    return partitioned