src/nt_composition/make_nt_correlation.py: modification of...

src/nt_composition/make_nt_correlation.py: modification of get_project_colocalisation function by adding a parameter level

src/nt_composition/make_nt_correlation.py

@@ -61,9 +61,10 @@ def get_project_colocalisation(cnx: sqlite3.Connection, project: str,
                                global_weight: int, same_gene: bool,
                                get_weight: bool = False,
                                exon_bc: Optional[np.array] = None,
-                               inter_chr: bool = False) -> np.array:
+                               inter_chr: bool = False,
+                               level: str = "exon") -> np.array:
     """
-    Get the interactions in project `project`
+    Get the interactions in project `project` for gene and exons
 
     :param cnx: Connection to chia-pet database
     :param project: The project of interest
@@ -78,22 +79,27 @@ def get_project_colocalisation(cnx: sqlite3.Connection, project: str,
     :param exon_bc: exons in big communities
     :param inter_chr: True to only get inter-chromosomal interactions \
     False else
+    :param the feature that we want to recover
     :return: The table containing the number of interaction by projects
     """
     logging.debug(f'Recovering interaction ({os.getpid()})')
     select_add = get_select_addition(global_weight, get_weight, same_gene)
     inter_str = " AND level = 'inter'" if inter_chr else ""
     inter_str_t = " AND t1.level = 'inter'" if inter_chr else ""
+    if level == "gene" and not same_gene:
+        logging.warning(f"When level is equal to gene then same_gene must "
+                        f"be true. Setting it to True.")
+        same_gene = True
     if exon_bc is None:
         filter_exons = filter_exons_t = ''
     else:
         tmp = tuple(exon_bc)
-        filter_exons = f' AND exon1 IN {tmp} AND exon2 IN {tmp}'
-        filter_exons_t = f' AND t1.exon1 IN {tmp} AND t1.exon2 IN {tmp}'
+        filter_exons = f' AND {level}1 IN {tmp} AND {level}2 IN {tmp}'
+        filter_exons_t = f' AND t1.{level}1 IN {tmp} AND t1.{level}2 IN {tmp}'
     if global_weight == 0:
         if same_gene:
-            query = f"SELECT exon1, exon2{select_add} " \
-                    "FROM cin_exon_interaction " \
+            query = f"SELECT {level}1, {level}2{select_add} " \
+                    f"FROM cin_{level}_interaction " \
                     f"WHERE weight >= {weight} " \
                     f"AND id_project = '{project}'" \
                     + filter_exons \
@@ -111,13 +117,13 @@ def get_project_colocalisation(cnx: sqlite3.Connection, project: str,
     else:
         good_projects = tuple(ConfigNt.good_projects)
         if same_gene:
-            query = f"SELECT exon1, exon2{select_add} " \
-                    f"FROM cin_exon_interaction " \
+            query = f"SELECT {level}1, {level}2{select_add} " \
+                    f"FROM cin_{level}_interaction " \
                     f"WHERE weight >= {weight} " \
                     f"AND id_project IN {good_projects} " \
                     f"{filter_exons} " \
                     f"{inter_str} " \
-                    f"GROUP BY exon1, exon2 " \
+                    f"GROUP BY {level}1, {level}2 " \
                     f"HAVING COUNT(*) >= {global_weight}"
         else:
             query = f"""SELECT t1.exon1, t1.exon2{select_add}
@@ -134,7 +140,7 @@ def get_project_colocalisation(cnx: sqlite3.Connection, project: str,
 
     c = cnx.cursor()
     c.execute(query)
-    result = np.asarray(c.fetchall())
+    result = np.asarray(c.fetchall(), dtype=str)
     if len(result) == 0:
         msg = f'No interaction found in project {project}'
         logging.exception(msg)