def test_add_subject_information(self):
patients_1, _ = self.create_spark_df({"patientID": [1, 2]})
events_1, _ = self.create_spark_df({
"patientID": [1, 2],
"value": ["fracture", "fracture"]
})
input1 = Cohort("liberal_fractures", "liberal_fractures", patients_1,
events_1)
patients_2, _ = self.create_spark_df({
"patientID": [1, 2],
"gender": [1, 1],
"birthDate": [
pd.to_datetime("1993-10-09"),
pd.to_datetime("1992-03-14"),
],
"deathDate": [
pd.to_datetime("1993-10-09"),
pd.to_datetime("1992-03-14"),
],
})
base_cohort1 = Cohort("patients", "patients", patients_2)
input1.add_subject_information(base_cohort1, "error")
self.assertTrue(input1.has_subject_information()
and input1.subjects.count() == 2)
patients_3, _ = self.create_spark_df({
"patientID": [1],
"gender": [1],
"birthDate": [pd.to_datetime("1993-10-09")],
"deathDate": [pd.to_datetime("1993-10-09")],
})
base_cohort2 = Cohort("liberal_fractures", "liberal_fractures",
patients_3, None)
input2 = Cohort("liberal_fractures", "liberal_fractures", patients_1,
events_1)
input2.add_subject_information(base_cohort2, "omit")
self.assertTrue(input2.has_subject_information()
and input2.subjects.count() == 1
and input2.events.count() == 2)
input3 = Cohort("liberal_fractures", "liberal_fractures", patients_1,
events_1)
input3.add_subject_information(base_cohort2, "omit_all")
self.assertTrue(input3.has_subject_information()
and input3.subjects.count() == 1
and input3.events.count() == 1)
def test_has_subject_information(self):
patients_1, _ = self.create_spark_df({"patientID": [1, 2]})
cohort1 = Cohort("liberal_fractures", "liberal_fractures", patients_1)
patients_2, _ = self.create_spark_df({
"patientID": [1, 2],
"gender": [1, 1],
"birthDate": [
pd.to_datetime("1993-10-09"),
pd.to_datetime("1992-03-14"),
],
"deathDate": [
pd.to_datetime("1993-10-09"),
pd.to_datetime("1992-03-14"),
],
})
cohort2 = Cohort("liberal_fractures", "liberal_fractures", patients_2,
None)
self.assertFalse(cohort1.has_subject_information())
self.assertTrue(cohort2.has_subject_information())
def _find_subjects_with_age_inconsistent_w_age_groups(
self, cohort: Cohort) -> Cohort:
"""Check if min and max age_groups are consistent with subjects ages."""
if not cohort.has_subject_information():
raise ValueError("Cohort should have subject information.")
duplicate = copy(cohort)
duplicate.add_age_information(
self.age_reference_date) # add starting age
study_length = (np.ceil(
(self.study_end - self.study_start).days /
365.25) if self.is_using_longitudinal_age_groups else 0)
min_starting_age = min(self.age_groups)
max_starting_age = max(self.age_groups) - np.ceil(study_length)
invalid_subjects = duplicate.subjects.where(
~sf.col("age").between(min_starting_age, max_starting_age))
return Cohort(
cohort.name + "_inconsistent_w_ages_and_age_groups",
"subjects inconsistent with age groups",
invalid_subjects,
)
def _compute_longitudinal_age_groups(
self, cohort: Cohort,
col_offset: int) -> Tuple[DataFrame, List[str]]:
"""
Parameters
----------
cohort: Cohort
cohort on which the age groups should be computed
col_offset: int
number of columns used by lagged exposure features
Returns
-------
(age_features, mapping): Tuple(DataFrame, List(str))
a dataframe containing the age features in aij format and a mapping giving
the correspondence between column number and age group.
"""
# This implementation is suboptimal, but we need to have something
# working with inconsistent python versions across the cluster.
assert (cohort.has_subject_information(
)), "Cohort subjects should have gender and birthdate information"
subjects = cohort.subjects.select("patientID", "gender", "birthDate")
bucket_ids = sf.array([sf.lit(i) for i in range(self.n_buckets)])
subjects = (subjects.withColumn("bucketID", bucket_ids).select(
"PatientID",
"gender",
"birthDate",
sf.explode("bucketID").alias("bucket"),
).withColumn("dateShift",
sf.col("bucket") * self.bucket_size).withColumn(
"referenceDate", sf.lit(self.age_reference_date)))
# Longitudinal age is based on referenceDate instead of minDate to
# be consistent with cohort definition.
time_references = sf.expr("date_add(referenceDate, dateShift)")
longitudinal_age = sf.floor(
sf.months_between(time_references, sf.col("birthdate")) / 12)
subjects = subjects.withColumn("longitudinalAge",
longitudinal_age).select(
"patientID", "gender", "birthDate",
"bucket", "longitudinalAge")
subjects, n_age_groups, mapping = self._bucketize_age_column(
subjects, "longitudinalAge", "longitudinalAgeBucket")
assert n_age_groups == self.n_age_groups, (
"Computed number of age groups is different from the number of specified"
" age groups at initialization. There might be empty age_groups,"
" you should investigate this.")
age_features = subjects.select(
sf.col("patientID"),
sf.col("bucket").alias("rowIndex"),
(sf.col("longitudinalAgeBucket") + col_offset).alias("colIndex"),
)
# Remove "age events" which are not in follow-up
fup_events = self.followups.intersection(self.final_cohort).events
fup_events = self._discretize_start_end(fup_events)
fup_events = rename_df_columns(fup_events, prefix="fup_")
age_features_columns = age_features.columns
age_features = age_features.join(fup_events, on="patientID")
age_features = age_features.where(
sf.col("rowIndex").between(sf.col("fup_startBucket"),
sf.col("fup_endBucket")))
age_features = age_features.select(*age_features_columns)
return age_features, mapping