cehrbert-data 0.0.1__py3-none-any.whl

cehrbert_data/apps/generate_concept_similarity_table.py

@@ -0,0 +1,423 @@
+"""This module provides functionality to extract patient event data from domain tables,.
+
+compute information content and semantic similarity for concepts, and calculate concept
+similarity scores.
+
+Functions: extract_data: Extract data from specified domain tables. compute_information_content:
+Compute the information content for concepts based on frequency.
+compute_information_content_similarity: Compute the similarity between concepts based on
+information content. compute_semantic_similarity: Compute the semantic similarity between concept
+pairs. main: Main function to orchestrate the extraction, processing, and saving of concept
+similarity data.
+"""
+
+import datetime
+import logging
+import os
+from typing import List
+
+from pyspark.sql import DataFrame, SparkSession
+from pyspark.sql import functions as F
+
+from cehrbert_data.config.output_names import CONCEPT_SIMILARITY_PATH, QUALIFIED_CONCEPT_LIST_PATH
+from cehrbert_data.const.common import CONCEPT, CONCEPT_ANCESTOR
+from cehrbert_data.utils.spark_utils import join_domain_tables, preprocess_domain_table, validate_table_names
+
+
+def extract_data(spark: SparkSession, input_folder: str, domain_table_list: List[str]):
+    """
+    Extract patient event data from the specified domain tables.
+
+    Args:
+        spark (SparkSession): The Spark session to use for processing.
+        input_folder (str): Path to the input folder containing domain tables.
+        domain_table_list (List[str]): List of domain table names to extract data from.
+
+    Returns:
+        DataFrame: A DataFrame containing extracted and processed patient event data.
+    """
+    domain_tables = []
+    for domain_table_name in domain_table_list:
+        domain_tables.append(preprocess_domain_table(spark, input_folder, domain_table_name))
+    patient_event = join_domain_tables(domain_tables)
+    # Remove all concept_id records
+    patient_event = patient_event.where("standard_concept_id <> 0")
+
+    return patient_event
+
+
+def compute_information_content(patient_event: DataFrame, concept_ancestor: DataFrame):
+    """
+    Calculate the information content using the frequency of each concept and the graph.
+
+    :param patient_event:
+    :param concept_ancestor:
+    :return:
+    """
+    # Get the total count
+    total_count = patient_event.distinct().count()
+    # Count the frequency of each concept
+    concept_frequency = patient_event.distinct().groupBy("standard_concept_id").count()
+    # left join b/w descendent_concept_id and the standard_concept_id in the concept freq table
+    freq_df = (
+        concept_frequency.join(
+            concept_ancestor,
+            F.col("descendant_concept_id") == F.col("standard_concept_id"),
+        )
+        .groupBy("ancestor_concept_id")
+        .sum("count")
+        .withColumnRenamed("ancestor_concept_id", "concept_id")
+        .withColumnRenamed("sum(count)", "count")
+    )
+    # Calculate information content for each concept
+    information_content = freq_df.withColumn("information_content", (-F.log(F.col("count") / total_count))).withColumn(
+        "probability", F.col("count") / total_count
+    )
+
+    return information_content
+
+
+def compute_information_content_similarity(
+    concept_pair: DataFrame, information_content: DataFrame, concept_ancestor: DataFrame
+):
+    """
+    Compute the similarity between concept pairs based on their information content.
+
+    Args:
+        concept_pair (DataFrame): A DataFrame containing pairs of concepts.
+        information_content (DataFrame): A DataFrame with information content for concepts.
+        concept_ancestor (DataFrame): A DataFrame containing concept ancestor relationships.
+
+    Returns:
+        DataFrame: A DataFrame containing various similarity measures for concept pairs.
+    """
+    # Extract the pairs of concepts from the training data and join to the information content table
+    information_content_concept_pair = (
+        concept_pair.select("concept_id_1", "concept_id_2")
+        .join(
+            information_content,
+            F.col("concept_id_1") == F.col("concept_id"),
+            "left_outer",
+        )
+        .select(
+            F.col("concept_id_1"),
+            F.col("concept_id_2"),
+            F.col("information_content").alias("information_content_1"),
+        )
+        .join(
+            information_content,
+            F.col("concept_id_2") == F.col("concept_id"),
+            "left_outer",
+        )
+        .select(
+            F.col("concept_id_1"),
+            F.col("concept_id_2"),
+            F.col("information_content_1"),
+            F.col("information_content").alias("information_content_2"),
+        )
+    )
+
+    # Join to get all the ancestors of concept_id_1
+    concept_id_1_ancestor = information_content_concept_pair.join(
+        concept_ancestor, F.col("concept_id_1") == F.col("descendant_concept_id")
+    ).select("concept_id_1", "concept_id_2", "ancestor_concept_id")
+
+    # Join to get all the ancestors of concept_id_2
+    concept_id_2_ancestor = concept_pair.join(
+        concept_ancestor, F.col("concept_id_2") == F.col("descendant_concept_id")
+    ).select("concept_id_1", "concept_id_2", "ancestor_concept_id")
+
+    # Compute the summed information content of all ancestors of concept_id_1 and concept_id_2
+    union_sum = (
+        concept_id_1_ancestor.union(concept_id_2_ancestor)
+        .distinct()
+        .join(information_content, F.col("ancestor_concept_id") == F.col("concept_id"))
+        .groupBy("concept_id_1", "concept_id_2")
+        .agg(F.sum("information_content").alias("ancestor_union_ic"))
+    )
+
+    # Compute the summed information content of common ancestors of concept_id_1 and concept_id_2
+    intersection_sum = (
+        concept_id_1_ancestor.intersect(concept_id_2_ancestor)
+        .join(information_content, F.col("ancestor_concept_id") == F.col("concept_id"))
+        .groupBy("concept_id_1", "concept_id_2")
+        .agg(F.sum("information_content").alias("ancestor_intersection_ic"))
+    )
+
+    # Compute the information content and probability of the most informative common ancestor (MICA)
+    mica_ancestor = (
+        concept_id_1_ancestor.intersect(concept_id_2_ancestor)
+        .join(information_content, F.col("ancestor_concept_id") == F.col("concept_id"))
+        .groupBy("concept_id_1", "concept_id_2")
+        .agg(
+            F.max("information_content").alias("mica_information_content"),
+            F.max("probability").alias("mica_probability"),
+        )
+    )
+
+    # Join the MICA to pairs of concepts
+    features = information_content_concept_pair.join(
+        mica_ancestor,
+        (information_content_concept_pair["concept_id_1"] == mica_ancestor["concept_id_1"])
+        & (information_content_concept_pair["concept_id_2"] == mica_ancestor["concept_id_2"]),
+        "left_outer",
+    ).select(
+        [information_content_concept_pair[f] for f in information_content_concept_pair.schema.fieldNames()]
+        + [F.col("mica_information_content"), F.col("mica_probability")]
+    )
+
+    # Compute the lin measure
+    features = features.withColumn(
+        "lin_measure",
+        2 * F.col("mica_information_content") / (F.col("information_content_1") * F.col("information_content_2")),
+    )
+
+    # Compute the jiang measure
+    features = features.withColumn(
+        "jiang_measure",
+        1 - (F.col("information_content_1") + F.col("information_content_2") - 2 * F.col("mica_information_content")),
+    )
+
+    # Compute the information coefficient
+    features = features.withColumn(
+        "information_coefficient",
+        F.col("lin_measure") * (1 - 1 / (1 + F.col("mica_information_content"))),
+    )
+
+    # Compute the relevance_measure
+    features = features.withColumn("relevance_measure", F.col("lin_measure") * (1 - F.col("mica_probability")))
+
+    # Join to get the summed information content of the common ancestors of concept_id_1 and
+    # concept_id_2
+    features = features.join(
+        intersection_sum,
+        (features["concept_id_1"] == intersection_sum["concept_id_1"])
+        & (features["concept_id_2"] == intersection_sum["concept_id_2"]),
+        "left_outer",
+    ).select([features[f] for f in features.schema.fieldNames()] + [F.col("ancestor_intersection_ic")])
+
+    # Join to get the summed information content of the common ancestors of concept_id_1 and
+    # concept_id_2
+    features = features.join(
+        union_sum,
+        (features["concept_id_1"] == union_sum["concept_id_1"])
+        & (features["concept_id_2"] == union_sum["concept_id_2"]),
+        "left_outer",
+    ).select([features[f] for f in features.schema.fieldNames()] + [F.col("ancestor_union_ic")])
+
+    # Compute the graph information content measure
+    features = features.withColumn(
+        "graph_ic_measure",
+        F.col("ancestor_intersection_ic") / F.col("ancestor_union_ic"),
+    )
+
+    return features.select(
+        [
+            F.col("concept_id_1"),
+            F.col("concept_id_2"),
+            F.col("mica_information_content"),
+            F.col("lin_measure"),
+            F.col("jiang_measure"),
+            F.col("information_coefficient"),
+            F.col("relevance_measure"),
+            F.col("graph_ic_measure"),
+        ]
+    )
+
+
+def compute_semantic_similarity(spark, patient_event, concept, concept_ancestor):
+    required_concept = (
+        patient_event.distinct()
+        .select("standard_concept_id")
+        .join(concept, F.col("standard_concept_id") == F.col("concept_id"))
+        .select("standard_concept_id", "domain_id")
+    )
+
+    concept_ancestor.createOrReplaceTempView("concept_ancestor")
+    required_concept.createOrReplaceTempView("required_concept")
+
+    concept_pair = spark.sql(
+        """
+        WITH concept_pair AS (
+            SELECT
+                c1.standard_concept_id AS concept_id_1,
+                c2.standard_concept_id AS concept_id_2,
+                c1.domain_id
+            FROM required_concept AS c1
+            JOIN required_concept AS c2
+                ON c1.domain_id = c2.domain_id
+            WHERE c1.standard_concept_id <> c2.standard_concept_id
+        )
+        SELECT
+            cp.concept_id_1,
+            cp.concept_id_2,
+            ca_1.ancestor_concept_id AS common_ancestor_concept_id,
+            ca_1.min_levels_of_separation AS distance_1,
+            ca_2.min_levels_of_separation AS distance_2
+        FROM concept_pair AS cp
+        JOIN concept_ancestor AS ca_1
+            ON cp.concept_id_1 = ca_1.descendant_concept_id
+        JOIN concept_ancestor AS ca_2
+            ON cp.concept_id_2 = ca_2.descendant_concept_id
+        WHERE ca_1.ancestor_concept_id = ca_2.ancestor_concept_id
+    """
+    )
+
+    # Find the root concepts
+    root_concept = (
+        concept_ancestor.groupBy("descendant_concept_id")
+        .count()
+        .where("count = 1")
+        .withColumnRenamed("descendant_concept_id", "root_concept_id")
+    )
+    # Retrieve all ancestor descendant relationships for the root concepts
+    root_concept_relationship = (
+        root_concept.join(
+            concept_ancestor,
+            root_concept["root_concept_id"] == concept_ancestor["ancestor_concept_id"],
+        )
+        .select(
+            concept_ancestor["ancestor_concept_id"],
+            concept_ancestor["descendant_concept_id"],
+            concept_ancestor["max_levels_of_separation"].alias("root_distance"),
+        )
+        .where("ancestor_concept_id <> descendant_concept_id")
+    )
+
+    # Join to get all root concepts and their corresponding root_distance
+    concept_pair = concept_pair.join(
+        root_concept_relationship,
+        F.col("common_ancestor_concept_id") == F.col("descendant_concept_id"),
+    ).select("concept_id_1", "concept_id_2", "distance_1", "distance_2", "root_distance")
+
+    # Compute the semantic similarity
+    concept_pair_similarity = concept_pair.withColumn(
+        "semantic_similarity",
+        2 * F.col("root_distance") / (2 * F.col("root_distance") + F.col("distance_1") + F.col("distance_2")),
+    )
+    # Find the maximum semantic similarity
+    concept_pair_similarity = concept_pair_similarity.groupBy("concept_id_1", "concept_id_2").agg(
+        F.max("semantic_similarity").alias("semantic_similarity")
+    )
+
+    return concept_pair_similarity
+
+
+def main(
+    input_folder: str,
+    output_folder: str,
+    domain_table_list: List[str],
+    date_filter: str,
+    include_concept_list: bool,
+):
+    """
+    Main function to generate the concept similarity table.
+
+    Args:
+        input_folder (str): The path to the input folder containing raw data.
+        output_folder (str): The path to the output folder to store the results.
+        domain_table_list (List[str]): List of domain tables to process.
+        date_filter (str): Date filter to apply to the data.
+        include_concept_list (bool): Whether to include a filtered concept list.
+    """
+
+    spark = SparkSession.builder.appName("Generate the concept similarity table").getOrCreate()
+
+    logger = logging.getLogger(__name__)
+    logger.info(
+        "input_folder: %s\noutput_folder: %s\ndomain_table_list: %s\ndate_filter: " "%s\ninclude_concept_list: %s",
+        input_folder,
+        output_folder,
+        domain_table_list,
+        date_filter,
+        include_concept_list,
+    )
+
+    concept = preprocess_domain_table(spark, input_folder, CONCEPT)
+    concept_ancestor = preprocess_domain_table(spark, input_folder, CONCEPT_ANCESTOR)
+
+    # Extract all data points from specified domains
+    patient_event = extract_data(spark, input_folder, domain_table_list)
+
+    # Calculate information content using unfiltered the patient event dataframe
+    information_content = compute_information_content(patient_event, concept_ancestor)
+
+    # Filter out concepts that are not required in the required concept_list
+    if include_concept_list and patient_event:
+        # Filter out concepts
+        qualified_concepts = F.broadcast(preprocess_domain_table(spark, input_folder, QUALIFIED_CONCEPT_LIST_PATH))
+
+        patient_event = patient_event.join(qualified_concepts, "standard_concept_id").select("standard_concept_id")
+
+    concept_pair_similarity = compute_semantic_similarity(spark, patient_event, concept, concept_ancestor)
+
+    # Compute the information content based similarity scores
+    concept_pair_ic_similarity = compute_information_content_similarity(
+        concept_pair_similarity, information_content, concept_ancestor
+    )
+
+    concept_pair_similarity_columns = [concept_pair_similarity[f] for f in concept_pair_similarity.schema.fieldNames()]
+    concept_pair_ic_similarity_columns = [
+        f for f in concept_pair_ic_similarity.schema.fieldNames() if "concept_id" not in f
+    ]
+
+    # Join two dataframes to get the final result
+    concept_pair_similarity = concept_pair_similarity.join(
+        concept_pair_ic_similarity,
+        (concept_pair_similarity["concept_id_1"] == concept_pair_ic_similarity["concept_id_1"])
+        & (concept_pair_similarity["concept_id_2"] == concept_pair_ic_similarity["concept_id_2"]),
+    ).select(concept_pair_similarity_columns + concept_pair_ic_similarity_columns)
+
+    concept_pair_similarity.write.mode("overwrite").parquet(os.path.join(output_folder, CONCEPT_SIMILARITY_PATH))
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Arguments for generate Concept Similarity Table")
+    parser.add_argument(
+        "-i",
+        "--input_folder",
+        dest="input_folder",
+        action="store",
+        help="The path for your input_folder where the raw data is",
+        required=True,
+    )
+    parser.add_argument(
+        "-o",
+        "--output_folder",
+        dest="output_folder",
+        action="store",
+        help="The path for your output_folder",
+        required=True,
+    )
+    parser.add_argument(
+        "-tc",
+        "--domain_table_list",
+        dest="domain_table_list",
+        nargs="+",
+        action="store",
+        help="The list of domain tables you want to download",
+        type=validate_table_names,
+        required=True,
+    )
+    parser.add_argument(
+        "-d",
+        "--date_filter",
+        dest="date_filter",
+        type=lambda s: datetime.datetime.strptime(s, "%Y-%m-%d"),
+        action="store",
+        required=False,
+        default="2018-01-01",
+    )
+    parser.add_argument("--include_concept_list", dest="include_concept_list", action="store_true")
+
+    ARGS = parser.parse_args()
+
+    main(
+        ARGS.input_folder,
+        ARGS.output_folder,
+        ARGS.domain_table_list,
+        ARGS.date_filter,
+        ARGS.include_concept_list,
+    )

cehrbert_data/apps/generate_hierarchical_bert_training_data.py

@@ -0,0 +1,238 @@
+"""
+This module generates hierarchical BERT training data based on domain tables from OMOP EHR data.
+
+It processes patient event data, joins multiple domain tables, filters concepts based on a
+minimum number of patients, and creates hierarchical sequence data for BERT training.
+
+Key Functions:
+    - preprocess_domain_table: Preprocesses domain tables for data extraction.
+    - process_measurement: Handles special processing for measurement data.
+    - join_domain_tables: Joins multiple domain tables into a unified DataFrame.
+    - create_hierarchical_sequence_data: Generates hierarchical sequence data for training.
+
+Command-line Arguments:
+    - input_folder: Path to the directory containing input data.
+    - output_folder: Path to the directory where the output will be saved.
+    - domain_table_list: List of domain tables to process.
+    - date_filter: Optional filter for processing the data based on date.
+    - max_num_of_visits_per_person: Maximum number of visits per patient to include.
+    - min_observation_period: Minimum observation period in days for patients to be included.
+    - include_concept_list: Whether to apply a filter to retain certain concepts.
+    - include_incomplete_visit: Whether to include incomplete visit records in the training data.
+"""
+
+import datetime
+import logging
+import os
+
+from pyspark.sql import SparkSession
+from pyspark.sql import functions as F
+
+from cehrbert_data.config.output_names import PARQUET_DATA_PATH, QUALIFIED_CONCEPT_LIST_PATH
+from cehrbert_data.const.common import MEASUREMENT, OBSERVATION_PERIOD, PERSON, REQUIRED_MEASUREMENT, VISIT_OCCURRENCE
+from cehrbert_data.utils.spark_utils import (
+    create_hierarchical_sequence_data,
+    join_domain_tables,
+    preprocess_domain_table,
+    process_measurement,
+    validate_table_names,
+)
+
+
+def main(
+    input_folder,
+    output_folder,
+    domain_table_list,
+    date_filter,
+    max_num_of_visits_per_person,
+    min_observation_period: int = 360,
+    include_concept_list: bool = True,
+    include_incomplete_visit: bool = True,
+):
+    """
+    Main function to generate hierarchical BERT training data from domain tables.
+
+    Args:
+        input_folder (str): The path to the input folder containing raw data.
+        output_folder (str): The path to the output folder for storing the training data.
+        domain_table_list (list): A list of domain tables to process (e.g., condition_occurrence).
+        date_filter (str): Date filter for processing data, default is '2018-01-01'.
+        max_num_of_visits_per_person (int): The maximum number of visits to include per person.
+        min_observation_period (int, optional): Minimum observation period in days. Default is 360.
+        include_concept_list (bool, optional): Whether to filter by concept list. Default is True.
+        include_incomplete_visit (bool, optional): Whether to include incomplete visits. Default is
+        True.
+
+    This function preprocesses domain tables, filters and processes measurement data,
+    and generates hierarchical sequence data for training BERT models on EHR records.
+    """
+    spark = SparkSession.builder.appName("Generate Hierarchical Bert Training Data").getOrCreate()
+
+    logger = logging.getLogger(__name__)
+    logger.info(
+        "input_folder: %s\n"
+        "output_folder: %s\n"
+        "domain_table_list: %s\n"
+        "date_filter: %s\n"
+        "max_num_of_visits_per_person: %s\n"
+        "min_observation_period: %s\n"
+        "include_concept_list: %s\n"
+        "include_incomplete_visit: %s",
+        input_folder,
+        output_folder,
+        domain_table_list,
+        date_filter,
+        max_num_of_visits_per_person,
+        min_observation_period,
+        include_concept_list,
+        include_incomplete_visit,
+    )
+
+    domain_tables = []
+    # Exclude measurement from domain_table_list if exists because we need to process measurement
+    # in a different way
+    for domain_table_name in domain_table_list:
+        if domain_table_name != MEASUREMENT:
+            domain_tables.append(preprocess_domain_table(spark, input_folder, domain_table_name))
+
+    observation_period = (
+        preprocess_domain_table(spark, input_folder, OBSERVATION_PERIOD)
+        .withColumn(
+            "observation_period_start_date",
+            F.col("observation_period_start_date").cast("date"),
+        )
+        .withColumn(
+            "observation_period_end_date",
+            F.col("observation_period_end_date").cast("date"),
+        )
+        .withColumn(
+            "period",
+            F.datediff("observation_period_end_date", "observation_period_start_date"),
+        )
+        .where(F.col("period") >= min_observation_period)
+        .select("person_id")
+    )
+
+    visit_occurrence = preprocess_domain_table(spark, input_folder, VISIT_OCCURRENCE)
+    person = preprocess_domain_table(spark, input_folder, PERSON)
+
+    # Filter for the persons that have enough observation period
+    person = person.join(observation_period, "person_id").select([person[f] for f in person.schema.fieldNames()])
+
+    # Union all domain table records
+    patient_events = join_domain_tables(domain_tables)
+
+    column_names = patient_events.schema.fieldNames()
+
+    if include_concept_list and patient_events:
+        # Filter out concepts
+        qualified_concepts = F.broadcast(preprocess_domain_table(spark, input_folder, QUALIFIED_CONCEPT_LIST_PATH))
+        # The select is necessary to make sure the order of the columns is the same as the
+        # original dataframe
+        patient_events = patient_events.join(qualified_concepts, "standard_concept_id").select(column_names)
+
+    # Process the measurement table if exists
+    if MEASUREMENT in domain_table_list:
+        measurement = preprocess_domain_table(spark, input_folder, MEASUREMENT)
+        required_measurement = preprocess_domain_table(spark, input_folder, REQUIRED_MEASUREMENT)
+        # The select is necessary to make sure the order of the columns is the same as the
+        # original dataframe, otherwise the union might use the wrong columns
+        scaled_measurement = process_measurement(spark, measurement, required_measurement).select(column_names)
+
+        if patient_events:
+            # Union all measurement records together with other domain records
+            patient_events = patient_events.union(scaled_measurement)
+        else:
+            patient_events = scaled_measurement
+
+    # cohort_member_id is the same as the person_id
+    patient_events = patient_events.withColumn("cohort_member_id", F.col("person_id"))
+
+    sequence_data = create_hierarchical_sequence_data(
+        person,
+        visit_occurrence,
+        patient_events,
+        date_filter=date_filter,
+        max_num_of_visits_per_person=max_num_of_visits_per_person,
+        include_incomplete_visit=include_incomplete_visit,
+    )
+
+    sequence_data.write.mode("overwrite").parquet(os.path.join(output_folder, PARQUET_DATA_PATH))
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Arguments for generate training data for Hierarchical Bert")
+    parser.add_argument(
+        "-i",
+        "--input_folder",
+        dest="input_folder",
+        action="store",
+        help="The path for your input_folder where the raw data is",
+        required=True,
+    )
+    parser.add_argument(
+        "-o",
+        "--output_folder",
+        dest="output_folder",
+        action="store",
+        help="The path for your output_folder",
+        required=True,
+    )
+    parser.add_argument(
+        "-tc",
+        "--domain_table_list",
+        dest="domain_table_list",
+        nargs="+",
+        action="store",
+        help="The list of domain tables you want to download",
+        type=validate_table_names,
+        required=True,
+    )
+    parser.add_argument(
+        "-d",
+        "--date_filter",
+        dest="date_filter",
+        type=lambda s: datetime.datetime.strptime(s, "%Y-%m-%d"),
+        action="store",
+        required=False,
+        default="2018-01-01",
+    )
+    parser.add_argument(
+        "--max_num_of_visits",
+        dest="max_num_of_visits",
+        action="store",
+        type=int,
+        default=200,
+        help="Max no.of visits per patient to be included",
+        required=False,
+    )
+    parser.add_argument(
+        "--min_observation_period",
+        dest="min_observation_period",
+        action="store",
+        type=int,
+        default=1,
+        help="Minimum observation period in days",
+        required=False,
+    )
+    parser.add_argument("--include_concept_list", dest="include_concept_list", action="store_true")
+    parser.add_argument(
+        "--include_incomplete_visit",
+        dest="include_incomplete_visit",
+        action="store_true",
+    )
+
+    ARGS = parser.parse_args()
+
+    main(
+        input_folder=ARGS.input_folder,
+        output_folder=ARGS.output_folder,
+        domain_table_list=ARGS.domain_table_list,
+        date_filter=ARGS.date_filter,
+        max_num_of_visits_per_person=ARGS.max_num_of_visits,
+        min_observation_period=ARGS.min_observation_period,
+        include_concept_list=ARGS.include_concept_list,
+        include_incomplete_visit=ARGS.include_incomplete_visit,
+    )