<commit_before># encoding: utf-8
# cython: profile=False, boundscheck=False, wraparound=False

from ._bp cimport BP
from . import to_skbio_treenode
import pandas as pd
import json
import skbio
cimport cython


# see the comment in _insert_setup. Avoid the use of invalidate_caches as it
# is very expensive for tree mutation operations
class TreeNodeNoValidate(skbio.TreeNode):
    def invalidate_caches(self):
        pass


# our noop used when monkey patching invalidate_caches
def noop(*arg, **kwargs):
    pass


# pandas apply functions for preallocation of objects in bulk
def _preallocate_fragment(r):
    return TreeNodeNoValidate(name=r['fragment'], length=r['pendant_length'])


def _preallocate_empty(r):
    return TreeNodeNoValidate()


def _insert_setup(placements, bptree, insert_type):
    # insertion setup addresses:
    # * treenode caching
    # * placement ordering
    # * preallocation of objects where "easy"

    sktree = to_skbio_treenode(bptree)
    node_lookup = {n.edge_num: n for n in sktree.traverse(include_self=True)}

    # mutation operations against TreeNode is expensive as every append or
    # remove triggers a call to invalidate caches, which requires a traversal
    # to find the root (and do other stuff). so let's monkey patch the method
    # to force a noop
    for node in sktree.traverse(include_self=True):
        node.invalidate_caches = noop

    if insert_type == 'multifurcating':
        placements = placements.sort_values(['edge_num', 'pendant_length'])
    elif insert_type == 'fully_resolved':
        placements = placements.sort_values(['edge_num', 'distal_length'],
                                            ascending=[True, False])
    else:
        raise ValueError()

    placements['node'] = placements.apply(_preallocate_fragment, axis=1)

    if insert_type == 'fully_resolved':
        placements['parent'] = placements.apply(_preallocate_empty, axis=1)

    return placements, sktree, node_lookup


# pd.DataFrame is not a resolved type so we cannot use it here for cython
def insert_fully_resolved(object placements, BP bptree):
    """Update the backbone, fully resolving edges with multiple queries

    Parameters
    ----------
    placements : pd.DataFrame
        jplace data represented as a DataFrame
    bptree : bp.BP
        An instance of a BP tree, this is expected to contain edge numbers
        and correspond to the backbone for the jplace data

    Returns
    -------
    skbio.TreeNode
        A tree with the fragments placed
    """
    # TODO: profile, type and re-profile
    placements, sktree, node_lookup = \
        _insert_setup(placements, bptree, 'fully_resolved')

    for edge, edge_grp in placements.groupby('edge_num'):
        existing_node = node_lookup[edge]
        current_parent = existing_node.parent

        for _, fragment in edge_grp.iterrows():
            distal_length = fragment['distal_length']
            fragment_node = fragment['node']
            fragment_parent = fragment['parent']

            # break the edge
            current_parent.remove(existing_node)
            existing_node.parent = None

            # update branch lengths
            fragment_parent.length = existing_node.length - distal_length
            existing_node.length = distal_length

            # attach the nodes
            fragment_parent.append(fragment_node)
            fragment_parent.append(existing_node)
            current_parent.append(fragment_parent)

            # update
            current_parent = fragment_parent
<commit_msg>minor tweak to inner loop<commit_after># encoding: utf-8
# cython: profile=False, boundscheck=False, wraparound=False

from ._bp cimport BP
from . import to_skbio_treenode
import pandas as pd
import json
import skbio
cimport cython


# see the comment in _insert_setup. Avoid the use of invalidate_caches as it
# is very expensive for tree mutation operations
class TreeNodeNoValidate(skbio.TreeNode):
    def invalidate_caches(self):
        pass


# our noop used when monkey patching invalidate_caches
def noop(*arg, **kwargs):
    pass


# pandas apply functions for preallocation of objects in bulk
def _preallocate_fragment(r):
    return TreeNodeNoValidate(name=r['fragment'], length=r['pendant_length'])


def _preallocate_empty(r):
    return TreeNodeNoValidate()


def _insert_setup(placements, bptree, insert_type):
    # insertion setup addresses:
    # * treenode caching
    # * placement ordering
    # * preallocation of objects where "easy"

    sktree = to_skbio_treenode(bptree)
    node_lookup = {n.edge_num: n for n in sktree.traverse(include_self=True)}

    # mutation operations against TreeNode is expensive as every append or
    # remove triggers a call to invalidate caches, which requires a traversal
    # to find the root (and do other stuff). so let's monkey patch the method
    # to force a noop
    for node in sktree.traverse(include_self=True):
        node.invalidate_caches = noop

    if insert_type == 'multifurcating':
        placements = placements.sort_values(['edge_num', 'pendant_length'])
    elif insert_type == 'fully_resolved':
        placements = placements.sort_values(['edge_num', 'distal_length'],
                                            ascending=[True, False])
    else:
        raise ValueError()

    placements['node'] = placements.apply(_preallocate_fragment, axis=1)

    if insert_type == 'fully_resolved':
        placements['parent'] = placements.apply(_preallocate_empty, axis=1)

    return placements, sktree, node_lookup


# pd.DataFrame is not a resolved type so we cannot use it here for cython
def insert_fully_resolved(object placements, BP bptree):
    """Update the backbone, fully resolving edges with multiple queries

    Parameters
    ----------
    placements : pd.DataFrame
        jplace data represented as a DataFrame
    bptree : bp.BP
        An instance of a BP tree, this is expected to contain edge numbers
        and correspond to the backbone for the jplace data

    Returns
    -------
    skbio.TreeNode
        A tree with the fragments placed
    """
    # TODO: profile, type and re-profile
    placements, sktree, node_lookup = \
        _insert_setup(placements, bptree, 'fully_resolved')

    for edge, edge_grp in placements.groupby('edge_num'):
        existing_node = node_lookup[edge]
        current_parent = existing_node.parent
        
        # break the edge
        current_parent.remove(existing_node)
        existing_node.parent = None
        existing_length = existing_node.length

        for _, fragment in edge_grp.iterrows():
            distal_length = fragment['distal_length']
            fragment_node = fragment['node']
            fragment_parent = fragment['parent']

            # update branch lengths
            fragment_parent.length = existing_length - distal_length
            existing_length = distal_length

            # attach the nodes
            fragment_parent.append(fragment_node)
            current_parent.append(fragment_parent)

            # update
            current_parent = fragment_parent
        
        existing_node.length = existing_length
        current_parent.append(existing_node)
        existing_node.length = distal_length

{'fraction_non_alphanumeric': 0.07555494278229698, 'fraction_numerical': 0.0008272438990762443, 'mean_word_length': 2.9682713347921226, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5472327', 'n_tokens_mistral': 2038, 'n_tokens_neox': 1961, 'n_words': 610}