ESGPT support

justin13601 · Aug 9, 2024 · e9af047 · e9af047
1 parent e209857
commit e9af047
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git a/sample_configs/inhospital_mortality.yaml b/sample_configs/inhospital_mortality.yaml
@@ -1,8 +1,10 @@
 # Task: 24-hour In-hospital Mortality Prediction
 predicates:
   admission:
-    code: event_type//ADMISSION
-    random_col: foo
+    code:
+      regex:
+      any: [event_type//ADMISSION, event_type//EMERGENCY, event_type//ELECTIVE]
+      all:
   discharge:
     code: event_type//DISCHARGE
   death:

diff --git a/src/aces/config.py b/src/aces/config.py
@@ -129,6 +129,9 @@ def ESGPT_eval_expr(self, values_column: str | None = None) -> pl.Expr:
             >>> expr = PlainPredicateConfig("BP").ESGPT_eval_expr()
             >>> print(expr) # doctest: +NORMALIZE_WHITESPACE
             col("BP").is_not_null()
+            >>> expr = PlainPredicateConfig("BP//systole", other_cols={"chamber": "atrial"}).ESGPT_eval_expr()
+            >>> print(expr) # doctest: +NORMALIZE_WHITESPACE
+            [(col("BP")) == (String(systole))].all_horizontal([[(col("chamber")) == (String(atrial))]])
             >>> expr = PlainPredicateConfig("BP//systolic", value_min=120).ESGPT_eval_expr()
             Traceback (most recent call last):
                 ...
@@ -176,6 +179,9 @@ def ESGPT_eval_expr(self, values_column: str | None = None) -> pl.Expr:
             else:
                 criteria.append(pl.col(values_column) < self.value_max)
 
+        if self.other_cols:
+            criteria.extend([pl.col(col) == value for col, value in self.other_cols.items()])
+
         if len(criteria) == 1:
             return criteria[0]
         else: