Merge pull request #19 from newgene/add-date-range-to-lineage-by-sub-…

…admin-most-recent Add min_date, max_date to lineage-by-sub-admin-most-recent
newgene · Mar 8, 2023 · fc6bcc9 · fc6bcc9
2 parents ff49e0c + 88f4f38
commit fc6bcc9
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/web/handlers/genomics/util.py b/web/handlers/genomics/util.py
@@ -124,7 +124,7 @@ def compute_cumulative(grp, cols):
                 grp.loc[:, "cum_{}".format(i)] = 0
         return grp.tail(1)
 
-def transform_prevalence_by_location_and_tiime(flattened_response, ndays = None, query_detected = False):
+def transform_prevalence_by_location_and_tiime(flattened_response, ndays = None, query_detected = False, min_date=None, max_date=None):
     df_response = (
         pd.DataFrame(flattened_response)
         .assign(date = lambda x: pd.to_datetime(x["date"], format="%Y-%m-%d"))
@@ -133,9 +133,14 @@ def transform_prevalence_by_location_and_tiime(flattened_response, ndays = None,
     grps = []
     dict_response = {}
     if not query_detected:
-        if ndays is not None:
-            date_limit = dt.today() - timedelta(days = ndays)
-            df_response = df_response[df_response["date"] >= date_limit]
+        if not max_date:
+            max_date = dt.today()
+        if not min_date and ndays is not None:
+            min_date = max_date - timedelta(days=ndays)
+        if min_date:
+            df_response = df_response[df_response["date"] >= min_date]
+        df_response = df_response[df_response["date"] <= max_date]
+
         if df_response.shape[0] == 0:
             return []
         df_response =  df_response.groupby("name").apply(compute_cumulative, ["total_count", "lineage_count"])

diff --git a/web/handlers/v2/genomics/cumulative_prevalence_by_location.py b/web/handlers/v2/genomics/cumulative_prevalence_by_location.py
@@ -1,3 +1,5 @@
+from tornado.web import HTTPError
+
 from web.handlers.genomics.base import BaseHandler
 from web.handlers.genomics.util import (
     create_iterator,
@@ -269,6 +271,8 @@ class CumulativePrevalenceByLocationHandler(BaseHandler):
         "mutations": {"type": str, "default": None},
         "location_id": {"type": str, "default": None},
         "ndays": {"type": int, "default": None, "min": 1},
+        "min_date": {"type": str, "default": None, "date_format": "%Y-%m-%d"},
+        "max_date": {"type": str, "default": None, "date_format": "%Y-%m-%d"},
     }
 
     async def _get(self):
@@ -281,6 +285,10 @@ async def _get(self):
         query_location = self.args.location_id
         query_mutations = query_mutations.split(" AND ") if query_mutations is not None else []
         query_ndays = self.args.ndays
+        min_date = self.args.min_date
+        max_date = self.args.max_date
+        if max_date and min_date and max_date < min_date:
+            raise HTTPError(400, reason="The max_date must greate or equal than the min_date")
         results = {}
         for query_lineage, query_mutation in create_iterator(
             query_pangolin_lineage, query_mutations
@@ -372,7 +380,11 @@ async def _get(self):
                         rec["id"] = "_".join([query_location, i["key"]["sub_id"]])
                     flattened_response.append(rec)
                 dict_response = transform_prevalence_by_location_and_tiime(
-                    flattened_response, query_ndays, query_detected
+                    flattened_response,
+                    query_ndays,
+                    query_detected,
+                    min_date=min_date,
+                    max_date=max_date,
                 )
             res_key = None
             if (