Merge remote-tracking branch 'origin/coa_analysis' into coa_analysis

# Conflicts:
#	coa_detection/AbfData.py

coa_detection/AbfData.py

@@ -62,9 +62,11 @@ class AbfData:
         plt.hist(event_lengths, bins=100)
         plt.show()
 
-    def get_pos(self, width, unfiltered=False):
-        while True:
-            event = random.choice(self.pos_events)
+    def get_pos(self, width, unfiltered=False, take_one=True):
+        pos_list = []
+        if take_one:
+            random.shuffle(self.pos_events)
+        for event in self.pos_events:
             start_idx = event[0]
             end_idx = event[-1]
             event_length = end_idx - start_idx
@@ -74,18 +76,26 @@ class AbfData:
                 start_idx -= random_offset
                 end_idx += room_left - random_offset
                 assert (end_idx - start_idx) == width
-                if unfiltered:
+                if take_one:
                     return self.unfiltered_raw[start_idx: end_idx]
-                return self.raw[start_idx: end_idx]
+                if unfiltered:
+                    pos_list.append(self.unfiltered_raw[start_idx: end_idx])
+                else:
+                    pos_list.append(self.raw[start_idx: end_idx])
+        return pos_list
 
-    def get_neg(self, width, unfiltered=False):
-        while True:
+    def get_neg(self, width, nb_neg, unfiltered=False):
+        neg_list = []
+        for i in range(nb_neg * 100): # take 100 attempts for each neg
             random_idx = random.randint(0, len(self.raw)-width)
             candidate_indices = np.arange(random_idx, random_idx+width)
             if not np.any(np.isin(candidate_indices, self.flat_pos_indices)):
                 if unfiltered:
-                    return self.unfiltered_raw[candidate_indices]
-                return self.raw[candidate_indices]
+                    neg_list.append(self.unfiltered_raw[candidate_indices])
+                else:
+                    neg_list.append(self.raw[candidate_indices])
+            if len(neg_list) >= nb_neg:
+                return neg_list
 
 
 if __name__ == '__main__':

coa_detection/CoaExampleDb.py

@@ -32,12 +32,11 @@ class ExampleDb(object):
     def target(self, target):
         self._target = target
 
-    def add_training_read(self, training_read, uncenter_kmer):
+    def add_training_read(self, training_read, non_target=None):
         """Add training read with positive and negative read examples of the
         target k-mer
 
         :param training_read: Object containing a training read
-        :type training_read: TrainingRead
         :param uncenter_kmer: If set to true, will extract reads where the k-mer
                               is randomly places somewhere in the read. Not
                               always in the center
@@ -45,28 +44,36 @@ class ExampleDb(object):
         """
         with self._db.transaction() as conn:
             # --- add positive examples (if any) ---
-            pos_examples = training_read.get_pos(self.target, self.width, uncenter_kmer)
-            for i, ex in enumerate(pos_examples):
-                conn.root.pos[len(conn.root.pos)] = ex
+            pos_examples = training_read.get_pos(self.width)
+            if not non_target:
+                for i, ex in enumerate(pos_examples):
+                    conn.root.pos[len(conn.root.pos)] = ex
+            else:
+                self.nb_neg += self.add_non_target(pos_examples, non_target, conn)
             nb_new_positives = len(pos_examples)
+
             # --- update record nb positive examples ---
             if self._db_empty:
                 if nb_new_positives > 0:
                     self._db_empty = False
             if not self._db_empty:
                 self.nb_pos = conn.root.pos.maxKey()
+
             # --- add negative examples ---
-            neg_examples, neg_kmers = training_read.get_neg(self.target, self.width, len(pos_examples) * 10)  # arbitrarily adding 10x as many neg examples
-            for i, ex in enumerate(neg_examples):
-                if neg_kmers[i] in self.neg_kmers:  # CL: kept the name 'kmer', even though we're not detecting k-mers anymore
-                    self.neg_kmers[neg_kmers[i]].append(self.nb_neg + i)
-                else:
-                    self.neg_kmers[neg_kmers[i]] = [self.nb_neg + i]
-                conn.root.neg[len(conn.root.neg)] = ex
-            self.nb_neg += len(neg_examples)
+            neg_examples = training_read.get_neg(self.width, len(pos_examples))  # add as many pos examples
+            self.nb_neg += self.add_non_target(neg_examples, 'bg', conn)
             conn.root.neg_kmers = self.neg_kmers
             return nb_new_positives
 
+    def add_non_target(self, examples, non_target_name, conn):
+        for i, ex in enumerate(examples):
+            if non_target_name in self.neg_kmers:  # CL: kept the name 'kmer', even though we're not detecting k-mers anymore
+                self.neg_kmers[non_target_name].append(self.nb_neg + i)
+            else:
+                self.neg_kmers[non_target_name] = [self.nb_neg + i]
+            conn.root.neg[len(conn.root.neg)] = ex
+        return len(examples)
+
     def get_training_set(self, size=None, includes=None):
         """
         Return a balanced subset of reads from the DB

coa_detection/__main__.py

@@ -4,9 +4,9 @@ import sys
 
 import coa_argparse_dicts as argparse_dicts
 
-import train_coa_nn  # todo
-import build_coa_db
-import compile_coa_model  # todo
+import build_coa_db  # todo
+import train_coa_nn
+import compile_coa_model
 import run_inference # todo
 
 import run_production_pipeline
@@ -48,10 +48,7 @@ def main(args=None):
     )
 
     for cmd, hlp, ap, fnc in commands:
-        try:
-            subparser = subparsers.add_parser(cmd, add_help=False, parents=[ap, ])
-        except:
-            cp=1
+        subparser = subparsers.add_parser(cmd, add_help=False, parents=[ap, ])
         subparser.set_defaults(func=fnc)
     args = parser.parse_args(args)
     args.func(args)

coa_detection/build_coa_db.py

-import sys
+import sys, re
 import numpy as np
 import pandas as pd
-import h5py
 from os.path import isdir, dirname, basename, splitext
 from shutil import rmtree
 from pathlib import Path
-from random import shuffle
-from AbfData import TrainingRead
+from AbfData import AbfData
 from CoaExampleDb import ExampleDb
 
 __location__ = dirname(Path(__file__).resolve())
@@ -26,32 +24,35 @@ def main(args):
         else:  # test
             file_list = list(read_index_df.query(f'fold == False').fn)
     else:
-        file_list = parse_input_path(args.fast5_in, pattern='*.fast5')
-    if args.randomize: shuffle(file_list)
+        file_list = parse_input_path(args.abf_in, pattern='*.abf')
     db_name = out_path+'db.fs'
     error_fn = out_path+'failed_reads.txt'
     npz_path = out_path + 'test_squiggles/'
     npz_path = parse_output_path(npz_path)
 
     db = ExampleDb(db_name=db_name, target=args.target, width=args.width)
-    nb_files = len(file_list)
-    count_pct_lim = 5
     nb_example_reads = 0
+
+    # split out abf files in target and non-target
+    target_idx = np.argwhere([args.target in basename(fl) for fl in file_list])[:,0]
+    if not len(target_idx):
+        raise ValueError('None of the abf files seem of target source (must have target in their file name)')
+
+    # --- process abf files ---
     for i, file in enumerate(file_list):
         try:
-            tr = TrainingRead(abf_fn=file, normalization=args.normalization)
-            nb_pos = db.add_training_read(training_read=tr, uncenter_kmer=args.uncenter_kmer)
+            non_target_name = None if i in target_idx else re.search('cOA[0-9]+', basename(file)).group(0)  # todo: counts on specific naming of files!
+            tr = AbfData(abf_fn=file, normalization=args.normalization, lowpass_freq=80)
+            nb_pos = db.add_training_read(training_read=tr, non_target=non_target_name)
             if nb_example_reads < args.nb_example_reads and nb_pos > 0:
-                np.savez(npz_path + splitext(basename(file))[0], base_labels=tr.events, raw=tr.raw)
-            if not i+1 % 10:  # Every 10 reads remove history of transactions ('pack' the database) to reduce size
-                db.pack_db()
+                event_name = non_target_name if non_target_name is not None else args.target
+                tr_labels = np.repeat(np.array('bg', dtype=f'<U{len(event_name)}'), len(tr.raw))
+                tr_labels[tr.flat_pos_indices] = event_name
+                np.savez(npz_path + splitext(basename(file))[0], base_labels=tr_labels, raw=tr.raw)
+            db.pack_db()
             if db.nb_pos > args.max_nb_examples:
                 print('Max number of examples reached')
                 break
-            percentage_processed = int((i+1) / nb_files * 100)
-            if not args.silent and percentage_processed >= count_pct_lim:
-                print(f'{percentage_processed}% of reads processed, {db.nb_pos} positives in DB')
-                count_pct_lim += 5
         except (KeyError, ValueError) as e:
             with open(error_fn, 'a') as efn:
                 efn.write('{fn}\t{err}\n'.format(err=e, fn=basename(file)))

coa_detection/coa_argparse_dicts.py

@@ -9,7 +9,6 @@ __location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file
 
 # --- GENERAL ARGUMENTS ---
 # --- inputs ---
-
 abf_in = ('--abf-in', {
     'type': lambda x: check_input_path(x),
     'required': True,
@@ -94,7 +93,6 @@ cores = ('--cores', {
 })
 
 
-
 # --- PARSERS ---
 def get_run_production_pipeline_parser():
     coa_list = ('--coa-list', {
@@ -179,11 +177,6 @@ def get_build_db_parser():
         'help': 'Filter width, should match width in the NN param file!'
     })
 
-    randomize = ('--randomize', {
-        'action': 'store_true',
-        'help': 'Randomize read order before adding examples'
-    })
-
     nb_example_reads = ('--nb-example-reads', {
         'type': int,
         'default': 100,
@@ -197,8 +190,8 @@ def get_build_db_parser():
     })
 
     for arg in (abf_in, db_dir, normalization, target, width,
-                uncenter_target, read_index, db_type, silent, nb_example_reads,
-                max_nb_examples, randomize):
+                read_index, db_type, nb_example_reads,
+                max_nb_examples):
         parser.add_argument(arg[0], **arg[1])
     return parser
 
@@ -238,7 +231,7 @@ def get_compile_model_parser():
         'required': False
     })
 
-    for arg in (nn_directory, out_model):
+    for arg in (nn_directory, out_model, parameter_file):
         parser.add_argument(arg[0], **arg[1])
     return parser
 

coa_detection/compile_coa_model.py

-import os, re, yaml, subprocess, warnings, h5py, pickle
+import os, yaml, h5py
 
 from pathlib import Path
-from tempfile import TemporaryDirectory
-
-import pandas as pd
 import tensorflow as tf
 import tensorflow.keras.backend as K
 import numpy as np
 
-from run_production_pipeline import main as rpp
 __location__ = os.path.realpath(os.path.join(os.getcwd(), os.path.dirname(__file__)))
 
-COMPLEMENT_DICT = {'A': 'T', 'T': 'A', 'C': 'G', 'G': 'C'}
-kmer_freqs_df = pd.read_parquet(f'{__location__}/../data/ncbi_16S_bacteria_archaea_kmer_counts.parquet')
-kmer_freqs_dict = dict(kmer_freqs_df.sum(axis=0))
-with open(f'{__location__}/../data/list_of_kmers_after_19_rounds.csv', 'r') as fh:
-    kmer_good_candidate_list = [km.strip() for km in fh.readlines() if len(km.strip())]
-
-
-def reverse_complement(km):
-    return ''.join([COMPLEMENT_DICT[k] for k in km][::-1])
-
-
-def fa2kmers(fa_fn, kmer_size):
-    """
-    take multifasta file, return dict of lists of k-mers present in each sequence
-    """
-    with open(fa_fn, 'r') as fh:
-        targets_txt = fh.read()
-    targets_list = ['>' + t for t in targets_txt.split('>') if t]
-    targets_dict = {re.search('(?<=>)[^\n]+', t).group(0): t for t in targets_list}
-    kmer_dict = {}
-    with TemporaryDirectory() as tdo:
-        for td in targets_dict:
-            with open(tdo + '/target.fasta', 'w') as fh:
-                fh.write(targets_dict[td])
-            subprocess.run(
-                f'jellyfish count -m {kmer_size} -s {4 ** kmer_size} -C -o {tdo}/mer_counts.jf  {tdo}/target.fasta',
-                shell=True)
-            kmer_dump = subprocess.run(f'jellyfish dump -c {tdo}/mer_counts.jf', shell=True, capture_output=True)
-            kmer_list = [km.split(' ')[0] for km in kmer_dump.stdout.decode('utf-8').split('\n')]
-            kmer_list = [km for km in kmer_list if len(km) == kmer_size and km in kmer_good_candidate_list]  # retain k-mers found to score well in discernability
-            kmer_list_curated = []  # remove reverse complements
-            for km in kmer_list:
-                if reverse_complement(km) in kmer_list_curated: continue
-                kmer_list_curated.append(km)
-            kmer_dict[td] = kmer_list_curated
-    return kmer_dict
-
-
-def get_kmer_candidates_16S(kmer_candidates_dict, min_nb_kmers=5, threshold=0.0001, filter_list=None):
-    """
-    get smallest set of 16S k-mers for which FDR<threshold
-    """
-    kmer_candidates_selection_dict = {}
-    for kc in kmer_candidates_dict:
-        p = 1.0
-        kmer_list = sorted(kmer_candidates_dict[kc], key=kmer_freqs_dict.get, reverse=True)  # sort by frequency in 16S genomes
-        kmer_candidates_selection_dict[kc] = kmer_list
-        # kmer_candidates_selection_dict[kc] = []
-        # sub_df = kmer_freqs_df.copy()
-        # sub_freqs_dict = copy(kmer_freqs_dict)
-        # cc=0
-        # while cc <= len(kmer_list) and (p >= threshold or len(kmer_candidates_selection_dict[kc]) < min_nb_kmers):
-        #     km = kmer_list.pop()  # pop least frequent k-mer
-        #     kmer_candidates_selection_dict[kc].append(km)
-        #
-        #     # Select entries positive for k-mer or its reverse-complement
-        #     km_rc = reverse_complement(km)
-        #     hit_found = False
-        #     if not len(sub_df):  # avoid making new queries if number of accessions is already 0, but min number of models has not been reached
-        #         hit_found=True
-        #     elif km in sub_df.columns and km_rc in sub_df.columns:
-        #         sub_df = sub_df.query(f'{km} or {km_rc}').copy()
-        #         hit_found=True
-        #     elif km in sub_df.columns:
-        #         sub_df.query(km, inplace=True)
-        #         hit_found = True
-        #     elif km_rc in sub_df.columns:
-        #         sub_df.query(km_rc, inplace=True)
-        #         hit_found = True
-        #     if hit_found:
-        #         hits = len(sub_df)
-        #     else:
-        #         hits = 1
-        #     hits = max(hits, 1)
-        #     p *= (hits - 1) / hits  # update FDR
-        #     kmer_list.sort(key=sub_freqs_dict.get, reverse=True)  # re-sort kmer list
-        #     sub_freqs_dict = dict(sub_df.sum(axis=0))
-        #     cc += 1
-        if filter_list:
-            kmer_candidates_selection_dict[kc] = [km for km in kmer_candidates_selection_dict[kc] if km in filter_list]
-    kc_count_dict = {kc: len(kmer_candidates_selection_dict[kc]) for kc in kmer_candidates_selection_dict}
-    selected_id = min(kc_count_dict, key=kc_count_dict.get)
-    return kmer_candidates_selection_dict[selected_id]
-
-
-def compile_model(kmer_dict, filter_width, filter_stride, threshold, k_frac):
-    k_threshold = round(len(kmer_dict) * k_frac)
-
-    input = tf.keras.Input(shape=(None, 1), ragged=True)
-    input_strided = tf.signal.frame(input.to_tensor(default_value=np.nan), frame_length=filter_width, frame_step=filter_stride, axis=1)
-    input_strided = tf.keras.layers.Masking(mask_value=np.nan)(input_strided)
-    ht_list = []
-    for km in kmer_dict:
-        mod = tf.keras.models.load_model(f'{kmer_dict[km]}/nn.h5',compile=False)
-        mod._name = km
-        h = tf.keras.layers.TimeDistributed(mod)(input_strided)
-        h = K.max(h, axis=1)
-        ht_list.append(K.cast_to_floatx(K.greater(h, threshold)))
-    nb_hits = tf.math.reduce_sum(tf.keras.layers.concatenate(ht_list), axis=1)
-    output = K.greater(nb_hits, k_threshold)
-    meta_mod = tf.keras.Model(inputs=input, outputs=output)
-    meta_mod.compile()
-    return meta_mod
-
 
-def compile_model_abundance(kmer_dict, filter_width, filter_stride, threshold):
+def compile_model_abundance(mod_dict, filter_width, filter_stride, threshold):
     input = tf.keras.Input(shape=(None, 1), ragged=True)
     input_strided = tf.signal.frame(input.to_tensor(default_value=np.nan), frame_length=filter_width, frame_step=filter_stride, axis=1)
     input_strided = tf.keras.layers.Masking(mask_value=np.nan)(input_strided)
     ht_list = []
-    for km in kmer_dict:
-        mod = tf.keras.models.load_model(f'{kmer_dict[km]}/nn.h5',compile=False)
+    for km in mod_dict:
+        mod = tf.keras.models.load_model(f'{mod_dict[km]}/nn.h5',compile=False)
         mod._name = km
         h = tf.keras.layers.TimeDistributed(mod)(input_strided)
         h = K.cast_to_floatx(K.greater(h, threshold))
@@ -135,84 +27,17 @@ def compile_model_abundance(kmer_dict, filter_width, filter_stride, threshold):
     return meta_mod
 
 
-def train_on_the_fly(kmer_list, available_mod_dict, args):
-    kmers_no_models = [km for km in kmer_list if km not in available_mod_dict]
-    if len(kmers_no_models):  # train additional models, if required
-        print(f'No models found for {len(kmers_no_models)} kmers, training on the fly!')
-        args.kmer_list = kmers_no_models
-        rpp(args)
-        # add newly generated models to available model list
-        for km in kmers_no_models:
-            if os.path.exists(f'{args.out_dir}nns/{km}/nn.h5'):  # Check to filter out failed models
-                available_mod_dict[km] = f'{args.out_dir}nns/{km}'
-            else:
-                warnings.warn(
-                    f'model generation failed for {km}, see {args.out_dir}logs. Continuing compilation without it.')
-                kmer_list.remove(km)
-    out_dict = {km: available_mod_dict.get(km, None) for km in kmer_list if km in available_mod_dict}
-    return out_dict
-
-
-def filter_accuracy(kmer_dict, acc_threshold):
-    out_dict = {}
-    for kmd in kmer_dict:
-        perf_fn = kmer_dict[kmd] + '/performance.pkl'
-        if not os.path.isfile(perf_fn): continue
-        with open(perf_fn, 'rb') as fh: perf_dict = pickle.load(fh)
-        if perf_dict['val_binary_accuracy'][-1] > acc_threshold:
-            out_dict[kmd] = kmer_dict[kmd]
-    return out_dict
-
-
 def main(args):
 
-    # List for which k-mers models are available
-    if args.nn_directory:
-        available_mod_dict = {pth.name: str(pth) for pth in Path(args.nn_directory).iterdir() if pth.is_dir()}
-    elif args.target_16S:
-        available_mod_dict = {pth.name: str(pth) for pth in Path(f'{__location__}/../16S_db/').iterdir() if pth.is_dir()}
-    else:
-        available_mod_dict = {}
+    # List for which compounds models are available
+    mod_dict = {pth.name: str(pth) for pth in Path(args.nn_directory).iterdir() if pth.is_dir()}
 
     # Parse target k-mers
-    if args.kmer_list:  # parse a given list of kmers
-        with open(args.kmer_list, 'r') as fh:
-            requested_kmer_list = [km.strip() for km in fh.readlines()]
-        requested_kmer_list = [km for km in requested_kmer_list if len(km)]
-        if args.train_required:
-            target_kmer_dict = train_on_the_fly(requested_kmer_list, available_mod_dict, args)
-        else:
-            target_kmer_dict = {km: available_mod_dict.get(km, None) for km in requested_kmer_list if km in available_mod_dict}
-    elif args.target_16S:  # estimate salient set of kmers from given 16S sequence
-        kmer_size = 8
-        requested_kmer_dict = fa2kmers(args.target_16S, kmer_size)  # Collect k-mers per sequence in target fasta marked as recognizable
-        if args.train_required:
-            target_kmer_list = get_kmer_candidates_16S(requested_kmer_dict, args.min_nb_models, 0.0001)
-            target_kmer_dict = train_on_the_fly(target_kmer_list, available_mod_dict, args)
-        else:  # filter out k-mers for which no stored model exists
-            target_kmer_list = get_kmer_candidates_16S(requested_kmer_dict, args.min_nb_models, 0.0001, filter_list=list(available_mod_dict))
-            target_kmer_dict = {km: available_mod_dict.get(km, None) for km in target_kmer_list if km in available_mod_dict}
-        if not len(target_kmer_dict):
-            raise ValueError('Sequences do not contain any of available models!')
-    else:
-        raise ValueError('Either provide --nn-directory or --target-16S')
-
-    if args.accuracy_threshold:
-        target_kmer_dict = filter_accuracy(target_kmer_dict, args.accuracy_threshold)
-
     with open(args.parameter_file, 'r') as fh:
         param_dict = yaml.load(fh, yaml.FullLoader)
-    if args.model_type == 'binary':
-        mod = compile_model(target_kmer_dict,
-                            param_dict['filter_width'], param_dict['filter_stride'],
-                            param_dict['threshold'], param_dict['k_frac'])
-    elif args.model_type == 'abundance':
-        mod = compile_model_abundance(target_kmer_dict,
-                                      param_dict['filter_width'], param_dict['filter_stride'],
-                                      param_dict['threshold'])
-    else:
-        raise ValueError(f'--model-type {args.model_type} not implemented')
+    mod = compile_model_abundance(mod_dict,
+                                  param_dict['filter_width'], param_dict['filter_stride'],
+                                  param_dict['threshold'])
     mod.save(args.out_model)
     with h5py.File(args.out_model, 'r+') as fh:
-        fh.attrs['model_type'] = args.model_type
-        fh.attrs['kmer_list'] = ','.join(list(target_kmer_dict))
+        fh.attrs['compound_list'] = ','.join(list(mod_dict))

coa_detection/train_coa_nn.py

+import os, sys, re
 import random
-import re
-from collections import Counter
 
 import yaml
 import pickle
@@ -12,6 +11,8 @@ from bokeh.io import save, output_file
 from os.path import basename, splitext
 from datetime import datetime
 
+sys.path.append(f'{os.path.dirname(__file__)}/..')
+
 import reader
 from helper_functions import load_db, parse_output_path, plot_timeseries
 
@@ -56,67 +57,41 @@ def train(parameter_file, training_data, test_data, plots_path=None,
                   cp_callback=cp_callback, tb_callback=tb_callback)
 
     # Start training
-    worst_kmers = []
     x_val, y_val = test_db.get_training_set(nb_examples)
+
     for epoch_index in range(1, params['num_kmer_switches'] + 1):
-        x_train, y_train = train_db.get_training_set(nb_examples, worst_kmers)  # todo worst_kmers mechanism errors out, fix
-        nn.train(x_train, y_train, x_val, y_val,
-                 eps_per_kmer_switch=params['eps_per_kmer_switch'], quiet=quiet)
-
-        # if params['num_kmer_switches'] > 1:
-        #     # Run on whole training reads for selection of top 5 wrongly classified k-mers
-        #     random.shuffle(train_npzs)
-        #     predicted_pos_list = []
-        #     squiggle_count = 0
-        #     for npz in train_npzs:
-        #         x, sequence_length, kmers = reader.npz_to_tf_and_kmers(npz, params['max_sequence_length'], target_kmer=train_db.target)
-        #         if x.shape[0] > nn.filter_width and np.any(np.in1d(train_db.target, kmers)):
-        #             y_hat = nn.predict(x, clean_signal=True)
-        #             predicted_pos_list.append(kmers[y_hat])
-        #             squiggle_count += 1
-        #             if squiggle_count == params['batch_size']:
-        #                 break
-        #     if not len(predicted_pos_list):
-        #         predicted_pos = np.array([])
-        #     else:
-        #         predicted_pos = np.concatenate(predicted_pos_list)
-        #     fp_kmers = Counter(predicted_pos)
-        #     _ = fp_kmers.pop(train_db.target, None)
-        #     nb_top = 5
-        #     if len(fp_kmers) < nb_top:
-        #         nb_top = len(fp_kmers)
-        #     highest_error_rates = sorted(fp_kmers.values(), reverse=True)[:nb_top]
-        #     worst_kmers = [k for k in fp_kmers if fp_kmers[k] in highest_error_rates]
+        x_train, y_train = train_db.get_training_set(nb_examples, [])
+        nn.train(x_train, y_train, x_val, y_val, eps_per_kmer_switch=params['eps_per_kmer_switch'], quiet=quiet)
 
         # Predict classification for a random trace and plot
         if plots_path:
             random.shuffle(ts_npzs)
             for i, npz in enumerate(ts_npzs):
-                x, sequence_length, kmers = reader.npz_to_tf_and_kmers(npz, target_kmer=train_db.target)
+                x, sequence_length, labels = reader.npz_to_tf_and_kmers(npz, target_kmer=train_db.target)
                 # x, sequence_length, kmers = reader.npz_to_tf_and_kmers(npz, params['max_sequence_length'], )
-                if x.shape[0] > nn.filter_width and (np.any(np.in1d(train_db.target, kmers)) or i+1 == len(ts_npzs)):
+                if x.shape[0] > nn.filter_width and (np.any(np.in1d(train_db.target, labels)) or i + 1 == len(ts_npzs)):
                     tr_fn = splitext(basename(npz))[0]
-                    start_idx = np.argwhere(np.in1d(kmers, train_db.target)).min()
+                    start_idx = np.argwhere(np.in1d(labels, train_db.target)).min()
                     oh = params['max_sequence_length'] // 2
-                    sidx = np.arange(max(0,start_idx-oh), min(start_idx+oh, len(kmers)))
+                    sidx = np.arange(max(0,start_idx-oh), min(start_idx + oh, len(labels)))
                     x = x[sidx, :]
-                    kmers = kmers[sidx]
+                    labels = labels[sidx]
                     y_hat, posterior = nn.predict(x, clean_signal=True, return_probs=True)
                     ts_predict_name = ("{path}pred_ep{ep}_ex{ex}.npz".format(path=ts_predict_path,
                                                                              ep=epoch_index,
                                                                              ex=tr_fn))
-                    if any(np.in1d(train_db.target, kmers)):
-                        graph_start = np.min(np.where(np.in1d(kmers, train_db.target))) - 30
+                    if any(np.in1d(train_db.target, labels)):
+                        graph_start = np.min(np.where(np.in1d(labels, train_db.target))) - 30
                     else:
                         graph_start = 0
                     ts_plot = plot_timeseries(raw=x,
-                                              base_labels=kmers,
+                                              base_labels=labels,
                                               posterior=posterior,
                                               y_hat=y_hat,
                                               start=graph_start,
                                               nb_classes=2)
                     output_file(sample_predictions_path + "pred_ep%s_ex%s.html" % (epoch_index, tr_fn))
-                    reader.add_to_npz(npz, ts_predict_name, [y_hat, posterior, x, kmers, nn.target], ['labels_predicted', 'posterior', 'raw_excerpt', 'base_labels_excerpt', 'target'])
+                    reader.add_to_npz(npz, ts_predict_name, [y_hat, posterior, x, labels, nn.target], ['labels_predicted', 'posterior', 'raw_excerpt', 'base_labels_excerpt', 'target'])
                     save(ts_plot)
                     break
                 if (i+1) == len(ts_npzs):