<commit_before>from SquiggleRead import *
from NanoUtils import *
from Clustal import *
from ctypes import *
from collections import defaultdict, namedtuple
import sys

StrandAnchor = namedtuple('StrandAnchor', ['idstr', 'row_id', 'sr_idx', 'strand', 'event_idx', 'rc', 'diff'])

class Anchor:
    def __init__(self, column, cons_kidx):
        self.column = column
        self.cons_kidx = cons_kidx
        self.strands = list()
        self.n_outliers = 0
        self.n_samples = 0
        self.sum_diff = 0

    def add_strand(self, sa):
        self.strands.append(sa)

# Convenience type for working with ctypes
c_double_p = POINTER(c_double)

class CPoreModelInterface(Structure):
    _fields_ = [('n_states', c_int),
                ('pore_model_level_mean', c_double_p),
                ('pore_model_level_stdv', c_double_p),
                ('pore_model_sd_mean', c_double_p),
                ('pore_model_sd_stdv', c_double_p),
                ('pore_model_scale', c_double),
                ('pore_model_shift', c_double),
                ('pore_model_drift', c_double),
                ('pore_model_var', c_double)]

class CEventSequenceInterface(Structure):
    _fields_ = [('n_events', c_int),
                ('level', c_double_p),
                ('stdv', c_double_p),
                ('time', c_double_p)]

class CSquiggleReadInterface(Structure):
    _fields_ = [('pore_model', CPoreModelInterface * 2),
                ('events', CEventSequenceInterface * 2)]

class CReadStateInterface(Structure):
    _fields_ = [('read_idx', c_int),
                ('event_start_idx', c_int),
                ('event_stop_idx', c_int),
                ('strand', c_int),
                ('stride', c_int),
                ('rc', c_int)]

class CReadAnchorInterface(Structure):
    _fields_ = [('event_idx', c_int),
                ('rc', c_int)]

# Generate all the paths of length (k+1)
# starting from the given k-mer
def generate_kmer_paths(kmer):
    k = len(kmer)
    all_kmers = generate_mers(dna_alphabet, k)

    out = []
    for a in all_kmers:
        path_str = kmer + a
        out.append(str2kmers(kmer + a, k))
    return out

# get the next observed event in the range
def get_next_event(sr, start, stop, stride, strand):
    for i in xrange(start, stop, stride):

        events = sr.event_map['2D'][i]
        if len(events) == 0:
            continue

        ei = -1

        if strand == 't':
            ei = events[0][0]
        else:
            ei = events[0][1]
        
        if ei != -1:
            return ei
    return -1

# Get the closest event to the indicated
# 2D kmer for the given squiggle read
def get_closest_event_to(sr, kidx, strand, kmer):
    e_level = sr.get_expected_level(kmer, strand)
    
    first_event_before = -1
    first_event_after = -1
    m = 1000
    
    # before
    stop = kidx - m
    if stop < 0:
        stop = -1 # end is exclusive
    
    first_event_before = get_next_event(sr, kidx, stop, -1, strand)

    # after
    stop = kidx + m
    if stop >= len(sr.event_map['2D']):
        stop = len(sr.event_map['2D'])

    first_event_after = get_next_event(sr, kidx, stop, 1, strand)

    if first_event_before is None or first_event_after is None:
        return -1

    # evaluate levels
    b_level = sr.get_drift_corrected_event_level(first_event_before, strand)
    a_level = sr.get_drift_corrected_event_level(first_event_after, strand)
    #print 'SEARCH', first_event_before, first_event_after, b_level, a_level, e_level

    if abs(b_level - e_level) < abs(a_level - e_level):
        return first_event_before
    else:
        return first_event_after


def calculate_diff_to_expected(sr, kidx, strand, kmer):
    
    ei = get_closest_event_to(sr, kidx, strand, kmer)
    if ei == -1:
        return (float("-inf"), ei)

    # Observed event   
    level = sr.get_drift_corrected_event_level(ei, strand)

    k_level = sr.get_expected_level(kmer, strand)
    k_sd = sr.get_expected_sd(kmer, strand)
    res = (level - k_level) / k_sd
    return (res, ei)

# For the pair of input Anchors, make a list
# of StrandAnchors thare are anchored at both inputs
def pair_anchor_strands(a1, a2):

    out = list()

    # Index the first anchor by ID
    anchor_rows = dict()

    a1_by_id = dict()
    for sa_1 in a1.strands:
        a1_by_id[sa_1.idstr] = sa_1
    
    for sa_2 in a2.strands:
        if sa_2.idstr in a1_by_id:
            
            # This read strand is anchored on both ends
            sa_1 = a1_by_id[sa_2.idstr]

            out.append( (sa_1, sa_2) )

    return out
            
def add_read_state_from_anchor_strands(lib, sa_1, sa_2):
    
    start_ei = sa_1.event_idx
    end_ei = sa_2.event_idx
    stride = 1

    if start_ei > end_ei:
        stride = -1

    strand_idx = 0
    if sa_1.strand == 'c':
        strand_idx = 1

    rs = CReadStateInterface(sa_1.sr_idx, start_ei, end_ei, strand_idx, stride, sa_1.rc)
    lib.add_read_state(rs)

# Import squiggle reads into C code
def load_reads(lib, squiggle_reads):

    for sr in squiggle_reads:
        if sr is None:
            continue
        
        pm_params = []
        event_params = []

        # Setup pore model and event sequence data for each strand of the read
        for s in ('t', 'c'):

            # Pore Model
            level_mean = sr.pm[s].model_level_mean.ctypes.data_as(c_double_p)
            level_stdv = sr.pm[s].model_level_stdv.ctypes.data_as(c_double_p)
            sd_mean = sr.pm[s].model_sd_mean.ctypes.data_as(c_double_p)
            sd_stdv = sr.pm[s].model_sd_stdv.ctypes.data_as(c_double_p)
            
            scale = sr.pm[s].scale
            shift = sr.pm[s].shift
            drift = sr.pm[s].drift
            var = sr.pm[s].var
            
            pm_params.append(CPoreModelInterface(1024, level_mean, level_stdv, sd_mean, sd_stdv, scale, shift, drift, var))

            # Events
            n_events = len(sr.event_level[s])
            level = sr.event_level[s].ctypes.data_as(c_double_p)
            stdv = sr.event_stdv[s].ctypes.data_as(c_double_p)
            time = sr.event_time[s].ctypes.data_as(c_double_p)
            event_params.append(CEventSequenceInterface(n_events, level, stdv, time))

        #
        params = CSquiggleReadInterface((pm_params[0], pm_params[1]), 
                                        (event_params[0], event_params[1]))
        lib.add_read(params)

#
def call_consensus_for_file(input_filename, num_threads):

    if num_threads <= 0:
        sys.stderr.write("Invalid number of threads")
        sys.exit(1)

    #
    # Build the map from read indices to fast5 files
    #
    fast5_fofn_fn = 'consensus_fast5.fofn'
    fast5_fofn_fh = open(fast5_fofn_fn)

    f5_files = []
    for l in fast5_fofn_fh:
        f5_files.append(l.rstrip())

    #
    # Load squiggle reads
    #
    use_poabase_signals = False
    do_training = False

    clustal = Clustal(input_filename)

    print "!!!!!!!! using first read as consensus row !!!!!!!"

    cons_row = 0
    read_rows = clustal.get_read_rows()

    # Initialize traversal
    n_reads = len(read_rows)

    # Parse the clustal row names to generate POAReads
    poa_reads = list()
    row_to_poa_read_idx = dict()
    for rr in read_rows:
        pa = unpack_poa_id(clustal.alignment[rr].id)

        # Append read and update map
        poa_reads.append(pa)
        
        idx = len(poa_reads) - 1
        row_to_poa_read_idx[rr] = idx

    sys.stderr.write("Loading squigglereads\n")
    squiggle_reads = list()
    poa_idx_to_sr_idx = dict()

    seen_ids = dict()

    for (poa_idx, pr) in enumerate(poa_reads):

        # A read may appear in multiple rows of the MA
        # if the aligner emitted a few local alignments
        # Only load a squiggleread once
        if pr.read_id in seen_ids:
            poa_idx_to_sr_idx[poa_idx] = seen_ids[pr.read_id]
            continue

        print 'Loading', pr.read_id, f5_files[pr.read_id]
        squiggle_reads.append(SquiggleRead(f5_files[pr.read_id]))

        sr_idx = len(squiggle_reads) - 1
        poa_idx_to_sr_idx[poa_idx] = sr_idx
        seen_ids[pr.read_id] = sr_idx

    sys.stderr.write("Done loading squigglereads\n")

    #
    # Build anchors
    #

    DEBUG=False
    ANCHOR_DISTANCE = 50 # in bases along the consensus row

    # this list stores the number of actual bases seen in each row up to current column
    n_bases = [0] * n_reads
    anchors = list()

    cons_kidx = 0
    consensus_sequence = str(clustal.alignment[cons_row].seq)

    n_cols = len(clustal.alignment[cons_row])
    col = 0

    #
    while col < n_cols:

        # Is there a consensus base in this column?
        cons_base = clustal.alignment[cons_row][col]
        cons_kmer = clustal.get_kmer(cons_row, col, 5)

        # Should we make an anchor here?
        if cons_kidx % ANCHOR_DISTANCE == 0 and cons_base != '-' and len(cons_kmer) == 5:

            print 'ANCHOR', cons_kidx
            current_anchor = Anchor(col, cons_kidx)

            # Make anchor
            for rr in read_rows:
                b = clustal.alignment[rr][col]

                # Read metadata
                poa_read = poa_reads[rr]

                if poa_read.is_base and not use_poabase_signals:
                    continue

                poa_idx = row_to_poa_read_idx[rr]
                sr_idx = poa_idx_to_sr_idx[poa_idx]
                sr = squiggle_reads[sr_idx]

                read_kidx = n_bases[rr] + poa_read.start
                
                last_read_kidx = poa_read.stop - 5
                if poa_read.is_base:
                    last_read_kidx = sr.get_2D_length() - 5

                print "\tROW", rr, poa_reads[rr].read_id, read_kidx, last_read_kidx

                # skip reads that are not aligned in this range
                if (read_kidx <= poa_read.start and b == '-') or (read_kidx >= last_read_kidx and b == '-'):<commit_msg>Fix anchoring for reads near the end of the alignment<commit_after>from SquiggleRead import *
from NanoUtils import *
from Clustal import *
from ctypes import *
from collections import defaultdict, namedtuple
import sys

StrandAnchor = namedtuple('StrandAnchor', ['idstr', 'row_id', 'sr_idx', 'strand', 'event_idx', 'rc', 'diff'])

class Anchor:
    def __init__(self, column, cons_kidx):
        self.column = column
        self.cons_kidx = cons_kidx
        self.strands = list()
        self.n_outliers = 0
        self.n_samples = 0
        self.sum_diff = 0

    def add_strand(self, sa):
        self.strands.append(sa)

# Convenience type for working with ctypes
c_double_p = POINTER(c_double)

class CPoreModelInterface(Structure):
    _fields_ = [('n_states', c_int),
                ('pore_model_level_mean', c_double_p),
                ('pore_model_level_stdv', c_double_p),
                ('pore_model_sd_mean', c_double_p),
                ('pore_model_sd_stdv', c_double_p),
                ('pore_model_scale', c_double),
                ('pore_model_shift', c_double),
                ('pore_model_drift', c_double),
                ('pore_model_var', c_double)]

class CEventSequenceInterface(Structure):
    _fields_ = [('n_events', c_int),
                ('level', c_double_p),
                ('stdv', c_double_p),
                ('time', c_double_p)]

class CSquiggleReadInterface(Structure):
    _fields_ = [('pore_model', CPoreModelInterface * 2),
                ('events', CEventSequenceInterface * 2)]

class CReadStateInterface(Structure):
    _fields_ = [('read_idx', c_int),
                ('event_start_idx', c_int),
                ('event_stop_idx', c_int),
                ('strand', c_int),
                ('stride', c_int),
                ('rc', c_int)]

class CReadAnchorInterface(Structure):
    _fields_ = [('event_idx', c_int),
                ('rc', c_int)]

# Generate all the paths of length (k+1)
# starting from the given k-mer
def generate_kmer_paths(kmer):
    k = len(kmer)
    all_kmers = generate_mers(dna_alphabet, k)

    out = []
    for a in all_kmers:
        path_str = kmer + a
        out.append(str2kmers(kmer + a, k))
    return out

# get the next observed event in the range
def get_next_event(sr, start, stop, stride, strand):
    for i in xrange(start, stop, stride):

        events = sr.event_map['2D'][i]
        if len(events) == 0:
            continue

        ei = -1

        if strand == 't':
            ei = events[0][0]
        else:
            ei = events[0][1]
        
        if ei != -1:
            return ei
    return -1

# Get the closest event to the indicated
# 2D kmer for the given squiggle read
def get_closest_event_to(sr, kidx, strand, kmer):
    e_level = sr.get_expected_level(kmer, strand)
    
    first_event_before = -1
    first_event_after = -1
    m = 1000
    
    # before
    stop = kidx - m
    if stop < 0:
        stop = -1 # end is exclusive
    
    first_event_before = get_next_event(sr, kidx, stop, -1, strand)

    # after
    stop = kidx + m
    if stop >= len(sr.event_map['2D']):
        stop = len(sr.event_map['2D'])

    first_event_after = get_next_event(sr, kidx, stop, 1, strand)

    if first_event_before is None or first_event_after is None:
        return -1

    # evaluate levels
    b_level = sr.get_drift_corrected_event_level(first_event_before, strand)
    a_level = sr.get_drift_corrected_event_level(first_event_after, strand)
    #print 'SEARCH', first_event_before, first_event_after, b_level, a_level, e_level

    if abs(b_level - e_level) < abs(a_level - e_level):
        return first_event_before
    else:
        return first_event_after


def calculate_diff_to_expected(sr, kidx, strand, kmer):
    
    ei = get_closest_event_to(sr, kidx, strand, kmer)
    if ei == -1:
        return (float("-inf"), ei)

    # Observed event   
    level = sr.get_drift_corrected_event_level(ei, strand)

    k_level = sr.get_expected_level(kmer, strand)
    k_sd = sr.get_expected_sd(kmer, strand)
    res = (level - k_level) / k_sd
    return (res, ei)

# For the pair of input Anchors, make a list
# of StrandAnchors thare are anchored at both inputs
def pair_anchor_strands(a1, a2):

    out = list()

    # Index the first anchor by ID
    anchor_rows = dict()

    a1_by_id = dict()
    for sa_1 in a1.strands:
        a1_by_id[sa_1.idstr] = sa_1
    
    for sa_2 in a2.strands:
        if sa_2.idstr in a1_by_id:
            
            # This read strand is anchored on both ends
            sa_1 = a1_by_id[sa_2.idstr]

            out.append( (sa_1, sa_2) )

    return out
            
def add_read_state_from_anchor_strands(lib, sa_1, sa_2):
    
    start_ei = sa_1.event_idx
    end_ei = sa_2.event_idx
    stride = 1

    if start_ei > end_ei:
        stride = -1

    strand_idx = 0
    if sa_1.strand == 'c':
        strand_idx = 1

    rs = CReadStateInterface(sa_1.sr_idx, start_ei, end_ei, strand_idx, stride, sa_1.rc)
    lib.add_read_state(rs)

# Import squiggle reads into C code
def load_reads(lib, squiggle_reads):

    for sr in squiggle_reads:
        if sr is None:
            continue
        
        pm_params = []
        event_params = []

        # Setup pore model and event sequence data for each strand of the read
        for s in ('t', 'c'):

            # Pore Model
            level_mean = sr.pm[s].model_level_mean.ctypes.data_as(c_double_p)
            level_stdv = sr.pm[s].model_level_stdv.ctypes.data_as(c_double_p)
            sd_mean = sr.pm[s].model_sd_mean.ctypes.data_as(c_double_p)
            sd_stdv = sr.pm[s].model_sd_stdv.ctypes.data_as(c_double_p)
            
            scale = sr.pm[s].scale
            shift = sr.pm[s].shift
            drift = sr.pm[s].drift
            var = sr.pm[s].var
            
            pm_params.append(CPoreModelInterface(1024, level_mean, level_stdv, sd_mean, sd_stdv, scale, shift, drift, var))

            # Events
            n_events = len(sr.event_level[s])
            level = sr.event_level[s].ctypes.data_as(c_double_p)
            stdv = sr.event_stdv[s].ctypes.data_as(c_double_p)
            time = sr.event_time[s].ctypes.data_as(c_double_p)
            event_params.append(CEventSequenceInterface(n_events, level, stdv, time))

        #
        params = CSquiggleReadInterface((pm_params[0], pm_params[1]), 
                                        (event_params[0], event_params[1]))
        lib.add_read(params)

#
def call_consensus_for_file(input_filename, num_threads):

    if num_threads <= 0:
        sys.stderr.write("Invalid number of threads")
        sys.exit(1)

    #
    # Build the map from read indices to fast5 files
    #
    fast5_fofn_fn = 'consensus_fast5.fofn'
    fast5_fofn_fh = open(fast5_fofn_fn)

    f5_files = []
    for l in fast5_fofn_fh:
        f5_files.append(l.rstrip())

    #
    # Load squiggle reads
    #
    use_poabase_signals = False
    do_training = False

    clustal = Clustal(input_filename)

    print "!!!!!!!! using first read as consensus row !!!!!!!"

    cons_row = 0
    read_rows = clustal.get_read_rows()

    # Initialize traversal
    n_reads = len(read_rows)

    # Parse the clustal row names to generate POAReads
    poa_reads = list()
    row_to_poa_read_idx = dict()
    for rr in read_rows:
        pa = unpack_poa_id(clustal.alignment[rr].id)

        # Append read and update map
        poa_reads.append(pa)
        
        idx = len(poa_reads) - 1
        row_to_poa_read_idx[rr] = idx

    sys.stderr.write("Loading squigglereads\n")
    squiggle_reads = list()
    poa_idx_to_sr_idx = dict()

    seen_ids = dict()

    for (poa_idx, pr) in enumerate(poa_reads):

        # A read may appear in multiple rows of the MA
        # if the aligner emitted a few local alignments
        # Only load a squiggleread once
        if pr.read_id in seen_ids:
            poa_idx_to_sr_idx[poa_idx] = seen_ids[pr.read_id]
            continue

        print 'Loading', pr.read_id, f5_files[pr.read_id]
        squiggle_reads.append(SquiggleRead(f5_files[pr.read_id]))

        sr_idx = len(squiggle_reads) - 1
        poa_idx_to_sr_idx[poa_idx] = sr_idx
        seen_ids[pr.read_id] = sr_idx

    sys.stderr.write("Done loading squigglereads\n")

    #
    # Build anchors
    #

    DEBUG=False
    ANCHOR_DISTANCE = 50 # in bases along the consensus row

    # this list stores the number of actual bases seen in each row up to current column
    n_bases = [0] * n_reads
    anchors = list()

    cons_kidx = 0
    consensus_sequence = str(clustal.alignment[cons_row].seq)

    n_cols = len(clustal.alignment[cons_row])
    col = 0

    #
    while col < n_cols:

        # Is there a consensus base in this column?
        cons_base = clustal.alignment[cons_row][col]
        cons_kmer = clustal.get_kmer(cons_row, col, 5)

        # Should we make an anchor here?
        if cons_kidx % ANCHOR_DISTANCE == 0 and cons_base != '-' and len(cons_kmer) == 5:

            #print 'ANCHOR', cons_kidx
            current_anchor = Anchor(col, cons_kidx)

            # Make anchor
            for rr in read_rows:
                b = clustal.alignment[rr][col]

                # Read metadata
                poa_read = poa_reads[rr]

                if poa_read.is_base and not use_poabase_signals:
                    continue

                poa_idx = row_to_poa_read_idx[rr]
                sr_idx = poa_idx_to_sr_idx[poa_idx]
                sr = squiggle_reads[sr_idx]

                read_kidx = n_bases[rr] + poa_read.start
                
                last_read_kidx = poa_read.stop - 5
                if poa_read.is_base:
                    last_read_kidx = sr.get_2D_length() - 5

                #print "\tROW", rr, poa_reads[rr].read_id, read_kidx, last_read_kidx

                # skip reads that are not aligned in this range
                if (read_kidx <= poa_read.start and b == '-') or (read_kidx >= last_read_kidx):
{'fraction_non_alphanumeric': 0.09821473472466684, 'fraction_numerical': 0.00935378425949208, 'mean_word_length': 2.221448242345699, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3562726', 'n_tokens_mistral': 7209, 'n_tokens_neox': 6774, 'n_words': 1567}