Add option for SBT creation to be localized

sourmash/cli/index.py

@@ -75,6 +75,13 @@ def subparser(subparsers):
         '--scaled', metavar='FLOAT', type=float, default=0,
         help='downsample signatures to the specified scaled factor'
     )
+    subparser.add_argument(
+        '--not-localized', action='store_true', default=False,
+        help='Do not create a localized SBT index. A localized index guarantees that '
+             'any two leaves sharing a parent are more similar to leaves not sharing '
+             'a parent. Localized indices are required for building a nearest neighbor'
+             ' graph from the SBT.'
+    )
     add_moltype_args(subparser)
 
 

sourmash/commands.py

@@ -345,10 +345,11 @@ def index(args):
     set_quiet(args.quiet)
     moltype = sourmash_args.calculate_moltype(args)
 
-    if args.append:
-        tree = load_sbt_index(args.sbt_name)
+    if args.traverse_directory:
+        inp_files = list(sourmash_args.traverse_find_sigs(args.signatures,
+                                                          args.force))
     else:
-        tree = create_sbt_index(args.bf_size, n_children=args.n_children)
+        inp_files = list(args.signatures)
 
     if args.sparseness < 0 or args.sparseness > 1.0:
         error('sparseness must be in range [0.0, 1.0].')
@@ -357,31 +358,34 @@ def index(args):
         args.scaled = int(args.scaled)
         notify('downsampling signatures to scaled={}', args.scaled)
 
-    inp_files = list(args.signatures)
-    if args.from_file:
-        more_files = sourmash_args.load_file_list_of_signatures(args.from_file)
-        inp_files.extend(more_files)
+    notify('loading {} files into SBT', len(inp_files))
 
-    if not inp_files:
-        error("ERROR: no files to index!? Supply on command line or use --from-file")
-        sys.exit(-1)
+    tree, n = load_matching_signatures_into_tree(
+        inp_files, args.ksize, moltype, args.scaled, args.append, args.sbt_name,
+        return_n=True)
 
-    notify('loading {} files into SBT', len(inp_files))
+    notify('loaded {} sigs; saving SBT under "{}"', n, args.sbt_name)
+    tree.save(args.sbt_name, sparseness=args.sparseness)
 
-    progress = sourmash_args.SignatureLoadingProgress()
+
+def load_matching_signatures_into_tree(filenames, ksize, moltype, scaled=0,
+                                       append=False, sbt_name=None, bf_size=1e5,
+                                       n_children=2, return_n=False):
+    if append:
+        tree = load_sbt_index(sbt_name)
+    else:
+        tree = create_sbt_index(bf_size, n_children=n_children)
 
     n = 0
     ksizes = set()
     moltypes = set()
     nums = set()
     scaleds = set()
-    for f in inp_files:
-        siglist = sourmash_args.load_file_as_signatures(f,
-                                                        ksize=args.ksize,
-                                                        select_moltype=moltype,
-                                                        traverse=args.traverse_directory,
-                                                        yield_all_files=args.force,
-                                                        progress=progress)
+    for f in filenames:
+        if n % 100 == 0:
+            notify('\r...reading from {} ({} signatures so far)', f, n, end='')
+        siglist = sig.load_signatures(f, ksize=ksize,
+                                      select_moltype=moltype)
 
         # load all matching signatures in this file
         ss = None
@@ -390,8 +394,8 @@ def index(args):
             moltypes.add(sourmash_args.get_moltype(ss))
             nums.add(ss.minhash.num)
 
-            if args.scaled:
-                ss.minhash = ss.minhash.downsample_scaled(args.scaled)
+            if scaled:
+                ss.minhash = ss.minhash.downsample_scaled(scaled)
             scaleds.add(ss.minhash.scaled)
 
             tree.insert(ss)
@@ -400,32 +404,35 @@ def index(args):
         if not ss:
             continue
 
-        # check to make sure we aren't loading incompatible signatures
-        if len(ksizes) > 1 or len(moltypes) > 1:
-            error('multiple k-mer sizes or molecule types present; fail.')
-            error('specify --dna/--protein and --ksize as necessary')
-            error('ksizes: {}; moltypes: {}',
-                  ", ".join(map(str, ksizes)), ", ".join(moltypes))
-            sys.exit(-1)
-
-        if nums == { 0 } and len(scaleds) == 1:
-            pass # good
-        elif scaleds == { 0 } and len(nums) == 1:
-            pass # also good
-        else:
-            error('trying to build an SBT with incompatible signatures.')
-            error('nums = {}; scaleds = {}', repr(nums), repr(scaleds))
-            sys.exit(-1)
-
+        check_signature_compatibilty_to_tree(ksizes, moltypes, nums, scaleds)
     notify('')
 
     # did we load any!?
     if n == 0:
         error('no signatures found to load into tree!? failing.')
         sys.exit(-1)
+    if return_n:
+        return tree, n
+    else:
+        return tree
 
-    notify('loaded {} sigs; saving SBT under "{}"', n, args.sbt_name)
-    tree.save(args.sbt_name, sparseness=args.sparseness)
+
+def check_signature_compatibilty_to_tree(ksizes, moltypes, nums, scaleds):
+    # check to make sure we aren't loading incompatible signatures
+    if len(ksizes) > 1 or len(moltypes) > 1:
+        error('multiple k-mer sizes or molecule types present; fail.')
+        error('specify --dna/--protein and --ksize as necessary')
+        error('ksizes: {}; moltypes: {}',
+              ", ".join(map(str, ksizes)), ", ".join(moltypes))
+        sys.exit(-1)
+    if nums == {0} and len(scaleds) == 1:
+        pass  # good
+    elif scaleds == {0} and len(nums) == 1:
+        pass  # also good
+    else:
+        error('trying to build an SBT with incompatible signatures.')
+        error('nums = {}; scaleds = {}', repr(nums), repr(scaleds))
+        sys.exit(-1)
 
 
 def search(args):

sourmash/lca/lca_db.py

@@ -4,27 +4,11 @@
 import json
 import gzip
 from collections import OrderedDict, defaultdict, Counter
-import functools
 
-import sourmash
 from sourmash._minhash import get_max_hash_for_scaled
-from sourmash.logging import notify, error, debug
+from sourmash.logging import notify, debug
 from sourmash.index import Index
-
-
-def cached_property(fun):
-    """A memoize decorator for class properties."""
-    @functools.wraps(fun)
-    def get(self):
-        try:
-            return self._cache[fun]
-        except AttributeError:
-            self._cache = {}
-        except KeyError:
-            pass
-        ret = self._cache[fun] = fun(self)
-        return ret
-    return property(get)
+from sourmash.utils import cached_property
 
 
 class LCA_Database(Index):
@@ -166,7 +150,6 @@ def __repr__(self):
 
     def signatures(self):
         "Return all of the signatures in this LCA database."
-        from sourmash import SourmashSignature
         for v in self._signatures.values():
             yield v