<jupyter_start><jupyter_text>Try on ALL signatures<jupyter_code>%%time

tree = create_sbt_index(bf_size, n_children=n_children)

folder = '/home/olga/pureScratch/olgabot-maca/facs/sourmash'

inp_files = sourmash_args.traverse_find_sigs([folder])

n = 0
ksizes = set()
moltypes = set()
nums = set()
scaleds = set()

for filename in inp_files:
    notify('loading {}', filename, end='\r')
    siglist = sig.load_signatures(filename,
                                 ksize=ksize,
                                 select_moltype=moltype)
    siglist = list(siglist)
    if not siglist:
        notify('\nwarning: no signatures loaded at given ksize/molecule type from {}', filename)

    # load all matching signatures in this file
    ss = None
    for ss in siglist:
        ksizes.add(ss.minhash.ksize)
        moltypes.add(sourmash_args.get_moltype(ss))
        nums.add(ss.minhash.num)

        if scaled:
            ss.minhash = ss.minhash.downsample_scaled(args.scaled)
        scaleds.add(ss.minhash.scaled)

        leaf = SigLeaf(ss.md5sum(), ss)
        tree.add_node(leaf)
        n += 1

    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)
# %%time

# adjacencies = tree.nearest_neighbor_adjacencies(n_neighbors=5, ignore_abundance=False, 
#                                                 downsample=False)
%time tree.save("/home/olga/pureScratch/olgabot-maca/facs/tabula-muris-k21.sbt.json")
def get_leaves_under(tree, node_position):
    visited, queue = set(), [node_position]
    leaves = []
    
    while queue:
        position = queue.pop(0)
        node = tree.nodes.get(position, None)
        
        if isinstance(node, Leaf):
            yield node
        else:
            queue.extend(c.pos for c in tree.children(position))
#     return leaves
list(get_leaves_under(tree, 37482))
%%time

# initialize search queue with top node of tree

n_neighbors = 5
adjacencies = []
matches = []
visited, queue = set(), [0]

ignore_abundance = False
downsample = False



# while the queue is not empty, load each node and apply search
# function.
while queue:
    position = queue.pop(0)
    node = tree.nodes.get(position, None)
    
    # repair while searching.
    if node is None:
        print("repairing...")
        if position in tree.missing_nodes:
            tree._rebuild_node(node)
            node = tree.nodes[position]
        else:
            continue
            
    # if we have not visited this node before,
    if position not in visited:
        visited.add(position)
        
    # Add 
    if isinstance(node, Leaf):
#         print(node.data)
        n = 1
        upper_internal_node = tree.parent(position)
        while n < n_parent_levels:
            upper_internal_node = tree.parent(upper_internal_node.pos)
            n += 1
#         print("upper_internal_node:", upper_internal_node)
        leaves = get_leaves_under(tree, upper_internal_node.pos)

        
        similarities = []
        for leaf in leaves:
            # Ignore self-simililarity
            if leaf == node:
                continue
#             print(leaf.data)
            similarity = node.data.similarity(leaf.data, 
                                              ignore_abundance=ignore_abundance, 
                                              downsample=downsample)
            similarities.append([node.data.name(), leaf.data.name(), similarity])
        adjacent = sorted(similarities, key=lambda x: x[1])[-n_neighbors:]
#         print(adjacent)
        adjacencies.extend(adjacent)
            
#         break
        
    else:
        queue.extend(c.pos for c in tree.children(position))
    
#     print(node)
#     print(node.data)
#     print()
    visited.add(node)

print(len(adjacencies))
tree
import igraph as ig
len(tree)
n_leaves = sum(1 for _ in tree.leaves())
n_leaves

some_leaves = [leaf for i, leaf in zip(range(5), tree.leaves())]
some_leaves


first_sig = some_leaves[0].data
first_sig

merged_minhashes = copy.copy(first_sig.minhash)
merged_minhashes.get_mins()[:20]
first_sig.minhash.get_mins()
for leaf in some_leaves[1:]:
    print(leaf.data.d)
    merged_minhashes = merged_minhashes.merge(leaf.data.minhash)
merged_minhashes.get_mins()[:20]

for leaf in some_leaves:
    print(leaf.data.minhash.jaccard(merged_minhashes))

def get_leaves(tree):
    for i, node in tree.nodes.items():
        if isinstance(node, (Leaf, SigLeaf)):
            yield i, node
knn_indices = []
knn_dists = []

leaf_to_index = dict((node.data.name(), i) for (i, (position, node)) in enumerate(get_leaves(tree)))
leaf_to_tree_position = dict((node.data.name(), position) for (i, (position, node)) in enumerate(get_leaves(tree)))

index_to_leaf = dict(zip(leaf_to_index.values(), leaf_to_index.keys()))

len(leaf_to_index)

for u, items in itertools.groupby(adjacencies, key=lambda x: x[0]):
    knn_indices_line = []
    knn_dists_line = []
    for u, v, similarity in items:
        knn_indices_line.append(leaf_to_index[v])
        knn_dists_line.append(1-similarity)
    knn_indices.append(knn_indices_line)
    knn_dists.append(knn_dists_line)
print(f"knn_indices: {knn_indices[-10:]}")
print(f"knn_dists: {knn_dists[-10:]}")
knn_dists = np.array(knn_dists)
knn_indices = np.array(knn_indices)
leaf_to_index = dict((node.data.name(), i) for i, (position, node) in enumerate(get_leaves(tree)))
leaf_to_position = dict((node.data.name(), position) for position, node in get_leaves(tree))
index_to_leaf = dict(zip(leaf_to_index.values(), leaf_to_index.keys()))

len(leaf_to_index)
# Default in UMAP
# https://github.com/lmcinnes/umap/blob/master///umap/umap_.py#L1246
local_connectivity = 1
set_op_mix_ratio = 1

sigmas, rhos = smooth_knn_dist(
        knn_dists, n_neighbors, local_connectivity=local_connectivity
    )

rows, cols, vals = compute_membership_strengths(
        knn_indices, knn_dists, sigmas, rhos
    )

result = scipy.sparse.coo_matrix(
        (vals, (rows, cols)), shape=(knn_dists.shape[0], knn_dists.shape[0])
    )
result.eliminate_zeros()

transpose = result.transpose()

prod_matrix = result.multiply(transpose)

result = (
    set_op_mix_ratio * (result + transpose - prod_matrix)
    + (1.0 - set_op_mix_ratio) * prod_matrix
)

result.eliminate_zeros()


graph_ = result
n_epochs = 0

graph = graph_.tocoo()
graph.sum_duplicates()
n_vertices = graph.shape[1]

if n_epochs <= 0:
    # For smaller datasets we can use more epochs
    if graph.shape[0] <= 10000:
        n_epochs = 500
    else:
        n_epochs = 200
n_epochs
graph.data[graph.data < (graph.data.max() / float(n_epochs))] = 0.0
graph.eliminate_zeros()
n_connected_components, labels = scipy.sparse.csgraph.connected_components(graph)

label = 0
np.where(labels == label)
graph_indexes = np.where(labels == label)[0]
signatures = [tree.nodes.get(leaf_to_position[index_to_leaf[i]]).data for i in graph_indexes]
signatures
signatures
from sourmash import signature as sig
%%time

centroid_signatures = []

for label in range(n_connected_components):
    graph_indexes = np.where(labels == label)[0]
    signatures = [tree.nodes.get(leaf_to_position[index_to_leaf[i]]).data for i in graph_indexes]
    
    first_sig = signatures[0]
    merged_signatures = sig.SourmashSignature(first_sig.minhash, name=label)
    for s in signatures[1:]:
        merged_signatures.minhash = merged_signatures.minhash.merge(s.minhash)
    centroid_signatures.append(merged_signatures)

centroid_signatures
merged_signatures.__dict__
# sig = signatures[0]
# sig
s.minhash.merge
type(s)
from sourmash.compare import compare_all_pairs
# centroid_signatures
len(centroid_signatures)
for i, sig1 in enumerate(centroid_signatures[:3]):
    for j, sig2 in enumerate(centroid_signatures[:3]):
        print(i, j)
        print('\t', sig1.minhash.similarity(sig2.minhash))
%time component_similarity_matrix = compare_all_pairs(centroid_signatures, ignore_abundance=ignore_abundance, n_jobs=32)
component_similarity_matrix.shape
%%time

component_embedding = SpectralEmbedding
    n_components=2, affinity="precomputed"
).fit_transform(component_similarity_matrix)
component_embedding /= component_embedding.max()<jupyter_output>CPU times: user 0 ns, sys: 43min 54s, total: 43min 54s
Wall time: 38 s<jupyter_text>Convert component embedding to sample embedding<jupyter_code>meta_embedding = component_embedding

component_labels = labels

dim = 2
result = np.empty((graph.shape[0], dim), dtype=np.float32)

for label in range(n_connected_components):
    component_graph = graph.tocsr()[labels == label, :].tocsc()
    component_graph = component_graph[:, labels == label].tocoo()

    distances = pairwise_distances([meta_embedding[label]], meta_embedding)
    data_range = distances[distances > 0.0].min() / 2.0

    if component_graph.shape[0] < 2 * dim:
        result[component_labels == label] = (
                random_state.uniform(
                    low=-data_range,
                    high=data_range,
                    size=(component_graph.shape[0], dim),
                )
                + meta_embedding[label]
        )
        continue

    diag_data = np.asarray(component_graph.sum(axis=0))
    # standard Laplacian
    # D = scipy.sparse.spdiags(diag_data, 0, graph.shape[0], graph.shape[0])
    # L = D - graph
    # Normalized Laplacian
    I = scipy.sparse.identity(component_graph.shape[0], dtype=np.float64)
    D = scipy.sparse.spdiags(
        1.0 / np.sqrt(diag_data),
        0,
        component_graph.shape[0],
        component_graph.shape[0],
    )
    L = I - D * component_graph * D

    k = dim + 1
    num_lanczos_vectors = max(2 * k + 1,
                              int(np.sqrt(component_graph.shape[0])))
    try:
        eigenvalues, eigenvectors = scipy.sparse.linalg.eigsh(
            L,
            k,
            which="SM",
            ncv=num_lanczos_vectors,
            tol=1e-4,
            v0=np.ones(L.shape[0]),
            maxiter=graph.shape[0] * 5,
        )
        order = np.argsort(eigenvalues)[1:k]
        component_embedding = eigenvectors[:, order]
        expansion = data_range / np.max(np.abs(component_embedding))
        component_embedding *= expansion
        result[component_labels == label] = (
                component_embedding + meta_embedding[label]
        )
    except scipy.sparse.linalg.ArpackError:
        warn(
            "WARNING: spectral initialisation failed! The eigenvector "
            "solver\nfailed. This is likely due to too small an eigengap."
            " Consider\nadding some noise or jitter to your data.\n\n"
            "Falling back to random initialisation!"
        )
        result[component_labels == label] = (
                random_state.uniform(
                    low=-data_range,
                    high=data_range,
                    size=(component_graph.shape[0], dim),
                )
                + meta_embedding[label]
        )
expansion = 10.0 / np.abs(component_embedding).max()
spectral_embedding = (component_embedding * expansion).astype(
    np.float32
) + random_state.normal(
    scale=0.0001, size=[graph.shape[0], 2]
).astype(
    np.float32
)
spectral_embedding.shape<jupyter_output><empty_output><jupyter_text>Optimize embedding<jupyter_code>a, b = find_ab_params(1, 0)

%%time


epochs_per_sample = make_epochs_per_sample(graph.data, n_epochs)

head = graph.row
tail = graph.col

random_embedding_optimized = optimize_layout(
        random_embedding,
        random_embedding,
        head,
        tail,
        n_epochs,
        n_vertices,
        epochs_per_sample,
        a,
        b,
        rng_state,
        gamma,
        initial_alpha,
        negative_sample_rate,
        verbose=verbose,
    )<jupyter_output><empty_output><jupyter_text>Do random embedding<jupyter_code>random_state = check_random_state(0)

random_embedding = random_state.uniform(
            low=-10.0, high=10.0, size=(graph.shape[0], dim)
        ).astype(np.float32)
%%time

INT32_MIN = np.iinfo(np.int32).min + 1
INT32_MAX = np.iinfo(np.int32).max - 1

rng_state = random_state.randint(INT32_MIN, INT32_MAX, 3).astype(np.int64)
gamma = 1.0
initial_alpha = 1.0
negative_sample_rate = 0.5
verbose = True
epochs_per_sample = make_epochs_per_sample(graph.data, n_epochs)

head = graph.row
tail = graph.col

random_embedding_optimized = optimize_layout(
        random_embedding,
        random_embedding,
        head,
        tail,
        n_epochs,
        n_vertices,
        epochs_per_sample,
        a,
        b,
        rng_state,
        gamma,
        initial_alpha,
        negative_sample_rate,
        verbose=verbose,
    )
random_embedding_optimized
random_embedding_optimized.shape
fig, ax = plt.subplots(figsize=(6, 4))

x = random_embedding_optimized[:, 0]
y = random_embedding_optimized[:, 1]
sns.scatterplot(x, y)
first_sig.minhash.jaccard(merged_minhashes)
g = ig.Graph(directed=True)
g.add_vertices([u for (u, v, weight) in adjacencies])  # this adds n_leaves vertices
g.add_edges([(u, v) for (u, v, weight) in adjacencies])
g.es['weight'] = [weight for (u, v, weight) in adjacencies]
g
layout = g.layout_fruchterman_reingold()
layout<jupyter_output><empty_output><jupyter_text>Do Leiden clustering<jupyter_code>partition leidenalg.find_partition(g, partition_type, **partition_kwargs)
# store output into adata.obs
groups = np.array(part.membership)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11268912360833572, 'fraction_numerical': 0.012917499286326007, 'mean_word_length': 2.9274103139013454, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 23, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '584975', 'license': 'MIT', 'n_tokens_mistral': 4918, 'n_tokens_neox': 4730, 'n_words': 926, 'path': 'notebooks/sbt2knn_playing_around.ipynb', 'repo_name': 'czbiohub/kh-analysis', 'size': '14012'}