Routine ICU data, made computable.

1def aki_urine_output(var, cohort):
2    # The goal is to calculate time weighted hourly urine outputs
3    # e.g. urine output is charted from 05:00 to 05:30 with 50ml and from 05:30 to 07:30 with 200 ml
4    # the output should be: 05:00 50ml, 05:00 50ml, 06:00 100ml, 07:00 50 ml
5    # we can then further group_by and aggregate on the hour
6
7    # Load urine_output, select required columns and rename them
8    urine_output: pl.LazyFrame = var.required_vars["urine_output"].data.lazy()
9    pl_urine_data = urine_output.select(
10        cohort.primary_key, "recordtime", "value"
11    ).rename({"recordtime": "end_time"})
12
13    # Sort and calculate start_time and duration
14    pl_urine_data = pl_urine_data.sort([cohort.primary_key, "end_time"])
15
16    pl_urine_data = (
17        pl_urine_data.with_columns(
18            pl.col("end_time").shift(1).over(cohort.primary_key).alias("start_time")
19        )
20        .with_columns(
21            ((pl.col("end_time") - pl.col("start_time")).dt.total_seconds()).alias(
22                "collection_duration_s"
23            )
24        )
25        .with_columns(
26            (pl.col("collection_duration_s") / 3600).alias("collection_duration_h")
27        )
28        .with_columns(
29            pl.datetime_ranges(
30                pl.col("start_time").dt.truncate("1h"),
31                "end_time",
32                "1h",
33            ).alias("start_hour")
34        )
35        .filter(pl.col("collection_duration_s") > 0)
36        .with_columns(pl.col("collection_duration_h").gt(24).alias("gt_24h"))
37    )
38
39    # Filter out first rows where start_time is null
40    pl_redistributable = pl_urine_data.filter(pl.col("start_time").is_not_null())
41
42    pl_hourly_urine = (
43        pl_redistributable.explode("start_hour")
44        .with_columns(
45            pl.col("start_hour").dt.offset_by("1h").alias("end_hour"),
46            pl.max_horizontal("start_time", "start_hour").alias("start_overlap"),
47        )
48        .with_columns(pl.min_horizontal("end_time", "end_hour").alias("end_overlap"))
49        .with_columns(
50            (
51                (pl.col("end_overlap") - pl.col("start_overlap")).dt.total_seconds()
52            ).alias("overlap_duration_s")
53        )
54        .with_columns(
55            (
56                pl.col("value")
57                * pl.col("overlap_duration_s")
58                / pl.col("collection_duration_s")
59            ).alias("urine_share_ml")
60        )

1def aki_crea_based(var: VariableContext, cohort: Cohort) -> pl.DataFrame:
2    # Extract creatinine data and sort
3    crea_data_all = var.required_vars["blood_creatinine_all"].data.lazy()
4
5    # Select case admission and discharge timestamp from cohort
6    obs_with_times = utils.add_time_window_expr(cohort.obs, time_window=var.time_window)
7    case_bounds: pl.LazyFrame = obs_with_times.lazy().select(
8        "patient_id",
9        *set(["case_id", cohort.primary_key]),
10        "hospital_admission",
11        "case_tmin",
12        "hospital_discharge",
13        "case_tmax",
14    )
15
16    # Merge admission info into creatinine data on hospital level
17    # Also compute date bounds
18    # NOTE: Change for experimental join_where if to slow
19    crea_merged = (
20        crea_data_all.drop("case_id")
21        .join(
22            case_bounds.select(
23                "patient_id", "case_id", "hospital_admission", "hospital_discharge"
24            ),
25            on="patient_id",
26        )
27        .drop("patient_id")
28        .with_columns(
29            (pl.col("hospital_admission") - pl.duration(days=7)).alias(
30                "admission_sub_7d"
31            ),
32            (pl.col("hospital_admission") - pl.duration(days=365)).alias(
33                "admission_sub_365d"
34            ),
35        )
36    )
37
38    # 1. admission crea
39    crea_admission = (
40        crea_merged.filter(pl.col("recordtime") >= pl.col("hospital_admission"))
41        .sort("case_id", "recordtime")
42        .group_by("case_id", "hospital_admission", "hospital_discharge")
43        .agg(
44            pl.col("value").first().alias("crea_admission"),
45            # pl.col("recordtime").first().alias("crea_admission_time"),
46        )
47    )
48
49    # 2. Filter for 7–365 days before admission and extract
50    # 2a. Most recent creatinine 7–365d pre-admission (unused)
51    # 2b. Mean creatinine 7–365d pre-admission (used as baseline if not null)
52    crea_recent_mean_pre = (
53        crea_merged.filter(
54            pl.col("recordtime").is_between(
55                "admission_sub_365d", "admission_sub_7d", closed="left"
56            )
57        )
58        .sort("case_id", "recordtime")
59        .group_by("case_id")
60        .agg(

1def blood_pressure_mean(var, cohort):
2    dfs: list[pl.DataFrame] = []
3    for name, var in var.required_vars.items():
4        df: pl.DataFrame = var.data
5        dfs.append(
6            df.select(
7                cohort.primary_key,
8                "recordtime",
9                "value",
10                "value_unit",
11                pl.struct(pl.lit(name).alias("source")).alias("attributes"),
12            )
13        )
14
15    df = pl.concat(dfs)
16
17    df = df.sort("recordtime")
18
19    return df

1def pf_ratio(
2    var: VariableContext,
3    cohort: Cohort,
4    plusminus: pd.Timedelta = pd.Timedelta(hours=24),
5    strategy: AsofJoinStrategy = "nearest",
6) -> pd.DataFrame:
7    """
8    Calculate the PF ratio by matching each PaO2 value to the closest FiO2 value within a specified time window,
9    and filter PF ratios to keep only values between 10 and 1000.
10
11    Args:
12        var: VariableContext container with required data for PaO2 and FiO2.
13        cohort: Cohort object containing the primary key.
14        plusminus (pd.Timedelta): Optional time range within which to find the closest FiO2 for each PaO2.
15                                  Defaults to 24 hours.
16        strategy (str): Strategy for asof join - "nearest", "backward", or "forward". Defaults to "nearest".
17
18    Returns:
19        pd.DataFrame: DataFrame with the primary key, PaO2, closest FiO2, and calculated, cleaned PF ratio.
20    """
21
22    merged_df = utils.attach_asof_value(
23        var.required_vars["blood_pao2_arterial"].data,
24        var.required_vars["vent_fio2"].data,
25        name="fio2",
26        by=cohort.primary_key,
27        strategy=strategy,
28        tolerance=plusminus,
29    )
30
31    # Calculate PF ratio and filter
32    merged_df = (
33        merged_df.with_columns(pl.col("value").truediv(pl.col("fio2")))
34        .filter(pl.col("value").is_between(10, 1000))
35        .select(cohort.primary_key, "recordtime", "value")
36    )
37
38    return merged_df

1def blood_pao2_arterial(var, cohort):
2    df = var.required_vars["blood_gas"].data
3
4    out = (
5        df.filter(pl.col("attributes").struct.field("abnahmeort") == "Arterial")
6        .with_columns(
7            [
8                pl.col("attributes")
9                .struct.field("po2")
10                .cast(pl.Float64, strict=False)
11                .alias("value")
12            ]
13        )
14        .select([cohort.primary_key, "recordtime", "value"])
15    )
16    return out

1def vent_fio2(var, cohort):
2    dfs: list[pl.DataFrame] = []
3    for name, var in var.required_vars.items():
4        df: pl.DataFrame = var.data
5        dfs.append(
6            df.select(
7                cohort.primary_key,
8                "recordtime",
9                "value",
10                "value_unit",
11                pl.struct(pl.lit(name).alias("source")).alias("attributes"),
12            )
13        )
14
15    fio2 = pl.concat(dfs)
16
17    fio2 = fio2.with_columns(
18        pl.when(pl.col("value") > 1)
19        .then(pl.col("value") / 100)
20        .otherwise(pl.col("value"))
21        .alias("value")
22    )
23    return fio2

1def norepinephrine(var, cohort):
2    """Process norepinephrine medication data at ICU stay level."""
3    original_fields = var.data["attributes"].struct.fields
4    keep_fields = [
5        field
6        for field in original_fields
7        if field not in ["drugname", "genericname", "drugid"]
8    ]
9    var.data = var.data.unnest("attributes")
10
11    # Cast numeric columns
12    var.data = var.data.with_columns(pl.col("value").cast(pl.Float64))
13
14    # Track if original unit was µg
15    var.data = var.data.with_columns((pl.col("value_unit") == "µg").alias("was_ug"))
16
17    # Convert µg → mg for value AND concentration
18    if "concentration" in var.data.columns:
19        var.data = var.data.with_columns(pl.col("concentration").cast(pl.Float64))
20        var.data = var.data.with_columns(
21            pl.when(pl.col("was_ug"))
22            .then(pl.col("concentration") / 1000)
23            .otherwise(pl.col("concentration"))
24            .alias("concentration")
25        )
26
27    # Now convert value and unit
28    var.data = var.data.with_columns(
29        [
30            pl.when(pl.col("was_ug"))
31            .then(pl.col("value") / 1000)
32            .otherwise(pl.col("value"))
33            .alias("value"),
34            pl.when(pl.col("was_ug"))
35            .then(pl.lit("mg"))
36            .otherwise(pl.col("value_unit"))
37            .alias("value_unit"),
38        ]
39    )
40
41    # Filter concentration ≤21 mg/ml and value ≠ 0
42    if "concentration" in var.data.columns:
43        var.data = var.data.filter(
44            (pl.col("concentration") <= 21) & (pl.col("value") != 0)
45        )
46    else:
47        var.data = var.data.filter(pl.col("value") != 0)
48
49    # Clean up temporary column
50    var.data = var.data.drop("was_ug")
51
52    # Flag bolus applications
53    var.data = var.data.with_columns(
54        (pl.col("recordtime") == pl.col("recordtime_end")).alias("bolus")
55    )
56
57    # Add uniform drug name
58    var.data = var.data.with_columns(pl.lit("norepinephrine").alias("drug_name"))
59
60    # Create attributes struct

1def endotracheal_tube(var, cohort):
2    df = var.required_vars["airway_device"].data
3    tube_types = [
4        "Endotracheal tube",
5        "Laryngektomietubus",
6        "Microcuff Endotracheal tube",
7    ]
8
9    out = df.filter(pl.col("value").is_in(tube_types)).select(
10        [cohort.primary_key, "recordtime", "recordtime_end", "value", "attributes"]
11    )
12    return out

1def fluid_balance(var, cohort):
2    # Unnest attributes
3    data = var.data.unnest("attributes")
4    # Ensure numeric columns are floats
5    data = data.with_columns(
6        [
7            pl.col("input").cast(pl.Float64),
8            pl.col("output").cast(pl.Float64),
9            pl.col("value").cast(pl.Float64),
10        ]
11    )
12
13    data = data.filter((pl.col("output") >= -5000) & (pl.col("input") <= 10000))
14
15    # Add record_date (date part only)
16    data = data.with_columns(pl.col("recordtime").dt.date().alias("record_date"))
17
18    # Group by patient + date and sum numeric columns
19    daily_sum = data.group_by(["case_id", "record_date"], maintain_order=True).agg(
20        [
21            pl.col("input").sum().alias("input"),
22            pl.col("output").sum().alias("output"),
23            pl.col("value").sum().alias("value"),
24        ]
25    )
26
27    # Optional: filter extreme daily sums
28    daily_sum = daily_sum.filter(
29        (pl.col("value") >= -10000) & (pl.col("value") <= 20000)
30    )
31
32    # Rename column
33    daily_sum = (
34        daily_sum.rename(
35            {
36                "record_date": "recordtime",
37            }
38        )
39        .with_columns(
40            [
41                pl.lit("ml").alias("value_unit"),
42                pl.struct(["input", "output"]).alias("attributes"),
43            ]
44        )
45        .drop(["input", "output"])
46    )
47
48    # Cast to datetime with time at 00:00:00
49    daily_sum = daily_sum.with_columns(
50        pl.col("recordtime").cast(pl.Datetime("ms")).dt.truncate("1d")
51    )
52
53    return daily_sum

1def nu_desc(var, cohort):
2    return var.data.with_columns(
3        pl.col("value")
4        .cast(pl.Utf8)
5        .str.replace_all(r"\D", "")
6        .replace("", None)
7        .cast(pl.Int64)
8    ).drop_nulls("value")