cehrgpt 0.0.1__py3-none-any.whl

cehrgpt/analysis/privacy/reid_inference.py

@@ -0,0 +1,407 @@
+import os
+
+from pyspark.sql import SparkSession
+from pyspark.sql import functions as f
+from pyspark.sql import types as t
+from tqdm import tqdm
+
+VOCAB_SIZE = 17870
+L = 0.001
+SENSITIVE_MATCH_THRESHOLD = int(VOCAB_SIZE * L)
+
+# This number came from the tokenizer, essentially all numeric tokens that represent valid OMOP concept ids
+COMMON_ATTRIBUTES = [
+    "320128",
+    "200219",
+    "77670",
+    "432867",
+    "254761",
+    "312437",
+    "378253",
+]
+MAX_LEVELS = [1, 0, 0] + [1] * len(COMMON_ATTRIBUTES)
+# The age groups are generated using the following bases e.g. 1 indicates using the age as is
+AGE_GENERALIZATION_LEVELS = [10, 1]
+
+
+def generate_lattice_bfs(top_gen_levels):  # BFS
+    """Came from https://github.com/yy6linda/synthetic-ehr-benchmarking/blob/main/privacy_evaluation/Synthetic_risk_model_reid.py#L81."""
+    visited = [top_gen_levels]
+    queue = [top_gen_levels]
+    lattice = []
+    while queue:
+        gen_levels = queue.pop(0)
+        lattice.append(gen_levels)
+        for i in range(len(gen_levels)):
+            if gen_levels[i] != 0:
+                gen_levels_new = gen_levels.copy()
+                gen_levels_new[i] -= 1
+                if not gen_levels_new in visited:
+                    visited.append(gen_levels_new)
+                    queue.append(gen_levels_new)
+    return lattice
+
+
+def update_dataset(reid_data, real_sample, synthetic_reid_data, config):
+    reid_data_dup = reid_data.alias("reid_data_dup")
+    real_sample_dup = real_sample.alias("real_sample_dup")
+    synthetic_reid_data_dup = synthetic_reid_data.alias("synthetic_reid_data_dup")
+    common_attributes_to_remove = []
+
+    for i in range(len(config)):
+        # age group
+        if i == 0:
+            age_level_index = config[i]
+            age_group_base = AGE_GENERALIZATION_LEVELS[age_level_index]
+            reid_data_dup = reid_data_dup.withColumn(
+                "age", f.ceil(f.col("age").cast("int") / age_group_base)
+            )
+            synthetic_reid_data_dup = synthetic_reid_data_dup.withColumn(
+                "age", f.ceil(f.col("age").cast("int") / age_group_base)
+            )
+            real_sample_dup = real_sample_dup.withColumn(
+                "age", f.ceil(f.col("age").cast("int") / age_group_base)
+            )
+        elif i in [1, 2]:
+            # gender and race are not generalized
+            continue
+        else:
+            # this indicates that the common attribute should be removed
+            if config[i] == 0:
+                common_attributes_to_remove.append(COMMON_ATTRIBUTES[i - 3])
+
+    def remove_common_attributes(concept_ids):
+        comm_atts = sorted(set(concept_ids) - set(common_attributes_to_remove))
+        if comm_atts:
+            return "-".join(comm_atts)
+        else:
+            return "empty"
+
+    extract_common_attributes_udf = f.udf(
+        lambda concept_ids: remove_common_attributes(concept_ids), t.StringType()
+    )
+
+    reid_data_dup = reid_data_dup.withColumn(
+        "common_attributes", extract_common_attributes_udf("common_attributes")
+    )
+
+    real_sample_dup = real_sample_dup.withColumn(
+        "common_attributes", extract_common_attributes_udf("common_attributes")
+    )
+
+    synthetic_reid_data_dup = synthetic_reid_data_dup.withColumn(
+        "common_attributes", extract_common_attributes_udf("common_attributes")
+    )
+    return reid_data_dup, real_sample_dup, synthetic_reid_data_dup
+
+
+def calculate_reid_risk_score(
+    real_sample_dup,
+    reid_data_dup,
+    synthetic_reid_data_dup,
+    lower_n,
+    cap_n,
+    lambda_val=0.23,
+    num_salts=20,
+):
+    real_sample_dup = real_sample_dup.withColumn(
+        "salt", (f.rand() * num_salts).cast("int")
+    )
+
+    reid_data_stats = reid_data_dup.groupby(
+        "age", "gender", "race", "common_attributes"
+    ).count()
+
+    real_to_population_matches = real_sample_dup.join(
+        reid_data_stats,
+        (real_sample_dup["age"] == reid_data_stats["age"])
+        & (real_sample_dup["gender"] == reid_data_stats["gender"])
+        & (real_sample_dup["race"] == reid_data_stats["race"])
+        & (
+            real_sample_dup["common_attributes"] == reid_data_stats["common_attributes"]
+        ),
+    ).select("person_id", "count")
+
+    # Alias the DataFrame for self join
+    real_sample_stats = real_sample_dup.groupby(
+        "age", "gender", "race", "common_attributes"
+    ).count()
+
+    real_to_real_matches = real_sample_dup.join(
+        real_sample_stats,
+        (real_sample_dup["age"] == real_sample_stats["age"])
+        & (real_sample_dup["gender"] == real_sample_stats["gender"])
+        & (real_sample_dup["race"] == real_sample_stats["race"])
+        & (
+            real_sample_dup["common_attributes"]
+            == real_sample_stats["common_attributes"]
+        ),
+    ).select("person_id", "count")
+
+    reid_data_step_one = (
+        reid_data_dup.join(real_to_population_matches, "person_id")
+        .withColumnRenamed("count", "upper_F_s")
+        .join(real_to_real_matches, "person_id")
+        .withColumnRenamed("count", "lower_f_s")
+    )
+
+    synthetic_reid_data_dup = synthetic_reid_data_dup.withColumn(
+        "salt", f.explode(f.array([f.lit(x) for x in range(num_salts)]))
+    )
+
+    real_sample_dup = real_sample_dup.where(
+        f.size("sensitive_attributes") >= SENSITIVE_MATCH_THRESHOLD
+    )
+    synthetic_reid_data_dup = synthetic_reid_data_dup.where(
+        f.size("sensitive_attributes") >= SENSITIVE_MATCH_THRESHOLD
+    )
+
+    real_to_synthetic = (
+        real_sample_dup.join(
+            synthetic_reid_data_dup,
+            (real_sample_dup["age"] == synthetic_reid_data_dup["age"])
+            & (real_sample_dup["gender"] == synthetic_reid_data_dup["gender"])
+            & (real_sample_dup["race"] == synthetic_reid_data_dup["race"])
+            & (
+                real_sample_dup["common_attributes"]
+                == synthetic_reid_data_dup["common_attributes"]
+            )
+            & (real_sample_dup["salt"] == synthetic_reid_data_dup["salt"]),
+        )
+        .select(
+            real_sample_dup["person_id"],
+            real_sample_dup["sensitive_attributes"].alias("real_sensitive_attributes"),
+            synthetic_reid_data_dup["sensitive_attributes"].alias(
+                "synthetic_sensitive_attributes"
+            ),
+        )
+        .withColumn(
+            "n_of_sensitive_matches",
+            f.size(
+                f.array_intersect(
+                    "real_sensitive_attributes", "synthetic_sensitive_attributes"
+                )
+            ),
+        )
+        .drop("real_sensitive_attributes", "synthetic_sensitive_attributes")
+    )
+
+    real_to_synthetic = (
+        real_to_synthetic.groupby("person_id")
+        .agg(f.max("n_of_sensitive_matches").alias("n_of_sensitive_matches"))
+        .withColumn(
+            "new_info",
+            (f.col("n_of_sensitive_matches") > f.lit(SENSITIVE_MATCH_THRESHOLD)).cast(
+                "int"
+            ),
+        )
+    )
+
+    reid_data_step_two = reid_data_step_one.join(
+        real_to_synthetic, "person_id", how="left_outer"
+    ).withColumn("new_info", f.coalesce(real_to_synthetic["new_info"], f.lit(0)))
+
+    return (
+        reid_data_step_two.withColumn(
+            "A_term",
+            1 / f.col("lower_f_s") * f.lit((1 + lambda_val) / 2) * f.col("new_info"),
+        )
+        .withColumn(
+            "B_term",
+            1 / f.col("upper_f_s") * f.lit((1 + lambda_val) / 2) * f.col("new_info"),
+        )
+        .select(
+            (f.sum("A_term") / f.lit(cap_n)).alias("A_term"),
+            (f.sum("B_term") / f.lit(lower_n)).alias("B_term"),
+        )
+    )
+
+
+def main(args):
+    all_configs = generate_lattice_bfs(MAX_LEVELS)
+
+    N = None
+    n = None
+    excluded_sensitive_attributes = None
+
+    for config in tqdm(all_configs, total=len(all_configs)):
+
+        spark = SparkSession.builder.appName(
+            f'Generate REID {"".join(map(str, config))}'
+        ).getOrCreate()
+
+        if sum(config) == 0:
+            continue
+
+        experiment_output_name = os.path.join(
+            args.output_folder, f'{"".join(map(str, config))}.parquet'
+        )
+        if os.path.exists(experiment_output_name):
+            continue
+
+        @f.udf(t.ArrayType(t.StringType()))
+        def extract_omop_concepts_udf(concept_ids):
+            return list(sorted(set([_ for _ in concept_ids if str.isnumeric(_)])))
+
+        def extract_common_attributes(concept_ids):
+            commn_atts = set([c for c in concept_ids if c in COMMON_ATTRIBUTES])
+            return list(commn_atts)
+
+        extract_common_attributes_udf = f.udf(
+            lambda concept_ids: extract_common_attributes(concept_ids),
+            t.ArrayType(t.StringType()),
+        )
+
+        def extract_sensitive_attributes(concept_ids):
+            return list(
+                sorted(set([c for c in concept_ids if c not in COMMON_ATTRIBUTES]))
+            )
+
+        extract_sensitive_attributes_udf = f.udf(
+            lambda concept_ids: extract_sensitive_attributes(concept_ids),
+            t.ArrayType(t.StringType()),
+        )
+
+        training_data = spark.read.parquet(args.training_data_folder)
+        evaluation_data = spark.read.parquet(args.evaluation_data_folder)
+        training_data = training_data.select(
+            "concept_ids", "num_of_concepts", "person_id"
+        ).withColumn("is_real", f.lit(True))
+        evaluation_data = evaluation_data.select(
+            "concept_ids", "num_of_concepts", "person_id"
+        ).withColumn("is_real", f.lit(False))
+        data = training_data.unionByName(evaluation_data)
+
+        reid_data = (
+            data.withColumn("age", f.col("concept_ids")[1])
+            .withColumn("age", f.split("age", ":")[1])
+            .withColumn("gender", f.col("concept_ids")[2])
+            .withColumn("race", f.col("concept_ids")[3])
+            .withColumn("concept_ids", extract_omop_concepts_udf("concept_ids"))
+            .withColumn(
+                "common_attributes", extract_common_attributes_udf("concept_ids")
+            )
+            .withColumn(
+                "sensitive_attributes", extract_sensitive_attributes_udf("concept_ids")
+            )
+            .drop("concept_ids")
+        )
+
+        real_sample = reid_data.where("is_real")
+
+        if not N:
+            N = reid_data.count()
+        if not n:
+            n = real_sample.count()
+
+        if not excluded_sensitive_attributes:
+            excluded_sensitive_attributes_df = (
+                real_sample.select(
+                    f.explode("sensitive_attributes").alias("sensitive_attribute")
+                )
+                .groupby("sensitive_attribute")
+                .count()
+                .withColumn("sensitive_attribute_prevalence", f.col("count") / n)
+                .withColumn(
+                    "is_majority", f.col("sensitive_attribute_prevalence") >= 0.5
+                )
+                .where("is_majority")
+            )
+            excluded_sensitive_attributes = [
+                row.sensitive_attribute
+                for row in excluded_sensitive_attributes_df.select(
+                    "sensitive_attribute"
+                ).collect()
+            ]
+
+        def filter_sensitive_attributes(concept_ids):
+            return [c for c in concept_ids if c not in excluded_sensitive_attributes]
+
+        filter_sensitive_attributes_udf = f.udf(
+            lambda concept_ids: filter_sensitive_attributes(concept_ids),
+            t.ArrayType(t.StringType()),
+        )
+
+        real_sample = real_sample.withColumn(
+            "sensitive_attributes",
+            filter_sensitive_attributes_udf("sensitive_attributes"),
+        )
+
+        synthetic_data = spark.read.parquet(args.synthetic_data_folder)
+
+        synthetic_reid_data = (
+            synthetic_data.withColumn("age", f.col("concept_ids")[1])
+            .withColumn("age", f.split("age", ":")[1])
+            .withColumn("gender", f.col("concept_ids")[2])
+            .withColumn("race", f.col("concept_ids")[3])
+            .withColumn("concept_ids", extract_omop_concepts_udf("concept_ids"))
+            .withColumn(
+                "common_attributes", extract_common_attributes_udf("concept_ids")
+            )
+            .withColumn(
+                "sensitive_attributes", extract_sensitive_attributes_udf("concept_ids")
+            )
+            .drop("concept_ids")
+        )
+
+        reid_data_dup, real_sample_dup, synthetic_reid_data_dup = update_dataset(
+            reid_data, real_sample, synthetic_reid_data, config=config
+        )
+        experiment_risk_score = calculate_reid_risk_score(
+            real_sample_dup,
+            reid_data_dup,
+            synthetic_reid_data_dup,
+            lower_n=n,
+            cap_n=N,
+            lambda_val=0.23,
+            num_salts=args.num_salts,
+        )
+        experiment_risk_score.toPandas().to_parquet(experiment_output_name)
+
+        spark.stop()
+
+
+def create_argparser():
+    import argparse
+
+    parser = argparse.ArgumentParser(
+        description="Arguments for re-identification risk evaluation"
+    )
+    parser.add_argument(
+        "--training_data_folder",
+        dest="training_data_folder",
+        action="store",
+        required=True,
+    )
+    parser.add_argument(
+        "--evaluation_data_folder",
+        dest="evaluation_data_folder",
+        action="store",
+        required=True,
+    )
+    parser.add_argument(
+        "--synthetic_data_folder",
+        dest="synthetic_data_folder",
+        action="store",
+        required=True,
+    )
+    parser.add_argument(
+        "--output_folder",
+        dest="output_folder",
+        action="store",
+        help="The output folder for storing the results",
+        required=True,
+    )
+    parser.add_argument(
+        "--num_salts",
+        dest="num_salts",
+        action="store",
+        type=int,
+        default=40,
+        required=False,
+    )
+    return parser
+
+
+if __name__ == "__main__":
+    main(create_argparser().parse_args())

cehrgpt/analysis/privacy/utils.py

@@ -0,0 +1,255 @@
+import logging
+
+import numpy as np
+from sklearn.preprocessing import OneHotEncoder
+from tqdm import tqdm
+
+LOG = logging.getLogger(__name__)
+
+RANDOM_SEE = 42
+NUM_OF_GENDERS = 3
+NUM_OF_RACES = 21
+
+
+def get_demographics(concept_ids):
+    year_token, age_token, gender, race = concept_ids[0:4]
+    try:
+        year = int(year_token[5:])
+    except ValueError:
+        LOG.error(
+            f"{year_token[5:]} cannot be converted to an integer, use the default value 1900"
+        )
+        year = 1900
+
+    try:
+        age = int(age_token[4:])
+    except ValueError:
+        LOG.error(
+            f"{age_token[4:]} cannot be converted to an integer, use the default value 1900"
+        )
+        age = -1
+
+    return year, age, gender, race
+
+
+def create_race_encoder(
+    train_data_sample, evaluation_data_sample, synthetic_data_sample
+):
+    race_encoder = OneHotEncoder()
+    all_unique_races = np.unique(
+        np.concatenate(
+            [
+                train_data_sample.race.unique(),
+                evaluation_data_sample.race.unique(),
+                synthetic_data_sample.race.unique(),
+            ],
+            axis=0,
+        )
+    )
+    race_encoder.fit(all_unique_races[:, np.newaxis])
+    return race_encoder
+
+
+def create_gender_encoder(
+    train_data_sample, evaluation_data_sample, synthetic_data_sample
+):
+    gender_encoder = OneHotEncoder()
+    all_unique_genders = np.unique(
+        np.concatenate(
+            [
+                train_data_sample.gender.unique(),
+                evaluation_data_sample.gender.unique(),
+                synthetic_data_sample.gender.unique(),
+            ],
+            axis=0,
+        )
+    )
+    gender_encoder.fit(all_unique_genders[:, np.newaxis])
+    return gender_encoder
+
+
+def extract_medical_concepts(concept_ids):
+    concept_ids = [_ for _ in concept_ids[4:] if str.isnumeric(_)]
+    return list(set(concept_ids))
+
+
+def create_binary_format(concept_ids, concept_tokenizer):
+    indices = np.array(concept_tokenizer.encode(concept_ids)).flatten().astype(int)
+    embeddings = np.zeros(concept_tokenizer.vocab_size)
+    embeddings.put(indices, 1)
+    return embeddings
+
+
+def extract_common_sensitive_concepts(
+    dataset, concept_tokenizer, common_attributes, sensitive_attributes
+):
+    common_embeddings_list = []
+    sensitive_embeddings_list = []
+    for _, pat_seq in dataset.concept_ids.items():
+        concept_ids = extract_medical_concepts(pat_seq)
+        common_concept_ids = [_ for _ in concept_ids if _ in common_attributes]
+        sensitive_concept_ids = [_ for _ in concept_ids if _ in sensitive_attributes]
+
+        common_embeddings_list.append(
+            create_binary_format(common_concept_ids, concept_tokenizer)
+        )
+        sensitive_embeddings_list.append(
+            create_binary_format(sensitive_concept_ids, concept_tokenizer)
+        )
+    return np.array(common_embeddings_list), np.array(sensitive_embeddings_list)
+
+
+def transform_concepts(dataset, concept_tokenizer):
+    embedding_list = []
+    for _, pat_seq in dataset.concept_ids.items():
+        embedding_list.append(
+            create_binary_format(extract_medical_concepts(pat_seq), concept_tokenizer)
+        )
+    return np.asarray(embedding_list)
+
+
+def scale_age(dataset):
+    ages = dataset.concept_ids.apply(
+        lambda concept_list: get_demographics(concept_list)[1]
+    )
+    dataset = dataset[ages >= 0]
+    ages = ages[ages >= 0]
+    max_age = ages.max()
+    dataset["scaled_age"] = ages / max_age
+    return dataset
+
+
+def create_demographics(dataset):
+    genders = dataset.concept_ids.apply(
+        lambda concept_list: get_demographics(concept_list)[2]
+    )
+    races = dataset.concept_ids.apply(
+        lambda concept_list: get_demographics(concept_list)[3]
+    )
+    dataset["gender"] = genders
+    dataset["race"] = races
+    return dataset
+
+
+def create_vector_representations(
+    dataset, concept_tokenizer, gender_encoder, race_encoder
+):
+    concept_vectors = transform_concepts(dataset, concept_tokenizer)
+    gender_vectors = gender_encoder.transform(
+        dataset.gender.to_numpy()[:, np.newaxis]
+    ).todense()
+    race_vectors = race_encoder.transform(
+        dataset.race.to_numpy()[:, np.newaxis]
+    ).todense()
+    age_vectors = dataset.scaled_age.to_numpy()[:, np.newaxis]
+
+    pat_vectors = np.concatenate(
+        [age_vectors, gender_vectors, race_vectors, concept_vectors], axis=-1
+    )
+
+    return np.asarray(pat_vectors)
+
+
+def create_vector_representations_for_attribute(
+    dataset,
+    concept_tokenizer,
+    gender_encoder,
+    race_encoder,
+    common_attributes,
+    sensitive_attributes,
+):
+    common_concept_vectors, sensitive_concept_vectors = (
+        extract_common_sensitive_concepts(
+            dataset, concept_tokenizer, common_attributes, sensitive_attributes
+        )
+    )
+    gender_vectors = gender_encoder.transform(
+        dataset.gender.to_numpy()[:, np.newaxis]
+    ).todense()
+    race_vectors = race_encoder.transform(
+        dataset.race.to_numpy()[:, np.newaxis]
+    ).todense()
+    age_vectors = dataset.scaled_age.to_numpy()[:, np.newaxis]
+
+    common_pat_vectors = np.concatenate(
+        [age_vectors, gender_vectors, race_vectors, common_concept_vectors], axis=-1
+    )
+
+    return np.asarray(common_pat_vectors), np.asarray(sensitive_concept_vectors)
+
+
+def batched_pairwise_euclidean_distance_indices(A, B, batch_size, self_exclude=False):
+    # Initialize arrays to hold the minimum distances and indices for each point in A
+    min_distances = np.full((A.shape[0],), np.inf)
+    min_indices = np.full((A.shape[0],), -1, dtype=int)
+
+    # Iterate over A in batches
+    for i in tqdm(range(0, A.shape[0], batch_size), total=A.shape[0] // batch_size + 1):
+        end_i = i + batch_size
+        A_batch = A[i:end_i]
+
+        # Adjust the identity matrix size based on the actual batch size
+        actual_batch_size = A_batch.shape[0]
+
+        # Iterate over B in batches
+        for j in range(0, B.shape[0], batch_size):
+            end_j = j + batch_size
+            B_batch = B[j:end_j]
+
+            # Compute distances between the current batches of A and B
+            distances = np.sqrt(
+                np.sum(
+                    (A_batch[:, np.newaxis, :] - B_batch[np.newaxis, :, :]) ** 2, axis=2
+                )
+            )
+
+            # Apply the identity matrix to exclude self-matches if required
+            if self_exclude and i == j:
+                identity_matrix = np.eye(actual_batch_size) * 10e8
+                distances += identity_matrix
+
+            # Find the minimum distance and corresponding indices for the A batch
+            min_batch_indices = np.argmin(distances, axis=1) + j
+            min_batch_distances = np.min(distances, axis=1)
+
+            # Update the minimum distances and indices if the current batch distances are smaller
+            update_mask = min_batch_distances < min_distances[i:end_i]
+            min_distances[i:end_i][update_mask] = min_batch_distances[update_mask]
+            min_indices[i:end_i][update_mask] = min_batch_indices[update_mask]
+
+    return min_indices
+
+
+def find_match(source, target, return_index: bool = False):
+    a = np.sum(target**2, axis=1).reshape(target.shape[0], 1) + np.sum(
+        source.T**2, axis=0
+    )
+    b = np.dot(target, source.T) * 2
+    distance_matrix = a - b
+    return (
+        np.argmin(distance_matrix, axis=0)
+        if return_index
+        else np.min(distance_matrix, axis=0)
+    )
+
+
+def find_match_self(source, target, return_index: bool = False):
+    a = np.sum(target**2, axis=1).reshape(target.shape[0], 1) + np.sum(
+        source.T**2, axis=0
+    )
+    b = np.dot(target, source.T) * 2
+    distance_matrix = a - b
+    n_col = np.shape(distance_matrix)[1]
+
+    if return_index:
+        min_indices = np.zeros(n_col, dtype=int)
+        for i in range(n_col):
+            sorted_indices = np.argsort(distance_matrix[:, i])
+            min_indices[i] = sorted_indices[1]  # Get index of second smallest value
+        return min_indices
+    else:
+        min_distance = np.zeros(n_col)
+        for i in range(n_col):
+            sorted_column = np.sort(distance_matrix[:, i])
+            min_distance[i] = sorted_column[1]
+        return min_distance