<commit_before>"""
Align multiple sequences from FASTA.
"""

import os
from shutil import copyfile
import numpy as np
from Bio import AlignIO, SeqIO, Seq, Align
from .utils import run_shell_command, nthreads_value, shquote

class AlignmentError(Exception):
    # TODO: this exception should potentially be renamed and made augur-wide
    # thus allowing any module to raise it and have the message printed & augur
    # exit with code 1
    pass

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, nargs="+", metavar="FASTA", help="sequences to align")
    parser.add_argument('--output', '-o', default="alignment.fasta", help="output file (default: %(default)s)")
    parser.add_argument('--nthreads', type=nthreads_value, default=1,
                                help="number of threads to use; specifying the value 'auto' will cause the number of available CPU cores on your system, if determinable, to be used")
    parser.add_argument('--method', default='mafft', choices=["mafft"], help="alignment program to use")
    parser.add_argument('--reference-name', metavar="NAME", type=str, help="strip insertions relative to reference sequence; use if the reference is already in the input sequences")
    parser.add_argument('--reference-sequence', metavar="PATH", type=str, help="Add this reference sequence to the dataset & strip insertions relative to this. Use if the reference is NOT already in the input sequences")
    parser.add_argument('--remove-reference', action="store_true", default=False, help="remove reference sequence from the alignment")
    parser.add_argument('--fill-gaps', action="store_true", default=False, help="If gaps represent missing data rather than true indels, replace by N after aligning.")
    parser.add_argument('--existing-alignment', metavar="FASTA", default=False, help="An existing alignment to which the sequences will be added. The ouput alignment will be the same length as this existing alignment.")
    parser.add_argument('--debug', action="store_true", default=False, help="Produce extra files (e.g. pre- and post-aligner files) which can help with debugging poor alignments.")

def prepare(sequences, existing_aln_fname, output, ref_name, ref_seq_fname):
    seqs = read_sequences(*sequences)
    seqs_to_align_fname = output + "to_align.fasta"

    # Load existing alignment
    existing_aln = None
    if existing_aln_fname:
        existing_aln = read_alignment(existing_aln_fname)

    # Load reference alignment
    ref_seq = None
    if ref_name:
        ensure_reference_strain_present(ref_name, existing_aln, seqs)
    elif ref_seq_fname:
        ref_seq = read_reference(ref_seq_fname)
        ref_name = ref_seq.id

    if existing_aln:
        # Strip the existing sequences from the new sequences, add the reference to the alignment
        seqs_to_align_fname = output + "new_seqs_to_align.fasta"
        seqs = prune_seqs_matching_alignment(seqs, existing_aln)
        if ref_seq:
            if len(ref_seq) != existing_aln.get_alignment_length():
                raise AlignmentError("ERROR: Provided existing alignment ({}bp) is not the same length as the reference sequence ({}bp)".format(existing_aln.get_alignment_length(), len(ref_seq)))
            existing_aln_fname = existing_aln_fname + ".ref.fasta"
            existing_aln.append(ref_seq)
            write_seqs(existing_aln, existing_aln_fname)
    elif ref_seq: # Got a reference sequence but no existing alignment
        # reference sequence needs to be the first one for auto direction
        # adjustment (auto reverse-complement)
        seqs.insert(0, ref_seq)
    write_seqs(seqs, seqs_to_align_fname)

    check_duplicates(existing_aln, seqs)
    return existing_aln_fname, seqs_to_align_fname, ref_name

def run(args):
    '''
    Parameters
    ----------
    args : namespace
        arguments passed in via the command-line from augur

    Returns
    -------
    int
        returns 0 for success, 1 for general error
    '''
    temp_files_to_remove = []

    try:
        check_arguments(args)
        existing_aln_fname, seqs_to_align_fname, ref_name = prepare(args.sequences, args.existing_alignment, args.output, args.reference_name, args.reference_sequence)
        # -- existing_aln_fname, seqs_to_align_fname, ref_name --

        # before aligning, make a copy of the data that the aligner receives as input (very useful for debugging purposes)
        if args.debug and not existing_aln_fname:
            copyfile(seqs_to_align_fname, args.output+".pre_aligner.fasta")

        # generate alignment command & run
        cmd = generate_alignment_cmd(args.method, args.nthreads, existing_aln_fname, seqs_to_align_fname, args.output, args.output+".log")
        success = run_shell_command(cmd)
        if not success:
            raise AlignmentError("Error during alignment")

        # after aligning, make a copy of the data that the aligner produced (useful for debugging)
        if args.debug:
            copyfile(args.output, args.output+".post_aligner.fasta")

        # -- ref_name --
        # reads the new alignment
        seqs = read_alignment(args.output)

        # convert the aligner output to upper case and remove auto reverse-complement prefix
        prettify_alignment(seqs)

        # if we've specified a reference, strip out all the columns not present in the reference
        # this will overwrite the alignment file
        if ref_name:
            seqs = strip_non_reference(seqs, ref_name, keep_reference=not args.remove_reference)
        if args.fill_gaps:
            make_gaps_ambiguous(seqs)

        # write the modified sequences back to the alignment file
        write_seqs(seqs, args.output)

    except AlignmentError as e:
        print(str(e))
        return 1

    # finally, remove any temporary files
    for fname in temp_files_to_remove:
        os.remove(fname)

#####################################################################################################

def read_sequences(*fnames):
    """return list of sequences from all fnames"""
    seqs = {}
    try:
        for fname in fnames:
            for record in SeqIO.parse(fname, 'fasta'):
                if record.name in seqs and record.seq != seqs[record.name].seq:
                    raise AlignmentError("Detected duplicate input strains \"%s\" but the sequences are different." % record.name)
                    # if the same sequence then we can proceed (and we only take one)
                seqs[record.name] = record
    except FileNotFoundError:
        raise AlignmentError("\nCannot read sequences -- make sure the file %s exists and contains sequences in fasta format" % fname)
    except ValueError as error:
        raise AlignmentError("\nERROR: Problem reading in {}: {}".format(fname, str(error)))
    return list(seqs.values())

def check_arguments(args):
    # Simple error checking related to a reference name/sequence
    if args.reference_name and args.reference_sequence:
        raise AlignmentError("ERROR: You cannot provide both --reference-name and --reference-sequence")
    if args.remove_reference and not (args.reference_name or args.reference_sequence):
        raise AlignmentError("ERROR: You've asked to remove the reference but haven't specified one!")

def read_alignment(fname):
    try:
        return AlignIO.read(fname, 'fasta')
    except Exception as error:
        raise AlignmentError("\nERROR: Problem reading in {}: {}".format(fname, str(error)))

def ensure_reference_strain_present(ref_name, existing_alignment, seqs):
    if existing_alignment:
        if ref_name not in {x.name for x in existing_alignment}:
            raise AlignmentError("ERROR: Specified reference name %s (via --reference-name) is not in the supplied alignment."%ref_name)
    else:
        if ref_name not in {x.name for x in seqs}:
            raise AlignmentError("ERROR: Specified reference name %s (via --reference-name) is not in the sequence sample."%ref_name)

    # align
    # if args.method=='mafft':
    #     shoutput = shquote(output)
    #     shname = shquote(seq_fname)
    #     cmd = "mafft --reorder --anysymbol --thread %d %s 1> %s 2> %s.log"%(args.nthreads, shname, shoutput, shoutput)
def read_reference(ref_fname):
    if not os.path.isfile(ref_fname):
        raise AlignmentError("ERROR: Cannot read reference sequence."
                             "\n\tmake sure the file \"%s\" exists"%ref_fname)
    try:
        ref_seq = SeqIO.read(ref_fname, 'genbank' if ref_fname.split('.')[-1] in ['gb', 'genbank'] else 'fasta')
    except:
        raise AlignmentError("ERROR: Cannot read reference sequence."
                "\n\tmake sure the file %s contains one sequence in genbank or fasta format"%ref_fname)
    return ref_seq

def generate_alignment_cmd(method, nthreads, existing_aln_fname, seqs_to_align_fname, aln_fname, log_fname):
    if method=='mafft':
        if existing_aln_fname:
            cmd = "mafft --add %s --keeplength --reorder --anysymbol --nomemsave --adjustdirection --thread %d %s 1> %s 2> %s"%(shquote(seqs_to_align_fname), nthreads, shquote(existing_aln_fname), shquote(aln_fname), shquote(log_fname))
        else:
            cmd = "mafft --reorder --anysymbol --nomemsave --adjustdirection --thread %d %s 1> %s 2> %s"%(nthreads, shquote(seqs_to_align_fname), shquote(aln_fname), shquote(log_fname))
        print("\nusing mafft to align via:\n\t" + cmd +
              " \n\n\tKatoh et al, Nucleic Acid Research, vol 30, issue 14"
              "\n\thttps://doi.org/10.1093%2Fnar%2Fgkf436\n")
    else:
        raise AlignmentError('ERROR: alignment method %s not implemented'%method)
    return cmd

def strip_non_reference(aln, reference, keep_reference=False):
    '''
    return sequences that have all insertions relative to the reference
    removed. The alignment is read from file and returned as list of sequences.

    Parameters
    ----------
    aln : MultipleSeqAlign
        Biopython Alignment
    reference : str
        name of reference sequence, assumed to be part of the alignment
    keep_reference : bool, optional
        by default, the reference sequence is removed after stripping
        non-reference sequence. To keep the reference, use keep_reference=True

    Returns
    -------
    list
        list of trimmed sequences, effectively a multiple alignment

    Tests
    -----
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "with_gaps", keep_reference=False)]
    Trimmed gaps in with_gaps from the alignment
    ['no_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "with_gaps", keep_reference=True)]
    Trimmed gaps in with_gaps from the alignment
    ['with_gaps', 'no_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "no_gaps", keep_reference=True)]
    No gaps in alignment to trim (with respect to the reference, no_gaps)
    ['with_gaps', 'no_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "no_gaps", keep_reference=False)]
    No gaps in alignment to trim (with respect to the reference, no_gaps)
    ['with_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "missing", keep_reference=False)]
    Traceback (most recent call last):
      ...
    augur.align.AlignmentError: ERROR: reference missing not found in alignment
    '''
    seqs = {s.name:s for s in aln}
    if reference in seqs:
        ref_array = np.array(seqs[reference])
        if "-" not in ref_array:
            print("No gaps in alignment to trim (with respect to the reference, %s)"%reference)
            return [seq for seq in aln if (keep_reference or seq.name != reference)]
        ungapped = ref_array!='-'
        ref_aln_array = np.array(aln)[:,ungapped]
    else:
        raise AlignmentError("ERROR: reference %s not found in alignment"%reference)

    out_seqs = []
    for seq, seq_array in zip(aln, ref_aln_array):
        seq.seq = Seq.Seq(''.join(seq_array))
        if keep_reference or seq.name!=reference:
            out_seqs.append(seq)

    print("Trimmed gaps in", reference, "from the alignment")
    return out_seqs

def prettify_alignment(aln):
    '''
    Converts all bases to uppercase and removes auto reverse-complement prefix (_R_).
    This modifies the alignment in place.

    Parameters
    ----------
    aln : MultipleSeqAlign
        Biopython Alignment
    '''
    for seq in aln:
        seq.seq = seq.seq.upper()
        # AlignIO.read repeats the ID in the name and description field
        if seq.id.startswith("_R_"):
            seq.id = seq.id[3:]
            print("Sequence \"{}\" was reverse-complemented by the alignment program.".format(seq.id))
        if seq.name.startswith("_R_"):
            seq.name = seq.name[3:]
        if seq.description.startswith("_R_"):
            seq.description = seq.description[3:]

def make_gaps_ambiguous(aln):
    '''
    replace all gaps by 'N' in all sequences in the alignment. TreeTime will treat them
    as fully ambiguous and replace then with the most likely state. This modifies the
    alignment in place.

    Parameters
    ----------
    aln : MultipleSeqAlign
        Biopython Alignment
    '''
    for seq in aln:
        _seq = str(seq.seq)
        _seq = _seq.replace('-', 'N')
        seq.seq = Seq.Seq(_seq, alphabet=seq.seq.alphabet)

def check_duplicates(*values):
    names = set()
    def add(name):
        if name in names:
            raise AlignmentError("Duplicate strains of \"{}\" detected".format(name))
        names.add(name)
    for sample in values:
        if not sample:
            # allows false-like values (e.g. always provide existing_alignment, allowing
            # the default which is `False`)
            continue
        elif isinstance(sample, (list, Align.MultipleSeqAlignment)):
            for s in sample:
                add(s.name)
        elif isinstance(sample, str):
            add(sample)
        else:
            raise TypeError()

def write_seqs(seqs, fname):
    """A wrapper around SeqIO.write with error handling"""
    try:
        SeqIO.write(seqs, fname, 'fasta')
    except FileNotFoundError:
        raise AlignmentError('ERROR: Couldn\'t write "{}" -- perhaps the directory doesn\'t exist?'.format(fname))

def prune_seqs_matching_alignment(seqs, aln):
    """
    Return a set of seqs excluding those set via `exclude` & print a warning
    message for each sequence which is exluded.
    """
    ret = []
    exclude_names = {s.name for s in aln}
    for seq in seqs:
        if seq.name in exclude_names:
            print("Excluding {} as it is already present in the alignment".format(name))
        else:
            ret.append(seq)
    return ret
<commit_msg>Cleanup temp files<commit_after>"""
Align multiple sequences from FASTA.
"""

import os
from shutil import copyfile
import numpy as np
from Bio import AlignIO, SeqIO, Seq, Align
from .utils import run_shell_command, nthreads_value, shquote

class AlignmentError(Exception):
    # TODO: this exception should potentially be renamed and made augur-wide
    # thus allowing any module to raise it and have the message printed & augur
    # exit with code 1
    pass

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, nargs="+", metavar="FASTA", help="sequences to align")
    parser.add_argument('--output', '-o', default="alignment.fasta", help="output file (default: %(default)s)")
    parser.add_argument('--nthreads', type=nthreads_value, default=1,
                                help="number of threads to use; specifying the value 'auto' will cause the number of available CPU cores on your system, if determinable, to be used")
    parser.add_argument('--method', default='mafft', choices=["mafft"], help="alignment program to use")
    parser.add_argument('--reference-name', metavar="NAME", type=str, help="strip insertions relative to reference sequence; use if the reference is already in the input sequences")
    parser.add_argument('--reference-sequence', metavar="PATH", type=str, help="Add this reference sequence to the dataset & strip insertions relative to this. Use if the reference is NOT already in the input sequences")
    parser.add_argument('--remove-reference', action="store_true", default=False, help="remove reference sequence from the alignment")
    parser.add_argument('--fill-gaps', action="store_true", default=False, help="If gaps represent missing data rather than true indels, replace by N after aligning.")
    parser.add_argument('--existing-alignment', metavar="FASTA", default=False, help="An existing alignment to which the sequences will be added. The ouput alignment will be the same length as this existing alignment.")
    parser.add_argument('--debug', action="store_true", default=False, help="Produce extra files (e.g. pre- and post-aligner files) which can help with debugging poor alignments.")

def prepare(sequences, existing_aln_fname, output, ref_name, ref_seq_fname):
    seqs = read_sequences(*sequences)
    seqs_to_align_fname = output + ".to_align.fasta"

    # Load existing alignment
    existing_aln = None
    if existing_aln_fname:
        existing_aln = read_alignment(existing_aln_fname)

    # Load reference alignment
    ref_seq = None
    if ref_name:
        ensure_reference_strain_present(ref_name, existing_aln, seqs)
    elif ref_seq_fname:
        ref_seq = read_reference(ref_seq_fname)
        ref_name = ref_seq.id

    if existing_aln:
        # Strip the existing sequences from the new sequences, add the reference to the alignment
        seqs_to_align_fname = output + ".new_seqs_to_align.fasta"
        seqs = prune_seqs_matching_alignment(seqs, existing_aln)
        if ref_seq:
            if len(ref_seq) != existing_aln.get_alignment_length():
                raise AlignmentError("ERROR: Provided existing alignment ({}bp) is not the same length as the reference sequence ({}bp)".format(existing_aln.get_alignment_length(), len(ref_seq)))
            existing_aln_fname = existing_aln_fname + ".ref.fasta"
            existing_aln.append(ref_seq)
            write_seqs(existing_aln, existing_aln_fname)
    elif ref_seq: # Got a reference sequence but no existing alignment
        # reference sequence needs to be the first one for auto direction
        # adjustment (auto reverse-complement)
        seqs.insert(0, ref_seq)
    write_seqs(seqs, seqs_to_align_fname)

    check_duplicates(existing_aln, seqs)
    return existing_aln_fname, seqs_to_align_fname, ref_name

def run(args):
    '''
    Parameters
    ----------
    args : namespace
        arguments passed in via the command-line from augur

    Returns
    -------
    int
        returns 0 for success, 1 for general error
    '''
    temp_files_to_remove = []

    try:
        check_arguments(args)
        existing_aln_fname, seqs_to_align_fname, ref_name = prepare(args.sequences, args.existing_alignment, args.output, args.reference_name, args.reference_sequence)
        temp_files_to_remove.append(seqs_to_align_fname)
        if existing_aln_fname != args.existing_alignment:
            temp_files_to_remove.append(existing_aln_fname)
        # -- existing_aln_fname, seqs_to_align_fname, ref_name --

        # before aligning, make a copy of the data that the aligner receives as input (very useful for debugging purposes)
        if args.debug and not existing_aln_fname:
            copyfile(seqs_to_align_fname, args.output+".pre_aligner.fasta")

        # generate alignment command & run
        cmd = generate_alignment_cmd(args.method, args.nthreads, existing_aln_fname, seqs_to_align_fname, args.output, args.output+".log")
        success = run_shell_command(cmd)
        if not success:
            raise AlignmentError("Error during alignment")

        # after aligning, make a copy of the data that the aligner produced (useful for debugging)
        if args.debug:
            copyfile(args.output, args.output+".post_aligner.fasta")

        # -- ref_name --
        # reads the new alignment
        seqs = read_alignment(args.output)

        # convert the aligner output to upper case and remove auto reverse-complement prefix
        prettify_alignment(seqs)

        # if we've specified a reference, strip out all the columns not present in the reference
        # this will overwrite the alignment file
        if ref_name:
            seqs = strip_non_reference(seqs, ref_name, keep_reference=not args.remove_reference)
        if args.fill_gaps:
            make_gaps_ambiguous(seqs)

        # write the modified sequences back to the alignment file
        write_seqs(seqs, args.output)

    except AlignmentError as e:
        print(str(e))
        return 1

    # finally, remove any temporary files
    for fname in temp_files_to_remove:
        os.remove(fname)

#####################################################################################################

def read_sequences(*fnames):
    """return list of sequences from all fnames"""
    seqs = {}
    try:
        for fname in fnames:
            for record in SeqIO.parse(fname, 'fasta'):
                if record.name in seqs and record.seq != seqs[record.name].seq:
                    raise AlignmentError("Detected duplicate input strains \"%s\" but the sequences are different." % record.name)
                    # if the same sequence then we can proceed (and we only take one)
                seqs[record.name] = record
    except FileNotFoundError:
        raise AlignmentError("\nCannot read sequences -- make sure the file %s exists and contains sequences in fasta format" % fname)
    except ValueError as error:
        raise AlignmentError("\nERROR: Problem reading in {}: {}".format(fname, str(error)))
    return list(seqs.values())

def check_arguments(args):
    # Simple error checking related to a reference name/sequence
    if args.reference_name and args.reference_sequence:
        raise AlignmentError("ERROR: You cannot provide both --reference-name and --reference-sequence")
    if args.remove_reference and not (args.reference_name or args.reference_sequence):
        raise AlignmentError("ERROR: You've asked to remove the reference but haven't specified one!")

def read_alignment(fname):
    try:
        return AlignIO.read(fname, 'fasta')
    except Exception as error:
        raise AlignmentError("\nERROR: Problem reading in {}: {}".format(fname, str(error)))

def ensure_reference_strain_present(ref_name, existing_alignment, seqs):
    if existing_alignment:
        if ref_name not in {x.name for x in existing_alignment}:
            raise AlignmentError("ERROR: Specified reference name %s (via --reference-name) is not in the supplied alignment."%ref_name)
    else:
        if ref_name not in {x.name for x in seqs}:
            raise AlignmentError("ERROR: Specified reference name %s (via --reference-name) is not in the sequence sample."%ref_name)

    # align
    # if args.method=='mafft':
    #     shoutput = shquote(output)
    #     shname = shquote(seq_fname)
    #     cmd = "mafft --reorder --anysymbol --thread %d %s 1> %s 2> %s.log"%(args.nthreads, shname, shoutput, shoutput)
def read_reference(ref_fname):
    if not os.path.isfile(ref_fname):
        raise AlignmentError("ERROR: Cannot read reference sequence."
                             "\n\tmake sure the file \"%s\" exists"%ref_fname)
    try:
        ref_seq = SeqIO.read(ref_fname, 'genbank' if ref_fname.split('.')[-1] in ['gb', 'genbank'] else 'fasta')
    except:
        raise AlignmentError("ERROR: Cannot read reference sequence."
                "\n\tmake sure the file %s contains one sequence in genbank or fasta format"%ref_fname)
    return ref_seq

def generate_alignment_cmd(method, nthreads, existing_aln_fname, seqs_to_align_fname, aln_fname, log_fname):
    if method=='mafft':
        if existing_aln_fname:
            cmd = "mafft --add %s --keeplength --reorder --anysymbol --nomemsave --adjustdirection --thread %d %s 1> %s 2> %s"%(shquote(seqs_to_align_fname), nthreads, shquote(existing_aln_fname), shquote(aln_fname), shquote(log_fname))
        else:
            cmd = "mafft --reorder --anysymbol --nomemsave --adjustdirection --thread %d %s 1> %s 2> %s"%(nthreads, shquote(seqs_to_align_fname), shquote(aln_fname), shquote(log_fname))
        print("\nusing mafft to align via:\n\t" + cmd +
              " \n\n\tKatoh et al, Nucleic Acid Research, vol 30, issue 14"
              "\n\thttps://doi.org/10.1093%2Fnar%2Fgkf436\n")
    else:
        raise AlignmentError('ERROR: alignment method %s not implemented'%method)
    return cmd

def strip_non_reference(aln, reference, keep_reference=False):
    '''
    return sequences that have all insertions relative to the reference
    removed. The alignment is read from file and returned as list of sequences.

    Parameters
    ----------
    aln : MultipleSeqAlign
        Biopython Alignment
    reference : str
        name of reference sequence, assumed to be part of the alignment
    keep_reference : bool, optional
        by default, the reference sequence is removed after stripping
        non-reference sequence. To keep the reference, use keep_reference=True

    Returns
    -------
    list
        list of trimmed sequences, effectively a multiple alignment

    Tests
    -----
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "with_gaps", keep_reference=False)]
    Trimmed gaps in with_gaps from the alignment
    ['no_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "with_gaps", keep_reference=True)]
    Trimmed gaps in with_gaps from the alignment
    ['with_gaps', 'no_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "no_gaps", keep_reference=True)]
    No gaps in alignment to trim (with respect to the reference, no_gaps)
    ['with_gaps', 'no_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "no_gaps", keep_reference=False)]
    No gaps in alignment to trim (with respect to the reference, no_gaps)
    ['with_gaps', 'some_other_seq', '_R_crick_strand']
    >>> [s.name for s in strip_non_reference(read_alignment("tests/data/align/test_aligned_sequences.fasta"), "missing", keep_reference=False)]
    Traceback (most recent call last):
      ...
    augur.align.AlignmentError: ERROR: reference missing not found in alignment
    '''
    seqs = {s.name:s for s in aln}
    if reference in seqs:
        ref_array = np.array(seqs[reference])
        if "-" not in ref_array:
            print("No gaps in alignment to trim (with respect to the reference, %s)"%reference)
            return [seq for seq in aln if (keep_reference or seq.name != reference)]
        ungapped = ref_array!='-'
        ref_aln_array = np.array(aln)[:,ungapped]
    else:
        raise AlignmentError("ERROR: reference %s not found in alignment"%reference)

    out_seqs = []
    for seq, seq_array in zip(aln, ref_aln_array):
        seq.seq = Seq.Seq(''.join(seq_array))
        if keep_reference or seq.name!=reference:
            out_seqs.append(seq)

    print("Trimmed gaps in", reference, "from the alignment")
    return out_seqs

def prettify_alignment(aln):
    '''
    Converts all bases to uppercase and removes auto reverse-complement prefix (_R_).
    This modifies the alignment in place.

    Parameters
    ----------
    aln : MultipleSeqAlign
        Biopython Alignment
    '''
    for seq in aln:
        seq.seq = seq.seq.upper()
        # AlignIO.read repeats the ID in the name and description field
        if seq.id.startswith("_R_"):
            seq.id = seq.id[3:]
            print("Sequence \"{}\" was reverse-complemented by the alignment program.".format(seq.id))
        if seq.name.startswith("_R_"):
            seq.name = seq.name[3:]
        if seq.description.startswith("_R_"):
            seq.description = seq.description[3:]

def make_gaps_ambiguous(aln):
    '''
    replace all gaps by 'N' in all sequences in the alignment. TreeTime will treat them
    as fully ambiguous and replace then with the most likely state. This modifies the
    alignment in place.

    Parameters
    ----------
    aln : MultipleSeqAlign
        Biopython Alignment
    '''
    for seq in aln:
        _seq = str(seq.seq)
        _seq = _seq.replace('-', 'N')
        seq.seq = Seq.Seq(_seq, alphabet=seq.seq.alphabet)

def check_duplicates(*values):
    names = set()
    def add(name):
        if name in names:
            raise AlignmentError("Duplicate strains of \"{}\" detected".format(name))
        names.add(name)
    for sample in values:
        if not sample:
            # allows false-like values (e.g. always provide existing_alignment, allowing
            # the default which is `False`)
            continue
        elif isinstance(sample, (list, Align.MultipleSeqAlignment)):
            for s in sample:
                add(s.name)
        elif isinstance(sample, str):
            add(sample)
        else:
            raise TypeError()

def write_seqs(seqs, fname):
    """A wrapper around SeqIO.write with error handling"""
    try:
        SeqIO.write(seqs, fname, 'fasta')
    except FileNotFoundError:
        raise AlignmentError('ERROR: Couldn\'t write "{}" -- perhaps the directory doesn\'t exist?'.format(fname))

def prune_seqs_matching_alignment(seqs, aln):
    """
    Return a set of seqs excluding those set via `exclude` & print a warning
    message for each sequence which is exluded.
    """
    ret = []
    exclude_names = {s.name for s in aln}
    for seq in seqs:
        if seq.name in exclude_names:
            print("Excluding {} as it is already present in the alignment".format(name))
        else:
            ret.append(seq)
    return ret
<|endoftext|>


<commit_before>"""Count occurrence of bases in a set of sequences.
"""

from Bio import SeqIO
from itertools import combinations
import Bio.Seq
import Bio.SeqRecord
import sys
import csv

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, help="sequences in fasta format")
    parser.add_argument('--output', '-o', help="tab-delimited file containing the number of bases per sequence in the given file. Output columns include strain, length, and counts for A, C, G, T, N, other valid IUPAC characters, ambiguous characters ('?' and '-'), and other invalid characters.", required=True)
    parser.add_argument('--verbose', '-v', action="store_true", help="print index statistics to stdout")

def index_sequence(sequence, values):
    """Count the number of nucleotides for a given sequence record.

    Parameters
    ----------
    sequence : Bio.SeqRecord.SeqRecord
        sequence record to index.

    values : list of sets of str
        values to count; sets must be non-overlapping and contain only
        single-character, lowercase strings

    Returns
    -------
    list :
        summary of the given sequence's strain name, length, nucleotide counts
        for the given values, and a final column with the number of characters
        that didn't match any of those in the given values.

    >>> other_IUPAC = {'r', 'y', 's', 'w', 'k', 'm', 'd', 'h', 'b', 'v'}
    >>> values = [{'a'},{'c'},{'g'},{'t'},{'n'}, other_IUPAC, {'-'}, {'?'}]
    >>> sequence_a = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTGN-?XWN"), id="seq_A")
    >>> index_sequence(sequence_a, values)
    ['seq_A', 10, 1, 1, 1, 1, 2, 1, 1, 1, 1]

    >>> sequence_b = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTGACTG"), id="seq_B")
    >>> index_sequence(sequence_b, values)
    ['seq_B', 8, 2, 2, 2, 2, 0, 0, 0, 0, 0]

    Characters in the given sequence that are not in the given list of values to
    count get counted in the final column.

    >>> sequence_c = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTG%@!!!NN"), id="seq_C")
    >>> index_sequence(sequence_c, values)
    ['seq_C', 11, 1, 1, 1, 1, 2, 0, 0, 0, 5]

    The list of value sets must not overlap.

    >>> sequence_d = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("A!C!TGXN"), id="seq_D")
    >>> index_sequence(sequence_d, [set('actg'), set('xn'), set('n')]) # doctest: +ELLIPSIS
    Traceback (most recent call last):
      ...
    ValueError: character sets ... and {'n'} overlap: {'n'}

    Value sets must contain only single-character, lowercase strings.

    >>> index_sequence(sequence_d, [{'a'}, {'c'}, {'T'}, {'g'}])
    Traceback (most recent call last):
      ...
    ValueError: character set {'T'} contains a non-lowercase character: 'T'

    >>> index_sequence(sequence_d, [{'actg'}])
    Traceback (most recent call last):
      ...
    ValueError: character set {'actg'} contains a multi-character (or maybe zero-length) string: 'actg'

    >>> index_sequence(sequence_d, [{'a', 'c'}, {0, 1}])
    Traceback (most recent call last):
      ...
    ValueError: character set {0, 1} contains a non-string element: 0
    """
    # Sets must be non-overlapping and contain only single-character, lowercase
    # strings, otherwise our assumptions are broken.  This is not much hardship
    # on the caller.
    for v in values:
        for c in v:
            if not isinstance(c, str):
                raise ValueError(f"character set {v!r} contains a non-string element: {c!r}")
            if len(c) != 1:
                raise ValueError(f"character set {v!r} contains a multi-character (or maybe zero-length) string: {c!r}")
            if c != c.lower():
                raise ValueError(f"character set {v!r} contains a non-lowercase character: {c!r}")

    for a, b in combinations(values, 2):
        if not a.isdisjoint(b):
            raise ValueError(f"character sets {a!r} and {b!r} overlap: {a & b!r}")

    counts = []
    seq = sequence.seq.lower()
    l = len(seq)

    for v in values:
        counts.append(sum(map(lambda x: seq.count(x), v)))

    invalid_nucleotides = l-sum(counts)
    row = [sequence.id, l]+counts+[invalid_nucleotides]

    return row

def index_sequences(sequences_path, sequence_index_path):
    """Count the number of A, C, T, G, N, other IUPAC nucleotides, ambiguous bases
    ("?" and "-"), and other invalid characters in a set of sequences and write
    the composition as a data frame to the given sequence index path.

    Parameters
    ----------
    sequences_path : str or Path-like
        path to a sequence file to index.

    sequence_index_path : str or Path-like
        path to a tab-delimited file containing the composition details for each
        sequence in the given input file.

    Returns
    -------
    int :
        number of sequences indexed

    int :
        total length of sequences indexed

    """
    #read in files
    try:
        seqs = SeqIO.parse(sequences_path, 'fasta')
    except ValueError as error:
        print("ERROR: Problem reading in {}:".format(sequences_path), file=sys.stderr)
        print(error, file=sys.stderr)
        return 1

    other_IUPAC = {'r', 'y', 's', 'w', 'k', 'm', 'd', 'h', 'b', 'v'}
    values = [{'a'},{'c'},{'g'},{'t'},{'n'},other_IUPAC,{'-'},{'?'}]
    labels = ['A','C','G','T','N','other_IUPAC','-','?']

    tot_length = 0
    num_of_seqs = 0

    with open(sequence_index_path, 'wt') as out_file:
        tsv_writer = csv.writer(out_file, delimiter = '\t')

        #write header i output file
        header = ['strain', 'length']+labels+['invalid_nucleotides']
        tsv_writer.writerow(header)

        for record in seqs:
            #index the sequence and write row in output file
            row = index_sequence(record, values)
            tsv_writer.writerow(row)

            tot_length += row[1]
            num_of_seqs += 1

    return num_of_seqs, tot_length

def run(args):
    '''
    runs index_sequences which counts the number of A, C, T, G, N, other IUPAC nucleotides, ambiguous bases
    ("?" and "-"), and other invalid characters in a set of sequences and write
    the composition as a data frame to the given sequence index path.
    '''
    num_of_seqs, tot_length = index_sequences(args.sequences, args.output)

    if args.verbose:
        print("Analysed %i sequences with an average length of %i nucleotides." % (num_of_seqs, int(tot_length / num_of_seqs)))
<commit_msg>Support compressed inputs/outputs for index<commit_after>"""Count occurrence of bases in a set of sequences.
"""

from Bio import SeqIO
from itertools import combinations
import Bio.Seq
import Bio.SeqRecord
import sys
import csv

from .io import open_file, read_sequences

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, help="sequences in fasta format")
    parser.add_argument('--output', '-o', help="tab-delimited file containing the number of bases per sequence in the given file. Output columns include strain, length, and counts for A, C, G, T, N, other valid IUPAC characters, ambiguous characters ('?' and '-'), and other invalid characters.", required=True)
    parser.add_argument('--verbose', '-v', action="store_true", help="print index statistics to stdout")

def index_sequence(sequence, values):
    """Count the number of nucleotides for a given sequence record.

    Parameters
    ----------
    sequence : Bio.SeqRecord.SeqRecord
        sequence record to index.

    values : list of sets of str
        values to count; sets must be non-overlapping and contain only
        single-character, lowercase strings

    Returns
    -------
    list :
        summary of the given sequence's strain name, length, nucleotide counts
        for the given values, and a final column with the number of characters
        that didn't match any of those in the given values.

    >>> other_IUPAC = {'r', 'y', 's', 'w', 'k', 'm', 'd', 'h', 'b', 'v'}
    >>> values = [{'a'},{'c'},{'g'},{'t'},{'n'}, other_IUPAC, {'-'}, {'?'}]
    >>> sequence_a = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTGN-?XWN"), id="seq_A")
    >>> index_sequence(sequence_a, values)
    ['seq_A', 10, 1, 1, 1, 1, 2, 1, 1, 1, 1]

    >>> sequence_b = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTGACTG"), id="seq_B")
    >>> index_sequence(sequence_b, values)
    ['seq_B', 8, 2, 2, 2, 2, 0, 0, 0, 0, 0]

    Characters in the given sequence that are not in the given list of values to
    count get counted in the final column.

    >>> sequence_c = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTG%@!!!NN"), id="seq_C")
    >>> index_sequence(sequence_c, values)
    ['seq_C', 11, 1, 1, 1, 1, 2, 0, 0, 0, 5]

    The list of value sets must not overlap.

    >>> sequence_d = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("A!C!TGXN"), id="seq_D")
    >>> index_sequence(sequence_d, [set('actg'), set('xn'), set('n')]) # doctest: +ELLIPSIS
    Traceback (most recent call last):
      ...
    ValueError: character sets ... and {'n'} overlap: {'n'}

    Value sets must contain only single-character, lowercase strings.

    >>> index_sequence(sequence_d, [{'a'}, {'c'}, {'T'}, {'g'}])
    Traceback (most recent call last):
      ...
    ValueError: character set {'T'} contains a non-lowercase character: 'T'

    >>> index_sequence(sequence_d, [{'actg'}])
    Traceback (most recent call last):
      ...
    ValueError: character set {'actg'} contains a multi-character (or maybe zero-length) string: 'actg'

    >>> index_sequence(sequence_d, [{'a', 'c'}, {0, 1}])
    Traceback (most recent call last):
      ...
    ValueError: character set {0, 1} contains a non-string element: 0
    """
    # Sets must be non-overlapping and contain only single-character, lowercase
    # strings, otherwise our assumptions are broken.  This is not much hardship
    # on the caller.
    for v in values:
        for c in v:
            if not isinstance(c, str):
                raise ValueError(f"character set {v!r} contains a non-string element: {c!r}")
            if len(c) != 1:
                raise ValueError(f"character set {v!r} contains a multi-character (or maybe zero-length) string: {c!r}")
            if c != c.lower():
                raise ValueError(f"character set {v!r} contains a non-lowercase character: {c!r}")

    for a, b in combinations(values, 2):
        if not a.isdisjoint(b):
            raise ValueError(f"character sets {a!r} and {b!r} overlap: {a & b!r}")

    counts = []
    seq = sequence.seq.lower()
    l = len(seq)

    for v in values:
        counts.append(sum(map(lambda x: seq.count(x), v)))

    invalid_nucleotides = l-sum(counts)
    row = [sequence.id, l]+counts+[invalid_nucleotides]

    return row

def index_sequences(sequences_path, sequence_index_path):
    """Count the number of A, C, T, G, N, other IUPAC nucleotides, ambiguous bases
    ("?" and "-"), and other invalid characters in a set of sequences and write
    the composition as a data frame to the given sequence index path.

    Parameters
    ----------
    sequences_path : str or Path-like
        path to a sequence file to index.

    sequence_index_path : str or Path-like
        path to a tab-delimited file containing the composition details for each
        sequence in the given input file.

    Returns
    -------
    int :
        number of sequences indexed

    int :
        total length of sequences indexed

    """
    seqs = read_sequences(sequences_path)

    other_IUPAC = {'r', 'y', 's', 'w', 'k', 'm', 'd', 'h', 'b', 'v'}
    values = [{'a'},{'c'},{'g'},{'t'},{'n'},other_IUPAC,{'-'},{'?'}]
    labels = ['A','C','G','T','N','other_IUPAC','-','?']

    tot_length = 0
    num_of_seqs = 0

    with open_file(sequence_index_path, 'wt') as out_file:
        tsv_writer = csv.writer(out_file, delimiter = '\t')

        #write header i output file
        header = ['strain', 'length']+labels+['invalid_nucleotides']
        tsv_writer.writerow(header)

        for record in seqs:
            #index the sequence and write row in output file
            row = index_sequence(record, values)
            tsv_writer.writerow(row)

            tot_length += row[1]
            num_of_seqs += 1

    return num_of_seqs, tot_length

def run(args):
    '''
    runs index_sequences which counts the number of A, C, T, G, N, other IUPAC nucleotides, ambiguous bases
    ("?" and "-"), and other invalid characters in a set of sequences and write
    the composition as a data frame to the given sequence index path.
    '''
    try:
        num_of_seqs, tot_length = index_sequences(args.sequences, args.output)
    except ValueError as error:
        print("ERROR: Problem reading in {}:".format(sequences_path), file=sys.stderr)
        print(error, file=sys.stderr)
        return 1

    if args.verbose:
        print("Analysed %i sequences with an average length of %i nucleotides." % (num_of_seqs, int(tot_length / num_of_seqs)))
<|endoftext|>


<commit_before># Basic idea is:
# Always scan H and W separately, i.e. consider as lines (1d > 2d)
# 0. An array of pixels and scores (-x = static, 0 = average/confused, x = moving)
# 1. Remove borders
# 2. Attempt to break into regions as follows:
#		a. For each line, require the next 5 (or so) lines to be excluded.
#		b. If including the next line (and excluding line n+6) would improve or not change your score, iterate.
#		c. When including the next line has lower score (presumably line n+6 has no additional cost), stop
#		d. Region is through n, restart process from 1, unless region is empty.
# 3. If larger range is broken into >1 region, pass regions into the alternate direction scan method.
#	4. If not, pass into border-stripping method, and then save final region into a list.
#	5. Stitch regions together (?)

# Removes pixels which are "bad" from the edge. Note: Trimming twice may remove "good" pixels, since there is no real buffer to use as comparison
def trim(xmin, xmax, ymin, ymax):
	v_line_scores = [0]*(xmax-xmin)
	h_line_scores = [0]*(ymax-ymin)
	for x in range(xmin, xmax):
		for y in range(ymin, ymax):
			v_line_scores[x-xmin] += regiondata[x][y]
			h_line_scores[y-ymin] += regiondata[x][y]

	# Unfortunately, the order of these loops matter, because this loop references xmin, which may be changed by the second loop.
	for x in range(xmax-1, xmin, -1):
		if v_line_scores[-1] < v_line_scores[x-xmin]:
			xmax = x+1
			break

	for x in range(xmin, xmax):
		if v_line_scores[0] < v_line_scores[x-xmin]:
			xmin = x
			break

	for y in range(ymax-1, ymin, -1):
		if h_line_scores[-1] < h_line_scores[y-ymin]:
			ymax = y+1
			break

	for y in range(ymin, ymax):
		if h_line_scores[0] < h_line_scores[y-ymin]:
			ymin = y
			break

	return (xmin, xmax, ymin, ymax)

def debug(xmin, xmax, ymin, ymax):
	for x in range(xmin, xmax):
		print regiondata[x][ymin:ymax]
	print

if __name__ == '__main__':
	# This should be generated by taking screenshots and comparing pixels. Somehow...
	global regiondata
	regiondata = []

	for line in '''
	00000000000000000000000000
	00000000000000000000000000
	00000000111100000000000000
	00000000111100000000000000
	00000000111100011110000000
	00000000111100011110000000
	00000000000000011110000000
	00000000000000011110000000
	00000000000000000000000000
	00000000000000000000000000
	00000000000000000000000000
	'''.strip().split('\n'):
		data = []
		for char in line.strip():
			if char == '0':
				data.append(-5)
			elif char == '1':
				data.append(5)
		regiondata.append(data)

	xmin, xmax, ymin, ymax = (0, len(regiondata), 0, len(regiondata[0]))
	debug(xmin, xmax, ymin, ymax)
	xmin, xmax, ymin, ymax = trim(xmin, xmax, ymin, ymax)
	debug(xmin, xmax, ymin, ymax)
<commit_msg>Added basic detection algorithms, and a merge algorithm.<commit_after># How to generate?
#		Take 10 screenshots.
#		For each pixel, count how many screenshots are different (multiplicitively, i.e. all different = 45)

# Deal with mush (i.e. possible matches)
# Deal with missing corners

from threading import Thread
global MIN_VALIDITY
MIN_VALIDITY = .2

# Removes pixels which are "bad" from the edge. Note: Trimming twice may remove "good" pixels, since there is no real buffer to use as comparison
def trim(xmin, xmax, ymin, ymax):
	v_line_scores = [0]*(xmax-xmin)
	h_line_scores = [0]*(ymax-ymin)
	for x in range(xmin, xmax):
		for y in range(ymin, ymax):
			v_line_scores[x-xmin] += regiondata[x][y]
			h_line_scores[y-ymin] += regiondata[x][y]

	# Unfortunately, the order of these loops matter, because the first loop references xmin, which may be changed by the second loop.
	for x in range(xmax-1, xmin, -1):
		if MIN_VALIDITY * len(v_line_scores) < v_line_scores[x-xmin]:
			xmax = x+1
			break
	for x in range(xmin, xmax):
		if MIN_VALIDITY * len(v_line_scores) < v_line_scores[x-xmin]:
			xmin = x
			break

	for y in range(ymax-1, ymin, -1):
		if MIN_VALIDITY * len(h_line_scores) < h_line_scores[y-ymin]:
			ymax = y+1
			break
	for y in range(ymin, ymax):
		if MIN_VALIDITY * len(h_line_scores) < h_line_scores[y-ymin]:
			ymin = y
			break

	return (xmin, xmax, ymin, ymax)

global regions
regions = []

# Attempts to split the given region along vertical lines.
def h_split(xmin, xmax, ymin, ymax):
	xmin, xmax, ymin, ymax = trim(xmin, xmax, ymin, ymax)
	h_line_scores = [0]*(xmax-xmin)
	for x in range(xmin, xmax):
		for y in range(ymin, ymax):
			h_line_scores[x-xmin] += regiondata[x][y]
	splits = []
	for x in range(xmin, xmax-1):
		if (1-MIN_VALIDITY) * h_line_scores[x-xmin] > h_line_scores[x-xmin+1]: # Appears to be the last line in a block
			splits.append(x+1)
	splits = [xmin] + splits + [xmax]
	regions = []
	threads = []
	for i in range(len(splits)-1):
		regions += v_split(splits[i], splits[i+1], ymin, ymax)
		thread = Thread(target=v_split, kwargs={
		'xmin': splits[i],
		'xmax': splits[i+1],
		'ymin': ymin,
		'ymax': ymax})
		thread.start()
		threads.append(thread)
	for thread in threads:
		thread.join()
	return regions

# Attempts to split the given region along horizontal lines.
def v_split(xmin, xmax, ymin, ymax):
	xmin, xmax, ymin, ymax = trim(xmin, xmax, ymin, ymax)
	v_line_scores = [0]*(ymax-ymin)
	for x in range(xmin, xmax):
		for y in range(ymin, ymax):
			v_line_scores[y-ymin] += regiondata[x][y]
	splits = []
	for y in range(ymin, ymax-1):
		if (1-MIN_VALIDITY) * v_line_scores[y-ymin] > v_line_scores[y-ymin+1]: # Appears to be the last line in a block
			splits.append(y+1)
	regions = []
	if len(splits) > 0: # If we split into >1 part, then recurse. Otherwise, terminate.
		splits = [ymin] + splits + [ymax]
		threads = []
		for i in range(len(splits)-1):
			regions += h_split(xmin, xmax, splits[i], splits[i+1])
			thread = Thread(target=h_split, kwargs={
			'xmin': xmin,
			'xmax': xmax,
			'ymin': splits[i],
			'ymax': splits[i+1]})
			thread.start()
		for thread in threads:
		 	thread.join()
		return regions
	else: # Region was already split vertically, if there are no horizontal splits, then it's finalized.
		return [[xmin, xmax, ymin, ymax]]

def merge(regions):
	new_region = None
	for i in range(len(regions)):
		for j in range(i):
			if (regions[i][0], regions[i][1]) == (regions[j][0], regions[j][1]):
				# i[xmin]=j[xmin], i[xmax]=j[xmax]
				if regions[i][3] == regions[j][2]:
					# i[ymax] == j[ymin]
					new_region = [regions[i][0], regions[i][1], regions[i][2], regions[j][3]]
				elif regions[j][3] == regions[i][2]:
					# j[ymax] == i[ymin]
					new_region = [regions[i][0], regions[i][1], regions[j][2], regions[i][3]]
			elif (regions[i][2], regions[i][3]) == (regions[j][2], regions[j][3]):
				# i[ymin]=j[ymin], i[ymax]=j[ymax]
				if regions[i][1] == regions[j][0]:
					# i[xmax] == j[xmin]
					new_region = [regions[i][0], regions[j][1], regions[i][2], regions[i][3]]
				elif regions[j][1] == regions[i][0]:
					# j[xmax] == i[xmin]
					new_region = [regions[j][0], regions[i][1], regions[i][2], regions[i][3]]
			if new_region != None:
				regions.append(new_region)
				regions.pop(i)
				regions.pop(j)
				return merge(regions)
	return regions

def debug(xmin, xmax, ymin, ymax):
	for x in range(xmin, xmax):
		line = ''
		for char in regiondata[x][ymin:ymax]:
			line += str(char)
		print line
	print

if __name__ == '__main__':
	# This should be generated by taking screenshots and comparing pixels, somehow...
	global regiondata
	regiondata = []

	for line in '''
	0000000000000000000
	0000000000000000000
	0000000005555500000
	0000000005555500000
	0000000005555500000
	0044444005555500000
	0044444005555500000
	0044444000000000000
	0044444000000000000
	0044444000066666600
	0000000000066666600
	0000000000066666600
	0000777770066666600
	0000777770066666600
	0000777770000000000
	0000777770000000000
	0000777770000000000
	0000000000000000000
	'''.strip().split('\n'):
		data = []
		from random import random
		for char in line.strip():
			# if char == '0' and random()*20 < 1:
			# 	char = '1'
			data.append(int(char))
		regiondata.append(data)

	xmin, xmax, ymin, ymax = (0, len(regiondata), 0, len(regiondata[0]))
	debug(xmin, xmax, ymin, ymax)
	regions = h_split(xmin, xmax, ymin, ymax)
	regions = merge(regions)
	for region in regions:
		debug(region[0], region[1], region[2], region[3])
<|endoftext|>


<commit_before>"""Autoencoders, denoising autoencoders, and stacked DAEs."""
# Standard library imports
from itertools import izip

# Third-party imports
import numpy
import theano
from theano import tensor
from pylearn.gd.sgd import sgd_updates

# Local imports
from base import Block, Trainer

def safe_update(dict_to, dict_from):
    """
    Like dict_to.update(dict_from), except don't overwrite any keys.
    """
    for key, val in dict(dict_from).iteritems():
        if key in dict_to:
            raise KeyError(key)
        dict_to[key] = val
    return dict_to

theano.config.warn.sum_div_dimshuffle_bug = False
floatX = theano.config.floatX
sharedX = lambda X, name: theano.shared(numpy.asarray(X, dtype=floatX), name=name)
if 0:
    print 'WARNING: using SLOW rng'
    RandomStreams = tensor.shared_randomstreams.RandomStreams
else:
    import theano.sandbox.rng_mrg
    RandomStreams = theano.sandbox.rng_mrg.MRG_RandomStreams

class DenoisingAutoencoder(Block):
    """
    A denoising autoencoder learns a representation of the input by
    reconstructing a noisy version of it.
    """
    def __init__(self, **kwargs):
        # TODO: Do we need anything else here?
        super(DenoisingAutoencoder, self).__init__(**kwargs)

    @classmethod
    def alloc(cls, corruptor, conf, rng=None):
        """Allocate a denoising autoencoder object."""
        if not hasattr(rng, 'randn'):
            rng = numpy.random.RandomState(rng)
        self = cls()
        self.corruptor = corruptor
        self.visbias = sharedX(
            numpy.zeros(conf['n_vis']),
            name='vb'
        )
        self.hidbias = sharedX(
            numpy.zeros(conf['n_hid']),
            name='hb'
        )
        self.weights = sharedX(
            .5 * rng.rand(conf['n_vis'], conf['n_hid']) * conf['irange'],
            name='W'
        )
        seed = int(rng.randint(2**30))
        self.s_rng = RandomStreams(seed)
        if conf['tied_weights']:
            self.w_prime = self.weights.T
        else:
            self.w_prime = sharedX(
                .5 * rng.rand(conf['n_hid'], conf['n_vis']) * conf['irange'],
                name='Wprime'
            )

        def _resolve_callable(conf_attr):
            if conf[conf_attr] is None:
                # The identity function, for linear layers.
                return lambda x: x
            # If it's a callable, use it directly.
            if hasattr(conf[conf_attr], '__call__'):
                return conf[conf_attr]
            elif hasattr(tensor.nnet, conf[conf_attr]):
                return getattr(tensor.nnet, conf[conf_attr])
            elif hasattr(tensor, conf[conf_attr]):
                return getattr(tensor, conf[conf_attr])
            else:
                raise ValueError("Couldn't interpret %s value: '%s'" %
                                 (conf_attr, conf[conf_attr]))

        self.act_enc = _resolve_callable('act_enc')
        self.act_dec = _resolve_callable('act_dec')
        self.conf = conf
        self._params = [
            self.visbias,
            self.hidbias,
            self.weights,
        ]
        if not conf['tied_weights']:
            self._params.append(self.w_prime)
        return self

    def _hidden_activation(self, x):
        """Single input pattern/minibatch activation function."""
        return self.act_enc(self.hidbias + tensor.dot(x, self.weights))

    def hidden_repr(self, inputs):
        """Hidden unit activations for each set of inputs."""
        return [self._hidden_activation(v) for v in inputs]

    def reconstruction(self, inputs):
        """Reconstructed inputs after corruption."""
        corrupted = self.corruptor(inputs)
        hiddens = self.hidden_repr(corrupted)
        return [
            self.act_dec(self.visbias + tensor.dot(h, self.w_prime))
            for h in hiddens
        ]

    def __call__(self, inputs):
        """
        Forward propagate (symbolic) input through this module, obtaining
        a representation to pass on to layers above.
        """
        return self.hidden_repr(inputs)

    def mse(self, inputs):
        """
        Symbolic expression for mean-squared error between the input and the
        denoised reconstruction.
        """
        pairs = izip(inputs, self.reconstruction(inputs))
        return [((inp - rec)**2).sum(axis=-1).mean() for inp, rec in pairs]

    def cross_entropy(self, inputs):
        """
        Symbolic expression for elementwise cross-entropy between input
        and reconstruction. Use for binary-valued features (but not for,
        e.g., one-hot codes).
        """
        pairs = izip(inputs, self.reconstruction(inputs))
        ce = lambda x, z: x * tensor.log(z) + (1 - x) * tensor.log(1 - z)
        return [ce(inp, rec).sum(axis=1).mean() for inp, rec in pairs]

class StackedDA(Block):
    """
    A class representing a stacked model. Forward propagation passes
    (symbolic) input through each layer sequentially.
    """
    def __init__(self, **kwargs):
        # TODO: Do we need anything else here?
        super(StackedDA, self).__init__(**kwargs)

    @classmethod
    def alloc(cls, corruptors, conf, rng=None):
        """Allocate a stacked denoising autoencoder object."""
        if not hasattr(rng, 'randn'):
            rng = numpy.random.RandomState(rng)
        self = cls()
        self._layers = []
        _local = {}
        # Make sure that if we have a sequence of encoder/decoder activations
        # or corruptors, that we have exactly as many as len(conf['n_hid'])
        if hasattr(corruptors, '__len__'):
            assert len(conf['n_hid']) == len(corruptors)
        else:
            corruptors = [corruptors] * len(conf['n_hid'])
        for c in ['act_enc', 'act_dec']:
            if type(conf[c]) is not str and hasattr(conf[c], '__len__'):
                assert len(conf['n_hid']) == len(conf[c])
                _local[c] = conf[c]
            else:
                _local[c] = [conf[c]] * len(conf['n_hid'])
        n_hids = conf['n_hid']
        # The number of visible units in each layer is the initial input
        # size and the first k-1 hidden unit sizes.
        n_viss = [conf['n_vis']] + conf['n_hid'][:-1]
        seq = izip(
            xrange(len(n_hids)),
            n_hids,
            n_viss,
            _local['act_enc'],
            _local['act_dec'],
            corruptors
        )
        # Create each layer.
        for k, n_hid, n_vis, act_enc, act_dec, corr in seq:
            # Create a local configuration dictionary for this layer.
            lconf = {
                'n_hid': n_hid,
                'n_vis': n_vis,
                'act_enc': act_enc,
                'act_dec': act_dec,
            }
            da = DenoisingAutoencoder.alloc(corr, lconf, rng)
            self._layers.append(da)
        return self

    def layers(self):
        """
        The layers of this model: the individual denoising autoencoder
        objects, which can be individually pre-trained.
        """
        return list(self._layers)

    def params(self):
        """
        The parameters that are learned in this model, i.e. the concatenation
        of all the layers' weights and biases.
        """
        return sum([l.params() for l in self._layers], [])

    def __call__(self, inputs):
        """
        Forward propagate (symbolic) input through this module, obtaining
        a representation to pass on to layers above.
        """
        transformed = inputs
        # Pass the input through each layer of the hierarchy.
        for layer in self._layers:
            transformed = layer(transformed)
        return transformed

class DATrainer(Trainer):
    @classmethod
    def alloc(cls, model, cost, minibatch, conf):
        """
        Takes a DenoisingAutoencoder object, a (symbolic) cost function
        which takes the input and turns it into a scalar (somehow),
        a (symbolic) minibatch, and a configuration dictionary.
        """
        # Take care of learning rate scales for individual parameters
        learning_rates = {}
        for parameter in model.params():
            lr_name = '%s_lr' % parameter.name
            try:
                thislr = conf[lr_name]
            except:
                thislr = 1.
            learning_rates[parameter] = sharedX(thislr, lr_name)

        # A shared variable for storing the iteration number.
        iteration = sharedX(0, 'iter')

        # A shared variable for storing the annealed base learning rate, used
        # to lower the learning rate gradually after a certain amount of time.
        annealed = sharedX(conf['base_lr'], 'annealed')

        # Instantiate the class, finally.
        self = cls(model=model, cost=cost, conf=conf,
                   learning_rates=learning_rates,
                   iteration=iteration, minibatch=minibatch)

    def updates(self):
        """Compute the updates for each of the parameter variables."""
        ups = {}
        # Base learning rate per example.
        base_lr = numpy.asarray(self.conf['base_lr'], dtype=floatX)

        # Annealing coefficient. Here we're using a formula of
        # base_lr * min(0.0, max(base_lr, lr_anneal_start / (iteration + 1))
        frac = self.conf['lr_anneal_start'] / (self.iteration + 1.)
        annealed = tensor.clip(
            tensor.cast(frac, floatX),
            0.0,    # minimum learning rate
            base_lr # maximum learning rate
        )

        # Update the shared variable for the annealed learning rate.
        ups[self.annealed] = annealed
        ups[self.iteration] = self.iteration + 1

        # Calculate the learning rates for each parameter, in the order
        # they appear in model.params()
        learn_rates = [annealed * self.lr_dict[p] for p in self.model.params()]
        # Get the gradient w.r.t. cost of each parameter.
        grads = [
            tensor.grad(self.cost(self.minibatch), p)
            for p in self.model.params()
        ]
        # Get the updates from sgd_updates, a PyLearn library function.
        updates = dict(sgd_updates(self.model.params(), grads, learn_rates))

        # Return the updates dictionary.
        return ups

<commit_msg>Pass along initialization range and tied_weights configuration variables to layers of a stacked network.<commit_after>"""Autoencoders, denoising autoencoders, and stacked DAEs."""
# Standard library imports
from itertools import izip

# Third-party imports
import numpy
import theano
from theano import tensor
from pylearn.gd.sgd import sgd_updates

# Local imports
from base import Block, Trainer

def safe_update(dict_to, dict_from):
    """
    Like dict_to.update(dict_from), except don't overwrite any keys.
    """
    for key, val in dict(dict_from).iteritems():
        if key in dict_to:
            raise KeyError(key)
        dict_to[key] = val
    return dict_to

theano.config.warn.sum_div_dimshuffle_bug = False
floatX = theano.config.floatX
sharedX = lambda X, name: theano.shared(numpy.asarray(X, dtype=floatX), name=name)
if 0:
    print 'WARNING: using SLOW rng'
    RandomStreams = tensor.shared_randomstreams.RandomStreams
else:
    import theano.sandbox.rng_mrg
    RandomStreams = theano.sandbox.rng_mrg.MRG_RandomStreams

class DenoisingAutoencoder(Block):
    """
    A denoising autoencoder learns a representation of the input by
    reconstructing a noisy version of it.
    """
    def __init__(self, **kwargs):
        # TODO: Do we need anything else here?
        super(DenoisingAutoencoder, self).__init__(**kwargs)

    @classmethod
    def alloc(cls, corruptor, conf, rng=None):
        """Allocate a denoising autoencoder object."""
        if not hasattr(rng, 'randn'):
            rng = numpy.random.RandomState(rng)
        self = cls()
        self.corruptor = corruptor
        self.visbias = sharedX(
            numpy.zeros(conf['n_vis']),
            name='vb'
        )
        self.hidbias = sharedX(
            numpy.zeros(conf['n_hid']),
            name='hb'
        )
        self.weights = sharedX(
            .5 * rng.rand(conf['n_vis'], conf['n_hid']) * conf['irange'],
            name='W'
        )
        seed = int(rng.randint(2**30))
        self.s_rng = RandomStreams(seed)
        if conf['tied_weights']:
            self.w_prime = self.weights.T
        else:
            self.w_prime = sharedX(
                .5 * rng.rand(conf['n_hid'], conf['n_vis']) * conf['irange'],
                name='Wprime'
            )

        def _resolve_callable(conf_attr):
            if conf[conf_attr] is None:
                # The identity function, for linear layers.
                return lambda x: x
            # If it's a callable, use it directly.
            if hasattr(conf[conf_attr], '__call__'):
                return conf[conf_attr]
            elif hasattr(tensor.nnet, conf[conf_attr]):
                return getattr(tensor.nnet, conf[conf_attr])
            elif hasattr(tensor, conf[conf_attr]):
                return getattr(tensor, conf[conf_attr])
            else:
                raise ValueError("Couldn't interpret %s value: '%s'" %
                                 (conf_attr, conf[conf_attr]))

        self.act_enc = _resolve_callable('act_enc')
        self.act_dec = _resolve_callable('act_dec')
        self.conf = conf
        self._params = [
            self.visbias,
            self.hidbias,
            self.weights,
        ]
        if not conf['tied_weights']:
            self._params.append(self.w_prime)
        return self

    def _hidden_activation(self, x):
        """Single input pattern/minibatch activation function."""
        return self.act_enc(self.hidbias + tensor.dot(x, self.weights))

    def hidden_repr(self, inputs):
        """Hidden unit activations for each set of inputs."""
        return [self._hidden_activation(v) for v in inputs]

    def reconstruction(self, inputs):
        """Reconstructed inputs after corruption."""
        corrupted = self.corruptor(inputs)
        hiddens = self.hidden_repr(corrupted)
        return [
            self.act_dec(self.visbias + tensor.dot(h, self.w_prime))
            for h in hiddens
        ]

    def __call__(self, inputs):
        """
        Forward propagate (symbolic) input through this module, obtaining
        a representation to pass on to layers above.
        """
        return self.hidden_repr(inputs)

    def mse(self, inputs):
        """
        Symbolic expression for mean-squared error between the input and the
        denoised reconstruction.
        """
        pairs = izip(inputs, self.reconstruction(inputs))
        return [((inp - rec)**2).sum(axis=-1).mean() for inp, rec in pairs]

    def cross_entropy(self, inputs):
        """
        Symbolic expression for elementwise cross-entropy between input
        and reconstruction. Use for binary-valued features (but not for,
        e.g., one-hot codes).
        """
        pairs = izip(inputs, self.reconstruction(inputs))
        ce = lambda x, z: x * tensor.log(z) + (1 - x) * tensor.log(1 - z)
        return [ce(inp, rec).sum(axis=1).mean() for inp, rec in pairs]

class StackedDA(Block):
    """
    A class representing a stacked model. Forward propagation passes
    (symbolic) input through each layer sequentially.
    """
    def __init__(self, **kwargs):
        # TODO: Do we need anything else here?
        super(StackedDA, self).__init__(**kwargs)

    @classmethod
    def alloc(cls, corruptors, conf, rng=None):
        """Allocate a stacked denoising autoencoder object."""
        if not hasattr(rng, 'randn'):
            rng = numpy.random.RandomState(rng)
        self = cls()
        self._layers = []
        _local = {}
        # Make sure that if we have a sequence of encoder/decoder activations
        # or corruptors, that we have exactly as many as len(conf['n_hid'])
        if hasattr(corruptors, '__len__'):
            assert len(conf['n_hid']) == len(corruptors)
        else:
            corruptors = [corruptors] * len(conf['n_hid'])
        for c in ['act_enc', 'act_dec']:
            if type(conf[c]) is not str and hasattr(conf[c], '__len__'):
                assert len(conf['n_hid']) == len(conf[c])
                _local[c] = conf[c]
            else:
                _local[c] = [conf[c]] * len(conf['n_hid'])
        n_hids = conf['n_hid']
        # The number of visible units in each layer is the initial input
        # size and the first k-1 hidden unit sizes.
        n_viss = [conf['n_vis']] + conf['n_hid'][:-1]
        seq = izip(
            xrange(len(n_hids)),
            n_hids,
            n_viss,
            _local['act_enc'],
            _local['act_dec'],
            corruptors
        )
        # Create each layer.
        for k, n_hid, n_vis, act_enc, act_dec, corr in seq:
            # Create a local configuration dictionary for this layer.
            lconf = {
                'n_hid': n_hid,
                'n_vis': n_vis,
                'act_enc': act_enc,
                'act_dec': act_dec,
                'irange': conf['irange'],
                'tied_weights': conf['tied_weights'],
            }
            da = DenoisingAutoencoder.alloc(corr, lconf, rng)
            self._layers.append(da)
        return self

    def layers(self):
        """
        The layers of this model: the individual denoising autoencoder
        objects, which can be individually pre-trained.
        """
        return list(self._layers)

    def params(self):
        """
        The parameters that are learned in this model, i.e. the concatenation
        of all the layers' weights and biases.
        """
        return sum([l.params() for l in self._layers], [])

    def __call__(self, inputs):
        """
        Forward propagate (symbolic) input through this module, obtaining
        a representation to pass on to layers above.
        """
        transformed = inputs
        # Pass the input through each layer of the hierarchy.
        for layer in self._layers:
            transformed = layer(transformed)
        return transformed

class DATrainer(Trainer):
    @classmethod
    def alloc(cls, model, cost, minibatch, conf):
        """
        Takes a DenoisingAutoencoder object, a (symbolic) cost function
        which takes the input and turns it into a scalar (somehow),
        a (symbolic) minibatch, and a configuration dictionary.
        """
        # Take care of learning rate scales for individual parameters
        learning_rates = {}
        for parameter in model.params():
            lr_name = '%s_lr' % parameter.name
            try:
                thislr = conf[lr_name]
            except:
                thislr = 1.
            learning_rates[parameter] = sharedX(thislr, lr_name)

        # A shared variable for storing the iteration number.
        iteration = sharedX(0, 'iter')

        # A shared variable for storing the annealed base learning rate, used
        # to lower the learning rate gradually after a certain amount of time.
        annealed = sharedX(conf['base_lr'], 'annealed')

        # Instantiate the class, finally.
        self = cls(model=model, cost=cost, conf=conf,
                   learning_rates=learning_rates,
                   iteration=iteration, minibatch=minibatch)

    def updates(self):
        """Compute the updates for each of the parameter variables."""
        ups = {}
        # Base learning rate per example.
        base_lr = numpy.asarray(self.conf['base_lr'], dtype=floatX)

        # Annealing coefficient. Here we're using a formula of
        # base_lr * min(0.0, max(base_lr, lr_anneal_start / (iteration + 1))
        frac = self.conf['lr_anneal_start'] / (self.iteration + 1.)
        annealed = tensor.clip(
            tensor.cast(frac, floatX),
            0.0,    # minimum learning rate
            base_lr # maximum learning rate
        )

        # Update the shared variable for the annealed learning rate.
        ups[self.annealed] = annealed
        ups[self.iteration] = self.iteration + 1

        # Calculate the learning rates for each parameter, in the order
        # they appear in model.params()
        learn_rates = [annealed * self.lr_dict[p] for p in self.model.params()]
        # Get the gradient w.r.t. cost of each parameter.
        grads = [
            tensor.grad(self.cost(self.minibatch), p)
            for p in self.model.params()
        ]
        # Get the updates from sgd_updates, a PyLearn library function.
        updates = dict(sgd_updates(self.model.params(), grads, learn_rates))

        # Return the updates dictionary.
        return ups

<|endoftext|>


<commit_before>
class GPIO:  # GPIOをTurtleに
    import RPi.GPIO as IO
    import go_f
    import math

    gof = go_f.getdistance
    gof.goal = []
    gof.goal.sppend([lon1, lat1])
    gof.goal.sppend([lon2, lat2])
    gof.goal.sppend([lon3, lat3])
    gps = gpsnavi.gpsparser(goal=gof.goal)

    def __init__(self, leftmotor=[15, 16], rightmotor=[17, 22]):
        self.leftmotor = leftmotor
        self.rightmotor = rightmotor
        self.IO.setmode(self.IO.BCM)
        self.IO.setup(self.leftmotor[0], self.IO.OUT)
        self.IO.setup(self.leftmotor[1], self.IO.OUT)
        self.IO.setup(self.rightmotor[0], self.IO.OUT)
        self.IO.setup(self.rightmotor[1], self.IO.OUT)

    def forward(self):  # 前進
        self.IO.output(self.leftmotor[0], False)
        self.IO.output(self.leftmotor[1], True)
        self.IO.output(self.rightmotor[0], False)
        self.IO.output(self.rightmotor[1], True)
        self.time.sleep(0.01)

    def back(self):  # 後進
        self.IO.output(self.leftmotor[0], True)
        self.IO.output(self.leftmotor[1], False)
        self.IO.output(self.rightmotor[0], True)
        self.IO.output(self.rightmotor[1], False)
        self.time.sleep(0.01)

    def left(self):  # 左旋回
        self.IO.output(self.leftmotor[0], True)
        self.IO.output(self.leftmotor[1], False)
        self.IO.output(self.rightmotor[0], False)
        self.IO.output(self.rightmotor[1], True)
        self.time.sleep(0.01)

    def right(self):  # 右旋回
        self.IO.output(self.leftmotor[0], False)
        self.IO.output(self.leftmotor[1], True)
        self.IO.output(self.rightmotor[0], True)
        self.IO.output(self.rightmotor[1], False)

    def y_coordinates(self):
        return self.gps.goalverticaldistance

    def turn(self):
        self.y_coord = []

        while len(gps) > 0:  # If "len" is 0,roba is stop.
            gps.pop([0])

            while self.goaldistance > 30:

                while True:
                    self.y_coord.pop([0])
                    self.y_coord.append([self.y_coordinates])
                    self.forward(50)

                    if self.y_coord = self.y_coordinates:
                        self.back(50)
                        self.left(30)
                        self.time.sleep(0.01)
                    elif:
                        break

                self.time.sleep(0.01)

                if 0 <= self.gps.goalszimath < 90 or 270 < self.gps.goalszimath <= 360  and self.y_coordinates < self.y_coord :    #1st quadrant , 4th quadrant , south direction
                    self.right(self.gps.goalszimath)
                elif 0 <= self.gps.goalszimath < 90 or 270 < self.gps.goalszimath <= 360 and self.y_coordinates > self.y_coord :   #1nd quadrant , 4rd quadrant , north direction
                    self.left(self.gps.goalszimath)
                elif 90 <= self.gps.goalszimath <= 270 and self.y_coordinates < self.y_coord :     #2nd quadrant , 3rd quadrant , south direction
                    self.left(self.gps.goalszimath)
                elif 90 <= self.gps.goalszimath <= 270 and self.y_coordinates > self.y_coord :     #2nd quadrant , 3rd quadrant , north direction
                    self.right(self.gps.goalszimath)

                self.time.sleep(0.01)
                self.forward(self.gps.getdistance)
                self.time.sleep(0.01)<commit_msg>go_GPIO.py<commit_after>
class GPIO:  # GPIOをTurtleに
    import RPi.GPIO as IO
    import go_f
    import math

    def __init__(self, leftmotor=[15, 16], rightmotor=[17, 22]):
        self.leftmotor = leftmotor
        self.rightmotor = rightmotor
        self.IO.setmode(self.IO.BCM)
        self.IO.setup(self.leftmotor[0], self.IO.OUT)
        self.IO.setup(self.leftmotor[1], self.IO.OUT)
        self.IO.setup(self.rightmotor[0], self.IO.OUT)
        self.IO.setup(self.rightmotor[1], self.IO.OUT)

    def forward(self):  # 前進
        self.IO.output(self.leftmotor[0], False)
        self.IO.output(self.leftmotor[1], True)
        self.IO.output(self.rightmotor[0], False)
        self.IO.output(self.rightmotor[1], True)
        self.time.sleep(0.01)

    def back(self):  # 後進
        self.IO.output(self.leftmotor[0], True)
        self.IO.output(self.leftmotor[1], False)
        self.IO.output(self.rightmotor[0], True)
        self.IO.output(self.rightmotor[1], False)
        self.time.sleep(0.01)

    def left(self):  # 左旋回
        self.IO.output(self.leftmotor[0], True)
        self.IO.output(self.leftmotor[1], False)
        self.IO.output(self.rightmotor[0], False)
        self.IO.output(self.rightmotor[1], True)
        self.time.sleep(0.01)

    def right(self):  # 右旋回
        self.IO.output(self.leftmotor[0], False)
        self.IO.output(self.leftmotor[1], True)
        self.IO.output(self.rightmotor[0], True)
        self.IO.output(self.rightmotor[1], False)

    def y_coordinates(self):
        return self.gps.goalverticaldistance

    def turn(self):
        gof = go_f.getdistance
        goal = []
        goal.sppend([lon1, lat1])
        goal.sppend([lon2, lat2])
        goal.sppend([lon3, lat3])
        gps = gpsnavi.gpsparser(goal=gof.goal)

        self.y_coord = []

        while len(gps) > 0:  # If "len" is 0,roba is stop.
            gps.pop([0])

            while self.goaldistance > 30:

                while True:
                    self.y_coord.pop([0])
                    self.y_coord.append([self.y_coordinates])
                    self.forward(50)

                    if self.y_coord = self.y_coordinates:
                        self.back(50)
                        self.left(30)
                        self.time.sleep(0.01)
                    elif:
                        break

                self.time.sleep(0.01)

                if 0 <= self.gps.goalszimath < 90 or 270 < self.gps.goalszimath <= 360  and self.y_coordinates < self.y_coord :     #1st quadrant , 4th quadrant , south direction
                    self.right(self.gps.goalszimath)
                elif 0 <= self.gps.goalszimath < 90 or 270 < self.gps.goalszimath <= 360 and self.y_coordinates > self.y_coord :    #1nd quadrant , 4rd quadrant , north direction
                    self.left(self.gps.goalszimath)
                elif 90 <= self.gps.goalszimath <= 270 and self.y_coordinates < self.y_coord :      #2nd quadrant , 3rd quadrant , south direction
                    self.left(self.gps.goalszimath)
                elif 90 <= self.gps.goalszimath <= 270 and self.y_coordinates > self.y_coord :      #2nd quadrant , 3rd quadrant , north direction
                    self.right(self.gps.goalszimath)

                self.time.sleep(0.01)
                self.forward(self.gps.getdistance)
<|endoftext|>


<commit_before>import unittest

from MockMockMock import Mock, MockException

groupMakers = {
    "o" : lambda m: m.ordered,
    "u" : lambda m: m.unordered,
    "a" : lambda m: m.atomic,
    "p" : lambda m: m.optional,
    "l" : lambda m: m.alternative,
    "r" : lambda m: m.repeated,
}

def makeExpectations( mock, groups ):
    if len( groups ) > 0:
        group = groups[ 0 ]
        with groupMakers[ group ]( mock ):
            mock.expect.foobar( group + "a" )
            mock.expect.foobar( group + "b" )
            makeExpectations( mock, groups[ 1: ] )
            mock.expect.foobar( group + "c" )
            mock.expect.foobar( group + "d" )

def testAllowedOrder( allowedOrder ):
    def test( self ):
        for argument in allowedOrder:
            self.mock.object.foobar( argument )
        self.mock.tearDown()
    return test

def testForbidenOrder( forbidenOrder ):
    def test( self ):
        for argument in forbidenOrder[ : -1 ]:
            self.mock.object.foobar( argument )
        with self.assertRaises( MockException ):
                self.mock.object.foobar( forbidenOrder[ -1 ] )
    return test

def testTearDownError( forbidenOrder ):
    def test( self ):
        for argument in forbidenOrder:
            self.mock.object.foobar( argument )
        with self.assertRaises( MockException ):
            self.mock.tearDown()
    return test

def makeTestCase( groups, allowedOrders, forbidenOrders, tearDownErrors ):
    groups = list( groups )
    allowedOrders = [
        [
            allowedOrder[ 2 * i : 2 * i + 2 ]
            for i in range( len( allowedOrder ) / 2 )
        ]
        for allowedOrder in allowedOrders
    ]
    forbidenOrders = [
        [
            forbidenOrder[ 2 * i : 2 * i + 2 ]
            for i in range( len( forbidenOrder ) / 2 )
        ]
        for forbidenOrder in forbidenOrders
    ]
    tearDownErrors = [
        [
            tearDownError[ 2 * i : 2 * i + 2 ]
            for i in range( len( tearDownError ) / 2 )
        ]
        for tearDownError in tearDownErrors
    ]

    class TestCase( unittest.TestCase ):
        def setUp( self ):
            unittest.TestCase.setUp( self )
            self.mock = Mock( "MyMock" )
            makeExpectations( self.mock, groups )

    for allowedOrder in allowedOrders:
        setattr( TestCase, "test_" + "".join( allowedOrder ), testAllowedOrder( allowedOrder ) )

    for forbidenOrder in forbidenOrders:
        setattr( TestCase, "test_" + "".join( forbidenOrder ), testForbidenOrder( forbidenOrder ) )

    for tearDownError in tearDownErrors:
        setattr( TestCase, "test_" + "".join( tearDownError ), testTearDownError( tearDownError ) )

    TestCase.__name__ = "TestCase_" + "".join( groups )

    return TestCase

UnorderedGroup = makeTestCase(
    "u",
    [
        # Completed in any order
        "uaubucud",
        "ubuauduc",
        "ubucuaud",
    ],
    [
        # Bad argument
        "uaxx",
    ],
    [
        # Not completed
        "uaubuc",
        "uaub",
        "ua",
        "",
    ]
)

OrderedGroup = makeTestCase(
    "o",
    [
        # Completed in good order
        "oaobocod",
    ],
    [
        # Wrong order
        "ob",
    ],
    [
        # Not completed
        "oaoboc",
        "oaob",
        "oa",
        "",
    ]
)

OptionalGroup = makeTestCase(
    "p",
    [
        # Completed in good order
        "papbpcpd",
        # Not completed
        "papbpc",
        "papb",
        "pa",
        "",
    ],
    [
        # Wrong order
        "pb",
    ],
    [
    ]
)

OrderedInUnorderedGroup = makeTestCase(
    "uo",
    [
        # Original order
        "uauboaobocoducud",
        # Other possible orders
        #  ordered group at once
        "oaobocoduaubucud",
        "uaubucudoaobocod",
        #  ordered group in pieces
        "oauaobubocucodud",
        "oaudobucocubodua",
    ],
    [
        # Ordered group in wrong order
        "uaubob",
    ],
    [
    ]
)

# AtomicInUnorderedGroup = makeTestCase(
    # "ua",
    # [
        # # Original order
        # "uaubaaabacaducud",
        # # Other possible orders
        # #  atomic group at once
        # "aaabacaduaubucud",
        # "uaubucudaaabacad",
    # ],
    # [
        # # Atomic group in wrong order
        # "uaubabaaacaducud",
        # # Atomic group in pieces
        # "aauaabubacucadud",
        # "aaudabucacubadua",
    # ],
    # [
    # ]
# )

unittest.main()
<commit_msg>Improve test readability<commit_after>import unittest

from MockMockMock import Mock, MockException

groupMakers = {
    "o" : lambda m: m.ordered,
    "u" : lambda m: m.unordered,
    "a" : lambda m: m.atomic,
    "p" : lambda m: m.optional,
    "l" : lambda m: m.alternative,
    "r" : lambda m: m.repeated,
}

def makeExpectations( mock, groups ):
    if len( groups ) > 0:
        group = groups[ 0 ]
        with groupMakers[ group ]( mock ):
            mock.expect.foobar( group + "A" )
            mock.expect.foobar( group + "B" )
            makeExpectations( mock, groups[ 1: ] )
            mock.expect.foobar( group + "C" )
            mock.expect.foobar( group + "D" )

def testAllowedOrder( allowedOrder ):
    def test( self ):
        for argument in allowedOrder:
            self.call( argument )
        self.mock.tearDown()
    return test

def testForbidenOrder( forbidenOrder ):
    def test( self ):
        for argument in forbidenOrder[ : -1 ]:
            self.call( argument )
        with self.assertRaises( MockException ):
            self.call( forbidenOrder[ -1 ] )
    return test

def testTearDownError( forbidenOrder ):
    def test( self ):
        for argument in forbidenOrder:
            self.call( argument )
        with self.assertRaises( MockException ):
            self.mock.tearDown()
    return test

def makeTestCase( groups, allowedOrders, forbidenOrders, tearDownErrors ):
    groups = list( groups )
    allowedOrders = [
        [
            allowedOrder[ 2 * i : 2 * i + 2 ]
            for i in range( len( allowedOrder ) / 2 )
        ]
        for allowedOrder in allowedOrders
    ]
    forbidenOrders = [
        [
            forbidenOrder[ 2 * i : 2 * i + 2 ]
            for i in range( len( forbidenOrder ) / 2 )
        ]
        for forbidenOrder in forbidenOrders
    ]
    tearDownErrors = [
        [
            tearDownError[ 2 * i : 2 * i + 2 ]
            for i in range( len( tearDownError ) / 2 )
        ]
        for tearDownError in tearDownErrors
    ]

    class TestCase( unittest.TestCase ):
        def setUp( self ):
            unittest.TestCase.setUp( self )
            self.mock = Mock( "MyMock" )
            makeExpectations( self.mock, groups )

        def call( self, argument ):
            assert( len( argument ) == 2 )
            assert( argument[ 0 ] in groupMakers )
            assert( argument[ 1 ] in "ABCDX" )
            self.mock.object.foobar( argument )

    for allowedOrder in allowedOrders:
        setattr( TestCase, "test_" + "".join( allowedOrder ), testAllowedOrder( allowedOrder ) )

    for forbidenOrder in forbidenOrders:
        setattr( TestCase, "test_" + "".join( forbidenOrder ), testForbidenOrder( forbidenOrder ) )

    for tearDownError in tearDownErrors:
        setattr( TestCase, "test_" + "".join( tearDownError ), testTearDownError( tearDownError ) )

    TestCase.__name__ = "TestCase_" + "".join( groups )

    return TestCase

UnorderedGroup = makeTestCase(
    "u",
    [
        # Completed in any order
        "uAuBuCuD",
        "uBuAuDuC",
        "uBuCuAuD",
    ],
    [
        # BaD argument
        "uAuX",
    ],
    [
        # Not completed
        "uAuBuC",
        "uAuB",
        "uA",
        "",
    ]
)

OrderedGroup = makeTestCase(
    "o",
    [
        # Completed in good order
        "oAoBoCoD",
    ],
    [
        # Wrong order
        "oB",
    ],
    [
        # Not completed
        "oAoBoC",
        "oAoB",
        "oA",
        "",
    ]
)

OptionalGroup = makeTestCase(
    "p",
    [
        # Completed in good order
        "pApBpCpD",
        # Not completed
        "pApBpC",
        "pApB",
        "pA",
        "",
    ],
    [
        # Wrong order
        "pB",
    ],
    [
    ]
)

OrderedInUnorderedGroup = makeTestCase(
    "uo",
    [
        # Original order
        "uAuBoAoBoCoDuCuD",
        # Other possible orders
        #  ordered group at once
        "oAoBoCoDuAuBuCuD",
        "uAuBuCuDoAoBoCoD",
        #  ordered group in pieces
        "oAuAoBuBoCuCoDuD",
        "oAuDoBuCoCuBoDuA",
    ],
    [
        # Ordered group in wrong order
        "uAuBoB",
    ],
    [
    ]
)

# AtomicInUnorderedGroup = makeTestCase(
    # "ua",
    # [
        # # Original order
        # "uAuBaAaBaCaDuCuD",
        # # Other possible orders
        # #  atomic group at once
        # "aAaBaCaDuAuBuCuD",
        # "uAuBuCuDaAaBaCaD",
    # ],
    # [
        # # Atomic group in wrong order
        # "uAuBaBaAaCaDuCuD",
        # # Atomic group in pieces
        # "aAuaaBuBaCuCaDuD",
        # "aAuDaBuCaCuBaDuA",
    # ],
    # [
    # ]
# )

unittest.main()
<|endoftext|>