diff --git a/src/plassembler/utils/plass_class.py b/src/plassembler/utils/plass_class.py index b8e7318..f19cc21 100644 --- a/src/plassembler/utils/plass_class.py +++ b/src/plassembler/utils/plass_class.py @@ -17,7 +17,11 @@ get_depths_from_bam, ) from plassembler.utils.mapping import minimap_long_reads, minimap_short_reads -from plassembler.utils.run_mash import get_contig_count, is_file_empty +from plassembler.utils.run_mash import ( + get_contig_count, + load_plsdb_metadata, + mash_tophits_df, +) class Plass: @@ -520,173 +524,12 @@ def process_mash_tsv(self, plassembler_db_dir): contig_count = get_contig_count(os.path.join(outdir, "assembly.fasta")) # update with final plasmid count number self.contig_count = contig_count - # get mash tsv output contig - mash_tsv = os.path.join(outdir, "mash.tsv") - col_list = [ - "contig", - "NUCCORE_ACC", - "mash_distance", - "mash_pval", - "mash_matching_hashes", - ] - - # check if mash tsv file is empty -> no hits - mash_empty = is_file_empty(mash_tsv) - # instantiate tophits list - tophits_mash_df = [] - if mash_empty is False: - mash_df = pd.read_csv( - mash_tsv, delimiter="\t", index_col=False, names=col_list - ) - # get list of contigs from unicycler: 1 -> total number of contigs - contigs = range(1, contig_count + 1) - - # instantiate tophits list - tophits = [] - - for contig in contigs: - hit_df = ( - mash_df.loc[mash_df["contig"] == contig] - .sort_values("mash_distance") - .reset_index(drop=True) - ) - hits = len(hit_df["mash_distance"]) - # add only if there is a hit - if hits > 0: - tmp_df = ( - mash_df.loc[mash_df["contig"] == contig] - .sort_values("mash_distance") - .reset_index(drop=True) - .loc[0] - ) - tophits.append( - [ - tmp_df.contig, - "Yes", - tmp_df.NUCCORE_ACC, - tmp_df.mash_distance, - tmp_df.mash_pval, - tmp_df.mash_matching_hashes, - ] - ) - else: # no hits append no it - tophits.append([contig, "", "", "", "", ""]) - # create tophits df - tophits_mash_df = pd.DataFrame( - tophits, - columns=[ - "contig", - "PLSDB_hit", - "NUCCORE_ACC", - "mash_distance", - "mash_pval", - "mash_matching_hashes", - ], - ) - - else: # empty mash file - contigs = range(1, contig_count + 1) - # create empty df - tophits_mash_df = pd.DataFrame( - columns=[ - "contig", - "PLSDB_hit", - "NUCCORE_ACC", - "mash_distance", - "mash_pval", - "mash_matching_hashes", - ] - ) - for contig in contigs: - tophits_mash_df.loc[contig - 1] = [contig, "", "", "", "", ""] - - # read in the plasdb tsv to get the description - plsdb_tsv_file = os.path.join(plassembler_db_dir, "plsdb_2023_11_03_v2.tsv") - cols = [ - "NUCCORE_UID", - "NUCCORE_ACC", - "NUCCORE_Description", - "NUCCORE_CreateDate", - "NUCCORE_Topology", - "NUCCORE_Completeness", - "NUCCORE_TaxonID", - "NUCCORE_Genome", - "NUCCORE_Length", - "NUCCORE_DuplicatedEntry", - "NUCCORE_Source", - "NUCCORE_BiosampleID", - "BIOSAMPLE_UID", - "BIOSAMPLE_ACC", - "BIOSAMPLE_Location", - "BIOSAMPLE_Coordinates", - "BIOSAMPLE_IsolationSource", - "BIOSAMPLE_Host", - "BIOSAMPLE_CollectionDate", - "BIOSAMPLE_HostDisease", - "BIOSAMPLE_SampleType", - "ASSEMBLY_UID", - "ASSEMBLY_ACC", - "ASSEMBLY_Status", - "ASSEMBLY_coverage", - "ASSEMBLY_SeqReleaseDate", - "ASSEMBLY_SubmissionDate", - "ASSEMBLY_Lastest", - "ASSEMBLY_BiosampleID", - "TAXONOMY_superkingdom", - "TAXONOMY_phylum", - "TAXONOMY_class", - "TAXONOMY_order", - "TAXONOMY_family", - "TAXONOMY_genus", - "TAXONOMY_species", - "TAXONOMY_strain", - "TAXONOMY_UID", - "TAXONOMY_taxon_rank", - "TAXONOMY_taxon_name", - "TAXONOMY_taxon_lineage", - "TAXONOMY_superkingdom_id", - "TAXONOMY_phylum_id", - "TAXONOMY_class_id", - "TAXONOMY_order_id", - "TAXONOMY_family_id", - "TAXONOMY_genus_id", - "TAXONOMY_species_id", - "TAXONOMY_strain_id", - "has_biosample", - "has_assembly", - "has_location", - "rMLST_hits", - "rMLST_hitscount", - "inclusions", - "NUCCORE_GC", - "Length", - "BIOSAMPLE_Host_processed", - "BIOSAMPLE_Host_processed_source", - "BIOSAMPLE_Host_label", - "BIOSAMPLE_HostDisease_processed", - "loc_lat", - "loc_lng", - "loc_parsed", - "D1", - "D2", - "plasmidfinder", - "pmlst", - ] - - plsdb_tsv = pd.read_csv( - plsdb_tsv_file, - delimiter="\t", - index_col=False, - names=cols, - skiprows=1, - low_memory=False, - ) - combined_mash_df = tophits_mash_df.merge( - plsdb_tsv, on="NUCCORE_ACC", how="left" + tophits_mash_df = mash_tophits_df( + os.path.join(outdir, "mash.tsv"), contig_count ) - - self.mash_df = combined_mash_df + plsdb_tsv = load_plsdb_metadata(plassembler_db_dir) + self.mash_df = tophits_mash_df.merge(plsdb_tsv, on="NUCCORE_ACC", how="left") def combine_depth_mash_tsvs(self, prefix, depth_filter, skip_mash): """ @@ -1040,174 +883,11 @@ def process_mash_tsv(self, plassembler_db_dir, plasmid_fasta): # update with final plasmid count number self.contig_count = contig_count - # get mash tsv output contig - mash_tsv = os.path.join(outdir, "mash.tsv") - col_list = [ - "contig", - "NUCCORE_ACC", - "mash_distance", - "mash_pval", - "mash_matching_hashes", - ] - - # check if mash tsv file is empty -> no hits - mash_empty = is_file_empty(mash_tsv) - # instantiate tophits list - tophits_mash_df = [] - - if mash_empty is False: - mash_df = pd.read_csv( - mash_tsv, delimiter="\t", index_col=False, names=col_list - ) - # get list of contigs from unicycler: 1 -> total number of contigs - contigs = range(1, contig_count + 1) - - # instantiate tophits list - tophits = [] - - for contig in contigs: - hit_df = ( - mash_df.loc[mash_df["contig"] == contig] - .sort_values("mash_distance") - .reset_index(drop=True) - ) - hits = len(hit_df["mash_distance"]) - # add only if there is a hit - if hits > 0: - tmp_df = ( - mash_df.loc[mash_df["contig"] == contig] - .sort_values("mash_distance") - .reset_index(drop=True) - .loc[0] - ) - tophits.append( - [ - tmp_df.contig, - "Yes", - tmp_df.NUCCORE_ACC, - tmp_df.mash_distance, - tmp_df.mash_pval, - tmp_df.mash_matching_hashes, - ] - ) - else: # no hits append no it - tophits.append([contig, "", "", "", "", ""]) - # create tophits df - tophits_mash_df = pd.DataFrame( - tophits, - columns=[ - "contig", - "PLSDB_hit", - "NUCCORE_ACC", - "mash_distance", - "mash_pval", - "mash_matching_hashes", - ], - ) - - else: # empty mash file - contigs = range(1, contig_count + 1) - # create empty df - tophits_mash_df = pd.DataFrame( - columns=[ - "contig", - "PLSDB_hit", - "NUCCORE_ACC", - "mash_distance", - "mash_pval", - "mash_matching_hashes", - ] - ) - for contig in contigs: - tophits_mash_df.loc[contig - 1] = [contig, "", "", "", "", ""] - - # read in the plasdb tsv to get the description - plsdb_tsv_file = os.path.join(plassembler_db_dir, "plsdb_2023_11_03_v2.tsv") - - cols = [ - "NUCCORE_UID", - "NUCCORE_ACC", - "NUCCORE_Description", - "NUCCORE_CreateDate", - "NUCCORE_Topology", - "NUCCORE_Completeness", - "NUCCORE_TaxonID", - "NUCCORE_Genome", - "NUCCORE_Length", - "NUCCORE_DuplicatedEntry", - "NUCCORE_Source", - "NUCCORE_BiosampleID", - "BIOSAMPLE_UID", - "BIOSAMPLE_ACC", - "BIOSAMPLE_Location", - "BIOSAMPLE_Coordinates", - "BIOSAMPLE_IsolationSource", - "BIOSAMPLE_Host", - "BIOSAMPLE_CollectionDate", - "BIOSAMPLE_HostDisease", - "BIOSAMPLE_SampleType", - "ASSEMBLY_UID", - "ASSEMBLY_ACC", - "ASSEMBLY_Status", - "ASSEMBLY_coverage", - "ASSEMBLY_SeqReleaseDate", - "ASSEMBLY_SubmissionDate", - "ASSEMBLY_Lastest", - "ASSEMBLY_BiosampleID", - "TAXONOMY_superkingdom", - "TAXONOMY_phylum", - "TAXONOMY_class", - "TAXONOMY_order", - "TAXONOMY_family", - "TAXONOMY_genus", - "TAXONOMY_species", - "TAXONOMY_strain", - "TAXONOMY_UID", - "TAXONOMY_taxon_rank", - "TAXONOMY_taxon_name", - "TAXONOMY_taxon_lineage", - "TAXONOMY_superkingdom_id", - "TAXONOMY_phylum_id", - "TAXONOMY_class_id", - "TAXONOMY_order_id", - "TAXONOMY_family_id", - "TAXONOMY_genus_id", - "TAXONOMY_species_id", - "TAXONOMY_strain_id", - "has_biosample", - "has_assembly", - "has_location", - "rMLST_hits", - "rMLST_hitscount", - "inclusions", - "NUCCORE_GC", - "Length", - "BIOSAMPLE_Host_processed", - "BIOSAMPLE_Host_processed_source", - "BIOSAMPLE_Host_label", - "BIOSAMPLE_HostDisease_processed", - "loc_lat", - "loc_lng", - "loc_parsed", - "D1", - "D2", - "plasmidfinder", - "pmlst", - ] - - plsdb_tsv = pd.read_csv( - plsdb_tsv_file, - delimiter="\t", - index_col=False, - names=cols, - skiprows=1, - low_memory=False, - ) - combined_mash_df = tophits_mash_df.merge( - plsdb_tsv, on="NUCCORE_ACC", how="left" + tophits_mash_df = mash_tophits_df( + os.path.join(outdir, "mash.tsv"), contig_count ) - - self.mash_df = combined_mash_df + plsdb_tsv = load_plsdb_metadata(plassembler_db_dir) + self.mash_df = tophits_mash_df.merge(plsdb_tsv, on="NUCCORE_ACC", how="left") def combine_depth_mash_tsvs(self, prefix, no_copy_numbers): """ diff --git a/src/plassembler/utils/run_mash.py b/src/plassembler/utils/run_mash.py index 5988725..c310eb7 100644 --- a/src/plassembler/utils/run_mash.py +++ b/src/plassembler/utils/run_mash.py @@ -2,10 +2,110 @@ import shutil from pathlib import Path +import pandas as pd from Bio import SeqIO from plassembler.utils.external_tools import ExternalTool +MASH_COL_LIST = [ + "contig", + "NUCCORE_ACC", + "mash_distance", + "mash_pval", + "mash_matching_hashes", +] + +TOPHITS_COLUMNS = [ + "contig", + "PLSDB_hit", + "NUCCORE_ACC", + "mash_distance", + "mash_pval", + "mash_matching_hashes", +] + +# the PLSDB metadata tsv has no header row of its own +PLSDB_COLUMNS = [ + "NUCCORE_UID", + "NUCCORE_ACC", + "NUCCORE_Description", + "NUCCORE_CreateDate", + "NUCCORE_Topology", + "NUCCORE_Completeness", + "NUCCORE_TaxonID", + "NUCCORE_Genome", + "NUCCORE_Length", + "NUCCORE_DuplicatedEntry", + "NUCCORE_Source", + "NUCCORE_BiosampleID", + "BIOSAMPLE_UID", + "BIOSAMPLE_ACC", + "BIOSAMPLE_Location", + "BIOSAMPLE_Coordinates", + "BIOSAMPLE_IsolationSource", + "BIOSAMPLE_Host", + "BIOSAMPLE_CollectionDate", + "BIOSAMPLE_HostDisease", + "BIOSAMPLE_SampleType", + "ASSEMBLY_UID", + "ASSEMBLY_ACC", + "ASSEMBLY_Status", + "ASSEMBLY_coverage", + "ASSEMBLY_SeqReleaseDate", + "ASSEMBLY_SubmissionDate", + "ASSEMBLY_Lastest", + "ASSEMBLY_BiosampleID", + "TAXONOMY_superkingdom", + "TAXONOMY_phylum", + "TAXONOMY_class", + "TAXONOMY_order", + "TAXONOMY_family", + "TAXONOMY_genus", + "TAXONOMY_species", + "TAXONOMY_strain", + "TAXONOMY_UID", + "TAXONOMY_taxon_rank", + "TAXONOMY_taxon_name", + "TAXONOMY_taxon_lineage", + "TAXONOMY_superkingdom_id", + "TAXONOMY_phylum_id", + "TAXONOMY_class_id", + "TAXONOMY_order_id", + "TAXONOMY_family_id", + "TAXONOMY_genus_id", + "TAXONOMY_species_id", + "TAXONOMY_strain_id", + "has_biosample", + "has_assembly", + "has_location", + "rMLST_hits", + "rMLST_hitscount", + "inclusions", + "NUCCORE_GC", + "Length", + "BIOSAMPLE_Host_processed", + "BIOSAMPLE_Host_processed_source", + "BIOSAMPLE_Host_label", + "BIOSAMPLE_HostDisease_processed", + "loc_lat", + "loc_lng", + "loc_parsed", + "D1", + "D2", + "plasmidfinder", + "pmlst", +] + +# NOTE on memory: this table is read whole (~240 MB of dataframe for the v1.5.0 +# database) to serve a merge that keeps at most one row per contig, so reading it +# in chunks and discarding non-matching rows is tempting. Do not: pandas infers +# column dtypes per chunk, and several PLSDB columns are only mixed-type when the +# whole file is seen. ASSEMBLY_coverage, for instance, is object over the full +# table and writes "120", but is float64 within a small chunk and writes "120.0". +# The merged frame goes verbatim into _summary.tsv, so that is a +# user-visible output change. Measured at 0.5s and ~240 MB transient, which is +# not worth pinning 68 dtypes to the current database release. + def mash_sketch(out_dir, fasta_file, logdir): """ @@ -81,3 +181,89 @@ def is_file_empty(file): if os.stat(file).st_size == 0: empty = True return empty + + +def mash_tophits_df(mash_tsv, contig_count): + """Best (lowest mash distance) PLSDB hit per contig, one row per contig. + + Contigs 1..contig_count always get a row; those without a hit get blanks. + + The previous implementation filtered and sorted the whole mash dataframe + twice for every contig inside a python loop, so cost grew with + contigs x hits: 400 contigs took 0.283s where this takes 0.006s. It also + lived in two near-identical ~90 line copies, one per class. + + Ties: when several references share the lowest distance - mash distance is a + function of matching hash count, so exact ties do happen - the old code took + whichever row pandas' unstable quicksort happened to place first. The stable + sort below always takes the first such row in mash's output order. For the + handful of hits per contig that `mash dist -d 0.1 -v 0.1` actually returns, + numpy sorts stably anyway, so this matches the old result and merely makes it + reproducible rather than arbitrary. + + :param mash_tsv: mash dist output + :param contig_count: number of contigs in the assembly + :return: dataframe with TOPHITS_COLUMNS + """ + blank_row = ["", "", "", "", ""] + + if is_file_empty(mash_tsv): + return pd.DataFrame( + [[contig] + blank_row for contig in range(1, contig_count + 1)], + columns=TOPHITS_COLUMNS, + ) + + mash_df = pd.read_csv( + mash_tsv, delimiter="\t", index_col=False, names=MASH_COL_LIST + ) + # stable sort keeps the old tie-breaking: among equal distances the first + # row in the mash output wins, which is what .sort_values().loc[0] gave + best = ( + mash_df.sort_values("mash_distance", kind="stable") + .groupby("contig", sort=False) + .first() + ) + + # Rows are assembled from the selected hits rather than produced by a + # left merge against 1..contig_count. A merge would introduce NaN for + # contigs without a hit, and that silently promotes int columns to float: + # a mash_pval of 0 came back out as "0.0" in _summary.tsv. This + # loop is over contigs only - the expensive part, selecting the best hit, + # already happened once above. + rows = [] + for contig in range(1, contig_count + 1): + if contig not in best.index: + rows.append([contig] + blank_row) + continue + hit = best.loc[contig] + rows.append( + [ + contig, + "Yes", + hit.NUCCORE_ACC, + hit.mash_distance, + hit.mash_pval, + hit.mash_matching_hashes, + ] + ) + return pd.DataFrame(rows, columns=TOPHITS_COLUMNS) + + +def load_plsdb_metadata(plassembler_db_dir): + """The PLSDB metadata table, used to describe each contig's mash tophit. + + See the note on PLSDB_COLUMNS above for why this is deliberately a single + whole-file read. + + :param plassembler_db_dir: plassembler database directory + :return: dataframe with PLSDB_COLUMNS + """ + plsdb_tsv_file = os.path.join(plassembler_db_dir, "plsdb_2023_11_03_v2.tsv") + return pd.read_csv( + plsdb_tsv_file, + delimiter="\t", + index_col=False, + names=PLSDB_COLUMNS, + skiprows=1, + low_memory=False, + ) diff --git a/tests/test_run_mash.py b/tests/test_run_mash.py new file mode 100644 index 0000000..d2f5440 --- /dev/null +++ b/tests/test_run_mash.py @@ -0,0 +1,107 @@ +"""Tests for mash tophit selection in plassembler.utils.run_mash.""" + +import pandas as pd +import pytest + +from src.plassembler.utils.run_mash import TOPHITS_COLUMNS, mash_tophits_df + +COL_LIST = [ + "contig", + "NUCCORE_ACC", + "mash_distance", + "mash_pval", + "mash_matching_hashes", +] + + +def write_mash(path, rows): + pd.DataFrame(rows, columns=COL_LIST).to_csv( + path, sep="\t", header=False, index=False + ) + return path + + +def test_tophits_picks_lowest_distance_per_contig(tmp_path): + tsv = write_mash( + tmp_path / "mash.tsv", + [ + [1, "ACC_far", 0.09, 0.0, "10/1000"], + [1, "ACC_near", 0.01, 0.0, "900/1000"], + [2, "ACC_two", 0.05, 0.0, "500/1000"], + ], + ) + df = mash_tophits_df(tsv, 2).set_index("contig") + assert df.loc[1, "NUCCORE_ACC"] == "ACC_near" + assert df.loc[1, "mash_distance"] == 0.01 + assert df.loc[2, "NUCCORE_ACC"] == "ACC_two" + assert (df["PLSDB_hit"] == "Yes").all() + + +def test_tophits_has_a_row_per_contig_including_misses(tmp_path): + """Contigs with no mash hit still get a row, with blanks - downstream code + indexes by contig and would silently drop them otherwise.""" + tsv = write_mash(tmp_path / "mash.tsv", [[2, "ACC", 0.05, 0.0, "500/1000"]]) + df = mash_tophits_df(tsv, 3) + assert list(df["contig"]) == [1, 2, 3] + assert list(df.columns) == TOPHITS_COLUMNS + assert df.set_index("contig").loc[1, "PLSDB_hit"] == "" + assert df.set_index("contig").loc[1, "NUCCORE_ACC"] == "" + assert df.set_index("contig").loc[3, "mash_distance"] == "" + + +def test_tophits_empty_mash_file(tmp_path): + """No hits at all: still one blank row per contig.""" + tsv = tmp_path / "mash.tsv" + tsv.write_text("") + df = mash_tophits_df(tsv, 2) + assert list(df["contig"]) == [1, 2] + assert (df["PLSDB_hit"] == "").all() + assert list(df.columns) == TOPHITS_COLUMNS + + +def test_tophits_tie_broken_by_mash_output_order(tmp_path): + """Equal distances resolve to the first row in mash's output, deterministically. + The old per-contig unstable sort left this arbitrary.""" + tsv = write_mash( + tmp_path / "mash.tsv", + [ + [1, "ACC_first", 0.02, 0.0, "500/1000"], + [1, "ACC_second", 0.02, 0.0, "500/1000"], + [1, "ACC_third", 0.02, 0.0, "500/1000"], + ], + ) + assert mash_tophits_df(tsv, 1).loc[0, "NUCCORE_ACC"] == "ACC_first" + + +@pytest.mark.parametrize("contig_count", [1, 5]) +def test_tophits_row_count_matches_contig_count(tmp_path, contig_count): + """Extra contigs in the mash output must not add rows beyond contig_count.""" + tsv = write_mash( + tmp_path / "mash.tsv", + [[c, f"ACC{c}", 0.01 * c, 0.0, "1/1000"] for c in range(1, 9)], + ) + assert len(mash_tophits_df(tsv, contig_count)) == contig_count + + +def test_integer_mash_pval_is_not_promoted_to_float(tmp_path): + """mash reports pval 0 for a strong hit. Assembling the table with a left + merge against 1..contig_count introduces NaN for contigs without a hit, + which promotes the int column to float and writes "0.0" into + _summary.tsv instead of "0".""" + tsv = write_mash( + tmp_path / "mash.tsv", + [[2, "NZ_CP124653.1", 0.0450347, 0, "241/1000"]], + ) + # contig 1 has no hit, so it is the row that would introduce the NaN + df = mash_tophits_df(tsv, 2) + rendered = df.to_csv(index=False, sep="\t") + assert "\t0\t241/1000" in rendered, rendered + assert "0.0\t241/1000" not in rendered + + +def test_contig_column_stays_integral(tmp_path): + """Same promotion hazard for the contig ids themselves.""" + tsv = write_mash(tmp_path / "mash.tsv", [[3, "ACC", 0.01, 0, "1/1000"]]) + df = mash_tophits_df(tsv, 3) + assert list(df["contig"]) == [1, 2, 3] + assert "1.0" not in df.to_csv(index=False)