<commit_before>"""
Celery tasks.

.. moduleauthor:: Martijn Vermaat <martijn@vermaat.name>

.. Licensed under the MIT license, see the LICENSE file.
"""

from __future__ import division

from collections import defaultdict
from contextlib import contextmanager
import hashlib
import os
import time
import uuid

from celery import current_task, current_app, Task
from celery.utils.log import get_task_logger
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm.exc import NoResultFound
from vcf.parser import _Info as VcfInfo, field_counts as vcf_field_counts
import vcf

from . import db, celery
from .models import Annotation, Coverage, DataSource, DataUnavailable, Observation, Sample, Region, Variant, Variation
from .region_binning import all_bins
from .utils import digest, normalize_variant, normalize_chromosome, normalize_region, ReferenceMismatch

logger = get_task_logger(__name__)

class ReadError(Exception):
    """
    Exception thrown on failed data reading.
    """
    pass

class TaskError(Exception):
    """
    Exception thrown on failed task execution.
    """
    def __init__(self, code, message):
        self.code = code
        self.message = message
        super(TaskError, self).__init__(code, message)

class CleanTask(Task):
    """
    Ordinary Celery task, but with a way of registering cleanup routines that
    are executed after the task failed (on the worker node).
    """
    abstract = True

    # We maintain a list of cleanups per task id since the worker node only
    # instantiates the task class once (not for every run). Just storing one
    # list of cleanups in this instance would mean they are shared between
    # different task runs.
    _cleanups = defaultdict(list)

    def register_cleanup(self, task_id, cleanup):
        self._cleanups[task_id].append(cleanup)

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        for cleanup in reversed(self._cleanups[task_id]):
            cleanup()
        del self._cleanups[task_id]

def annotate_variants(original_variants, annotated_variants, original_filetype='vcf', annotated_filetype='vcf', ignore_sample_ids=None, original_records=1):
    """
    .. todo:: Merge back population study annotation (see implementation in
        the old-population-study branch).

    .. todo:: Calculate frequencies instead of counts, and take chromosome
        into account (config.CHROMOSOMES).
    """
    ignore_sample_ids = ignore_sample_ids or []

    if original_filetype != 'vcf':
        raise ReadError('Original data must be in VCF format')

    if annotated_filetype != 'vcf':
        raise ReadError('Annotated data must be in VCF format')

    reader = vcf.Reader(original_variants)

    reader.infos['OBS'] = VcfInfo('OBS', vcf_field_counts['A'], 'Integer',
        'Samples with variant (out of %i)' % Sample.query.count())
    reader.infos['COV'] = VcfInfo('COV', vcf_field_counts['A'], 'Integer',
        'Samples with coverage (out of %i)' % Sample.query.count())
    writer = vcf.Writer(annotated_variants, reader, lineterminator='\n')

    old_percentage = -1
    for i, record in enumerate(reader):
        # Task progress is updated in whole percentages, so for a maximum of
        # 100 times per task.
        percentage = min(int(i / original_records * 100), 99)
        if percentage > old_percentage:
            current_task.update_state(state='PROGRESS', meta={'percentage': percentage})
            old_percentage = percentage

        observations = []
        coverage = []
        for index, allele in enumerate(record.ALT):
            try:
                chromosome, position, reference, observed = normalize_variant(record.CHROM, record.POS, record.REF, str(allele))
            except ReferenceMismatch as e:
                raise ReadError(str(e))

            end_position = position + max(1, len(reference)) - 1
            bins = all_bins(position, end_position)

            try:
                variant = Variant.query.filter_by(chromosome=chromosome, position=position, reference=reference, observed=observed).one()
                observations.append(variant.observations.join(Variation).filter(~Variation.sample_id.in_(ignore_sample_ids)).count())
            except NoResultFound:
                observations.append(0)
            coverage.append(Region.query.join(Coverage).filter(Region.chromosome == chromosome,
                                                               Region.begin <= position,
                                                               Region.end >= end_position,
                                                               Region.bin.in_(bins),
                                                               ~Coverage.sample_id.in_(ignore_sample_ids)).count())

        record.add_info('OBS', observations)
        record.add_info('COV', coverage)
        writer.write_record(record)

def read_observations(observations, filetype='vcf'):
    # Todo: Merge back population study importing (see implementation in the
    #     old-population-study branch).
    if filetype != 'vcf':
        raise ReadError('Data must be in VCF format')

    reader = vcf.Reader(observations)

    for record in reader:
        if 'SV' in record.INFO:
            # For now we ignore these, reference is likely to be larger than
            # the maximum of 200 by the database schema.
            # Example use of this type are large deletions in 1000 Genomes.
            continue

        for index, allele in enumerate(record.ALT):
            try:
                chromosome, position, reference, observed = normalize_variant(
                    record.CHROM, record.POS, record.REF, str(allele))
            except ReferenceMismatch as e:
                logger.info('Reference mismatch: %s', str(e))
                if current_app.conf['REFERENCE_MISMATCH_ABORT']:
                    raise ReadError(str(e))
                continue

            # Variant support is defined by the number of samples in which a
            # variant allele was called, ignoring homo-/heterozygocity.
            # Todo: This check can break if index > 9.
            try:
                support = sum(1 for sample in record.samples if str(index + 1) in sample['GT'])
            except AttributeError:
                support = 1

            #if 'SF' in record.INFO and False:
            #    # Todo: Per alt allele?
            #    support = len(record.INFO['SF'])
            #elif 'AC' in record.INFO and False:
            #    # Todo: Per alt allele?
            #    support = record.INFO['AC'][0]
            #else:
            #    support = 1

            yield chromosome, position, reference, observed, support

def read_regions(regions, filetype='bed'):
    # Todo: Use pybedtools to parse BED file?
    if filetype != 'bed':
        raise ReadError('Data must be in BED format')

    for line in regions:
        fields = line.split()
        if len(fields) < 1 or fields[0] == 'track':
            continue
        try:
            chromosome, begin, end = normalize_region(
                fields[0], int(fields[1]), int(fields[2]))
        except (IndexError, ValueError):
            raise ReadError('Invalid line in BED file: "%s"' % line)
        except ReferenceMismatch as e:
            logger.info('Reference mismatch: %s', str(e))
            if current_app.conf['REFERENCE_MISMATCH_ABORT']:
                raise ReadError(str(e))
            continue
        yield chromosome, begin + 1, end

@celery.task(base=CleanTask)
def import_variation(variation_id):
    """
    Import variation as observations.
    """
    logger.info('Started task: import_variation(%d)', variation_id)

    variation = Variation.query.get(variation_id)
    if variation is None:
        raise TaskError('variation_not_found', 'Variation not found')

    if variation.imported:
        raise TaskError('variation_imported', 'Variation already imported')

    if variation.import_task_uuid:
        # Todo: Check somehow if the importing task is still running.
        # http://stackoverflow.com/questions/9824172/find-out-whether-celery-task-exists
        raise TaskError('variation_importing', 'Variation is being imported')

    # Todo: This has a possible race condition, but I'm not bothered to fix it
    #     at the moment. Reading and setting import_task_uuid should be an
    #     atomic action.
    variation.import_task_uuid = current_task.request.id
    db.session.commit()

    data_source = variation.data_source

    # Calculate data digest if it is not yet known.
    # Todo: Can we somehow factor this out into a separate (singleton) task,
    #     on which we wait?
    if not data_source.checksum:
        with data_source.data() as data:
            data_source.checksum, data_source.records = digest(data)
        db.session.commit()

    # Check if checksum is not in imported data sources.
    if DataSource.query.filter_by(checksum=data_source.checksum).join(Variation).filter_by(imported=True).count() > 0:
        raise TaskError('duplicate_data_source', 'Identical data source already imported')

    # Note: Since we are dealing with huge numbers of entries here, we commit
    # after each INSERT and manually rollback. Using builtin session rollback
    # would fill up all our memory.
    def delete_observations():
        variation.observations.delete()
        db.session.commit()
    current_task.register_cleanup(current_task.request.id, delete_observations)

    try:
        data = data_source.data()
    except DataUnavailable as e:
        raise TaskError(e.code, e.message)

    with data as observations:
        try:
            old_percentage = -1
            for i, (chromosome, position, reference, observed, support) in enumerate(read_observations(observations, filetype=data_source.filetype)):
                # Task progress is updated in whole percentages, so for a
                # maximum of 100 times per task.
                percentage = min(int(i / data_source.records * 100), 99)
                if percentage > old_percentage:
                    current_task.update_state(state='PROGRESS', meta={'percentage': percentage})
                    old_percentage = percentage
                    time.sleep(1)
                # SQLAlchemy doesn't seem to have anything like INSERT IGNORE
                # or INSERT ... ON DUPLICATE KEY UPDATE, so we have to work
                # our way around the situation.
                try:
                    # Todo: Check for errors (binning on begin/end may fail,
                    #     sequences might be too long).
                    variant = Variant(chromosome, position, reference, observed)
                    db.session.add(variant)
                    db.session.commit()
                except IntegrityError:
                    db.session.rollback()
                    try:
                        variant = Variant.query.filter_by(chromosome=chromosome, position=position, reference=reference, observed=observed).one()
                    except NoResultFound:
                        # Should never happen.
                        raise TaskError('database_inconsistency', 'Unrecoverable inconsistency of the database observed')
                observation = Observation(variant, variation, support=support)
                db.session.add(observation)
                db.session.commit()
        except ReadError as e:
            raise TaskError('invalid_observations', str(e))

    current_task.update_state(state='PROGRESS', meta={'percentage': 100})
    variation.imported = True
    db.session.commit()

    logger.info('Finished task: import_variation(%d)', variation_id)

@celery.task(base=CleanTask)
def import_coverage(coverage_id):
    """
    Import coverage as regions.
    """
    logger.info('Started task: import_coverage(%d)', coverage_id)

    coverage = Coverage.query.get(coverage_id)
    if coverage is None:
        raise TaskError('coverage_not_found', 'Coverage not found')

    if coverage.imported:
        raise TaskError('coverage_imported', 'Coverage already imported')

    if coverage.import_task_uuid is not None:
        # Todo: Check somehow if the importing task is still running.
        # http://stackoverflow.com/questions/9824172/find-out-whether-celery-task-exists
        raise TaskError('coverage_importing', 'Coverage is being imported')

    # Todo: This has a possible race condition, but I'm not bothered to fix it
    #     at the moment. Reading and setting import_task_uuid should be an
    #     atomic action.
    coverage.import_task_uuid = current_task.request.id
    db.session.commit()

    data_source = coverage.data_source

    # Calculate data digest if it is not yet known.
    if not data_source.checksum:
        with data_source.data() as data:
            data_source.checksum, data_source.records = digest(data)
        db.session.commit()

    # Check if checksum is not in imported data sources.
    if DataSource.query.filter_by(checksum=data_source.checksum).join(Coverage).filter_by(imported=True).count() > 0:
        raise TaskError('duplicate_data_source', 'Identical data source already imported')

    try:
        data = data_source.data()
    except DataUnavailable as e:
        raise TaskError(e.code, e.message)

    # Note: Since we are dealing with huge numbers of entries here, we commit
    # after each INSERT and manually rollback. Using builtin session rollback
    # would fill up all our memory.
    def delete_regions():
        coverage.regions.delete()
        db.session.commit()
    current_task.register_cleanup(current_task.request.id, delete_regions)

    with data as regions:
        try:
            old_percentage = -1
            for i, (chromosome, begin, end) in enumerate(read_regions(regions, filetype=data_source.filetype)):
                # Task progress is updated in whole percentages, so for a
                # maximum of 100 times per task.
                percentage = min(int(i / data_source.records * 100), 99)
                if percentage > old_percentage:
                    current_task.update_state(state='PROGRESS', meta={'percentage': percentage})
                    old_percentage = percentage
                db.session.add(Region(coverage, chromosome, begin, end))
                db.session.commit()
        except ReadError as e:
            raise TaskError('invalid_regions', str(e))

    current_task.update_state(state='PROGRESS', meta={'percentage': 100})
    coverage.imported = True
    db.session.commit()

    logger.info('Finished task: import_coverage(%d)', coverage_id)

@celery.task(base=CleanTask)
def write_annotation(annotation_id, ignore_sample_ids=None):
    """
    Annotate variants with frequencies from the database.
    """
    logger.info('Started task: write_annotation(%d)', annotation_id)

    ignore_sample_ids = ignore_sample_ids or []

    annotation = Annotation.query.get(annotation_id)
    if annotation is None:
        raise TaskError('annotation_not_found', 'Annotation not found')

    if annotation.written:
        raise TaskError('annotation_written', 'Annotation already written')

    if annotation.write_task_uuid is not None:
        # Todo: Check somehow if the writing task is still running, it might
        #     also be a failed task.
        # http://stackoverflow.com/questions/9824172/find-out-whether-celery-task-exists
        raise TaskError('annotation_writing', 'Annotation is being written')

    # Todo: This has a possible race condition, but I'm not bothered to fix it
    #     at the moment. Reading and setting write_task_uuid should be an
    #     atomic action.
    annotation.write_task_uuid = current_task.request.id
    db.session.commit()

    original_data_source = annotation.original_data_source
    annotated_data_source = annotation.annotated_data_source

    # Calculate data digest if it is not yet known.
    if not original_data_source.checksum:
        with original_data_source.data() as data:
            original_data_source.checksum, original_data_source.records = digest(data)
        db.session.commit()

    try:
        original_data = original_data_source.data()
        annotated_data = annotated_data_source.data_writer()
    except DataUnavailable as e:
        raise TaskError(e.code, e.message)

    with original_data as original_variants, annotated_data as annotated_variants:
        try:
            annotate_variants(original_variants, annotated_variants,
                              original_filetype=original_data_source.filetype,
                              annotated_filetype=annotated_data_source.filetype,
                              ignore_sample_ids=ignore_sample_ids,
                              original_records=original_data_source.records)
        except ReadError as e:
            raise TaskError('invalid_observations', str(e))

    current_task.update_state(state='PROGRESS', meta={'percentage': 100})
    annotation.written = True
    db.session.commit()

    logger.info('Finished task: write_annotation(%d)', annotation_id)

@celery.task
def ping():
    """
    Ping-pong task usefull for testing purposes.
    """
    logger.info('Started task: ping')
    logger.info('Finished task: ping')
    return 'pong'
<commit_msg>Comment<commit_after>"""
Celery tasks.

.. moduleauthor:: Martijn Vermaat <martijn@vermaat.name>

.. Licensed under the MIT license, see the LICENSE file.
"""

from __future__ import division

from collections import defaultdict
from contextlib import contextmanager
import hashlib
import os
import time
import uuid

from celery import current_task, current_app, Task
from celery.utils.log import get_task_logger
from sqlalchemy.exc import IntegrityError
from sqlalchemy.orm.exc import NoResultFound
from vcf.parser import _Info as VcfInfo, field_counts as vcf_field_counts
import vcf

from . import db, celery
from .models import Annotation, Coverage, DataSource, DataUnavailable, Observation, Sample, Region, Variant, Variation
from .region_binning import all_bins
from .utils import digest, normalize_variant, normalize_chromosome, normalize_region, ReferenceMismatch

logger = get_task_logger(__name__)

class ReadError(Exception):
    """
    Exception thrown on failed data reading.
    """
    pass

class TaskError(Exception):
    """
    Exception thrown on failed task execution.
    """
    def __init__(self, code, message):
        self.code = code
        self.message = message
        super(TaskError, self).__init__(code, message)

class CleanTask(Task):
    """
    Ordinary Celery task, but with a way of registering cleanup routines that
    are executed after the task failed (on the worker node).
    """
    abstract = True

    # We maintain a list of cleanups per task id since the worker node only
    # instantiates the task class once (not for every run). Just storing one
    # list of cleanups in this instance would mean they are shared between
    # different task runs.
    _cleanups = defaultdict(list)

    def register_cleanup(self, task_id, cleanup):
        self._cleanups[task_id].append(cleanup)

    def on_failure(self, exc, task_id, args, kwargs, einfo):
        for cleanup in reversed(self._cleanups[task_id]):
            cleanup()
        del self._cleanups[task_id]

def annotate_variants(original_variants, annotated_variants, original_filetype='vcf', annotated_filetype='vcf', ignore_sample_ids=None, original_records=1):
    """
    .. todo:: Merge back population study annotation (see implementation in
        the old-population-study branch).

    .. todo:: Calculate frequencies instead of counts, and take chromosome
        into account (config.CHROMOSOMES).
    """
    ignore_sample_ids = ignore_sample_ids or []

    if original_filetype != 'vcf':
        raise ReadError('Original data must be in VCF format')

    if annotated_filetype != 'vcf':
        raise ReadError('Annotated data must be in VCF format')

    reader = vcf.Reader(original_variants)

    reader.infos['OBS'] = VcfInfo('OBS', vcf_field_counts['A'], 'Integer',
        'Samples with variant (out of %i)' % Sample.query.count())
    reader.infos['COV'] = VcfInfo('COV', vcf_field_counts['A'], 'Integer',
        'Samples with coverage (out of %i)' % Sample.query.count())
    writer = vcf.Writer(annotated_variants, reader, lineterminator='\n')

    old_percentage = -1
    for i, record in enumerate(reader):
        # Task progress is updated in whole percentages, so for a maximum of
        # 100 times per task.
        percentage = min(int(i / original_records * 100), 99)
        if percentage > old_percentage:
            current_task.update_state(state='PROGRESS', meta={'percentage': percentage})
            old_percentage = percentage

        observations = []
        coverage = []
        for index, allele in enumerate(record.ALT):
            try:
                chromosome, position, reference, observed = normalize_variant(record.CHROM, record.POS, record.REF, str(allele))
            except ReferenceMismatch as e:
                raise ReadError(str(e))

            end_position = position + max(1, len(reference)) - 1
            bins = all_bins(position, end_position)

            try:
                variant = Variant.query.filter_by(chromosome=chromosome, position=position, reference=reference, observed=observed).one()
                observations.append(variant.observations.join(Variation).filter(~Variation.sample_id.in_(ignore_sample_ids)).count())
            except NoResultFound:
                observations.append(0)
            coverage.append(Region.query.join(Coverage).filter(Region.chromosome == chromosome,
                                                               Region.begin <= position,
                                                               Region.end >= end_position,
                                                               Region.bin.in_(bins),
                                                               ~Coverage.sample_id.in_(ignore_sample_ids)).count())

        record.add_info('OBS', observations)
        record.add_info('COV', coverage)
        writer.write_record(record)

def read_observations(observations, filetype='vcf'):
    # Todo: Merge back population study importing (see implementation in the
    #     old-population-study branch).
    if filetype != 'vcf':
        raise ReadError('Data must be in VCF format')

    reader = vcf.Reader(observations)

    for record in reader:
        if 'SV' in record.INFO:
            # For now we ignore these, reference is likely to be larger than
            # the maximum of 200 by the database schema.
            # Example use of this type are large deletions in 1000 Genomes.
            continue

        for index, allele in enumerate(record.ALT):
            try:
                chromosome, position, reference, observed = normalize_variant(
                    record.CHROM, record.POS, record.REF, str(allele))
            except ReferenceMismatch as e:
                logger.info('Reference mismatch: %s', str(e))
                if current_app.conf['REFERENCE_MISMATCH_ABORT']:
                    raise ReadError(str(e))
                continue

            # Variant support is defined by the number of samples in which a
            # variant allele was called, ignoring homo-/heterozygocity.
            # Todo: This check can break if index > 9.
            try:
                support = sum(1 for sample in record.samples if str(index + 1) in sample['GT'])
            except AttributeError:
                support = 1

            #if 'SF' in record.INFO and False:
            #    # Todo: Per alt allele?
            #    support = len(record.INFO['SF'])
            #elif 'AC' in record.INFO and False:
            #    # Todo: Per alt allele?
            #    support = record.INFO['AC'][0]
            #else:
            #    support = 1

            yield chromosome, position, reference, observed, support

def read_regions(regions, filetype='bed'):
    # Todo: Use pybedtools to parse BED file?
    if filetype != 'bed':
        raise ReadError('Data must be in BED format')

    for line in regions:
        fields = line.split()
        if len(fields) < 1 or fields[0] == 'track':
            continue
        try:
            chromosome, begin, end = normalize_region(
                fields[0], int(fields[1]), int(fields[2]))
        except (IndexError, ValueError):
            raise ReadError('Invalid line in BED file: "%s"' % line)
        except ReferenceMismatch as e:
            logger.info('Reference mismatch: %s', str(e))
            if current_app.conf['REFERENCE_MISMATCH_ABORT']:
                raise ReadError(str(e))
            continue
        yield chromosome, begin + 1, end

@celery.task(base=CleanTask)
def import_variation(variation_id):
    """
    Import variation as observations.
    """
    logger.info('Started task: import_variation(%d)', variation_id)

    variation = Variation.query.get(variation_id)
    if variation is None:
        raise TaskError('variation_not_found', 'Variation not found')

    if variation.imported:
        raise TaskError('variation_imported', 'Variation already imported')

    if variation.import_task_uuid:
        # Todo: Check somehow if the importing task is still running.
        # http://stackoverflow.com/questions/9824172/find-out-whether-celery-task-exists
        raise TaskError('variation_importing', 'Variation is being imported')

    # Todo: This has a possible race condition, but I'm not bothered to fix it
    #     at the moment. Reading and setting import_task_uuid should be an
    #     atomic action.
    variation.import_task_uuid = current_task.request.id
    db.session.commit()

    data_source = variation.data_source

    # Calculate data digest if it is not yet known.
    # Todo: Can we somehow factor this out into a separate (singleton) task,
    #     on which we wait?
    if not data_source.checksum:
        with data_source.data() as data:
            data_source.checksum, data_source.records = digest(data)
        db.session.commit()

    # Check if checksum is not in imported data sources.
    if DataSource.query.filter_by(checksum=data_source.checksum).join(Variation).filter_by(imported=True).count() > 0:
        raise TaskError('duplicate_data_source', 'Identical data source already imported')

    # Note: Since we are dealing with huge numbers of entries here, we commit
    # after each INSERT and manually rollback. Using builtin session rollback
    # would fill up all our memory.
    def delete_observations():
        variation.observations.delete()
        db.session.commit()
    current_task.register_cleanup(current_task.request.id, delete_observations)

    try:
        data = data_source.data()
    except DataUnavailable as e:
        raise TaskError(e.code, e.message)

    with data as observations:
        try:
            old_percentage = -1
            for i, (chromosome, position, reference, observed, support) in enumerate(read_observations(observations, filetype=data_source.filetype)):
                # Task progress is updated in whole percentages, so for a
                # maximum of 100 times per task.
                # Todo: Increment ``i`` per record instead of per yielded
                #     observation.
                percentage = min(int(i / data_source.records * 100), 99)
                if percentage > old_percentage:
                    current_task.update_state(state='PROGRESS', meta={'percentage': percentage})
                    old_percentage = percentage
                    time.sleep(1)
                # SQLAlchemy doesn't seem to have anything like INSERT IGNORE
                # or INSERT ... ON DUPLICATE KEY UPDATE, so we have to work
                # our way around the situation.
                try:
                    # Todo: Check for errors (binning on begin/end may fail,
                    #     sequences might be too long).
                    variant = Variant(chromosome, position, reference, observed)
                    db.session.add(variant)
                    db.session.commit()
                except IntegrityError:
                    db.session.rollback()
                    try:
                        variant = Variant.query.filter_by(chromosome=chromosome, position=position, reference=reference, observed=observed).one()
                    except NoResultFound:
                        # Should never happen.
                        raise TaskError('database_inconsistency', 'Unrecoverable inconsistency of the database observed')
                observation = Observation(variant, variation, support=support)
                db.session.add(observation)
                db.session.commit()
        except ReadError as e:
            raise TaskError('invalid_observations', str(e))

    current_task.update_state(state='PROGRESS', meta={'percentage': 100})
    variation.imported = True
    db.session.commit()

    logger.info('Finished task: import_variation(%d)', variation_id)

@celery.task(base=CleanTask)
def import_coverage(coverage_id):
    """
    Import coverage as regions.
    """
    logger.info('Started task: import_coverage(%d)', coverage_id)

    coverage = Coverage.query.get(coverage_id)
    if coverage is None:
        raise TaskError('coverage_not_found', 'Coverage not found')

    if coverage.imported:
        raise TaskError('coverage_imported', 'Coverage already imported')

    if coverage.import_task_uuid is not None:
        # Todo: Check somehow if the importing task is still running.
        # http://stackoverflow.com/questions/9824172/find-out-whether-celery-task-exists
        raise TaskError('coverage_importing', 'Coverage is being imported')

    # Todo: This has a possible race condition, but I'm not bothered to fix it
    #     at the moment. Reading and setting import_task_uuid should be an
    #     atomic action.
    coverage.import_task_uuid = current_task.request.id
    db.session.commit()

    data_source = coverage.data_source

    # Calculate data digest if it is not yet known.
    if not data_source.checksum:
        with data_source.data() as data:
            data_source.checksum, data_source.records = digest(data)
        db.session.commit()

    # Check if checksum is not in imported data sources.
    if DataSource.query.filter_by(checksum=data_source.checksum).join(Coverage).filter_by(imported=True).count() > 0:
        raise TaskError('duplicate_data_source', 'Identical data source already imported')

    try:
        data = data_source.data()
    except DataUnavailable as e:
        raise TaskError(e.code, e.message)

    # Note: Since we are dealing with huge numbers of entries here, we commit
    # after each INSERT and manually rollback. Using builtin session rollback
    # would fill up all our memory.
    def delete_regions():
        coverage.regions.delete()
        db.session.commit()
    current_task.register_cleanup(current_task.request.id, delete_regions)

    with data as regions:
        try:
            old_percentage = -1
            for i, (chromosome, begin, end) in enumerate(read_regions(regions, filetype=data_source.filetype)):
                # Task progress is updated in whole percentages, so for a
                # maximum of 100 times per task.
                percentage = min(int(i / data_source.records * 100), 99)
                if percentage > old_percentage:
                    current_task.update_state(state='PROGRESS', meta={'percentage': percentage})
                    old_percentage = percentage
                db.session.add(Region(coverage, chromosome, begin, end))
                db.session.commit()
        except ReadError as e:
            raise TaskError('invalid_regions', str(e))

    current_task.update_state(state='PROGRESS', meta={'percentage': 100})
    coverage.imported = True
    db.session.commit()

    logger.info('Finished task: import_coverage(%d)', coverage_id)

@celery.task(base=CleanTask)
def write_annotation(annotation_id, ignore_sample_ids=None):
    """
    Annotate variants with frequencies from the database.
    """
    logger.info('Started task: write_annotation(%d)', annotation_id)

    ignore_sample_ids = ignore_sample_ids or []

    annotation = Annotation.query.get(annotation_id)
    if annotation is None:
        raise TaskError('annotation_not_found', 'Annotation not found')

    if annotation.written:
        raise TaskError('annotation_written', 'Annotation already written')

    if annotation.write_task_uuid is not None:
        # Todo: Check somehow if the writing task is still running, it might
        #     also be a failed task.
        # http://stackoverflow.com/questions/9824172/find-out-whether-celery-task-exists
        raise TaskError('annotation_writing', 'Annotation is being written')

    # Todo: This has a possible race condition, but I'm not bothered to fix it
    #     at the moment. Reading and setting write_task_uuid should be an
    #     atomic action.
    annotation.write_task_uuid = current_task.request.id
    db.session.commit()

    original_data_source = annotation.original_data_source
    annotated_data_source = annotation.annotated_data_source

    # Calculate data digest if it is not yet known.
    if not original_data_source.checksum:
        with original_data_source.data() as data:
            original_data_source.checksum, original_data_source.records = digest(data)
        db.session.commit()

    try:
        original_data = original_data_source.data()
        annotated_data = annotated_data_source.data_writer()
    except DataUnavailable as e:
        raise TaskError(e.code, e.message)

    with original_data as original_variants, annotated_data as annotated_variants:
        try:
            annotate_variants(original_variants, annotated_variants,
                              original_filetype=original_data_source.filetype,
                              annotated_filetype=annotated_data_source.filetype,
                              ignore_sample_ids=ignore_sample_ids,
                              original_records=original_data_source.records)
        except ReadError as e:
            raise TaskError('invalid_observations', str(e))

    current_task.update_state(state='PROGRESS', meta={'percentage': 100})
    annotation.written = True
    db.session.commit()

    logger.info('Finished task: write_annotation(%d)', annotation_id)

@celery.task
def ping():
    """
    Ping-pong task usefull for testing purposes.
    """
    logger.info('Started task: ping')
    logger.info('Finished task: ping')
    return 'pong'
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#   Debian Changes Bot
#   Copyright (C) 2008 Chris Lamb <chris@chris-lamb.co.uk>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as
#   published by the Free Software Foundation, either version 3 of the
#   License, or (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

import unittest

import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from DebianDevelChangesBot.mailparsers import BugSubmittedParser as p

class TestMailParserBugSubmitted(unittest.TestCase):

    def setUp(self):
        self.headers = {
            'Subject': 'Bug#123456: Bug title',
            'From': 'Submitter Name <name@host.tld>',
            'List-Id': '<debian-bugs-dist.lists.debian.org>',
        }

        self.body = [
            "Package: package-name",
            "Version: version-here",
            "",
            "Description"
        ]

    def testSimple(self):
        msg = p.parse(self.headers, self.body)

        self.assert_(msg)
        self.assertEqual(msg.package, 'package-name')
        self.assertEqual(msg.version, 'version-here')
        self.assertEqual(msg.by, 'Submitter Name <name@host.tld>')

    def testVersionWithSpaces(self):
        self.body[1] = "Version: version with spaces"
        msg = p.parse(self.headers, self.body)

        self.failIf(msg.version)

if __name__ == "__main__":
    unittest.main()
<commit_msg>Shorten some more email addresses used in tests<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#   Debian Changes Bot
#   Copyright (C) 2008 Chris Lamb <chris@chris-lamb.co.uk>
#
#   This program is free software: you can redistribute it and/or modify
#   it under the terms of the GNU Affero General Public License as
#   published by the Free Software Foundation, either version 3 of the
#   License, or (at your option) any later version.
#
#   This program is distributed in the hope that it will be useful,
#   but WITHOUT ANY WARRANTY; without even the implied warranty of
#   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#   GNU Affero General Public License for more details.
#
#   You should have received a copy of the GNU Affero General Public License
#   along with this program.  If not, see <http://www.gnu.org/licenses/>.

import unittest

import os, sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from DebianDevelChangesBot.mailparsers import BugSubmittedParser as p

class TestMailParserBugSubmitted(unittest.TestCase):

    def setUp(self):
        self.headers = {
            'Subject': 'Bug#123456: Bug title',
            'From': 'Submitter Name <name@t.com>',
            'List-Id': '<debian-bugs-dist.lists.debian.org>',
        }

        self.body = [
            "Package: package-name",
            "Version: version-here",
            "",
            "Description"
        ]

    def testSimple(self):
        msg = p.parse(self.headers, self.body)

        self.assert_(msg)
        self.assertEqual(msg.package, 'package-name')
        self.assertEqual(msg.version, 'version-here')
        self.assertEqual(msg.by, 'Submitter Name <name@t.com>')

    def testVersionWithSpaces(self):
        self.body[1] = "Version: version with spaces"
        msg = p.parse(self.headers, self.body)

        self.failIf(msg.version)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before>import sys
import os
import numpy as np
import numpy.ma as ma
import argparse
import csv
import time
from condor_kmeans.vector import VectorStream
from condor_kmeans.utils import make_directory

FIND_CLUSTER_MAP_HEADER = '''universe = vanilla
Executable=/lusr/bin/python
Requirements = InMastodon
+Group   = "GRAD"
+Project = "AI_ROBOTICS"
+ProjectDescription = "Distributed K-means"
getenv = True

'''

FIND_CLUSTER_MAP_JOB = '''Arguments = {python_filepath} worker {worker_id} {step} {data_filename} {start} {end} {weights_filename} {centroids_filename} {assignments_outfile} {mindistance_outfile} {partial_centroids_outfile} {partial_centroids_counts_outfile}
Output = {output_filename}
Error = {error_filename}
Queue 1

'''

class CondorKmeansPool(object):
    def __init__(self, username, num_workers, working_dir, polling_delay_in_seconds=30, timeout_in_seconds=60*60*24):
        self._username = username
        self._num_workers = num_workers
        self._base_dir = make_directory(working_dir, 'condor')
        self._data_dir = make_directory(self._base_dir, 'data')
        self._polling_delay_in_seconds = polling_delay_in_seconds

    def map_find_nearest_cluster(self, step, data, weights, centroids, assignments, min_distances, assign=True):
        # Create a map for use in pretty string formatting for the jobs file
        dargs = {'username': self._username, 'step': step, 'base_dir': self._base_dir, 'python_filepath': os.path.abspath(__file__) }

        # Get the location of the data
        if data is VectorStream:
            dargs['data_filename'] = data._filename
        else:
            # If we were given data in memory instead of in a stream, save it to file so it can be streamed
            dargs['data_filename'] = self._data_dir + 'data.csv'
            np.savetxt(dargs['data_filename'], data, delimiter=',')

        # Setup the file structure
        dargs['output_dir'] = make_directory(self._base_dir, 'output')
        dargs['error_dir'] = make_directory(self._base_dir, 'error')
        dargs['assignments_dir'] = make_directory(self._data_dir, 'assignments')
        dargs['mindistance_dir'] = make_directory(self._data_dir, 'mindistance')
        dargs['partial_centroids_dir'] = make_directory(self._data_dir, 'partial_centroids')
        dargs['partial_centroids_counts_dir'] = make_directory(self._data_dir, 'partial_centroids_counts')

        # Write the weights to file
        dargs['weights_filename'] = self._data_dir + 'weights.csv'
        np.savetxt(dargs['weights_filename'], weights, delimiter=',')

        # Write the centroids to file
        dargs['centroids_filename'] = self._data_dir + 'step{step}_centroids.csv'.format(**dargs)
        np.savetxt(dargs['centroids_filename'], centroids, delimiter=',')

        # Figure out how many data points each worker should be working on
        num_workers = min(len(data), self._num_workers)
        worker_load = len(data) / num_workers

        # Each worker is given a range in the form (inclusive_start, exclusive_end)
        worker_ranges = [(i*worker_load, (i+1)*worker_load) for i in xrange(num_workers)]
        worker_ranges[-1] = (worker_ranges[-1][0], len(data))

        # Open up a jobs file
        dargs['jobs_filename'] = '{base_dir}find_nearest_cluster_jobs'.format(**dargs)
        with open(dargs['jobs_filename'], 'wb') as f:
            # Write the header to the top of the file
            f.write(FIND_CLUSTER_MAP_HEADER.format(**dargs))

            # Write each worker's job to file
            for i, (start, end) in enumerate(worker_ranges):
                dargs['worker_id'] = i
                dargs['start'] = start
                dargs['end'] = end
                dargs['assignments_outfile'] = '{assignments_dir}{worker_id}.csv'.format(**dargs)
                dargs['mindistance_outfile'] = '{mindistance_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_outfile'] = '{partial_centroids_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_counts_outfile'] = '{partial_centroids_counts_dir}{worker_id}.csv'.format(**dargs)
                dargs['output_filename'] = '{output_dir}find_nearest_cluster_{worker_id}.out'.format(**dargs)
                dargs['error_filename'] = '{error_dir}find_nearest_cluster_{worker_id}.err'.format(**dargs)
                f.write(FIND_CLUSTER_MAP_JOB.format(**dargs))

        # Submit the job to condor. Fuck using subprocess at the moment
        os.system('condor_submit {jobs_filename}'.format(**dargs))

        # Flag to keep waiting
        poll = True

        # Poll and chill
        # TODO: implement timeout counter
        while poll:
            # Chill
            time.sleep(self._polling_delay_in_seconds)

            poll = False

            # Check to see if each file exists
            for i, (start, end) in enumerate(worker_ranges):
                dargs['worker_id'] = i
                dargs['start'] = start
                dargs['end'] = end
                dargs['assignments_outfile'] = '{assignments_dir}{worker_id}.csv'.format(**dargs)
                dargs['mindistance_outfile'] = '{mindistance_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_outfile'] = '{partial_centroids_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_counts_outfile'] = '{partial_centroids_counts_dir}{worker_id}.csv'.format(**dargs)
                dargs['output_filename'] = '{output_dir}find_nearest_cluster_{worker_id}.out'.format(**dargs)
                dargs['error_filename'] = '{error_dir}find_nearest_cluster_{worker_id}.err'.format(**dargs)

                # If the worker hasn't even started yet, chill some more
                if not os.path.exists(dargs['output_filename']):
                    poll = True
                    break

                # Check the output of the file
                with open(dargs['output_filename'], 'rb') as f:
                    lines = f.readlines()

                    # If the file hasn't been written to yet, chill some more
                    if len(lines) == 0:
                        poll = True
                        break

                    # Get the last non-empty line of the output file
                    text = lines[-1].strip('\n')
                    while text == '' and len(lines) > 0:
                        lines = lines[:-1]
                        text = lines[-1].strip('\n')

                    # If the worker failed, raise an exception in the master and point to the debug files
                    if text == 'Quit due to error':
                        os.system('condor_rm {username}'.format(**dargs))
                        raise Exception('Worker failure occurred. Worker #{worker_id} range: [{start}-{end}] output file: {output_filename} error file: {error_filename}'.format(**dargs))
                    elif text == 'Success!':
                        # If the worker succeeded, move on to the next file
                        continue
                    else:
                        # If the worker hasn't finished writing its output file, chill some more
                        poll = True
                        break

        if assign:
            updated_centroids = np.ma.zeros(centroids.shape, mask=np.zeros(centroids.shape, dtype=int))
            updated_centroids_counts = np.zeros(centroids.shape[0])

        # All the workers are finished. Merge the results back
        for i, (start, end) in enumerate(worker_ranges):
            dargs['worker_id'] = i
            dargs['start'] = start
            dargs['end'] = end
            dargs['assignments_outfile'] = '{assignments_dir}{worker_id}.csv'.format(**dargs)
            dargs['mindistance_outfile'] = '{mindistance_dir}{worker_id}.csv'.format(**dargs)
            dargs['partial_centroids_outfile'] = '{partial_centroids_dir}{worker_id}.csv'.format(**dargs)
            dargs['partial_centroids_counts_outfile'] = '{partial_centroids_counts_dir}{worker_id}.csv'.format(**dargs)
            dargs['output_filename'] = '{output_dir}find_nearest_cluster_{worker_id}.out'.format(**dargs)
            dargs['error_filename'] = '{error_dir}find_nearest_cluster_{worker_id}.err'.format(**dargs)

            # Check the output of the file
            min_distances[start:end] = np.loadtxt(dargs['mindistance_outfile'], delimiter=',')
            if assign:
                assignments[start:end] = np.loadtxt(dargs['assignments_outfile'], delimiter=',')
                partial_centroids = np.loadtxt(dargs['partial_centroids_outfile'], delimiter=',')
                partial_centroids_counts = np.loadtxt(dargs['partial_centroids_counts_outfile'], delimiter=',')
                # Calculate a running mean for the cluster centers
                next_counts = updated_centroids_counts + partial_centroids_counts
                updated_centroids = (updated_centroids * (updated_centroids_counts / next_counts.clip(1))[:,np.newaxis]
                                     + partial_centroids * (partial_centroids_counts / next_counts.clip(1))[:,np.newaxis])
                updated_centroids_counts = next_counts

            # Clean up
            os.remove(dargs['assignments_outfile'])
            os.remove(dargs['mindistance_outfile'])
            os.remove(dargs['partial_centroids_outfile'])
            os.remove(dargs['partial_centroids_counts_outfile'])
            os.remove(dargs['output_filename'])
            os.remove(dargs['error_filename'])

        os.remove(dargs['jobs_filename'])

        # If we had to write the data to file, delete the temp file
        if data is not VectorStream:
            os.remove(dargs['data_filename'])

        if assign:
            return updated_centroids

def condor_find_nearest_cluster(condor_username, data, weights, centroids, assignments, num_workers, step, assign=True, polling_delay=30):
    min_distances = np.zeros(assignments.shape)

    # Create the thread pool to process things in parallel on condor
    pool = CondorKmeansPool(condor_username, num_workers, os.getcwd(), polling_delay_in_seconds=polling_delay)

    # Map the jobs to condor workers
    pool.map_find_nearest_cluster(step, data, weights, centroids, assignments, min_distances, assign=assign)

    return min_distances

def worker_main():
    worker_id = int(sys.argv[2])
    step = int(sys.argv[3])
    data_filename = sys.argv[4]
    start = int(sys.argv[5])
    end = int(sys.argv[6])
    weights_filename = sys.argv[7]
    centroids_filename = sys.argv[8]
    assignments_outfile = sys.argv[9]
    mindistance_outfile = sys.argv[10]
    partial_centroids_outfile = sys.argv[11]
    partial_centroids_counts_outfile = sys.argv[12]

    try:
        data = VectorStream(data_filename)
        weights = np.loadtxt(weights_filename, delimiter=',')
        centroids = np.loadtxt(centroids_filename, delimiter=',')
        assignments = np.zeros(end-start, dtype=int) - 1

        if len(centroids.shape) == 1:
            centroids = centroids[:, np.newaxis].T
    except Exception as ex:
        print 'Error loading files: {0}'.format(ex)
        print 'Quit due to error'
        exit(1)

    try:
        # Reuse the local kmeans nearest-cluster routine
        from diet2vec.clustering.kmeans import find_nearest_cluster
        all_params = (step, data, start, end, weights, centroids, assignments, worker_id)
        (worker_id, assignments, min_distances) = find_nearest_cluster(all_params)
    except Exception as ex:
        print 'Error finding nearest cluster: {0}'.format(ex)
        print 'Quit due to error'
        exit(1)

    if assign:
        try:
            # Calculate the cluster centroids and counts
            partial_centroids = np.zeros(centroids.shape)
            centroid_counts = np.zeros(end - start)
            assign_idx = 0
            for i, x in enumerate(data):
                if i < start:
                    continue
                if i >= end:
                    break
                centroid_idx = assignments[assign_idx]
                partial_centroids[centroid_idx] += x
                centroid_counts[centroid_idx] += 1
                assign_idx += 1
            partial_centroids = partial_centroids / centroid_counts[:,np.newaxis]
        except Exception as ex:
            print 'Error calculating partial centroids: {0}'.format(ex)
            print 'Quit due to error'
            exit(1)

    try:
        np.savetxt(assignments_outfile, assignments, delimiter=',')
        np.savetxt(mindistance_outfile, min_distances, delimiter=',')
        np.savetxt(partial_centroids_outfile, partial_centroids, delimiter=',')
        np.savetxt(partial_centroids_counts_outfile, centroid_counts, delimiter=',')
    except Exception as ex:
        print 'Error saving files: {0}'.format(ex)
        print 'Quit due to error'
        exit(1)

    print 'Success!'

if __name__ == '__main__':
    if sys.argv[1] == 'worker':
        worker_main()
    elif sys.argv[1] == 'aggregate':
        aggregate_main()

<commit_msg>Refactoring to remove need for constant process on submit node<commit_after>import sys
import os
import numpy as np
import numpy.ma as ma
import argparse
import csv
import time
from condor_kmeans.vector import VectorStream
from condor_kmeans.utils import make_directory

FIND_CLUSTER_MAP_HEADER = '''universe = vanilla
Executable=/lusr/bin/python
Requirements = InMastodon
+Group   = "GRAD"
+Project = "AI_ROBOTICS"
+ProjectDescription = "Distributed K-means"
getenv = True

'''

FIND_CLUSTER_MAP_JOB = '''Arguments = {python_filepath} worker {worker_id} {step} {data_filename} {start} {end} {weights_filename} {centroids_filename} {assignments_outfile} {mindistance_outfile} {partial_centroids_outfile} {partial_centroids_counts_outfile}
Output = {output_filename}
Error = {error_filename}
Queue 1

'''

class CondorKmeansPool(object):
    def __init__(self, username, num_workers, working_dir, polling_delay_in_seconds=30, timeout_in_seconds=60*60*24):
        self._username = username
        self._num_workers = num_workers
        self._base_dir = make_directory(working_dir, 'condor')
        self._data_dir = make_directory(self._base_dir, 'data')
        self._polling_delay_in_seconds = polling_delay_in_seconds

    def map_find_nearest_cluster(self, step, data, weights, centroids, assignments, min_distances, assign=True):
        # Create a map for use in pretty string formatting for the jobs file
        dargs = {'username': self._username, 'step': step, 'base_dir': self._base_dir, 'python_filepath': os.path.abspath(__file__) }

        # Get the location of the data
        if data is VectorStream:
            dargs['data_filename'] = data._filename
        else:
            # If we were given data in memory instead of in a stream, save it to file so it can be streamed
            dargs['data_filename'] = self._data_dir + 'data.csv'
            np.savetxt(dargs['data_filename'], data, delimiter=',')

        # Setup the file structure
        dargs['output_dir'] = make_directory(self._base_dir, 'output')
        dargs['error_dir'] = make_directory(self._base_dir, 'error')
        dargs['assignments_dir'] = make_directory(self._data_dir, 'assignments')
        dargs['mindistance_dir'] = make_directory(self._data_dir, 'mindistance')
        dargs['partial_centroids_dir'] = make_directory(self._data_dir, 'partial_centroids')
        dargs['partial_centroids_counts_dir'] = make_directory(self._data_dir, 'partial_centroids_counts')

        # Write the weights to file
        dargs['weights_filename'] = self._data_dir + 'weights.csv'
        np.savetxt(dargs['weights_filename'], weights, delimiter=',')

        # Write the centroids to file
        dargs['centroids_filename'] = self._data_dir + 'step{step}_centroids.csv'.format(**dargs)
        np.savetxt(dargs['centroids_filename'], centroids, delimiter=',')

        # Figure out how many data points each worker should be working on
        num_workers = min(len(data), self._num_workers)
        worker_load = len(data) / num_workers

        # Each worker is given a range in the form (inclusive_start, exclusive_end)
        worker_ranges = [(i*worker_load, (i+1)*worker_load) for i in xrange(num_workers)]
        worker_ranges[-1] = (worker_ranges[-1][0], len(data))

        # Open up a jobs file
        dargs['jobs_filename'] = '{base_dir}find_nearest_cluster_jobs'.format(**dargs)
        with open(dargs['jobs_filename'], 'wb') as f:
            # Write the header to the top of the file
            f.write(FIND_CLUSTER_MAP_HEADER.format(**dargs))

            # Write each worker's job to file
            for i, (start, end) in enumerate(worker_ranges):
                dargs['worker_id'] = i
                dargs['start'] = start
                dargs['end'] = end
                dargs['assignments_outfile'] = '{assignments_dir}{worker_id}.csv'.format(**dargs)
                dargs['mindistance_outfile'] = '{mindistance_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_outfile'] = '{partial_centroids_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_counts_outfile'] = '{partial_centroids_counts_dir}{worker_id}.csv'.format(**dargs)
                dargs['output_filename'] = '{output_dir}find_nearest_cluster_{worker_id}.out'.format(**dargs)
                dargs['error_filename'] = '{error_dir}find_nearest_cluster_{worker_id}.err'.format(**dargs)
                f.write(FIND_CLUSTER_MAP_JOB.format(**dargs))

        # Submit the job to condor. Fuck using subprocess at the moment
        os.system('condor_submit {jobs_filename}'.format(**dargs))

        print 'Jobs submitted. Polling for completion...'

        # Flag to keep waiting
        poll = True

        # Poll and chill
        # TODO: implement timeout counter
        while poll:
            # Chill
            time.sleep(self._polling_delay_in_seconds)

            poll = False

            # Check to see if each file exists
            for i, (start, end) in enumerate(worker_ranges):
                dargs['worker_id'] = i
                dargs['start'] = start
                dargs['end'] = end
                dargs['assignments_outfile'] = '{assignments_dir}{worker_id}.csv'.format(**dargs)
                dargs['mindistance_outfile'] = '{mindistance_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_outfile'] = '{partial_centroids_dir}{worker_id}.csv'.format(**dargs)
                dargs['partial_centroids_counts_outfile'] = '{partial_centroids_counts_dir}{worker_id}.csv'.format(**dargs)
                dargs['output_filename'] = '{output_dir}find_nearest_cluster_{worker_id}.out'.format(**dargs)
                dargs['error_filename'] = '{error_dir}find_nearest_cluster_{worker_id}.err'.format(**dargs)

                # If the worker hasn't even started yet, chill some more
                if not os.path.exists(dargs['output_filename']):
                    poll = True
                    break

                # Check the output of the file
                with open(dargs['output_filename'], 'rb') as f:
                    lines = f.readlines()

                    # If the file hasn't been written to yet, chill some more
                    if len(lines) == 0:
                        poll = True
                        break

                    # Get the last non-empty line of the output file
                    text = lines[-1].strip('\n')
                    while text == '' and len(lines) > 0:
                        lines = lines[:-1]
                        text = lines[-1].strip('\n')

                    # If the worker failed, raise an exception in the master and point to the debug files
                    if text == 'Quit due to error':
                        os.system('condor_rm {username}'.format(**dargs))
                        raise Exception('Worker failure occurred. Worker #{worker_id} range: [{start}-{end}] output file: {output_filename} error file: {error_filename}'.format(**dargs))
                    elif text == 'Success!':
                        # If the worker succeeded, move on to the next file
                        continue
                    else:
                        # If the worker hasn't finished writing its output file, chill some more
                        poll = True
                        break

        print 'Workers finished. Aggregating results.'

        if assign:
            updated_centroids = np.ma.zeros(centroids.shape, mask=np.zeros(centroids.shape, dtype=int))
            updated_centroids_counts = np.zeros(centroids.shape[0])

        # All the workers are finished. Merge the results back
        for i, (start, end) in enumerate(worker_ranges):
            dargs['worker_id'] = i
            dargs['start'] = start
            dargs['end'] = end
            dargs['assignments_outfile'] = '{assignments_dir}{worker_id}.csv'.format(**dargs)
            dargs['mindistance_outfile'] = '{mindistance_dir}{worker_id}.csv'.format(**dargs)
            dargs['partial_centroids_outfile'] = '{partial_centroids_dir}{worker_id}.csv'.format(**dargs)
            dargs['partial_centroids_counts_outfile'] = '{partial_centroids_counts_dir}{worker_id}.csv'.format(**dargs)
            dargs['output_filename'] = '{output_dir}find_nearest_cluster_{worker_id}.out'.format(**dargs)
            dargs['error_filename'] = '{error_dir}find_nearest_cluster_{worker_id}.err'.format(**dargs)

            # Check the output of the file
            min_distances[start:end] = np.loadtxt(dargs['mindistance_outfile'], delimiter=',')
            if assign:
                assignments[start:end] = np.loadtxt(dargs['assignments_outfile'], delimiter=',')
                partial_centroids = np.loadtxt(dargs['partial_centroids_outfile'], delimiter=',')
                partial_centroids_counts = np.loadtxt(dargs['partial_centroids_counts_outfile'], delimiter=',')
                # Calculate a running mean for the cluster centers
                next_counts = updated_centroids_counts + partial_centroids_counts
                updated_centroids = (updated_centroids * (updated_centroids_counts / next_counts.clip(1))[:,np.newaxis]
                                     + partial_centroids * (partial_centroids_counts / next_counts.clip(1))[:,np.newaxis])
                updated_centroids_counts = next_counts

            # Clean up
            os.remove(dargs['assignments_outfile'])
            os.remove(dargs['mindistance_outfile'])
            os.remove(dargs['partial_centroids_outfile'])
            os.remove(dargs['partial_centroids_counts_outfile'])
            os.remove(dargs['output_filename'])
            os.remove(dargs['error_filename'])

        os.remove(dargs['jobs_filename'])

        # If we had to write the data to file, delete the temp file
        if data is not VectorStream:
            os.remove(dargs['data_filename'])

        if assign:
            return updated_centroids

def condor_find_nearest_cluster(condor_username, data, weights, centroids, assignments, num_workers, step, assign=True, polling_delay=30):
    min_distances = np.zeros(assignments.shape)

    # Create the thread pool to process things in parallel on condor
    pool = CondorKmeansPool(condor_username, num_workers, os.getcwd(), polling_delay_in_seconds=polling_delay)

    # Map the jobs to condor workers
    pool.map_find_nearest_cluster(step, data, weights, centroids, assignments, min_distances, assign=assign)

    return min_distances

def worker_main():
    worker_id = int(sys.argv[2])
    step = int(sys.argv[3])
    data_filename = sys.argv[4]
    start = int(sys.argv[5])
    end = int(sys.argv[6])
    weights_filename = sys.argv[7]
    centroids_filename = sys.argv[8]
    assignments_outfile = sys.argv[9]
    mindistance_outfile = sys.argv[10]
    partial_centroids_outfile = sys.argv[11]
    partial_centroids_counts_outfile = sys.argv[12]

    try:
        data = VectorStream(data_filename)
        weights = np.loadtxt(weights_filename, delimiter=',')
        centroids = np.loadtxt(centroids_filename, delimiter=',')
        assignments = np.zeros(end-start, dtype=int) - 1

        if len(centroids.shape) == 1:
            centroids = centroids[:, np.newaxis].T
    except Exception as ex:
        print 'Error loading files: {0}'.format(ex)
        print 'Quit due to error'
        exit(1)

    try:
        # Reuse the local kmeans nearest-cluster routine
        from diet2vec.clustering.kmeans import find_nearest_cluster
        all_params = (step, data, start, end, weights, centroids, assignments, worker_id)
        (worker_id, assignments, min_distances) = find_nearest_cluster(all_params)
    except Exception as ex:
        print 'Error finding nearest cluster: {0}'.format(ex)
        print 'Quit due to error'
        exit(1)

    if assign:
        try:
            # Calculate the cluster centroids and counts
            print 'Calculating partial centroids'
            partial_centroids = np.zeros(centroids.shape)
            centroid_counts = np.zeros(end - start)
            assign_idx = 0
            for i, x in enumerate(data):
                if i < start:
                    continue
                if i >= end:
                    break
                centroid_idx = assignments[assign_idx]
                partial_centroids[centroid_idx] += x
                centroid_counts[centroid_idx] += 1
                assign_idx += 1
            partial_centroids = partial_centroids / centroid_counts[:,np.newaxis]
        except Exception as ex:
            print 'Error calculating partial centroids: {0}'.format(ex)
            print 'Quit due to error'
            exit(1)

    try:
        np.savetxt(assignments_outfile, assignments, delimiter=',')
        np.savetxt(mindistance_outfile, min_distances, delimiter=',')
        np.savetxt(partial_centroids_outfile, partial_centroids, delimiter=',')
        np.savetxt(partial_centroids_counts_outfile, centroid_counts, delimiter=',')
    except Exception as ex:
        print 'Error saving files: {0}'.format(ex)
        print 'Quit due to error'
        exit(1)

    print 'Success!'

if __name__ == '__main__':
    if sys.argv[1] == 'worker':
        worker_main()
    elif sys.argv[1] == 'aggregate':
        aggregate_main()

<|endoftext|>


<commit_before>#! /usr/bin/python
"""
Varmint Pest - Attempts to frighten away varmints such as raccoons
Created 7/04/17 by Greg Griffes
"""
import time, sys, os
import logging

from automat import MethodicalMachine

##import RPi.GPIO as GPIO
##import pigpio # http://abyz.co.uk/rpi/pigpio/python.html
##import spidev
##import Adafruit_ADS1x15
##from collections import deque
##import threading

##import datetime
##from picamera import PiCamera

# for graphics
##from webcolors import name_to_rgb
##import pygame
##from pygame.locals import *

# Global variables
##global flow_count
##flow_count = 0

#camera = PiCamera()

# Provides a logging solution for troubleshooting
logging.basicConfig(level=logging.DEBUG, format='(%(threadName)-10s) %(message)s')

# -----------------------------------
# Constants
# -----------------------------------
NIGHT = 0
DAY = 1

# -----------------------------------
# Initialize
# -----------------------------------
# =======================================
# Local classes
# =======================================
class _daylight_sensor(object):
    # Create the daylight sensor class
    def __init__(self):
        self.night = NIGHT
        self.day = DAY

    def __del__(self):
        logging.debug('daylight sensor class closed')

    def read(self):
        logging.debug("Daylight sensor indicates night time")
        return self.night

##########################
class varmintpest_fsm(object):
##########################
# Create the state machine class from Automat
    _machine = MethodicalMachine()

    #=======================================
    # Inputs
    #=======================================
    @_machine.input()
    def day_detected(self):
        "Light sensor indicates day time"

    @_machine.input()
    def night_detected(self):
        "Light sensor indicates night time"

    #=======================================
    # Outputs
    #=======================================
    @_machine.output()
    def _play_bird_sounds(self):
        logging.debug('playing bird sounds')
        time.sleep(2.0)

    @_machine.output()
    def _varmint_hunting(self):
        logging.debug('hunting varmints')
        time.sleep(2.0)

    #=======================================
    # States
    #=======================================
    @_machine.state(initial=True)
    def daytime(self):
        "In this state, do daytime things"

    @_machine.state()
    def nighttime(self):
        "in this state, do night time things"

    #=======================================
    # Transition logic
    #=======================================

    daytime.upon(night_detected, enter=nighttime, outputs=[_varmint_hunting])
    nighttime.upon(day_detected, enter=daytime, outputs=[_play_bird_sounds])

###############################################################
# Main program
###############################################################
if __name__ == '__main__':

# -----------------------------------
# Initialize variables
# -----------------------------------

# -----------------------------------
# Exception handler
# -----------------------------------
    try:

        logging.debug('debug logging is on')

# -----------------------------------
# Create objects
# -----------------------------------
        varmintpest = varmintpest_fsm()
        daylight_sensor = _daylight_sensor()

# -----------------------------------
# Create and start daemons
# -----------------------------------

# -----------------------------------
# Create and start graphics
# -----------------------------------

###############################################################
# Main loop
###############################################################

        loop = 0    # used to determine the first time through each valve_status change

        daylight_status = DAY

        while True:         #Demo loop

# intial state is daytime
            if daylight_sensor.read() == NIGHT and daylight_status == DAY:
                daylight_status = NIGHT
                varmintpest.night_detected()    # move to night time state
            elif daylight_sensor.read() == DAY and daylight_status == NIGHT:
                daylight_status = DAY
                varmintpest.day_detected()      # move to day time state
            else:
                time.sleep(3.0)

###########################################################
# END
###########################################################
    except KeyboardInterrupt:
        logging.debug("Keyboard Interrupt exception!")
        exit()

    except BaseException as e:
        logging.error('General exception!: ' + str(e))

    # normal exit
<commit_msg>Full state machine up and running<commit_after>#! /usr/bin/python
"""
Varmint Pest - Attempts to frighten away varmints such as raccoons
Created 7/04/17 by Greg Griffes
"""
import time, sys, os
import logging

from automat import MethodicalMachine

##import RPi.GPIO as GPIO
##import pigpio # http://abyz.co.uk/rpi/pigpio/python.html
##import spidev
##import Adafruit_ADS1x15
##from collections import deque
##import threading

##import datetime
##from picamera import PiCamera

# for graphics
##from webcolors import name_to_rgb
##import pygame
##from pygame.locals import *

# Global variables
##global flow_count
##flow_count = 0

#camera = PiCamera()

# Provides a logging solution for troubleshooting
logging.basicConfig(level=logging.DEBUG, format='(%(threadName)-10s) %(message)s')

# -----------------------------------
# Constants
# -----------------------------------
NIGHT = 0
DAY = 1

# -----------------------------------
# Initialize
# -----------------------------------
# =======================================
# Local classes
# =======================================
class _daylight_sensor(object):
    # Create the daylight sensor class
    def __init__(self):
        self.night = NIGHT
        self.day = DAY

    def __del__(self):
        logging.debug('daylight sensor class closed')

    def read(self):
        logging.debug("Daylight sensor indicates night time")
        return self.night

##########################
class varmintpest_fsm(object):
##########################
    """varmintpest_fsm is a state machine using the Automat MethodicalMachine
    There are three states:
    S1: Daytime (the initial state)
    S2: Nighttime-motion (waiting for motion)
    S3: Nighttime-IR (upon IR signature, set off the varmint blast)

    State: S1 (Daytime) parameters:
        S1I1: Input 1: daylight-yes
        S1I2: Input 2: daylight-no

        S1T1: Transition 1: upon daylight-yes
            S1O1: Output 1: Play bird songs
            S1T1ns: Next state: enter daytime
        S1T2: Transition 2: upon daylight-no
            S1O2: Output 2: message entering nighttime and test for motion
            S1T2ns: Next state: enter nighttime-motion

    State: S2 (nighttime-motion) parameters:
        S2I1: Input 1: daylight-yes
        S2I2: Input 2: daylight-no
        S2I3: Input 3: motion-detected
        S2I4: Input 4: motion-not-detected

        S2T1: Transition 1: upon daylight-yes
            S2O1: message start of daytime
            S2T1ns: enter S1 daytime
        S2T2: Transition 2: upon daylight-no
            S2T2ns: enter S2 nighttime-motion
        S2T3: Transition 3: upon motion-detected
            S2O2: message motion detected
            S2T3ns: enter S3 nighttime-IR
        S2T4: Transition 4: upon motion-not-detected
            S2T4ns: enter S2 nighttime-motion

     State: S3 (Nighttime-IR) parameters:
        S3I1: Input 1: IR-detected
        S3I2: Input 2: IR-not-detected

        S3T1: Transition 1: upon IR-detected
            S3O1: Output 1: Perform Varmint-blast
            S3T1ns: Next state: enter nighttime-motion
        S3T2: Transition 2: upon IR-not-detected
            S3O2: Output 2: log false-motion-detection
            S3T2ns: Next state: enter S2 nighttime-motion

    """

# Create the state machine class from Automat
    _machine = MethodicalMachine()

#=======================================
# Inputs
#=======================================
    @_machine.input()
    def daylightYes(self):
        "Light sensor indicates day time"

    @_machine.input()
    def daylightNo(self):
        "Light sensor indicates night time"

    @_machine.input()
    def motionDetected(self):
        "Motion - possible varmint"

    @_machine.input()
    def motionNotDetected(self):
        "Nothing found, keep looking"

    @_machine.input()
    def IRDetected(self):
        "IR signature indicates varmint"

    @_machine.input()
    def IRNotDetected(self):
        "Nothing found, keep looking"

#=======================================
# Outputs
#=======================================
    @_machine.output()
    def _play_bird_sounds(self):
        logging.debug('playing bird sounds')
        time.sleep(1.0)

    @_machine.output()
    def _nighttime(self):
        logging.debug('entering nighttime-motion state')
        time.sleep(1.0)

    @_machine.output()
    def _messageStartOfDaytime(self):
        logging.debug('entering daytime state')
        time.sleep(1.0)

    @_machine.output()
    def _messageMotionDetected(self):
        logging.debug('Motion detected! Entering motion-IR state')
        time.sleep(1.0)

    @_machine.output()
    def _performVarmintBlast(self):
        logging.debug('IR detected! Performing Varmint Blast!')
        time.sleep(1.0)

    @_machine.output()
    def _messageFalseMotionDetected(self):
        logging.debug('Motion detected but no IR signature')
        time.sleep(1.0)

#=======================================
# S1 State
#=======================================
    @_machine.state(initial=True)
    def daytime(self):
        "wait for nighttime"

#=======================================
# S2 State - nighttime-motion
#=======================================
    @_machine.state()
    def nighttimeMotion(self):
        "wait for motion or daytime"

#=======================================
# S3 State - nighttime-IR
#=======================================
    @_machine.state()
    def nighttimeIR(self):
        "wait for IR signature"

#=======================================
# S1 Transition logic
#=======================================
    daytime.upon(daylightYes, enter=daytime, outputs=[_play_bird_sounds])
    daytime.upon(daylightNo, enter=nighttimeMotion, outputs=[_nighttime])

#=======================================
# S2 Transition logic
#=======================================
    nighttimeMotion.upon(daylightYes, enter=daytime, outputs=[_messageStartOfDaytime])
    nighttimeMotion.upon(daylightNo, enter=nighttimeMotion, outputs=[])
    nighttimeMotion.upon(motionDetected, enter=nighttimeIR, outputs=[_messageMotionDetected])
    nighttimeMotion.upon(motionNotDetected, enter=nighttimeMotion, outputs=[])

#=======================================
# S3 Transition logic
#=======================================
    nighttimeIR.upon(IRDetected, enter=nighttimeMotion, outputs=[_performVarmintBlast])
    nighttimeIR.upon(IRNotDetected, enter=nighttimeMotion, outputs=[_messageFalseMotionDetected])

###############################################################
# Main program
###############################################################
if __name__ == '__main__':

# -----------------------------------
# Initialize variables
# -----------------------------------

# -----------------------------------
# Exception handler
# -----------------------------------
    try:

        logging.debug('debug logging is on')

# -----------------------------------
# Create objects
# -----------------------------------
        varmintpest = varmintpest_fsm()
#        daylight_sensor = _daylight_sensor()

# -----------------------------------
# Create and start daemons
# -----------------------------------

# -----------------------------------
# Create and start graphics
# -----------------------------------

###############################################################
# Main program
###############################################################

        varmintpest.daylightYes()       # move to day time state
        varmintpest.daylightNo()        # move to night time state

        varmintpest.motionNotDetected() # move to day time state
        varmintpest.motionNotDetected() # move to day time state
        varmintpest.motionNotDetected() # move to day time state
        varmintpest.motionNotDetected() # move to day time state

        varmintpest.motionDetected() # move to day time state

        varmintpest.IRNotDetected() # move to day time state

        varmintpest.motionDetected() # move to day time state

        varmintpest.IRDetected() # move to day time state

        varmintpest.motionNotDetected() # move to day time state
        varmintpest.motionNotDetected() # move to day time state
        varmintpest.motionNotDetected() # move to day time state
        varmintpest.motionNotDetected() # move to day time state

        varmintpest.daylightYes() # move to day time state

        time.sleep(3.0)

###########################################################
# END
###########################################################
    except KeyboardInterrupt:
        logging.debug("Keyboard Interrupt exception!")
        exit()

    except BaseException as e:
        logging.error('General exception!: ' + str(e))

    # normal exit
<|endoftext|>


<commit_before># Time:  O(n)
# Space: O(1)

# Rotate an array of n elements to the right by k steps.
#
# For example, with n = 7 and k = 3, the array [1,2,3,4,5,6,7] is rotated to [5,6,7,1,2,3,4].
#
# Note:
# Try to come up as many solutions as you can, there are at least 3 different ways to solve this problem.

class Solution(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """

    def rotate(self, nums, k):
        k %= len(nums)
        self.reverse(nums, 0, len(nums))
        self.reverse(nums, 0, k)
        self.reverse(nums, k, len(nums))

    def reverse(self, nums, start, end):
        while start < end:
            nums[start], nums[end - 1] = nums[end - 1], nums[start]
            start += 1
            end -= 1

# Time:  O(n)
# Space: O(1)
from fractions import gcd

class Solution2(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """

    def rotate(self, nums, k):
        def apply_cycle_permutation(k, offset, cycle_len, nums):
            tmp = nums[offset]
            for i in xrange(1, cycle_len):
                nums[(offset + i * k) % len(nums)], tmp = tmp, nums[(offset + i * k) % len(nums)]
            nums[offset] = tmp

        k %= len(nums)
        num_cycles = gcd(len(nums), k)
        cycle_len = len(nums) / num_cycles
        for i in xrange(num_cycles):
            apply_cycle_permutation(k, i, cycle_len, nums)

# Time:  O(n)
# Space: O(1)
class Solution3(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """

    def rotate(self, nums, k):
        count = 0
        start = 0
        while count < len(nums):
            curr = start
            prev = nums[curr]
            while True:
                idx = (curr + k) % len(nums)
                nums[idx], prev = prev, nums[idx]
                curr = idx
                count += 1
                if start == curr:
                    break
            start += 1

# Time:  O(n)
# Space: O(n)
class Solution4(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """
    def rotate(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: void Do not return anything, modify nums in-place instead.
        """
        nums[:] = nums[len(nums) - k:] + nums[:len(nums) - k]

# Time:  O(n^2)
# Space: O(1)
class Solution5(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """
    def rotate(self, nums, k):
        while k > 0:
            nums.insert(0, nums.pop())
            k -= 1

if __name__ == '__main__':
    nums = [1, 2, 3, 4, 5, 6, 7]
    Solution().rotate(nums, 3)
    print nums
<commit_msg>Update rotate-array.py<commit_after># Time:  O(n)
# Space: O(1)

# Rotate an array of n elements to the right by k steps.
#
# For example, with n = 7 and k = 3, the array [1,2,3,4,5,6,7] is rotated to [5,6,7,1,2,3,4].
#
# Note:
# Try to come up as many solutions as you can, there are at least 3 different ways to solve this problem.

class Solution(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """

    def rotate(self, nums, k):
        k %= len(nums)
        self.reverse(nums, 0, len(nums))
        self.reverse(nums, 0, k)
        self.reverse(nums, k, len(nums))

    def reverse(self, nums, start, end):
        while start < end:
            nums[start], nums[end - 1] = nums[end - 1], nums[start]
            start += 1
            end -= 1

# Time:  O(n)
# Space: O(1)
from fractions import gcd

class Solution2(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """

    def rotate(self, nums, k):
        def apply_cycle_permutation(k, offset, cycle_len, nums):
            tmp = nums[offset]
            for i in xrange(1, cycle_len):
                nums[(offset + i * k) % len(nums)], tmp = tmp, nums[(offset + i * k) % len(nums)]
            nums[offset] = tmp

        k %= len(nums)
        num_cycles = gcd(len(nums), k)
        cycle_len = len(nums) / num_cycles
        for i in xrange(num_cycles):
            apply_cycle_permutation(k, i, cycle_len, nums)

# Time:  O(n)
# Space: O(1)
class Solution3(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """

    def rotate(self, nums, k):
        count = 0
        start = 0
        while count < len(nums):
            curr = start
            prev = nums[curr]
            while True:
                idx = (curr + k) % len(nums)
                nums[idx], prev = prev, nums[idx]
                curr = idx
                count += 1
                if start == curr:
                    break
            start += 1

# Time:  O(n)
# Space: O(n)
class Solution4(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """
    def rotate(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: void Do not return anything, modify nums in-place instead.
        """
        nums[:] = nums[len(nums) - k:] + nums[:len(nums) - k]

# Time:  O(k * n)
# Space: O(1)
class Solution5(object):
    """
    :type nums: List[int]
    :type k: int
    :rtype: void Do not return anything, modify nums in-place instead.
    """
    def rotate(self, nums, k):
        while k > 0:
            nums.insert(0, nums.pop())
            k -= 1

if __name__ == '__main__':
    nums = [1, 2, 3, 4, 5, 6, 7]
    Solution().rotate(nums, 3)
    print nums
<|endoftext|>