<commit_before>#!/usr/bin/env python
'''Search & Fetch records from NCBI's Sequence Read Archive.

Use http://www.ncbi.nlm.nih.gov/sra/advanced to build a query and the script
will output a CSV file with the results. Example:

./sra.py -s 'agalma[Organism]' -m 3 -o sra_output -e your@email.com

A more fine grained filtering based on the read length, library layout, number
of spots, platform, etc... needs to be done on a separate script importing this
program as a module. Check fetch_sra.py for an example on how to filter.
'''

import argparse
import os
import pandas as pd
import re

from Bio import Entrez

class SRADatabase:
    '''General information about SRA Database'''

    def __init__(self):
        einfo_handle = Entrez.einfo(db='sra')
        einfo = Entrez.read(einfo_handle, validate=False)

        # Define attributes.
        self.count = einfo['DbInfo']['Count']
        self.last_update = einfo['DbInfo']['LastUpdate']
        self.menu_name = einfo['DbInfo']['MenuName']
        self.description = einfo['DbInfo']['Description']
        self.link_list = einfo['DbInfo']['LinkList']
        self.field_list = einfo['DbInfo']['FieldList']
        self.einfo = einfo

        print('\n' + self.description + ' (%s)' % self.menu_name)
        print(self.count + ' entries, last updated in ' + self.last_update)
        print('Available search fields: ')
        for field in self.field_list:
            print('\t[%s] %s | %s (%s terms)' % (field['Name'], field['FullName'], field['Description'], field['TermCount']))

    # TODO Use def __unicode__ or __str__ to identify class objects.

class SRASearch:
    '''Perform search and keep IDs of SRA packages.

    Example of query:

        ((((("strategy rna seq"[Properties]) AND "platform illumina"[Properties])
        AND metazoa[Organism]) NOT vertebrata[Organism]) NOT insects[Organism]) AND
        ("2000/01/01"[Modification Date] : "3000"[Modification Date])
    '''

    def __init__(self, query, retmax, email):
        # Required arguments.
        # TODO Cache search based on query. Store list of IDs!
        self.query = query
        if int(retmax) > 100000:
            # Limit defined by Entrez.
            self.retmax = 100000
        else:
            self.retmax = retmax
        Entrez.email = email

        # Search metadata.
        self.count = None
        self.retstart = None
        self.query_translation = None
        self.idlist = None

        # Additional attributes.
        self.results = None
        self.database = SRADatabase()
        # TODO Add timestamp.

    def esearch(self):
        '''Search SRA packages with Entrez using query.'''
        handle = Entrez.esearch(db='sra', term=self.query, retmax=self.retmax)
        self.results = Entrez.read(handle)
        self.parse_results()
        print('\nSuccess! %s results found.' % self.count)
        print('Your query was: %s' % self.query_translation)
        print('Returned IDs (max=%s): %s' % (self.retmax, ', '.join(self.idlist)))
        return self.results

    def parse_results(self):
        '''Populate class attributes by parsing results.'''
        self.count = self.results['Count']
        self.retstart = self.results['RetStart']
        self.query_translation = self.results['QueryTranslation']
        self.idlist = self.results['IdList']
        print('\nFetched %d package IDs.' % len(self.idlist))

class SRAPackage:
    '''Fetch and store metadata from a SRA package.'''

    def __init__(self, sra_id):
        self.id = sra_id
        self.record = None
        self.cached_filepath = os.path.join('.cache', self.id)

        self.accession = None
        self.title = None
        self.study_title = None
        self.library_strategy = None
        self.library_layout = None
        self.instrument_model = None
        self.taxon_id = None
        self.scientific_name = None
        self.lineage = None
        self.run_accession = None
        self.nreads = None
        self.read_average = None
        self.total_spots = None
        self.total_bases = None
        self.size = None
        self.published = None

        # Define header section for CSV. Must match self.metadata.
        self.header = ['id', 'accession', 'title', 'lineage',
                       'taxon_id', 'scientific_name',
                       'library_strategy', 'library_layout',
                       'instrument_model', 'run_accession',
                       'nreads', 'read_average', 'total_spots',
                       'total_bases', 'size', 'published']

        # Do the actual metadata fetching.
        self.efetch()

        # Retrieve whole lineage by taxon ID.
        self.get_lineage()

        # Fill metadata set for later processing.
        self.metadata = (self.id, self.accession, self.title, self.lineage,
                         self.taxon_id, self.scientific_name,
                         self.library_strategy, self.library_layout,
                         self.instrument_model, self.run_accession,
                         self.nreads, self.read_average, self.total_spots,
                         self.total_bases, self.size, self.published,)

        print('Done!')

    def efetch(self):
        '''Fetch package metadata from Entrez'''
        print('\nProcessing ID=%s' % self.id)
        cached_file = self.cache()
        if cached_file:
            self.record = cached_file.read()
            cached_file.close()
        else:
            print('Record not in cache. Fetching...')
            handle = Entrez.efetch(db='sra', id=self.id)
            self.record = handle.read()
            # Write cache file.
            new_cache = open(self.cached_filepath, 'w')
            new_cache.write(self.record)
            new_cache.close()
        self.extract()

    def cache(self):
        '''Write and read cache files.'''
        # Make sure folder exists.
        cache_folder = '.cache'
        if not os.path.isdir(cache_folder):
            os.mkdir(cache_folder)

        # Try to get cache file.
        try:
            cached = open(self.cached_filepath)
            return cached
        except:
            return None

    def extract(self):
        '''Extract relevant fields from summary.'''

        # Fields with attributes.
        fields = {}

        # Fields to be parsed.
        regexes = {
            'accession': '<EXPERIMENT\s+.*?accession="(?P<accession>.*?)".*?>',
            'title': '<EXPERIMENT\s+.*?>.*?<TITLE>(?P<title>.*?)<\/TITLE>',
            'study_title': '<STUDY_TITLE>(?P<study_title>.*?)<\/STUDY_TITLE>',
            'library_strategy': '<LIBRARY_STRATEGY>(?P<library_strategy>.*?)<\/LIBRARY_STRATEGY>',
            'library_layout': '<LIBRARY_LAYOUT>\s*<(?P<library_layout>SINGLE|PAIRED)',
            'instrument_model': '<INSTRUMENT_MODEL>(?P<instrument_model>.*?)<\/INSTRUMENT_MODEL>',
            'taxon_id': '<TAXON_ID>(?P<taxon_id>.*?)<\/TAXON_ID>',
            'scientific_name': '<SCIENTIFIC_NAME>(?P<scientific_name>.*?)<\/SCIENTIFIC_NAME>',
            'run_accession': '<RUN\s+.*?accession="(?P<run_accession>.*?)"\s+.*?total_spots="(?P<total_spots>.*?)"\s+.*?total_bases="(?P<total_bases>.*?)"\s+.*?size="(?P<size>.*?)"\s+.*?published="(?P<published>.*?)"\s+.*?>',
            'nreads': '<Statistics\s+.*?nreads="(?P<nreads>.*?)"\s+.*?>',
            'read_average': '<Read\s+.*?average="(?P<read_average>.*?)"\s+.*?\/>',
        }

        # Iterate over regexes to parse attributes.
        # TODO handle multiple matches like "runs", "nreads", and "average"?
        # Right now it only gets the first run accession, nreads and
        # read_average. This is OK for now, since it is only a primary filter.
        for field, regex in regexes.iteritems():
            re_search = re.search(regex, self.record)
            if re_search:
                re_groups = re_search.groupdict()
                if re_groups:
                    for k, v in re_groups.iteritems():
                        fields[k] = v
                else:
                    if field in ['taxon_id', 'nreads', 'read_average',
                                 'total_spots', 'total_bases', 'size']:
                        fields[field] = 0
                    else:
                        fields[field] = ''
            else:
                if field in ['taxon_id', 'nreads', 'read_average',
                                'total_spots', 'total_bases', 'size']:
                    fields[field] = 0
                else:
                    fields[field] = ''

        self.accession = fields['accession']
        self.title = fields['title']
        if not self.title:
            self.title = fields['study_title']
        self.library_strategy = fields['library_strategy']
        self.library_layout = fields['library_layout']
        self.instrument_model = fields['instrument_model']
        self.taxon_id = int(fields['taxon_id'])
        self.scientific_name = fields['scientific_name']
        self.nreads = int(fields['nreads'])
        self.read_average = int(float(fields['read_average']))
        self.run_accession = fields['run_accession']
        if self.run_accession:
            self.total_spots = int(fields['total_spots'])
            self.total_bases = int(fields['total_bases'])
            self.size = int(fields['size'])
            self.published = fields['published']
        else:
            self.total_spots = 0
            self.total_bases = 0
            self.size = 0
            self.published = ''

    def get_lineage(self):
        '''Fetch hierarchy from NCBI's Taxonomy database.'''
        # Open taxa cache file.
        try:
            cached_taxa = pd.DataFrame.from_csv('.cache/taxa.csv')
        except:
            cached_taxa = pd.DataFrame(columns=['taxon_id', 'lineage', 'scientific_name'])

        # Fetch row with taxon_id.
        taxon_row = cached_taxa[cached_taxa.taxon_id == self.taxon_id]
        if not taxon_row.empty:
            self.scientific_name = taxon_row.scientific_name.values[0]
            self.lineage = taxon_row.lineage.values[0]
        else:
            print('Taxon %d not in cache. Fetching...' % self.taxon_id)
            handle = Entrez.efetch(db='taxonomy', id=str(self.taxon_id))
            taxon = Entrez.read(handle)
            self.scientific_name = taxon[0]['ScientificName']
            self.lineage = taxon[0]['Lineage'] + '; ' + self.scientific_name
            new_cache = cached_taxa.append([{'taxon_id': self.taxon_id,
                                             'lineage': self.lineage,
                                             'scientific_name': self.scientific_name}])
            new_cache.to_csv('.cache/taxa.csv')

class FilterPackages:
    '''Build data frame with package metadata for filtering.'''

    def __init__(self, packages, filter=None):
        self.packages = packages
        self.data_frame = None
        self.build_data_frame()
        self.filtered_data_frame = self.data_frame

    def build_data_frame(self):
        '''Get metadata from each package and save to data frame'''
        data = []
        index_ids = []
        header = []
        for package in self.packages:
            data.append(package.metadata)
            index_ids.append(package.metadata[0])
            header = package.header
        self.data_frame = pd.DataFrame(data, index=index_ids, columns=header)

    def write_csv(self, basename):
        '''Write CSV file from data frame.'''
        self.data_frame.to_csv(basename + '_unfiltered' + '.csv', index=False)
        self.filtered_data_frame.to_csv(basename + '.csv', index=False)
        print('\n%d of %d packages written to "%s" after filtering.\n' % (self.filtered_data_frame.index.size, self.data_frame.index.size, basename + '.csv'))

def main():
    '''Parse arguments and call SRA search.

    Main function simply parses arguments from command line input and assures
    everything is ok to instantiate the SRA search class.
    '''

    # Parse arguments.
    parser = argparse.ArgumentParser(description='Search & Fetch records from NCBI\'s Sequence Read Archive.',
                                     epilog='Work out those reads, dude.')
    parser.add_argument('-s', '--search',
                        help='put search terms between "quotes"',
                        type=str, required=True)
    parser.add_argument('-m', '--maximum',
                        help='maximum number of records to be retrieved',
                        default='20')
    parser.add_argument('-o', '--output',
                        help='indicate output CSV file',
                        required=True)
    parser.add_argument('-e', '--email',
                        help='an email address is required for Entrez',
                        required=True)
    args = parser.parse_args()

    # Instantiate search object.
    sra_search = SRASearch(query=args.search, retmax=args.maximum,
                           email=args.email)

    # Execute search itself.
    sra_search.esearch()

    # Fetch metadata from packages.
    packages = [SRAPackage(sra_id) for sra_id in sra_search.idlist]

    # Store packages in data frame for filtering.
    packages_to_filter = FilterPackages(packages)

    # Write CSV out.
    packages_to_filter.write_csv(args.output)

if __name__ == '__main__':
    main()
<commit_msg>Catch variable nreads, fix #1<commit_after>#!/usr/bin/env python
'''Search & Fetch records from NCBI's Sequence Read Archive.

Use http://www.ncbi.nlm.nih.gov/sra/advanced to build a query and the script
will output a CSV file with the results. Example:

./sra.py -s 'agalma[Organism]' -m 3 -o sra_output -e your@email.com

A more fine grained filtering based on the read length, library layout, number
of spots, platform, etc... needs to be done on a separate script importing this
program as a module. Check fetch_sra.py for an example on how to filter.
'''

import argparse
import os
import pandas as pd
import re

from Bio import Entrez

class SRADatabase:
    '''General information about SRA Database'''

    def __init__(self):
        einfo_handle = Entrez.einfo(db='sra')
        einfo = Entrez.read(einfo_handle, validate=False)

        # Define attributes.
        self.count = einfo['DbInfo']['Count']
        self.last_update = einfo['DbInfo']['LastUpdate']
        self.menu_name = einfo['DbInfo']['MenuName']
        self.description = einfo['DbInfo']['Description']
        self.link_list = einfo['DbInfo']['LinkList']
        self.field_list = einfo['DbInfo']['FieldList']
        self.einfo = einfo

        print('\n' + self.description + ' (%s)' % self.menu_name)
        print(self.count + ' entries, last updated in ' + self.last_update)
        print('Available search fields: ')
        for field in self.field_list:
            print('\t[%s] %s | %s (%s terms)' % (field['Name'], field['FullName'], field['Description'], field['TermCount']))

    # TODO Use def __unicode__ or __str__ to identify class objects.

class SRASearch:
    '''Perform search and keep IDs of SRA packages.

    Example of query:

        ((((("strategy rna seq"[Properties]) AND "platform illumina"[Properties])
        AND metazoa[Organism]) NOT vertebrata[Organism]) NOT insects[Organism]) AND
        ("2000/01/01"[Modification Date] : "3000"[Modification Date])
    '''

    def __init__(self, query, retmax, email):
        # Required arguments.
        # TODO Cache search based on query. Store list of IDs!
        self.query = query
        if int(retmax) > 100000:
            # Limit defined by Entrez.
            self.retmax = 100000
        else:
            self.retmax = retmax
        Entrez.email = email

        # Search metadata.
        self.count = None
        self.retstart = None
        self.query_translation = None
        self.idlist = None

        # Additional attributes.
        self.results = None
        self.database = SRADatabase()
        # TODO Add timestamp.

    def esearch(self):
        '''Search SRA packages with Entrez using query.'''
        handle = Entrez.esearch(db='sra', term=self.query, retmax=self.retmax)
        self.results = Entrez.read(handle)
        self.parse_results()
        print('\nSuccess! %s results found.' % self.count)
        print('Your query was: %s' % self.query_translation)
        print('Returned IDs (max=%s): %s' % (self.retmax, ', '.join(self.idlist)))
        return self.results

    def parse_results(self):
        '''Populate class attributes by parsing results.'''
        self.count = self.results['Count']
        self.retstart = self.results['RetStart']
        self.query_translation = self.results['QueryTranslation']
        self.idlist = self.results['IdList']
        print('\nFetched %d package IDs.' % len(self.idlist))

class SRAPackage:
    '''Fetch and store metadata from a SRA package.'''

    def __init__(self, sra_id):
        self.id = sra_id
        self.record = None
        self.cached_filepath = os.path.join('.cache', self.id)

        self.accession = None
        self.title = None
        self.study_title = None
        self.library_strategy = None
        self.library_layout = None
        self.instrument_model = None
        self.taxon_id = None
        self.scientific_name = None
        self.lineage = None
        self.run_accession = None
        self.nreads = None
        self.read_average = None
        self.total_spots = None
        self.total_bases = None
        self.size = None
        self.published = None

        # Define header section for CSV. Must match self.metadata.
        self.header = ['id', 'accession', 'title', 'lineage',
                       'taxon_id', 'scientific_name',
                       'library_strategy', 'library_layout',
                       'instrument_model', 'run_accession',
                       'nreads', 'read_average', 'total_spots',
                       'total_bases', 'size', 'published']

        # Do the actual metadata fetching.
        self.efetch()

        # Retrieve whole lineage by taxon ID.
        self.get_lineage()

        # Fill metadata set for later processing.
        self.metadata = (self.id, self.accession, self.title, self.lineage,
                         self.taxon_id, self.scientific_name,
                         self.library_strategy, self.library_layout,
                         self.instrument_model, self.run_accession,
                         self.nreads, self.read_average, self.total_spots,
                         self.total_bases, self.size, self.published,)

        print('Done!')

    def efetch(self):
        '''Fetch package metadata from Entrez'''
        print('\nProcessing ID=%s' % self.id)
        cached_file = self.cache()
        if cached_file:
            self.record = cached_file.read()
            cached_file.close()
        else:
            print('Record not in cache. Fetching...')
            handle = Entrez.efetch(db='sra', id=self.id)
            self.record = handle.read()
            # Write cache file.
            new_cache = open(self.cached_filepath, 'w')
            new_cache.write(self.record)
            new_cache.close()
        self.extract()

    def cache(self):
        '''Write and read cache files.'''
        # Make sure folder exists.
        cache_folder = '.cache'
        if not os.path.isdir(cache_folder):
            os.mkdir(cache_folder)

        # Try to get cache file.
        try:
            cached = open(self.cached_filepath)
            return cached
        except:
            return None

    def extract(self):
        '''Extract relevant fields from summary.'''

        # Fields with attributes.
        fields = {}

        # Fields to be parsed.
        regexes = {
            'accession': '<EXPERIMENT\s+.*?accession="(?P<accession>.*?)".*?>',
            'title': '<EXPERIMENT\s+.*?>.*?<TITLE>(?P<title>.*?)<\/TITLE>',
            'study_title': '<STUDY_TITLE>(?P<study_title>.*?)<\/STUDY_TITLE>',
            'library_strategy': '<LIBRARY_STRATEGY>(?P<library_strategy>.*?)<\/LIBRARY_STRATEGY>',
            'library_layout': '<LIBRARY_LAYOUT>\s*<(?P<library_layout>SINGLE|PAIRED)',
            'instrument_model': '<INSTRUMENT_MODEL>(?P<instrument_model>.*?)<\/INSTRUMENT_MODEL>',
            'taxon_id': '<TAXON_ID>(?P<taxon_id>.*?)<\/TAXON_ID>',
            'scientific_name': '<SCIENTIFIC_NAME>(?P<scientific_name>.*?)<\/SCIENTIFIC_NAME>',
            'run_accession': '<RUN\s+.*?accession="(?P<run_accession>.*?)"\s+.*?total_spots="(?P<total_spots>.*?)"\s+.*?total_bases="(?P<total_bases>.*?)"\s+.*?size="(?P<size>.*?)"\s+.*?published="(?P<published>.*?)"\s+.*?>',
            'nreads': '<Statistics\s+.*?nreads="(?P<nreads>.*?)"\s+.*?>',
            'read_average': '<Read\s+.*?average="(?P<read_average>.*?)"\s+.*?\/>',
        }

        # Iterate over regexes to parse attributes.
        # TODO handle multiple matches like "runs", "nreads", and "average"?
        # Right now it only gets the first run accession, nreads and
        # read_average. This is OK for now, since it is only a primary filter.
        for field, regex in regexes.iteritems():
            re_search = re.search(regex, self.record)
            if re_search:
                re_groups = re_search.groupdict()
                if re_groups:
                    for k, v in re_groups.iteritems():
                        fields[k] = v
                else:
                    if field in ['taxon_id', 'nreads', 'read_average',
                                 'total_spots', 'total_bases', 'size']:
                        fields[field] = 0
                    else:
                        fields[field] = ''
            else:
                if field in ['taxon_id', 'nreads', 'read_average',
                                'total_spots', 'total_bases', 'size']:
                    fields[field] = 0
                else:
                    fields[field] = ''

        self.accession = fields['accession']
        self.title = fields['title']
        if not self.title:
            self.title = fields['study_title']
        self.library_strategy = fields['library_strategy']
        self.library_layout = fields['library_layout']
        self.instrument_model = fields['instrument_model']
        self.taxon_id = int(fields['taxon_id'])
        self.scientific_name = fields['scientific_name']
        # Catch nreads="variable". See https://github.com/nelas/sra.py/issues/1
        try:
            self.nreads = int(fields['nreads'])
        except:
            self.nreads = 1
        self.read_average = int(float(fields['read_average']))
        self.run_accession = fields['run_accession']
        if self.run_accession:
            self.total_spots = int(fields['total_spots'])
            self.total_bases = int(fields['total_bases'])
            self.size = int(fields['size'])
            self.published = fields['published']
        else:
            self.total_spots = 0
            self.total_bases = 0
            self.size = 0
            self.published = ''

    def get_lineage(self):
        '''Fetch hierarchy from NCBI's Taxonomy database.'''
        # Open taxa cache file.
        try:
            cached_taxa = pd.DataFrame.from_csv('.cache/taxa.csv')
        except:
            cached_taxa = pd.DataFrame(columns=['taxon_id', 'lineage', 'scientific_name'])

        # Fetch row with taxon_id.
        taxon_row = cached_taxa[cached_taxa.taxon_id == self.taxon_id]
        if not taxon_row.empty:
            self.scientific_name = taxon_row.scientific_name.values[0]
            self.lineage = taxon_row.lineage.values[0]
        else:
            print('Taxon %d not in cache. Fetching...' % self.taxon_id)
            handle = Entrez.efetch(db='taxonomy', id=str(self.taxon_id))
            taxon = Entrez.read(handle)
            self.scientific_name = taxon[0]['ScientificName']
            self.lineage = taxon[0]['Lineage'] + '; ' + self.scientific_name
            new_cache = cached_taxa.append([{'taxon_id': self.taxon_id,
                                             'lineage': self.lineage,
                                             'scientific_name': self.scientific_name}])
            new_cache.to_csv('.cache/taxa.csv')

class FilterPackages:
    '''Build data frame with package metadata for filtering.'''

    def __init__(self, packages, filter=None):
        self.packages = packages
        self.data_frame = None
        self.build_data_frame()
        self.filtered_data_frame = self.data_frame

    def build_data_frame(self):
        '''Get metadata from each package and save to data frame'''
        data = []
        index_ids = []
        header = []
        for package in self.packages:
            data.append(package.metadata)
            index_ids.append(package.metadata[0])
            header = package.header
        self.data_frame = pd.DataFrame(data, index=index_ids, columns=header)

    def write_csv(self, basename):
        '''Write CSV file from data frame.'''
        self.data_frame.to_csv(basename + '_unfiltered' + '.csv', index=False)
        self.filtered_data_frame.to_csv(basename + '.csv', index=False)
        print('\n%d of %d packages written to "%s" after filtering.\n' % (self.filtered_data_frame.index.size, self.data_frame.index.size, basename + '.csv'))

def main():
    '''Parse arguments and call SRA search.

    Main function simply parses arguments from command line input and assures
    everything is ok to instantiate the SRA search class.
    '''

    # Parse arguments.
    parser = argparse.ArgumentParser(description='Search & Fetch records from NCBI\'s Sequence Read Archive.',
                                     epilog='Work out those reads, dude.')
    parser.add_argument('-s', '--search',
                        help='put search terms between "quotes"',
                        type=str, required=True)
    parser.add_argument('-m', '--maximum',
                        help='maximum number of records to be retrieved',
                        default='20')
    parser.add_argument('-o', '--output',
                        help='indicate output CSV file',
                        required=True)
    parser.add_argument('-e', '--email',
                        help='an email address is required for Entrez',
                        required=True)
    args = parser.parse_args()

    # Instantiate search object.
    sra_search = SRASearch(query=args.search, retmax=args.maximum,
                           email=args.email)

    # Execute search itself.
    sra_search.esearch()

    # Fetch metadata from packages.
    packages = [SRAPackage(sra_id) for sra_id in sra_search.idlist]

    # Store packages in data frame for filtering.
    packages_to_filter = FilterPackages(packages)

    # Write CSV out.
    packages_to_filter.write_csv(args.output)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>"""
This module handles initial database propagation, which is only run the first
time the game starts. It will create some default channels, objects, and
other things.

Everything starts at handle_setup()
"""

import django
from django.conf import settings
from django.utils.translation import ugettext as _
from evennia.players.models import PlayerDB
from evennia.server.models import ServerConfig
from evennia.utils import create

def create_config_values():
    """
    Creates the initial config values.
    """
    ServerConfig.objects.conf("site_name", settings.SERVERNAME)
    ServerConfig.objects.conf("idle_timeout", settings.IDLE_TIMEOUT)

def get_god_player():
    """
    Creates the god user and don't take no for an answer.
    """
    god_player = None
    while not god_player:
        try:
            god_player = PlayerDB.objects.get(id=1)
        except PlayerDB.DoesNotExist:
            print "\nCreate a superuser below. The superuser is Player #1, the 'owner' account of the server.\n"
            django.core.management.call_command("createsuperuser", interactive=True)
    return god_player

def create_objects():
    """
    Creates the #1 player and Limbo room.
    """

    print " Creating objects (Player #1 and Limbo room) ..."

    # Set the initial User's account object's username on the #1 object.
    # This object is pure django and only holds name, email and password.
    god_player = get_god_player()

    # Create a Player 'user profile' object to hold eventual
    # mud-specific settings for the PlayerDB object.
    player_typeclass = settings.BASE_PLAYER_TYPECLASS

    # run all creation hooks on god_player (we must do so manually
    # since the manage.py command does not)
    god_player.swap_typeclass(player_typeclass, clean_attributes=True)
    god_player.basetype_setup()
    god_player.at_player_creation()
    god_player.locks.add("examine:perm(Immortals);edit:false();delete:false();boot:false();msg:all()")
    # this is necessary for quelling to work correctly.
    god_player.permissions.add("Immortals")

    # Limbo is the default "nowhere" starting room

    # Create the in-game god-character for player #1 and set
    # it to exist in Limbo.
    character_typeclass = settings.BASE_CHARACTER_TYPECLASS
    god_character = create.create_object(character_typeclass,
                                         key=god_player.username,
                                         nohome=True)

    god_character.id = 1
    god_character.db.desc = _('This is User #1.')
    god_character.locks.add("examine:perm(Immortals);edit:false();delete:false();boot:false();msg:all();puppet:false()")
    god_character.permissions.add("Immortals")
    god_character.save()

    god_player.attributes.add("_first_login", True)
    god_player.attributes.add("_last_puppet", god_character)
    god_player.db._playable_characters.append(god_character)

    room_typeclass = settings.BASE_ROOM_TYPECLASS
    limbo_obj = create.create_object(room_typeclass, _('Limbo'), nohome=True)
    limbo_obj.id = 2
    string = \
        "Welcome to your new {wEvennia{n-based game. From here you are ready " \
        "to begin development. Visit http://evennia.com if you should need " \
        "help or would like to participate in community discussions. If you " \
        "are logged in as user #1 you can create a demo/tutorial area with " \
        "{w@batchcommand contrib.tutorial_world.build{n. Use {w@quell{n or login " \
        "as normal player to play the demo properly."
    string = _(string)
    limbo_obj.db.desc = string
    limbo_obj.save()

    # Now that Limbo exists, try to set the user up in Limbo (unless
    # the creation hooks already fixed this).
    if not god_character.location:
        god_character.location = limbo_obj
    if not god_character.home:
        god_character.home = limbo_obj

def create_channels():
    """
    Creates some sensible default channels.
    """
    print " Creating default channels ..."

    # public channel
    key1, aliases, desc, locks = settings.CHANNEL_PUBLIC
    pchan = create.create_channel(key1, aliases, desc, locks=locks)
    # mudinfo channel
    key2, aliases, desc, locks = settings.CHANNEL_MUDINFO
    ichan = create.create_channel(key2, aliases, desc, locks=locks)
    # connectinfo channel
    key3, aliases, desc, locks = settings.CHANNEL_CONNECTINFO
    cchan = create.create_channel(key3, aliases, desc, locks=locks)

    # TODO: postgresql-psycopg2 has a strange error when trying to
    # connect the user to the default channels. It works fine from inside
    # the game, but not from the initial startup. We are temporarily bypassing
    # the problem with the following fix. See Evennia Issue 151.
    if ((".".join(str(i) for i in django.VERSION) < "1.2"
                    and settings.DATABASE_ENGINE == "postgresql_psycopg2")
        or (hasattr(settings, 'DATABASES')
            and settings.DATABASES.get("default", {}).get('ENGINE', None)
            == 'django.db.backends.postgresql_psycopg2')):
        warning = """
        PostgreSQL-psycopg2 compatability fix:
        The in-game channels %s, %s and %s were created,
        but the superuser was not yet connected to them. Please use in
        game commands to onnect Player #1 to those channels when first
        logging in.
        """ % (key1, key2, key3)
        print warning
        return

    # connect the god user to all these channels by default.
    goduser = get_god_player()
    pchan.connect(goduser)
    ichan.connect(goduser)
    cchan.connect(goduser)

def create_system_scripts():
    """
    Setup the system repeat scripts. They are automatically started
    by the create_script function.
    """
    from evennia.scripts import scripts

    print " Creating and starting global scripts ..."

    # check so that all sessions are alive.
    script1 = create.create_script(scripts.CheckSessions)
    # validate all scripts in script table.
    script2 = create.create_script(scripts.ValidateScripts)
    # update the channel handler to make sure it's in sync
    script3 = create.create_script(scripts.ValidateChannelHandler)
    # flush the idmapper cache
    script4 = create.create_script(scripts.ValidateIdmapperCache)

    if not script1 or not script2 or not script3 or not script4:
        print " Error creating system scripts."

def start_game_time():
    """
    This starts a persistent script that keeps track of the
    in-game time (in whatever accelerated reference frame), but also
    the total run time of the server as well as its current uptime
    (the uptime can also be found directly from the server though).
    """
    print " Starting in-game time ..."
    from evennia.utils import gametime
    gametime.init_gametime()

def at_initial_setup():
    """
    Custom hook for users to overload some or all parts of the initial
    setup. Called very last in the sequence. It tries to import and
    srun a module settings.AT_INITIAL_SETUP_HOOK_MODULE and will fail
    silently if this does not exist or fails to load.
    """
    modname = settings.AT_INITIAL_SETUP_HOOK_MODULE
    if not modname:
        return
    try:
        mod = __import__(modname, fromlist=[None])
    except (ImportError, ValueError):
        return
    print " Running at_initial_setup() hook."
    if mod.__dict__.get("at_initial_setup", None):
        mod.at_initial_setup()

def reset_server():
    """
    We end the initialization by resetting the server. This
    makes sure the first login is the same as all the following
    ones, particularly it cleans all caches for the special objects.
    It also checks so the warm-reset mechanism works as it should.
    """
    from evennia.server.sessionhandler import SESSIONS
    print " Initial setup complete. Restarting Server once."
    SESSIONS.server.shutdown(mode='reset')

def handle_setup(last_step):
    """
    Main logic for the module. It allows for restarting
    the initialization at any point if one of the modules
    should crash.
    """

    if last_step < 0:
        # this means we don't need to handle setup since
        # it already ran sucessfully once.
        return
    elif last_step is None:
        # config doesn't exist yet. First start of server
        last_step = 0

    # setting up the list of functions to run
    setup_queue = [
        create_config_values,
        create_objects,
        create_channels,
        create_system_scripts,
        start_game_time,
        at_initial_setup,
        reset_server
        ]

    #print " Initial setup: %s steps." % (len(setup_queue))

    # step through queue, from last completed function
    for num, setup_func in enumerate(setup_queue[last_step:]):
        # run the setup function. Note that if there is a
        # traceback we let it stop the system so the config
        # step is not saved.
        #print "%s..." % num

        try:
            setup_func()
        except Exception:
            if last_step + num == 2:
                from evennia.objects.models import ObjectDB
                for obj in ObjectDB.objects.all():
                    obj.delete()
            elif last_step + num == 3:
                from evennia.comms.models import ChannelDB
                ChannelDB.objects.all().delete()
            raise
        ServerConfig.objects.conf("last_initial_setup_step", last_step + num + 1)
    # We got through the entire list. Set last_step to -1 so we don't
    # have to run this again.
    ServerConfig.objects.conf("last_initial_setup_step", -1)
<commit_msg>Fixes to initial_startup.<commit_after>"""
This module handles initial database propagation, which is only run the first
time the game starts. It will create some default channels, objects, and
other things.

Everything starts at handle_setup()
"""

import django
from django.conf import settings
from django.utils.translation import ugettext as _
from evennia.players.models import PlayerDB
from evennia.server.models import ServerConfig
from evennia.utils import create

def create_config_values():
    """
    Creates the initial config values.
    """
    ServerConfig.objects.conf("site_name", settings.SERVERNAME)
    ServerConfig.objects.conf("idle_timeout", settings.IDLE_TIMEOUT)

def get_god_player():
    """
    Creates the god user and don't take no for an answer.
    """
    try:
        god_player = PlayerDB.objects.get(id=1)
    except PlayerDB.DoesNotExist:
        txt = "\n\nNo superuser exists yet. The superuser is the 'owner'\n" \
              "account on the Evennia server. Create a new superuser using\n" \
              "the command\n\n" \
              " evennia createsuperuser\n\n" \
              "Follow the prompts, then restart the server."
        raise Exception(txt)
    return god_player

def create_objects():
    """
    Creates the #1 player and Limbo room.
    """

    print " Creating objects (Player #1 and Limbo room) ..."

    # Set the initial User's account object's username on the #1 object.
    # This object is pure django and only holds name, email and password.
    god_player = get_god_player()

    # Create a Player 'user profile' object to hold eventual
    # mud-specific settings for the PlayerDB object.
    player_typeclass = settings.BASE_PLAYER_TYPECLASS

    # run all creation hooks on god_player (we must do so manually
    # since the manage.py command does not)
    god_player.swap_typeclass(player_typeclass, clean_attributes=True)
    god_player.basetype_setup()
    god_player.at_player_creation()
    god_player.locks.add("examine:perm(Immortals);edit:false();delete:false();boot:false();msg:all()")
    # this is necessary for quelling to work correctly.
    god_player.permissions.add("Immortals")

    # Limbo is the default "nowhere" starting room

    # Create the in-game god-character for player #1 and set
    # it to exist in Limbo.
    character_typeclass = settings.BASE_CHARACTER_TYPECLASS
    god_character = create.create_object(character_typeclass,
                                         key=god_player.username,
                                         nohome=True)

    god_character.id = 1
    god_character.db.desc = _('This is User #1.')
    god_character.locks.add("examine:perm(Immortals);edit:false();delete:false();boot:false();msg:all();puppet:false()")
    god_character.permissions.add("Immortals")
    god_character.save()

    god_player.attributes.add("_first_login", True)
    god_player.attributes.add("_last_puppet", god_character)
    god_player.db._playable_characters.append(god_character)

    room_typeclass = settings.BASE_ROOM_TYPECLASS
    limbo_obj = create.create_object(room_typeclass, _('Limbo'), nohome=True)
    limbo_obj.id = 2
    string = \
        "Welcome to your new {wEvennia{n-based game. From here you are ready " \
        "to begin development. Visit http://evennia.com if you should need " \
        "help or would like to participate in community discussions. If you " \
        "are logged in as user #1 you can create a demo/tutorial area with " \
        "{w@batchcommand contrib.tutorial_world.build{n. Use {w@quell{n or login " \
        "as normal player to play the demo properly."
    string = _(string)
    limbo_obj.db.desc = string
    limbo_obj.save()

    # Now that Limbo exists, try to set the user up in Limbo (unless
    # the creation hooks already fixed this).
    if not god_character.location:
        god_character.location = limbo_obj
    if not god_character.home:
        god_character.home = limbo_obj

def create_channels():
    """
    Creates some sensible default channels.
    """
    print " Creating default channels ..."

    # public channel
    key1, aliases, desc, locks = settings.CHANNEL_PUBLIC
    pchan = create.create_channel(key1, aliases, desc, locks=locks)
    # mudinfo channel
    key2, aliases, desc, locks = settings.CHANNEL_MUDINFO
    ichan = create.create_channel(key2, aliases, desc, locks=locks)
    # connectinfo channel
    key3, aliases, desc, locks = settings.CHANNEL_CONNECTINFO
    cchan = create.create_channel(key3, aliases, desc, locks=locks)

    # TODO: postgresql-psycopg2 has a strange error when trying to
    # connect the user to the default channels. It works fine from inside
    # the game, but not from the initial startup. We are temporarily bypassing
    # the problem with the following fix. See Evennia Issue 151.
    if ((".".join(str(i) for i in django.VERSION) < "1.2"
                    and settings.DATABASE_ENGINE == "postgresql_psycopg2")
        or (hasattr(settings, 'DATABASES')
            and settings.DATABASES.get("default", {}).get('ENGINE', None)
            == 'django.db.backends.postgresql_psycopg2')):
        warning = """
        PostgreSQL-psycopg2 compatability fix:
        The in-game channels %s, %s and %s were created,
        but the superuser was not yet connected to them. Please use in
        game commands to onnect Player #1 to those channels when first
        logging in.
        """ % (key1, key2, key3)
        print warning
        return

    # connect the god user to all these channels by default.
    goduser = get_god_player()
    pchan.connect(goduser)
    ichan.connect(goduser)
    cchan.connect(goduser)

def create_system_scripts():
    """
    Setup the system repeat scripts. They are automatically started
    by the create_script function.
    """
    from evennia.scripts import scripts

    print " Creating and starting global scripts ..."

    # check so that all sessions are alive.
    script1 = create.create_script(scripts.CheckSessions)
    # validate all scripts in script table.
    script2 = create.create_script(scripts.ValidateScripts)
    # update the channel handler to make sure it's in sync
    script3 = create.create_script(scripts.ValidateChannelHandler)
    # flush the idmapper cache
    script4 = create.create_script(scripts.ValidateIdmapperCache)

    if not script1 or not script2 or not script3 or not script4:
        print " Error creating system scripts."

def start_game_time():
    """
    This starts a persistent script that keeps track of the
    in-game time (in whatever accelerated reference frame), but also
    the total run time of the server as well as its current uptime
    (the uptime can also be found directly from the server though).
    """
    print " Starting in-game time ..."
    from evennia.utils import gametime
    gametime.init_gametime()

def at_initial_setup():
    """
    Custom hook for users to overload some or all parts of the initial
    setup. Called very last in the sequence. It tries to import and
    srun a module settings.AT_INITIAL_SETUP_HOOK_MODULE and will fail
    silently if this does not exist or fails to load.
    """
    modname = settings.AT_INITIAL_SETUP_HOOK_MODULE
    if not modname:
        return
    try:
        mod = __import__(modname, fromlist=[None])
    except (ImportError, ValueError):
        return
    print " Running at_initial_setup() hook."
    if mod.__dict__.get("at_initial_setup", None):
        mod.at_initial_setup()

def reset_server():
    """
    We end the initialization by resetting the server. This
    makes sure the first login is the same as all the following
    ones, particularly it cleans all caches for the special objects.
    It also checks so the warm-reset mechanism works as it should.
    """
    from evennia.server.sessionhandler import SESSIONS
    print " Initial setup complete. Restarting Server once."
    SESSIONS.server.shutdown(mode='reset')

def handle_setup(last_step):
    """
    Main logic for the module. It allows for restarting
    the initialization at any point if one of the modules
    should crash.
    """

    if last_step < 0:
        # this means we don't need to handle setup since
        # it already ran sucessfully once.
        return
    elif last_step is None:
        # config doesn't exist yet. First start of server
        last_step = 0

    # setting up the list of functions to run
    setup_queue = [
        create_config_values,
        create_objects,
        create_channels,
        create_system_scripts,
        start_game_time,
        at_initial_setup,
        reset_server
        ]

    #print " Initial setup: %s steps." % (len(setup_queue))

    # step through queue, from last completed function
    for num, setup_func in enumerate(setup_queue[last_step:]):
        # run the setup function. Note that if there is a
        # traceback we let it stop the system so the config
        # step is not saved.
        #print "%s..." % num

        try:
            setup_func()
        except Exception:
            if last_step + num == 2:
                from evennia.objects.models import ObjectDB
                for obj in ObjectDB.objects.all():
                    obj.delete()
            elif last_step + num == 3:
                from evennia.comms.models import ChannelDB
                ChannelDB.objects.all().delete()
            raise
        ServerConfig.objects.conf("last_initial_setup_step", last_step + num + 1)
    # We got through the entire list. Set last_step to -1 so we don't
    # have to run this again.
    ServerConfig.objects.conf("last_initial_setup_step", -1)
<|endoftext|>


<commit_before>import sys
import numpy
import gzip
import math
import itertools
import time
import codecs
import argparse
from index_data import DataProcessor
from onto_attention import OntoAttentionLSTM
from keras.models import Model
from keras.layers import Input, Dropout, LSTM, Dense
from keras.layers.wrappers import TimeDistributed
from keras.callbacks import EarlyStopping
from keras_extensions import HigherOrderEmbedding

class SentenceModel(object):
  def __init__(self, word_dim=50, num_senses=2, num_hyps=5):
    self.dp = DataProcessor(word_syn_cutoff=num_senses, syn_path_cutoff=num_hyps)
    self.num_hyps = num_hyps
    self.num_senses = num_senses
    self.numpy_rng = numpy.random.RandomState(12345)
    self.word_dim = word_dim
    self.model = None

  def read_sentences(self, tagged_sentences, sentlenlimit=None):
    num_sentences = len(tagged_sentences)
    all_words = []
    all_pos_tags = []
    maxsentlen = 0
    for tagged_sentence in tagged_sentences:
      words = []
      pos_tags = []
      # Expects each token to be a "_" separated combination of word and POS tag.
      tagged_words = tagged_sentence.split(" ")
      if sentlenlimit is not None:
        tagged_words = tagged_words[:sentlenlimit]
      for word_tag in tagged_words:
        parts = word_tag.split("_")
        tag = parts[-1]
        word = "_".join(parts[:-1]).lower()
        words.append(word)
        pos_tags.append(tag)
      if len(words) > maxsentlen:
        maxsentlen = len(words)
      all_words.append(words)
      all_pos_tags.append(pos_tags)
    if not sentlenlimit:
      sentlenlimit = maxsentlen
    C_ind = numpy.zeros((num_sentences, sentlenlimit, self.num_senses, self.num_hyps), dtype='int32')
    S_ind = numpy.zeros((num_sentences, sentlenlimit), dtype='int32')
    for i, (words, pos_tags) in enumerate(zip(all_words, all_pos_tags)):
      sentlen = len(words)
      word_inds, syn_inds = self.dp.index_sentence(words, pos_tags)
      S_ind[i][-sentlen:] = word_inds
      for j in range(sentlen):
        sense_syn_ind = syn_inds[j]
        sense_syn_ind_len = len(sense_syn_ind)
        for k, syn_ind in enumerate(sense_syn_ind):
          C_ind[i][-sentlen+j][-sense_syn_ind_len+k][-len(syn_ind):] = syn_ind
    return S_ind, C_ind

  def factor_target_indices(self, Y_inds, vocab_size=None):
    if vocab_size is None:
      vocab_size = Y_inds.max() + 1
    factor_size = int(math.ceil(math.sqrt(vocab_size)))
    Y_inds_1 = numpy.asarray(Y_inds/factor_size, dtype='int32')
    Y_inds_2 = numpy.asarray(Y_inds%factor_size, dtype='int32')
    # Making one-hot vectors out of Y_inds_1 and Y_inds_2
    Y_1 = numpy.zeros((Y_inds_1.shape + (factor_size,)))
    for inds in itertools.product(*[numpy.arange(s) for s in Y_inds_1.shape]):
      Y_1[inds+(Y_inds_1[inds],)] = 1
    Y_2 = numpy.zeros((Y_inds_2.shape + (factor_size,)))
    for inds in itertools.product(*[numpy.arange(s) for s in Y_inds_2.shape]):
      Y_2[inds+(Y_inds_2[inds],)] = 1
    return Y_1, Y_2

  def train(self, S_ind, C_ind, use_onto_lstm=True, use_attention=True, num_epochs=20, S_ind_test=None, C_ind_test=None):
    # Predict next word from current synsets
    X = C_ind[:,:-1] if use_onto_lstm else S_ind[:,:-1] # remove the last words' hyps in all sentences
    Y_inds = S_ind[:,1:] # remove the first words in all sentences
    do_test = False
    if S_ind_test is not None and C_ind_test is not None:
      do_test = True
      X_test = C_ind_test[:,:-1] if use_onto_lstm else S_ind_test[:,:-1] # remove the last words' hyps in all sentences
      Y_inds_test = S_ind_test[:,1:]
      vocab_size = max(Y_inds.max(), Y_inds_test.max()) + 1
      Y_1_test, Y_2_test = self.factor_target_indices(Y_inds_test, vocab_size)
      Y_1, Y_2 = self.factor_target_indices(Y_inds, vocab_size)
    else:
      Y_1, Y_2 = self.factor_target_indices(Y_inds)
    length = Y_inds.shape[1]
    lstm_outdim = self.word_dim

    num_words = len(self.dp.word_index)
    num_syns = len(self.dp.synset_index)
    input = Input(shape=X.shape[1:], dtype='int32')
    embed_input_dim = num_syns if use_onto_lstm else num_words
    sent_rep = HigherOrderEmbedding(name='embedding', input_dim=embed_input_dim, output_dim=self.word_dim, input_shape=X.shape[1:])(input)
    reg_sent_rep = Dropout(0.5)(sent_rep)
    if use_onto_lstm:
      lstm_out = OntoAttentionLSTM(name='sent_lstm', input_dim=self.word_dim, output_dim=lstm_outdim, input_length=length, num_senses=self.num_senses, num_hyps=self.num_hyps, return_sequences=True, use_attention=True)(reg_sent_rep)
    else:
      lstm_out = LSTM(name='sent_lstm', input_dim=self.word_dim, output_dim=lstm_outdim, input_length=length, return_sequences=True)(reg_sent_rep)
    softmax_1 = TimeDistributed(Dense(input_dim=lstm_outdim, output_dim=Y_1.shape[-1], activation='softmax'))(lstm_out)
    softmax_2 = TimeDistributed(Dense(input_dim=lstm_outdim, output_dim=Y_2.shape[-1], activation='softmax'))(lstm_out)

    model = Model(input=input, output=[softmax_1, softmax_2])
    print >>sys.stderr, model.summary()
    early_stopping = EarlyStopping()
    precompile_time = time.time()
    model.compile(loss='categorical_crossentropy', optimizer='adam')
    postcompile_time = time.time()
    print >>sys.stderr, "Model compilation took %d s"%(postcompile_time - precompile_time)
    model.fit(X, [Y_1, Y_2], nb_epoch=num_epochs, validation_split=0.1, callbacks=[early_stopping])
    posttrain_time = time.time()
    print >>sys.stderr, "Training took %d s"%(posttrain_time - postcompile_time)
    if do_test:
      print >>sys.stderr, "Evaluating model on test data"
      test_loss = model.evaluate(X_test, [Y_1_test, Y_2_test])
      print >>sys.stderr, "Test loss: %.4f"%test_loss
    concept_reps = model.layers[1].get_weights()
    self.model = model
    return concept_reps

  def get_attention(self, C_ind):
    if not self.model:
      raise RuntimeError, "Model not trained!"
    model_embedding = None
    model_weights = None
    for layer in self.model.layers:
      if layer.name.lower() == "embedding":
        model_embedding = layer
      if layer.name.lower() == "sent_lstm":
        model_lstm = layer
    if model_embedding is None or model_lstm is None:
      raise RuntimeError, "Did not find expected layers"
    lstm_weights = model_lstm.get_weights()
    embedding_weights = model_embedding.get_weights()
    embed_in_dim, embed_out_dim = embedding_weights[0].shape
    att_embedding = HigherOrderEmbedding(input_dim=embed_in_dim, output_dim=embed_out_dim, weights=embedding_weights)
    onto_lstm = OntoAttentionLSTM(input_dim=embed_out_dim, output_dim=embed_out_dim, input_length=model_lstm.input_length, num_senses=self.num_senses, num_hyps=self.num_hyps, use_attention=True, return_attention=True, weights=lstm_weights)
    att_input = Input(shape=C_ind.shape[1:], dtype='int32')
    att_sent_rep = att_embedding(att_input)
    att_output = onto_lstm(att_sent_rep)
    att_model = Model(input=att_input, output=att_output)
    att_model.compile(optimizer='adam', loss='mse') # optimizer and loss are not needed since we are not going to train this model.
    C_att = att_model.predict(C_ind)
    print >>sys.stderr, "Got attention values. Input, output shapes:", C_ind.shape, C_att.shape
    return C_att

if __name__ == '__main__':
  argparser = argparse.ArgumentParser(description="Model sentences using ontoLSTM")
  argparser.add_argument('train_file', metavar='TRAIN-FILE', type=str, help="One sentence per line, POS tagged")
  argparser.add_argument('--test_file', type=str, help="Test file for which attention values will be printed. One sentence per line, POS tagged")
  argparser.add_argument('--dim', type=int, help="Word/synset dimensionality", default=50)
  argparser.add_argument('--num_senses', type=int, help="Number of senses per word if using OntoLSTM (default 2)", default=2)
  argparser.add_argument('--num_hyps', type=int, help="Number of hypernyms per sense if using OntoLSTM (default 5)", default=5)
  argparser.add_argument('--use_onto_lstm', help="If this flag is not set, will use traditional LSTM", action='store_true')
  argparser.add_argument('--use_attention', help="Use attention in ontoLSTM. If this flag is not set, will use average concept representations", action='store_true')
  argparser.add_argument('--show_attention', help="Print attention values for sentences in the test file (printed in <testfilename>.att_out)", action='store_true')
  argparser.add_argument('--synset_embedding_output', type=str, help="Print learned synset representations in the given file")
  argparser.add_argument('--num_epochs', type=int, help="Number of epochs (default 20)", default=20)
  args = argparser.parse_args()
  sm = SentenceModel(word_dim=args.dim, num_senses=args.num_senses, num_hyps=args.num_hyps)
  ts = [x.strip() for x in codecs.open(args.train_file, "r", "utf-8").readlines()]
  if args.show_attention and (not args.use_attention or not args.use_onto_lstm):
    raise RuntimeError, "Use OntoLSTM with attention to print attention values of the test file"
  print >>sys.stderr, "Reading training data"
  S_ind, C_ind = sm.read_sentences(ts)
  _, train_sent_len, _, _ = C_ind.shape
  if args.test_file is not None:
    print >>sys.stderr, "Reading test data"
    ts_test = [x.strip() for x in open(args.test_file).readlines()]
    S_ind_test, C_ind_test = sm.read_sentences(ts_test, sentlenlimit=train_sent_len)
    concept_reps = sm.train(S_ind, C_ind, use_onto_lstm=args.use_onto_lstm, use_attention=args.use_attention, num_epochs=args.num_epochs, S_ind_test=S_ind_test, C_ind_test=C_ind_test)
  else:
    concept_reps = sm.train(S_ind, C_ind, use_onto_lstm=args.use_onto_lstm, use_attention=args.use_attention, num_epochs=args.num_epochs)
  if args.synset_embedding_output is not None:
    concrepfile = open(args.synset_embedding_output, "w")
    for syn in sm.dp.synset_index:
      print >>concrepfile, syn, " ".join(["%s"%x for x in concept_reps[0][sm.dp.synset_index[syn]]])
  rev_synset_ind = {ind: syn for (syn, ind) in sm.dp.synset_index.items()}
  if args.show_attention:
    if args.test_file is None:
      raise RuntimeError, "Test file not provided!"
    outfile_name = args.test_file.split("/")[-1] + ".att_out"
    outfile = open(outfile_name, "w")
    C_att = sm.get_attention(C_ind_test)
    for i, (sent, sent_c_inds, sent_c_atts) in enumerate(zip(ts_test, C_ind_test, C_att)):
      print >>outfile, "SENT %d: %s"%(i, sent)
      words = sent.split()[:train_sent_len]
      word_id = 0
      for word_c_inds, word_c_atts in zip(sent_c_inds, sent_c_atts):
        if word_c_inds.sum() == 0:
          continue
        sense_id = 0
        #print >>outfile, "Attention for %s"%(words[word_id])
        best_sense = ""
        max_sense_prob = 0.0
        for s_h_ind, s_h_att in zip(word_c_inds, word_c_atts):
          if sum(s_h_ind) == 0:
            continue
          #print >>outfile, "\nSense %d"%(sense_id)
          sense_id += 1
          if s_h_att[-1] > max_sense_prob:
            max_sense_prob = s_h_att[-1]
            best_sense = rev_synset_ind[s_h_ind[-1]]
          #for h_ind, h_att in zip(s_h_ind, s_h_att):
          #  print >>outfile, rev_synset_ind[h_ind], h_att
        print >>outfile, words[word_id], best_sense
        word_id += 1
        #print >>outfile
      print >>outfile
<commit_msg>Unicode problems. Sigh!<commit_after>import sys
import numpy
import gzip
import math
import itertools
import time
import codecs
import argparse
from index_data import DataProcessor
from onto_attention import OntoAttentionLSTM
from keras.models import Model
from keras.layers import Input, Dropout, LSTM, Dense
from keras.layers.wrappers import TimeDistributed
from keras.callbacks import EarlyStopping
from keras_extensions import HigherOrderEmbedding

class SentenceModel(object):
  def __init__(self, word_dim=50, num_senses=2, num_hyps=5):
    self.dp = DataProcessor(word_syn_cutoff=num_senses, syn_path_cutoff=num_hyps)
    self.num_hyps = num_hyps
    self.num_senses = num_senses
    self.numpy_rng = numpy.random.RandomState(12345)
    self.word_dim = word_dim
    self.model = None

  def read_sentences(self, tagged_sentences, sentlenlimit=None):
    num_sentences = len(tagged_sentences)
    all_words = []
    all_pos_tags = []
    maxsentlen = 0
    for tagged_sentence in tagged_sentences:
      words = []
      pos_tags = []
      # Expects each token to be a "_" separated combination of word and POS tag.
      tagged_words = tagged_sentence.split(" ")
      if sentlenlimit is not None:
        tagged_words = tagged_words[:sentlenlimit]
      for word_tag in tagged_words:
        parts = word_tag.split("_")
        tag = parts[-1]
        word = "_".join(parts[:-1]).lower()
        words.append(word)
        pos_tags.append(tag)
      if len(words) > maxsentlen:
        maxsentlen = len(words)
      all_words.append(words)
      all_pos_tags.append(pos_tags)
    if not sentlenlimit:
      sentlenlimit = maxsentlen
    C_ind = numpy.zeros((num_sentences, sentlenlimit, self.num_senses, self.num_hyps), dtype='int32')
    S_ind = numpy.zeros((num_sentences, sentlenlimit), dtype='int32')
    for i, (words, pos_tags) in enumerate(zip(all_words, all_pos_tags)):
      sentlen = len(words)
      word_inds, syn_inds = self.dp.index_sentence(words, pos_tags)
      S_ind[i][-sentlen:] = word_inds
      for j in range(sentlen):
        sense_syn_ind = syn_inds[j]
        sense_syn_ind_len = len(sense_syn_ind)
        for k, syn_ind in enumerate(sense_syn_ind):
          C_ind[i][-sentlen+j][-sense_syn_ind_len+k][-len(syn_ind):] = syn_ind
    return S_ind, C_ind

  def factor_target_indices(self, Y_inds, vocab_size=None):
    if vocab_size is None:
      vocab_size = Y_inds.max() + 1
    factor_size = int(math.ceil(math.sqrt(vocab_size)))
    Y_inds_1 = numpy.asarray(Y_inds/factor_size, dtype='int32')
    Y_inds_2 = numpy.asarray(Y_inds%factor_size, dtype='int32')
    # Making one-hot vectors out of Y_inds_1 and Y_inds_2
    Y_1 = numpy.zeros((Y_inds_1.shape + (factor_size,)))
    for inds in itertools.product(*[numpy.arange(s) for s in Y_inds_1.shape]):
      Y_1[inds+(Y_inds_1[inds],)] = 1
    Y_2 = numpy.zeros((Y_inds_2.shape + (factor_size,)))
    for inds in itertools.product(*[numpy.arange(s) for s in Y_inds_2.shape]):
      Y_2[inds+(Y_inds_2[inds],)] = 1
    return Y_1, Y_2

  def train(self, S_ind, C_ind, use_onto_lstm=True, use_attention=True, num_epochs=20, S_ind_test=None, C_ind_test=None):
    # Predict next word from current synsets
    X = C_ind[:,:-1] if use_onto_lstm else S_ind[:,:-1] # remove the last words' hyps in all sentences
    Y_inds = S_ind[:,1:] # remove the first words in all sentences
    do_test = False
    if S_ind_test is not None and C_ind_test is not None:
      do_test = True
      X_test = C_ind_test[:,:-1] if use_onto_lstm else S_ind_test[:,:-1] # remove the last words' hyps in all sentences
      Y_inds_test = S_ind_test[:,1:]
      vocab_size = max(Y_inds.max(), Y_inds_test.max()) + 1
      Y_1_test, Y_2_test = self.factor_target_indices(Y_inds_test, vocab_size)
      Y_1, Y_2 = self.factor_target_indices(Y_inds, vocab_size)
    else:
      Y_1, Y_2 = self.factor_target_indices(Y_inds)
    length = Y_inds.shape[1]
    lstm_outdim = self.word_dim

    num_words = len(self.dp.word_index)
    num_syns = len(self.dp.synset_index)
    input = Input(shape=X.shape[1:], dtype='int32')
    embed_input_dim = num_syns if use_onto_lstm else num_words
    sent_rep = HigherOrderEmbedding(name='embedding', input_dim=embed_input_dim, output_dim=self.word_dim, input_shape=X.shape[1:])(input)
    reg_sent_rep = Dropout(0.5)(sent_rep)
    if use_onto_lstm:
      lstm_out = OntoAttentionLSTM(name='sent_lstm', input_dim=self.word_dim, output_dim=lstm_outdim, input_length=length, num_senses=self.num_senses, num_hyps=self.num_hyps, return_sequences=True, use_attention=True)(reg_sent_rep)
    else:
      lstm_out = LSTM(name='sent_lstm', input_dim=self.word_dim, output_dim=lstm_outdim, input_length=length, return_sequences=True)(reg_sent_rep)
    softmax_1 = TimeDistributed(Dense(input_dim=lstm_outdim, output_dim=Y_1.shape[-1], activation='softmax'))(lstm_out)
    softmax_2 = TimeDistributed(Dense(input_dim=lstm_outdim, output_dim=Y_2.shape[-1], activation='softmax'))(lstm_out)

    model = Model(input=input, output=[softmax_1, softmax_2])
    print >>sys.stderr, model.summary()
    early_stopping = EarlyStopping()
    precompile_time = time.time()
    model.compile(loss='categorical_crossentropy', optimizer='adam')
    postcompile_time = time.time()
    print >>sys.stderr, "Model compilation took %d s"%(postcompile_time - precompile_time)
    model.fit(X, [Y_1, Y_2], nb_epoch=num_epochs, validation_split=0.1, callbacks=[early_stopping])
    posttrain_time = time.time()
    print >>sys.stderr, "Training took %d s"%(posttrain_time - postcompile_time)
    if do_test:
      print >>sys.stderr, "Evaluating model on test data"
      test_loss = model.evaluate(X_test, [Y_1_test, Y_2_test])
      print >>sys.stderr, "Test loss: %.4f"%test_loss
    concept_reps = model.layers[1].get_weights()
    self.model = model
    return concept_reps

  def get_attention(self, C_ind):
    if not self.model:
      raise RuntimeError, "Model not trained!"
    model_embedding = None
    model_weights = None
    for layer in self.model.layers:
      if layer.name.lower() == "embedding":
        model_embedding = layer
      if layer.name.lower() == "sent_lstm":
        model_lstm = layer
    if model_embedding is None or model_lstm is None:
      raise RuntimeError, "Did not find expected layers"
    lstm_weights = model_lstm.get_weights()
    embedding_weights = model_embedding.get_weights()
    embed_in_dim, embed_out_dim = embedding_weights[0].shape
    att_embedding = HigherOrderEmbedding(input_dim=embed_in_dim, output_dim=embed_out_dim, weights=embedding_weights)
    onto_lstm = OntoAttentionLSTM(input_dim=embed_out_dim, output_dim=embed_out_dim, input_length=model_lstm.input_length, num_senses=self.num_senses, num_hyps=self.num_hyps, use_attention=True, return_attention=True, weights=lstm_weights)
    att_input = Input(shape=C_ind.shape[1:], dtype='int32')
    att_sent_rep = att_embedding(att_input)
    att_output = onto_lstm(att_sent_rep)
    att_model = Model(input=att_input, output=att_output)
    att_model.compile(optimizer='adam', loss='mse') # optimizer and loss are not needed since we are not going to train this model.
    C_att = att_model.predict(C_ind)
    print >>sys.stderr, "Got attention values. Input, output shapes:", C_ind.shape, C_att.shape
    return C_att

if __name__ == '__main__':
  argparser = argparse.ArgumentParser(description="Model sentences using ontoLSTM")
  argparser.add_argument('train_file', metavar='TRAIN-FILE', type=str, help="One sentence per line, POS tagged")
  argparser.add_argument('--test_file', type=str, help="Test file for which attention values will be printed. One sentence per line, POS tagged")
  argparser.add_argument('--dim', type=int, help="Word/synset dimensionality", default=50)
  argparser.add_argument('--num_senses', type=int, help="Number of senses per word if using OntoLSTM (default 2)", default=2)
  argparser.add_argument('--num_hyps', type=int, help="Number of hypernyms per sense if using OntoLSTM (default 5)", default=5)
  argparser.add_argument('--use_onto_lstm', help="If this flag is not set, will use traditional LSTM", action='store_true')
  argparser.add_argument('--use_attention', help="Use attention in ontoLSTM. If this flag is not set, will use average concept representations", action='store_true')
  argparser.add_argument('--show_attention', help="Print attention values for sentences in the test file (printed in <testfilename>.att_out)", action='store_true')
  argparser.add_argument('--synset_embedding_output', type=str, help="Print learned synset representations in the given file")
  argparser.add_argument('--num_epochs', type=int, help="Number of epochs (default 20)", default=20)
  args = argparser.parse_args()
  sm = SentenceModel(word_dim=args.dim, num_senses=args.num_senses, num_hyps=args.num_hyps)
  ts = [x.strip() for x in codecs.open(args.train_file, "r", "utf-8").readlines()]
  if args.show_attention and (not args.use_attention or not args.use_onto_lstm):
    raise RuntimeError, "Use OntoLSTM with attention to print attention values of the test file"
  print >>sys.stderr, "Reading training data"
  S_ind, C_ind = sm.read_sentences(ts)
  _, train_sent_len, _, _ = C_ind.shape
  if args.test_file is not None:
    print >>sys.stderr, "Reading test data"
    ts_test = [x.strip() for x in codecs.open(args.test_file, "r", "utf-8").readlines()]
    S_ind_test, C_ind_test = sm.read_sentences(ts_test, sentlenlimit=train_sent_len)
    concept_reps = sm.train(S_ind, C_ind, use_onto_lstm=args.use_onto_lstm, use_attention=args.use_attention, num_epochs=args.num_epochs, S_ind_test=S_ind_test, C_ind_test=C_ind_test)
  else:
    concept_reps = sm.train(S_ind, C_ind, use_onto_lstm=args.use_onto_lstm, use_attention=args.use_attention, num_epochs=args.num_epochs)
  if args.synset_embedding_output is not None:
    concrepfile = open(args.synset_embedding_output, "w")
    for syn in sm.dp.synset_index:
      print >>concrepfile, syn, " ".join(["%s"%x for x in concept_reps[0][sm.dp.synset_index[syn]]])
  rev_synset_ind = {ind: syn for (syn, ind) in sm.dp.synset_index.items()}
  if args.show_attention:
    if args.test_file is None:
      raise RuntimeError, "Test file not provided!"
    outfile_name = args.test_file.split("/")[-1] + ".att_out"
    outfile = open(outfile_name, "w")
    C_att = sm.get_attention(C_ind_test)
    for i, (sent, sent_c_inds, sent_c_atts) in enumerate(zip(ts_test, C_ind_test, C_att)):
      print >>outfile, "SENT %d: %s"%(i, sent)
      words = sent.split()[:train_sent_len]
      word_id = 0
      for word_c_inds, word_c_atts in zip(sent_c_inds, sent_c_atts):
        if word_c_inds.sum() == 0:
          continue
        sense_id = 0
        #print >>outfile, "Attention for %s"%(words[word_id])
        best_sense = ""
        max_sense_prob = 0.0
        for s_h_ind, s_h_att in zip(word_c_inds, word_c_atts):
          if sum(s_h_ind) == 0:
            continue
          #print >>outfile, "\nSense %d"%(sense_id)
          sense_id += 1
          if s_h_att[-1] > max_sense_prob:
            max_sense_prob = s_h_att[-1]
            best_sense = rev_synset_ind[s_h_ind[-1]]
          #for h_ind, h_att in zip(s_h_ind, s_h_att):
          #  print >>outfile, rev_synset_ind[h_ind], h_att
        print >>outfile, words[word_id], best_sense
        word_id += 1
        #print >>outfile
      print >>outfile
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# Copyright (c) 2014-2020 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test logic for skipping signature validation on old blocks.

Test logic for skipping signature validation on blocks which we've assumed
valid (https://github.com/syscoin/syscoin/pull/9484)

We build a chain that includes and invalid signature for one of the
transactions:

    0:        genesis block
    1:        block 1 with coinbase transaction output.
    2-101:    bury that block with 100 blocks so the coinbase transaction
              output can be spent
    102:      a block containing a transaction spending the coinbase
              transaction output. The transaction has an invalid signature.
    103-2202: bury the bad block with just over two weeks' worth of blocks
              (2100 blocks)

Start three nodes:

    - node0 has no -assumevalid parameter. Try to sync to block 2202. It will
      reject block 102 and only sync as far as block 101
    - node1 has -assumevalid set to the hash of block 102. Try to sync to
      block 2202. node1 will sync all the way to block 2202.
    - node2 has -assumevalid set to the hash of block 102. Try to sync to
      block 200. node2 will reject block 102 since it's assumed valid, but it
      isn't buried by at least two weeks' work.
"""
import time

from test_framework.blocktools import (create_block, create_coinbase)
from test_framework.key import ECKey
from test_framework.messages import (
    CBlockHeader,
    COutPoint,
    CTransaction,
    CTxIn,
    CTxOut,
    msg_block,
    msg_headers,
)
from test_framework.mininode import P2PInterface
from test_framework.script import (CScript, OP_TRUE)
from test_framework.test_framework import SyscoinTestFramework
from test_framework.util import assert_equal

class BaseNode(P2PInterface):
    def send_header_for_blocks(self, new_blocks):
        headers_message = msg_headers()
        headers_message.headers = [CBlockHeader(b) for b in new_blocks]
        self.send_message(headers_message)

class AssumeValidTest(SyscoinTestFramework):
    def set_test_params(self):
        self.setup_clean_chain = True
        self.num_nodes = 3
        self.rpc_timeout = 120
        # Must set '-dip3params=9000:9000' to create pre-dip3 blocks only
        self.extra_args = [['-dip3params=9000:9000'],['-dip3params=9000:9000'],['-dip3params=9000:9000']]

    def setup_network(self):
        self.add_nodes(3, extra_args=self.extra_args)
        # Start node0. We don't start the other nodes yet since
        # we need to pre-mine a block with an invalid transaction
        # signature so we can pass in the block hash as assumevalid.
        self.start_node(0)

    def send_blocks_until_disconnected(self, p2p_conn):
        """Keep sending blocks to the node until we're disconnected."""
        for i in range(len(self.blocks)):
            if not p2p_conn.is_connected:
                break
            try:
                p2p_conn.send_message(msg_block(self.blocks[i]))
            except IOError:
                assert not p2p_conn.is_connected
                break

    def assert_blockchain_height(self, node, height):
        """Wait until the blockchain is no longer advancing and verify it's reached the expected height."""
        last_height = node.getblock(node.getbestblockhash())['height']
        timeout = 10
        while True:
            time.sleep(0.25)
            current_height = node.getblock(node.getbestblockhash())['height']
            if current_height != last_height:
                last_height = current_height
                if timeout < 0:
                    assert False, "blockchain too short after timeout: %d" % current_height
                timeout - 0.25
                continue
            elif current_height > height:
                assert False, "blockchain too long: %d" % current_height
            elif current_height == height:
                break

    def run_test(self):
        p2p0 = self.nodes[0].add_p2p_connection(BaseNode())

        # Build the blockchain
        self.tip = int(self.nodes[0].getbestblockhash(), 16)
        self.block_time = self.nodes[0].getblock(self.nodes[0].getbestblockhash())['time'] + 1

        self.blocks = []

        # Get a pubkey for the coinbase TXO
        coinbase_key = ECKey()
        coinbase_key.generate()
        coinbase_pubkey = coinbase_key.get_pubkey().get_bytes()

        # Create the first block with a coinbase output to our key
        height = 1
        block = create_block(self.tip, create_coinbase(height, coinbase_pubkey), self.block_time)
        self.blocks.append(block)
        self.block_time += 1
        block.solve()
        # Save the coinbase for later
        self.block1 = block
        self.tip = block.sha256
        height += 1

        # Bury the block 100 deep so the coinbase output is spendable
        for i in range(100):
            block = create_block(self.tip, create_coinbase(height), self.block_time)
            block.solve()
            self.blocks.append(block)
            self.tip = block.sha256
            self.block_time += 1
            height += 1

        # Create a transaction spending the coinbase output with an invalid (null) signature
        tx = CTransaction()
        tx.vin.append(CTxIn(COutPoint(self.block1.vtx[0].sha256, 0), scriptSig=b""))
        tx.vout.append(CTxOut(49 * 100000000, CScript([OP_TRUE])))
        tx.calc_sha256()

        block102 = create_block(self.tip, create_coinbase(height), self.block_time)
        self.block_time += 1
        block102.vtx.extend([tx])
        block102.hashMerkleRoot = block102.calc_merkle_root()
        block102.rehash()
        block102.solve()
        self.blocks.append(block102)
        self.tip = block102.sha256
        self.block_time += 1
        height += 1

        # Bury the assumed valid block 2100 deep
        for i in range(2100):
            block = create_block(self.tip, create_coinbase(height), self.block_time)
            block.set_base_version(4)
            block.solve()
            self.blocks.append(block)
            self.tip = block.sha256
            self.block_time += 1
            height += 1

        self.nodes[0].disconnect_p2ps()

        # SYSCOIN Start node1 and node2 with assumevalid so they accept a block with a bad signature.
        self.start_node(1, extra_args=self.extra_args + ["-assumevalid=" + hex(block102.sha256)])
        self.start_node(2, extra_args=self.extra_args + ["-assumevalid=" + hex(block102.sha256)])

        p2p0 = self.nodes[0].add_p2p_connection(BaseNode())
        p2p1 = self.nodes[1].add_p2p_connection(BaseNode())
        p2p2 = self.nodes[2].add_p2p_connection(BaseNode())

        # send header lists to all three nodes
        p2p0.send_header_for_blocks(self.blocks[0:2000])
        p2p0.send_header_for_blocks(self.blocks[2000:])
        p2p1.send_header_for_blocks(self.blocks[0:2000])
        p2p1.send_header_for_blocks(self.blocks[2000:])
        p2p2.send_header_for_blocks(self.blocks[0:200])

        # Send blocks to node0. Block 102 will be rejected.
        self.send_blocks_until_disconnected(p2p0)
        self.assert_blockchain_height(self.nodes[0], 101)

        # Send all blocks to node1. All blocks will be accepted.
        for i in range(2202):
            p2p1.send_message(msg_block(self.blocks[i]))
        # Syncing 2200 blocks can take a while on slow systems. Give it plenty of time to sync.
        p2p1.sync_with_ping(960)
        assert_equal(self.nodes[1].getblock(self.nodes[1].getbestblockhash())['height'], 2202)

        # Send blocks to node2. Block 102 will be rejected.
        self.send_blocks_until_disconnected(p2p2)
        self.assert_blockchain_height(self.nodes[2], 101)

if __name__ == '__main__':
    AssumeValidTest().main()
<commit_msg>fix test<commit_after>#!/usr/bin/env python3
# Copyright (c) 2014-2020 The Bitcoin Core developers
# Distributed under the MIT software license, see the accompanying
# file COPYING or http://www.opensource.org/licenses/mit-license.php.
"""Test logic for skipping signature validation on old blocks.

Test logic for skipping signature validation on blocks which we've assumed
valid (https://github.com/syscoin/syscoin/pull/9484)

We build a chain that includes and invalid signature for one of the
transactions:

    0:        genesis block
    1:        block 1 with coinbase transaction output.
    2-101:    bury that block with 100 blocks so the coinbase transaction
              output can be spent
    102:      a block containing a transaction spending the coinbase
              transaction output. The transaction has an invalid signature.
    103-2202: bury the bad block with just over two weeks' worth of blocks
              (2100 blocks)

Start three nodes:

    - node0 has no -assumevalid parameter. Try to sync to block 2202. It will
      reject block 102 and only sync as far as block 101
    - node1 has -assumevalid set to the hash of block 102. Try to sync to
      block 2202. node1 will sync all the way to block 2202.
    - node2 has -assumevalid set to the hash of block 102. Try to sync to
      block 200. node2 will reject block 102 since it's assumed valid, but it
      isn't buried by at least two weeks' work.
"""
import time

from test_framework.blocktools import (create_block, create_coinbase)
from test_framework.key import ECKey
from test_framework.messages import (
    CBlockHeader,
    COutPoint,
    CTransaction,
    CTxIn,
    CTxOut,
    msg_block,
    msg_headers,
)
from test_framework.mininode import P2PInterface
from test_framework.script import (CScript, OP_TRUE)
from test_framework.test_framework import SyscoinTestFramework
from test_framework.util import assert_equal

class BaseNode(P2PInterface):
    def send_header_for_blocks(self, new_blocks):
        headers_message = msg_headers()
        headers_message.headers = [CBlockHeader(b) for b in new_blocks]
        self.send_message(headers_message)

class AssumeValidTest(SyscoinTestFramework):
    def set_test_params(self):
        self.setup_clean_chain = True
        self.num_nodes = 3
        self.rpc_timeout = 120
        # Must set '-dip3params=9000:9000' to create pre-dip3 blocks only
        self.extra_args = ['-dip3params=9000:9000']

    def setup_network(self):
        self.add_nodes(3)
        # Start node0. We don't start the other nodes yet since
        # we need to pre-mine a block with an invalid transaction
        # signature so we can pass in the block hash as assumevalid.
        self.start_node(0, extra_args=self.extra_args)

    def send_blocks_until_disconnected(self, p2p_conn):
        """Keep sending blocks to the node until we're disconnected."""
        for i in range(len(self.blocks)):
            if not p2p_conn.is_connected:
                break
            try:
                p2p_conn.send_message(msg_block(self.blocks[i]))
            except IOError:
                assert not p2p_conn.is_connected
                break

    def assert_blockchain_height(self, node, height):
        """Wait until the blockchain is no longer advancing and verify it's reached the expected height."""
        last_height = node.getblock(node.getbestblockhash())['height']
        timeout = 10
        while True:
            time.sleep(0.25)
            current_height = node.getblock(node.getbestblockhash())['height']
            if current_height != last_height:
                last_height = current_height
                if timeout < 0:
                    assert False, "blockchain too short after timeout: %d" % current_height
                timeout - 0.25
                continue
            elif current_height > height:
                assert False, "blockchain too long: %d" % current_height
            elif current_height == height:
                break

    def run_test(self):
        p2p0 = self.nodes[0].add_p2p_connection(BaseNode())

        # Build the blockchain
        self.tip = int(self.nodes[0].getbestblockhash(), 16)
        self.block_time = self.nodes[0].getblock(self.nodes[0].getbestblockhash())['time'] + 1

        self.blocks = []

        # Get a pubkey for the coinbase TXO
        coinbase_key = ECKey()
        coinbase_key.generate()
        coinbase_pubkey = coinbase_key.get_pubkey().get_bytes()

        # Create the first block with a coinbase output to our key
        height = 1
        block = create_block(self.tip, create_coinbase(height, coinbase_pubkey), self.block_time)
        self.blocks.append(block)
        self.block_time += 1
        block.solve()
        # Save the coinbase for later
        self.block1 = block
        self.tip = block.sha256
        height += 1

        # Bury the block 100 deep so the coinbase output is spendable
        for i in range(100):
            block = create_block(self.tip, create_coinbase(height), self.block_time)
            block.solve()
            self.blocks.append(block)
            self.tip = block.sha256
            self.block_time += 1
            height += 1

        # Create a transaction spending the coinbase output with an invalid (null) signature
        tx = CTransaction()
        tx.vin.append(CTxIn(COutPoint(self.block1.vtx[0].sha256, 0), scriptSig=b""))
        tx.vout.append(CTxOut(49 * 100000000, CScript([OP_TRUE])))
        tx.calc_sha256()

        block102 = create_block(self.tip, create_coinbase(height), self.block_time)
        self.block_time += 1
        block102.vtx.extend([tx])
        block102.hashMerkleRoot = block102.calc_merkle_root()
        block102.rehash()
        block102.solve()
        self.blocks.append(block102)
        self.tip = block102.sha256
        self.block_time += 1
        height += 1

        # Bury the assumed valid block 2100 deep
        for i in range(2100):
            block = create_block(self.tip, create_coinbase(height), self.block_time)
            block.set_base_version(4)
            block.solve()
            self.blocks.append(block)
            self.tip = block.sha256
            self.block_time += 1
            height += 1

        self.nodes[0].disconnect_p2ps()

        # SYSCOIN Start node1 and node2 with assumevalid so they accept a block with a bad signature.
        self.start_node(1, extra_args=self.extra_args + ["-assumevalid=" + hex(block102.sha256)])
        self.start_node(2, extra_args=self.extra_args + ["-assumevalid=" + hex(block102.sha256)])

        p2p0 = self.nodes[0].add_p2p_connection(BaseNode())
        p2p1 = self.nodes[1].add_p2p_connection(BaseNode())
        p2p2 = self.nodes[2].add_p2p_connection(BaseNode())

        # send header lists to all three nodes
        p2p0.send_header_for_blocks(self.blocks[0:2000])
        p2p0.send_header_for_blocks(self.blocks[2000:])
        p2p1.send_header_for_blocks(self.blocks[0:2000])
        p2p1.send_header_for_blocks(self.blocks[2000:])
        p2p2.send_header_for_blocks(self.blocks[0:200])

        # Send blocks to node0. Block 102 will be rejected.
        self.send_blocks_until_disconnected(p2p0)
        self.assert_blockchain_height(self.nodes[0], 101)

        # Send all blocks to node1. All blocks will be accepted.
        for i in range(2202):
            p2p1.send_message(msg_block(self.blocks[i]))
        # Syncing 2200 blocks can take a while on slow systems. Give it plenty of time to sync.
        p2p1.sync_with_ping(960)
        assert_equal(self.nodes[1].getblock(self.nodes[1].getbestblockhash())['height'], 2202)

        # Send blocks to node2. Block 102 will be rejected.
        self.send_blocks_until_disconnected(p2p2)
        self.assert_blockchain_height(self.nodes[2], 101)

if __name__ == '__main__':
    AssumeValidTest().main()
<|endoftext|>


<commit_before>"""nostray_prototype URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/1.11/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  url(r'^$', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  url(r'^$', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.conf.urls import url, include
    2. Add a URL to urlpatterns:  url(r'^blog/', include('blog.urls'))
"""
from django.conf.urls import url, include
from django.contrib import admin
from django.views.generic import TemplateView
from django.views.static import serve

from register.views import LoginView, RegisteView, ActiveUserView, ForgetPwdView, ResetPwdView, ModifyPwdView
from index.views import index
from nostray_prototype.settings import MEDIA_ROOT

urlpatterns = [
    url(r'^admin/', admin.site.urls),
    url(r'^captcha/', include('captcha.urls')),
    url(r'^$', index, name='index'),
    url(r'^login/$', LoginView.as_view(), name='login'),
    url(r'^registe/$', RegisteView.as_view(), name='registe'),
    url(r'^community/', include('community.urls', namespace='community')),
    url(r'^adopt/', include('adopt.urls', namespace='adopt')),
    url(r'^active/(?P<active_code>.*)/$', ActiveUserView.as_view(), name='user_active'),
    url(r'^forget/$', ForgetPwdView.as_view(), name='forget_pwd'),
    url(r'^reset/(?P<active_code>.*)/$', ResetPwdView.as_view(), name='reset_pwd'),
    url(r'^modify_pwd/$', ModifyPwdView.as_view(), name='modify_pwd'),
    url(r'^media/(?P<path>.*)$', serve, {"document_root":MEDIA_ROOT}),
]
<commit_msg>setup center url<commit_after>"""nostray_prototype URL Configuration

The `urlpatterns` list routes URLs to views. For more information please see:
    https://docs.djangoproject.com/en/1.11/topics/http/urls/
Examples:
Function views
    1. Add an import:  from my_app import views
    2. Add a URL to urlpatterns:  url(r'^$', views.home, name='home')
Class-based views
    1. Add an import:  from other_app.views import Home
    2. Add a URL to urlpatterns:  url(r'^$', Home.as_view(), name='home')
Including another URLconf
    1. Import the include() function: from django.conf.urls import url, include
    2. Add a URL to urlpatterns:  url(r'^blog/', include('blog.urls'))
"""
from django.conf.urls import url, include
from django.contrib import admin
from django.views.generic import TemplateView
from django.views.static import serve

from register.views import LoginView, RegisteView, ActiveUserView, ForgetPwdView, ResetPwdView, ModifyPwdView
from index.views import index
from nostray_prototype.settings import MEDIA_ROOT

urlpatterns = [
    url(r'^admin/', admin.site.urls),
    url(r'^captcha/', include('captcha.urls')),
    url(r'^$', index, name='index'),
    url(r'^login/$', LoginView.as_view(), name='login'),
    url(r'^registe/$', RegisteView.as_view(), name='registe'),
    url(r'^community/', include('community.urls', namespace='community')),
    url(r'^adopt/', include('adopt.urls', namespace='adopt')),
    url(r'^center/', include('center.urls', namespace='center')),
    url(r'^active/(?P<active_code>.*)/$', ActiveUserView.as_view(), name='user_active'),
    url(r'^forget/$', ForgetPwdView.as_view(), name='forget_pwd'),
    url(r'^reset/(?P<active_code>.*)/$', ResetPwdView.as_view(), name='reset_pwd'),
    url(r'^modify_pwd/$', ModifyPwdView.as_view(), name='modify_pwd'),
    url(r'^media/(?P<path>.*)$', serve, {"document_root":MEDIA_ROOT}),
]
<|endoftext|>


<commit_before># from utils import encode_predictions
from params import args
from datetime import datetime
import time
import numpy as np
import os
import threading
import queue
import pandas as pd
from scipy import ndimage

def rle_encode(mask_image):
    pixels = mask_image.flatten()
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2
    runs[1::2] = runs[1::2] - runs[:-1:2]
    return runs

def rle_to_string(runs):
    return ' '.join(str(x) for x in runs)

def encoder(in_queue, threshold, generated_masks, time_counts):
    while True:
        img_name, mask_img_path = in_queue.get()

        if img_name is None:
            break

        t0 = time.clock()
        mask_img = ndimage.imread(mask_img_path, mode='L')
        mask_img[mask_img <= threshold] = 0
        mask_img[mask_img > threshold] = 1
        time_counts['time_read'].append(time.clock() - t0)

        t0 = time.clock()
        rle = rle_encode(mask_img)
        time_counts['time_rle'].append(time.clock() - t0)

        t0 = time.clock()
        rle_string = rle_to_string(rle)
        time_counts['time_stringify'].append(time.clock() - t0)

        generated_masks.append((img_name, rle_string))

def encode_predictions(predicted_dir, filenames, output_path, n_threads=None, threshold=127):
    print('Predicting RLE encoding masks using {} threads...'.format(n_threads))

    num_masks = 0

    # @TODO: Lacks proper concurrency primitives, refactor
    time_counts = {
        'time_read': [],
        'time_rle': [],
        'time_stringify': []
    }

    filepaths_queue = queue.Queue()
    generated_masks = []
    threads = []

    for i in range(n_threads):
        threads.append(threading.Thread(target=encoder,
                                        name='Encoder',
                                        args=(filepaths_queue,
                                              threshold,
                                              generated_masks,
                                              time_counts)))

    for f in os.listdir(predicted_dir):
        num_masks += 1
        mask_img_path = os.path.join(predicted_dir, f)
        img_name = mask_img_path.split("/")[-1]

        filepaths_queue.put((img_name, mask_img_path))

    for thread in threads:
        filepaths_queue.put((None, None))

    print('{} masks found. Reading from {}, saving to {}'.format(num_masks, predicted_dir, output_path))

    for thread in threads:
        thread.start()

    for thread in threads:
        thread.join()

    time_read = sum(time_counts['time_read'])
    time_rle = sum(time_counts['time_rle'])
    time_stringify = sum(time_counts['time_stringify'])

    print('Time spent reading mask images:', time_read, 's, =>', 1000 * (time_read / num_masks), 'ms per mask.')
    print('Time spent RLE encoding masks:', time_rle, 's, =>', 1000 * (time_rle / num_masks), 'ms per mask.')
    print('Time spent stringifying RLEs:', time_stringify, 's, =>', 1000 * (time_stringify / num_masks), 'ms per mask.')

    output_df = pd.DataFrame(generated_masks, columns=['img', 'rle_mask'])
    output_df.to_csv(output_path, compression='gzip', index=False)

def main():
    input_df = pd.read_csv(args.pred_sample_csv)
    input_filenames = input_df['img']
    output_filename = 'submission-{:%Y-%m-%d_%H-%M-%S}.csv.gz'.format(datetime.now())
    output_path = os.path.join(args.submissions_dir, output_filename)
    encode_predictions(args.pred_mask_dir, input_filenames, output_path, n_threads=args.pred_threads)

if __name__ == '__main__':
    main()
<commit_msg>Fix rle submission generation (correctly specify filename)<commit_after># from utils import encode_predictions
from params import args
from datetime import datetime
import time
import numpy as np
import os
import threading
import queue
import pandas as pd
from scipy import ndimage

def rle_encode(mask_image):
    pixels = mask_image.flatten()
    runs = np.where(pixels[1:] != pixels[:-1])[0] + 2
    runs[1::2] = runs[1::2] - runs[:-1:2]
    return runs

def rle_to_string(runs):
    return ' '.join(str(x) for x in runs)

def encoder(in_queue, threshold, generated_masks, time_counts):
    while True:
        img_name, mask_img_path = in_queue.get()

        if img_name is None:
            break

        t0 = time.clock()
        mask_img = ndimage.imread(mask_img_path, mode='L')
        mask_img[mask_img <= threshold] = 0
        mask_img[mask_img > threshold] = 1
        time_counts['time_read'].append(time.clock() - t0)

        t0 = time.clock()
        rle = rle_encode(mask_img)
        time_counts['time_rle'].append(time.clock() - t0)

        t0 = time.clock()
        rle_string = rle_to_string(rle)
        time_counts['time_stringify'].append(time.clock() - t0)

        generated_masks.append((img_name, rle_string))

def encode_predictions(predicted_dir, filenames, n_threads=None, threshold=127):
    print('Predicting RLE encoding masks using {} threads...'.format(n_threads))

    num_masks = 0

    # @TODO: Lacks proper concurrency primitives, refactor
    time_counts = {
        'time_read': [],
        'time_rle': [],
        'time_stringify': []
    }

    filepaths_queue = queue.Queue()
    generated_masks = []
    threads = []

    for i in range(n_threads):
        threads.append(threading.Thread(target=encoder,
                                        name='Encoder',
                                        args=(filepaths_queue,
                                              threshold,
                                              generated_masks,
                                              time_counts)))

    for img_name in filenames:
        mask_img_path = os.path.join(predicted_dir, img_name)

        filepaths_queue.put((img_name, mask_img_path))

    for thread in threads:
        filepaths_queue.put((None, None))

    for thread in threads:
        thread.start()

    for thread in threads:
        thread.join()

    time_read = sum(time_counts['time_read'])
    time_rle = sum(time_counts['time_rle'])
    time_stringify = sum(time_counts['time_stringify'])

    num_masks = len(filenames)
    print('Time spent reading mask images:', time_read, 's, =>', 1000 * (time_read / num_masks), 'ms per mask.')
    print('Time spent RLE encoding masks:', time_rle, 's, =>', 1000 * (time_rle / num_masks), 'ms per mask.')
    print('Time spent stringifying RLEs:', time_stringify, 's, =>', 1000 * (time_stringify / num_masks), 'ms per mask.')

    return generated_masks

def main():
    input_df = pd.read_csv(args.pred_sample_csv)
    input_filenames = input_df['img'].str.replace('.jpg', '.png')

    output_filename = 'submission-{:%Y-%m-%d_%H-%M-%S}.csv.gz'.format(datetime.now())
    output_path = os.path.join(args.submissions_dir, output_filename)
    print('{} masks found. Reading from {}, saving to {}'.format(len(input_filenames), args.pred_mask_dir, output_path))
    generated_masks = encode_predictions(args.pred_mask_dir, input_filenames, n_threads=args.pred_threads)

    output_df = pd.DataFrame(generated_masks, columns=['img', 'rle_mask'])
    output_df['img'] = output_df.img.str.replace('.png', '.jpg')
    output_df.to_csv(output_path, compression='gzip', index=False)

if __name__ == '__main__':
    main()
<|endoftext|>