slow data partitioning in workflow/make_dataset

[flowers9]

I attempted to run workflow/make_dataset on a set of four moderately sized genomes. This was still running after two days. After looking at the code and trying some variations, it looks like this came down to a simple issue: speed of access to the genome data in gpn/data.py:Genome::get_seq(). And this could be changed in two lines by switching the underlying storage from a string to a numpy view of the raw sequence bytes (although there may be knock-on effects I missed):

--- a/gpn/data.py
+++ b/gpn/data.py
@@ -22,7 +22,7 @@ def load_fasta(path, subset_chroms=None):
     with gzip.open(path, "rt") if path.endswith(".gz") else open(path) as handle:
         genome = pd.Series(
             {
-                rec.id: str(rec.seq)
+                rec.id: np.frombuffer(bytes(rec.seq), dtype="S1")
                 for rec in SeqIO.parse(handle, "fasta")
                 if subset_chroms is None or rec.id in subset_chroms
             }
@@ -95,7 +95,7 @@ class Genome:
 
     def get_seq(self, chrom, start, end, strand="+"):
         chrom_size = self.chrom_sizes[chrom]
-        seq = self._genome[chrom][max(start, 0) : min(end, chrom_size)]
+        seq = str(self._genome[chrom][max(start, 0) : min(end, chrom_size)], encoding="ascii")
 
         if start < 0:
             seq = "N" * (-start) + seq  # left padding

At least on my system, this went from processing 20-30 items/second to ~87k/s, when taking intervals and creating dataset_assemblies.

[flowers9]

Found a spot I missed in data.py:

@@ -138,7 +138,7 @@ class Genome:
                 dict(
                     start=np.where(
                         ~np.isin(
-                            np.frombuffer(chrom.seq.encode("ascii"), dtype="S1"),
+                            chrom.seq,
                             symbols,
                         )
                     )[0]