Skip to content

generate_analysis_schema

This file defines the output schema for generating trajectories using zero-shot models.

Example Usage:

At the end of generate_trajectory script:

validated_df = validate_generated_data(df)

generation_analysis_schema()

Generate the schema for trajectory analysis data with consistent field ordering.

Returns:

Type Description

pyarrow.Schema: The schema for trajectory analysis
Source code in meds_torch/schemas/generate_analysis_schema.py 
26
27
28
29
30
31
32
33
def generation_analysis_schema():
    """
    Generate the schema for trajectory analysis data with consistent field ordering.

    Returns:
        pyarrow.Schema: The schema for trajectory analysis
    """
    return pa.schema([(name, dtype) for name, dtype in MEDS_TRAJECTORY_SCHEMA.items()])

validate_generated_data(df)

Validate generated data against the schema and ensure correct ordering.

Parameters:

Name Type Description Default
df

polars.DataFrame to validate

 required

Returns:

Type Description

pyarrow.Table: Validated data cast to the correct schema with ordered columns

Example:

import polars as pl from datetime import datetime df = pl.DataFrame({ … ‘time’: [datetime(2024, 7, 18, 16, 21, 41)] * 3, … ‘code’: [‘A1’, ‘A2’, ‘A3’], … ‘code/vocab_index’: [55, 59, 61], … ‘numeric_value’: [0.625, 0.625, float(‘nan’)], … ‘subject_id’: [109767, 109767, 109767], … ‘prediction_time’: [datetime(2024, 7, 18, 16, 21, 41)] * 3, … ‘TRAJECTORY_TYPE’: [‘INPUT_DATA’] * 3 … }) validate_generated_data(df) pyarrow.Table subject_id: int64 prediction_time: timestamp[ns] time: timestamp[ns] code: string code/vocab_index: int64 numeric_value: double TRAJECTORY_TYPE: string

subject_id: [[109767,109767,109767]] prediction_time: [[2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.00…]] time: [[2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.000000000]] code: [[“A1”,”A2”,”A3”]] code/vocab_index: [[55,59,61]] numeric_value: [[0.625,0.625,nan]] TRAJECTORY_TYPE: [[“INPUT_DATA”,”INPUT_DATA”,”INPUT_DATA”]]

Source code in meds_torch/schemas/generate_analysis_schema.py 
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
def validate_generated_data(df):
    """
    Validate generated data against the schema and ensure correct ordering.

    Args:
        df: polars.DataFrame to validate

    Returns:
        pyarrow.Table: Validated data cast to the correct schema with ordered columns

    Example:
    >>> import polars as pl
    >>> from datetime import datetime
    >>> df = pl.DataFrame({
    ...     'time': [datetime(2024, 7, 18, 16, 21, 41)] * 3,
    ...     'code': ['A1', 'A2', 'A3'],
    ...     'code/vocab_index': [55, 59, 61],
    ...     'numeric_value': [0.625, 0.625, float('nan')],
    ...     'subject_id': [109767, 109767, 109767],
    ...     'prediction_time': [datetime(2024, 7, 18, 16, 21, 41)] * 3,
    ...     'TRAJECTORY_TYPE': ['INPUT_DATA'] * 3
    ... })
    >>> validate_generated_data(df)
    pyarrow.Table
    subject_id: int64
    prediction_time: timestamp[ns]
    time: timestamp[ns]
    code: string
    code/vocab_index: int64
    numeric_value: double
    TRAJECTORY_TYPE: string
    ----
    subject_id: [[109767,109767,109767]]
    prediction_time: [[2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.00...]]
    time: [[2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.000000000,2024-07-18 16:21:41.000000000]]
    code: [["A1","A2","A3"]]
    code/vocab_index: [[55,59,61]]
    numeric_value: [[0.625,0.625,nan]]
    TRAJECTORY_TYPE: [["INPUT_DATA","INPUT_DATA","INPUT_DATA"]]
    """
    # Get the validated schema
    validated_schema = generation_analysis_schema()

    # Convert to arrow
    arrow_table = df.to_arrow()

    # Ensure columns are in the correct order
    expected_columns = validated_schema.names
    arrow_table = arrow_table.select(expected_columns)

    # Cast to the validated schema (this will handle struct field ordering)
    return arrow_table.cast(validated_schema)