<commit_before># -*- coding: utf-8 -*-
import logging

from path import path

from puzzle.models import (Case, Compound, Variant, Gene, Genotype, Transcript)
from puzzle.utils import (get_most_severe_consequence, get_hgnc_symbols,
                          get_omim_number, get_ensembl_id)

from puzzle.plugins import Plugin

from vcftoolbox import (get_variant_dict, HeaderParser, get_info_dict,
                        get_vep_dict)

logger = logging.getLogger(__name__)

class VcfPlugin(Plugin):
    """docstring for Plugin"""

    def init_app(self, app):
        """Initialize plugin via Flask."""
        self.root_path = app.config['PUZZLE_ROOT']
        self.pattern = app.config['PUZZLE_PATTERN']

    def _find_vcfs(self, pattern='*.vcf'):
        """Walk subdirectories and return VCF files."""
        return path(self.root_path).walkfiles(pattern)

    def cases(self, pattern=None):
        """Return all VCF file paths."""
        pattern = pattern or self.pattern
        vcfs = self._find_vcfs(pattern)
        case_objs = (Case(case_id=vcf.replace('/', '|'),
                          name=vcf.basename()) for vcf in vcfs)
        return case_objs

    def _add_compounds(self, variant, info_dict):
        """Check if there are any compounds and add them to the variant

        """
        compound_entry = info_dict.get('Compounds')
        if compound_entry:
            for family_annotation in compound_entry.split(','):
                compounds = family_annotation.split(':')[-1].split('|')
                for compound in compounds:
                    splitted_compound = compound.split('>')

                    compound_score = None
                    if len(splitted_compound) > 1:
                        compound_id = splitted_compound[0]
                        compound_score = splitted_compound[-1]

                    variant.add_compound(Compound(
                        variant_id=compound_id,
                        combined_score=compound_score
                    ))

    def _add_genes(self, variant):
        """Add the genes for a variant"""

        hgnc_symbols = get_hgnc_symbols(
            transcripts = variant['transcripts']
        )

        for hgnc_id in hgnc_symbols:
            variant.add_gene(Gene(
                symbol=hgnc_id,
                omim_number=get_omim_number(hgnc_id),
                ensembl_id=get_ensembl_id(hgnc_id)
            ))

    def _add_transcripts(self, variant, vep_dict):
        """Add the transcripts for a variant"""
        for allele in vep_dict:
            for transcript_info in vep_dict[allele]:
                variant.add_transcript(Transcript(
                    SYMBOL = transcript_info.get('SYMBOL'),
                    Feature = transcript_info.get('Feature'),
                    BIOTYPE = transcript_info.get('BIOTYPE'),
                    Consequence = transcript_info.get('Consequence'),
                    STRAND = transcript_info.get('STRAND'),
                    SIFT = transcript_info.get('SIFT'),
                    PolyPhen = transcript_info.get('PolyPhen'),
                    EXON = transcript_info.get('EXON'),
                    HGVSc = transcript_info.get('HGVSc'),
                    HGVSp = transcript_info.get('HGVSp')
                ))

    def _variants(self, vcf_file_path):
        head = HeaderParser()
        # Parse the header
        with open(vcf_file_path, 'r') as variant_file:
            for line in variant_file:
                line = line.rstrip()
                if line.startswith('#'):
                    if line.startswith('##'):
                        head.parse_meta_data(line)
                    else:
                        head.parse_header_line(line)
                else:
                    break

        header_line = head.header
        individuals = head.individuals

        variant_columns = ['CHROM', 'POS', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER']

        vep_header = head.vep_columns

        with open(vcf_file_path, 'r') as vcf_file:
            index = 0
            for variant_line in vcf_file:
                if not variant_line.startswith('#'):
                    index += 1
                    #Create a variant dict:
                    variant_dict =  get_variant_dict(
                        variant_line = variant_line,
                        header_line = header_line
                    )
                    #Crreate a info dict:
                    info_dict = get_info_dict(
                        info_line = variant_dict['INFO']
                    )
                    vep_string = info_dict.get('CSQ')

                    if vep_string:
                        vep_dict = get_vep_dict(
                            vep_string = vep_string,
                            vep_header = vep_header
                        )

                    variant = Variant(
                        **{column: variant_dict.get(column, '.')
                            for column in variant_columns}
                        )
                    variant['frequencies'] = []
                    logger.debug("Creating a variant object of variant {0}".format(
                        variant.get('variant_id')))

                    variant['index'] = index
                    logger.debug("Updating index to: {0}".format(
                        index))

                    variant['start'] = int(variant_dict['POS'])

                    variant['stop'] = int(variant_dict['POS']) + \
                        (len(variant_dict['REF']) - len(variant_dict['ALT']))

                    # It would be easy to update these keys...
                    thousand_g = info_dict.get('1000GAF')
                    if thousand_g:
                        logger.debug("Updating thousand_g to: {0}".format(
                            thousand_g))
                        variant['thousand_g'] = float(thousand_g)
                    variant.add_frequency('1000GAF', variant.get('thousand_g'))

                    cadd_score = info_dict.get('CADD')
                    if cadd_score:
                        logger.debug("Updating cadd_score to: {0}".format(
                            cadd_score))
                        variant['cadd_score'] = float(cadd_score)

                    rank_score_entry = info_dict.get('RankScore')
                    if rank_score_entry:
                        for family_annotation in rank_score_entry.split(','):
                            rank_score = family_annotation.split(':')[-1]
                        logger.debug("Updating rank_score to: {0}".format(
                            rank_score))
                        variant['rank_score'] = float(rank_score)

                    genetic_models_entry = info_dict.get('GeneticModels')
                    if genetic_models_entry:
                        genetic_models = []
                        for family_annotation in genetic_models_entry.split(','):
                            for genetic_model in family_annotation.split(':')[-1].split('|'):
                                genetic_models.append(genetic_model)
                        logger.debug("Updating rank_score to: {0}".format(
                            rank_score))
                        variant['genetic_models'] = genetic_models

                    #Add genotype calls:
                    if individuals:
                        for sample_id in individuals:
                            raw_call = dict(zip(
                                variant_dict['FORMAT'].split(':'),
                                variant_dict[sample_id].split(':'))
                            )
                            variant.add_individual(Genotype(
                                sample_id = sample_id,
                                genotype = raw_call.get('GT', './.'),
                                ref_depth = raw_call.get('AD', ',').split(',')[0],
                                alt_depth = raw_call.get('AD', ',').split(',')[1],
                                genotype_quality = raw_call.get('GQ', '.'),
                                depth = raw_call.get('DP', '.')
                            ))

                    # Add transcript information:
                    if vep_string:
                        self._add_transcripts(variant, vep_dict)

                    variant['most_severe_consequence'] = get_most_severe_consequence(
                        variant['transcripts']
                    )
                    self._add_genes(variant)
                    self._add_compounds(variant=variant, info_dict=info_dict)
                    yield variant

    def variants(self, case_id, skip=0, count=30, gene_list=None,
                 thousand_g=None):
        """Return all variants in the VCF.

            Args:
                case_id (str): Path to a vcf file (for this adapter)
                skip (int): Skip first variants
                count (int): The number of variants to return
                gene_list (list): A list of genes
                thousand_g (float): filter variants based on frequency
        """
        vcf_path = case_id.replace('|', '/')
        limit = count + skip

        filtered_variants = self._variants(vcf_path)
        if gene_list:
            gene_list = set(gene_list)
            filtered_variants = (variant for variant in filtered_variants
                                 if (set(gene['symbol'] for gene in variant['genes'])
                                     .intersection(gene_list)))
        if thousand_g:
            filtered_variants = (variant for variant in filtered_variants
                                 if variant['thousand_g'] <= thousand_g)

        for index, variant_obj in enumerate(filtered_variants):
            if index >= skip:
                if index <= limit:
                    yield variant_obj
                else:
                    break

    def variant(self, case_id, variant_id):
        """Return a specific variant.

            Args:
                case_id (str): Path to vcf file
                variant_id (str): A variant id

            Returns:
                variant (Variant): The variant object for the given id
        """
        for variant_obj in self.variants(case_id, count=float('inf')):
            if variant_obj['variant_id'] == variant_id:
                return variant_obj
        return None

if __name__ == '__main__':
    import sys
    from pprint import pprint as pp
    vcf_file = sys.argv[1]
    plugin = Plugin()
    for variant in plugin.variants(case_id=vcf_file):
        print(variant)
        print('')
<commit_msg>Updated vcf_plugin to inherit form Plugin<commit_after># -*- coding: utf-8 -*-
import logging

from path import path

from puzzle.models import (Case, Compound, Variant, Gene, Genotype, Transcript)
from puzzle.utils import (get_most_severe_consequence, get_hgnc_symbols,
                          get_omim_number, get_ensembl_id)

from puzzle.plugins import Plugin

from vcftoolbox import (get_variant_dict, HeaderParser, get_info_dict,
                        get_vep_dict)

logger = logging.getLogger(__name__)

class VcfPlugin(Plugin):
    """docstring for Plugin"""

    def __init__(self):
        super(VcfPlugin, self).__init__()

    def init_app(self, app):
        """Initialize plugin via Flask."""
        self.root_path = app.config['PUZZLE_ROOT']
        self.pattern = app.config['PUZZLE_PATTERN']

    def _find_vcfs(self, pattern='*.vcf'):
        """Walk subdirectories and return VCF files."""
        return path(self.root_path).walkfiles(pattern)

    def cases(self, pattern=None):
        """Return all VCF file paths."""
        pattern = pattern or self.pattern
        vcfs = self._find_vcfs(pattern)
        case_objs = (Case(case_id=vcf.replace('/', '|'),
                          name=vcf.basename()) for vcf in vcfs)
        return case_objs

    def _add_compounds(self, variant, info_dict):
        """Check if there are any compounds and add them to the variant

        """
        compound_entry = info_dict.get('Compounds')
        if compound_entry:
            for family_annotation in compound_entry.split(','):
                compounds = family_annotation.split(':')[-1].split('|')
                for compound in compounds:
                    splitted_compound = compound.split('>')

                    compound_score = None
                    if len(splitted_compound) > 1:
                        compound_id = splitted_compound[0]
                        compound_score = splitted_compound[-1]

                    variant.add_compound(Compound(
                        variant_id=compound_id,
                        combined_score=compound_score
                    ))

    def _add_genes(self, variant):
        """Add the genes for a variant"""

        hgnc_symbols = get_hgnc_symbols(
            transcripts = variant['transcripts']
        )

        for hgnc_id in hgnc_symbols:
            variant.add_gene(Gene(
                symbol=hgnc_id,
                omim_number=get_omim_number(hgnc_id),
                ensembl_id=get_ensembl_id(hgnc_id)
            ))

    def _add_transcripts(self, variant, vep_dict):
        """Add the transcripts for a variant"""
        for allele in vep_dict:
            for transcript_info in vep_dict[allele]:
                variant.add_transcript(Transcript(
                    SYMBOL = transcript_info.get('SYMBOL'),
                    Feature = transcript_info.get('Feature'),
                    BIOTYPE = transcript_info.get('BIOTYPE'),
                    Consequence = transcript_info.get('Consequence'),
                    STRAND = transcript_info.get('STRAND'),
                    SIFT = transcript_info.get('SIFT'),
                    PolyPhen = transcript_info.get('PolyPhen'),
                    EXON = transcript_info.get('EXON'),
                    HGVSc = transcript_info.get('HGVSc'),
                    HGVSp = transcript_info.get('HGVSp')
                ))

    def _variants(self, vcf_file_path):
        head = HeaderParser()
        # Parse the header
        with open(vcf_file_path, 'r') as variant_file:
            for line in variant_file:
                line = line.rstrip()
                if line.startswith('#'):
                    if line.startswith('##'):
                        head.parse_meta_data(line)
                    else:
                        head.parse_header_line(line)
                else:
                    break

        header_line = head.header
        individuals = head.individuals

        variant_columns = ['CHROM', 'POS', 'ID', 'REF', 'ALT', 'QUAL', 'FILTER']

        vep_header = head.vep_columns

        with open(vcf_file_path, 'r') as vcf_file:
            index = 0
            for variant_line in vcf_file:
                if not variant_line.startswith('#'):
                    index += 1
                    #Create a variant dict:
                    variant_dict =  get_variant_dict(
                        variant_line = variant_line,
                        header_line = header_line
                    )
                    #Crreate a info dict:
                    info_dict = get_info_dict(
                        info_line = variant_dict['INFO']
                    )
                    vep_string = info_dict.get('CSQ')

                    if vep_string:
                        vep_dict = get_vep_dict(
                            vep_string = vep_string,
                            vep_header = vep_header
                        )

                    variant = Variant(
                        **{column: variant_dict.get(column, '.')
                            for column in variant_columns}
                        )

                    logger.debug("Creating a variant object of variant {0}".format(
                        variant.get('variant_id')))

                    variant['index'] = index
                    logger.debug("Updating index to: {0}".format(
                        index))

                    variant['start'] = int(variant_dict['POS'])

                    variant['stop'] = int(variant_dict['POS']) + \
                        (len(variant_dict['REF']) - len(variant_dict['ALT']))

                    # It would be easy to update these keys...
                    thousand_g = info_dict.get('1000GAF')
                    if thousand_g:
                        logger.debug("Updating thousand_g to: {0}".format(
                            thousand_g))
                        variant['thousand_g'] = float(thousand_g)
                    variant.add_frequency('1000GAF', variant.get('thousand_g'))

                    cadd_score = info_dict.get('CADD')
                    if cadd_score:
                        logger.debug("Updating cadd_score to: {0}".format(
                            cadd_score))
                        variant['cadd_score'] = float(cadd_score)

                    rank_score_entry = info_dict.get('RankScore')
                    if rank_score_entry:
                        for family_annotation in rank_score_entry.split(','):
                            rank_score = family_annotation.split(':')[-1]
                        logger.debug("Updating rank_score to: {0}".format(
                            rank_score))
                        variant['rank_score'] = float(rank_score)

                    genetic_models_entry = info_dict.get('GeneticModels')
                    if genetic_models_entry:
                        genetic_models = []
                        for family_annotation in genetic_models_entry.split(','):
                            for genetic_model in family_annotation.split(':')[-1].split('|'):
                                genetic_models.append(genetic_model)
                        logger.debug("Updating rank_score to: {0}".format(
                            rank_score))
                        variant['genetic_models'] = genetic_models

                    #Add genotype calls:
                    if individuals:
                        for sample_id in individuals:
                            raw_call = dict(zip(
                                variant_dict['FORMAT'].split(':'),
                                variant_dict[sample_id].split(':'))
                            )
                            variant.add_individual(Genotype(
                                sample_id = sample_id,
                                genotype = raw_call.get('GT', './.'),
                                ref_depth = raw_call.get('AD', ',').split(',')[0],
                                alt_depth = raw_call.get('AD', ',').split(',')[1],
                                genotype_quality = raw_call.get('GQ', '.'),
                                depth = raw_call.get('DP', '.')
                            ))

                    # Add transcript information:
                    if vep_string:
                        self._add_transcripts(variant, vep_dict)

                    variant['most_severe_consequence'] = get_most_severe_consequence(
                        variant['transcripts']
                    )
                    self._add_genes(variant)
                    self._add_compounds(variant=variant, info_dict=info_dict)
                    yield variant

    def variants(self, case_id, skip=0, count=30, gene_list=None,
                 thousand_g=None):
        """Return all variants in the VCF.

            Args:
                case_id (str): Path to a vcf file (for this adapter)
                skip (int): Skip first variants
                count (int): The number of variants to return
                gene_list (list): A list of genes
                thousand_g (float): filter variants based on frequency
        """
        vcf_path = case_id.replace('|', '/')
        limit = count + skip

        filtered_variants = self._variants(vcf_path)
        if gene_list:
            gene_list = set(gene_list)
            filtered_variants = (variant for variant in filtered_variants
                                 if (set(gene['symbol'] for gene in variant['genes'])
                                     .intersection(gene_list)))
        if thousand_g:
            filtered_variants = (variant for variant in filtered_variants
                                 if variant['thousand_g'] <= thousand_g)

        for index, variant_obj in enumerate(filtered_variants):
            if index >= skip:
                if index <= limit:
                    yield variant_obj
                else:
                    break

    def variant(self, case_id, variant_id):
        """Return a specific variant.

            Args:
                case_id (str): Path to vcf file
                variant_id (str): A variant id

            Returns:
                variant (Variant): The variant object for the given id
        """
        for variant_obj in self.variants(case_id, count=float('inf')):
            if variant_obj['variant_id'] == variant_id:
                return variant_obj
        return None

if __name__ == '__main__':
    import sys
    from pprint import pprint as pp
    vcf_file = sys.argv[1]
    plugin = VcfPlugin()
    for variant in plugin.variants(case_id=vcf_file):
        print(variant)
        print('')
<|endoftext|>


<commit_before>import collections

def read_line_to_list(as_type=int):
    lst = raw_input().strip().split(' ')
    if as_type is None:
        return lst
    return map(as_type, lst)

class TripPlanner(object):
    hours_per_day = 8

    def __init__(self, budget, price_per_day, num_attr, attractions):
        self.budget = budget
        self.price_per_day = price_per_day
        self.num_attractions = num_attr
        self.attractions = [(score, -visit_time, -cost)
                            for cost, visit_time, score in attractions]
        self.attractions = sorted(self.attractions)
        for i in self.attractions:
            print i

    def max_score_sum(self):
        return self.knapsack(self.budget, 0, 0, self.num_attractions)

    def knapsack(self, budget, time_cnt, day_cnt, n):
        if n == 0 or budget == 0:
            return 0

        info = self.attractions[n - 1]
        score, visit_time, cost = info[0], -info[1], -info[2]
        pass_night = int(time_cnt + visit_time / self.hours_per_day) > day_cnt
        if pass_night:
            cost += self.price_per_day
            add_additional_day = 1
        else:
            add_additional_day = 0

        # Check if we can include the n-th attraction
        score_if_not = self.knapsack(budget, time_cnt, day_cnt, n - 1)

        if budget < cost:
            return score_if_not
        else:
            score_if_visit_current = score + self.knapsack(
                budget - cost, time_cnt + visit_time,
                day_cnt + add_additional_day,
                n - 1
            )
            return max(score_if_visit_current, score_if_not)

B, P, N = read_line_to_list()
ATTR = [map(int, read_line_to_list(None)[1:]) for _ in range(N)]

planner = TripPlanner(B, P, N, ATTR)
print planner.max_score_sum()<commit_msg>add tripplanner - fail #16 timeout #19<commit_after>import collections

def read_line_to_list(as_type=int):
    lst = raw_input().strip().split(' ')
    if as_type is None:
        return lst
    return map(as_type, lst)

class TripPlanner(object):
    hours_per_day = 8

    def __init__(self, budget, price_per_day, num_attr, attractions):
        self.budget = budget
        self.price_per_day = price_per_day
        self.num_attractions = num_attr
        self.attractions = attractions

    def max_score_sum(self):
        max_sum = collections.defaultdict(int)
        max_sum[(self.budget, 0)] = 0

        max_sum_of_all = 0
        for cost, visit_time, score in self.attractions:
            new_max_sum = {}
            # print "\n", max_sum
            for (budget, time_cnt), score_sum in max_sum.items():
                day_cnt, past_time = divmod(time_cnt, self.hours_per_day)
                remaining_time = self.hours_per_day - past_time
                remaining_budget = budget - cost
                new_time_cnt = time_cnt + visit_time
                if visit_time > remaining_time:
                    remaining_budget -= self.price_per_day
                    new_time_cnt += remaining_time
                elif past_time == 0:
                    remaining_budget -= self.price_per_day

                if remaining_budget < 0:
                    continue
                new_key = (remaining_budget, time_cnt + visit_time)
                this_max_sum = max(max_sum[new_key], score_sum + score)
                new_max_sum[new_key] = this_max_sum
                max_sum_of_all = max(max_sum_of_all, this_max_sum)
            max_sum.update(new_max_sum)
        return max_sum_of_all

B, P, N = read_line_to_list()
ATTR = [map(int, read_line_to_list(None)[1:]) for _ in range(N)]

planner = TripPlanner(B, P, N, ATTR)
print planner.max_score_sum()
<|endoftext|>


<commit_before>'''
Created on Jan 18, 2017

@author: Daniel
'''

from tf.fabric import Fabric, Locality as L, Text as T
from bibleutils.versification import Identifier
from inspect import currentframe

class __Corpus(Identifier):
    '''Defines the corpus identifiers. Essentially this is just Hebrew or Greek
    right.
    '''
    # FIXME would like to distinguish work and corpus eventually but not
    # necessary now.

    def __init__(self):
        super().__init__({'GREEK' : 1,
                          'HEBREW' : 2})

    @property
    def GREEK(self):
        return self._map.get(currentframe().f_code.co_name)

    @property
    def HEBREW(self):
        return self._map.get(currentframe().f_code.co_name)

Corpus = __Corpus()

def get_words(*refs, work=Corpus.HEBREW):
    '''
    Get a list of unique words in the specified work, book, chapter and verses.

    Parameters
        refs is a list of tuples, each tuple being book, chapter and verse.
            books are string names as defined in Text-Fabric for English.
            chapter and verse are integers.
    '''

    if work == Corpus.HEBREW:
        work_home = 'hebrew/etcbc4c'
        word_feature = 'g_word_utf8'
    elif work == Corpus.GREEK:
        work_home = 'greek/sblgnt'
        word_feature = 'Unicode_no_punctuations'

    TF = Fabric(modules=[work_home], silent=True)
    api = TF.load(word_feature, silent=True)
    api.makeAvailableIn(globals())
    lex_list = []
    for ref in refs:
        verse_nodes = T.nodeFromSection(ref)
        slots = L.d(verse_nodes, 'word')
        lex_list.extend([api.Fs(word_feature).v(s) for s in slots])
    return lex_list
<commit_msg>update for g_word without punctuation<commit_after>'''
Created on Jan 18, 2017

@author: Daniel
'''

from tf.fabric import Fabric, Locality as L, Text as T
from bibleutils.versification import Identifier
from inspect import currentframe

class __Corpus(Identifier):
    '''Defines the corpus identifiers. Essentially this is just Hebrew or Greek
    right.
    '''
    # FIXME would like to distinguish work and corpus eventually but not
    # necessary now.

    def __init__(self):
        super().__init__({'GREEK' : 1,
                          'HEBREW' : 2})

    @property
    def GREEK(self):
        return self._map.get(currentframe().f_code.co_name)

    @property
    def HEBREW(self):
        return self._map.get(currentframe().f_code.co_name)

Corpus = __Corpus()

def get_words(*refs, work=Corpus.HEBREW):
    '''
    Get a list of unique words in the specified work, book, chapter and verses.

    Parameters
        refs is a list of tuples, each tuple being book, chapter and verse.
            books are string names as defined in Text-Fabric for English.
            chapter and verse are integers.
    '''

    if work == Corpus.HEBREW:
        work_home = 'hebrew/etcbc4c'
        word_feature = 'g_word_utf8'
    elif work == Corpus.GREEK:
        work_home = 'greek/sblgnt'
        word_feature = 'g_word'

    TF = Fabric(modules=[work_home], silent=True)
    api = TF.load(word_feature, silent=True)
    api.makeAvailableIn(globals())
    lex_list = []
    for ref in refs:
        verse_nodes = T.nodeFromSection(ref)
        slots = L.d(verse_nodes, 'word')
        lex_list.extend([api.Fs(word_feature).v(s) for s in slots])
    return lex_list
<|endoftext|>


<commit_before>from pygments.lexer import RegexLexer
from pygments.lexer import include
from pygments.token import Comment
from pygments.token import Name
from pygments.token import Number
from pygments.token import Punctuation
from pygments.token import String
from pygments.token import Text

__all__ = ["PwntoolsLexer"]

class PwntoolsLexer(RegexLexer):
    """
    Fork from pwntools
    https://github.com/Gallopsled/pwntools/blob/7860eecf025135380b137dd9df85dd02a2fd1667/pwnlib/lexer.py

    Edit:
        * Remove Objdump rules
        * Merge pygments-arm (https://github.com/heia-fr/pygments-arm)
    """

    name = "PwntoolsLexer"
    filenames = ["*.s", "*.S", "*.asm"]

    #: optional Comment or Whitespace
    string = r'"(\\"|[^"])*"'
    char = r"[\w$.@-]"
    identifier = r"(?:[a-zA-Z$_]" + char + r"*|\." + char + r"+|or)"
    number = r"(?:0[xX][a-zA-Z0-9]+|\d+)"
    memory = r"(?:[\]\[])"

    eol = r"[\r\n]+"

    tokens = {
        "root": [
            include("whitespace"),
            # Label
            (identifier + ":", Name.Label),
            (number + ":", Name.Label),
            # AT&T directive
            (r"\." + identifier, Name.Attribute, "directive-args"),
            (r"lock|rep(n?z)?|data\d+", Name.Attribute),
            # Instructions
            (identifier, Name.Function, "instruction-args"),
            (r"[\r\n]+", Text),
        ],
        "directive-args": [
            (identifier, Name.Constant),
            (string, String),
            ("@" + identifier, Name.Attribute),
            (number, Number.Integer),
            (eol, Text, "#pop"),
            (r"#.*?$", Comment, "#pop"),
            include("punctuation"),
            include("whitespace"),
        ],
        "instruction-args": [
            # Fun things
            (r"([\]\[]|BYTE|DWORD|PTR|\+|\-|}|{|\^|>>|<<|&)", Text),
            # Address constants
            (identifier, Name.Constant),
            ("=" + identifier, Name.Constant),  # ARM symbol
            (number, Number.Integer),
            # Registers
            ("%" + identifier, Name.Variable),
            ("$" + identifier, Name.Variable),
            # Numeric constants
            ("$" + number, Number.Integer),
            ("#" + number, Number.Integer),
            # ARM predefined constants
            ("#" + identifier, Name.Constant),
            (r"$'(.|\\')'", String.Char),
            (eol, Text, "#pop"),
            include("punctuation"),
            include("whitespace"),
        ],
        "whitespace": [
            (r"\n", Text),
            (r"\s+", Text),
            # Block comments
            # /* */ (AT&T)
            (r"/\*.*?\*/", Comment),
            # Line comments
            # //    (AArch64)
            # #     (AT&T)
            # ;     (NASM/intel, LLVM)
            # @     (ARM)
            (r"(//|[#;@]).*$", Comment.Single),
        ],
        "punctuation": [(r"[-*,.():]+", Punctuation)],
    }
<commit_msg>Fix lexer for coloring negative numbers in asm (#1367)<commit_after>from pygments.lexer import RegexLexer
from pygments.lexer import include
from pygments.token import Comment
from pygments.token import Name
from pygments.token import Number
from pygments.token import Punctuation
from pygments.token import String
from pygments.token import Text

__all__ = ["PwntoolsLexer"]

class PwntoolsLexer(RegexLexer):
    """
    Fork from pwntools
    https://github.com/Gallopsled/pwntools/blob/7860eecf025135380b137dd9df85dd02a2fd1667/pwnlib/lexer.py

    Edit:
        * Remove Objdump rules
        * Merge pygments-arm (https://github.com/heia-fr/pygments-arm)
    """

    name = "PwntoolsLexer"
    filenames = ["*.s", "*.S", "*.asm"]

    #: optional Comment or Whitespace
    string = r'"(\\"|[^"])*"'
    char = r"[\w$.@-]"
    identifier = r"(?:[a-zA-Z$_]" + char + r"*|\." + char + r"+|or)"
    number = r"(?:\-?0[xX][a-zA-Z0-9]+|\d+)"
    memory = r"(?:[\]\[])"

    eol = r"[\r\n]+"

    tokens = {
        "root": [
            include("whitespace"),
            # Label
            (identifier + ":", Name.Label),
            (number + ":", Name.Label),
            # AT&T directive
            (r"\." + identifier, Name.Attribute, "directive-args"),
            (r"lock|rep(n?z)?|data\d+", Name.Attribute),
            # Instructions
            (identifier, Name.Function, "instruction-args"),
            (r"[\r\n]+", Text),
        ],
        "directive-args": [
            (identifier, Name.Constant),
            (string, String),
            ("@" + identifier, Name.Attribute),
            (number, Number.Integer),
            (eol, Text, "#pop"),
            (r"#.*?$", Comment, "#pop"),
            include("punctuation"),
            include("whitespace"),
        ],
        "instruction-args": [
            # Fun things
            (r"([\]\[]|BYTE|DWORD|PTR|\+|\-|}|{|\^|>>|<<|&)", Text),
            # Address constants
            (identifier, Name.Constant),
            ("=" + identifier, Name.Constant),  # ARM symbol
            (number, Number.Integer),
            # Registers
            ("%" + identifier, Name.Variable),
            ("$" + identifier, Name.Variable),
            # Numeric constants
            ("$" + number, Number.Integer),
            ("#" + number, Number.Integer),
            # ARM predefined constants
            ("#" + identifier, Name.Constant),
            (r"$'(.|\\')'", String.Char),
            (eol, Text, "#pop"),
            include("punctuation"),
            include("whitespace"),
        ],
        "whitespace": [
            (r"\n", Text),
            (r"\s+", Text),
            # Block comments
            # /* */ (AT&T)
            (r"/\*.*?\*/", Comment),
            # Line comments
            # //    (AArch64)
            # #     (AT&T)
            # ;     (NASM/intel, LLVM)
            # @     (ARM)
            (r"(//|[#;@]).*$", Comment.Single),
        ],
        "punctuation": [(r"[-*,.():]+", Punctuation)],
    }
<|endoftext|>


<commit_before>"""Plan future DESI observations.
"""
from __future__ import print_function, division

import numpy as np

import astropy.io.fits as fits
import astropy.units as u
import astropy.table
import astropy.time

import specsim.atmosphere

import desimodel.io

import desiutil.log

import desisurvey.ephemerides
import desisurvey.exposurecalc
import desisurvey.config
import desisurvey.utils

def initialize(ephem, start_date=None, stop_date=None, step_size=5.*u.min,
               healpix_nside=16, output_name='planner.fits'):
    """Calculate exposure-time factors over a grid of times and pointings.

    Takes about 6 minutes to run and writes a 1.3Gb output file with the
    default parameters.

    Requires that healpy is installed.

    Parameters
    ----------
    ephem : `desisurvey.ephem.Ephemerides`
        Tabulated ephemerides data to use for planning.
    start_date : date or None
        Survey planning starts on the evening of this date. Must be convertible
        to a date using :func:`desisurvey.utils.get_date`.  Use the first night
        of the ephemerides when None.
    stop_date : date or None
        Survey planning stops on the morning of this date. Must be convertible
        to a date using :func:`desisurvey.utils.get_date`.  Use the first night
        of the ephemerides when None.
    step_size : astropy.units.Quantity
        Exposure-time factors are tabulated at this interval during each night.
    healpix_nside : int
        Healpix NSIDE parameter to use for binning the sky. Must be a power of
        two.  Values larger than 16 will lead to holes in the footprint with
        the current implementation.
    output_name : str
        Name of the FITS output file where results are saved. A relative path
        refers to the :meth:`configuration output path
        <desisurvey.config.Configuration.get_path>`.
    """
    import healpy

    log = desiutil.log.get_logger()

    config = desisurvey.config.Configuration()
    output_name = config.get_path(output_name)

    start_date = desisurvey.utils.get_date(start_date or ephem.start)
    stop_date = desisurvey.utils.get_date(stop_date or ephem.stop)
    if start_date >= stop_date:
        raise ValueError('Expected start_date < stop_date.')
    mjd = ephem._table['noon']
    sel = ((mjd >= desisurvey.utils.local_noon_on_date(start_date).mjd) &
           (mjd < desisurvey.utils.local_noon_on_date(stop_date).mjd))
    t = ephem._table[sel]
    num_nights = len(t)

    # Build a grid of elapsed time relative to local midnight during each night.
    midnight = t['noon'] + 0.5
    t_edges = desisurvey.ephemerides.get_grid(step_size)
    t_centers = 0.5 * (t_edges[1:] + t_edges[:-1])
    num_points = len(t_centers)

    # Create an empty HDU0 with header info.
    header = fits.Header()
    header['START'] = str(start_date)
    header['STOP'] = str(stop_date)
    header['NSIDE'] = healpix_nside
    header['NPOINTS'] = num_points
    header['STEP'] = step_size.to(u.min).value
    hdus = fits.HDUList()
    hdus.append(fits.ImageHDU(header=header))

    # Save time grid.
    hdus.append(fits.ImageHDU(name='GRID', data=t_edges))

    # Load the list of tiles to observe.
    tiles = astropy.table.Table(
        desimodel.io.load_tiles(onlydesi=True, extra=False))

    # Build the footprint as a healpix map of the requested size.
    # The footprint includes any pixel containing at least one tile center.
    npix = healpy.nside2npix(healpix_nside)
    footprint = np.zeros(npix, bool)
    pixels = healpy.ang2pix(
            healpix_nside, np.radians(90 - tiles['DEC'].data),
            np.radians(tiles['RA'].data))
    footprint[np.unique(pixels)] = True
    footprint_pixels = np.where(footprint)[0]
    num_footprint = len(footprint_pixels)
    log.info('Footprint contains {0} pixels.'.format(num_footprint))

    # Sort pixels in order of increasing phi + 60deg so that the north and south
    # galactic caps are contiguous in the arrays we create below.
    pix_theta, pix_phi = healpy.pix2ang(healpix_nside, footprint_pixels)
    pix_dphi = np.fmod(pix_phi + np.pi / 3, 2 * np.pi)
    sort_order = np.argsort(pix_dphi)
    footprint_pixels = footprint_pixels[sort_order]

    # Average E(B-V) for all tiles falling into a pixel.
    tiles_per_pixel = np.bincount(pixels, minlength=npix)
    EBV = np.bincount(pixels, weights=tiles['EBV_MED'], minlength=npix)
    EBV[footprint] /= tiles_per_pixel[footprint]

    # Calculate dust extinction exposure-time factor.
    Ag = 3.303 * EBV # Use g-band
    f_EBV = 10.0 ** (-2.0 * Ag / 2.5)

    # Save HDU with the footprint and static dust exposure map.
    table = astropy.table.Table()
    table['pixel'] = footprint_pixels
    table['dust'] = f_EBV[footprint_pixels]
    hdu = fits.table_to_hdu(table)
    hdu.name = 'STATIC'
    hdus.append(hdu)

    # Prepare a table of calendar data.
    calendar = astropy.table.Table()
    calendar['midnight'] = midnight
    calendar['monsoon'] = np.zeros(num_nights, bool)
    calendar['fullmoon'] = np.zeros(num_nights, bool)

    # Prepare a table of ephemeris data.
    etable = astropy.table.Table()
    etable['program'] = np.zeros(num_nights * num_points, dtype=np.int16)
    etable['moon_frac'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['moon_ra'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['moon_dec'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['zenith_ra'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['zenith_dec'] = np.zeros(num_nights * num_points, dtype=np.float32)

    # Build sky coordinates for each pixel in the footprint.
    pix_theta, pix_phi = healpy.pix2ang(healpix_nside, footprint_pixels)
    pix_ra, pix_dec = np.degrees(pix_phi), 90 - np.degrees(pix_theta)
    pix_sky = astropy.coordinates.ICRS(pix_ra * u.deg, pix_dec * u.deg)

    # Initialize exposure factor calculations.
    alt, az = np.full(num_points, 90.) * u.deg, np.zeros(num_points) * u.deg
    fexp = np.zeros((num_nights * num_points, num_footprint), dtype=np.float32)
    vband_extinction = 0.15154
    one = np.ones((num_points, num_footprint))

    # Loop over nights.
    for i in range(num_nights):
        night = ephem.get_night(midnight[i])
        date = desisurvey.utils.get_date(midnight[i])
        if date.day == 1:
            log.info('Starting {0} (completed {1}/{2} nights)'
                     .format(date.strftime('%b %Y'), i, num_nights))
        # Initialize the slice of the fexp[] time index for this night.
        sl = slice(i * num_points, (i + 1) * num_points)
        # Skip monsoon and full moon.
        calendar[i]['monsoon'] = desisurvey.utils.is_monsoon(midnight[i])
        calendar[i]['fullmoon'] = ephem.is_full_moon(midnight[i])
        if calendar[i]['monsoon'] or calendar[i]['fullmoon']:
            continue
        # Calculate the program during this night.
        mjd = midnight[i] + t_centers
        dark, gray, bright = ephem.get_program(mjd)
        etable['program'][sl][dark] = 1
        etable['program'][sl][gray] = 2
        etable['program'][sl][bright] = 3
        # Zero the exposure factor whenever we are not oberving.
        fexp[sl] = (dark | gray | bright)[:, np.newaxis]
        # Transform the local zenith to (ra,dec).
        times = astropy.time.Time(mjd, format='mjd')
        zenith = desisurvey.utils.get_observer(
            times, alt=alt, az=az).transform_to(astropy.coordinates.ICRS)
        etable['zenith_ra'][sl] = zenith.ra.to(u.deg).value
        etable['zenith_dec'][sl] = zenith.dec.to(u.deg).value
        # Calculate zenith angles to each pixel in the footprint.
        pix_sep = pix_sky.separation(zenith[:, np.newaxis])
        # Zero the exposure factor for pixels below the horizon.
        visible = pix_sep < 90 * u.deg
        fexp[sl][~visible] = 0.
        # Calculate the airmass exposure-time penalty
        # as 1/X**1.25 with X = 1/cos(pix_sep).
        fexp[sl][visible] *= np.cos(pix_sep[visible]) ** 1.25
        # Loop over objects we need to avoid.
        for name in config.avoid_bodies.keys:
            f_obj = desisurvey.ephemerides.get_object_interpolator(night, name)
            # Calculate this object's (dec,ra) path during the night.
            obj_dec, obj_ra = f_obj(times.mjd)
            sky_obj = astropy.coordinates.ICRS(
                ra=obj_ra[:, np.newaxis] * u.deg,
                dec=obj_dec[:, np.newaxis] * u.deg)
            # Calculate the separation angles to each pixel in the footprint.
            obj_sep = pix_sky.separation(sky_obj)
            if name == 'moon':
                etable['moon_ra'][sl] = obj_ra
                etable['moon_dec'][sl] = obj_dec
                # Calculate moon altitude during the night.
                moon_alt, _ = desisurvey.ephemerides.get_object_interpolator(
                    night, 'moon', altaz=True)(times.mjd)
                moon_zenith = (90 - moon_alt[:,np.newaxis]) * u.deg
                moon_up = moon_alt > 0
                assert np.all(moon_alt[gray] > 0)
                # Calculate the moon illuminated fraction during the night.
                moon_frac = ephem.get_moon_illuminated_fraction(times.mjd)
                etable['moon_frac'][sl] = moon_frac
                # Convert to temporal moon phase.
                moon_phase = np.arccos(2 * moon_frac[:,np.newaxis] - 1) / np.pi
                # Calculate scattered moon V-band brightness at each pixel.
                V = specsim.atmosphere.krisciunas_schaefer(
                    pix_sep, moon_zenith, obj_sep,
                    moon_phase, vband_extinction).value
                # Estimate the exposure time factor from V.
                X = np.dstack((one, np.exp(-V), 1/V, 1/V**2, 1/V**3))
                T = X.dot(desisurvey.exposurecalc._moonCoefficients)
                # No penalty when the moon is below the horizon.
                T[moon_alt < 0, :] = 1.
                fexp[sl] *= 1. / T
            else:
                # Lookup the avoidance size for this object.
                size = getattr(config.avoid_bodies, name)()
                # Penalize the exposure-time with a factor
                # 1 - exp(-0.5*(obj_sep/size)**2)
                penalty = 1. - np.exp(-0.5 * (obj_sep / size) ** 2)
                fexp[sl] *= penalty

    # Save calendar table.
    hdu = fits.table_to_hdu(calendar)
    hdu.name = 'CALENDAR'
    hdus.append(hdu)

    # Save ephemerides table.
    hdu = fits.table_to_hdu(etable)
    hdu.name = 'EPHEM'
    hdus.append(hdu)

    # Save dynamic exposure-time factors.
    hdus.append(fits.ImageHDU(name='DYNAMIC', data=fexp))

    # Finalize the output file.
    hdus.writeto(output_name, overwrite=True)
    log.info('Plan initialization saved to {0}'.format(output_name))

if __name__ == '__main__':

    #stop = desisurvey.utils.get_date('2019-10-03')
    #stop = desisurvey.utils.get_date('2020-07-13')
    stop = None
    ephem = desisurvey.ephemerides.Ephemerides(stop_date=stop)
    initialize(ephem, stop_date=stop)
<commit_msg>Save some tile info in the planner init file<commit_after>"""Plan future DESI observations.
"""
from __future__ import print_function, division

import numpy as np

import astropy.io.fits as fits
import astropy.units as u
import astropy.table
import astropy.time

import specsim.atmosphere

import desimodel.io

import desiutil.log

import desisurvey.ephemerides
import desisurvey.exposurecalc
import desisurvey.config
import desisurvey.utils

def initialize(ephem, start_date=None, stop_date=None, step_size=5.*u.min,
               healpix_nside=16, output_name='planner.fits'):
    """Calculate exposure-time factors over a grid of times and pointings.

    Takes about 6 minutes to run and writes a 1.3Gb output file with the
    default parameters.

    Requires that healpy is installed.

    Parameters
    ----------
    ephem : `desisurvey.ephem.Ephemerides`
        Tabulated ephemerides data to use for planning.
    start_date : date or None
        Survey planning starts on the evening of this date. Must be convertible
        to a date using :func:`desisurvey.utils.get_date`.  Use the first night
        of the ephemerides when None.
    stop_date : date or None
        Survey planning stops on the morning of this date. Must be convertible
        to a date using :func:`desisurvey.utils.get_date`.  Use the first night
        of the ephemerides when None.
    step_size : astropy.units.Quantity
        Exposure-time factors are tabulated at this interval during each night.
    healpix_nside : int
        Healpix NSIDE parameter to use for binning the sky. Must be a power of
        two.  Values larger than 16 will lead to holes in the footprint with
        the current implementation.
    output_name : str
        Name of the FITS output file where results are saved. A relative path
        refers to the :meth:`configuration output path
        <desisurvey.config.Configuration.get_path>`.
    """
    import healpy

    log = desiutil.log.get_logger()

    config = desisurvey.config.Configuration()
    output_name = config.get_path(output_name)

    start_date = desisurvey.utils.get_date(start_date or ephem.start)
    stop_date = desisurvey.utils.get_date(stop_date or ephem.stop)
    if start_date >= stop_date:
        raise ValueError('Expected start_date < stop_date.')
    mjd = ephem._table['noon']
    sel = ((mjd >= desisurvey.utils.local_noon_on_date(start_date).mjd) &
           (mjd < desisurvey.utils.local_noon_on_date(stop_date).mjd))
    t = ephem._table[sel]
    num_nights = len(t)

    # Build a grid of elapsed time relative to local midnight during each night.
    midnight = t['noon'] + 0.5
    t_edges = desisurvey.ephemerides.get_grid(step_size)
    t_centers = 0.5 * (t_edges[1:] + t_edges[:-1])
    num_points = len(t_centers)

    # Create an empty HDU0 with header info.
    header = fits.Header()
    header['START'] = str(start_date)
    header['STOP'] = str(stop_date)
    header['NSIDE'] = healpix_nside
    header['NPOINTS'] = num_points
    header['STEP'] = step_size.to(u.min).value
    hdus = fits.HDUList()
    hdus.append(fits.ImageHDU(header=header))

    # Save time grid.
    hdus.append(fits.ImageHDU(name='GRID', data=t_edges))

    # Load the list of tiles to observe.
    tiles = astropy.table.Table(
        desimodel.io.load_tiles(onlydesi=True, extra=False))

    # Build the footprint as a healpix map of the requested size.
    # The footprint includes any pixel containing at least one tile center.
    npix = healpy.nside2npix(healpix_nside)
    footprint = np.zeros(npix, bool)
    pixels = healpy.ang2pix(
            healpix_nside, np.radians(90 - tiles['DEC'].data),
            np.radians(tiles['RA'].data))
    footprint[np.unique(pixels)] = True
    footprint_pixels = np.where(footprint)[0]
    num_footprint = len(footprint_pixels)
    log.info('Footprint contains {0} pixels.'.format(num_footprint))

    # Sort pixels in order of increasing phi + 60deg so that the north and south
    # galactic caps are contiguous in the arrays we create below.
    pix_theta, pix_phi = healpy.pix2ang(healpix_nside, footprint_pixels)
    pix_dphi = np.fmod(pix_phi + np.pi / 3, 2 * np.pi)
    sort_order = np.argsort(pix_dphi)
    footprint_pixels = footprint_pixels[sort_order]

    # Record per-tile info needed for planning.
    table = astropy.table.Table()
    table['tileid'] = tiles['TILEID']
    # Map each tile ID to the corresponding index in our spatial arrays.
    mapper = np.zeros(npix, int)
    mapper[footprint_pixels] = np.arange(len(footprint_pixels))
    table['map'] = mapper[pixels].astype(np.int16)
    # Use a small int to identify the program.
    table['program'] = np.zeros(len(tiles), np.int16)
    for i, program in enumerate(('DARK', 'GRAY', 'BRIGHT')):
        table['program'][tiles['PROGRAM'] == program] = i + 1
    assert np.all(table['program'] > 0)
    hdu = fits.table_to_hdu(table)
    hdu.name = 'TILES'
    hdus.append(hdu)

    # Average E(B-V) for all tiles falling into a pixel.
    tiles_per_pixel = np.bincount(pixels, minlength=npix)
    EBV = np.bincount(pixels, weights=tiles['EBV_MED'], minlength=npix)
    EBV[footprint] /= tiles_per_pixel[footprint]

    # Calculate dust extinction exposure-time factor.
    Ag = 3.303 * EBV # Use g-band
    f_EBV = 10.0 ** (-2.0 * Ag / 2.5)

    # Save HDU with the footprint and static dust exposure map.
    table = astropy.table.Table()
    table['pixel'] = footprint_pixels
    table['dust'] = f_EBV[footprint_pixels]
    hdu = fits.table_to_hdu(table)
    hdu.name = 'STATIC'
    hdus.append(hdu)

    # Prepare a table of calendar data.
    calendar = astropy.table.Table()
    calendar['midnight'] = midnight
    calendar['monsoon'] = np.zeros(num_nights, bool)
    calendar['fullmoon'] = np.zeros(num_nights, bool)

    # Prepare a table of ephemeris data.
    etable = astropy.table.Table()
    etable['program'] = np.zeros(num_nights * num_points, dtype=np.int16)
    etable['moon_frac'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['moon_ra'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['moon_dec'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['zenith_ra'] = np.zeros(num_nights * num_points, dtype=np.float32)
    etable['zenith_dec'] = np.zeros(num_nights * num_points, dtype=np.float32)

    # Build sky coordinates for each pixel in the footprint.
    pix_theta, pix_phi = healpy.pix2ang(healpix_nside, footprint_pixels)
    pix_ra, pix_dec = np.degrees(pix_phi), 90 - np.degrees(pix_theta)
    pix_sky = astropy.coordinates.ICRS(pix_ra * u.deg, pix_dec * u.deg)

    # Initialize exposure factor calculations.
    alt, az = np.full(num_points, 90.) * u.deg, np.zeros(num_points) * u.deg
    fexp = np.zeros((num_nights * num_points, num_footprint), dtype=np.float32)
    vband_extinction = 0.15154
    one = np.ones((num_points, num_footprint))

    # Loop over nights.
    for i in range(num_nights):
        night = ephem.get_night(midnight[i])
        date = desisurvey.utils.get_date(midnight[i])
        if date.day == 1:
            log.info('Starting {0} (completed {1}/{2} nights)'
                     .format(date.strftime('%b %Y'), i, num_nights))
        # Initialize the slice of the fexp[] time index for this night.
        sl = slice(i * num_points, (i + 1) * num_points)
        # Skip monsoon and full moon.
        calendar[i]['monsoon'] = desisurvey.utils.is_monsoon(midnight[i])
        calendar[i]['fullmoon'] = ephem.is_full_moon(midnight[i])
        if calendar[i]['monsoon'] or calendar[i]['fullmoon']:
            continue
        # Calculate the program during this night.
        mjd = midnight[i] + t_centers
        dark, gray, bright = ephem.get_program(mjd)
        etable['program'][sl][dark] = 1
        etable['program'][sl][gray] = 2
        etable['program'][sl][bright] = 3
        # Zero the exposure factor whenever we are not oberving.
        fexp[sl] = (dark | gray | bright)[:, np.newaxis]
        # Transform the local zenith to (ra,dec).
        times = astropy.time.Time(mjd, format='mjd')
        zenith = desisurvey.utils.get_observer(
            times, alt=alt, az=az).transform_to(astropy.coordinates.ICRS)
        etable['zenith_ra'][sl] = zenith.ra.to(u.deg).value
        etable['zenith_dec'][sl] = zenith.dec.to(u.deg).value
        # Calculate zenith angles to each pixel in the footprint.
        pix_sep = pix_sky.separation(zenith[:, np.newaxis])
        # Zero the exposure factor for pixels below the horizon.
        visible = pix_sep < 90 * u.deg
        fexp[sl][~visible] = 0.
        # Calculate the airmass exposure-time penalty
        # as 1/X**1.25 with X = 1/cos(pix_sep).
        fexp[sl][visible] *= np.cos(pix_sep[visible]) ** 1.25
        # Loop over objects we need to avoid.
        for name in config.avoid_bodies.keys:
            f_obj = desisurvey.ephemerides.get_object_interpolator(night, name)
            # Calculate this object's (dec,ra) path during the night.
            obj_dec, obj_ra = f_obj(times.mjd)
            sky_obj = astropy.coordinates.ICRS(
                ra=obj_ra[:, np.newaxis] * u.deg,
                dec=obj_dec[:, np.newaxis] * u.deg)
            # Calculate the separation angles to each pixel in the footprint.
            obj_sep = pix_sky.separation(sky_obj)
            if name == 'moon':
                etable['moon_ra'][sl] = obj_ra
                etable['moon_dec'][sl] = obj_dec
                # Calculate moon altitude during the night.
                moon_alt, _ = desisurvey.ephemerides.get_object_interpolator(
                    night, 'moon', altaz=True)(times.mjd)
                moon_zenith = (90 - moon_alt[:,np.newaxis]) * u.deg
                moon_up = moon_alt > 0
                assert np.all(moon_alt[gray] > 0)
                # Calculate the moon illuminated fraction during the night.
                moon_frac = ephem.get_moon_illuminated_fraction(times.mjd)
                etable['moon_frac'][sl] = moon_frac
                # Convert to temporal moon phase.
                moon_phase = np.arccos(2 * moon_frac[:,np.newaxis] - 1) / np.pi
                # Calculate scattered moon V-band brightness at each pixel.
                V = specsim.atmosphere.krisciunas_schaefer(
                    pix_sep, moon_zenith, obj_sep,
                    moon_phase, vband_extinction).value
                # Estimate the exposure time factor from V.
                X = np.dstack((one, np.exp(-V), 1/V, 1/V**2, 1/V**3))
                T = X.dot(desisurvey.exposurecalc._moonCoefficients)
                # No penalty when the moon is below the horizon.
                T[moon_alt < 0, :] = 1.
                fexp[sl] *= 1. / T
            else:
                # Lookup the avoidance size for this object.
                size = getattr(config.avoid_bodies, name)()
                # Penalize the exposure-time with a factor
                # 1 - exp(-0.5*(obj_sep/size)**2)
                penalty = 1. - np.exp(-0.5 * (obj_sep / size) ** 2)
                fexp[sl] *= penalty

    # Save calendar table.
    hdu = fits.table_to_hdu(calendar)
    hdu.name = 'CALENDAR'
    hdus.append(hdu)

    # Save ephemerides table.
    hdu = fits.table_to_hdu(etable)
    hdu.name = 'EPHEM'
    hdus.append(hdu)

    # Save dynamic exposure-time factors.
    hdus.append(fits.ImageHDU(name='DYNAMIC', data=fexp))

    # Finalize the output file.
    hdus.writeto(output_name, overwrite=True)
    log.info('Plan initialization saved to {0}'.format(output_name))

if __name__ == '__main__':

    #stop = desisurvey.utils.get_date('2019-10-03')
    #stop = desisurvey.utils.get_date('2020-07-13')
    stop = None
    ephem = desisurvey.ephemerides.Ephemerides(stop_date=stop)
    initialize(ephem, stop_date=stop)
<|endoftext|>


<commit_before>#!/usr/bin/env python
'''
A simple test of the toxml module, which allows extended RSON
syntax to be used to create XML.

This module contains some extended RSON, and the corresponding XML.
Executing the module verifies they are the same.

Copyright (c) 2010, Patrick Maupin.  All rights reserved.

See http://code.google.com/p/rson/source/browse/trunk/license.txt

'''

from toxml import rson2xml

source = r'''

# Comments can be used, but are not kept when the XML is generated.

# In general, blank lines are meaningless, except when inside a
# triple-quoted or equal delimited string block.

# XML syntax requires a single tag at the top of the file.  We'll
# call it root:

root:

   # Attribute blocks are inside JSON-style dictionaries
   # or using the extended syntax that allows you to
   # make an empty dict, then fill it.
   {}
      attribute1: hi
      attribute2: there

   # Unquoted text is allowed, but cannot have any of \{},"
   # as the first character, or : or = anywhere inside.

   Some text inside the root element.

   [Unlike regular RSON, the [] characters are not special at the top level.]
   [I wouldn't use them inside attributes ({}), though without quoting.]

   NormalTextTag =
     If I use the equal sign, everything after that that
     is indented is treated as text.  I can use any
     special symbol such as " or = or : or {} here with
     no problems.

           Additional indentation is passed on to the output.

     This is probably the best, most natural, way to create
     elements that only have text (no attributes or sub-elements).

   othertag:
       If I don't use the equal sign, I can have sub-tags
       or attributes here as well as text.  But the text
       on any given line cannot start with a RSON special
       character, or contain the colon or equal anywhere
       on the line (unless I quote that line with "").

   Some more text inside the root element.

   Yetanothertag
       Note the colon after the tag is optional if you indent.

       AMoreIndentedTag:

           This text is inside AMoreIndentedTag

           # The equal sign starts a text block.  The text block
           # continues until we unindent to the level of the
           # equal sign (or even farther left).

           = Lots of text here, all together in a block.
             Can have special characters like \[]{},:"=
             if I want them.

             Blank lines embedded in the block are kept,
             but blank lines at the end are not.

           # Attribute for indented tag could be anywhere
           # at the same indentation level as the rest of
           # the data inside the tag.  Also, they do not
           # all need to be declared together.

           {att1:whatever}
           {att2: something else}
           {}
               att3: Lots of ways to provide attributes
               att4: if you feel you really need them
           {att5: "but really, do you?  That's the question"}

       """I can also use triple-quoted strings,
if I want weird characters inside my data.  With triple-quoted strings, the
only escaping that occurs is that \\""" has the \ removed from the front of it. """

       "And, of course, I can use JSON double-quoted strings for all sorts of character escapes."

       = XML special characters like &, <, and > are no problem, either.
         (I just put them inside the equal block quote because the comma is an
          RSON delimiter character.)  Here, I'll show them without any quoting whatsoever:

       & < >

   ATagWithOnlyAttributes:
       {HereIsAnAttr: and a value, AndAnother: value too}
       {EscapeQuotes: " I said \"Hi!\""}

   TinyTag: On a single line.
   FinalTag:  Indentation is optional if you use a colon.
'''

dest = r'''

<root attribute1="hi" attribute2="there">
    Some text inside the root element.
    [Unlike regular RSON, the [] characters are not special at the top level.]
    [I wouldn't use them inside attributes ({}), though without quoting.]
    <NormalTextTag>
        If I use the equal sign, everything after that that
        is indented is treated as text.  I can use any
        special symbol such as " or = or : or {} here with
        no problems.

              Additional indentation is passed on to the output.

        This is probably the best, most natural, way to create
        elements that only have text (no attributes or sub-elements).
    </NormalTextTag>
    <othertag>
        If I don't use the equal sign, I can have sub-tags
        or attributes here as well as text.  But the text
        on any given line cannot start with a RSON special
        character, or contain the colon or equal anywhere
        on the line (unless I quote that line with "").
    </othertag>
    Some more text inside the root element.
    <Yetanothertag>
        Note the colon after the tag is optional if you indent.
        <AMoreIndentedTag att1="whatever" att2="something else" att3="Lots of ways to provide attributes" att4="if you feel you really need them" att5="but really, do you?  That's the question">
            This text is inside AMoreIndentedTag
            Lots of text here, all together in a block.
            Can have special characters like \[]{},:"=
            if I want them.

            Blank lines embedded in the block are kept,
            but blank lines at the end are not.
        </AMoreIndentedTag>
        I can also use triple-quoted strings,
        if I want weird characters inside my data.  With triple-quoted strings, the
        only escaping that occurs is that \""" has the \ removed from the front of it.
        And, of course, I can use JSON double-quoted strings for all sorts of character escapes.
        XML special characters like &, <, and > are no problem, either.
        (I just put them inside the equal block quote because the comma is an
         RSON delimiter character.)  Here, I'll show them without any quoting whatsoever:
        & < >
    </Yetanothertag>
    <ATagWithOnlyAttributes HereIsAnAttr="and a value" AndAnother="value too" EscapeQuotes=' I said "Hi!"'/>
    <TinyTag>On a single line.</TinyTag>
    <FinalTag>
        Indentation is optional if you use a colon.
    </FinalTag>
</root>

'''

answer = str(rson2xml(source))

print answer

answer = [x.rstrip() for x in answer.strip().splitlines()]
dest = [x.rstrip() for x in dest.strip().splitlines()]

if 0:
    for x, y in zip(answer, dest):
        if x != y:
            print (x,y)

assert answer == dest
<commit_msg>Simplified XML special chars discussion<commit_after>#!/usr/bin/env python
'''
A simple test of the toxml module, which allows extended RSON
syntax to be used to create XML.

This module contains some extended RSON, and the corresponding XML.
Executing the module verifies they are the same.

Copyright (c) 2010, Patrick Maupin.  All rights reserved.

See http://code.google.com/p/rson/source/browse/trunk/license.txt

'''

from toxml import rson2xml

source = r'''

# Comments can be used, but are not kept when the XML is generated.

# In general, blank lines are meaningless, except when inside a
# triple-quoted or equal delimited string block.

# XML syntax requires a single tag at the top of the file.  We'll
# call it root:

root:

   # Attribute blocks are inside JSON-style dictionaries
   # or using the extended syntax that allows you to
   # make an empty dict, then fill it.
   {}
      attribute1: hi
      attribute2: there

   # Unquoted text is allowed, but cannot have any of \{},"
   # as the first character, or : or = anywhere inside.

   Some text inside the root element.

   [Unlike regular RSON, the [] characters are not special at the top level.]
   [I wouldn't use them inside attributes ({}), though without quoting.]

   NormalTextTag =
     If I use the equal sign, everything after that that
     is indented is treated as text.  I can use any
     special symbol such as " or = or : or {} here with
     no problems.

           Additional indentation is passed on to the output.

     This is probably the best, most natural, way to create
     elements that only have text (no attributes or sub-elements).

   othertag:
       If I don't use the equal sign, I can have sub-tags
       or attributes here as well as text.  But the text
       on any given line cannot start with a RSON special
       character, or contain the colon or equal anywhere
       on the line (unless I quote that line with "").

   Some more text inside the root element.

   Yetanothertag
       Note the colon after the tag is optional if you indent.

       AMoreIndentedTag:

           This text is inside AMoreIndentedTag

           # The equal sign starts a text block.  The text block
           # continues until we unindent to the level of the
           # equal sign (or even farther left).

           = Lots of text here, all together in a block.
             Can have special characters like \[]{},:"=
             if I want them.

             Blank lines embedded in the block are kept,
             but blank lines at the end are not.

           # Attribute for indented tag could be anywhere
           # at the same indentation level as the rest of
           # the data inside the tag.  Also, they do not
           # all need to be declared together.

           {att1:whatever}
           {att2: something else}
           {}
               att3: Lots of ways to provide attributes
               att4: if you feel you really need them
           {att5: "but really, do you?  That's the question"}

       """I can also use triple-quoted strings,
if I want weird characters inside my data.  With triple-quoted strings, the
only escaping that occurs is that \\""" has the \ removed from the front of it. """

       "And, of course, I can use JSON double-quoted strings for all sorts of character escapes."

       XML special characters like &, <, and > are no problem, either.

   ATagWithOnlyAttributes:
       {HereIsAnAttr: and a value, AndAnother: value too}
       {EscapeQuotes: " I said \"Hi!\""}

   TinyTag: On a single line.
   FinalTag:  Indentation is optional if you use a colon.
'''

dest = r'''

<root attribute1="hi" attribute2="there">
    Some text inside the root element.
    [Unlike regular RSON, the [] characters are not special at the top level.]
    [I wouldn't use them inside attributes ({}), though without quoting.]
    <NormalTextTag>
        If I use the equal sign, everything after that that
        is indented is treated as text.  I can use any
        special symbol such as " or = or : or {} here with
        no problems.

              Additional indentation is passed on to the output.

        This is probably the best, most natural, way to create
        elements that only have text (no attributes or sub-elements).
    </NormalTextTag>
    <othertag>
        If I don't use the equal sign, I can have sub-tags
        or attributes here as well as text.  But the text
        on any given line cannot start with a RSON special
        character, or contain the colon or equal anywhere
        on the line (unless I quote that line with "").
    </othertag>
    Some more text inside the root element.
    <Yetanothertag>
        Note the colon after the tag is optional if you indent.
        <AMoreIndentedTag att1="whatever" att2="something else" att3="Lots of ways to provide attributes" att4="if you feel you really need them" att5="but really, do you?  That's the question">
            This text is inside AMoreIndentedTag
            Lots of text here, all together in a block.
            Can have special characters like \[]{},:"=
            if I want them.

            Blank lines embedded in the block are kept,
            but blank lines at the end are not.
        </AMoreIndentedTag>
        I can also use triple-quoted strings,
        if I want weird characters inside my data.  With triple-quoted strings, the
        only escaping that occurs is that \""" has the \ removed from the front of it.
        And, of course, I can use JSON double-quoted strings for all sorts of character escapes.
        XML special characters like &, <, and > are no problem, either.
    </Yetanothertag>
    <ATagWithOnlyAttributes HereIsAnAttr="and a value" AndAnother="value too" EscapeQuotes=' I said "Hi!"'/>
    <TinyTag>On a single line.</TinyTag>
    <FinalTag>
        Indentation is optional if you use a colon.
    </FinalTag>
</root>

'''

answer = str(rson2xml(source))

print answer

answer = [x.rstrip() for x in answer.strip().splitlines()]
dest = [x.rstrip() for x in dest.strip().splitlines()]

if 0:
    for x, y in zip(answer, dest):
        if x != y:
            print (x,y)

assert answer == dest
<|endoftext|>