Make some changes to cohorts to keep compatibility between Analysis a…

…nd Experience
X-DataInitiative · Jan 30, 2020 · 11a53b3 · 11a53b3
1 parent 5c713ad
commit 11a53b3
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/scalpel/core/cohort.py b/scalpel/core/cohort.py
@@ -284,6 +284,12 @@ def load(input: Dict) -> "Cohort":
     def from_description(description: str) -> "Cohort":
         raise NotImplementedError
 
+    def cache(self) -> "Cohort":
+        self.subjects = self.subjects.cache()
+        if self.events is not None:
+            self.events = self.events.cache()
+        return self
+
 
 def _union(a: Cohort, b: Cohort) -> Cohort:
     if a.events is None or b.events is None: