Skip to content

utils

TrajectoryBatch dataclass

Initialize a batch of trajectories.

Parameters:

Name Type Description Default
time Tensor

Tensor of shape (batch_size, sequence_length) containing days after prediction time. Values must be monotonically increasing within each sequence.

 required
code Tensor

Tensor of shape (batch_size, sequence_length) containing event code vocabulary indices

 required
mask Tensor

Tensor of shape (batch_size, sequence_length) indicates valid measurements/codes

 required
numeric_value Tensor

Tensor of shape (batch_size, sequence_length) containing numeric values

 required
numeric_value_mask Tensor

Tensor of shape (batch_size, sequence_length) indicating valid numeric values

 required
metadata_df DataFrame

DataFrame containing code vocabulary mapping with ‘code’ and ‘code/vocab_index’ columns

 required
time_scale str

scale of the time, by default it is ‘Y’ (years). Any numpy datetime units can be used, see https://numpy.org/doc/2.1/reference/arrays.datetime.html#datetime-units

 'Y'
Source code in meds_torch/models/components/utils.py 
 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
@dataclass
class TrajectoryBatch:
    """
    Initialize a batch of trajectories.

    Args:
        time (torch.Tensor): Tensor of shape (batch_size, sequence_length) containing days after prediction
            time. Values must be monotonically increasing within each sequence.
        code (torch.Tensor): Tensor of shape (batch_size, sequence_length) containing event code vocabulary
            indices
        mask (torch.Tensor): Tensor of shape (batch_size, sequence_length) indicates valid measurements/codes
        numeric_value (torch.Tensor): Tensor of shape (batch_size, sequence_length) containing numeric values
        numeric_value_mask (torch.Tensor): Tensor of shape (batch_size, sequence_length) indicating valid
            numeric values
        metadata_df (pl.DataFrame): DataFrame containing code vocabulary mapping with 'code' and
            'code/vocab_index' columns
        time_scale: scale of the time, by default it is 'Y' (years). Any numpy datetime units can be used,
            see https://numpy.org/doc/2.1/reference/arrays.datetime.html#datetime-units
    """

    time: torch.Tensor
    code: torch.Tensor
    mask: torch.Tensor
    numeric_value: torch.Tensor
    numeric_value_mask: torch.Tensor
    metadata_df: pl.DataFrame
    time_scale: str = "Y"

    def to_meds(self, prediction_time: list[datetime], subject_id: list[str | int]) -> pl.DataFrame:
        """Convert the trajectory batch to MEDS format.

        Args:
            prediction_time: List of prediction times for each trajectory in the batch
            subject_id: List of subject IDs for each trajectory in the batch

        Returns:
            pl.DataFrame: MEDS format DataFrame with columns:
                - time: Absolute timestamp of the event
                - code: The medical code string
                - numeric_value: The numeric value associated with the code (if any)
                - subject_id: ID of the subject
                - prediction_time: The prediction time for this trajectory

        Example:
        >>> metadata_df = pl.DataFrame({
        ...     'code': ['A1', 'A2', 'A3', 'A4', 'A5', 'A6'],
        ...     'code/vocab_index': [1, 2, 3, 4, 5, 6]
        ... })
        >>> batch = TrajectoryBatch(
        ...     time=torch.tensor([[0, .5, 2], [0, 3, 5]]),
        ...     code=torch.tensor([[1, 2, 3], [4, 5, 6]]),
        ...     mask=torch.tensor([[1, 1, 1], [1, 1, 0]]),
        ...     numeric_value=torch.tensor([[0.5, 1.0, 1.5], [2.0, 2.5, 3.0]]),
        ...     numeric_value_mask=torch.tensor([[1, 1, 0], [1, 0, 0]]),
        ...     metadata_df=metadata_df
        ... )
        >>> prediction_times = [datetime(2024, 1, 1), datetime(2024, 1, 1)]
        >>> subject_ids = [1, 2]
        >>> df = batch.to_meds(prediction_times, subject_ids)
        >>> df.sort("subject_id", "code")
        shape: (5, 6)
        ┌────────────┬─────────────────────┬─────────────────────┬──────┬──────────────────┬───────────────┐
        │ subject_id ┆ prediction_time     ┆ time                ┆ code ┆ code/vocab_index ┆ numeric_value │
        │ ---        ┆ ---                 ┆ ---                 ┆ ---  ┆ ---              ┆ ---           │
        │ i32        ┆ datetime[ns]        ┆ datetime[ns]        ┆ str  ┆ i64              ┆ f32           │
        ╞════════════╪═════════════════════╪═════════════════════╪══════╪══════════════════╪═══════════════╡
        │ 1          ┆ 2024-01-01 00:00:00 ┆ 2024-01-01 00:00:00 ┆ A1   ┆ 1                ┆ 0.5           │
        │ 1          ┆ 2024-01-01 00:00:00 ┆ 2024-07-01 14:54:36 ┆ A2   ┆ 2                ┆ 1.0           │
        │ 1          ┆ 2024-01-01 00:00:00 ┆ 2025-12-31 11:38:24 ┆ A3   ┆ 3                ┆ NaN           │
        │ 2          ┆ 2024-01-01 00:00:00 ┆ 2024-01-01 00:00:00 ┆ A4   ┆ 4                ┆ 2.0           │
        │ 2          ┆ 2024-01-01 00:00:00 ┆ 2026-12-31 17:27:36 ┆ A5   ┆ 5                ┆ NaN           │
        └────────────┴─────────────────────┴─────────────────────┴──────┴──────────────────┴───────────────┘
        """
        if len(prediction_time) != len(subject_id) or len(prediction_time) != self.time.shape[0]:
            raise ValueError("Number of prediction times and subject IDs must match batch size")
        schema = {
            "time": pl.Datetime,
            "code": self.metadata_df.schema["code/vocab_index"],
            "numeric_value": pl.Float32,
            "subject_id": pl.Int32,
            "prediction_time": pl.Datetime,
        }

        # Pre-filter masked data using torch operations for efficiency
        batch_indices, seq_indices = torch.where(self.mask)
        if len(batch_indices) == 0:
            return pl.DataFrame(schema=schema)

        # Gather only the valid data points using index tensors
        time_values = self.time[batch_indices, seq_indices]
        code_values = self.code[batch_indices, seq_indices]
        numeric_values = self.numeric_value[batch_indices, seq_indices]
        numeric_value_masks = self.numeric_value_mask[batch_indices, seq_indices]

        # Convert to numpy for faster processing
        time_array = time_values.numpy()
        code_array = code_values.numpy().astype(np.int32)
        numeric_value_array = numeric_values.numpy()
        numeric_value_mask_array = numeric_value_masks.numpy()
        batch_indices = batch_indices.numpy()

        # Create arrays for prediction times and subject IDs
        pred_times = np.array(prediction_time, dtype="datetime64[ns]")[batch_indices]
        subject_ids = np.array(subject_id)[batch_indices].astype(np.int32)

        # Parallel processing using numpy vectorization
        time_deltas = time_array * np.timedelta64(1, self.time_scale).astype("timedelta64[ns]")
        timestamps = pred_times + time_deltas

        # Create the final dictionary with only valid data
        data_dict = {
            "time": timestamps,
            "code": code_array,
            "numeric_value": np.where(
                numeric_value_mask_array, numeric_value_array.astype(np.float32), np.nan
            ),
            "subject_id": subject_ids,
            "prediction_time": pred_times,
        }

        # Create DataFrame directly from the efficient dictionary
        df = pl.from_dict(data_dict, schema=schema)

        # Convert code vocab indexes to strings using the metadata mapping
        df = df.join(
            self.metadata_df.select("code", "code/vocab_index"), left_on="code", right_on="code/vocab_index"
        ).rename({"code_right": "code", "code": "code/vocab_index"})

        return df["subject_id", "prediction_time", "time", "code", "code/vocab_index", "numeric_value"]

to_meds(prediction_time, subject_id)

Convert the trajectory batch to MEDS format.

Parameters:

Name Type Description Default
prediction_time list[datetime]

List of prediction times for each trajectory in the batch

 required
subject_id list[str | int]

List of subject IDs for each trajectory in the batch

 required

Returns:

Type Description
DataFrame

pl.DataFrame: MEDS format DataFrame with columns: - time: Absolute timestamp of the event - code: The medical code string - numeric_value: The numeric value associated with the code (if any) - subject_id: ID of the subject - prediction_time: The prediction time for this trajectory

Example:

metadata_df = pl.DataFrame({ … ‘code’: [‘A1’, ‘A2’, ‘A3’, ‘A4’, ‘A5’, ‘A6’], … ‘code/vocab_index’: [1, 2, 3, 4, 5, 6] … }) batch = TrajectoryBatch( … time=torch.tensor([[0, .5, 2], [0, 3, 5]]), … code=torch.tensor([[1, 2, 3], [4, 5, 6]]), … mask=torch.tensor([[1, 1, 1], [1, 1, 0]]), … numeric_value=torch.tensor([[0.5, 1.0, 1.5], [2.0, 2.5, 3.0]]), … numeric_value_mask=torch.tensor([[1, 1, 0], [1, 0, 0]]), … metadata_df=metadata_df … ) prediction_times = [datetime(2024, 1, 1), datetime(2024, 1, 1)] subject_ids = [1, 2] df = batch.to_meds(prediction_times, subject_ids) df.sort(“subject_id”, “code”) shape: (5, 6) ┌────────────┬─────────────────────┬─────────────────────┬──────┬──────────────────┬───────────────┐ │ subject_id ┆ prediction_time ┆ time ┆ code ┆ code/vocab_index ┆ numeric_value │ │ — ┆ — ┆ — ┆ — ┆ — ┆ — │ │ i32 ┆ datetime[ns] ┆ datetime[ns] ┆ str ┆ i64 ┆ f32 │ ╞════════════╪═════════════════════╪═════════════════════╪══════╪══════════════════╪═══════════════╡ │ 1 ┆ 2024-01-01 00:00:00 ┆ 2024-01-01 00:00:00 ┆ A1 ┆ 1 ┆ 0.5 │ │ 1 ┆ 2024-01-01 00:00:00 ┆ 2024-07-01 14:54:36 ┆ A2 ┆ 2 ┆ 1.0 │ │ 1 ┆ 2024-01-01 00:00:00 ┆ 2025-12-31 11:38:24 ┆ A3 ┆ 3 ┆ NaN │ │ 2 ┆ 2024-01-01 00:00:00 ┆ 2024-01-01 00:00:00 ┆ A4 ┆ 4 ┆ 2.0 │ │ 2 ┆ 2024-01-01 00:00:00 ┆ 2026-12-31 17:27:36 ┆ A5 ┆ 5 ┆ NaN │ └────────────┴─────────────────────┴─────────────────────┴──────┴──────────────────┴───────────────┘

Source code in meds_torch/models/components/utils.py 
 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
def to_meds(self, prediction_time: list[datetime], subject_id: list[str | int]) -> pl.DataFrame:
    """Convert the trajectory batch to MEDS format.

    Args:
        prediction_time: List of prediction times for each trajectory in the batch
        subject_id: List of subject IDs for each trajectory in the batch

    Returns:
        pl.DataFrame: MEDS format DataFrame with columns:
            - time: Absolute timestamp of the event
            - code: The medical code string
            - numeric_value: The numeric value associated with the code (if any)
            - subject_id: ID of the subject
            - prediction_time: The prediction time for this trajectory

    Example:
    >>> metadata_df = pl.DataFrame({
    ...     'code': ['A1', 'A2', 'A3', 'A4', 'A5', 'A6'],
    ...     'code/vocab_index': [1, 2, 3, 4, 5, 6]
    ... })
    >>> batch = TrajectoryBatch(
    ...     time=torch.tensor([[0, .5, 2], [0, 3, 5]]),
    ...     code=torch.tensor([[1, 2, 3], [4, 5, 6]]),
    ...     mask=torch.tensor([[1, 1, 1], [1, 1, 0]]),
    ...     numeric_value=torch.tensor([[0.5, 1.0, 1.5], [2.0, 2.5, 3.0]]),
    ...     numeric_value_mask=torch.tensor([[1, 1, 0], [1, 0, 0]]),
    ...     metadata_df=metadata_df
    ... )
    >>> prediction_times = [datetime(2024, 1, 1), datetime(2024, 1, 1)]
    >>> subject_ids = [1, 2]
    >>> df = batch.to_meds(prediction_times, subject_ids)
    >>> df.sort("subject_id", "code")
    shape: (5, 6)
    ┌────────────┬─────────────────────┬─────────────────────┬──────┬──────────────────┬───────────────┐
    │ subject_id ┆ prediction_time     ┆ time                ┆ code ┆ code/vocab_index ┆ numeric_value │
    │ ---        ┆ ---                 ┆ ---                 ┆ ---  ┆ ---              ┆ ---           │
    │ i32        ┆ datetime[ns]        ┆ datetime[ns]        ┆ str  ┆ i64              ┆ f32           │
    ╞════════════╪═════════════════════╪═════════════════════╪══════╪══════════════════╪═══════════════╡
    │ 1          ┆ 2024-01-01 00:00:00 ┆ 2024-01-01 00:00:00 ┆ A1   ┆ 1                ┆ 0.5           │
    │ 1          ┆ 2024-01-01 00:00:00 ┆ 2024-07-01 14:54:36 ┆ A2   ┆ 2                ┆ 1.0           │
    │ 1          ┆ 2024-01-01 00:00:00 ┆ 2025-12-31 11:38:24 ┆ A3   ┆ 3                ┆ NaN           │
    │ 2          ┆ 2024-01-01 00:00:00 ┆ 2024-01-01 00:00:00 ┆ A4   ┆ 4                ┆ 2.0           │
    │ 2          ┆ 2024-01-01 00:00:00 ┆ 2026-12-31 17:27:36 ┆ A5   ┆ 5                ┆ NaN           │
    └────────────┴─────────────────────┴─────────────────────┴──────┴──────────────────┴───────────────┘
    """
    if len(prediction_time) != len(subject_id) or len(prediction_time) != self.time.shape[0]:
        raise ValueError("Number of prediction times and subject IDs must match batch size")
    schema = {
        "time": pl.Datetime,
        "code": self.metadata_df.schema["code/vocab_index"],
        "numeric_value": pl.Float32,
        "subject_id": pl.Int32,
        "prediction_time": pl.Datetime,
    }

    # Pre-filter masked data using torch operations for efficiency
    batch_indices, seq_indices = torch.where(self.mask)
    if len(batch_indices) == 0:
        return pl.DataFrame(schema=schema)

    # Gather only the valid data points using index tensors
    time_values = self.time[batch_indices, seq_indices]
    code_values = self.code[batch_indices, seq_indices]
    numeric_values = self.numeric_value[batch_indices, seq_indices]
    numeric_value_masks = self.numeric_value_mask[batch_indices, seq_indices]

    # Convert to numpy for faster processing
    time_array = time_values.numpy()
    code_array = code_values.numpy().astype(np.int32)
    numeric_value_array = numeric_values.numpy()
    numeric_value_mask_array = numeric_value_masks.numpy()
    batch_indices = batch_indices.numpy()

    # Create arrays for prediction times and subject IDs
    pred_times = np.array(prediction_time, dtype="datetime64[ns]")[batch_indices]
    subject_ids = np.array(subject_id)[batch_indices].astype(np.int32)

    # Parallel processing using numpy vectorization
    time_deltas = time_array * np.timedelta64(1, self.time_scale).astype("timedelta64[ns]")
    timestamps = pred_times + time_deltas

    # Create the final dictionary with only valid data
    data_dict = {
        "time": timestamps,
        "code": code_array,
        "numeric_value": np.where(
            numeric_value_mask_array, numeric_value_array.astype(np.float32), np.nan
        ),
        "subject_id": subject_ids,
        "prediction_time": pred_times,
    }

    # Create DataFrame directly from the efficient dictionary
    df = pl.from_dict(data_dict, schema=schema)

    # Convert code vocab indexes to strings using the metadata mapping
    df = df.join(
        self.metadata_df.select("code", "code/vocab_index"), left_on="code", right_on="code/vocab_index"
    ).rename({"code_right": "code", "code": "code/vocab_index"})

    return df["subject_id", "prediction_time", "time", "code", "code/vocab_index", "numeric_value"]

get_last_token(output, mask)

Get the last non-masked token from the output tensor.

Args: output (torch.Tensor): The output tensor of shape (batch_size, seq_len, hidden_dim) mask (torch.Tensor): The mask tensor of shape (batch_size, seq_len) where True indicates masked positions

Returns: torch.Tensor: The last non-masked token for each sequence in the batch

Source code in meds_torch/models/components/utils.py 
 9
10
11
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
def get_last_token(output: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """Get the last non-masked token from the output tensor.

    Args: output (torch.Tensor): The output tensor of shape (batch_size, seq_len, hidden_dim) mask
    (torch.Tensor): The mask tensor of shape (batch_size, seq_len) where True indicates masked positions

    Returns: torch.Tensor: The last non-masked token for each sequence in the batch
    """
    # Find the last non-masked position
    last_non_masked = (~mask).float().cumsum(dim=1).argmax(dim=1)

    # Handle cases where all positions are masked
    all_masked = mask.all(dim=1)
    last_non_masked[all_masked] = 0

    if all_masked.any():
        raise ValueError(
            f"{all_masked.sum().item()} sequences have all positions masked. Mask should likely be negated"
        )

    # Create indices for gathering
    batch_size, _, hidden_dim = output.shape
    indices = last_non_masked.view(batch_size, 1, 1).expand(-1, 1, hidden_dim)

    # Gather the last non-masked tokens
    last_token = torch.gather(output, dim=1, index=indices).squeeze(1)

    return last_token

get_time_days_delta(pred_times, input_end_times, device=torch.device('cpu'))

Convert lists of prediction times and generation start times to the continuous number of days (including fractional days) between each pair.

Parameters:

Name Type Description Default
pred_times List[datetime]

List of prediction times

 required
input_end_times List[datetime]

List of generation start times, one for each prediction time

 required
device device

Device to place the output tensor on. Defaults to CPU.

 device('cpu')

Returns:

Type Description
Tensor

torch.Tensor: Tensor of shape (len(pred_times),) containing the number of days (including fractional parts) after each generation start time for each prediction time.

Examples:

>>> from datetime import datetime
>>> import torch
>>> pred_times = [
...     datetime(2022, 1, 1, 12, 0),  # Noon on Jan 1
...     datetime(2022, 1, 2, 0, 0)    # Midnight on Jan 2
... ]
>>> end_times = [
...     datetime(2022, 1, 1, 0, 0),   # Midnight on Jan 1
...     datetime(2022, 1, 1, 0, 0)    # Midnight on Jan 1
... ]
>>> get_time_days_delta(pred_times, end_times)
tensor([0.5000, 1.0000])

>>> pred_times = [
...     datetime(2022, 1, 1, 6, 0),   # 6 AM on Jan 1
...     datetime(2022, 1, 1, 18, 0)   # 6 PM on Jan 1
... ]
>>> end_times = [
...     datetime(2022, 1, 1, 0, 0),   # Midnight on Jan 1
...     datetime(2022, 1, 1, 12, 0)   # Noon on Jan 1
... ]
>>> get_time_days_delta(pred_times, end_times)
tensor([0.2500, 0.2500])
Source code in meds_torch/models/components/utils.py 
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def get_time_days_delta(
    pred_times: list[datetime], input_end_times: list[datetime], device: torch.device = torch.device("cpu")
) -> torch.Tensor:
    """
    Convert lists of prediction times and generation start times to the continuous number
    of days (including fractional days) between each pair.

    Args:
        pred_times (List[datetime]): List of prediction times
        input_end_times (List[datetime]): List of generation start times, one for each prediction time
        device (torch.device, optional): Device to place the output tensor on. Defaults to CPU.

    Returns:
        torch.Tensor: Tensor of shape (len(pred_times),) containing the number of days (including
            fractional parts) after each generation start time for each prediction time.

    Examples:
        >>> from datetime import datetime
        >>> import torch
        >>> pred_times = [
        ...     datetime(2022, 1, 1, 12, 0),  # Noon on Jan 1
        ...     datetime(2022, 1, 2, 0, 0)    # Midnight on Jan 2
        ... ]
        >>> end_times = [
        ...     datetime(2022, 1, 1, 0, 0),   # Midnight on Jan 1
        ...     datetime(2022, 1, 1, 0, 0)    # Midnight on Jan 1
        ... ]
        >>> get_time_days_delta(pred_times, end_times)
        tensor([0.5000, 1.0000])

        >>> pred_times = [
        ...     datetime(2022, 1, 1, 6, 0),   # 6 AM on Jan 1
        ...     datetime(2022, 1, 1, 18, 0)   # 6 PM on Jan 1
        ... ]
        >>> end_times = [
        ...     datetime(2022, 1, 1, 0, 0),   # Midnight on Jan 1
        ...     datetime(2022, 1, 1, 12, 0)   # Noon on Jan 1
        ... ]
        >>> get_time_days_delta(pred_times, end_times)
        tensor([0.2500, 0.2500])
    """
    if len(pred_times) != len(input_end_times):
        raise ValueError(
            f"Length mismatch: pred_times has {len(pred_times)} elements "
            f"but input_end_times has {len(input_end_times)} elements"
        )

    # Convert datetime lists to numpy arrays of timestamps
    pred_timestamps = np.array([dt.timestamp() for dt in pred_times])
    end_timestamps = np.array([dt.timestamp() for dt in input_end_times])

    # Calculate time differences in seconds
    time_deltas_seconds = pred_timestamps - end_timestamps

    # Convert to days (86400 seconds per day)
    days_delta = time_deltas_seconds / 86400

    return torch.tensor(days_delta, device=device, dtype=torch.float32)