patient_data_processing.py

import os
from os.path import dirname, join
from typing import Dict

import numpy as np
import pandas as pd
import yaml
from pandas import DataFrame
from pandas.core.groupby.generic import DataFrameGroupBy

from feature_selection import fill_age_height_weight

RECORDS_DIR = "records"
META_DIR = "records/meta"
ECG_DIR = "records/ecg"

KEY_ORDER = [
    "patient_id",
    "diagnosis",
    "ritmi",
    "age",
    "sex",
    "height",
    "weight",
    "ecg_id",
    "device",
    "recording_date",
    "report",
    "pacemaker",
    "nurse",
    "site",
    "scp_codes",
    "heart_axis",
    "infarction_stadium1",
    "infarction_stadium2",
    "validated_by",
    "second_opinion",
    "initial_autogenerated_report",
    "validated_by_human",
    "baseline_drift",
    "static_noise",
    "burst_noise",
    "electrodes_problems",
    "extra_beats",
    "strat_fold",
]


def open_patient_csv(csv_fn: str) -> DataFrame:
    """
    Reads a CSV file containing patient information and converts it into a pandas DataFrame.

    :param csv_fn: The filename of the CSV file to read.
    :type csv_fn: str
    :return: A DataFrame containing the patient information.
    :rtype: pd.DataFrame
    """
    # Load data from CSV file
    df = pd.read_csv(csv_fn, sep=";")

    # Replace sex codes with human-readable strings
    df["sex"].replace({0: "male", 1: "female", None: "unknown"}, inplace=True)

    # Replace pacemaker codes with boolean values
    pacemaker_map = {"ja": True, np.nan: False}
    df["pacemaker"] = df["pacemaker"].str.contains("ja").replace(pacemaker_map)

    # Replace diagnosis codes with human-readable strings
    diagnosi_map = {
        "SR": "Sinus Rhythm",
        "AFIB": "Atrial Fibrillation",
        "STACH": "Sinus Tachycardia",
        "SARRH": "Sinus Arrhythmia",
        "SBRAD": "Sinus Bradycardia",
        "PACE": "Normal Functioning Artificial Pacemaker",
        "SVARR": "Supraventricular Arrhythmia",
        "BIGU": "Bigeminal Pattern (Unknown Origin, SV or Ventricular)",
        "AFLT": "Atrial Flutter",
        "SVTAC": "Supraventricular Tachycardia",
        "PSVT": "Paroxysmal Supraventricular Tachycardia",
        "TRIGU": "Trigeminal Pattern (Unknown Origin, SV or Ventricular)",
    }
    df["diagnosi"] = df["diagnosi"].replace(diagnosi_map)
    df.rename(columns={"diagnosi": "diagnosis"}, inplace=True)

    # fill missing age, height, and weight
    df = fill_age_height_weight(df)

    # Replace NaN values with None
    df = df.replace(np.nan, None)

    return df


def patient_df_to_meta_files(df: DataFrame) -> None:
    """
    Writes metadata files in YAML format for each row in the given DataFrame.

    :param df: The DataFrame containing patient information.
    :type df: pd.DataFrame
    :return: None
    """
    if len(df):
        os.makedirs(META_DIR, exist_ok=True)
    print("writing meta files...")
    for i, row in df.iterrows():
        row.patient_id = f"patient_{i + 1}"
        data = {k: row[k] for k in KEY_ORDER}
        fn = join(META_DIR, f"patient_{i + 1}.yaml")
        with open(fn, "w", encoding="utf-8") as fp:
            yaml.dump(data, fp, sort_keys=False)
    print("finished")


def group_by_rhythm(df: DataFrame) -> DataFrameGroupBy:
    """
    Groups the given DataFrame by the "diagnosi" column and returns a DataFrameGroupBy object.

    :param df: The DataFrame containing patient information.
    :type df: pd.DataFrame
    :return: A DataFrameGroupBy object.
    :rtype: pd.core.groupby.generic.DataFrameGroupBy
    """
    grps = df.groupby("diagnosis")
    return grps


def count_rhythms(df: DataFrame) -> Dict:
    """
    Counts the number of patients with each diagnosis in the given DataFrame and returns a dictionary.

    :param df: The DataFrame containing patient information.
    :type df: pd.DataFrame
    :return: A dictionary containing the counts for each diagnosis.
    :rtype: Dict
    """
    rhythms = {}
    for rhythm, grp in group_by_rhythm(df):
        rhythms[rhythm] = len(grp)
    return rhythms


def write_rhythm_counts_to_file(df: DataFrame, fn: str) -> None:
    """
    Writes the counts of patients with each diagnosis to a YAML file.

    :param df: The DataFrame containing patient information.
    :type df: pd.DataFrame
    :param fn: The filename to write the YAML file to.
    :type fn: str
    :return: None
    """
    rhythms = count_rhythms(df)
    if rhythms:
        os.makedirs(dirname(fn), exist_ok=True)
    with open(fn, "w") as fp:
        yaml.dump(rhythms, fp)


def ecg_to_csvs(fn: str) -> None:
    """
    Reads an NPY file containing ECG data and writes CSV files for each patient's ECG data.

    :param fn: The filename of the NPY file containing ECG data.
    :type fn: str
    :return: None
    """
    columns = [
        "I",
        "II",
        "III",
        "aVF",
        "aVR",
        "aVL",
        "V1",
        "V2",
        "V3",
        "V4",
        "V5",
        "V6",
    ]
    print(f'loading ecg from "{fn}"...')
    all_ecg = np.load(fn)
    if all_ecg.any():
        os.makedirs(ECG_DIR, exist_ok=True)
    print("writing ecg files...")
    n_ecg = all_ecg.shape[0]
    for i in range(n_ecg):
        fn = join(ECG_DIR, f"patient_{i + 1}.csv")
        ecg = all_ecg[i, :, :]
        np.savetxt(fn, ecg, delimiter=",", header=",".join(columns), comments="")
    print("finished")


if __name__ == "__main__":
    # np_fn = "ecgeq-500hzsrfava.npy"
    # ecg_to_csvs(np_fn)

    csv_fn = "coorteeqsrafva_en.csv"
    df = open_patient_csv(csv_fn)
    write_rhythm_counts_to_file(df, join(RECORDS_DIR, "!rhythm_count.yaml"))
    patient_df_to_meta_files(df)