Merge pull request #12 from VACLab/add-progress-bar

Added progress bar to create_cohort() and get_concept_hierarchy()

Author: hyi

biasanalyzer/cohort.py

@@ -2,6 +2,7 @@
 import duckdb
 import pandas as pd
 from datetime import datetime
+from tqdm.auto import tqdm
 from pydantic import ValidationError
 from biasanalyzer.models import CohortDefinition, Cohort
 from biasanalyzer.config import load_cohort_creation_config
@@ -79,10 +80,18 @@ def create_cohort(self, cohort_name: str, description: str, query_or_yaml_file:
         and storing the result in BiasDatabase. The query can be passed in directly
         or built from a yaml file using a corresponding SQL query template
         """
+        stages = [
+            "Built query",
+            "Executed query on OMOP database to get cohort data",
+            "Inserted cohort data into DuckDB - Done"
+        ]
+        progress = tqdm(total=len(stages), desc="Cohort creation", unit="stage", dynamic_ncols=True, leave=True)
+
+        progress.set_postfix_str(stages[0])
         if query_or_yaml_file.endswith('.yaml') or query_or_yaml_file.endswith('.yml'):
             try:
                 cohort_config = load_cohort_creation_config(query_or_yaml_file)
-                print(f'configuration specified in {query_or_yaml_file} loaded successfully')
+                tqdm.write(f'configuration specified in {query_or_yaml_file} loaded successfully')
             except FileNotFoundError:
                 print('specified cohort creation configuration file does not exist. Make sure '
                       'the configuration file name with path is specified correctly.')
@@ -95,7 +104,9 @@ def create_cohort(self, cohort_name: str, description: str, query_or_yaml_file:
             query = self._query_builder.build_query(cohort_config)
         else:
             query = clean_string(query_or_yaml_file)
+        progress.update(1)
 
+        progress.set_postfix_str(stages[1])
         omop_session = self.omop_db.get_session()
         try:
             # Execute read-only query from OMOP CDM database
@@ -108,14 +119,18 @@ def create_cohort(self, cohort_name: str, description: str, query_or_yaml_file:
                 creation_info=clean_string(query),
                 created_by=created_by
             )
-            cohort_def_id = self.bias_db.create_cohort_definition(cohort_def)
+            cohort_def_id = self.bias_db.create_cohort_definition(cohort_def, progress_obj=tqdm)
+            progress.update(1)
 
+            progress.set_postfix_str(stages[2])
             # Store cohort_definition and cohort data into BiasDatabase
             cohort_df = pd.DataFrame(result)
             cohort_df['cohort_definition_id'] = cohort_def_id
             cohort_df = cohort_df.rename(columns={"person_id": "subject_id"})
             self.bias_db.create_cohort_in_bulk(cohort_df)
-            print(f"Cohort {cohort_name} successfully created.")
+            progress.update(1)
+
+            tqdm.write(f"Cohort {cohort_name} successfully created.")
             return CohortData(cohort_id=cohort_def_id, bias_db=self.bias_db, omop_db=self.omop_db)
         except duckdb.Error as e:
             print(f"Error executing query: {e}")

biasanalyzer/database.py

@@ -2,6 +2,7 @@
 import pandas as pd
 from typing import Optional
 from datetime import datetime
+from tqdm.auto import tqdm
 from sqlalchemy.orm import sessionmaker
 from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy import create_engine, text
@@ -91,7 +92,7 @@ def load_postgres_extension(self):
         self.conn.execute("INSTALL postgres_scanner;")
         self.conn.execute("LOAD postgres_scanner;")
 
-    def create_cohort_definition(self, cohort_definition: CohortDefinition):
+    def create_cohort_definition(self, cohort_definition: CohortDefinition, progress_obj=None):
         self.conn.execute('''
             INSERT INTO cohort_definition (name, description, created_date, creation_info, created_by)
             VALUES (?, ?, ?, ?, ?)
@@ -102,7 +103,10 @@ def create_cohort_definition(self, cohort_definition: CohortDefinition):
             cohort_definition.creation_info,
             cohort_definition.created_by
         ))
-        print("Cohort definition inserted successfully.")
+        if progress_obj is None:
+            print("Cohort definition inserted successfully.")
+        else:
+            progress_obj.write("Cohort definition inserted successfully.")
         self.conn.execute("SELECT id from cohort_definition ORDER BY id DESC LIMIT 1")
         created_cohort_id = self.conn.fetchone()[0]
         return created_cohort_id
@@ -407,6 +411,14 @@ def get_concept_hierarchy(self, concept_id: int):
         Retrieves the full concept hierarchy (ancestors and descendants) for a given concept_id
         and organizes it into a nested dictionary to represent the tree structure.
         """
+        stages = [
+            "Queried concept hierarchy",
+            "Fetched concept details",
+            "Built hierarchy tree"
+        ]
+        progress = tqdm(total=len(stages), desc="Concept Hierarchy", unit="stage")
+
+        progress.set_postfix_str(stages[0])
         query = """
                 WITH RECURSIVE concept_hierarchy AS (
                     SELECT ancestor_concept_id, descendant_concept_id, min_levels_of_separation
@@ -425,7 +437,9 @@ def get_concept_hierarchy(self, concept_id: int):
                 """
 
         results = self.execute_query(query, params={"concept_id": concept_id})
+        progress.update(1)
 
+        progress.set_postfix_str(stages[1])
         # Collect all unique concept IDs involved in the hierarchy using set comprehension
         concept_ids = {row['ancestor_concept_id'] for row in results} | {row['descendant_concept_id'] for row in results}
         # Fetch details of each concept
@@ -439,7 +453,9 @@ def get_concept_hierarchy(self, concept_id: int):
 
             result = self.execute_query(query, params={"concept_ids": tuple(concept_ids)})
             concept_details = {row['concept_id']: row for row in result}
+        progress.update(1)
 
+        progress.set_postfix_str(stages[2])
         # Build the hierarchy tree using a dictionary
         hierarchy = {}
         reverse_hierarchy = {}
@@ -458,6 +474,7 @@ def get_concept_hierarchy(self, concept_id: int):
             ancestor_entry_rev = reverse_hierarchy.setdefault(
                 ancestor_id, {"details": concept_details[ancestor_id], "parents": []})
             desc_entry_rev["parents"].append(ancestor_entry_rev)
+        progress.update(1)
 
         # Return the parent hierarchy and children hierarchy of the specified concept
         return reverse_hierarchy[concept_id], hierarchy[concept_id]

poetry.lock

@@ -21,6 +21,7 @@ psycopg2 = "^2.9.1"
 ipytree = "^0.2.2"
 ipywidgets = "^8.1.5"
 jinja2 = "3.1.6"
+tqdm = "4.67.1"
 
 [tool.poetry.dev-dependencies]
 pytest = "^8.3.3"