feat: build catalog and add qc report (#344) (#364)

Author: hunterckx

catalog/build/intermediate/genomes-from-ncbi.tsv

@@ -6,6 +6,8 @@
 
 GENOMES_OUTPUT_PATH = "catalog/build/intermediate/genomes-from-ncbi.tsv"
 
+QC_REPORT_PATH = "catalog/output/qc-report.md"
+
 TAXONOMIC_GROUPS_BY_TAXONOMY_ID = {
   2: "Bacteria",
   10239: "Viruses",
@@ -41,5 +43,6 @@
     GENOMES_OUTPUT_PATH,
     UCSC_ASSEMBLIES_URL,
     TAXANOMIC_LEVELS_FOR_TREE,
-    {"taxonomicGroup": TAXONOMIC_GROUPS_BY_TAXONOMY_ID}
+    {"taxonomicGroup": TAXONOMIC_GROUPS_BY_TAXONOMY_ID},
+    qc_report_path=QC_REPORT_PATH
   )

catalog/build/py/build-files-from-ncbi.py

@@ -7,6 +7,8 @@
 from bs4 import BeautifulSoup
 import logging
 
+MAX_NCBI_URL_LENGTH = 2000 # The actual limit seems to be a bit over 4000
+
 log = logging.getLogger(__name__)
 
 
@@ -31,6 +33,37 @@ def get_paginated_ncbi_results(base_url, query_description):
   return results
 
 
+def get_next_ncbi_url_batch(get_url, items, start_index):
+  # Do a binary search to find the longest possible URL
+  max_valid_end_index = start_index
+  min_invalid_end_index = len(items)
+  end_index = len(items)
+  while end_index != max_valid_end_index:
+    test_url = get_url(items[start_index:end_index])
+    if len(test_url) > MAX_NCBI_URL_LENGTH:
+      min_invalid_end_index = end_index
+    else:
+      max_valid_end_index = end_index
+    end_index = int((max_valid_end_index + min_invalid_end_index)/2)
+  return get_url(items[start_index:end_index]), end_index
+
+
+def get_batched_ncbi_urls(get_url, items):
+  urls = []
+  index = 0
+  while index < len(items):
+    url, index = get_next_ncbi_url_batch(get_url, items, index)
+    urls.append(url)
+  return urls
+
+
+def get_batched_ncbi_results(get_base_url, items, query_description):
+  results = []
+  for batch_index, batch_url in enumerate(get_batched_ncbi_urls(get_base_url, items)):
+    results.extend(get_paginated_ncbi_results(batch_url, f"{query_description} batch {batch_index + 1}"))
+  return results
+
+
 def match_taxonomic_group(tax_id, lineage, taxonomic_groups):
   if tax_id not in taxonomic_groups:
     return None
@@ -74,7 +107,11 @@ def get_species_row(taxon_info, taxonomic_group_sets, taxonomic_levels):
 
 
 def get_species_df(taxonomy_ids, taxonomic_group_sets, taxonomic_levels):
-  species_info = get_paginated_ncbi_results(f"https://api.ncbi.nlm.nih.gov/datasets/v2/taxonomy/taxon/{",".join([str(id) for id in taxonomy_ids])}/dataset_report", "taxa")
+  species_info = get_batched_ncbi_results(
+    lambda ids: f"https://api.ncbi.nlm.nih.gov/datasets/v2/taxonomy/taxon/{",".join(ids)}/dataset_report",
+    [str(id) for id in taxonomy_ids],
+    "taxa"
+  )
   return pd.DataFrame([get_species_row(info, taxonomic_group_sets, taxonomic_levels) for info in species_info])
 
 
@@ -107,20 +144,15 @@ def get_biosample_data(genome_info):
 
 
 def get_genomes_and_primarydata_df(accessions):
-  # Send accessions in batches of 100
-  # ncbi api produces a 414 error if there are too many accessions
-  batch_size = 100
-  genomes_info = []
-  biosamples_info = []
-  for i in range(0, len(accessions), batch_size):
-    batch = accessions[i:i+batch_size]
-    batch_genomes_info = get_paginated_ncbi_results(f"https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/{','.join(batch)}/dataset_report", "genomes")
-    genomes_info.extend([get_genome_row(info) for info in batch_genomes_info])
-    biosamples_info.extend([get_biosample_data(info) for info in batch_genomes_info if 'biosample' in info['assembly_info']])
+  genomes_info = get_batched_ncbi_results(
+    lambda a: f"https://api.ncbi.nlm.nih.gov/datasets/v2/genome/accession/{",".join(a)}/dataset_report",
+    accessions,
+    "genomes"
+  )
 
   return (
-          pd.DataFrame(data=genomes_info),
-          pd.DataFrame(data=biosamples_info))
+          pd.DataFrame(data=[get_genome_row(info) for info in genomes_info]),
+          pd.DataFrame(data=[get_biosample_data(info) for info in genomes_info if 'biosample' in info['assembly_info']]))
 
 
 def _id_to_gene_model_url(asm_id):
@@ -161,7 +193,7 @@ def report_missing_values_from(values_name, message_predicate, all_values_series
   present_values_mask[:] = False
   for series in partial_values_series:
     present_values_mask |= all_values_series.isin(series)
-  report_missing_values(values_name, message_predicate, all_values_series, present_values_mask)
+  return report_missing_values(values_name, message_predicate, all_values_series, present_values_mask)
 
 
 def report_missing_values(values_name, message_predicate, values_series, present_values_mask):
@@ -172,6 +204,15 @@ def report_missing_values(values_name, message_predicate, values_series, present
       print(f"Only {len(present_values)} of {len(values_series)} {values_name} {message_predicate}: {", ".join(present_values)}")
     else:
       print(f"{len(missing_values)} {values_name} not {message_predicate}: {", ".join(missing_values)}")
+  return missing_values
+
+
+def report_inconsistent_taxonomy_ids(df):
+  inconsistent_ids_series = df.groupby(["species", "strain"]).filter(lambda g: g["taxonomyId"].nunique() > 1).groupby(["species", "strain"])["taxonomyId"].apply(set)
+  inconsistent_ids_strings = [(f"{species} strain {strain}" if strain else species, ", ".join([str(id) for id in ids])) for (species, strain), ids in inconsistent_ids_series.items()]
+  if len(inconsistent_ids_strings) > 0:
+    print(f"Taxa with inconsistent taxonomy IDs: {", ".join([f"{taxon} ({ids})" for taxon, ids in inconsistent_ids_strings])}")
+  return inconsistent_ids_strings
 
 
 def fetch_sra_metadata(srs_ids, batch_size=20):
@@ -344,6 +385,14 @@ def fetch_url_data(url, counter=0, counter_limit=2, wait_time=2, num_retry=3):
   return data
 
 
+def make_qc_report(missing_ncbi_assemblies, inconsistent_taxonomy_ids, missing_ucsc_assemblies, missing_gene_model_urls=None):
+  ncbi_assemblies_text = "None" if len(missing_ncbi_assemblies) == 0 else "\n".join([f"- {accession}" for accession in missing_ncbi_assemblies])
+  ucsc_assemblies_text = "None" if len(missing_ucsc_assemblies) == 0 else "\n".join([f"- {accession}" for accession in missing_ucsc_assemblies])
+  gene_model_urls_text = "N/A" if missing_gene_model_urls is None else "None" if len(missing_gene_model_urls) == 0 else "\n".join([f"- {accession}" for accession in missing_gene_model_urls])
+  taxonomy_ids_text = "None" if len(inconsistent_taxonomy_ids) == 0 else "\n".join([f"- {taxon}: {ids}" for taxon, ids in inconsistent_taxonomy_ids])
+  return f"# Catalog QC report\n\n## Assemblies not found on NCBI\n\n{ncbi_assemblies_text}\n\n## Assemblies not found in UCSC list\n\n{ucsc_assemblies_text}\n\n## Assemblies with gene model URLs not found\n\n{gene_model_urls_text}\n\n## Species and strain combinations with multiple taxonomy IDs\n\n{taxonomy_ids_text}\n"
+
+
 def build_files(
   assemblies_path,
   genomes_output_path,
@@ -352,10 +401,13 @@ def build_files(
   taxonomic_group_sets={},
   do_gene_model_urls=True,
   extract_primary_data=False,
-  primary_output_path=None
+  primary_output_path=None,
+  qc_report_path=None
 ):
   print("Building files")
 
+  qc_report_params = {}
+
   source_list_df = read_assemblies(assemblies_path)
 
   base_genomes_df, primarydata_df = get_genomes_and_primarydata_df(source_list_df["accession"])
@@ -369,26 +421,29 @@ def build_files(
     sra_metadata = fetch_sra_metadata(sra_ids_list)
     sra_metadata_df = pd.DataFrame([sra_metadata[sra][srr] for sra in sra_metadata for srr in sra_metadata[sra]])
     primarydata_df = primarydata_df.merge(sra_metadata_df, how="left", left_on="sra_sample_acc", right_on="sra_sample_acc")
-  report_missing_values_from("accessions", "found on NCBI", source_list_df["accession"], base_genomes_df["accession"])
+  
+  qc_report_params["missing_ncbi_assemblies"] = report_missing_values_from("accessions", "found on NCBI", source_list_df["accession"], base_genomes_df["accession"])
 
   species_df = get_species_df(base_genomes_df["taxonomyId"], taxonomic_group_sets, taxonomic_levels_for_tree)
 
   report_missing_values_from("species", "found on NCBI", base_genomes_df["taxonomyId"], species_df["taxonomyId"])
 
   genomes_with_species_df = base_genomes_df.merge(species_df, how="left", on="taxonomyId")
 
+  qc_report_params["inconsistent_taxonomy_ids"] = report_inconsistent_taxonomy_ids(genomes_with_species_df)
+
   assemblies_df = pd.DataFrame(requests.get(ucsc_assemblies_url).json()["data"])[["ucscBrowser", "genBank", "refSeq"]]
 
   gen_bank_merge_df = genomes_with_species_df.merge(assemblies_df, how="left", left_on="accession", right_on="genBank")
   ref_seq_merge_df = genomes_with_species_df.merge(assemblies_df, how="left", left_on="accession", right_on="refSeq")
 
-  report_missing_values_from("accessions", "matched in assembly list", genomes_with_species_df["accession"], assemblies_df["genBank"], assemblies_df["refSeq"])
+  qc_report_params["missing_ucsc_assemblies"] = report_missing_values_from("accessions", "matched in assembly list", genomes_with_species_df["accession"], assemblies_df["genBank"], assemblies_df["refSeq"])
 
   genomes_df = gen_bank_merge_df.combine_first(ref_seq_merge_df)
 
   if do_gene_model_urls:
     genomes_df = add_gene_model_url(genomes_df)
-    report_missing_values("accessions", "matched with gene model URLs", genomes_df["accession"], genomes_df["geneModelUrl"].astype(bool))
+    qc_report_params["missing_gene_model_urls"] = report_missing_values("accessions", "matched with gene model URLs", genomes_df["accession"], genomes_df["geneModelUrl"].astype(bool))
   else:
     genomes_df["geneModelUrl"] = ""
 
@@ -400,3 +455,8 @@ def build_files(
     primarydata_df.to_csv(primary_output_path, index=False, sep="\t")
 
     print(f"Wrote to {primary_output_path}")
+  
+  if qc_report_path is not None:
+    qc_report_text = make_qc_report(**qc_report_params)
+    with open(qc_report_path, "w") as file:
+      file.write(qc_report_text)

catalog/build/py/package/catalog_build/build.py

@@ -2,7 +2,7 @@
 
 setup(
   name="catalog_build",
-  version="1.3.1",
+  version="1.4.0",
   packages=["catalog_build"],
   install_requires=["pandas", "requests", "PyYAML", "BeautifulSoup4", "lxml"],
 )