Add a function for reading a fasta into a dict of numpy arrays

allel/io/fasta.py

@@ -44,5 +44,47 @@ def write_fasta(path, sequences, names, mode='w', width=80):
             header = b'>' + name + b'\n'
             fasta.write(header)
             for i in range(0, sequence.size, width):
-                line = sequence[i:i+width].tostring() + b'\n'
+                line = sequence[i:i + width].tostring() + b'\n'
                 fasta.write(line)
+
+
+def read_fasta(path):
+    """Read nucleotide sequences from a FASTA file and return them as a dictionary
+    mapping a sequence name to a sequence.
+
+    Parameters
+    ----------
+
+    path : string
+        File path.
+
+    """
+    sequences = []
+    names = []
+
+    current_sequence = []
+
+    with open(path, 'r') as fasta:
+
+        # Skip any headers and get the first line
+        for line in fasta:
+            if line.startswith('>'):
+                name = line[1:].rstrip()
+                break
+
+        for line in fasta:
+            if line.startswith('>'):
+                names.append(name)
+                sequences.append(np.concatenate(current_sequence))
+                name = line[1:].rstrip()
+                current_sequence = []
+            elif line.startswith(';'):
+                continue
+            else:
+                current_sequence.append(np.array(list(line.strip()), dtype=np.dtype('S1')))
+        else:
+            if names:
+                names.append(name)
+                sequences.append(np.concatenate(current_sequence))
+
+    return dict(zip(names, sequences))

allel/test/data/test.fa

@@ -0,0 +1,45 @@
+>crab_anapl ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDITIHNPLIRRPLFSWLAPSRIFDQIFGEHLQESELLPASPSLSPFLMR
+SPIFRMPSWLETGLSEMRLEKDKFSVNLDVKHFSPEELKVKVLGDMVEIH
+GKHEERQDEHGFIAREFNRKYRIPADVDPLTITSSLSLDGVLTVSAPRKQ
+SDVPERSIPITREEKPAIAGAQRK
+>crab_bovin ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDIAIHHPWIRRPFFPFHSPSRLFDQFFGEHLLESDLFPASTSLSPFYLR
+PPSFLRAPSWIDTGLSEMRLEKDRFSVNLDVKHFSPEELKVKVLGDVIEV
+HGKHEERQDEHGFISREFHRKYRIPADVDPLAITSSLSSDGVLTVNGPRK
+QASGPERTIPITREEKPAVTAAPKK
+>crab_chick ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDITIHNPLVRRPLFSWLTPSRIFDQIFGEHLQESELLPTSPSLSPFLMR
+SPFFRMPSWLETGLSEMRLEKDKFSVNLDVKHFSPEELKVKVLGDMIEIH
+GKHEERQDEHGFIAREFSRKYRIPADVDPLTITSSLSLDGVLTVSAPRKQ
+SDVPERSIPITREEKPAIAGSQRK
+>crab_human ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDIAIHHPWIRRPFFPFHSPSRLFDQFFGEHLLESDLFPTSTSLSPFYLR
+PPSFLRAPSWFDTGLSEMRLEKDRFSVNLDVKHFSPEELKVKVLGDVIEV
+HGKHEERQDEHGFISREFHRKYRIPADVDPLTITSSLSSDGVLTVNGPRK
+QVSGPERTIPITREEKPAVTAAPKK
+>crab_mesau ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDIAIHHPWIRRPFFPFHSPSRLFDQFFGEHLLESDLFSTATSLSPFYLR
+PPSFLRAPSWIDTGLSEMRMEKDRFSVNLDVKHFSPEELKVKVLGDVVEV
+HGKHEERQDEHGFISREFHRKYRIPADVDPLTITSSLSSDGVLTVNGPRK
+QASGPERTIPITREEKPAVTAAPKK
+>crab_mouse ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN) (P23).
+MDIAIHHPWIRRPFFPFHSPSRLFDQFFGEHLLESDLFSTATSLSPFYLR
+PPSFLRAPSWIDTGLSEMRLEKDRFSVNLDVKHFSPEELKVKVLGDVIEV
+HGKHEERQDEHGFISREFHRKYRIPADVDPLAITSSLSSDGVLTVNGPRK
+QVSGPERTIPITREEKPAVAAAPKK
+>crab_rabit ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDIAIHHPWIRRPFFPFHSPSRLFDQFFGEHLLESDLFPTSTSLSPFYLR
+PPSFLRAPSWIDTGLSEMRLEKDRFSVNLDVKHFSPEELKVKVLGDVIEV
+HGKHEERQDEHGFISREFHRKYRIPADVDPLTITSSLSSDGVLTVNGPRK
+QAPGPERTIPITREEKPAVTAAPKK
+>crab_rat ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDIAIHHPWIRRPFFPFHSPSRLFDQFFGEHLLESDLFSTATSLSPFYLR
+PPSFLRAPSWIDTGLSEMRMEKDRFSVNLDVKHFSPEELKVKVLGDVIEV
+HGKHEERQDEHGFISREFHRKYRIPADVDPLTITSSLSSDGVLTVNGPRK
+QASGPERTIPITREEKPAVTAAPKK
+>crab_squac ALPHA CRYSTALLIN B CHAIN (ALPHA(B)-CRYSTALLIN).
+MDIAIQHPWLRRPLFPSSIFPSRIFDQNFGEHFDPDLFPSFSSMLSPFYW
+RMGAPMARMPSWAQTGLSELRLDKDKFAIHLDVKHFTPEELRVKILGDFI
+EVQAQHEERQDEHGYVSREFHRKYKVPAGVDPLVITCSLSADGVLTITGP
+RKVADVPERSVPISRDEKPAVAGPQQK

allel/test/data/test_comments.fa

@@ -0,0 +1,10 @@
+;Comment line
+>sequence A
+ggtaagtcctctagtacaaacacccccaatattgtgatataattaaaattatattcatat
+tctgttgccagaaaaaacacttttaggctatattagagccatcttctttgaagcgttgtc
+;Multiple comment lines
+;Another comment line
+>sequence B
+ggtaagtgctctagtacaaacacccccaatattgtgatataattaaaattatattcatat
+tctgttgccagattttacacttttaggctatattagagccatcttctttgaagcgttgtc
+tatgcatcgatcgacgactg

allel/test/io/test_fasta_read.py

@@ -0,0 +1,21 @@
+import os
+
+from allel.io.fasta import read_fasta
+
+
+def fixture_path(fn):
+    return os.path.join(os.path.dirname(__file__), os.pardir, 'data', fn)
+
+
+def test_fasta_read():
+    path = fixture_path('test.fa')
+    data = read_fasta(path)
+    assert len(data.keys()) == 9
+
+
+def test_fasta_read_comments():
+    path = fixture_path('test_comments.fa')
+    data = read_fasta(path)
+    assert len(data.keys()) == 2
+    assert 'sequence A' in data.keys()
+    assert 'sequence B' in data.keys()