added attribute dict to Cohort class

limebit · Oct 9, 2024 · 5b7103a · 5b7103a
1 parent 9ed5ef5
commit 5b7103a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 27 deletions.
diff --git a/medmodels/medrecord/types.py b/medmodels/medrecord/types.py
@@ -89,7 +89,7 @@
 PandasEdgeDataFrameInput: TypeAlias = Tuple[pd.DataFrame, str, str]
 
 AttributeInfo: TypeAlias = Union[
-    TemporalAttributeInfo, NumericAttributeInfo, StringAttributeInfo
+    "TemporalAttributeInfo", "NumericAttributeInfo", "StringAttributeInfo"
 ]
 
 

diff --git a/medmodels/statistic_evaluations/comparer/data_comparer.pyi b/medmodels/statistic_evaluations/comparer/data_comparer.pyi
@@ -10,7 +10,6 @@ from medmodels.medrecord.types import (
     Group,
     GroupInputList,
     MedRecordAttribute,
-    MedRecordAttributeInputList,
     NodeIndex,
 )
 
@@ -26,22 +25,24 @@ class Cohort:
     name: str
     cohort_group: Group
     time_attribute: MedRecordAttribute
-    concepts_groups: GroupInputList
+    attributes: Optional[Dict[str, MedRecordAttribute]]
+    concepts_groups: Optional[GroupInputList]
 
     def __init__(
         self,
         medrecord: MedRecord,
         name: str,
         cohort_group: Union[Group, NodeOperation] = "patients",
         time_attribute: MedRecordAttribute = "time",
+        attributes: Optional[Dict[str, MedRecordAttribute]] = None,
         concepts_groups: Optional[GroupInputList] = None,
     ) -> None: ...
 
 class CohortSummary(TypedDict):
     """Dictionary for the cohort summary."""
 
     attribute_info: Dict[Group, AttributeSummary]
-    top_concepts: Dict[Group, List[NodeIndex]]
+    top_k_concepts: Dict[Group, List[NodeIndex]]
 
 class DistanceSummary(TypedDict):
     """Dictonary for the Jensen-Shannon-Divergence and normalized distance between
@@ -53,8 +54,8 @@ class DistanceSummary(TypedDict):
 class ComparerSummary(TypedDict):
     """Dictionary for the comparing results."""
 
-    attribute_tests: Dict[MedRecordAttribute, TestSummary]
-    concepts_tests: Dict[Group, TestSummary]
+    attribute_tests: Dict[MedRecordAttribute, List[TestSummary]]
+    concepts_tests: Dict[Group, List[TestSummary]]
     concepts_distance: Dict[Group, DistanceSummary]
 
 class TestSummary(TypedDict):
@@ -64,45 +65,44 @@ class TestSummary(TypedDict):
     p_value: float
 
 class DataComparer:
+    @staticmethod
+    def compare_cohort_attribute(
+        cohorts_attribute: List[Tuple[Cohort, MedRecordAttribute]],
+    ) -> Dict[str, AttributeInfo]: ...
+    @staticmethod
+    def test_difference_attribute(
+        cohorts_attribute: List[Tuple[Cohort, MedRecordAttribute]],
+        significance_level: float,
+    ) -> List[TestSummary]: ...
     @staticmethod
     def compare_cohorts(
         cohorts: List[Cohort],
-    ) -> Dict[Group, CohortSummary]: ...
+    ) -> Dict[str, CohortSummary]: ...
     @staticmethod
-    def compare_cohort_attributes(
-        medrecords: List[Cohort],
-        attributes: MedRecordAttributeInputList,
-    ) -> Dict[
-        Group,
-        Dict[MedRecordAttribute, AttributeInfo],
-    ]: ...
+    def test_difference_cohort_attributes(
+        cohorts: List[Cohort],
+        significance_level: float,
+    ) -> Dict[str, List[TestSummary]]: ...
     @staticmethod
     def calculate_absolute_relative_difference(
-        control_cohort: Cohort,
-        cohorts: List[Cohort],
-        attributes: MedRecordAttributeInputList,
+        control_group: Cohort,
+        case_groups: List[Cohort],
     ) -> Tuple[float, Dict[MedRecordAttribute, float]]: ...
     @staticmethod
-    def test_difference_attributes(
-        cohorts: List[Cohort],
-        attributes: MedRecordAttributeInputList,
-        significance_level: float,
-    ) -> Dict[MedRecordAttribute, TestSummary]: ...
+    def get_concept_counts(
+        cohort: Cohort,
+    ) -> List[Tuple[NodeIndex, int]]: ...
     @staticmethod
     def get_top_k_concepts(
         cohort: Cohort,
         top_k: int,
     ) -> List[NodeIndex]: ...
     @staticmethod
-    def get_concept_counts(
-        cohort: Cohort,
-    ) -> List[Tuple[NodeIndex, int]]: ...
-    @staticmethod
     def test_difference_top_k_concepts(
         cohorts: List[Cohort],
         top_k: int,
         significance_level: float,
-    ) -> Dict[Group, TestSummary]: ...
+    ) -> Dict[Group, List[TestSummary]]: ...
     @staticmethod
     def calculate_distance_concepts(
         cohorts: List[Cohort],